diff --git a/CHANGELOG.md b/CHANGELOG.md index 890c4eb7..7b5cca28 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/scripts/create_component/create_python_metric.sh b/scripts/create_component/create_python_metric.sh index d36bc7a9..23b97b02 100755 --- a/scripts/create_component/create_python_metric.sh +++ b/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 diff --git a/src/metrics/sbee/config.vsh.yaml b/src/metrics/sbee/config.vsh.yaml new file mode 100644 index 00000000..1d44d8ed --- /dev/null +++ b/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] diff --git a/src/metrics/sbee/script.py b/src/metrics/sbee/script.py new file mode 100644 index 00000000..aca2feef --- /dev/null +++ b/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')