diff --git a/.gitignore b/.gitignore index 48c565d..bcafdbc 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,13 @@ +data/ +local_outputs/ +work/ +*.log* +.nextflow* +*.html +*.h5ad +*.csv +*.ckpt + # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/README.md b/README.md index 4827c2e..e179ebe 100644 --- a/README.md +++ b/README.md @@ -1,21 +1,115 @@ -# Cellarium Workflows Example -This code contains helper functions and example to run scripts in Vertex AI platform (powered by Kubeflow) +# Nextflow pipelines for `cellarium-ml` + +Nextflow pipelines for running [cellarium-ml](https://github.com/cellarium-ai/cellarium-ml) workflows at scale on Google Cloud Batch. + +## The general idea + +`cellarium-ml` can scale algorithms to handle large scRNA-seq datasets. This repository helps orchestrate the compute, and composes `cellarium-ml` models, hooking up outputs of one model to inputs of the next. + +Here is an example. Run PCA on an arbitrarily large dataset (ideally created by Cellarium Nexus, but not necessarily) using google cloud hardware: + +```bash +nextflow run cellarium/workflows/nextflow/pca_workflow.nf \ + --h5ad_bucket 'gs://my-bucket/data/extract_files' \ + --output_bucket 'gs://my-bucket/outputs/pca' +``` + +Includes highly-variable gene selection and gene z-scoring. Takes maybe 9 hours on 100M cells. The output bucket will contain a trained PCA model checkpoint as well as the per-cell PCs in sharded CSV files. + +## Workflows + +Descriptions of current workflows can be found [here](cellarium/workflows/nextflow/README.md) + +--- + +## Installing Nextflow + +Nextflow requires Java 17 or later. + +### Conda install + +You can [install nextflow in a conda environment](https://docs.seqera.io/nextflow/install#conda). + +```bash +conda create -n nextflow bioconda::nextflow +conda activate nextflow +nextflow info +``` + +### System-wide install + +The [recommended install method](https://docs.seqera.io/nextflow/install) is via the standalone installer. + +If you do not yet have java 17 - 26 (`java -version`), then this is recommended: -## Quick start -* [Install gcloud CLI](https://cloud.google.com/sdk/docs/install) -* [Authenticate gcloud CLI util](https://cloud.google.com/docs/authentication/gcloud) -* Install project requirements like: ```bash -$ pip isntall -r requirements/base.txt +curl -s https://get.sdkman.io | bash +sdk install java 17.0.10-tem ``` -## Example -Go to example dir +Confirm the java installation using `java -version` and then do + ```bash -$ cd cellarium/workflows/example +curl -s https://get.nextflow.io | bash +chmod +x nextflow +mkdir -p $HOME/.local/bin/ +mv nextflow $HOME/.local/bin/ ``` -Submit an example pipeline: + +In either case, you can verify your installation with this command: ```bash -$ python submit_example_pipeline.py --project_id dsp-cell-annotation-service --location us-central1 --display_name test --component_1_config gs://test-bucket/test-config-1.yaml --component_2_config gs://test-bucket/test-config-2.yaml -``` \ No newline at end of file +nextflow info +``` + +You need at least **Nextflow 23.10** for the Google Batch executor used here. + +--- + +## nf-google plugin + +The `nf-google` plugin provides the Google Batch executor and GCS file staging — it is what lets Nextflow copy your `.h5ad` files from GCS onto each task VM without any `gcloud` CLI in your process container. + +**You do not need to install it manually.** `nextflow.config` declares it: + +```groovy +plugins { + id 'nf-google' +} +``` + +On first run Nextflow downloads it automatically from [plugins.nextflow.io](https://plugins.nextflow.io) and caches it in `~/.nextflow/plugins/`. Subsequent runs are offline. + +If you are in an air-gapped environment or want to pin a specific version: + +```groovy +plugins { + id 'nf-google@1.14.0' +} +``` + +--- + +## GCP prerequisites + +Before running on Google Cloud Batch you need: + +1. **APIs enabled** in your GCP project: + ```bash + gcloud services enable batch.googleapis.com storage.googleapis.com \ + artifactregistry.googleapis.com logging.googleapis.com + ``` + +2. **A service account** (or your user account via ADC) with these roles: + - `roles/batch.jobsEditor` + - `roles/batch.agentReporter` + - `roles/storage.objectAdmin` (on your work bucket and output bucket) + - `roles/logging.logWriter` + - `roles/artifactregistry.reader` (to pull the container image) + +3. **Application Default Credentials** active on the machine where you run `nextflow`: + ```bash + gcloud auth application-default login + ``` + +4. **A GCS work bucket** — Nextflow uses this for task staging and intermediate files. It must be in the same region as your Batch jobs (`us-central1` by default). Separate from your output bucket. diff --git a/cellarium/workflows/__init__.py b/cellarium/workflows/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/cellarium/workflows/configs/hvg_seurat_v3.yaml.j2 b/cellarium/workflows/configs/hvg_seurat_v3.yaml.j2 new file mode 100644 index 0000000..32dbf0b --- /dev/null +++ b/cellarium/workflows/configs/hvg_seurat_v3.yaml.j2 @@ -0,0 +1,46 @@ +seed_everything: true +trainer: + accelerator: {{ accelerator }} + strategy: + class_path: lightning.pytorch.strategies.DDPStrategy + init_args: + broadcast_buffers: false + max_epochs: 2 + default_root_dir: ./outputs +model: + model: + class_path: cellarium.ml.models.HVGSeuratV3 + init_args: + n_top_genes: {{ n_top_genes }} + use_batch_key: {{ 'true' if batch_index_n else 'false' }} + flavor: {{ flavor }} + output_path: ./outputs/seurat_v3_hvg_genes.csv +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: +{% include 'partials/_dadc_filenames_limits.yaml.j2' %} + indices_strict: true +{% if batch_index_n %} + obs_columns_to_validate: + - {{ batch_index_n }} +{% else %} + obs_columns_to_validate: [] +{% endif %} + batch_keys: + x_ng: + attr: X + convert_fn: cellarium.ml.utilities.data.densify +{% include 'partials/_var_names_g.yaml.j2' %} +{% if batch_index_n %} + batch_index_n: + attr: obs + key: {{ batch_index_n }} + convert_fn: cellarium.ml.utilities.data.categories_to_codes +{% endif %} + batch_size: {{ batch_size }} + iteration_strategy: cache_efficient + num_workers: {{ num_workers }} + prefetch_factor: {{ prefetch_factor if prefetch_factor is not none else 'null' }} + shuffle: true +ckpt_path: null \ No newline at end of file diff --git a/cellarium/workflows/configs/incremental_pca.yaml.j2 b/cellarium/workflows/configs/incremental_pca.yaml.j2 new file mode 100644 index 0000000..8031064 --- /dev/null +++ b/cellarium/workflows/configs/incremental_pca.yaml.j2 @@ -0,0 +1,42 @@ +seed_everything: true +trainer: + accelerator: {{ accelerator }} + strategy: + class_path: lightning.pytorch.strategies.DDPStrategy + dict_kwargs: + broadcast_buffers: false + max_epochs: 1 + logger: null + default_root_dir: ./outputs + callbacks: + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + dirpath: ./outputs/checkpoints/ + save_last: true +model: +{% include 'partials/_filter_cpu_transforms.yaml.j2' %} + transforms: + - cellarium.ml.transforms.Densify +{% include 'partials/_transforms_pflogpf.yaml.j2' %} +{% include 'partials/_zscore_transform.yaml.j2' %} + model: + class_path: cellarium.ml.models.IncrementalPCA + init_args: + var_names_g: null + n_components: {{ n_components }} +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: +{% include 'partials/_dadc_filenames_limits.yaml.j2' %} + obs_columns_to_validate: [] + batch_keys: + x_ng: + attr: X + convert_fn: cellarium.ml.utilities.data.keep_sparse +{% include 'partials/_var_names_g.yaml.j2' %} + batch_size: {{ batch_size }} + num_workers: {{ num_workers }} + prefetch_factor: {{ prefetch_factor if prefetch_factor is not none else 'null' }} + persistent_workers: false +ckpt_path: null \ No newline at end of file diff --git a/cellarium/workflows/configs/incremental_pca_predict.yaml.j2 b/cellarium/workflows/configs/incremental_pca_predict.yaml.j2 new file mode 100644 index 0000000..0fdd210 --- /dev/null +++ b/cellarium/workflows/configs/incremental_pca_predict.yaml.j2 @@ -0,0 +1,45 @@ +seed_everything: true +trainer: + accelerator: {{ accelerator }} + strategy: + class_path: lightning.pytorch.strategies.DDPStrategy + dict_kwargs: + broadcast_buffers: false + max_epochs: 1 + logger: null + default_root_dir: ./outputs + callbacks: + - class_path: cellarium.ml.callbacks.PredictionWriter + init_args: + output_dir: ./outputs/predictions + inference_mode: true +model: +{% include 'partials/_filter_cpu_transforms.yaml.j2' %} + transforms: + - cellarium.ml.transforms.Densify +{% include 'partials/_transforms_pflogpf.yaml.j2' %} +{% include 'partials/_zscore_transform.yaml.j2' %} + model: + class_path: cellarium.ml.models.IncrementalPCA + init_args: + var_names_g: null + n_components: {{ n_components }} +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: +{% include 'partials/_dadc_filenames_limits.yaml.j2' %} + obs_columns_to_validate: [] + batch_keys: + x_ng: + attr: X + convert_fn: cellarium.ml.utilities.data.keep_sparse + obs_names_n: + attr: obs_names +{% include 'partials/_var_names_g.yaml.j2' %} + batch_size: {{ batch_size }} + num_workers: {{ num_workers }} + prefetch_factor: {{ prefetch_factor if prefetch_factor is not none else 'null' }} + persistent_workers: false +return_predictions: false +ckpt_path: {{ pca_model }} diff --git a/cellarium/workflows/configs/onepass_mean_var_std.yaml.j2 b/cellarium/workflows/configs/onepass_mean_var_std.yaml.j2 new file mode 100644 index 0000000..b7c3067 --- /dev/null +++ b/cellarium/workflows/configs/onepass_mean_var_std.yaml.j2 @@ -0,0 +1,47 @@ +seed_everything: true +trainer: + accelerator: {{ accelerator }} + strategy: + class_path: lightning.pytorch.strategies.DDPStrategy + dict_kwargs: + broadcast_buffers: false + max_epochs: 1 + logger: null + default_root_dir: ./outputs +model: + transforms: +{% if sparse_dataloader %} + - cellarium.ml.transforms.Densify +{% endif %} +{% include 'partials/_transforms_pflogpf.yaml.j2' %} + model: + class_path: cellarium.ml.models.OnePassMeanVarStd + init_args: + algorithm: shifted_data + output_path: ./outputs/onepass_mean_var_std.csv +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: +{% include 'partials/_dadc_filenames_limits.yaml.j2' %} + obs_columns_to_validate: [] + batch_keys: + x_ng: + attr: X +{% if sparse_dataloader %} + convert_fn: cellarium.ml.utilities.data.to_torch_sparse_csr +{% else %} + convert_fn: cellarium.ml.utilities.data.densify +{% endif %} +{% include 'partials/_var_names_g.yaml.j2' %} +{% if total_mrna_umis_key %} + total_mrna_umis_n: + attr: obs + key: {{ total_mrna_umis_key }} +{% endif %} + batch_size: {{ batch_size }} + iteration_strategy: cache_efficient + num_workers: {{ num_workers }} + prefetch_factor: {{ prefetch_factor if prefetch_factor is not none else 'null' }} + shuffle: true +ckpt_path: null \ No newline at end of file diff --git a/cellarium/workflows/configs/partials/_dadc_filenames_limits.yaml.j2 b/cellarium/workflows/configs/partials/_dadc_filenames_limits.yaml.j2 new file mode 100644 index 0000000..73539b8 --- /dev/null +++ b/cellarium/workflows/configs/partials/_dadc_filenames_limits.yaml.j2 @@ -0,0 +1,14 @@ +{% if filenames %} + filenames: +{% for f in filenames %} + - {{ f }} +{% endfor %} + limits: +{% for l in limits %} + - {{ l }} +{% endfor %} +{% else %} + filenames: [] + limits: [] +{% endif %} + max_cache_size: {{ max_cache_size }} diff --git a/cellarium/workflows/configs/partials/_filter_cpu_transforms.yaml.j2 b/cellarium/workflows/configs/partials/_filter_cpu_transforms.yaml.j2 new file mode 100644 index 0000000..823b434 --- /dev/null +++ b/cellarium/workflows/configs/partials/_filter_cpu_transforms.yaml.j2 @@ -0,0 +1,10 @@ + cpu_transforms: + - class_path: cellarium.ml.transforms.Filter + init_args: + filter_list: + !FileLoader + file_path: {{ hvg_csv }} + loader_fn: pandas.read_csv + attr: "gene" + convert_fn: pandas.Series.to_list + ordering: true diff --git a/cellarium/workflows/configs/partials/_transforms_pflogpf.yaml.j2 b/cellarium/workflows/configs/partials/_transforms_pflogpf.yaml.j2 new file mode 100644 index 0000000..0de4aec --- /dev/null +++ b/cellarium/workflows/configs/partials/_transforms_pflogpf.yaml.j2 @@ -0,0 +1,12 @@ +{% if use_pflogpf %} + - cellarium.ml.transforms.PFlogPF +{% else %} +{% if apply_normalize_total %} + - class_path: cellarium.ml.transforms.NormalizeTotal + init_args: + target_count: {{ target_count }} +{% endif %} +{% if apply_log1p %} + - cellarium.ml.transforms.Log1p +{% endif %} +{% endif %} diff --git a/cellarium/workflows/configs/partials/_var_names_g.yaml.j2 b/cellarium/workflows/configs/partials/_var_names_g.yaml.j2 new file mode 100644 index 0000000..05ba129 --- /dev/null +++ b/cellarium/workflows/configs/partials/_var_names_g.yaml.j2 @@ -0,0 +1,7 @@ + var_names_g: +{% if var_names_key %} + attr: var + key: {{ var_names_key }} +{% else %} + attr: var_names +{% endif %} diff --git a/cellarium/workflows/configs/partials/_zscore_transform.yaml.j2 b/cellarium/workflows/configs/partials/_zscore_transform.yaml.j2 new file mode 100644 index 0000000..cd13a1d --- /dev/null +++ b/cellarium/workflows/configs/partials/_zscore_transform.yaml.j2 @@ -0,0 +1,17 @@ +{% if zscore_genes %} + - class_path: cellarium.ml.transforms.ZScore + init_args: + !FileMultiLoader + file_path: {{ onepass_csv }} + loader_fn: pandas.read_csv + fields: + mean_g: + attr: mean_g.values + convert_fn: torch.FloatTensor + std_g: + attr: std_g.values + convert_fn: torch.FloatTensor + var_names_g: + attr: var_names_g + convert_fn: pandas.Series.to_numpy +{% endif %} diff --git a/cellarium/workflows/configs/scvi.yaml.j2 b/cellarium/workflows/configs/scvi.yaml.j2 new file mode 100644 index 0000000..03b8c08 --- /dev/null +++ b/cellarium/workflows/configs/scvi.yaml.j2 @@ -0,0 +1,91 @@ +seed_everything: true +trainer: + accelerator: {{ accelerator }} + max_epochs: 50 + log_every_n_steps: 100 + gradient_clip_algorithm: norm + gradient_clip_val: 0.5 + inference_mode: false + default_root_dir: ./outputs + callbacks: + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + dirpath: ./outputs/checkpoints/ + every_n_epochs: 1 + save_last: true + - class_path: lightning.pytorch.callbacks.LearningRateMonitor + init_args: + logging_interval: step +model: +{% include 'partials/_filter_cpu_transforms.yaml.j2' %} + model: + class_path: cellarium.ml.models.SingleCellVariationalInference + init_args: + var_names_g: null + n_batch: null + batch_embedded: true + batch_representation_sampled: true + n_cats_per_cov: null + kl_warmup_epochs: 20 + batch_kl_weight_max: 0.0 + use_batch_norm: both + encoder: + hidden_layers: + - class_path: torch.nn.Linear + init_args: + out_features: 512 + final_layer: + class_path: torch.nn.Linear + init_args: {} + decoder: + hidden_layers: + - class_path: cellarium.ml.models.scvi.LinearWithBatch + init_args: + out_features: 512 + label_to_bias_hidden_layers: [] + final_layer: + class_path: torch.nn.Linear + init_args: {} + final_additive_bias: false + n_latent: {{ n_latent }} + optim_fn: torch.optim.AdamW + optim_kwargs: + lr: 1e-4 +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: +{% include 'partials/_dadc_filenames_limits.yaml.j2' %} +{% if batch_key %} + obs_columns_to_validate: + - {{ batch_index_n }} +{% for k in categorical_covariate_keys %} + - {{ k }} +{% endfor %} +{% else %} + obs_columns_to_validate: [] +{% endif %} + batch_keys: + x_ng: + attr: X + convert_fn: cellarium.ml.utilities.data.densify +{% include 'partials/_var_names_g.yaml.j2' %} + batch_index_n: + attr: obs + key: {{ batch_key }} + convert_fn: cellarium.ml.utilities.data.categories_to_codes +{% if categorical_covariate_keys %} + categorical_covariate_index_nd: + attr: obs + keys: +{% for k in categorical_covariate_keys %} + - {{ k }} +{% endfor %} + convert_fn: cellarium.ml.utilities.data.categories_to_codes +{% endif %} + batch_size: {{ batch_size }} + num_workers: {{ num_workers }} + prefetch_factor: {{ prefetch_factor if prefetch_factor is not none else 'null' }} + persistent_workers: false + shuffle: true +ckpt_path: null \ No newline at end of file diff --git a/cellarium/workflows/example/__init__.py b/cellarium/workflows/example/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/cellarium/workflows/example/components.py b/cellarium/workflows/example/components.py deleted file mode 100644 index 89289ea..0000000 --- a/cellarium/workflows/example/components.py +++ /dev/null @@ -1,19 +0,0 @@ -from kfp import dsl - - -@dsl.component(base_image="python:3.10-slim") -def example_component_1(config: str): - """ - Test Example component - """ - print(config) - print("Example component is being executed....") - - -@dsl.component(base_image="python:3.10-slim") -def example_component_2(config: str): - """ - Test Example component - """ - print(config) - print("Second example component is being executed...") diff --git a/cellarium/workflows/example/pipelines.py b/cellarium/workflows/example/pipelines.py deleted file mode 100644 index 0a3fda5..0000000 --- a/cellarium/workflows/example/pipelines.py +++ /dev/null @@ -1,31 +0,0 @@ -from kfp import dsl -from cellarium.workflows import kfp_helpers -from cellarium.workflows.example import components - - -@dsl.pipeline() -def example_pipeline(component_1_config: str, component_2_config: str): - """ - KFP pipeline to run PCA train pipeline. - - """ - component_job_1 = kfp_helpers.create_job( - component_func=components.example_component_1, - display_name="Example Job 1", - config=component_1_config - ) - - component_job_2 = kfp_helpers.create_job( - component_func=components.example_component_2, - display_name="Example Job 2", - config=component_2_config - ) - - task_1 = component_job_1() - task_2 = component_job_2() - - task_2.after(task_1) - - -if __name__ == '__main__': - example_pipeline() diff --git a/cellarium/workflows/example/submit_example_pipeline.py b/cellarium/workflows/example/submit_example_pipeline.py deleted file mode 100644 index fce4e90..0000000 --- a/cellarium/workflows/example/submit_example_pipeline.py +++ /dev/null @@ -1,25 +0,0 @@ -import click - -from cellarium.workflows.example import pipelines -from cellarium.workflows import kfp_helpers - - -@click.command() -@click.option("--project_id") -@click.option("--location") -@click.option("--display_name") -@click.option("--component_1_config") -@click.option("--component_2_config") -def submit_example(project_id: str, location: str, display_name: str, component_1_config: str, component_2_config: str): - kfp_helpers.submit_pipeline( - pipeline_func=pipelines.example_pipeline, - project_id=project_id, - location=location, - pipeline_display_name=display_name, - pipeline_kwargs={"component_1_config": component_1_config, "component_2_config": component_2_config} - ) - print("Submitted pipeline!") - - -if __name__ == "__main__": - submit_example() diff --git a/cellarium/workflows/kfp_helpers.py b/cellarium/workflows/kfp_helpers.py deleted file mode 100644 index d18db3c..0000000 --- a/cellarium/workflows/kfp_helpers.py +++ /dev/null @@ -1,79 +0,0 @@ -import tempfile -import typing as t -import os - -from kfp import compiler -from kfp.components import BaseComponent -from google.cloud import aiplatform -from google_cloud_pipeline_components.v1.custom_job import create_custom_training_job_from_component - - -def create_job( - component_func: t.Callable[..., t.Any], - config: str, - display_name: str = "", - replica_count: int = 1, - machine_type: str = "n1-standard-4", - accelerator_type: str = "", - accelerator_count: int = 1, - boot_disk_size_gb: int = 100 -) -> t.Callable[[], t.Any]: - """ - Create a custom training Google Vertex AI job for running a custom training component. - - :param component_func: Custom training component. - :param config: Config file path on GCS. - :param display_name: Display name of component in Vertex AI - :param replica_count: The count of instances in the cluster. - :param machine_type: The type of the machine to run the CustomJob. - :param accelerator_type: The type of accelerator(s) that may be attached to the machine per `accelerator_count`. - :param accelerator_count: The number of accelerators to attach to the machine. - :param boot_disk_size_gb: Size in GB of the boot disk - - :return: Callable custom training job. - """ - job = create_custom_training_job_from_component( - component_func, - display_name=display_name, - replica_count=replica_count, - machine_type=machine_type, - accelerator_type=accelerator_type, - accelerator_count=accelerator_count, - boot_disk_size_gb=boot_disk_size_gb, - ) - - return lambda: job(config=config) - - -def submit_pipeline( - pipeline_func: t.Union[BaseComponent, t.Callable], - pipeline_display_name: str, - pipeline_kwargs: t.Dict[str, t.Any], - project_id: str, - location: str, -) -> None: - """ - Create and run a pipeline on Vertex AI Pipelines. Use a temporary file to compile the pipeline config, - then run the pipeline job and delete the temporary file. - - :param pipeline_func: Pipeline function, must be a function wrapped :func:`kfp.dsl.pipeline` decorator. - :param pipeline_display_name: A name displayed in the Vertex AI Pipelines UI. - :param pipeline_kwargs: Keyword arguments to pass to the pipeline function. - :param project_id: Google Cloud Project ID - :param location: Datacenter location of Google Cloud Platform to run the pipeline job. - """ - temp_file = tempfile.NamedTemporaryFile(suffix=".yaml") - os.environ["GRPC_DNS_RESOLVER"] = "native" - - aiplatform.init(project=project_id, location=location) - - compiler.Compiler().compile(pipeline_func=pipeline_func, package_path=temp_file.name) - - job = aiplatform.PipelineJob( - display_name=pipeline_display_name, - template_path=temp_file.name, - parameter_values=pipeline_kwargs, - ) - - job.submit() - temp_file.close() diff --git a/cellarium/workflows/nextflow/README.md b/cellarium/workflows/nextflow/README.md new file mode 100644 index 0000000..197b53e --- /dev/null +++ b/cellarium/workflows/nextflow/README.md @@ -0,0 +1,175 @@ +# The workflows + +- [PCA workflow](#pca-workflow-pca_workflownf) — full PCA training + prediction pipeline +- [HVG workflow](#hvg-workflow-hvg_workflownf) — standalone highly variable gene selection +- [Onepass workflow](#onepass-workflow-onepass_workflownf) — standalone mean/variance/std pass + +`nextflow.config` (in this directory) is auto-detected by all workflows. Run with `-profile gcp` for Google Cloud Batch or `-profile local` for local execution via a conda environment. + +--- + +# PCA workflow (`pca_workflow.nf`) + +Runs a four-step [cellarium-ml](https://github.com/cellarium-ai/cellarium-ml) PCA training and prediction workflow on Google Cloud Batch. + +See the [repo README](../../../README.md) for Nextflow installation, the nf-google plugin, and GCP prerequisites. + +## DAG + +``` +ONEPASS_MEAN_VAR ──┐ + ├──► INCREMENTAL_PCA ──► INCREMENTAL_PCA_PREDICT +HIGHLY_VARIABLE ───┘ + GENES +``` + +1. **`ONEPASS_MEAN_VAR`** — single pass over all cells to compute per-gene mean, variance, and std. Output: `onepass_mean_var_std.ckpt`. Runs in parallel with step 2. +2. **`HIGHLY_VARIABLE_GENES`** — selects HVGs (Seurat v3 flavor) per batch key. Output: `hvg.csv`. Runs in parallel with step 1. +3. **`INCREMENTAL_PCA`** — fits an incremental PCA model using ZScore normalization (from step 1) and HVG filtering (from step 2). Output: `pca_final.ckpt`. +4. **`INCREMENTAL_PCA_PREDICT`** — projects all cells through the trained PCA model. Output: `batch*.csv.gz`. + +Config templates for each step live in `cellarium/workflows/configs/` as `.yaml.j2` files and are rendered at runtime with `bin/render_config.py`. + +## Parameters + +All parameters can be overridden on the command line with `--param value`. + +**Infrastructure** (`nextflow.config`): + +| Parameter | Default | Description | +|---|---|---| +| `google_project` | `dsp-cellarium` | GCP project ID | +| `google_region` | `us-central1` | Region for Batch jobs and GCS buckets | +| `work_bucket` | `gs://cellarium-dev-central/workflows/tmp` | GCS work dir for Nextflow staging | +| `spot` | `false` | Use preemptible VMs (`true` saves ~70%; evictions auto-retried) | +| `disk_size` | `750 GB` | pd-ssd disk per task VM | +| `container` | `cellarium-ml:0.0.8` | Container image for `-profile gcp` runs | + +**Pipeline** (`pca_workflow.nf`): + +| Parameter | Default | Description | +|---|---|---| +| `dataset_dir` | *(see pca_workflow.nf)* | GCS directory (or local path) of input `.h5ad` files | +| `outdir` | *(see pca_workflow.nf)* | GCS path (or local dir) where outputs are published | +| `n_components` | `64` | Number of PCA components | +| `n_top_genes` | `8000` | Number of highly variable genes to select | +| `flavor` | `seurat_v3` | HVG selection method | +| `batch_index_n` | `assay_suspension_type` | Obs column for batch-aware HVG selection; pass `''` to disable | + +## Running on Google Cloud Batch + +```bash +cd cellarium/workflows/nextflow + +# Override dataset and output locations +nextflow run pca_workflow.nf \ + -profile gcp \ + --run_label my_pca_workflow_name \ + --dataset_dir 'gs://my-bucket/data/extract_files' \ + --outdir 'gs://my-bucket/outputs/pca_run_001' + +# Use spot VMs for lower cost (~70% savings; evictions auto-retried) +nextflow run pca_workflow.nf -profile gcp --spot true + +# Specify 2000 hvgs and 20 PCs +nextflow run pca_workflow.nf -profile gcp --n_top_genes 2000 --n_components 20 + +# Resume a failed or interrupted run from where it left off +nextflow run pca_workflow.nf -profile gcp -resume +``` + +Nextflow prints a run name (e.g. `festive_curie`). Monitor progress: + +```bash +# Live log +nextflow log festive_curie + +# Poll until complete +watch -n 30 nextflow log festive_curie -f 'process,status,exit,duration' +``` + +Batch jobs are also visible in the [Google Cloud Batch console](https://console.cloud.google.com/batch/jobs). + +## Running locally + +Activate your cellarium conda environment first. Processes launched by Nextflow inherit the environment nextflow was launched from, so `render_config.py` and `cellarium-ml` are found automatically. + +```bash +conda activate cellarium +cd cellarium/workflows/nextflow + +nextflow run pca_workflow.nf -profile local \ + --dataset_dir /path/to/local/h5ad/dir \ + --outdir ./local_outputs +``` + +--- + +# HVG workflow (`hvg_workflow.nf`) + +Runs just the highly variable gene selection step in isolation. Useful for tuning HVG parameters independently before a full PCA run. + +**`HIGHLY_VARIABLE_GENES`** — fits `HVGSeuratV3` over the dataset, computing per-gene variability scores per batch key. Output: a CSV of selected gene IDs published to `outdir/hvg_seurat_v3/`. + +## Parameters + +| Parameter | Default | Description | +|---|---|---| +| `dataset_dir` | *(see hvg_workflow.nf)* | Input `.h5ad` directory | +| `outdir` | *(see hvg_workflow.nf)* | Output directory | +| `n_top_genes` | `8000` | Number of HVGs to select | +| `flavor` | `seurat_v3` | HVG selection method | +| `batch_index_n` | `assay_suspension_type` | Obs column for batch-aware selection; pass `''` to disable | + +## Running + +```bash +cd cellarium/workflows/nextflow + +# GCP +nextflow run hvg_workflow.nf \ + -profile gcp \ + --dataset_dir 'gs://my-bucket/data/extract_files' \ + --outdir 'gs://my-bucket/outputs/hvg_run_001' + +# Local +nextflow run hvg_workflow.nf \ + -profile local \ + --dataset_dir /path/to/local/h5ad/dir \ + --outdir ./local_outputs + +# Choose the number of HVGs +nextflow run hvg_workflow.nf -profile gcp --n_top_genes 2000 +``` + +--- + +# Onepass workflow (`onepass_workflow.nf`) + +Runs just the mean/variance/std pass in isolation. Produces statistics used downstream by `INCREMENTAL_PCA` for ZScore normalization. + +**`ONEPASS_MEAN_VAR`** — computes per-gene mean, variance, and std over all cells using `NormalizeTotal` → `Log1p` transforms. Output: `onepass_mean_var_std.ckpt` published to `outdir/onepass_mean_var_std/`. + +## Parameters + +| Parameter | Default | Description | +|---|---|---| +| `dataset_dir` | *(see onepass_workflow.nf)* | Input `.h5ad` directory | +| `outdir` | *(see onepass_workflow.nf)* | Output directory | + +## Running + +```bash +cd cellarium/workflows/nextflow + +# GCP +nextflow run onepass_workflow.nf -profile gcp \ + --dataset_dir 'gs://my-bucket/data/extract_files' \ + --outdir 'gs://my-bucket/outputs/onepass_run_001' + +# Local +nextflow run onepass_workflow.nf -profile local \ + --dataset_dir /path/to/local/h5ad/dir \ + --outdir ./local_outputs +``` + diff --git a/cellarium/workflows/nextflow/bin/hvg_helper.py b/cellarium/workflows/nextflow/bin/hvg_helper.py new file mode 100755 index 0000000..a0b9153 --- /dev/null +++ b/cellarium/workflows/nextflow/bin/hvg_helper.py @@ -0,0 +1,50 @@ +#!/usr/bin/env python3 +"""Run highly variable gene selection on a dataset directory and output a CSV file.""" +import argparse +from pathlib import Path + +import pandas as pd +import torch + +from cellarium.ml.preprocessing import ( + kotliar_compute_highly_variable_genes, + seurat_compute_highly_variable_genes, +) + + +def main(): + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("method", help="Method for HVG selection in ['seurat', 'kotliar']") + parser.add_argument("n_top_genes", type=int, help="Number of top HVGs to select") + parser.add_argument("onepass_csv", type=str, help="Path to one-pass CSV file") + parser.add_argument("--output", default="hvgs.csv", help="Output path (default: hvgs.csv)") + args = parser.parse_args() + + if args.method == "seurat": + hvg_fun = seurat_compute_highly_variable_genes + elif args.method == "kotliar": + hvg_fun = kotliar_compute_highly_variable_genes + else: + raise ValueError(f"Error: Invalid method '{args.method}'. Must be 'seurat' or 'kotliar'.") + + onepass_df = pd.read_csv(args.onepass_csv) + + kwargs = { + "var_names_g": onepass_df["var_names_g"].values, + "mean_g": torch.from_numpy(onepass_df["mean_g"].values), + "var_g": torch.from_numpy(onepass_df["var_g"].values), + "n_top_genes": args.n_top_genes, + } + + hvg_df = hvg_fun(**kwargs) + hvg_df.index.name = "gene" + hvg_df = hvg_df[hvg_df["highly_variable"]].copy() + hvg_df.to_csv(args.output, index=True) + print(f"Highly variable genes written to {args.output}") + + +if __name__ == "__main__": + main() diff --git a/cellarium/workflows/nextflow/bin/maybe_pip_install.sh b/cellarium/workflows/nextflow/bin/maybe_pip_install.sh new file mode 100755 index 0000000..68cd7db --- /dev/null +++ b/cellarium/workflows/nextflow/bin/maybe_pip_install.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +# Install cellarium-ml from GitHub at a specific git ref (tag, branch, or SHA). +# Usage: maybe_pip_install.sh +# +# If is empty or the string "null", this script is a no-op and the +# version already present in the container image is used unchanged. + +set -euo pipefail + +ref="${1:-}" + +if [ -z "$ref" ] || [ "$ref" = "null" ]; then + exit 0 +fi + +pip install "git+https://github.com/cellarium-ai/cellarium-ml.git@${ref}" + +# no sudo, removing this: +# # Increase i/o read-ahead on the local SSD (mounted at /tmp by Google Batch). +# LOCAL_SSD_DEVICE=$(findmnt -n -o SOURCE /tmp 2>/dev/null || df -P /tmp 2>/dev/null | tail -1 | awk '{print $1}') +# if [ -n "$LOCAL_SSD_DEVICE" ]; then +# sudo blockdev --setra 8192 "$LOCAL_SSD_DEVICE" +# fi diff --git a/cellarium/workflows/nextflow/bin/render_config.py b/cellarium/workflows/nextflow/bin/render_config.py new file mode 100755 index 0000000..6ffcd34 --- /dev/null +++ b/cellarium/workflows/nextflow/bin/render_config.py @@ -0,0 +1,318 @@ +#!/usr/bin/env python3 +"""Render one or more Jinja2 YAML templates from a shared context. + +Usage (new multi-template form): + render_config.py [KEY=VALUE ...] \\ + [--context-file base.json] \\ + [--dump-context out.json] \\ + [--template OUTPUT:TEMPLATE ...] + +Usage (legacy single-template form, still supported): + render_config.py KEY=VALUE [KEY=VALUE ...] \\ + [--output run_config.yaml] + +Value coercion: + "None" / "null" -> None (falsy; omits conditional blocks) + "true" / "false" -> bool + digits-only -> int + "" (empty) -> "" (also falsy) + anything else -> str + +Special key: + dataset_dir= + Scans the directory for *.h5ad files (lexicographic order), reads + n_obs from each via h5py, and injects: + filenames - list of absolute file path strings + limits - cumulative sum of n_obs across files + If dataset_dir is absent, filenames=[] and limits=[] are injected + so preview renders (without real data) produce valid YAML. + +Flags: + --context-file FILE + Load FILE (JSON) as the base context. CLI KEY=VALUE args still + override individual keys. + --dump-context FILE + Write the fully coerced context (before dataset_dir expansion) to + FILE as JSON. Can be combined with --template to also render. + --template OUTPUT:TEMPLATE (repeatable, short: -t) + Render TEMPLATE to OUTPUT. All --template targets share one context + and one dataset_dir scan. The loader searches the directory of the + first template provided. + --output FILE + Output path for the legacy positional-template form (default: + run_config.yaml). + --patch OUTPUT:INPUT (repeatable, short: -p) + Load INPUT as a pre-rendered YAML, inject filenames+limits from a + dataset_dir scan (if dataset_dir KEY=VALUE is given), and apply any + other KEY=VALUE overrides by matching key names recursively in the + YAML tree. Writes the result to OUTPUT. No Jinja2 processing is + performed — this is a surgical patch of an already-rendered config. +""" +import argparse +import json +import sys +from pathlib import Path + + +# --------------------------------------------------------------------------- +# Value coercion +# --------------------------------------------------------------------------- + +def coerce(value: str): + """Convert a CLI string value to an appropriate Python type.""" + if value in ("None", "null"): + return None + if value.lower() == "true": + return True + if value.lower() == "false": + return False + try: + return int(value) + except ValueError: + return value + + +# --------------------------------------------------------------------------- +# Dataset scanning +# --------------------------------------------------------------------------- + +def n_obs(path: Path) -> int: + import h5py + with h5py.File(path, "r") as f: + obs = f["obs"] + index_key = obs.attrs.get("_index", obs.attrs.get("index", None)) + if index_key is not None: + return f["obs"][index_key].shape[0] + return obs.shape[0] + + +def dataset_dir_to_filenames_and_limits(dataset_dir: str): + files = sorted(Path(dataset_dir).glob("*.h5ad")) + if not files: + print(f"Error: no .h5ad files found in {dataset_dir}", file=sys.stderr) + sys.exit(1) + counts = [n_obs(f) for f in files] + filenames = [str(f) for f in files] + limits = [] + cumsum = 0 + for c in counts: + cumsum += c + limits.append(cumsum) + return filenames, limits + + +# --------------------------------------------------------------------------- +# YAML patching (pre-rendered configs — inject filenames/limits at runtime) +# --------------------------------------------------------------------------- + +def _inject_filenames_limits(node, filenames: list, limits: list) -> bool: + """Recursively find the first dict with a 'filenames' key and replace it.""" + if isinstance(node, dict): + if "filenames" in node: + node["filenames"] = filenames + node["limits"] = limits + return True + for v in node.values(): + if _inject_filenames_limits(v, filenames, limits): + return True + elif isinstance(node, list): + for item in node: + if _inject_filenames_limits(item, filenames, limits): + return True + return False + + +def _replace_keys(node, overrides: dict): + """Recursively replace values for any matching keys in the YAML tree.""" + if isinstance(node, dict): + for k in list(node.keys()): + if k in overrides: + node[k] = overrides[k] + else: + _replace_keys(node[k], overrides) + elif isinstance(node, list): + for item in node: + _replace_keys(item, overrides) + + +def patch_yaml(input_path: str, output_path: str, context: dict): + """Patch a pre-rendered YAML with runtime values and write to output_path.""" + try: + import yaml + except ImportError: + print("Error: PyYAML is not installed. Run: pip install pyyaml", file=sys.stderr) + sys.exit(1) + + with open(input_path) as f: + data = yaml.safe_load(f) + + if "dataset_dir" in context: + filenames, limits = dataset_dir_to_filenames_and_limits(context["dataset_dir"]) + if not _inject_filenames_limits(data, filenames, limits): + print(f"Warning: no 'filenames' key found in {input_path} to patch", + file=sys.stderr) + + other = {k: v for k, v in context.items() if k != "dataset_dir"} + if other: + _replace_keys(data, other) + + with open(output_path, "w") as f: + yaml.dump(data, f, default_flow_style=False, sort_keys=False) + print(f"Patched {input_path} -> {output_path}") + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +def main(): + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + # Legacy positional template (optional – new mode uses --template flags) + parser.add_argument( + "template_positional", nargs="?", default=None, metavar="template", + help="Template file path (legacy positional form)", + ) + parser.add_argument( + "variables", nargs="*", metavar="key=value", + help="Context variables; override --context-file values", + ) + parser.add_argument( + "--template", "-t", action="append", dest="extra_templates", + metavar="OUTPUT:TEMPLATE", + help="Render OUTPUT from TEMPLATE (repeatable)", + ) + parser.add_argument( + "--context-file", metavar="FILE", + help="JSON file to load as base context (CLI key=value overrides it)", + ) + parser.add_argument( + "--dump-context", metavar="FILE", + help="Write coerced context (before dataset scan) to FILE as JSON", + ) + parser.add_argument( + "--output", default="run_config.yaml", + help="Output path for legacy positional-template form", + ) + parser.add_argument( + "--patch", "-p", action="append", dest="patches", + metavar="OUTPUT:INPUT", + help="Patch INPUT (pre-rendered YAML) with dataset_dir scan and key overrides, write to OUTPUT (repeatable)", + ) + args = parser.parse_args() + + # ------------------------------------------------------------------ + # Distinguish legacy positional template from a stray key=value arg + # that argparse grabbed as template_positional. + # ------------------------------------------------------------------ + all_variables = list(args.variables) + legacy_template = None + + if args.template_positional is not None: + if "=" in args.template_positional: + # It's actually a key=value – prepend it to variables + all_variables = [args.template_positional] + all_variables + else: + legacy_template = args.template_positional + + # Build the full list of (output_path, template_path) pairs + templates_to_render = [] + if legacy_template: + templates_to_render.append((args.output, legacy_template)) + if args.extra_templates: + for spec in args.extra_templates: + output, _, template = spec.partition(":") + if not template: + print(f"Error: --template argument must be OUTPUT:TEMPLATE, got {spec!r}", + file=sys.stderr) + sys.exit(1) + templates_to_render.append((output, template)) + + patches_to_apply = [] + if args.patches: + for spec in args.patches: + output, _, input_path = spec.partition(":") + if not input_path: + print(f"Error: --patch argument must be OUTPUT:INPUT, got {spec!r}", + file=sys.stderr) + sys.exit(1) + patches_to_apply.append((output, input_path)) + + # ------------------------------------------------------------------ + # Build context + # ------------------------------------------------------------------ + context: dict = {} + + # Base: load from --context-file if provided + if args.context_file: + with open(args.context_file) as fh: + context.update(json.load(fh)) + + # Override / extend with CLI key=value pairs + for item in all_variables: + key, _, value = item.partition("=") + if key: + context[key] = coerce(value) + + # ------------------------------------------------------------------ + # Optionally dump context (before dataset expansion) to JSON + # ------------------------------------------------------------------ + if args.dump_context: + Path(args.dump_context).write_text(json.dumps(context, indent=2, default=str)) + print(f"Context written to {args.dump_context}") + + # ------------------------------------------------------------------ + # Early exit if nothing to do + # ------------------------------------------------------------------ + if not templates_to_render and not patches_to_apply: + return + + # ------------------------------------------------------------------ + # Jinja2 template rendering + # ------------------------------------------------------------------ + if templates_to_render: + render_context = dict(context) + if "dataset_dir" in render_context: + dataset_dir = render_context.pop("dataset_dir") + render_context["filenames"], render_context["limits"] = \ + dataset_dir_to_filenames_and_limits(dataset_dir) + else: + render_context.setdefault("filenames", []) + render_context.setdefault("limits", []) + + try: + from jinja2 import Environment, FileSystemLoader, StrictUndefined + except ImportError: + print("Error: jinja2 is not installed. Run: pip install jinja2", file=sys.stderr) + sys.exit(1) + + first_template_path = Path(templates_to_render[0][1]) + template_dir = str(first_template_path.parent) if first_template_path.parent != Path(".") else "." + env = Environment( + loader=FileSystemLoader(template_dir), + undefined=StrictUndefined, + trim_blocks=True, + lstrip_blocks=True, + keep_trailing_newline=True, + ) + + for output_path, template_path in templates_to_render: + tmpl = env.get_template(Path(template_path).name) + rendered = tmpl.render(**render_context) + Path(output_path).write_text(rendered) + print(f"Rendered {template_path} -> {output_path}") + print(rendered) + print() + + # ------------------------------------------------------------------ + # YAML patching (pre-rendered configs — inject filenames/limits and + # any other runtime key overrides without re-running Jinja2) + # ------------------------------------------------------------------ + for output_path, input_path in patches_to_apply: + patch_yaml(input_path, output_path, context) + + +if __name__ == "__main__": + main() diff --git a/cellarium/workflows/nextflow/bin/run_with_gpu_monitor.sh b/cellarium/workflows/nextflow/bin/run_with_gpu_monitor.sh new file mode 100755 index 0000000..6ad1246 --- /dev/null +++ b/cellarium/workflows/nextflow/bin/run_with_gpu_monitor.sh @@ -0,0 +1,19 @@ +#!/usr/bin/env bash +# Run a command while logging GPU utilization and memory to gpu_metrics.log. +# Usage: run_with_gpu_monitor.sh [args...] +# +# Columns: timestamp, utilization.gpu (%), memory.used (MiB), memory.total (MiB) +# +# If nvidia-smi is not available (e.g. local CPU-only runs), the monitor is +# skipped silently and the command runs normally. + +if command -v nvidia-smi &>/dev/null; then + nvidia-smi \ + --query-gpu=timestamp,utilization.gpu,memory.used,memory.total \ + --format=csv,noheader,nounits \ + --loop=1 >> gpu_metrics.log 2>&1 & + GPU_MONITOR_PID=$! + trap "kill $GPU_MONITOR_PID 2>/dev/null || true" EXIT +fi + +"$@" diff --git a/cellarium/workflows/nextflow/bin/stage_dataset.sh b/cellarium/workflows/nextflow/bin/stage_dataset.sh new file mode 100755 index 0000000..22e6f82 --- /dev/null +++ b/cellarium/workflows/nextflow/bin/stage_dataset.sh @@ -0,0 +1,37 @@ +#!/usr/bin/env bash +# Stage an h5ad dataset for a task, optionally downloading it from GCS first. +# Usage: stage_dataset.sh [smoke_test] +# +# "true" – copy *.h5ad files from GCS to /tmp/dataset and +# print /tmp/dataset as the resolved path. +# "false" – print unchanged (no network I/O). +# +# [smoke_test] "true" – when gcp_download=true, download only the first +# 10 files instead of the full dataset. +# (No effect when gcp_download=false.) +# +# The resolved dataset path is printed to stdout so callers can capture it: +# _dataset_dir=$(stage_dataset.sh "${params.gcp_download}" "${dataset_dir}" "${params.smoke_test}") + +set -euo pipefail + +gcp_download="${1:-false}" +dataset_dir="${2:?Usage: stage_dataset.sh [smoke_test]}" +smoke_test="${3:-false}" + +if [ "$gcp_download" = "true" ]; then + mkdir -p /tmp/dataset + if [ "$smoke_test" = "true" ]; then + echo "smoke_test=true: downloading first 10 .h5ad files from ${dataset_dir}" >&2 + gcloud storage ls "${dataset_dir}/*.h5ad" \ + | head -10 \ + | xargs -I{} gcloud storage cp {} /tmp/dataset/ + else + gcloud storage cp "${dataset_dir}/*.h5ad" /tmp/dataset/ + fi + echo /tmp/dataset +else + # Streaming mode: pass GCS (or local) path through unchanged. + # smoke_test is not supported in streaming mode (out of scope). + echo "$dataset_dir" +fi diff --git a/cellarium/workflows/nextflow/bin/write_metadata.py b/cellarium/workflows/nextflow/bin/write_metadata.py new file mode 100755 index 0000000..a83b581 --- /dev/null +++ b/cellarium/workflows/nextflow/bin/write_metadata.py @@ -0,0 +1,53 @@ +#!/usr/bin/env python3 +"""Write run provenance metadata as a JSON file. + +Usage: + write_metadata.py [--output run_metadata.json] [--params-file context.json] KEY=VALUE ... + +All KEY=VALUE pairs are stored as top-level string fields in the JSON. +If --params-file is given, its contents are embedded under the "params" key. + +Nextflow places this script on $PATH for every task when it lives in nextflow/bin/. +""" +import argparse +import json +import sys +from pathlib import Path + + +def main(): + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "fields", nargs="*", metavar="key=value", + help="Metadata fields to record (stored as strings)", + ) + parser.add_argument( + "--output", default="run_metadata.json", + help="Output JSON path (default: run_metadata.json)", + ) + parser.add_argument( + "--params-file", metavar="FILE", + help="JSON file whose contents are embedded under the 'params' key", + ) + args = parser.parse_args() + + metadata: dict = {} + + for item in args.fields: + key, _, value = item.partition("=") + if key: + metadata[key] = value # keep as strings for provenance clarity + + if args.params_file: + with open(args.params_file) as fh: + metadata["params"] = json.load(fh) + + Path(args.output).write_text(json.dumps(metadata, indent=2)) + print(f"Metadata written to {args.output}") + + +if __name__ == "__main__": + main() diff --git a/cellarium/workflows/nextflow/modules/combo_hvg_pca_predict.nf b/cellarium/workflows/nextflow/modules/combo_hvg_pca_predict.nf new file mode 100644 index 0000000..4c68c80 --- /dev/null +++ b/cellarium/workflows/nextflow/modules/combo_hvg_pca_predict.nf @@ -0,0 +1,93 @@ +process ONEPASS_HVG_INCREMENTAL_PCA_PLUS_PREDICTION { + publishDir "${params.run_outdir}/pca/", mode: 'copy' + + input: + val fit_dataset_dir + val predict_dataset_dir + val hvg_method + path onepass_config + path pca_fit_config + path pca_predict_config + path seurat_v3_hvg_config + + output: + path 'outputs/checkpoints/last.ckpt', emit: final_model + path 'outputs/predictions/batch*.csv.gz', emit: pcs + path 'gpu_metrics.log', optional: true, emit: gpu_metrics + path 'onepass_config.yaml', emit: onepass_config_yaml + path 'seurat_v3_hvg_config.yaml', optional: true, emit: seurat_v3_hvg_config_yaml + path 'outputs/seurat_hvgs.csv', emit: seurat_hvg_csv + path 'outputs/kotliar_hvgs.csv', emit: kotliar_hvg_csv + path 'outputs/seurat_v3_hvg_genes__top${params.n_top_genes}__hvg_only.csv', optional: true, emit: seurat_v3_hvg_csv + path 'pca_fit_config.yaml', emit: pca_fit_config_yaml + path 'pca_predict_config.yaml', emit: pca_predict_config_yaml + + script: + """ + mkdir -p outputs/checkpoints + mkdir -p outputs/predictions + + maybe_pip_install.sh "${params.cellarium_ml_ref}" + + echo "Staging training dataset..." + _train_dataset_dir=\$(stage_dataset.sh "${params.gcp_download}" "${fit_dataset_dir}" "${params.smoke_test}") + + echo "Patching and running onepass_mean_var_std..." + ${params.python3_bin} \$(which render_config.py) \ + --patch onepass_config.yaml:${onepass_config} \ + "dataset_dir=\${_train_dataset_dir}" + + run_with_gpu_monitor.sh cellarium-ml onepass_mean_var_std fit -c onepass_config.yaml + + echo "Running HVG gene selection: seurat..." + ${params.python3_bin} \$(which hvg_helper.py) \ + seurat \ + ${params.n_top_genes} \ + outputs/onepass_mean_var_std.csv \ + --output outputs/seurat_hvgs.csv + + echo "Running HVG gene selection: kotliar..." + ${params.python3_bin} \$(which hvg_helper.py) \ + kotliar \ + ${params.n_top_genes} \ + outputs/onepass_mean_var_std.csv \ + --output outputs/kotliar_hvgs.csv + + if [ "${hvg_method}" = "seurat_v3" ]; then + echo "Running HVG gene selection: seurat_v3..." + ${params.python3_bin} \$(which render_config.py) \ + --patch seurat_v3_hvg_config.yaml:${seurat_v3_hvg_config} \ + "dataset_dir=\${_train_dataset_dir}" + + run_with_gpu_monitor.sh cellarium-ml hvg_seurat_v3 fit -c seurat_v3_hvg_config.yaml + hvg_csv=outputs/hvg_genes__top${params.n_top_genes}__hvg_only.csv + else + if [ "${hvg_method}" = "seurat" ]; then + hvg_csv=outputs/seurat_hvgs.csv + else + hvg_csv=outputs/kotliar_hvgs.csv + fi + fi + + echo "Patching and running incremental_pca fit..." + ${params.python3_bin} \$(which render_config.py) \ + --patch pca_fit_config.yaml:${pca_fit_config} \ + "dataset_dir=\${_train_dataset_dir}" + + run_with_gpu_monitor.sh cellarium-ml incremental_pca fit -c pca_fit_config.yaml + + if [ "${predict_dataset_dir}" != "${fit_dataset_dir}" ]; then + rm -r \${_train_dataset_dir} + _predict_dataset_dir=\$(stage_dataset.sh "${params.gcp_download}" "${predict_dataset_dir}" "${params.smoke_test}") + else + _predict_dataset_dir=\${_train_dataset_dir} + fi + + echo "Patching and running incremental_pca predict..." + ${params.python3_bin} \$(which render_config.py) \ + --patch pca_predict_config.yaml:${pca_predict_config} \ + "dataset_dir=\${_predict_dataset_dir}" + + run_with_gpu_monitor.sh cellarium-ml incremental_pca predict -c pca_predict_config.yaml + """ +} diff --git a/cellarium/workflows/nextflow/modules/onepass.nf b/cellarium/workflows/nextflow/modules/onepass.nf new file mode 100644 index 0000000..eee839f --- /dev/null +++ b/cellarium/workflows/nextflow/modules/onepass.nf @@ -0,0 +1,26 @@ +process ONEPASS_MEAN_VAR { + publishDir "${params.run_outdir}/onepass_mean_var_std/", mode: 'copy' + + input: + val dataset_dir + path onepass_config + + output: + path 'outputs/onepass_mean_var_std.csv', emit: onepass_csv + path 'gpu_metrics.log', optional: true, emit: gpu_metrics + path 'run_config.yaml', emit: config_yaml + + script: + """ + mkdir -p outputs + + maybe_pip_install.sh "${params.cellarium_ml_ref}" + _dataset_dir=\$(stage_dataset.sh "${params.gcp_download}" "${dataset_dir}" "${params.smoke_test}") + + ${params.python3_bin} \$(which render_config.py) \ + --patch run_config.yaml:${onepass_config} \ + "dataset_dir=\${_dataset_dir}" + + run_with_gpu_monitor.sh cellarium-ml onepass_mean_var_std fit -c run_config.yaml + """ +} diff --git a/cellarium/workflows/nextflow/modules/onepass_with_hvgs.nf b/cellarium/workflows/nextflow/modules/onepass_with_hvgs.nf new file mode 100644 index 0000000..9c443e4 --- /dev/null +++ b/cellarium/workflows/nextflow/modules/onepass_with_hvgs.nf @@ -0,0 +1,41 @@ +process ONEPASS_MEAN_VAR_WITH_HVGS { + publishDir "${params.run_outdir}/onepass_mean_var_std_with_hvgs/", mode: 'copy' + + input: + val dataset_dir + path onepass_config + val n_top_genes + + output: + path 'outputs/onepass_mean_var_std.csv', emit: onepass_csv + path 'seurat_hvgs.csv', emit: seurat_hvg_csv + path 'kotliar_hvgs.csv', emit: kotliar_hvg_csv + path 'gpu_metrics.log', optional: true, emit: gpu_metrics + path 'run_config.yaml', emit: config_yaml + + script: + """ + mkdir -p outputs + + maybe_pip_install.sh "${params.cellarium_ml_ref}" + _dataset_dir=\$(stage_dataset.sh "${params.gcp_download}" "${dataset_dir}" "${params.smoke_test}") + + ${params.python3_bin} \$(which render_config.py) \ + --patch run_config.yaml:${onepass_config} \ + "dataset_dir=\${_dataset_dir}" + + run_with_gpu_monitor.sh cellarium-ml onepass_mean_var_std fit -c run_config.yaml + + ${params.python3_bin} \$(which hvg_helper.py) \ + seurat \ + ${n_top_genes} \ + outputs/onepass_mean_var_std.csv \ + --output seurat_hvgs.csv + + ${params.python3_bin} \$(which hvg_helper.py) \ + kotliar \ + ${n_top_genes} \ + outputs/onepass_mean_var_std.csv \ + --output kotliar_hvgs.csv + """ +} diff --git a/cellarium/workflows/nextflow/modules/pca.nf b/cellarium/workflows/nextflow/modules/pca.nf new file mode 100644 index 0000000..a2df8b5 --- /dev/null +++ b/cellarium/workflows/nextflow/modules/pca.nf @@ -0,0 +1,30 @@ +process INCREMENTAL_PCA { + publishDir "${params.run_outdir}/pca/", mode: 'copy' + + input: + val dataset_dir + path onepass_csv + path hvg_csv + path pca_fit_config + + output: + path 'outputs/checkpoints/last.ckpt', emit: final_model + path 'gpu_metrics.log', optional: true, emit: gpu_metrics + path 'run_config.yaml', emit: config_yaml + + script: + """ + mkdir -p outputs/checkpoints + + maybe_pip_install.sh "${params.cellarium_ml_ref}" + _dataset_dir=\$(stage_dataset.sh "${params.gcp_download}" "${dataset_dir}" "${params.smoke_test}") + + ${params.python3_bin} \$(which render_config.py) \ + --patch run_config.yaml:${pca_fit_config} \ + "dataset_dir=\${_dataset_dir}" \ + "onepass_csv=./${onepass_csv}" \ + "hvg_csv=./${hvg_csv}" + + run_with_gpu_monitor.sh cellarium-ml incremental_pca fit -c run_config.yaml + """ +} diff --git a/cellarium/workflows/nextflow/modules/pca_plus_predict.nf b/cellarium/workflows/nextflow/modules/pca_plus_predict.nf new file mode 100644 index 0000000..df38c1f --- /dev/null +++ b/cellarium/workflows/nextflow/modules/pca_plus_predict.nf @@ -0,0 +1,50 @@ +process INCREMENTAL_PCA_PLUS_PREDICTION { + publishDir "${params.run_outdir}/pca/", mode: 'copy' + + input: + val fit_dataset_dir + val predict_dataset_dir + path onepass_csv + path hvg_csv + path pca_fit_config + path pca_predict_config + + output: + path 'outputs/checkpoints/last.ckpt', emit: final_model + path 'outputs/predictions/batch*.csv.gz', emit: pcs + path 'gpu_metrics.log', optional: true, emit: gpu_metrics + path 'fit_config.yaml', emit: fit_config_yaml + path 'predict_config.yaml', emit: predict_config_yaml + + script: + """ + mkdir -p outputs/checkpoints + mkdir -p outputs/predictions + + maybe_pip_install.sh "${params.cellarium_ml_ref}" + _train_dataset_dir=\$(stage_dataset.sh "${params.gcp_download}" "${fit_dataset_dir}" "${params.smoke_test}") + + ${params.python3_bin} \$(which render_config.py) \ + --patch fit_config.yaml:${pca_fit_config} \ + "dataset_dir=\${_train_dataset_dir}" \ + "onepass_csv=./${onepass_csv}" \ + "hvg_csv=./${hvg_csv}" + + run_with_gpu_monitor.sh cellarium-ml incremental_pca fit -c fit_config.yaml + + if [ "${predict_dataset_dir}" != "${fit_dataset_dir}" ]; then + rm -r \${_train_dataset_dir} # remove training dataset before staging prediction dataset + _predict_dataset_dir=\$(stage_dataset.sh "${params.gcp_download}" "${predict_dataset_dir}" "${params.smoke_test}") + else + _predict_dataset_dir=\${_train_dataset_dir} + fi + + ${params.python3_bin} \$(which render_config.py) \ + --patch predict_config.yaml:${pca_predict_config} \ + "dataset_dir=\${_predict_dataset_dir}" \ + "onepass_csv=./${onepass_csv}" \ + "hvg_csv=./${hvg_csv}" + + run_with_gpu_monitor.sh cellarium-ml incremental_pca predict -c predict_config.yaml + """ +} diff --git a/cellarium/workflows/nextflow/modules/pca_predict.nf b/cellarium/workflows/nextflow/modules/pca_predict.nf new file mode 100644 index 0000000..732c24e --- /dev/null +++ b/cellarium/workflows/nextflow/modules/pca_predict.nf @@ -0,0 +1,33 @@ +process INCREMENTAL_PCA_PREDICT { + publishDir "${params.run_outdir}/pca/predictions/", mode: 'copy' + + input: + val dataset_dir + path pca_model + path onepass_csv + path hvg_csv + path pca_predict_config + + output: + // grab all the output files + path 'outputs/predictions/batch*.csv.gz', emit: pcs + path 'gpu_metrics.log', optional: true, emit: gpu_metrics + path 'run_config.yaml', emit: config_yaml + + script: + """ + mkdir -p outputs/predictions + + maybe_pip_install.sh "${params.cellarium_ml_ref}" + _dataset_dir=\$(stage_dataset.sh "${params.gcp_download}" "${dataset_dir}" "${params.smoke_test}") + + ${params.python3_bin} \$(which render_config.py) \ + --patch run_config.yaml:${pca_predict_config} \ + "dataset_dir=\${_dataset_dir}" \ + "pca_model=./${pca_model}" \ + "onepass_csv=./${onepass_csv}" \ + "hvg_csv=./${hvg_csv}" + + run_with_gpu_monitor.sh cellarium-ml incremental_pca predict -c run_config.yaml + """ +} diff --git a/cellarium/workflows/nextflow/modules/render_configs.nf b/cellarium/workflows/nextflow/modules/render_configs.nf new file mode 100644 index 0000000..2dc07a2 --- /dev/null +++ b/cellarium/workflows/nextflow/modules/render_configs.nf @@ -0,0 +1,319 @@ +/* + * render_configs.nf + * + * Local pre-render processes – one per workflow type. + * + * These processes run on the SUBMITTING machine (executor = 'local') before + * any remote GPU job is dispatched. They produce: + * + * run_context.json – coerced params, shared by all downstream GCP processes + * as a --context-file argument to render_config.py. + * run_metadata.json – human-readable provenance record: dataset paths, + * timestamps, git SHA, Nextflow run name/session, full + * params dump. + * preview_*.yaml – fully rendered configs with filenames: [] / limits: [] + * (placeholder – dataset not yet scanned). Archived to + * params.run_outdir/configs/ so you have a pre-run record. + * + * All processes are overridden to executor = 'local' in nextflow.config via: + * process { withName: ~/RENDER_.*_CONFIGS/ { executor = 'local' } } + */ + + +// ── PCA workflow (all-in-one combo: onepass + HVG + PCA fit + PCA predict) ── + +process RENDER_PCA_CONFIGS { + publishDir "${params.run_outdir}/configs/", mode: 'copy' + + input: + val fit_dataset_dir + val predict_dataset_dir + val hvg_method + val run_name + val session_id + path configs_dir + + output: + path 'run_context.json', emit: context + path 'run_metadata.json', emit: metadata + path 'onepass_config.yaml', emit: onepass_config + path 'pca_fit_config.yaml', emit: pca_fit_config + path 'pca_predict_config.yaml', emit: pca_predict_config + path 'seurat_v3_hvg_config.yaml', emit: seurat_v3_hvg_config + + script: + """ + # ── Compute hvg_csv path from hvg_method (pre-bake into context) ── + if [ "${hvg_method}" = "seurat_v3" ]; then + _hvg_csv="outputs/hvg_genes__top${params.n_top_genes}__hvg_only.csv" + elif [ "${hvg_method}" = "seurat" ]; then + _hvg_csv="outputs/seurat_hvgs.csv" + else + _hvg_csv="outputs/kotliar_hvgs.csv" + fi + + # ── Dump shared context (all model/transform params + fixed inter-step paths) ── + ${params.python3_bin} \$(which render_config.py) \ + --dump-context run_context.json \ + "num_workers=${params.num_workers}" \ + "prefetch_factor=${params.prefetch_factor}" \ + "var_names_key=${params.var_names_key}" \ + "total_mrna_umis_key=${params.total_mrna_umis_key}" \ + "apply_normalize_total=${params.apply_normalize_total}" \ + "target_count=${params.target_count}" \ + "apply_log1p=${params.apply_log1p}" \ + "use_pflogpf=${params.use_pflogpf}" \ + "sparse_dataloader=${params.sparse_dataloader}" \ + "accelerator=${params.accelerator}" \ + "batch_size=${params.batch_size}" \ + "max_cache_size=${params.max_cache_size}" \ + "zscore_genes=${params.zscore_genes}" \ + "n_components=${params.n_components}" \ + "n_top_genes=${params.n_top_genes}" \ + "batch_index_n=${params.batch_index_n}" \ + "seurat_v3_flavor=${params.seurat_v3_flavor}" \ + "cellarium_ml_ref=${params.cellarium_ml_ref}" \ + "onepass_csv=outputs/onepass_mean_var_std.csv" \ + "pca_model=outputs/checkpoints/last.ckpt" \ + "hvg_csv=\${_hvg_csv}" + + # ── Render runtime configs (filenames/limits default to [] — patched on VM) ── + ${params.python3_bin} \$(which render_config.py) \ + --context-file run_context.json \ + --template onepass_config.yaml:${configs_dir}/onepass_mean_var_std.yaml.j2 \ + --template pca_fit_config.yaml:${configs_dir}/incremental_pca.yaml.j2 \ + --template pca_predict_config.yaml:${configs_dir}/incremental_pca_predict.yaml.j2 \ + --template seurat_v3_hvg_config.yaml:${configs_dir}/hvg_seurat_v3.yaml.j2 + + # ── Write provenance metadata ── + _git_sha=\$(git -C "${projectDir}" rev-parse HEAD 2>/dev/null || echo unknown) + ${params.python3_bin} \$(which write_metadata.py) \ + --output run_metadata.json \ + --params-file run_context.json \ + "fit_dataset_dir=${fit_dataset_dir}" \ + "predict_dataset_dir=${predict_dataset_dir}" \ + "hvg_method=${hvg_method}" \ + "timestamp=\$(date -u +%Y-%m-%dT%H:%M:%SZ)" \ + "cellarium_ml_ref=${params.cellarium_ml_ref}" \ + "nextflow_run_name=${run_name}" \ + "nextflow_session_id=${session_id}" \ + "workflows_git_sha=\${_git_sha}" + """ +} + + +// ── Onepass-only workflow ────────────────────────────────────────────────── + +process RENDER_ONEPASS_CONFIGS { + publishDir "${params.run_outdir}/configs/", mode: 'copy' + + input: + val dataset_dir + val run_name + val session_id + path configs_dir + + output: + path 'run_context.json', emit: context + path 'run_metadata.json', emit: metadata + path 'onepass_config.yaml', emit: onepass_config + + script: + """ + ${params.python3_bin} \$(which render_config.py) \ + --dump-context run_context.json \ + "num_workers=${params.num_workers}" \ + "prefetch_factor=${params.prefetch_factor}" \ + "var_names_key=${params.var_names_key}" \ + "total_mrna_umis_key=${params.total_mrna_umis_key}" \ + "apply_normalize_total=${params.apply_normalize_total}" \ + "target_count=${params.target_count}" \ + "apply_log1p=${params.apply_log1p}" \ + "use_pflogpf=${params.use_pflogpf}" \ + "sparse_dataloader=${params.sparse_dataloader}" \ + "accelerator=${params.accelerator}" \ + "batch_size=${params.batch_size}" \ + "max_cache_size=${params.max_cache_size}" \ + "cellarium_ml_ref=${params.cellarium_ml_ref}" + + ${params.python3_bin} \$(which render_config.py) \ + --context-file run_context.json \ + --template onepass_config.yaml:${configs_dir}/onepass_mean_var_std.yaml.j2 + + _git_sha=\$(git -C "${projectDir}" rev-parse HEAD 2>/dev/null || echo unknown) + ${params.python3_bin} \$(which write_metadata.py) \ + --output run_metadata.json \ + --params-file run_context.json \ + "dataset_dir=${dataset_dir}" \ + "timestamp=\$(date -u +%Y-%m-%dT%H:%M:%SZ)" \ + "cellarium_ml_ref=${params.cellarium_ml_ref}" \ + "nextflow_run_name=${run_name}" \ + "nextflow_session_id=${session_id}" \ + "workflows_git_sha=\${_git_sha}" + """ +} + + +// ── Seurat-v3 HVG-only workflow ─────────────────────────────────────────── + +process RENDER_SEURAT_HVG_CONFIGS { + publishDir "${params.run_outdir}/configs/", mode: 'copy' + + input: + val dataset_dir + val run_name + val session_id + path configs_dir + + output: + path 'run_context.json', emit: context + path 'run_metadata.json', emit: metadata + path 'seurat_v3_hvg_config.yaml', emit: seurat_v3_hvg_config + + script: + """ + ${params.python3_bin} \$(which render_config.py) \ + --dump-context run_context.json \ + "num_workers=${params.num_workers}" \ + "prefetch_factor=${params.prefetch_factor}" \ + "var_names_key=${params.var_names_key}" \ + "accelerator=${params.accelerator}" \ + "batch_size=${params.batch_size}" \ + "max_cache_size=${params.max_cache_size}" \ + "n_top_genes=${params.n_top_genes}" \ + "flavor=${params.seurat_v3_flavor}" \ + "batch_index_n=${params.batch_index_n}" \ + "cellarium_ml_ref=${params.cellarium_ml_ref}" + + ${params.python3_bin} \$(which render_config.py) \ + --context-file run_context.json \ + --template seurat_v3_hvg_config.yaml:${configs_dir}/hvg_seurat_v3.yaml.j2 + + _git_sha=\$(git -C "${projectDir}" rev-parse HEAD 2>/dev/null || echo unknown) + ${params.python3_bin} \$(which write_metadata.py) \ + --output run_metadata.json \ + --params-file run_context.json \ + "dataset_dir=${dataset_dir}" \ + "timestamp=\$(date -u +%Y-%m-%dT%H:%M:%SZ)" \ + "cellarium_ml_ref=${params.cellarium_ml_ref}" \ + "nextflow_run_name=${run_name}" \ + "nextflow_session_id=${session_id}" \ + "workflows_git_sha=\${_git_sha}" + """ +} + + +// ── scVI workflow ────────────────────────────────────────────────────────── + +process RENDER_SCVI_CONFIGS { + publishDir "${params.run_outdir}/configs/", mode: 'copy' + + input: + val dataset_dir + val hvg_csv_path + val run_name + val session_id + path configs_dir + + output: + path 'run_context.json', emit: context + path 'run_metadata.json', emit: metadata + path 'scvi_config.yaml', emit: scvi_config + + script: + """ + ${params.python3_bin} \$(which render_config.py) \ + --dump-context run_context.json \ + "num_workers=${params.num_workers}" \ + "prefetch_factor=${params.prefetch_factor}" \ + "var_names_key=${params.var_names_key}" \ + "accelerator=${params.accelerator}" \ + "batch_size=${params.batch_size}" \ + "max_cache_size=${params.max_cache_size}" \ + "batch_key=${params.batch_key}" \ + "categorical_covariate_keys=${params.categorical_covariate_keys}" \ + "n_latent=${params.n_latent}" \ + "batch_index_n=${params.batch_index_n}" \ + "hvg_csv=${hvg_csv_path}" \ + "cellarium_ml_ref=${params.cellarium_ml_ref}" + + ${params.python3_bin} \$(which render_config.py) \ + --context-file run_context.json \ + --template scvi_config.yaml:${configs_dir}/scvi.yaml.j2 + + _git_sha=\$(git -C "${projectDir}" rev-parse HEAD 2>/dev/null || echo unknown) + ${params.python3_bin} \$(which write_metadata.py) \ + --output run_metadata.json \ + --params-file run_context.json \ + "dataset_dir=${dataset_dir}" \ + "hvg_csv=${hvg_csv_path}" \ + "timestamp=\$(date -u +%Y-%m-%dT%H:%M:%SZ)" \ + "cellarium_ml_ref=${params.cellarium_ml_ref}" \ + "nextflow_run_name=${run_name}" \ + "nextflow_session_id=${session_id}" \ + "workflows_git_sha=\${_git_sha}" + """ +} + + +// ── PCA-no-HVG workflow (pre-existing onepass_csv and hvg_csv) ──────────── + +process RENDER_PCA_NO_HVG_CONFIGS { + publishDir "${params.run_outdir}/configs/", mode: 'copy' + + input: + val train_dataset_dir + val predict_dataset_dir + val onepass_csv_path + val hvg_csv_path + val run_name + val session_id + path configs_dir + + output: + path 'run_context.json', emit: context + path 'run_metadata.json', emit: metadata + path 'pca_fit_config.yaml', emit: pca_fit_config + path 'pca_predict_config.yaml', emit: pca_predict_config + + script: + """ + ${params.python3_bin} \$(which render_config.py) \ + --dump-context run_context.json \ + "num_workers=${params.num_workers}" \ + "prefetch_factor=${params.prefetch_factor}" \ + "var_names_key=${params.var_names_key}" \ + "accelerator=${params.accelerator}" \ + "batch_size=${params.batch_size}" \ + "max_cache_size=${params.max_cache_size}" \ + "use_pflogpf=${params.use_pflogpf}" \ + "apply_normalize_total=${params.apply_normalize_total}" \ + "target_count=${params.target_count}" \ + "apply_log1p=${params.apply_log1p}" \ + "zscore_genes=${params.zscore_genes}" \ + "n_components=${params.n_components}" \ + "onepass_csv=${onepass_csv_path}" \ + "hvg_csv=${hvg_csv_path}" \ + "pca_model=outputs/checkpoints/last.ckpt" \ + "cellarium_ml_ref=${params.cellarium_ml_ref}" + + ${params.python3_bin} \$(which render_config.py) \ + --context-file run_context.json \ + --template pca_fit_config.yaml:${configs_dir}/incremental_pca.yaml.j2 \ + --template pca_predict_config.yaml:${configs_dir}/incremental_pca_predict.yaml.j2 + + _git_sha=\$(git -C "${projectDir}" rev-parse HEAD 2>/dev/null || echo unknown) + ${params.python3_bin} \$(which write_metadata.py) \ + --output run_metadata.json \ + --params-file run_context.json \ + "train_dataset_dir=${train_dataset_dir}" \ + "predict_dataset_dir=${predict_dataset_dir}" \ + "onepass_csv=${onepass_csv_path}" \ + "hvg_csv=${hvg_csv_path}" \ + "timestamp=\$(date -u +%Y-%m-%dT%H:%M:%SZ)" \ + "cellarium_ml_ref=${params.cellarium_ml_ref}" \ + "nextflow_run_name=${run_name}" \ + "nextflow_session_id=${session_id}" \ + "workflows_git_sha=\${_git_sha}" + """ +} diff --git a/cellarium/workflows/nextflow/modules/script_hvg.nf b/cellarium/workflows/nextflow/modules/script_hvg.nf new file mode 100644 index 0000000..6f0d733 --- /dev/null +++ b/cellarium/workflows/nextflow/modules/script_hvg.nf @@ -0,0 +1,20 @@ +process HIGHLY_VARIABLE_GENES { + publishDir "${params.run_outdir}/hvg/", mode: 'copy' + + input: + val onepass_csv + val method + val n_top_genes + + output: + path 'hvgs.csv', emit: hvg_csv + + script: + """ + mkdir -p outputs + + maybe_pip_install.sh "${params.cellarium_ml_ref}" + + ${params.python3_bin} \$(which hvg_helper.py) ${method} ${n_top_genes} ${onepass_csv} --output hvgs.csv + """ +} diff --git a/cellarium/workflows/nextflow/modules/scvi.nf b/cellarium/workflows/nextflow/modules/scvi.nf new file mode 100644 index 0000000..6c31b31 --- /dev/null +++ b/cellarium/workflows/nextflow/modules/scvi.nf @@ -0,0 +1,28 @@ +process SCVI { + publishDir "${params.run_outdir}/scvi/", mode: 'copy' + + input: + val dataset_dir + path hvg_csv + path scvi_config + + output: + path 'outputs/checkpoints/last.ckpt', emit: final_model + path 'gpu_metrics.log', optional: true, emit: gpu_metrics + path 'run_config.yaml', emit: config_yaml + + script: + """ + mkdir -p outputs/checkpoints + + maybe_pip_install.sh "${params.cellarium_ml_ref}" + _dataset_dir=\$(stage_dataset.sh "${params.gcp_download}" "${dataset_dir}" "${params.smoke_test}") + + ${params.python3_bin} \$(which render_config.py) \ + --patch run_config.yaml:${scvi_config} \ + "dataset_dir=\${_dataset_dir}" \ + "hvg_csv=./${hvg_csv}" + + run_with_gpu_monitor.sh cellarium-ml scvi fit -c run_config.yaml + """ +} diff --git a/cellarium/workflows/nextflow/modules/seurat_v3_hvg.nf b/cellarium/workflows/nextflow/modules/seurat_v3_hvg.nf new file mode 100644 index 0000000..dd4ff05 --- /dev/null +++ b/cellarium/workflows/nextflow/modules/seurat_v3_hvg.nf @@ -0,0 +1,26 @@ +process SEURAT_V3_HIGHLY_VARIABLE_GENES { + publishDir "${params.run_outdir}/hvg_seurat_v3/", mode: 'copy' + + input: + val dataset_dir + path seurat_v3_hvg_config + + output: + path 'outputs/hvg_genes__top*__hvg_only.csv', emit: hvg_csv + path 'gpu_metrics.log', optional: true, emit: gpu_metrics + path 'run_config.yaml', emit: config_yaml + + script: + """ + mkdir -p outputs + + maybe_pip_install.sh "${params.cellarium_ml_ref}" + _dataset_dir=\$(stage_dataset.sh "${params.gcp_download}" "${dataset_dir}" "${params.smoke_test}") + + ${params.python3_bin} \$(which render_config.py) \ + --patch run_config.yaml:${seurat_v3_hvg_config} \ + "dataset_dir=\${_dataset_dir}" + + run_with_gpu_monitor.sh cellarium-ml hvg_seurat_v3 fit -c run_config.yaml + """ +} diff --git a/cellarium/workflows/nextflow/nextflow.config b/cellarium/workflows/nextflow/nextflow.config new file mode 100644 index 0000000..a2982b2 --- /dev/null +++ b/cellarium/workflows/nextflow/nextflow.config @@ -0,0 +1,177 @@ +/* + * Nextflow configuration for cellarium-ml workflows. + * + * Usage: + * nextflow run .nf -profile gcp # Google Cloud Batch + * nextflow run .nf -profile local # local machine via conda + * + * Override any param on the command line, e.g.: + * nextflow run pca_workflow.nf -profile gcp --spot true + */ + +plugins { + id 'nf-google' +} + +// ── Shared parameters ────────────────────────────────────────────────────── +def _conda_prefix = System.getenv('CONDA_PREFIX') +params { + // Container image used by all processes on GCP + container = 'us-central1-docker.pkg.dev/broad-dsde-methods/cellarium-ai/cellarium-ml:latest' + + // Python interpreter for task scripts. + // When CONDA_PREFIX is set (local run with activated env), use that env's + // python3 directly — avoids JVM PATH inheritance issues. + // On GCP the container's python3 is always on PATH, so 'python3' suffices. + python3_bin = _conda_prefix ? "${_conda_prefix}/bin/python3" : 'python3' + + // Google Cloud settings (used by the gcp profile) + google_project = 'dsp-cellarium' + google_region = 'us-central1' + work_bucket = 'gs://cellarium-dev-central/workflows/tmp' + + // Use spot (preemptible) VMs. Set --spot true to cut costs ~70%. + // With maxRetries = 2, spot evictions are handled automatically. + spot = false + + // Local SSD is provisioned on each GCP task VM for fast dataset downloads. + // Must be a multiple of 375 GB. Mounted at /tmp by Google Batch. + disk_size = '750 GB' + + // Set to true in the gcp profile so module scripts pre-download the dataset + // from GCS to the local NVMe SSD before training begins. + gcp_download = false + + run_label = 'cellarium-workflow' + + // Optional: override the cellarium-ml version in the container by installing + // directly from GitHub before each task runs. Accepts any git ref: a tag, + // branch name, or full/short SHA (e.g. 'v1.2.3', 'my-branch', 'abc1234'). + // Set to null (the default) to skip the pip install and use whatever version + // is already baked into the container image. + cellarium_ml_ref = null + + // When true, stage_dataset.sh downloads only the first 10 .h5ad files + // (gcp_download=true runs only). Use to validate the pipeline quickly + // before committing to a full run. + smoke_test = false + + // Timestamped output root. Computed per-workflow-file from params.outdir + // and params.run_label. Override on the CLI to write to a specific path. + // When null the workflow file sets it automatically. + run_outdir = null +} + +// ── Process defaults (apply to all profiles) ────────────────────────────── +// process { +// maxRetries = 2 +// errorStrategy = 'retry' +// } + +// ── Profiles ─────────────────────────────────────────────────────────────── +profiles { + + // Run on local machine. Activate your cellarium conda env first: + // conda activate cellarium + // nextflow run .nf -profile local + local { + process.executor = 'local' + params.num_workers = 1 + params.prefetch_factor = 4 + } + + // Run on Google Cloud Batch with per-process hardware configuration. + gcp { + params.num_workers = 8 + params.prefetch_factor = 4 + params.python3_bin = 'python3' + workDir = params.work_bucket + + params.gcp_download = true + + process { + executor = 'google-batch' + container = params.container + // Mount a local NVMe SSD at /tmp. Module scripts download the + // dataset here before training to avoid gcsfuse read overhead. + disk = params.disk_size + diskType = 'local-ssd' + // --ipc=host shares the host IPC namespace, giving PyTorch + // DataLoader workers unrestricted shared memory. + containerOptions = '--ipc=host --privileged' + + withName: 'ONEPASS_MEAN_VAR' { + machineType = 'n1-standard-16' + cpus = 16 + memory = '60 GB' + accelerator = [request: 1, type: 'nvidia-tesla-t4'] + } + withName: 'ONEPASS_MEAN_VAR_WITH_HVGS' { + machineType = 'n1-standard-16' + cpus = 16 + memory = '60 GB' + accelerator = [request: 1, type: 'nvidia-tesla-t4'] + } + withName: 'SEURAT_V3_HIGHLY_VARIABLE_GENES' { + machineType = 'n1-standard-16' + cpus = 16 + memory = '60 GB' + accelerator = [request: 1, type: 'nvidia-tesla-t4'] + } + withName: 'HIGHLY_VARIABLE_GENES' { + machineType = 'e2-standard-2' + cpus = 2 + memory = '8 GB' + } + withName: 'INCREMENTAL_PCA' { + machineType = 'n1-standard-16' + cpus = 16 + memory = '60 GB' + accelerator = [request: 1, type: 'nvidia-tesla-t4'] + } + withName: 'INCREMENTAL_PCA_PREDICT' { + machineType = 'n1-standard-16' + cpus = 16 + memory = '60 GB' + accelerator = [request: 1, type: 'nvidia-tesla-t4'] + } + withName: 'INCREMENTAL_PCA_PLUS_PREDICTION' { + machineType = 'n1-standard-16' + cpus = 16 + memory = '60 GB' + accelerator = [request: 1, type: 'nvidia-tesla-t4'] + } + // RENDER_*_CONFIGS run on the submitting machine. + // scratch overrides the process working dir to a local POSIX path, + // bypassing the global GCS workDir which the local executor cannot use. + withName: 'RENDER_PCA_CONFIGS' { executor = 'local'; scratch = '/tmp/nf-local' } + withName: 'RENDER_ONEPASS_CONFIGS' { executor = 'local'; scratch = '/tmp/nf-local' } + withName: 'RENDER_SEURAT_HVG_CONFIGS' { executor = 'local'; scratch = '/tmp/nf-local' } + withName: 'RENDER_SCVI_CONFIGS' { executor = 'local'; scratch = '/tmp/nf-local' } + withName: 'RENDER_PCA_NO_HVG_CONFIGS' { executor = 'local'; scratch = '/tmp/nf-local' } + } + + google { + project = params.google_project + location = params.google_region + batch { + spot = params.spot + installGpuDrivers = true + labels = [pipeline: 'cellarium-ml', run: params.run_label] + } + } + } + +} + +report { + enabled = true + overwrite = true + file = 'nextflow_report.html' +} + +timeline { + enabled = true + overwrite = true + file = 'nextflow_timeline.html' +} diff --git a/cellarium/workflows/nextflow/nextflow_schema.json b/cellarium/workflows/nextflow/nextflow_schema.json new file mode 100644 index 0000000..707426b --- /dev/null +++ b/cellarium/workflows/nextflow/nextflow_schema.json @@ -0,0 +1,246 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema", + "$id": "https://raw.githubusercontent.com/cellarium-ai/cellarium-workflows/main/cellarium/workflows/nextflow/nextflow_schema.json", + "title": "Cellarium-ML Nextflow Workflow Parameters", + "description": "JSON Schema for the cellarium-ml Nextflow workflows. Enables --validate-params and nextflow run --help.", + "type": "object", + "definitions": { + "shared_compute": { + "title": "Compute settings", + "type": "object", + "properties": { + "num_workers": { + "type": "integer", + "minimum": 0, + "default": 8, + "description": "Number of DataLoader worker processes." + }, + "prefetch_factor": { + "type": "integer", + "minimum": 1, + "default": 4, + "description": "Number of batches to prefetch per DataLoader worker." + }, + "accelerator": { + "type": "string", + "default": "auto", + "description": "PyTorch Lightning accelerator. 'auto' selects GPU if available." + }, + "batch_size": { + "type": "integer", + "minimum": 1, + "default": 5000, + "description": "Number of cells per training batch." + }, + "max_cache_size": { + "type": "integer", + "minimum": 1, + "default": 4, + "description": "Maximum number of AnnData shards cached in memory." + }, + "sparse_dataloader": { + "type": "boolean", + "default": true, + "description": "Whether to use sparse tensors in the DataLoader." + } + } + }, + "shared_normalization": { + "title": "Data normalization", + "type": "object", + "properties": { + "var_names_key": { + "type": "string", + "default": "null", + "description": "obs column holding gene identifiers. Use 'null' to skip var filtering." + }, + "total_mrna_umis_key": { + "type": "string", + "default": "null", + "description": "obs column with total mRNA UMI counts for normalize_total." + }, + "apply_normalize_total": { + "type": "boolean", + "default": true, + "description": "Apply total-count normalization before log1p." + }, + "target_count": { + "type": "number", + "default": 10000, + "description": "Target sum for normalize_total." + }, + "apply_log1p": { + "type": "boolean", + "default": true, + "description": "Apply log1p transform after normalize_total." + }, + "use_pflogpf": { + "type": "boolean", + "default": false, + "description": "Use PFlogPF normalization instead of normalize_total+log1p." + } + } + }, + "shared_infra": { + "title": "Infrastructure / GCP settings", + "type": "object", + "properties": { + "container": { + "type": "string", + "description": "Container image used by all GCP processes." + }, + "python3_bin": { + "type": "string", + "description": "Python interpreter path for task scripts." + }, + "google_project": { + "type": "string", + "default": "dsp-cellarium", + "description": "GCP project ID." + }, + "google_region": { + "type": "string", + "default": "us-central1", + "description": "GCP region for Google Batch jobs." + }, + "work_bucket": { + "type": "string", + "description": "GCS bucket path used as Nextflow workDir." + }, + "spot": { + "type": "boolean", + "default": false, + "description": "Use spot (preemptible) VMs on GCP." + }, + "disk_size": { + "type": "string", + "default": "750 GB", + "description": "Local SSD size provisioned per GCP task VM (multiple of 375 GB)." + }, + "gcp_download": { + "type": "boolean", + "default": false, + "description": "When true, stage_dataset.sh downloads dataset to local NVMe before training." + }, + "smoke_test": { + "type": "boolean", + "default": false, + "description": "When true (and gcp_download=true), stage_dataset.sh downloads only the first 10 .h5ad files." + }, + "cellarium_ml_ref": { + "type": ["string", "null"], + "default": null, + "description": "Git ref (tag/branch/SHA) of cellarium-ml to pip-install before each task. null uses the container image version." + }, + "run_label": { + "type": "string", + "default": "cellarium-workflow", + "description": "Human-readable label for this run; used in output directory paths and GCP batch labels." + }, + "run_outdir": { + "type": ["string", "null"], + "default": null, + "description": "Timestamped output root directory. Computed automatically from outdir + run_label if null." + } + } + } + }, + "allOf": [ + {"$ref": "#/definitions/shared_compute"}, + {"$ref": "#/definitions/shared_normalization"}, + {"$ref": "#/definitions/shared_infra"} + ], + "properties": { + "outdir": { + "type": "string", + "description": "Base output directory (local path or gs:// URI)." + }, + + "dataset_dir": { + "type": "string", + "description": "[onepass / seurat_v3_hvg workflows] Path or gs:// URI to the sharded .h5ad dataset directory." + }, + "fit_dataset_dir": { + "type": "string", + "description": "[pca / pca_no_hvg workflows] Dataset used for training." + }, + "predict_dataset_dir": { + "type": "string", + "description": "[pca / pca_no_hvg workflows] Dataset used for prediction (may equal fit_dataset_dir)." + }, + + "n_components": { + "type": "integer", + "minimum": 1, + "default": 64, + "description": "Number of PCA components." + }, + "n_top_genes": { + "type": "integer", + "minimum": 1, + "default": 8000, + "description": "Number of highly-variable genes to select." + }, + "zscore_genes": { + "type": "boolean", + "default": true, + "description": "Z-score genes before PCA." + }, + + "hvg_method": { + "type": "string", + "enum": ["seurat_v3", "seurat", "kotliar"], + "default": "seurat", + "description": "HVG selection method. 'seurat_v3' runs the GPU seurat-v3 model; 'seurat' and 'kotliar' use the CPU hvg_helper.py script." + }, + "seurat_v3_flavor": { + "type": "string", + "enum": ["seurat_v3", "seurat_v3_paper"], + "default": "seurat_v3", + "description": "Flavor for seurat_v3 HVG selection." + }, + "batch_index_n": { + "type": "string", + "default": "null", + "description": "obs column to use as batch variable in HVG and scVI models." + }, + + "onepass_csv": { + "type": "string", + "description": "[pca_no_hvg workflow] Path or gs:// URI to a pre-computed onepass_mean_var_std.csv." + }, + "hvg_csv": { + "type": "string", + "description": "[pca_no_hvg / scvi workflows] Path or gs:// URI to a pre-computed HVG CSV." + }, + + "batch_key": { + "type": "string", + "description": "[scvi workflow] obs column to use as the scVI batch key." + }, + "categorical_covariate_keys": { + "type": "string", + "default": "null", + "description": "[scvi workflow] Comma-separated obs columns to use as categorical covariates." + }, + "n_latent": { + "type": "integer", + "minimum": 1, + "default": 128, + "description": "[scvi workflow] Dimensionality of the scVI latent space." + }, + + "config_onepass": { + "type": "string", + "description": "Path to the onepass_mean_var_std.yaml.j2 Jinja2 template." + }, + "config_hvg": { + "type": "string", + "description": "Path to the hvg_seurat_v3.yaml.j2 Jinja2 template." + }, + "config_scvi": { + "type": "string", + "description": "Path to the scvi.yaml.j2 Jinja2 template." + } + } +} diff --git a/cellarium/workflows/nextflow/onepass_workflow.nf b/cellarium/workflows/nextflow/onepass_workflow.nf new file mode 100644 index 0000000..93c7891 --- /dev/null +++ b/cellarium/workflows/nextflow/onepass_workflow.nf @@ -0,0 +1,43 @@ +#!/usr/bin/env nextflow + +params.dataset_dir = 'gs://cellarium-nexus-file-system-3293a8/pipeline/data-extracts/20260403_cas_pca_model_10x/extract_files' +params.outdir = 'gs://cellarium-dev-central/workflows/tmp' +params.config_onepass = "${projectDir}/../configs/onepass_mean_var_std.yaml.j2" +params.num_workers = 8 +params.prefetch_factor = 4 +params.var_names_key = 'null' +params.total_mrna_umis_key = 'null' +params.apply_normalize_total = true +params.target_count = 10000 +params.apply_log1p = true +params.use_pflogpf = false +params.sparse_dataloader = true +params.accelerator = 'auto' +params.batch_size = 5000 +params.max_cache_size = 4 + +def _run_ts = new java.text.SimpleDateFormat("yyyyMMdd_HHmmss").format(new Date()) +params.run_outdir = params.run_outdir ?: "${params.outdir}/${params.run_label}/${_run_ts}" + +include { RENDER_ONEPASS_CONFIGS } from './modules/render_configs.nf' +include { ONEPASS_MEAN_VAR } from './modules/onepass.nf' + +workflow { + dataset_ch = Channel.value( + params.dataset_dir.startsWith('gs://') + ? params.dataset_dir + : file(params.dataset_dir).toAbsolutePath().toString()) + configs_dir_ch = Channel.value(file(params.config_onepass).parent) + + render_out = RENDER_ONEPASS_CONFIGS( + dataset_dir = dataset_ch, + run_name = workflow.runName, + session_id = workflow.sessionId, + configs_dir = configs_dir_ch + ) + + ONEPASS_MEAN_VAR( + dataset_dir = dataset_ch, + onepass_config = render_out.onepass_config + ) +} \ No newline at end of file diff --git a/cellarium/workflows/nextflow/pca_no_hvg_workflow.nf b/cellarium/workflows/nextflow/pca_no_hvg_workflow.nf new file mode 100644 index 0000000..0bf998a --- /dev/null +++ b/cellarium/workflows/nextflow/pca_no_hvg_workflow.nf @@ -0,0 +1,69 @@ +#!/usr/bin/env nextflow + +// Train PCA on one dataset, run prediction on a (potentially different) dataset. +// ONEPASS and HVG come from some previous run and are not re-run here. +// INCREMENTAL_PCA trains on train_dataset_dir, then INCREMENTAL_PCA_PREDICT runs +// on predict_dataset_dir on a separate machine — avoiding co-localising both datasets. + +params.train_dataset_dir = 'gs://cellarium-nexus-file-system-3293a8/pipeline/data-extracts/20260403_cas_pca_model_10x/extract_files' +params.predict_dataset_dir = 'gs://cellarium-nexus-file-system-3293a8/pipeline/data-extracts/20260403_cas_pca_vsindex_10x/extract_files' +params.outdir = 'gs://cellarium-dev-central/workflows/nextflow_cas_pca' +params.config_onepass = "${projectDir}/../configs/onepass_mean_var_std.yaml.j2" +params.onepass_csv = 'gs://cellarium-dev-central/workflows/nextflow_cas_pca/onepass_mean_var_std/outputs/onepass.csv' // from ONEPASS step +params.hvg_csv = 'gs://cellarium-dev-central/workflows/nextflow_cas_pca/onepass_mean_var_std/seurat_hvg.csv' // from HVG step +params.n_components = 32 +params.n_top_genes = 4000 +params.batch_index_n = 'null' +params.num_workers = 8 +params.prefetch_factor = 4 +params.var_names_key = 'feature_id' +params.accelerator = 'auto' +params.batch_size = 5000 +params.max_cache_size = 4 +params.use_pflogpf = false // whether to use PFlogPF data normalization +params.zscore_genes = true // whether to z-score genes before PCA +params.total_mrna_umis_key = 'raw_sum' +params.apply_normalize_total = true +params.target_count = 10000 +params.apply_log1p = true +params.sparse_dataloader = true + +def _run_ts = new java.text.SimpleDateFormat("yyyyMMdd_HHmmss").format(new Date()) +params.run_outdir = params.run_outdir ?: "${params.outdir}/${params.run_label}/${_run_ts}" + +include { RENDER_PCA_NO_HVG_CONFIGS } from './modules/render_configs.nf' +include { INCREMENTAL_PCA_PLUS_PREDICTION } from './modules/pca_plus_predict.nf' + +workflow { + train_ch = Channel.value( + params.train_dataset_dir.startsWith('gs://') + ? params.train_dataset_dir + : file(params.train_dataset_dir).toAbsolutePath().toString()) + predict_ch = Channel.value( + params.predict_dataset_dir.startsWith('gs://') + ? params.predict_dataset_dir + : file(params.predict_dataset_dir).toAbsolutePath().toString()) + onepass_csv_ch = Channel.value(file(params.onepass_csv)) + hvg_csv_ch = Channel.value(file(params.hvg_csv)) + configs_dir_ch = Channel.value(file(params.config_onepass).parent) + + render_out = RENDER_PCA_NO_HVG_CONFIGS( + train_dataset_dir = train_ch, + predict_dataset_dir = predict_ch, + onepass_csv_path = params.onepass_csv, + hvg_csv_path = params.hvg_csv, + run_name = workflow.runName, + session_id = workflow.sessionId, + configs_dir = configs_dir_ch + ) + + // INCREMENTAL_PCA_PLUS_PREDICTION waits for both, runs prediction on same machine + INCREMENTAL_PCA_PLUS_PREDICTION( + fit_dataset_dir = train_ch, + predict_dataset_dir = predict_ch, + onepass_csv = onepass_csv_ch, + hvg_csv = hvg_csv_ch, + pca_fit_config = render_out.pca_fit_config, + pca_predict_config = render_out.pca_predict_config + ) +} diff --git a/cellarium/workflows/nextflow/pca_workflow.nf b/cellarium/workflows/nextflow/pca_workflow.nf new file mode 100644 index 0000000..4c5db6c --- /dev/null +++ b/cellarium/workflows/nextflow/pca_workflow.nf @@ -0,0 +1,73 @@ +#!/usr/bin/env nextflow + +params.fit_dataset_dir = 'gs://cellarium-nexus-file-system-3293a8/pipeline/data-extracts/20260403_cas_pca_model_10x/extract_files' +params.predict_dataset_dir = 'gs://cellarium-nexus-file-system-3293a8/pipeline/data-extracts/20260403_cas_pca_model_10x/extract_files' +params.outdir = 'gs://cellarium-dev-central/workflows/tmp' +params.config_onepass = "${projectDir}/../configs/onepass_mean_var_std.yaml.j2" +params.n_components = 64 +params.n_top_genes = 8000 +params.batch_index_n = 'null' +params.num_workers = 8 +params.prefetch_factor = 4 +params.var_names_key = 'null' +params.accelerator = 'auto' +params.batch_size = 5000 +params.max_cache_size = 4 +params.hvg_method = 'seurat' // or 'kotliar' or 'seurat_v3' +params.seurat_v3_flavor = 'seurat_v3' // or 'seurat_v3_paper' -- only relevant if hvg_method is seurat_v3 +params.use_pflogpf = false // whether to use PFlogPF data normalization +params.zscore_genes = true // whether to z-score genes before PCA +params.total_mrna_umis_key = 'raw_sum' +params.apply_normalize_total = true +params.target_count = 10000 +params.apply_log1p = true +params.sparse_dataloader = true + +// Compute a timestamped output root so successive runs don't overwrite each other. +// Override with --run_outdir gs://... to write to a specific path. +def _run_ts = new java.text.SimpleDateFormat("yyyyMMdd_HHmmss").format(new Date()) +params.run_outdir = params.run_outdir ?: "${params.outdir}/${params.run_label}/${_run_ts}" + +def VALID_HVG_METHODS = ['seurat_v3', 'seurat', 'kotliar'] +if (!(params.hvg_method in VALID_HVG_METHODS)) { + error "Invalid hvg_method '${params.hvg_method}'. Must be one of: ${VALID_HVG_METHODS.join(', ')}" +} + +include { RENDER_PCA_CONFIGS } from './modules/render_configs.nf' +include { ONEPASS_HVG_INCREMENTAL_PCA_PLUS_PREDICTION } from './modules/combo_hvg_pca_predict.nf' + +workflow { + fit_dataset_ch = Channel.value( + params.fit_dataset_dir.startsWith('gs://') + ? params.fit_dataset_dir + : file(params.fit_dataset_dir).toAbsolutePath().toString()) + predict_dataset_ch = Channel.value( + params.predict_dataset_dir.startsWith('gs://') + ? params.predict_dataset_dir + : file(params.predict_dataset_dir).toAbsolutePath().toString()) + + // configs_dir stages the entire configs/ directory into each process work dir, + // making all templates and partials available to render_config.py. + configs_dir_ch = Channel.value(file(params.config_onepass).parent) + + // Run locally first: render preview configs + write provenance record to GCS. + render_out = RENDER_PCA_CONFIGS( + fit_dataset_dir = fit_dataset_ch, + predict_dataset_dir = predict_dataset_ch, + hvg_method = params.hvg_method, + run_name = workflow.runName, + session_id = workflow.sessionId, + configs_dir = configs_dir_ch + ) + + // Dispatch the full pipeline to a single GCP GPU VM. + ONEPASS_HVG_INCREMENTAL_PCA_PLUS_PREDICTION( + fit_dataset_dir = fit_dataset_ch, + predict_dataset_dir = predict_dataset_ch, + hvg_method = params.hvg_method, + onepass_config = render_out.onepass_config, + pca_fit_config = render_out.pca_fit_config, + pca_predict_config = render_out.pca_predict_config, + seurat_v3_hvg_config = render_out.seurat_v3_hvg_config + ) +} diff --git a/cellarium/workflows/nextflow/scvi_workflow.nf b/cellarium/workflows/nextflow/scvi_workflow.nf new file mode 100644 index 0000000..966f448 --- /dev/null +++ b/cellarium/workflows/nextflow/scvi_workflow.nf @@ -0,0 +1,47 @@ +#!/usr/bin/env nextflow + +params.dataset_dir = 'gs://cellarium-nexus-file-system-3293a8/pipeline/data-extracts/20260403_cas_pca_model_10x/extract_files' +params.outdir = 'gs://cellarium-dev-central/workflows/tmp' +params.config_scvi = "${projectDir}/../configs/scvi.yaml.j2" +params.hvg_csv = 'hvg.csv' // from HVG step + +params.num_workers = 8 +params.prefetch_factor = 4 +params.var_names_key = 'null' +params.accelerator = 'auto' +params.batch_size = 5000 +params.max_cache_size = 4 + +params.batch_key = 'assay_dataset_id_donor_id_suspension_type' +params.batch_index_n = 'null' +params.categorical_covariate_keys = 'null' // for scVI, a comma-separated list of keys in the anndata obs +params.n_latent = 128 // for scVI, the dimensionality of the latent space + +def _run_ts = new java.text.SimpleDateFormat("yyyyMMdd_HHmmss").format(new Date()) +params.run_outdir = params.run_outdir ?: "${params.outdir}/${params.run_label}/${_run_ts}" + +include { RENDER_SCVI_CONFIGS } from './modules/render_configs.nf' +include { SCVI } from './modules/scvi.nf' + +workflow { + dataset_ch = Channel.value( + params.dataset_dir.startsWith('gs://') + ? params.dataset_dir + : file(params.dataset_dir).toAbsolutePath().toString()) + hvg_csv_ch = Channel.value(file(params.hvg_csv)) + configs_dir_ch = Channel.value(file(params.config_scvi).parent) + + render_out = RENDER_SCVI_CONFIGS( + dataset_dir = dataset_ch, + hvg_csv_path = params.hvg_csv, + run_name = workflow.runName, + session_id = workflow.sessionId, + configs_dir = configs_dir_ch + ) + + SCVI( + dataset_dir = dataset_ch, + hvg_csv = hvg_csv_ch, + scvi_config = render_out.scvi_config + ) +} diff --git a/cellarium/workflows/nextflow/seurat_v3_hvg_workflow.nf b/cellarium/workflows/nextflow/seurat_v3_hvg_workflow.nf new file mode 100644 index 0000000..bb55015 --- /dev/null +++ b/cellarium/workflows/nextflow/seurat_v3_hvg_workflow.nf @@ -0,0 +1,41 @@ +#!/usr/bin/env nextflow + +// params.dataset_dir = 'gs://cellarium-nexus-file-system-3293a8/pipeline/data-extracts/20260403_cas_pca_model_10x/extract_files' +params.dataset_dir = 'gs://cellarium-nexus-file-system-3293a8/pipeline/data-extracts/czi_human_primary_gr300genes_fullschema/extract_files' +params.outdir = 'gs://cellarium-dev-central/workflows/czi_human_primary_gr300genes_fullschema_hvg_4000' +params.config_hvg = "${projectDir}/../configs/hvg_seurat_v3.yaml.j2" +params.n_top_genes = 4000 +params.seurat_v3_flavor = 'seurat_v3' +params.batch_index_n = 'assay_suspension_type' +params.num_workers = 8 +params.prefetch_factor = 4 +params.var_names_key = 'null' +params.accelerator = 'auto' +params.batch_size = 5000 +params.max_cache_size = 4 + +def _run_ts = new java.text.SimpleDateFormat("yyyyMMdd_HHmmss").format(new Date()) +params.run_outdir = params.run_outdir ?: "${params.outdir}/${params.run_label}/${_run_ts}" + +include { RENDER_SEURAT_HVG_CONFIGS } from './modules/render_configs.nf' +include { SEURAT_V3_HIGHLY_VARIABLE_GENES } from './modules/seurat_v3_hvg.nf' + +workflow { + dataset_ch = Channel.value( + params.dataset_dir.startsWith('gs://') + ? params.dataset_dir + : file(params.dataset_dir).toAbsolutePath().toString()) + configs_dir_ch = Channel.value(file(params.config_hvg).parent) + + render_out = RENDER_SEURAT_HVG_CONFIGS( + dataset_dir = dataset_ch, + run_name = workflow.runName, + session_id = workflow.sessionId, + configs_dir = configs_dir_ch + ) + + SEURAT_V3_HIGHLY_VARIABLE_GENES( + dataset_dir = dataset_ch, + seurat_v3_hvg_config = render_out.seurat_v3_hvg_config + ) +} diff --git a/default_configs/scvi_train.yaml b/default_configs/scvi_train.yaml new file mode 100644 index 0000000..43e8524 --- /dev/null +++ b/default_configs/scvi_train.yaml @@ -0,0 +1,126 @@ +# lightning.pytorch==2.5.2 +seed_everything: true +trainer: + accelerator: gpu + devices: 4 + strategy: + class_path: lightning.pytorch.strategies.DDPStrategy + dict_kwargs: + find_unused_parameters: false + broadcast_buffers: false + callbacks: + - class_path: lightning.pytorch.callbacks.LearningRateMonitor + init_args: + logging_interval: step + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + every_n_epochs: 1 + max_epochs: 20 + log_every_n_steps: 1000 + val_check_interval: 10000 + gradient_clip_algorithm: norm + gradient_clip_val: 1.0 + inference_mode: false + default_root_dir: gs://cellarium-dev-central/workflows/20260514_human_scvi_glyco_validation_runs +model: + cpu_transforms: + - class_path: cellarium.ml.transforms.Filter + init_args: + filter_list: + !FileLoader + file_path: gs://cellarium-dev-central/workflows/20260514_human_scvi_glyco/20260514_hvg_plus_glyco_genes.csv + loader_fn: pandas.read_csv + attr: "feature_id" + convert_fn: pandas.Series.to_list + model: + class_path: cellarium.ml.models.SingleCellVariationalInference + init_args: + var_names_g: null + n_batch: null + n_latent_batch: 64 + batch_embedded: true + batch_representation_sampled: true + n_cats_per_cov: null + kl_warmup_epochs: null + kl_warmup_steps: 200_000 + batch_kl_weight_max: 0.0 + use_batch_norm: both + + # if using a validation data set + cell_type_categories: null + ontology_distance_matrix: + !FileLoader + file_path: https://github.com/obophenotype/cell-ontology/releases/download/v2025-07-30/cl-basic.owl + loader_fn: cellarium.ml.utilities.data.compute_cl_distance_matrix + val_cell_type_classifier_reservoir_size: 200_000 + + encoder: + hidden_layers: + - class_path: torch.nn.Linear + init_args: + out_features: 512 + # - class_path: torch.nn.Linear + # init_args: + # out_features: 512 + final_layer: + class_path: torch.nn.Linear + init_args: {} + decoder: + hidden_layers: + - class_path: cellarium.ml.models.scvi.LinearWithBatch + init_args: + out_features: 512 + label_to_bias_hidden_layers: [] + # - class_path: cellarium.ml.models.scvi.LinearWithBatch + # init_args: + # out_features: 512 + # label_to_bias_hidden_layers: [] + final_layer: + class_path: torch.nn.Linear + init_args: {} + final_additive_bias: false + n_latent: 64 + optim_fn: torch.optim.AdamW + optim_kwargs: + lr: 1e-3 + scheduler_fn: torch.optim.lr_scheduler.LinearLR + scheduler_kwargs: + start_factor: 0.01 + total_iters: 100000 +data: + dadc: + class_path: cellarium.ml.data.DistributedAnnDataCollection + init_args: + filenames: gs://cellarium-nexus-file-system-3293a8/pipeline/data-extracts/czi_human_primary_gr300genes_fullschema/extract_files/extract_{000000..001000}.h5ad + shard_size: 10_000 + last_shard_size: null + max_cache_size: 3 + batch_keys: + x_ng: + attr: X + convert_fn: cellarium.ml.utilities.data.densify + var_names_g: + attr: var + key: feature_id + batch_index_n: + attr: obs + key: assay_dataset_id_donor_id_suspension_type + convert_fn: cellarium.ml.utilities.data.categories_to_codes + # categorical_covariate_index_nd: + # attr: obs + # key: + # - suspension_type + # - assay + # convert_fn: cellarium.ml.utilities.data.categories_to_codes + # only used for validation data + validation_cell_type_index_n: + attr: obs + key: cell_type_ontology_term_id + convert_fn: cellarium.ml.utilities.data.categories_to_codes + batch_size: 4096 + shuffle: true + val_size: 0.02 + num_workers: 2 + prefetch_factor: 2 + persistent_workers: false +# ckpt_path: gs://cellarium-dev-central/workflows/20260401_human_scvi_our_hvg_copy_czi/lightning_logs/version_0/checkpoints/epoch=6-step=645771.ckpt diff --git a/requirements/base.txt b/requirements/base.txt deleted file mode 100644 index 0088d2c..0000000 --- a/requirements/base.txt +++ /dev/null @@ -1,4 +0,0 @@ -google-cloud-aiplatform==1.68.0 -google-cloud-pipeline-components==2.17.0 -kfp==2.7.0 -click==8.1.7 \ No newline at end of file diff --git a/requirements/dev.txt b/requirements/dev.txt deleted file mode 100644 index 2d54ff9..0000000 --- a/requirements/dev.txt +++ /dev/null @@ -1 +0,0 @@ -tox~=4.6 \ No newline at end of file diff --git a/requirements/test.txt b/requirements/test.txt deleted file mode 100644 index e382152..0000000 --- a/requirements/test.txt +++ /dev/null @@ -1,5 +0,0 @@ -pytest~=7.3 -coverage~=4.5 -click~=8.0 -mockito>=1.5.0 -parameterized>=0.9.0 \ No newline at end of file