Cross-species label transfer with bio foundation models

from pathlib import Path
import importlib.metadata
import json
import sys

import anndata as ad
import numpy as np
import pandas as pd
import scanpy as sc
import matplotlib.pyplot as plt
import seaborn as sns
from IPython.display import display, Markdown
from IPython import get_ipython

get_ipython().run_line_magic('matplotlib', 'inline')

DATA_DIR = Path('data')
INPUT = DATA_DIR / 'cross_species_balanced_10000.h5ad'
EMBEDDINGS = DATA_DIR / 'embeddings'
OUTPUT = Path('results')
FIGURES = OUTPUT / 'figures'
FIGURES.mkdir(parents=True, exist_ok=True)
SEED = 42
PRIMATES = ['Macaque', 'Mouse lemur']
pd.set_option('display.max_colwidth', 90)
sns.set_theme(style='white', context='notebook')

data = ad.read_h5ad(INPUT)
assert data.obs_names.is_unique, 'Repeated cell IDs invalidate alignment.'
print('Execution environment:', sys.version.split()[0])
print({package: importlib.metadata.version(package)
       for package in ['anndata', 'scanpy', 'numpy', 'pandas', 'scikit-learn', 'scib-metrics']})

display(Markdown(f"**Curation:** {data.uns['curation_description']}"))
display(pd.crosstab(data.obs['source_dataset'], data.obs['species']))

balance = pd.crosstab(data.obs['species'], data.obs['tissue'])
display(Markdown(f'**Cohort:** {data.n_obs:,} cells × {data.n_vars:,} common features; '
                 f'{len(balance)} species and {len(balance.columns)} tissues.'))
print('Every tissue present in every species:', bool((balance > 0).all().all()))
print('Identical per-tissue counts across species:', bool(balance.nunique(axis=0).eq(1).all()))
print('Smallest/largest species-tissue groups:', int(balance.min().min()), int(balance.max().max()))
fig, ax = plt.subplots(figsize=(max(7, len(balance.columns) * 0.7), max(3, len(balance) * 0.55)))
sns.heatmap(balance, annot=True, fmt='d', cmap='Blues', linewidths=0.5,
            cbar_kws={'label': 'Unique cells'}, ax=ax)
ax.set(title='Which species-tissue comparisons does the cohort support?', xlabel='Harmonized tissue', ylabel='Species')
plt.tight_layout()
plt.show()

embeddings = {}
for path in sorted(EMBEDDINGS.glob('*.npy')):
    matrix = np.load(path, allow_pickle=False)
    assert matrix.ndim == 2 and matrix.shape[0] == data.n_obs, f'{path.name}: one row per input cell is required.'
    assert np.isfinite(matrix).all(), f'{path.name}: non-finite values.'
    embeddings[path.stem] = matrix.astype(np.float32)
assert embeddings, f'No embeddings found in {EMBEDDINGS}.'
display(pd.DataFrame([{'model': name, 'cells': m.shape[0], 'dimensions': m.shape[1]}
                      for name, m in embeddings.items()]))

from scib_metrics.benchmark import Benchmarker, BioConservation, BatchCorrection

# Unintegrated baseline: 50 PCs on 3,000 batch-aware highly variable genes.
baseline = data.copy()
sc.pp.normalize_total(baseline, target_sum=1e4)
sc.pp.log1p(baseline)
sc.pp.highly_variable_genes(baseline, n_top_genes=3000, batch_key='species')
sc.tl.pca(baseline, n_comps=50, mask_var='highly_variable', random_state=SEED)
data.obsm['PCA'] = baseline.obsm['X_pca']
for name, matrix in embeddings.items():
    data.obsm[name] = matrix

benchmark = Benchmarker(
    data, batch_key='species', label_key='tissue',
    embedding_obsm_keys=['PCA', *embeddings],
    bio_conservation_metrics=BioConservation(isolated_labels=False),
    batch_correction_metrics=BatchCorrection(),
    pre_integrated_embedding_obsm_key='PCA', n_jobs=4, progress_bar=False,
)
benchmark.benchmark()
scib_results = benchmark.get_results(min_max_scale=False).drop(index='Metric Type').astype(float)
scib_results = scib_results.sort_values('Total', ascending=False).rename_axis('model').reset_index()
scib_results['n_dimensions'] = scib_results['model'].map(lambda m: data.obsm[m].shape[1])
summary_columns = {'model': 'Model', 'Total': 'scIB Total', 'Bio conservation': 'Tissue conservation',
                   'Batch correction': 'Species mixing', 'n_dimensions': 'Dimensions'}
