adds a new sort utility to order h5ad files for faster cell load

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "arc-state"
-version = "0.10.2"
+version = "0.10.3"
 description = "State is a machine learning model that predicts cellular perturbation response across diverse contexts."
 readme = "README.md"
 authors = [
@@ -13,7 +13,7 @@ authors = [
 requires-python = ">=3.10,<3.13"
 dependencies = [
     "anndata>=0.11.4",
-    "cell-load>=0.8.7",
+    "cell-load>=0.9.0",
     "numpy>=2.2.6",
     "pandas>=2.2.3",
     "pyyaml>=6.0.2",

src/state/__main__.py

@@ -14,6 +14,7 @@
     run_tx_infer,
     run_tx_predict,
     run_tx_preprocess_train,
+    run_tx_sort,
     run_tx_train,
 )
 
@@ -126,6 +127,8 @@ def main():
                 case "preprocess_train":
                     # Run preprocessing using argparse
                     run_tx_preprocess_train(args.adata, args.output, args.num_hvgs)
+                case "sort":
+                    run_tx_sort(args)
 
 
 if __name__ == "__main__":

src/state/_cli/__init__.py

@@ -4,6 +4,7 @@
     run_tx_infer,
     run_tx_predict,
     run_tx_preprocess_train,
+    run_tx_sort,
     run_tx_train,
 )
 
@@ -14,6 +15,7 @@
     "run_tx_predict",
     "run_tx_infer",
     "run_tx_preprocess_train",
+    "run_tx_sort",
     "run_emb_fit",
     "run_emb_query",
     "run_emb_transform",

src/state/_cli/_tx/__init__.py

@@ -3,13 +3,15 @@
 from ._infer import add_arguments_infer, run_tx_infer
 from ._predict import add_arguments_predict, run_tx_predict
 from ._preprocess_train import add_arguments_preprocess_train, run_tx_preprocess_train
+from ._sort import add_arguments_sort, run_tx_sort
 from ._train import add_arguments_train, run_tx_train
 
 __all__ = [
     "run_tx_train",
     "run_tx_predict",
     "run_tx_infer",
     "run_tx_preprocess_train",
+    "run_tx_sort",
     "add_arguments_tx",
 ]
 
@@ -21,3 +23,4 @@ def add_arguments_tx(parser: ap.ArgumentParser):
     add_arguments_predict(subparsers.add_parser("predict"))
     add_arguments_infer(subparsers.add_parser("infer"))
     add_arguments_preprocess_train(subparsers.add_parser("preprocess_train"))
+    add_arguments_sort(subparsers.add_parser("sort"))

src/state/_cli/_tx/_sort.py

@@ -0,0 +1,72 @@
+import argparse as ap
+
+
+def add_arguments_sort(parser: ap.ArgumentParser):
+    """Add arguments for the sort subcommand."""
+    parser.add_argument(
+        "--input",
+        type=str,
+        required=True,
+        help="Path to input AnnData file (.h5ad)",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        required=True,
+        help="Path to output sorted AnnData file (.h5ad)",
+    )
+    parser.add_argument(
+        "--context-col",
+        type=str,
+        required=True,
+        help="obs column to sort by context (e.g. cell type)",
+    )
+    parser.add_argument(
+        "--batch-col",
+        type=str,
+        required=False,
+        default=None,
+        help="optional obs column to sort by batch (if omitted, sorts by context + perturbation)",
+    )
+    parser.add_argument(
+        "--pert-col",
+        type=str,
+        required=True,
+        help="obs column to sort by perturbation",
+    )
+
+
+def run_tx_sort(args: ap.Namespace):
+    import logging
+    from pathlib import Path
+
+    import anndata as ad
+
+    logging.basicConfig(level=logging.INFO)
+    logger = logging.getLogger(__name__)
+
+    input_path = args.input
+    output_path = args.output
+    sort_cols = [args.context_col]
+    if args.batch_col:
+        sort_cols.append(args.batch_col)
+    sort_cols.append(args.pert_col)
+
+    logger.info("Loading AnnData from %s", input_path)
+    adata = ad.read_h5ad(input_path)
+
+    missing = [col for col in sort_cols if col not in adata.obs.columns]
+    if missing:
+        raise ValueError(f"Missing obs columns for sorting: {missing}")
+
+    logger.info("Sorting AnnData by columns: %s", sort_cols)
+    order = adata.obs.sort_values(by=sort_cols, kind="mergesort").index
+    adata_sorted = adata[order].copy()
+
+    output_dir = Path(output_path).parent
+    if output_dir and not output_dir.exists():
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+    logger.info("Writing sorted AnnData to %s", output_path)
+    adata_sorted.write_h5ad(output_path)
+    logger.info("Sort complete. Wrote %d cells.", adata_sorted.n_obs)