add new metric sBEE

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## New functionality
 
+* Added `metrics/sbee` component (PR #97)
 * Added `metrics/kbet_pg` and `metrics/kbet_pg_label` components (PR #52).
 * Added `methods/stacas` new method (PR #58).
     - Add non-supervised version of STACAS tool for integration of single-cell transcriptomics data. This functionality enables correction of batch effects while preserving biological variability without requiring prior cell type annotations.

scripts/create_component/create_python_metric.sh

@@ -3,6 +3,6 @@
 set -e
 
 common/scripts/create_component \
-  --name my_python_metric \
+  --name sbee \
   --language python \
   --type metric

src/metrics/sbee/config.vsh.yaml

@@ -0,0 +1,76 @@
+# The API specifies which type of component this is.
+# It contains specifications for:
+#   - The input/output files
+#   - Common parameters
+#   - A unit test
+__merge__: ../../api/comp_metric.yaml
+
+# A unique identifier for your component (required).
+# Can contain only lowercase letters or underscores.
+name: sbee
+
+
+
+# Metadata for your component
+info:
+  metrics:
+      # A unique identifier for your metric (required).
+      # Can contain only lowercase letters or underscores.
+    - name: sbee
+      # A relatively short label, used when rendering visualisarions (required)
+      label: sBEE
+      # A one sentence summary of how this metric works (required). Used when 
+      # rendering summary tables.
+      summary: "A unified metric that jointly evaluates cross-batch distance relationships and local neighborhood batch composition."
+      # A multi-line description of how this component works (required). Used
+      # when rendering reference documentation.
+      description: |
+        sBEE (single-cell Batch Effect Evaluator) is a per-cell batch integration metric that produces scores in [0, 1], where higher values indicate better batch mixing.
+        It combines two components via their harmonic mean.
+
+        Distance component checks whether a cell is geometrically closer to same-type cells from other batches than to same-type cells from its own batch.
+        When the ratio of intra-batch to inter-batch distance is 1 or above, the component is set to 1 (no penalty).
+        When the ratio is below 1, a penalty is applied that grows with the degree of separation.
+        Cells whose cell type appears in only one batch are assigned a perfect score, as batch correction is not applicable there.
+
+        Neighborhood composition component checks whether the local batch composition around a cell matches the global batch distribution for that cell type.
+        It compares batch proportions among same-type cells in the k-nearest neighborhood against global proportions using Jensen-Shannon distance. Smaller divergence gives a higher score.
+
+        The two components are combined via harmonic mean. A low score on either component pulls the overall score down.
+        Cell-type scores are computed by macro-averaging across batches so that each batch contributes equally regardless of its size.
+
+      references:
+        doi: 
+          - 10.64898/2026.04.22.720135
+
+      links:
+        # URL to the documentation for this metric (required).
+        documentation: https://github.com/tastanlab/sBEE
+        # URL to the code repository for this metric (required).
+        repository: https://github.com/tastanlab/sBEE
+      # The minimum possible value for this metric (required)
+      min: 0
+      # The maximum possible value for this metric (required)
+      max: 1
+      # Whether a higher value represents a 'better' solution (required)
+      maximize: true
+
+# Resources required to run the component
+resources:
+  # The script of your component (required)
+  - type: python_script
+    path: script.py
+  - path: /src/utils/read_anndata_partial.py
+
+engines:
+  # Specifications for the Docker image for this component.
+  - type: docker
+    image: openproblems/base_python:1.0.0
+
+runners:
+  # This platform allows running the component natively
+  - type: executable
+  # Allows turning the component into a Nextflow module / pipeline.
+  - type: nextflow
+    directives:
+      label: [midtime,midmem,midcpu]

src/metrics/sbee/script.py

@@ -0,0 +1,195 @@
+import anndata as ad
+import sys
+import numpy as np
+import pandas as pd
+from sklearn.neighbors import NearestNeighbors
+from scipy.spatial.distance import jensenshannon, cdist
+
+## 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_integrated': 'resources_test/.../integrated.h5ad',
+  'input_solution': 'resources_test/.../solution.h5ad',
+  'output': 'output.h5ad'
+}
+meta = {
+  'name': 'sbee'
+}
+## VIASH END
+
+sys.path.append(meta["resources_dir"])
+from read_anndata_partial import read_anndata
+
+
+def knn(df, k=90, metric="euclidean", include_self=True):
+    k = k if include_self else k + 1
+    X = df.to_numpy(dtype=np.float32, copy=False)
+    nn = NearestNeighbors(n_neighbors=k, metric=metric, algorithm='auto', n_jobs=-1).fit(X)
+    distances, indices = nn.kneighbors(X, return_distance=True)
+    if not include_self:
+        distances = distances[:, 1:]
+        indices = indices[:, 1:]
+    return indices, distances
+
+
+def js(p, q, epsilon=1e-10):
+    p = np.clip(p, a_min=epsilon, a_max=None)
+    q = np.clip(q, a_min=epsilon, a_max=None)
+    p /= p.sum()
+    q /= q.sum()
+    return jensenshannon(p, q, base=2.0)
+
+
+def build_distribution(adata, dknn_df, celltypes_df, batches_df, dist_type="count"):
+    k = len(dknn_df.columns)
+    celltype_counts = adata.obs['cell_type'].value_counts()
+
+    batches_neighbors_df = pd.DataFrame(index=dknn_df.index, columns=list(range(k)))
+    celltypes_neighbors_df = pd.DataFrame(index=dknn_df.index, columns=list(range(k)))
+
+    batches = adata[batches_neighbors_df.index].obs["batch"].values
+    celltypes = adata[celltypes_neighbors_df.index].obs["cell_type"].values
+
+    neighbor_indices = dknn_df.to_numpy()
+    batches_neighbors_df = pd.DataFrame(batches[neighbor_indices], index=dknn_df.index, columns=list(range(k)))
+    celltypes_neighbors_df = pd.DataFrame(celltypes[neighbor_indices], index=dknn_df.index, columns=list(range(k)))
+
+    unique_batches = np.unique(batches)
+    dist = pd.DataFrame(0, index=batches_df.index, columns=unique_batches, dtype=float)
+
+    if dist_type == "global":
+        for b in unique_batches:
+            for idx, cell_id in enumerate(celltypes_neighbors_df.index):
+                cell_type = celltypes_df.iloc[idx]['cell_type']
+                dist.loc[cell_id, b] = adata[(adata.obs.batch == b) & (adata.obs.cell_type == cell_type)].shape[0]
+        return dist
+
+    for b in unique_batches:
+        for idx, cell_id in enumerate(celltypes_neighbors_df.index):
+            cell_type = celltypes_df.iloc[idx]['cell_type']
+            k_adjusted = min(k, celltype_counts[cell_type])
+            neigh_celltypes = np.array(celltypes_neighbors_df.loc[cell_id].iloc[:k_adjusted])
+            neigh_batch_labels = np.array(batches_neighbors_df.loc[cell_id].iloc[:k_adjusted])
+            same_type_batch_mask = (neigh_celltypes == cell_type) & (neigh_batch_labels == b)
+            dist.loc[cell_id, b] = same_type_batch_mask.sum()
+
+    return dist
+
+
+def js_dist(scores, local_dist, global_dist):
+    scores["JS Dist"] = 0.
+    for cell_id in scores.index:
+        loc = np.array(local_dist.loc[cell_id])
+        glob = np.array(global_dist.loc[cell_id])
+        scores.loc[cell_id, "JS Dist"] = js(loc, glob)
+    return scores
+
+
+def compute_intra_inter_distances(adata, batch_key="batch", label_key="cell_type", agg="median"):
+    agg_fn = np.median if agg == "median" else np.mean
+    nan_agg_fn = np.nanmedian if agg == "median" else np.nanmean
+    col_suffix = agg
+    X = adata.obsm["X_emb"]
+    obs_df = adata.obs[[batch_key, label_key]].copy()
+    intra_col = f"intra_{col_suffix}"
+    inter_col = f"inter_{col_suffix}"
+    obs_df[intra_col] = 0.0
+    obs_df[inter_col] = 0.0
+
+    celltype_groups = obs_df.groupby(label_key).groups
+    batch_celltype_groups = obs_df.groupby([label_key, batch_key]).groups
+
+    for (ct, _), group_idx in batch_celltype_groups.items():
+        pos = adata.obs_names.get_indexer(group_idx)
+        X_group = X[pos]
+
+        # Intra: same (cell_type, batch), exclude self via NaN diagonal
+        if len(pos) > 1:
+            D_intra = cdist(X_group, X_group, metric="euclidean")
+            np.fill_diagonal(D_intra, np.nan)
+            obs_df.loc[group_idx, intra_col] = nan_agg_fn(D_intra, axis=1)
+
+        # Inter: same cell_type, different batch
+        other_idx = celltype_groups[ct].difference(group_idx)
+        if len(other_idx) > 0:
+            X_other = X[adata.obs_names.get_indexer(other_idx)]
+            obs_df.loc[group_idx, inter_col] = agg_fn(
+                cdist(X_group, X_other, metric="euclidean"), axis=1
+            )
+        else:
+            obs_df.loc[group_idx, inter_col] = obs_df.loc[group_idx, intra_col]
+
+    obs_df["intra_inter_ratio"] = np.where(
+        (obs_df[intra_col] == 0) | (obs_df[inter_col] == 0),
+        np.nan,
+        obs_df[intra_col] / obs_df[inter_col]
+    )
+    return obs_df
+
+
+def sbee(adata, js_dist_key="JS Dist", ratio_key="intra_inter_ratio", sensitivity=0.15):
+
+    print('Building kNN graph', flush=True)
+    emb_df = pd.DataFrame(adata.obsm['X_emb'], index=adata.obs_names)
+    indices, _ = knn(emb_df, k=90, include_self=False)
+    dknn_df = pd.DataFrame(indices, index=adata.obs_names)
+
+    print('Building distributions', flush=True)
+    celltypes_df = adata.obs[['cell_type']]
+    batches_df = adata.obs[['batch']]
+    local_dist = build_distribution(adata, dknn_df, celltypes_df, batches_df, dist_type='count')
+    global_dist = build_distribution(adata, dknn_df, celltypes_df, batches_df, dist_type='global')
+
+    print('Computing JS distances', flush=True)
+    scores = pd.DataFrame(index=adata.obs_names)
+    scores['cell_type'] = adata.obs['cell_type'].values
+    scores['batch'] = adata.obs['batch'].values
+    scores = js_dist(scores, local_dist, global_dist)
+
+    print('Computing intra/inter distances', flush=True)
+    ratio_df = compute_intra_inter_distances(adata, batch_key='batch', label_key='cell_type')
+    scores['intra_inter_ratio'] = ratio_df['intra_inter_ratio'].values
+
+    js_score = 1 - scores[js_dist_key]
+
+    effective_ratio = scores[ratio_key].clip(upper=1)
+    ratio_score = np.exp(-np.abs(1 - effective_ratio) / sensitivity)
+
+    scores["js_part"] = js_score
+    scores["ratio_part"] = ratio_score
+    scores['sBEE'] = 2 * js_score * ratio_score / (js_score + ratio_score)
+    return scores
+
+
+print('Reading input files', flush=True)
+adata = read_anndata(par['input_integrated'], obs='obs', obsm='obsm', uns='uns')
+adata.obs = read_anndata(par['input_solution'], obs='obs').obs
+adata.uns |= read_anndata(par['input_solution'], uns='uns').uns
+
+
+print('sBEE score', flush=True)
+scores = sbee(adata)
+
+# print(scores.groupby(['cell_type', 'batch'])['sBEE'].mean())
+
+# macro-average across batches per cell type, then macro-average across cell types
+score = (
+    scores.groupby(['cell_type', 'batch'])['sBEE'].mean()
+    .groupby(level='cell_type').mean()
+    .mean()
+)
+
+print('Create output AnnData object', flush=True)
+output = ad.AnnData(
+    uns={
+        'dataset_id': adata.uns['dataset_id'],
+        'normalization_id': adata.uns['normalization_id'],
+        'method_id': adata.uns['method_id'],
+        'metric_ids': [meta['name']],
+        'metric_values': [score]
+    }
+)
+
+print('Write output AnnData to file', flush=True)
+output.write_h5ad(par['output'], compression='gzip')