display(scib_results[list(summary_columns)].rename(columns=summary_columns).style.format(
    {'scIB Total': '{:.3f}', 'Tissue conservation': '{:.3f}', 'Species mixing': '{:.3f}'}).hide(axis='index'))

foundation_ranking = scib_results.loc[scib_results['model'].ne('PCA')]
best_model = foundation_ranking.iloc[0]['model']
if len(foundation_ranking) >= 2:
    gap = foundation_ranking.iloc[0]['Total'] - foundation_ranking.iloc[1]['Total']
    if gap < .01:
        display(Markdown(f"**Narrow winner:** `{best_model}` leads `{foundation_ranking.iloc[1]['model']}` "
                         f'by only **{gap:.4f}** scIB points.'))

ranking = scib_results.melt(id_vars='model', value_vars=['Total', 'Bio conservation', 'Batch correction'],
                            var_name='score', value_name='value')
fig, ax = plt.subplots(figsize=(9, max(4, 0.35 * len(scib_results))))
sns.barplot(data=ranking, y='model', x='value', hue='score', ax=ax)
ax.axvline(scib_results.loc[scib_results['model'].eq('PCA'), 'Total'].item(), color='grey', linestyle='--')
ax.set(title='scIB scores; dashed line is the PCA baseline Total', xlabel='', ylabel='')
plt.tight_layout(); plt.savefig(FIGURES / 'scib_ranking.png', dpi=150); plt.show()

display(Markdown(f'**Selected foundation model:** `{best_model}`'))
sc.pp.neighbors(data, use_rep=best_model, n_neighbors=15, metric='cosine', random_state=SEED)
sc.tl.umap(data, random_state=SEED)
sc.pl.umap(data, color=['species', 'tissue'], wspace=0.35, size=8, show=False)
plt.savefig(FIGURES / 'best_model_coembedding.png', dpi=150, bbox_inches='tight')
plt.show()

