cellarr 0.4.1.dev1__tar.gz → 0.4.1.dev2__tar.gz

{cellarr-0.4.1.dev1 → cellarr-0.4.1.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cellarr
-Version: 0.4.1.dev1
+Version: 0.4.1.dev2
 Summary: TileDB-based array storage for genomics data collections.
 Home-page: https://github.com/BiocPy/cellarr
 Author: Jayaram Kancherla
@@ -45,7 +45,7 @@ datasets but can be generalized to store any 2-dimensional experimental data.
 
 To get started, install the package from [PyPI](https://pypi.org/project/cellarr/)
 
-```bash
+```sh
 pip install cellarr
 
 ## to include optional dependencies
@@ -156,6 +156,47 @@ print(dataset)
 
 Check out the [documentation](https://biocpy.github.io/cellarr/tutorial.html) for more details.
 
+### Building on HPC environments with `slurm`
+
+To simplify building TileDB files on HPC environments that use `slurm`, there are a few steps you need to follow.
+
+- Step 1: Construct a manifest file
+A minimal manifest file (json) must contain the following fields
+- `"files"`: A list of file path to the input `h5ad` objects.
+- `"python_env"`: A set of commands to activate the Python environment containing this package and its dependencies.
+
+Here’s an example of the manifest file:
+```json
+manifest = {
+    "files": your/list/of/files,
+    "python_env": """
+ml Miniforge3
+conda activate cellarr
+
+python --version
+which python
+    """,
+    "matrix_options": [
+        {
+            "matrix_name": "non_zero_cells",
+            "dtype": "uint32"
+        },
+        {
+            "matrix_name": "pseudo_bulk_log_normed",
+            "dtype": "float32"
+        }
+    ],
+}
+```
+For more options, check out the [README](./src/cellarr/slurm/README.md).
+
+- Step 2: Submit the job
+Once your manifest file is ready, you can submit the necessary jobs using the `cellarr_build` CLI. Run the following command:
+
+```sh
+cellarr_build --input-manifest your/path/to/test_manifest.json --output-dir your/path/to/output --memory-per-job 8 --cpus-per-task 2
+```
+
 ### Query a `CellArrDataset`
 
 Users have the option to reuse the `dataset` object returned when building the dataset or by creating a `CellArrDataset` object by initializing it to the path where the files were created.

{cellarr-0.4.1.dev1 → cellarr-0.4.1.dev2}/README.md

@@ -12,7 +12,7 @@ datasets but can be generalized to store any 2-dimensional experimental data.
 
 To get started, install the package from [PyPI](https://pypi.org/project/cellarr/)
 
-```bash
+```sh
 pip install cellarr
 
 ## to include optional dependencies
@@ -123,6 +123,47 @@ print(dataset)
 
 Check out the [documentation](https://biocpy.github.io/cellarr/tutorial.html) for more details.
 
+### Building on HPC environments with `slurm`
+
+To simplify building TileDB files on HPC environments that use `slurm`, there are a few steps you need to follow.
+
+- Step 1: Construct a manifest file
+A minimal manifest file (json) must contain the following fields
+- `"files"`: A list of file path to the input `h5ad` objects.
+- `"python_env"`: A set of commands to activate the Python environment containing this package and its dependencies.
+
+Here’s an example of the manifest file:
+```json
+manifest = {
+    "files": your/list/of/files,
+    "python_env": """
+ml Miniforge3
+conda activate cellarr
+
+python --version
+which python
+    """,
+    "matrix_options": [
+        {
+            "matrix_name": "non_zero_cells",
+            "dtype": "uint32"
+        },
+        {
+            "matrix_name": "pseudo_bulk_log_normed",
+            "dtype": "float32"
+        }
+    ],
+}
+```
+For more options, check out the [README](./src/cellarr/slurm/README.md).
+
+- Step 2: Submit the job
+Once your manifest file is ready, you can submit the necessary jobs using the `cellarr_build` CLI. Run the following command:
+
+```sh
+cellarr_build --input-manifest your/path/to/test_manifest.json --output-dir your/path/to/output --memory-per-job 8 --cpus-per-task 2
+```
+
 ### Query a `CellArrDataset`
 
 Users have the option to reuse the `dataset` object returned when building the dataset or by creating a `CellArrDataset` object by initializing it to the path where the files were created.

{cellarr-0.4.1.dev1 → cellarr-0.4.1.dev2}/setup.cfg

@@ -49,6 +49,8 @@ testing =
 	pytorch-lightning
 
 [options.entry_points]
+console_scripts = 
+	cellarr_build = cellarr.slurm.build_cellarr_steps:main
 
 [tool:pytest]
 addopts = 

{cellarr-0.4.1.dev1 → cellarr-0.4.1.dev2}/src/cellarr/buildutils_tiledb_frame.py

@@ -76,7 +76,7 @@ def create_tiledb_frame_from_column_names(
     )
 
 
-def create_tiledb_frame_from_dataframe(tiledb_uri_path: str, frame: List[str], column_types=dict):
+def create_tiledb_frame_from_dataframe(tiledb_uri_path: str, frame: List[str], column_types: dict = None):
     """Create a TileDB file with the provided attributes to persistent storage.
 
     This will materialize the array directory and all

{cellarr-0.4.1.dev1 → cellarr-0.4.1.dev2}/src/cellarr/slurm/README.md

@@ -37,10 +37,17 @@
             "cell_type": "ascii",
             "quality_score": "float32"
         },
-        "metadata": {
-            "processing_date": ["2024-01-01", "2024-01-02", ...]
-        }
-    }
+    },
+    "python_env": """
+. /system/gredit/clientos/etc/profile
+
+ml Miniforge3
+conda activate biocpy_miniforge
+
+~/.conda/envs/biocpy_miniforge/bin/python --version
+which python
+python --version
+    """,
 }
 ```
 

{cellarr-0.4.1.dev1 → cellarr-0.4.1.dev2}/src/cellarr/slurm/build_cellarr_steps.py

@@ -4,7 +4,6 @@ import subprocess
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple
 
-import anndata
 import numpy as np
 
 from cellarr.buildutils_tiledb_array import create_tiledb_array
@@ -55,7 +54,12 @@ class SlurmBuilder:
         self.cpus_per_task = cpus_per_task
 
     def create_slurm_script(
-        self, job_name: str, python_script: str, args: Dict, dependencies: Optional[str] = None
+        self,
+        job_name: str,
+        python_script: str,
+        args: Dict,
+        dependencies: Optional[str] = None,
+        python_env: str = "",
     ) -> str:
         """Create a SLURM job submission script."""
         script = f"""#!/bin/bash
@@ -70,8 +74,7 @@ class SlurmBuilder:
             script += f"#SBATCH --dependency={dependencies}\n"
 
         script += f"""
-source ~/.bashrc
-conda activate cellarr_env  # MODIFY AS NEEDED
+{python_env}
 
 python {python_script} '{json.dumps(args)}'
 """
@@ -81,7 +84,13 @@ python {python_script} '{json.dumps(args)}'
         return script_path
 
     def create_array_script(
-        self, job_name: str, python_script: str, args: Dict, n_tasks: int, dependencies: Optional[str] = None
+        self,
+        job_name: str,
+        python_script: str,
+        args: Dict,
+        n_tasks: int,
+        dependencies: Optional[str] = None,
+        python_env: str = "",
     ) -> str:
         """Create a SLURM array job submission script."""
         script = f"""#!/bin/bash
@@ -97,8 +106,7 @@ python {python_script} '{json.dumps(args)}'
             script += f"#SBATCH --dependency={dependencies}\n"
 
         script += f"""
-source ~/.bashrc
-conda activate cellarr_env # MODIFY AS NEEDED
+{python_env}
 
 python {python_script} '{json.dumps(args)}'
 """
@@ -112,7 +120,7 @@ python {python_script} '{json.dumps(args)}'
         result = subprocess.run(["sbatch", script_path], capture_output=True, text=True)
         return result.stdout.strip().split()[-1]
 
-    def submit_gene_annotation_job(self, files: List[str], gene_options: Dict) -> str:
+    def submit_gene_annotation_job(self, files: List[str], gene_options: Dict, python_env: str) -> str:
         """Submit gene annotation processing job."""
         args = {
             "files": files,
@@ -121,12 +129,19 @@ python {python_script} '{json.dumps(args)}'
             "temp_dir": str(self.temp_dir / "gene_annotation"),
         }
 
+        parent_dir = str(Path(__file__).parent)
+
         script_path = self.create_slurm_script(
-            job_name="cellarr_gene_annot", python_script="process_gene_annotation.py", args=args
+            job_name="cellarr_gene_annot",
+            python_script=f"{parent_dir}/process_gene_annotation.py",
+            args=args,
+            python_env=python_env,
         )
         return self.submit_job(script_path)
 
-    def submit_sample_metadata_job(self, files: List[str], sample_options: Dict, dependency: str) -> str:
+    def submit_sample_metadata_job(
+        self, files: List[str], sample_options: Dict, dependency: str, python_env: str
+    ) -> str:
         """Submit sample metadata processing job."""
         args = {
             "files": files,
@@ -135,15 +150,18 @@ python {python_script} '{json.dumps(args)}'
             "temp_dir": str(self.temp_dir / "sample_metadata"),
         }
 
+        parent_dir = str(Path(__file__).parent)
+
         script_path = self.create_slurm_script(
             job_name="cellarr_sample_meta",
-            python_script="process_sample_metadata.py",
+            python_script=f"{parent_dir}/process_sample_metadata.py",
             args=args,
+            python_env=python_env,
             # dependencies=f"afterok:{dependency}",
         )
         return self.submit_job(script_path)
 
-    def submit_cell_metadata_job(self, files: List[str], cell_options: Dict, dependency: str) -> str:
+    def submit_cell_metadata_job(self, files: List[str], cell_options: Dict, dependency: str, python_env: str) -> str:
         """Submit cell metadata processing job."""
         args = {
             "files": files,
@@ -152,23 +170,21 @@ python {python_script} '{json.dumps(args)}'
             "temp_dir": str(self.temp_dir / "cell_metadata"),
         }
 
+        parent_dir = str(Path(__file__).parent)
+
         script_path = self.create_slurm_script(
             job_name="cellarr_cell_meta",
-            python_script="process_cell_metadata.py",
+            python_script=f"{parent_dir}/process_cell_metadata.py",
             args=args,
+            python_env=python_env,
             # dependencies=f"afterok:{dependency}",
         )
         return self.submit_job(script_path)
 
-    def submit_matrix_processing(self, files: List[str], matrix_options: Dict, dependency: str) -> Tuple[str, str]:
+    def submit_matrix_processing(
+        self, files: List[str], matrix_options: Dict, dependency: str, python_env: str
+    ) -> Tuple[str, str]:
         """Submit matrix processing as SLURM array job."""
-        # Calculate cell offsets for each file
-        offset = 0
-        file_infos = []
-        for file in files:
-            file_infos.append({"file": file, "offset": offset})
-            adata = anndata.read_h5ad(file, backed="r")
-            offset += adata.n_obs
 
         # Create matrix TileDB array
         matrix_uri = str(self.output_dir / "assays" / matrix_options["matrix_name"])
@@ -184,16 +200,19 @@ python {python_script} '{json.dumps(args)}'
             "temp_dir": str(self.temp_dir / f"matrix_{matrix_options['matrix_name']}"),
             "matrix_options": matrix_options,
             "gene_annotation_file": str(self.temp_dir / "gene_annotation/gene_set.json"),
-            "file_infos": file_infos,
+            "files": files,
         }
 
+        parent_dir = str(Path(__file__).parent)
+
         # Submit array job
         array_script = self.create_array_script(
             job_name=f"matrix_{matrix_options['matrix_name']}",
-            python_script="process_matrix.py",
+            python_script=f"{parent_dir}/process_matrix.py",
             args=array_args,
             n_tasks=len(files),
             dependencies=f"afterok:{dependency}",
+            python_env=python_env,
         )
         array_job_id = self.submit_job(array_script)
 
@@ -202,28 +221,32 @@ python {python_script} '{json.dumps(args)}'
             "output_dir": str(self.output_dir),
             "temp_dir": str(self.temp_dir / f"matrix_{matrix_options['matrix_name']}"),
             "matrix_options": matrix_options,
-            "file_infos": file_infos,
+            "files": files,
         }
 
         final_script = self.create_slurm_script(
             job_name=f"finalize_matrix_{matrix_options['matrix_name']}",
-            python_script="finalize_matrix.py",
+            python_script=f"{parent_dir}/finalize_matrix.py",
             args=final_args,
             dependencies=f"afterok:{array_job_id}",
+            python_env=python_env,
         )
         final_job_id = self.submit_job(final_script)
 
         return array_job_id, final_job_id
 
-    def submit_final_assembly(self, matrix_names: List[str], dependencies: List[str]) -> str:
+    def submit_final_assembly(self, matrix_names: List[str], dependencies: List[str], python_env: str) -> str:
         """Submit final assembly job."""
         args = {"input_dir": str(self.output_dir), "output_dir": str(self.output_dir), "matrix_names": matrix_names}
 
+        parent_dir = str(Path(__file__).parent)
+
         script_path = self.create_slurm_script(
             job_name="cellarr_final_assembly",
-            python_script="final_assembly.py",
+            python_script=f"{parent_dir}/final_assembly.py",
             args=args,
             dependencies=f"afterok:{','.join(dependencies)}",
+            python_env=python_env,
         )
         return self.submit_job(script_path)
 
@@ -242,8 +265,9 @@ def main():
     log_dir = base_dir / "logs"
     temp_dir = base_dir / "temp"
     final_dir = base_dir / "final"
+    assays_dir = base_dir / "final/assays"
 
-    for d in [log_dir, temp_dir, final_dir]:
+    for d in [log_dir, temp_dir, final_dir, assays_dir]:
         d.mkdir(parents=True, exist_ok=True)
 
     # Read manifest
@@ -261,13 +285,17 @@ def main():
     )
 
     # Submit jobs
-    gene_job_id = builder.submit_gene_annotation_job(manifest["files"], manifest.get("gene_options", {}))
+    gene_job_id = builder.submit_gene_annotation_job(
+        manifest["files"], manifest.get("gene_options", {}), manifest["python_env"]
+    )
 
     sample_job_id = builder.submit_sample_metadata_job(
-        manifest["files"], manifest.get("sample_options", {}), gene_job_id
+        manifest["files"], manifest.get("sample_options", {}), gene_job_id, manifest["python_env"]
     )
 
-    cell_job_id = builder.submit_cell_metadata_job(manifest["files"], manifest.get("cell_options", {}), sample_job_id)
+    cell_job_id = builder.submit_cell_metadata_job(
+        manifest["files"], manifest.get("cell_options", {}), sample_job_id, manifest["python_env"]
+    )
 
     # Process matrices
     matrix_options = manifest.get("matrix_options", [{"matrix_name": "counts"}])
@@ -276,11 +304,15 @@ def main():
 
     matrix_job_ids = []
     for matrix_opt in matrix_options:
-        _, final_id = builder.submit_matrix_processing(manifest["files"], matrix_opt, f"{cell_job_id},{gene_job_id}")
+        _, final_id = builder.submit_matrix_processing(
+            manifest["files"], matrix_opt, f"{cell_job_id},{gene_job_id}", manifest["python_env"]
+        )
         matrix_job_ids.append(final_id)
 
     # Submit final assembly
-    builder.submit_final_assembly([opt["matrix_name"] for opt in matrix_options], matrix_job_ids)
+    builder.submit_final_assembly(
+        [opt["matrix_name"] for opt in matrix_options], matrix_job_ids, manifest["python_env"]
+    )
 
 
 if __name__ == "__main__":

{cellarr-0.4.1.dev1 → cellarr-0.4.1.dev2}/src/cellarr/slurm/final_assembly.py

@@ -1,4 +1,5 @@
 import json
+import sys
 
 from cellarr import CellArrDataset
 
@@ -18,3 +19,11 @@ def final_assembly(args_json: str):
 
     with open(f"{args['output_dir']}/metadata.json", "w") as f:
         json.dump(metadata, f)
+
+
+if __name__ == "__main__":
+    if len(sys.argv) != 2:
+        print("Usage: python final_assembly.py '<json_args>'")
+        sys.exit(1)
+
+    final_assembly(sys.argv[1])

{cellarr-0.4.1.dev1 → cellarr-0.4.1.dev2}/src/cellarr/slurm/finalize_matrix.py

@@ -15,7 +15,7 @@ def finalize_matrix(args_json: str):
 
     # Verify all tasks completed
     completed_dir = Path(args["temp_dir"]) / "completed"
-    expected_tasks = len(args["file_infos"])
+    expected_tasks = len(args["files"])
     completed_tasks = len(list(completed_dir.glob("task_*.json")))
 
     if completed_tasks != expected_tasks:
@@ -30,7 +30,7 @@ def finalize_matrix(args_json: str):
         json.dump(
             {
                 "matrix_name": args["matrix_options"]["matrix_name"],
-                "files_processed": len(args["file_infos"]),
+                "files_processed": len(args["files"]),
                 "uri": matrix_uri,
             },
             f,

{cellarr-0.4.1.dev1 → cellarr-0.4.1.dev2}/src/cellarr/slurm/process_gene_annotation.py

@@ -1,4 +1,6 @@
 import json
+import sys
+from pathlib import Path
 
 import pandas as pd
 
@@ -15,7 +17,7 @@ def process_gene_annotation(args_json: str):
 
     # Extract gene information from all files
     files_cache = uad.extract_anndata_info(
-        args["files"], var_feature_column=args["gene_options"].get("feature_column", "index")
+        args["files"], var_feature_column=args.get("gene_options", {}).get("feature_column", "index")
     )
 
     # Scan for features
@@ -24,10 +26,20 @@ def process_gene_annotation(args_json: str):
 
     # Create gene annotation dataframe
     gene_annotation = pd.DataFrame({"cellarr_gene_index": gene_set}, index=gene_set)
+    gene_annotation.reset_index(drop=True, inplace=True)
+
+    # Create TileDB store
+    create_tiledb_frame_from_dataframe(f"{args['output_dir']}/gene_annotation", gene_annotation)
 
     # Save gene set for later use
+    Path(args["temp_dir"]).mkdir(parents=True, exist_ok=True)
     with open(f"{args['temp_dir']}/gene_set.json", "w") as f:
         json.dump(gene_set, f)
 
-    # Create TileDB store
-    create_tiledb_frame_from_dataframe(f"{args['output_dir']}/gene_annotation", gene_annotation)
+
+if __name__ == "__main__":
+    if len(sys.argv) != 2:
+        print("Usage: python process_gene_annotation.py '<json_args>'")
+        sys.exit(1)
+
+    process_gene_annotation(sys.argv[1])

{cellarr-0.4.1.dev1 → cellarr-0.4.1.dev2}/src/cellarr/slurm/process_matrix.py

@@ -3,6 +3,8 @@ import os
 import sys
 from pathlib import Path
 
+import tiledb
+
 from cellarr import utils_anndata as uad
 from cellarr.buildutils_tiledb_array import write_csr_matrix_to_tiledb
 
@@ -24,9 +26,13 @@ def process_matrix_file(args_json: str):
     gene_map = {gene: idx for idx, gene in enumerate(gene_set)}
 
     # Get file and offset for this task
-    file_info = args["file_infos"][task_id]
-    input_file = file_info["file"]
-    row_offset = file_info["offset"]
+    file_info = args["files"][task_id]
+    input_file = file_info
+
+    # get sample offset
+    sample_uri = tiledb.open(f"{args['output_dir']}/sample_metadata", "r")
+    sample_row = sample_uri.df[task_id]
+    row_offset = sample_row["cellarr_sample_start_index"]
 
     # Process the file
     matrix = uad.remap_anndata(input_file, gene_map, layer_matrix_name=args["matrix_options"]["matrix_name"])
@@ -36,6 +42,8 @@ def process_matrix_file(args_json: str):
     write_csr_matrix_to_tiledb(matrix_uri, matrix[args["matrix_options"]["matrix_name"]], row_offset=row_offset)
 
     # Save task completion marker
+    Path(args["temp_dir"]).mkdir(parents=True, exist_ok=True)
+    Path(args["temp_dir"] + "/completed").mkdir(parents=True, exist_ok=True)
     task_marker = Path(args["temp_dir"]) / "completed" / f"task_{task_id}.json"
     task_marker.parent.mkdir(exist_ok=True)
     with open(task_marker, "w") as f:

{cellarr-0.4.1.dev1 → cellarr-0.4.1.dev2}/src/cellarr/slurm/process_sample_metadata.py

@@ -1,5 +1,6 @@
 import json
 import sys
+from pathlib import Path
 
 import pandas as pd
 
@@ -64,6 +65,7 @@ def process_sample_metadata(args_json: str):
     # Save metadata for subsequent steps
     metadata = {"num_samples": len(sample_names), "total_cells": int(sample_metadata["cellarr_cell_counts"].sum())}
 
+    Path(args["temp_dir"]).mkdir(parents=True, exist_ok=True)
     with open(f"{args['temp_dir']}/sample_metadata.json", "w") as f:
         json.dump(metadata, f)
 

{cellarr-0.4.1.dev1 → cellarr-0.4.1.dev2}/src/cellarr.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cellarr
-Version: 0.4.1.dev1
+Version: 0.4.1.dev2
 Summary: TileDB-based array storage for genomics data collections.
 Home-page: https://github.com/BiocPy/cellarr
 Author: Jayaram Kancherla
@@ -45,7 +45,7 @@ datasets but can be generalized to store any 2-dimensional experimental data.
 
 To get started, install the package from [PyPI](https://pypi.org/project/cellarr/)
 
-```bash
+```sh
 pip install cellarr
 
 ## to include optional dependencies
@@ -156,6 +156,47 @@ print(dataset)
 
 Check out the [documentation](https://biocpy.github.io/cellarr/tutorial.html) for more details.
 
+### Building on HPC environments with `slurm`
+
+To simplify building TileDB files on HPC environments that use `slurm`, there are a few steps you need to follow.
+
+- Step 1: Construct a manifest file
+A minimal manifest file (json) must contain the following fields
+- `"files"`: A list of file path to the input `h5ad` objects.
+- `"python_env"`: A set of commands to activate the Python environment containing this package and its dependencies.
+
+Here’s an example of the manifest file:
+```json
+manifest = {
+    "files": your/list/of/files,
+    "python_env": """
+ml Miniforge3
+conda activate cellarr
+
+python --version
+which python
+    """,
+    "matrix_options": [
+        {
+            "matrix_name": "non_zero_cells",
+            "dtype": "uint32"
+        },
+        {
+            "matrix_name": "pseudo_bulk_log_normed",
+            "dtype": "float32"
+        }
+    ],
+}
+```
+For more options, check out the [README](./src/cellarr/slurm/README.md).
+
+- Step 2: Submit the job
+Once your manifest file is ready, you can submit the necessary jobs using the `cellarr_build` CLI. Run the following command:
+
+```sh
+cellarr_build --input-manifest your/path/to/test_manifest.json --output-dir your/path/to/output --memory-per-job 8 --cpus-per-task 2
+```
+
 ### Query a `CellArrDataset`
 
 Users have the option to reuse the `dataset` object returned when building the dataset or by creating a `CellArrDataset` object by initializing it to the path where the files were created.

{cellarr-0.4.1.dev1 → cellarr-0.4.1.dev2}/src/cellarr.egg-info/SOURCES.txt

@@ -42,13 +42,13 @@ src/cellarr/utils_anndata.py
 src/cellarr.egg-info/PKG-INFO
 src/cellarr.egg-info/SOURCES.txt
 src/cellarr.egg-info/dependency_links.txt
+src/cellarr.egg-info/entry_points.txt
 src/cellarr.egg-info/not-zip-safe
 src/cellarr.egg-info/requires.txt
 src/cellarr.egg-info/top_level.txt
 src/cellarr/slurm/README.md
 src/cellarr/slurm/__init__.py
 src/cellarr/slurm/build_cellarr_steps.py
-src/cellarr/slurm/build_cellarr_steps_bkp.py
 src/cellarr/slurm/final_assembly.py
 src/cellarr/slurm/finalize_matrix.py
 src/cellarr/slurm/process_cell_metadata.py

cellarr-0.4.1.dev2/src/cellarr.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+cellarr_build = cellarr.slurm.build_cellarr_steps:main

cellarr-0.4.1.dev1/src/cellarr/slurm/build_cellarr_steps_bkp.py

@@ -1,185 +0,0 @@
-import argparse
-import json
-import os
-import subprocess
-from pathlib import Path
-from typing import Dict, Optional
-
-
-__author__ = "Jayaram Kancherla"
-__copyright__ = "Jayaram Kancherla"
-__license__ = "MIT"
-
-
-def create_slurm_script(
-    job_name: str,
-    log_dir: str,
-    python_script: str,
-    args: Dict,
-    memory_gb: int = 64,
-    time_hours: int = 24,
-    cpus_per_task: int = 4,
-    dependencies: Optional[str] = None,
-) -> str:
-    """Create SLURM job submission script."""
-    script = f"""#!/bin/bash
-#SBATCH --job-name={job_name}
-#SBATCH --output={log_dir}/{job_name}_%j.out
-#SBATCH --error={log_dir}/{job_name}_%j.err
-#SBATCH --time={time_hours}:00:00
-#SBATCH --mem={memory_gb}G
-#SBATCH --cpus-per-task={cpus_per_task}
-"""
-    if dependencies:
-        script += f"#SBATCH --dependency={dependencies}\n"
-
-    script += """
-source ~/.bashrc
-conda activate cellarr_env  # Modify as needed for your environment
-
-python {python_script} '{json.dumps(args)}'
-"""
-    script_path = os.path.join(log_dir, f"{job_name}_submit.sh")
-    with open(script_path, "w") as f:
-        f.write(script)
-    return script_path
-
-
-def submit_slurm_job(script_path: str, dependency: Optional[str] = None) -> str:
-    """Submit SLURM job and return job ID."""
-    cmd = ["sbatch"]
-    if dependency:
-        cmd.extend(["--dependency", dependency])
-    cmd.append(script_path)
-
-    result = subprocess.run(cmd, capture_output=True, text=True)
-    job_id = result.stdout.strip().split()[-1]
-    return job_id
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Build CellArrDataset using SLURM steps")
-    parser.add_argument("--input-manifest", required=True, help="Path to JSON manifest file")
-    parser.add_argument("--output-dir", required=True, help="Base output directory")
-    parser.add_argument("--memory-per-job", type=int, default=64, help="Memory in GB per job")
-    parser.add_argument("--time-per-job", type=int, default=24, help="Time in hours per job")
-    parser.add_argument("--cpus-per-task", type=int, default=4, help="CPUs per task")
-    args = parser.parse_args()
-
-    # Create output directories
-    base_dir = Path(args.output_dir)
-    log_dir = base_dir / "logs"
-    temp_dir = base_dir / "temp"
-    final_dir = base_dir / "final"
-
-    for d in [log_dir, temp_dir, final_dir]:
-        d.mkdir(parents=True, exist_ok=True)
-
-    # Read manifest file
-    with open(args.input_manifest) as f:
-        manifest = json.load(f)
-
-    # Step 1: Gene Annotation Job
-    gene_args = {
-        "files": manifest["files"],
-        "output_dir": str(final_dir),
-        "gene_options": manifest.get("gene_options", {}),
-        "temp_dir": str(temp_dir / "gene_annotation"),
-    }
-
-    gene_script = create_slurm_script(
-        job_name="cellarr_gene_annot",
-        log_dir=str(log_dir),
-        python_script="process_gene_annotation.py",
-        args=gene_args,
-        memory_gb=args.memory_per_job,
-        cpus_per_task=args.cpus_per_task,
-    )
-    gene_job_id = submit_slurm_job(gene_script)
-
-    # Step 2: Sample Metadata Job
-    sample_args = {
-        "files": manifest["files"],
-        "output_dir": str(final_dir),
-        "sample_options": manifest.get("sample_options", {}),
-        "temp_dir": str(temp_dir / "sample_metadata"),
-    }
-
-    sample_script = create_slurm_script(
-        job_name="cellarr_sample_meta",
-        log_dir=str(log_dir),
-        python_script="process_sample_metadata.py",
-        args=sample_args,
-        memory_gb=args.memory_per_job,
-        cpus_per_task=args.cpus_per_task,
-        # dependencies=f"afterok:{gene_job_id}",
-    )
-    sample_job_id = submit_slurm_job(sample_script)
-
-    # Step 3: Cell Metadata Job
-    cell_args = {
-        "files": manifest["files"],
-        "output_dir": str(final_dir),
-        "cell_options": manifest.get("cell_options", {}),
-        "temp_dir": str(temp_dir / "cell_metadata"),
-    }
-
-    cell_script = create_slurm_script(
-        job_name="cellarr_cell_meta",
-        log_dir=str(log_dir),
-        python_script="process_cell_metadata.py",
-        args=cell_args,
-        memory_gb=args.memory_per_job,
-        cpus_per_task=args.cpus_per_task,
-        # dependencies=f"afterok:{sample_job_id}",
-    )
-    cell_job_id = submit_slurm_job(cell_script)
-
-    # Step 4: Matrix Jobs (one per matrix type)
-    matrix_job_ids = []
-    matrix_options = manifest.get("matrix_options", [{"matrix_name": "counts"}])
-    if not isinstance(matrix_options, list):
-        matrix_options = [matrix_options]
-
-    for matrix_opt in matrix_options:
-        matrix_args = {
-            "files": manifest["files"],
-            "output_dir": str(final_dir),
-            "matrix_options": matrix_opt,
-            "gene_annotation_file": str(temp_dir / "gene_annotation/gene_set.json"),
-            "temp_dir": str(temp_dir / f"matrix_{matrix_opt['matrix_name']}"),
-        }
-
-        matrix_script = create_slurm_script(
-            job_name=f"cellarr_matrix_{matrix_opt['matrix_name']}",
-            log_dir=str(log_dir),
-            python_script="process_matrix.py",
-            args=matrix_args,
-            memory_gb=args.memory_per_job * 2,  # More memory for matrix processing
-            cpus_per_task=args.cpus_per_task,
-            dependencies=f"afterok:{cell_job_id},{gene_job_id}",
-        )
-        job_id = submit_slurm_job(matrix_script)
-        matrix_job_ids.append(job_id)
-
-    # Step 5: Final Assembly Job
-    final_args = {
-        "input_dir": str(final_dir),
-        "output_dir": str(final_dir),
-        "matrix_names": [opt["matrix_name"] for opt in matrix_options],
-    }
-
-    final_script = create_slurm_script(
-        job_name="cellarr_final_assembly",
-        log_dir=str(log_dir),
-        python_script="final_assembly.py",
-        args=final_args,
-        memory_gb=args.memory_per_job,
-        cpus_per_task=args.cpus_per_task,
-        dependencies=f"afterok:{','.join(matrix_job_ids)}",
-    )
-    submit_slurm_job(final_script)
-
-
-if __name__ == "__main__":
-    main()