Add batch_key support for hvg step

Author: cakirb

README.md

@@ -76,6 +76,8 @@ The pipeline requires the following primary inputs, typically configured via com
 *   QC mode (`--qc_mode`): `original`, `multires`, or `combined` (default: `original`).
 *   CellTypist model (`--celltypist_model`): `gut` preset or a path to a custom model.
 *   Integration HVG strategy (`--hvg_strategy`): `cell_nuclei_union` or `global` (default: `cell_nuclei_union`).
+*   Optional batch-aware HVG selection (`--hvg_batch_key`): obs column passed to `scanpy.pp.highly_variable_genes(batch_key=...)`.
+*   HVG span (`--hvg_span`): span passed to Scanpy Seurat v3 HVG selection (default: `auto`, using `1.0` with `--hvg_batch_key` and `0.3` otherwise).
 *   (optional) Cell‑level metadata CSV to add after pooling (`--metadata`).
 *   (optional) CSV with custom QC thresholds for the GMM (`--metrics_csv`).
 *   CSV for manual QC cutoffs when running in `subset` mode (`--limits_csv`) (see `example_cutoffs.csv`).
@@ -455,11 +457,13 @@ This step requires the h5ad object from `finalize_qc` step.
 
 The integration script keeps memory lower by opening the input H5AD in backed mode, reading only the retained cells and selected HVGs for scVI training, then re-reading the original object at the end to write a full-gene integrated output. The output keeps all genes for retained cells and adds the scVI embedding, neighbours, UMAP, and HVG annotations.
 
-Note: integrated outputs may differ from previous scAutoQC releases because HVG selection is now configurable and defaults to `cell_nuclei_union`. Use `--hvg_strategy global` to select one HVG set across all retained cells.
+Note: integrated outputs may differ from previous scAutoQC releases because HVG selection is now configurable and defaults to `cell_nuclei_union`. Use `--hvg_strategy global` to select one HVG set across all retained cells. Set `--hvg_batch_key <obs_column>` to run Scanpy HVG selection with `batch_key=<obs_column>`. With the default `--hvg_span auto`, the pipeline uses `span=1.0` for batch-aware HVG selection and Scanpy's `span=0.3` default otherwise.
 
 The following preprocessing is applied before scVI training:
 * Stringent doublets are removed when `--from_scautoqc true`.
 * Cell-cycle genes are excluded from HVG selection when `--from_scautoqc true`.
+* If `--hvg_batch_key` is set, that obs column is passed to `scanpy.pp.highly_variable_genes` as `batch_key`; leave it empty to keep the existing non-batch-aware HVG behavior.
+* `--hvg_span` is passed to `scanpy.pp.highly_variable_genes` as `span`; `auto` uses `1.0` for batch-aware HVG selection to reduce LOESS near-singularity failures, while retaining Scanpy's `0.3` default when no HVG batch key is used.
 * HVGs are selected with `--hvg_strategy`:
   * `cell_nuclei_union` (default): calculates HVGs separately for observations annotated as `cell` and `nuclei` in `cell_or_nuclei`, then uses the union of both sets. `--n_top_genes` is applied per group, so the final union can contain more than `--n_top_genes` genes.
   * `global`: ignores `cell_or_nuclei` and calculates one HVG set from all retained cells. `--n_top_genes` is the total target number of HVGs.

bin/integration.py

@@ -52,6 +52,29 @@ def str_to_bool(value):
     return str(value).strip().lower() in {"1", "true", "t", "yes", "y"}
 
 
+def optional_obs_key(value):
+    """Return a non-empty obs key or None for unset Nextflow/CLI values."""
+    if value is None:
+        return None
+    value = str(value).strip()
+    return value or None
+
+
+def resolve_hvg_span(value, hvg_batch_key):
+    """Resolve an explicit HVG span or choose a batch-aware default."""
+    if value is None or str(value).strip().lower() in {"", "auto"}:
+        return 1.0 if hvg_batch_key is not None else 0.3
+
+    try:
+        hvg_span = float(value)
+    except ValueError as error:
+        raise ValueError("--hvg_span must be 'auto' or a number in (0, 1].") from error
+
+    if not 0 < hvg_span <= 1:
+        raise ValueError("--hvg_span must be greater than 0 and less than or equal to 1.")
+    return hvg_span
+
+
 def close_backed(adata):
     if getattr(adata, "isbacked", False):
         adata.file.close()