BROAD_CLASSES = {
    "T / NKT cell": ["T cell", "T cell/NKT cell", "CD4-positive, alpha-beta T cell", "CD8-positive, alpha-beta T cell", "naive thymus-derived CD4-positive, alpha-beta T cell", "naive thymus-derived CD8-positive, alpha-beta T cell", "regulatory T cell", "mature NK T cell"],
    "NK cell": ["natural killer cell"],
    "Innate lymphoid cell": ["innate lymphoid cell"],
    "B cell": ["B cell"],
    "Plasma cell": ["Plasma cell", "Plasma B cell", "plasma cell"],
    "Monocyte": ["Monocyte", "monocyte", "classical monocyte", "intermediate monocyte", "non-classical monocyte"],
    "Macrophage": ["Macrophage", "macrophage", "alveolar macrophage", "tissue-resident macrophage"],
    "Dendritic cell": ["Dendritic cell", "dendritic cell", "conventional dendritic cell", "myeloid dendritic cell"],
    "Plasmacytoid dendritic cell": ["plasmacytoid dendritic cell"],
    "Neutrophil": ["Neutrophil", "neutrophil"],
    "Mast cell / basophil": ["Mast cell", "mast cell", "basophil"],
    "Erythroid cell": ["Erythroid cell", "Erythroblast", "erythrocyte", "erythroid lineage cell", "erythroid progenitor cell"],
    "Hematopoietic precursor": ["hematopoietic precursor cell"],
    "Endothelial cell": ["Endothelial cell", "endothelial cell", "capillary endothelial cell", "cardiac endothelial cell", "endothelial cell of artery", "endothelial cell of hepatic sinusoid", "vein endothelial cell", "Liver endothelial cell 1", "Liver endothelial cell 2", "Digestive tract endothelial cell 2", "Epidermal endothelial cell"],
    "Lymphatic endothelial cell": ["endothelial cell of lymphatic vessel"],
    "Fibroblast / stromal cell": ["Fibroblast", "fibroblast", "fibroblast of cardiac tissue", "fibroblast of lung", "alveolar adventitial fibroblast", "adventitial cell", "Stromal cell", "stromal cell", "Uterus stromal cell", "Tongue stromal cell", "Digestive tract stromal cell", "Adipocyte progenitor/stromal cell", "Mesenchymal cell"],
    "Myofibroblast": ["myofibroblast cell"],
    "Pericyte": ["pericyte"],
    "Smooth muscle cell": ["Smooth muscle cell", "smooth muscle cell", "vascular associated smooth muscle cell", "bronchial smooth muscle cell", "Epidermal smooth muscle cell", "Male reproductive smooth muscle cell", "myometrial cell"],
    "Stellate cell": ["Stellate cell", "hepatic stellate cell", "pancreatic stellate cell"],
    "Mesothelial cell": ["Mesothelial cell", "mesothelial cell", "Viseral adipose mesothelial cell"],
    "Adipocyte": ["Adipocyte", "Adipocyte cell", "adipocyte"],
    "Cardiomyocyte": ["Cardiomyocyte", "cardiac muscle cell", "regular atrial cardiac myocyte", "regular ventricular cardiac myocyte", "ventricular cardiac muscle cell"],
    "Hepatocyte": ["Hepatocyte", "hepatocyte"],
    "Cholangiocyte": ["intrahepatic cholangiocyte", "Bile ductal cell"],
    "Alveolar type 1 cell": ["AT1 cell", "Alveolar epithelial type 1 cell", "pulmonary alveolar type 1 cell"],
    "Alveolar type 2 cell": ["AT2 cell", "Alveolar epithelial type 2 cell", "pulmonary alveolar type 2 cell"],
    "Ciliated cell": ["Ciliated cell", "ciliated epithelial cell", "lung multiciliated epithelial cell", "Respiratory ciliated cell", "Female reproductive ciliated cell"],
    "Club / goblet cell": ["Club cell", "club cell", "respiratory tract goblet cell", "Goblet/serous cell"],
    "Basal cell": ["Basal cell", "basal cell", "Male reproductive basal cell"],
    "Urothelial cell": ["Bladder urothelial cell", "bladder urothelial cell", "Bladder intermediate cell"],
    "Pancreatic acinar cell": ["Pancreas acinar cell", "pancreatic acinar cell", "Acinar cell"],
    "Pancreatic ductal cell": ["pancreatic ductal cell"],
    "Pancreatic endocrine cell": ["type B pancreatic cell", "pancreatic A cell", "Pancreatic endocrine cell"],
    "Uterine epithelial cell": ["epithelial cell of uterus", "Uterus luminal epithelial cell"],
    "Spermatogonium": ["Spermatogonia", "spermatogonium"],
    "Spermatocyte": ["Spermatocyte", "spermatocyte"],
    "Spermatid": ["Spermatid", "spermatid"],
    "Sertoli cell": ["Sertoli cell"],
    "Leydig cell": ["Leydig cell"],
    "Peritubular myoid cell": ["peritubular myoid cell"],
    "Neuron": ["Neuron", "neuron"],
}
LABEL_MAP = {label: broad for broad, labels in BROAD_CLASSES.items() for label in labels}
MIN_HUMAN_CELLS = 10

from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import GridSearchCV, StratifiedKFold
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

data.obs['cell_type_broad'] = data.obs['cell_type'].astype(str).map(LABEL_MAP)
species = data.obs['species'].astype(str).to_numpy()
broad = data.obs['cell_type_broad'].to_numpy(dtype=object)
human_counts = pd.Series(broad[(species == 'Human') & pd.notna(broad)]).value_counts()
training_classes = sorted(human_counts[human_counts >= MIN_HUMAN_CELLS].index)
train = (species == 'Human') & np.isin(broad, training_classes)
target = np.isin(species, PRIMATES)

display(Markdown(f'**Cell-type harmonization:** {data.obs["cell_type"].nunique()} source labels map to '
                 f'{len(BROAD_CLASSES)} broad classes; {len(training_classes)} have at least '
                 f'{MIN_HUMAN_CELLS} human cells and are used for training.'))
label_mapping = (data.obs.groupby(['cell_type', 'cell_type_broad'], observed=True, dropna=False)
                 .size().rename('cells').reset_index())
display(label_mapping)

coverage = pd.DataFrame({'species': species[target], 'tissue': data.obs['tissue'].astype(str).to_numpy()[target],
                         'supported': np.isin(broad[target], training_classes)})
display(Markdown('**Which target cells have a supported human reference label?**'))
display(coverage.groupby(['species', 'tissue'])['supported'].agg(['size', 'mean'])
        .rename(columns={'size': 'cells', 'mean': 'coverage'}).round(3))

