Add gene efficiency correction as a separate pipeline stage

scripts/run_benchmark/run_full_local.sh

@@ -64,9 +64,11 @@ celltype_annotation_methods:
   # - rctd
 expression_correction_methods:
   - no_correction
-  # - gene_efficiency_correction
   # - resolvi_correction
   # - split
+gene_efficiency_correction_methods:
+  - no_correction
+  # - gene_efficiency_correction
 method_parameters_yaml: /tmp/method_params.yaml
 HERE
 

scripts/run_benchmark/run_full_seqeracloud.sh

@@ -56,9 +56,11 @@ celltype_annotation_methods:
   - rctd
 expression_correction_methods:
   - no_correction
-  - gene_efficiency_correction
   - resolvi_correction
   - split
+gene_efficiency_correction_methods:
+  - no_correction
+  - gene_efficiency_correction
 method_parameters_yaml: /tmp/method_params.yaml
 HERE
 

scripts/run_benchmark/run_test_local.sh

@@ -59,9 +59,11 @@ celltype_annotation_methods:
   # - rctd
 expression_correction_methods:
   - no_correction
-  # - gene_efficiency_correction
   # - resolvi_correction
   # - split
+gene_efficiency_correction_methods:
+  - no_correction
+  # - gene_efficiency_correction
 method_parameters_yaml: /tmp/method_params.yaml
 HERE
 

scripts/run_benchmark/run_test_segger_nebius.sh

@@ -44,6 +44,9 @@ celltype_annotation_methods:
   - tacco
 expression_correction_methods:
   - no_correction
+gene_efficiency_correction_methods:
+  - no_correction
+  # - gene_efficiency_correction
 HERE
 
 # Write the parameters to file (input_states version, NOTE: enable `-entry_name auto` for this)

scripts/run_benchmark/run_test_seqeracloud.sh

@@ -55,9 +55,11 @@ celltype_annotation_methods:
   - rctd
 expression_correction_methods:
   - no_correction
-  - gene_efficiency_correction
   - resolvi_correction
   - split
+gene_efficiency_correction_methods:
+  - no_correction
+  - gene_efficiency_correction
 #method_parameters_yaml: /tmp/method_params.yaml
 HERE
 

src/api/comp_method_gene_efficiency_correction.yaml

@@ -0,0 +1,31 @@
+namespace: methods_gene_efficiency_correction
+info:
+  type: method
+  subtype: method_gene_efficiency_correction
+  type_info:
+    label: Gene efficiency correction
+    summary: Correcting per-gene detection efficiency in spatial data
+    description: >-
+      A gene efficiency correction method corrects for per-gene detection efficiency in the
+      (expression-)corrected spatial counts. It runs as a separate stage after expression
+      correction, consuming and producing the corrected-counts format.
+arguments:
+  - name: --input
+    required: true
+    direction: input
+    __merge__: /src/api/file_spatial_corrected_counts.yaml
+  - name: --input_scrnaseq_reference
+    required: false
+    direction: input
+    __merge__: /src/api/file_scrnaseq_reference.yaml
+  - name: --output
+    required: true
+    direction: output
+    __merge__: /src/api/file_spatial_corrected_counts.yaml
+test_resources:
+  - path: /resources_test/task_ist_preprocessing/mouse_brain_combined
+    dest: resources_test/task_ist_preprocessing/mouse_brain_combined
+  - type: python_script
+    path: /common/component_tests/run_and_check_output.py
+  - type: python_script
+    path: /common/component_tests/check_config.py

src/methods_expression_correction/gene_efficiency_correction/config.vsh.yaml → src/methods_gene_efficiency_correction/gene_efficiency_correction/config.vsh.yaml

@@ -1,4 +1,4 @@
-__merge__: /src/api/comp_method_expression_correction.yaml
+__merge__: /src/api/comp_method_gene_efficiency_correction.yaml
 
 name: gene_efficiency_correction
 label: "Gene Efficiency Correction"
@@ -27,7 +27,7 @@ resources:
 engines:
   - type: docker
     image: openproblems/base_python:1
-    __merge__: 
+    __merge__:
       - /src/base/setup_txsim_partial.yaml
     setup:
         - type: python
@@ -38,4 +38,4 @@ runners:
   - type: executable
   - type: nextflow
     directives:
-      label: [ midtime, lowcpu, midmem ]
+      label: [ midtime, lowcpu, midmem ]

src/methods_expression_correction/gene_efficiency_correction/script.py → src/methods_gene_efficiency_correction/gene_efficiency_correction/script.py