@@ -240,7 +263,9 @@ def get_training_masks(source_ad, from_scautoqc):
     return keep_cells, keep_genes
 
 
-def select_cell_nuclei_hvg_union(source_ad, keep_cells, keep_genes, n_top_genes):
+def select_cell_nuclei_hvg_union(
+    source_ad, keep_cells, keep_genes, n_top_genes, hvg_batch_key, hvg_span
+):
     """Select union HVGs and retain per-modality HVG annotations."""
     group_key = "cell_or_nuclei"
     if group_key not in source_ad.obs:
@@ -281,6 +306,8 @@ def select_cell_nuclei_hvg_union(source_ad, keep_cells, keep_genes, n_top_genes)
                 modality_ad,
                 flavor="seurat_v3",
                 n_top_genes=min(n_top_genes, modality_ad.n_vars),
+                batch_key=hvg_batch_key,
+                span=hvg_span,
                 subset=False,
                 inplace=False,
             )
@@ -317,7 +344,9 @@ def select_cell_nuclei_hvg_union(source_ad, keep_cells, keep_genes, n_top_genes)
     return hvg_genes, hvg_info
 
 
-def select_global_hvgs(source_ad, keep_cells, keep_genes, n_top_genes):
+def select_global_hvgs(
+    source_ad, keep_cells, keep_genes, n_top_genes, hvg_batch_key, hvg_span
+):
     """Select HVGs from all retained cells regardless of cell/nuclei status."""
     log("Global HVG selection will use all retained cells.")
     with log_step("Read matrix for global HVG selection"):
@@ -334,6 +363,8 @@ def select_global_hvgs(source_ad, keep_cells, keep_genes, n_top_genes):
             hvg_ad,
             flavor="seurat_v3",
             n_top_genes=min(n_top_genes, hvg_ad.n_vars),
+            batch_key=hvg_batch_key,
+            span=hvg_span,
             subset=False,
             inplace=False,
         )
@@ -354,23 +385,52 @@ def select_global_hvgs(source_ad, keep_cells, keep_genes, n_top_genes):
     return hvg_genes, hvg_info
 
 
-def select_hvgs(source_ad, keep_cells, keep_genes, n_top_genes, hvg_strategy):
+def select_hvgs(
+    source_ad,
+    keep_cells,
+    keep_genes,
+    n_top_genes,
+    hvg_strategy,
+    hvg_batch_key,
+    hvg_span,
+):
+    if hvg_batch_key is not None and hvg_batch_key not in source_ad.obs:
+        raise KeyError(f"Expected HVG batch key {hvg_batch_key!r} in adata.obs.")
+
+    if hvg_batch_key is None:
+        log("HVG selection will not use a batch key.")
+    else:
+        log(f"HVG selection will use batch key: {hvg_batch_key}")
+    log(f"HVG selection will use span: {hvg_span}")
+
     if hvg_strategy == "cell_nuclei_union":
-        return select_cell_nuclei_hvg_union(source_ad, keep_cells, keep_genes, n_top_genes)
+        return select_cell_nuclei_hvg_union(
+            source_ad, keep_cells, keep_genes, n_top_genes, hvg_batch_key, hvg_span
+        )
     if hvg_strategy == "global":
-        return select_global_hvgs(source_ad, keep_cells, keep_genes, n_top_genes)
+        return select_global_hvgs(
+            source_ad, keep_cells, keep_genes, n_top_genes, hvg_batch_key, hvg_span
+        )
     raise ValueError(f"Unknown HVG strategy: {hvg_strategy}")
 
 