x = data.obsm[best_model]
search = GridSearchCV(
    make_pipeline(StandardScaler(), LogisticRegression(class_weight='balanced', max_iter=1500, random_state=SEED)),
    {'logisticregression__C': [0.01, 0.1, 1.0]}, scoring='f1_macro',
    cv=StratifiedKFold(5, shuffle=True, random_state=SEED), n_jobs=-1,
)
search.fit(x[train], broad[train].astype(str))
classifier = search.best_estimator_
display(Markdown(f"**Human validation:** macro F1 {search.best_score_:.3f} at C = "
                 f"{search.best_params_['logisticregression__C']} (5-fold cross-validation on "
                 f'{train.sum():,} human cells; cell-level splits, so this does not test new donors).'))
classes = classifier.classes_
probabilities = classifier.predict_proba(x[target])
predicted = classes[probabilities.argmax(axis=1)]
print(f'Trained on {train.sum():,} human cells; predicted {target.sum():,} primate cells.')

from sklearn.metrics import roc_auc_score, balanced_accuracy_score, f1_score

truth = broad[target]
target_obs = data.obs.loc[target, ['species', 'tissue']].astype(str).reset_index(drop=True)

# Tissue-only baseline: human cell-type frequencies within each tissue, no embeddings.
human_tissue = pd.crosstab(data.obs['tissue'][train], broad[train]).reindex(columns=classes, fill_value=0) + 1
tissue_prior = (human_tissue.div(human_tissue.sum(axis=1), axis=0)
                .reindex(target_obs['tissue']).to_numpy())


def macro_auc(y, scores):
    aucs = [roc_auc_score(y == label, scores[:, k]) for k, label in enumerate(classes)
            if (y == label).any() and (y != label).any()]
    return float(np.mean(aucs)) if aucs else np.nan


def score_scope(mask):
    supported = mask & np.isin(truth, classes)
    y = truth[supported].astype(str)
    return {'n_target': int(mask.sum()), 'n_evaluable': int(supported.sum()),
            'label_coverage': supported.sum() / max(mask.sum(), 1),
            'auc_macro_ovr': macro_auc(y, probabilities[supported]),
            'tissue_prior_auc_macro_ovr': macro_auc(y, tissue_prior[supported]),
            'balanced_accuracy': balanced_accuracy_score(y, predicted[supported]) if supported.any() else np.nan,
            'macro_f1': f1_score(y, predicted[supported], average='macro') if supported.any() else np.nan}


rows = [{'scope': 'all_primates', 'species': 'All primates', 'tissue': 'All',
         **score_scope(np.ones(len(truth), bool))}]
for name in PRIMATES:
    rows.append({'scope': 'species', 'species': name, 'tissue': 'All',
                 **score_scope(target_obs['species'].eq(name).to_numpy())})
    for tissue in sorted(target_obs['tissue'].unique()):
        rows.append({'scope': 'species_tissue', 'species': name, 'tissue': tissue,
                     **score_scope((target_obs['species'].eq(name) & target_obs['tissue'].eq(tissue)).to_numpy())})
transfer_metrics = pd.DataFrame(rows)
metric_names = {'species': 'Species', 'tissue': 'Tissue', 'n_target': 'Target cells', 'n_evaluable': 'Scored cells',
                'label_coverage': 'Coverage', 'auc_macro_ovr': 'Macro AUC', 'tissue_prior_auc_macro_ovr': 'Tissue-only AUC',
                'balanced_accuracy': 'Balanced accuracy', 'macro_f1': 'Macro F1'}
formats = {'Coverage': '{:.1%}', 'Macro AUC': '{:.3f}', 'Tissue-only AUC': '{:.3f}',
           'Balanced accuracy': '{:.3f}', 'Macro F1': '{:.3f}'}
for scopes, title in [(['all_primates', 'species'], 'Overall and species-specific transfer'),
                      (['species_tissue'], 'Within matched tissues')]:
    display(Markdown('**' + title + '**'))
    display(transfer_metrics.loc[transfer_metrics['scope'].isin(scopes), list(metric_names)]
            .rename(columns=metric_names).style.format(formats, na_rep='Undefined').hide(axis='index'))

supported = np.isin(truth, classes)
per_class_auc = pd.DataFrame([
    {'cell_class': label, 'positives': int((truth[supported] == label).sum()),
     'auroc': roc_auc_score(truth[supported] == label, probabilities[supported, k])}
    for k, label in enumerate(classes)
    if (truth[supported] == label).any() and (truth[supported] != label).any()]).sort_values('auroc')