@@ -5,29 +5,33 @@
 # Note: this section is auto-generated by viash at runtime. To edit it, make changes
 # in config.vsh.yaml and then run `viash config inject config.vsh.yaml`.
 par = {
-  'input_spatial_with_cell_types': 'resources_test/task_ist_preprocessing/mouse_brain_combined/spatial_with_celltypes.h5ad',
+  'input': 'resources_test/task_ist_preprocessing/mouse_brain_combined/spatial_corrected.h5ad',
   'input_scrnaseq_reference': 'resources_test/task_ist_preprocessing/mouse_brain_combined/scrnaseq_reference.h5ad',
   'celltype_key': 'cell_type',
-  'output': 'spatial_corrected.h5ad',
+  'output': 'spatial_gene_efficiency_corrected.h5ad',
 }
 meta = {
   'name': 'gene_efficiency_correction',
 }
 ## VIASH END
 
-# Optional parameter check: For this specific correction method the par['input_sc'] is required
-assert par['input_scrnaseq_reference'] is not None, 'Single cell input is required for this expr correction method.'
-    
+# Optional parameter check: this correction method requires the scRNA-seq reference
+assert par['input_scrnaseq_reference'] is not None, 'Single cell input is required for this gene efficiency correction method.'
+
 # Read input
 print('Reading input files', flush=True)
-adata_sp = ad.read_h5ad(par['input_spatial_with_cell_types'])
+adata_sp = ad.read_h5ad(par['input'])
 adata_sc = ad.read_h5ad(par['input_scrnaseq_reference'])
-adata_sp.layers["normalized_uncorrected"] = adata_sp.layers["normalized"]
+# Preserve the pre-correction normalized layer. An upstream expression-correction stage
+# already sets it, so only initialise it here if it is missing (keeps it meaning "before
+# any correction").
+if "normalized_uncorrected" not in adata_sp.layers:
+    adata_sp.layers["normalized_uncorrected"] = adata_sp.layers["normalized"]
 adata_sp_reduced = adata_sp[:,adata_sc.var_names].copy()
 obs_sp = adata_sp.obs.copy()
 
 # Apply gene efficiency correction
-print('Annotating cell types', flush=True)
+print('Applying gene efficiency correction', flush=True)
 adata_sp_reduced = tx.preprocessing.gene_efficiency_correction(
     adata_sp_reduced, adata_sc, layer_key='normalized', ct_key=par['celltype_key']
 )
@@ -36,7 +40,7 @@
 print('Concatenating and reordering', flush=True)
 gene_order = adata_sp.var_names.tolist()
 gene_mask = ~adata_sp.var_names.isin(adata_sp_reduced.var_names)
-adata_sp = ad.concat([adata_sp[:,gene_mask], adata_sp_reduced], axis=1, join="outer") 
+adata_sp = ad.concat([adata_sp[:,gene_mask], adata_sp_reduced], axis=1, join="outer")
 adata_sp = adata_sp[:,gene_order]
 adata_sp.obs = obs_sp
 

src/methods_gene_efficiency_correction/no_correction/config.vsh.yaml

@@ -0,0 +1,31 @@
+__merge__: /src/api/comp_method_gene_efficiency_correction.yaml
+
+name: no_correction
+label: "No Correction"
+summary: "Dummy method that does not apply gene efficiency correction"
+description: >-
+  Dummy method that does not apply gene efficiency correction; passes the (expression-)corrected
+  input through unchanged.
+links:
+  documentation: "https://github.com/openproblems-bio/task_ist_preprocessing"
+  repository: "https://github.com/openproblems-bio/task_ist_preprocessing"
+references:
+  doi: "10.1101/2023.02.13.528102"
+
+resources:
+  - type: python_script
+    path: script.py
+
+engines:
+  - type: docker
+    image: openproblems/base_python:1
+    setup:
+      - type: python
+        pypi: [anndata]
+  - type: native
+
+runners:
+  - type: executable
+  - type: nextflow
+    directives:
+      label: [ midtime, lowcpu, lowmem ]

src/methods_gene_efficiency_correction/no_correction/script.py

