Address review comments on bbknn_ts

CHANGELOG.md

@@ -11,7 +11,7 @@
 
 * Added `ARI_batch` and `NMI_batch` to `metrics/clustering_overlap` (PR #68).
 
-* Added `methods/bbknn_ts` component (PR #84).
+* Added `methods/era_combat_bbknn` component (PR #84).
 
 * Added `metrics/cilisi` new metric component (PR #57).
     - ciLISI measures batch mixing in a cell type-aware manner by computing iLISI within each cell type and normalizing

src/methods/era_combat_bbknn/config.vsh.yaml

@@ -0,0 +1,64 @@
+__merge__: ../../api/comp_method.yaml
+
+name: era_combat_bbknn
+label: ERA ComBat + BBKNN
+summary: "A combination of ComBat and BBKNN discovered and implemented by the ERA AI system."
+description: |
+  Standard scRNA-seq preprocessing steps are applied to the raw counts, including
+  total count normalization, log-transformation and scaling of the gene expression
+  data. Batch effect correction is performed using `scanpy.pp.combat` directly on the
+  gene expression matrix (before dimensionality reduction). Dimensionality reduction
+  is then applied using PCA on the ComBat-corrected data, and this PCA embedding is
+  returned as the integrated embedding.
+
+  A custom batch-aware nearest neighbors graph is constructed based on this integrated
+  embedding: for each cell, neighbors are independently identified within its own batch
+  and other batches, up to `n_neighbors_per_batch`. These candidate neighbors are
+  merged, keeping the minimum distance for duplicate entries, and the top
+  `total_k_neighbors` are selected for each cell. Finally, a symmetric sparse distance
+  matrix and a binary connectivities matrix are generated to represent the integrated
+  neighborhood graph.
+
+  The integration code was entirely written by the AI system described in the
+  associated publication, and is included here unmodified.
+references:
+  # Aygün, E., Belyaeva, A., Comanici, G. et al.
+  # An AI system to help scientists write expert-level empirical software.
+  # Nature 654, 909-916 (2026). https://doi.org/10.1038/s41586-026-10658-6
+  doi: 10.1038/s41586-026-10658-6
+links:
+  documentation: https://google-research.github.io/era/
+  repository: https://github.com/google-research/era
+
+info:
+  method_types: [embedding, graph]
+  preferred_normalization: counts
+
+arguments:
+  - name: "--n_pca_components"
+    type: "integer"
+    default: 100
+    description: "Number of PCA components."
+  - name: "--n_neighbors_per_batch"
+    type: "integer"
+    default: 10
+    description: "Number of neighbors to use within each batch."
+  - name: "--total_k_neighbors"
+    type: "integer"
+    default: 50
+    description: "Total number of nearest neighbors to retain for the final graph."
+
+resources:
+  - type: python_script
+    path: script.py
+  - path: /src/utils/read_anndata_partial.py
+
+engines:
+  - type: docker
+    image: openproblems/base_python:1
+
+runners:
+  - type: executable
+  - type: nextflow
+    directives:
+      label: [hightime, highmem, midcpu]

src/methods/bbknn_ts/script.py → src/methods/era_combat_bbknn/script.py

@@ -1,17 +1,20 @@
-"""BBKNN (TS) code for removing batch effects.
+"""ComBat + BBKNN code for removing batch effects.
 
 This is the top-performing batch correction implementation discovered by the AI
-system described in https://arxiv.org/abs/2509.06503.
+system described in https://doi.org/10.1038/s41586-026-10658-6.
 """
 ## VIASH START
 # 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': 'resources_test/.../input.h5ad',
-  'output': 'output.h5ad'
+  'input': 'resources_test/task_batch_integration/cxg_immune_cell_atlas/dataset.h5ad',
+  'output': 'output.h5ad',
+  'n_pca_components': 100,
+  'n_neighbors_per_batch': 10,
+  'total_k_neighbors': 50
 }
 meta = {
-  'name': 'bbknn_ts'
+  'name': 'era_combat_bbknn'
 }
 ## VIASH END
 
@@ -235,14 +238,42 @@ def eliminate_batch_effect_fn(
 from read_anndata_partial import read_anndata
 
 print('Read input', flush=True)
-input_adata = read_anndata(
+adata = read_anndata(
     par['input'],
     X='layers/counts',
     obs='obs',
     var='var',
     uns='uns'
 )
 
-output = eliminate_batch_effect_fn(input_adata, config=config)
-output.uns['method_id'] = 'bbknn_ts'
+print('Run ComBat + BBKNN', flush=True)
+integrated = eliminate_batch_effect_fn(
+    adata,
+    config={
+        'n_pca_components': par['n_pca_components'],
+        'n_neighbors_per_batch': par['n_neighbors_per_batch'],
+        'total_k_neighbors': par['total_k_neighbors']
+    }
+)
+
+print('Store output', flush=True)
+output = ad.AnnData(
+    obs=adata.obs[[]],
+    var=adata.var[[]],
+    obsm={
+        'X_emb': integrated.obsm['X_emb'],
+    },
+    obsp={
+        'connectivities': integrated.obsp['connectivities'],
+        'distances': integrated.obsp['distances'],
+    },
+    uns={
+        'dataset_id': adata.uns['dataset_id'],
+        'normalization_id': adata.uns['normalization_id'],
+        'method_id': meta['name'],
+        'neighbors': integrated.uns['neighbors'],
+    }
+)
+
+print('Write output to file', flush=True)
 output.write_h5ad(par['output'], compression='gzip')

src/workflows/run_benchmark/config.vsh.yaml

@@ -91,8 +91,8 @@ dependencies:
   - name: methods/batchelor_fastmnn
   - name: methods/batchelor_mnn_correct
   - name: methods/bbknn
-  - name: methods/bbknn_ts
   - name: methods/combat
+  - name: methods/era_combat_bbknn
   - name: methods/geneformer
   - name: methods/harmony
   - name: methods/harmonypy

src/workflows/run_benchmark/main.nf

@@ -19,8 +19,8 @@ methods = [
   batchelor_fastmnn,
   batchelor_mnn_correct,
   bbknn,
-  bbknn_ts,
   combat,
+  era_combat_bbknn,
   geneformer,
   harmony,
   harmonypy,