display(Markdown('**Per-class discrimination and support**'))
display(per_class_auc.round(3))

confusion = pd.crosstab(pd.Series(truth[supported], name='Ground truth'),
                        pd.Series(predicted[supported], name='Prediction'), normalize='index')
fig, ax = plt.subplots(figsize=(14, 11))
sns.heatmap(confusion, cmap='Blues', vmin=0, vmax=1, ax=ax)
ax.set_title(f'{best_model}: human-trained predictions on macaque and mouse lemur (row-normalised)')
plt.tight_layout(); plt.savefig(FIGURES / 'label_transfer_confusion.png', dpi=150); plt.show()

predictions = pd.DataFrame({'cell_id': data.obs_names[target], **target_obs,
                            'cell_type_harmonized': truth, 'prediction': predicted})
palette = dict(zip(classes, sns.color_palette('husl', len(classes)).as_hex()))
coords = data.obsm['X_umap'][target]
fig, axes = plt.subplots(len(PRIMATES), 2, figsize=(16, 7 * len(PRIMATES)), squeeze=False)
for row, name in zip(axes, PRIMATES):
    in_species = predictions['species'].eq(name).to_numpy()
    for ax, column in zip(row, ['cell_type_harmonized', 'prediction']):
        values = predictions.loc[in_species, column].to_numpy(dtype=object)
        colours = [palette.get(v, '#d0d0d0') for v in values]
        ax.scatter(coords[in_species, 0], coords[in_species, 1], c=colours, s=6, linewidths=0)
        ax.set_title(f'{name}: {"ground truth" if column == "cell_type_harmonized" else "prediction"}')
        ax.set_aspect('equal'); ax.axis('off')
fig.legend(handles=[plt.Line2D([], [], marker='o', linestyle='', color=c, label=l) for l, c in palette.items()],
           loc='lower center', ncol=4, frameon=False)
fig.tight_layout(rect=(0, 0.12, 1, 1))
plt.savefig(FIGURES / 'label_transfer_umap.png', dpi=150, bbox_inches='tight'); plt.show()
display(Markdown(f'**Predictions:** {len(predictions):,} target cells; grey cells have no supported label.'))
display(predictions.head(12))

pooled = transfer_metrics.query("scope == 'all_primates'").iloc[0]
by_species = transfer_metrics.query("scope == 'species'").set_index('species')
pca_total = scib_results.loc[scib_results['model'].eq('PCA'), 'Total'].item()
best_total = scib_results.loc[scib_results['model'].eq(best_model), 'Total'].item()
findings = (
    f'**{best_model}** scores highest on scIB Total ({best_total:.3f}, against {pca_total:.3f} for the PCA baseline). '
    f"Trained on human cells only, it reaches a macro AUROC of **{pooled['auc_macro_ovr']:.3f}** on "
    f"{int(pooled['n_evaluable']):,} of {int(pooled['n_target']):,} primate cells "
    f"({pooled['label_coverage']:.1%} coverage), against {pooled['tissue_prior_auc_macro_ovr']:.3f} for the tissue-only baseline. "
    + ' '.join(f"{name}: AUROC {by_species.loc[name, 'auc_macro_ovr']:.3f} with "
               f"{by_species.loc[name, 'label_coverage']:.1%} coverage." for name in PRIMATES)
)
display(Markdown(findings))
(OUTPUT / 'findings.md').write_text(findings + '\n')

scib_results.to_csv(OUTPUT / 'scib_metrics.csv', index=False)
transfer_metrics.to_csv(OUTPUT / 'transfer_metrics.csv', index=False)
per_class_auc.to_csv(OUTPUT / 'per_class_auc.csv', index=False)
predictions.to_csv(OUTPUT / 'transfer_predictions.csv.gz', index=False)
label_mapping.to_csv(OUTPUT / 'label_mapping.csv', index=False)
(OUTPUT / 'analysis_settings.json').write_text(json.dumps({
    'seed': SEED, 'best_model': best_model, 'models': list(embeddings), 'training_classes': list(classes),
    'min_human_cells_per_class': MIN_HUMAN_CELLS, 'regularization_C': search.best_params_['logisticregression__C'],
}, indent=2))
print(f'Results and figures saved to {OUTPUT}')