@@ -0,0 +1,19 @@
+import anndata as ad
+
+## VIASH START
+par = {
+  'input': 'resources_test/task_ist_preprocessing/mouse_brain_combined/spatial_corrected.h5ad',
+  'output': 'spatial_gene_efficiency_corrected.h5ad',
+}
+## VIASH END
+
+# Read input
+print('Reading input files', flush=True)
+adata_sp = ad.read_h5ad(par['input'])
+# Preserve the pre-correction normalized layer only if an upstream stage hasn't set it.
+if "normalized_uncorrected" not in adata_sp.layers:
+    adata_sp.layers["normalized_uncorrected"] = adata_sp.layers["normalized"]
+
+# Write output
+print('Writing output', flush=True)
+adata_sp.write(par['output'])

src/workflows/run_benchmark/config.vsh.yaml

@@ -104,7 +104,13 @@ argument_groups:
           A list of expression correction methods to run.
         type: string
         multiple: true
-        default: "no_correction:gene_efficiency_correction:resolvi_correction:split"
+        default: "no_correction:resolvi_correction:split"
+      - name: "--gene_efficiency_correction_methods"
+        description: |
+          A list of gene efficiency correction methods to run.
+        type: string
+        multiple: true
+        default: "no_correction:gene_efficiency_correction"
   - name: Method parameters
     description: |
       Use these arguments to control the parameter sets that are run for each
@@ -175,9 +181,11 @@ dependencies:
   - name: methods_cell_type_annotation/singler
   - name: methods_cell_type_annotation/rctd
   - name: methods_expression_correction/no_correction
-  - name: methods_expression_correction/gene_efficiency_correction
   - name: methods_expression_correction/resolvi_correction
   - name: methods_expression_correction/split
+  - name: methods_gene_efficiency_correction/no_correction
+    alias: gene_eff_no_correction
+  - name: methods_gene_efficiency_correction/gene_efficiency_correction
   - name: methods_data_aggregation/aggregate_spatial_data
   - name: metrics/similarity
   - name: metrics/quality

src/workflows/run_benchmark/main.nf

@@ -415,11 +415,10 @@ workflow run_wf {
    ****************************************/
   expr_corr_methods = [
     no_correction,
-    gene_efficiency_correction,
     resolvi_correction,
     split
   ]
-  
+
   expr_corr_ch = cta_ch
     | expandChannelWithParameterSets(expr_corr_methods, "corr", "expression_correction_methods")
     | runEach(
@@ -445,6 +444,44 @@ workflow run_wf {
         ]
       }
     )
+
+
+  /****************************************
+   *      GENE EFFICIENCY CORRECTION      *
+   ****************************************/
+  // Runs after expression correction on its corrected counts. The output overwrites the
+  // `output_correction` state key, so `aggregate_spatial_data` and the similarity metric
+  // (which both read `output_correction`) consume the final corrected counts unchanged.
+  gene_eff_methods = [
+    gene_eff_no_correction,
+    gene_efficiency_correction
+  ]
+
+  gene_eff_ch = expr_corr_ch
+    | expandChannelWithParameterSets(gene_eff_methods, "gene_eff", "gene_efficiency_correction_methods")
+    | runEach(
+      components: gene_eff_methods,
+      filter: { id, state, comp ->
+        comp.config.name == state.current_method_id
+      },
+      fromState: { id, state, comp ->
+        [
+          input: state.output_correction,
+          input_scrnaseq_reference: state.input_sc
+        ] + state.current_method_args
+      },
+      toState: { id, out_dict, state, comp ->
+        removeKeys(state, ["current_method_id", "current_method_variant", "current_method_args"]) + [
+          steps: state.steps + [[
+            type: "gene_efficiency_correction",
+            component_id: state.current_method_id,
+            component_variant: state.current_method_variant,
+            run_id: id
+          ]],
+          output_correction: out_dict.output
+        ]
+      }
+    )
     // aggregate spatial data for quality metrics
     | aggregate_spatial_data.run(
       fromState: [
@@ -469,7 +506,7 @@ workflow run_wf {
    *          COMBINE WITH CONTROL        *
    ****************************************/
 
-  expr_corr_and_control_ch = expr_corr_ch.mix(control_ch)
+  expr_corr_and_control_ch = gene_eff_ch.mix(control_ch)
 
 
   /****************************************
@@ -586,7 +623,7 @@ workflow run_wf {
       def methods =
         segm_methods + segm_ass_methods + direct_ass_methods + count_aggr_methods +
         qc_filter_methods + cell_vol_methods + vol_norm_methods + direct_norm_methods +
-        cta_methods + expr_corr_methods
+        cta_methods + expr_corr_methods + gene_eff_methods
       def method_configs = methods.collect{it.config}
       def method_configs_yaml_blob = toYamlBlob(method_configs)
       def method_configs_file = tempFile("method_configs.yaml")