Workflows in nextflow

.gitignore

@@ -1,3 +1,13 @@
+data/
+local_outputs/
+work/
+*.log*
+.nextflow*
+*.html
+*.h5ad
+*.csv
+*.ckpt
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

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 
-```
+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 '[email protected]'
+}
+```
+
+---
+
+## 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.

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

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

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 }}

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

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 }}

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

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 %}

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 %}

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 %}