-def read_training_object(obj_path, from_scautoqc, n_top_genes, hvg_strategy):
+def read_training_object(
+    obj_path, from_scautoqc, n_top_genes, hvg_strategy, hvg_batch_key, hvg_span
+):
     """Load only retained cells and selected HVGs for scVI training."""
     with log_step(f"Open input object backed: {obj_path}"):
         source_ad = sc.read_h5ad(obj_path, backed="r")
     try:
         keep_cells, keep_genes = get_training_masks(source_ad, from_scautoqc)
         with log_step(f"Select HVGs using strategy: {hvg_strategy}"):
             hvg_genes, hvg_info = select_hvgs(
-                source_ad, keep_cells, keep_genes, n_top_genes, hvg_strategy
+                source_ad,
+                keep_cells,
+                keep_genes,
+                n_top_genes,
+                hvg_strategy,
+                hvg_batch_key,
+                hvg_span,
             )
         with log_step("Read final SCVI training matrix"):
             train_ad = read_minimal_view(source_ad, keep_cells, hvg_genes)
@@ -462,6 +522,22 @@ def main():
             "cell_or_nuclei and calculates one HVG set from all retained cells."
         ),
     )
+    parser.add_argument(
+        "--hvg_batch_key",
+        default="",
+        help=(
+            "optional obs column to pass as batch_key to "
+            "scanpy.pp.highly_variable_genes"
+        ),
+    )
+    parser.add_argument(
+        "--hvg_span",
+        default="auto",
+        help=(
+            "span to pass to scanpy.pp.highly_variable_genes; "
+            "auto uses 1.0 with --hvg_batch_key and 0.3 otherwise [default: auto]"
+        ),
+    )
     parser.add_argument(
         "--verbose",
         action="store_true",
@@ -478,7 +554,8 @@ def main():
     log(
         "integration.py started with "
         f"obj={args.obj}, batch={args.batch}, n_top_genes={args.n_top_genes}, "
-        f"hvg_strategy={args.hvg_strategy}, from_scautoqc={args.from_scautoqc}"
+        f"hvg_strategy={args.hvg_strategy}, hvg_batch_key={args.hvg_batch_key}, "
+        f"hvg_span={args.hvg_span}, from_scautoqc={args.from_scautoqc}"
     )
 
     arches_params = dict(
@@ -490,10 +567,17 @@ def main():
     )
 
     from_scautoqc = str_to_bool(args.from_scautoqc)
+    hvg_batch_key = optional_obs_key(args.hvg_batch_key)
+    hvg_span = resolve_hvg_span(args.hvg_span, hvg_batch_key)
 
     with log_step("Prepare training object"):
         train_ad, hvg_info = read_training_object(
-            args.obj, from_scautoqc, args.n_top_genes, args.hvg_strategy
+            args.obj,
+            from_scautoqc,
+            args.n_top_genes,
+            args.hvg_strategy,
+            hvg_batch_key,
+            hvg_span,
         )
     gc.collect()
 

main.nf

@@ -199,8 +199,9 @@ process integrate {
   path("*.pkl")
 
   script:
+  def hvgBatchArg = params.hvg_batch_key ? "--hvg_batch_key ${params.hvg_batch_key}" : ""
   """
-  python ${projectDir}/bin/integration.py --obj ${qc2_out} --batch ${params.batch_key} --n_top_genes ${params.n_top_genes} --hvg_strategy ${params.hvg_strategy} --from_scautoqc ${params.from_scautoqc}
+  python ${projectDir}/bin/integration.py --obj ${qc2_out} --batch ${params.batch_key} --n_top_genes ${params.n_top_genes} --hvg_strategy ${params.hvg_strategy} ${hvgBatchArg} --hvg_span ${params.hvg_span} --from_scautoqc ${params.from_scautoqc}
   """
 }
 

nextflow.config

@@ -78,6 +78,8 @@ params {
   from_scautoqc = true
   n_top_genes = 5000
   hvg_strategy = "cell_nuclei_union" // Options: cell_nuclei_union, global
+  hvg_batch_key = "" // Optional obs column for batch-aware HVG selection
+  hvg_span = "auto" // Uses 1.0 with hvg_batch_key and 0.3 otherwise
   batch_key = ""
 
   // Container image used by portable profiles.