Finding spatial niches with Nicheformer

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

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 scipy import sparse
from sklearn.preprocessing import normalize
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 / 'nicheformer_input.h5ad'
IMAGES = DATA_DIR / 'spatial_images.h5ad'
EMBEDDING = DATA_DIR / 'nicheformer.npy'
OUTPUT = Path('results')
OUTPUT.mkdir(parents=True, exist_ok=True)
SEED = 42
DEVELOPMENT = ['P10', 'P11', 'P15', 'P16', 'P17', 'P19']
HOLDOUT = ['P24', 'P25']
sc.settings.n_jobs = 8
sc.settings.set_figure_params(dpi=100, frameon=False)
sns.set_context('notebook')
print('Python:', sys.version.split()[0])
print({name: importlib.metadata.version(name) for name in
       ['anndata', 'scanpy', 'numpy', 'pandas', 'scikit-learn']})

data = ad.read_h5ad(INPUT)
assert data.obs_names.is_unique, 'Repeated spot IDs break the alignment with the embeddings.'
print(f'{data.n_obs:,} spots × {data.n_vars:,} genes; '
      f'{data.obs["sample_id"].nunique()} slides from {data.obs["patient_id"].nunique()} subjects.')

embedding = np.load(EMBEDDING, allow_pickle=False)
assert embedding.ndim == 2 and embedding.shape[0] == data.n_obs, 'One embedding row per input spot is required.'
assert np.isfinite(embedding).all() and (np.linalg.norm(embedding, axis=1) > 0).all()
data.obsm['X_nicheformer'] = embedding.astype(np.float32)
data.obs['patient_id'] = data.obs['sample_id'].astype(str).str.split('_').str[0]
data.obs['tissue_group'] = data.obs['tissue_type'].astype(str).map(
    {'tumour': 'Tumour', 'border': 'Background', 'healthy': 'Healthy'})
assert data.obs['tissue_group'].notna().all()
data.obs['split'] = np.select(
    [data.obs['patient_id'].isin(DEVELOPMENT), data.obs['patient_id'].isin(HOLDOUT)],
    ['Discovery', 'Held-out'], default='Healthy reference')
print(f'Attached {embedding.shape[0]:,} spots × {embedding.shape[1]} embedding dimensions.')

symbols = data.var['gene_symbols'].astype(str)
data.var['mt'] = symbols.str.upper().str.startswith('MT-').to_numpy()
sc.pp.calculate_qc_metrics(data, qc_vars=['mt'], percent_top=None, log1p=False, inplace=True)
data.obs['log1p_total_counts'] = np.log1p(data.obs['total_counts'])
slide_table = data.obs.groupby('sample_id', observed=True).agg(
    patient=('patient_id', 'first'), tissue=('tissue_group', 'first'),
    split=('split', 'first'), spots=('patient_id', 'size'),
    median_counts=('total_counts', 'median'), median_genes=('n_genes_by_counts', 'median'))
display(slide_table)
display(pd.crosstab(slide_table['patient'], slide_table['tissue']))
fig, axes = plt.subplots(1, 2, figsize=(11, 4))
sns.boxplot(data=slide_table, x='tissue', y='median_counts', ax=axes[0])
sns.boxplot(data=slide_table, x='tissue', y='median_genes', ax=axes[1])
fig.suptitle('Each observation in these plots is one slide')
plt.tight_layout(); plt.show()

from sklearn.cluster import KMeans

N_STATES, N_NICHES, PER_SLIDE = 8, 6, 1000
HEX_OFFSETS = [(0, 2), (0, -2), (1, 1), (1, -1), (-1, 1), (-1, -1)]


def balanced_sample(mask, obs, per_slide, rng):
    '''Up to `per_slide` spots from every slide in `mask`, so large sections don't dominate.'''
    index = np.flatnonzero(mask)
    slides = obs['sample_id'].astype(str).to_numpy()[index]
    return np.concatenate([rng.choice(group, min(per_slide, len(group)), replace=False)
                           for _, group in pd.Series(index).groupby(slides)])


def visium_adjacency(obs):
    '''Physical neighbours on the Visium hexagonal grid, within each slide.'''
    spots = pd.DataFrame({'slide': obs['sample_id'].astype(str).to_numpy(),
                          'row': obs['array_row'].to_numpy(), 'col': obs['array_col'].to_numpy(),
                          'i': np.arange(len(obs))})
    edges = []
    for d_row, d_col in HEX_OFFSETS:
        shifted = spots.assign(row=spots['row'] + d_row, col=spots['col'] + d_col)
        pairs = shifted.merge(spots, on=['slide', 'row', 'col'], suffixes=('', '_neighbour'))
        edges.append(pairs[['i', 'i_neighbour']].to_numpy())
    edges = np.concatenate(edges)
    return sparse.csr_matrix((np.ones(len(edges)), (edges[:, 0], edges[:, 1])),
                             shape=(len(obs), len(obs)))


def fit_spatial_states_and_niches(embedding, obs, development_patients, random_state):
    rng = np.random.default_rng(random_state)
    x = normalize(embedding)
    fit_index = balanced_sample(obs['patient_id'].isin(development_patients).to_numpy(), obs, PER_SLIDE, rng)

    state_model = KMeans(N_STATES, n_init=10, random_state=random_state).fit(x[fit_index])
    state_dist = state_model.transform(x)
    state_codes = state_dist.argmin(axis=1)
    state_distance = state_dist.min(axis=1)
    state_in_reference = state_distance <= np.quantile(state_distance[fit_index], 0.95)

    adjacency = visium_adjacency(obs)
    one_hot = np.eye(N_STATES)[state_codes]
    with_self = adjacency + sparse.identity(len(obs), format='csr')
    proportions = np.asarray(with_self @ one_hot)
    proportions /= proportions.sum(axis=1, keepdims=True)

    niche_model = KMeans(N_NICHES, n_init=10, random_state=random_state).fit(proportions[fit_index])
    niche_dist = niche_model.transform(proportions)
    niche_distance = niche_dist.min(axis=1)
    niche_covered = niche_distance <= np.quantile(niche_distance[fit_index], 0.95)

    state_names = [f'S{k + 1}' for k in range(N_STATES)]
    niche_names = [f'N{k + 1}' for k in range(N_NICHES)]
    assignments = pd.DataFrame({
        'niche_covered': niche_covered,
        'qc_pass': (obs['total_counts'].to_numpy() >= 500) & (obs['n_genes_by_counts'].to_numpy() >= 200),
        'n_neighbors': np.asarray(adjacency.sum(axis=1)).ravel().astype(int),
    }, index=obs.index)
    coverage = pd.DataFrame({'patient_id': obs['patient_id'].to_numpy(), 'split': obs['split'].to_numpy(),
                             'state_in_reference': state_in_reference, 'niche_covered': niche_covered})
    coverage = coverage.groupby(['patient_id', 'split'], observed=True).agg(
        spots=('state_in_reference', 'size'), state_coverage=('state_in_reference', 'mean'),
        niche_coverage=('niche_covered', 'mean')).reset_index()
    return {
        'states': np.array(state_names)[state_codes],
        'niches': np.array(niche_names)[niche_dist.argmin(axis=1)],
        'state_distance': state_distance, 'state_in_reference': state_in_reference,
        'assignments': assignments, 'neighborhood_proportions': proportions, 'adjacency': adjacency,
        'coverage': coverage,
        'niche_centroid_proportions': pd.DataFrame(niche_model.cluster_centers_, index=niche_names, columns=state_names),
        'models': {'states': state_model, 'niches': niche_model},
    }


discovery = fit_spatial_states_and_niches(
    data.obsm['X_nicheformer'], data.obs, development_patients=DEVELOPMENT, random_state=SEED,
)
data.obs['spot_state'] = pd.Categorical(discovery['states'])
data.obs['spatial_niche'] = pd.Categorical(discovery['niches'])
state_palette = dict(zip(data.obs['spot_state'].cat.categories, sns.color_palette('tab10', 8).as_hex()))
niche_palette = dict(zip(data.obs['spatial_niche'].cat.categories, sns.color_palette('Set2', 6).as_hex()))
data.uns['spot_state_colors'] = list(state_palette.values())
data.uns['spatial_niche_colors'] = list(niche_palette.values())
data.obs['state_distance'] = discovery['state_distance']
data.obs['state_in_reference'] = discovery['state_in_reference']
data.obs['niche_in_reference'] = discovery['assignments']['niche_covered'].to_numpy()
data.obs['qc_pass'] = discovery['assignments']['qc_pass'].to_numpy()
data.obs['spatial_degree'] = discovery['assignments']['n_neighbors'].to_numpy()
data.obs['isolated_spot'] = data.obs['spatial_degree'].eq(0)
data.obsm['state_neighborhood_proportions'] = discovery['neighborhood_proportions']
data.obsp['physical_neighbors'] = discovery['adjacency']
display(discovery['coverage'])
display(discovery['niche_centroid_proportions'].round(3))
display(data.obs.groupby('spot_state', observed=True).agg(
    spots=('sample_id', 'size'), qc_pass_fraction=('qc_pass', 'mean'),
    median_counts=('total_counts', 'median'), median_genes=('n_genes_by_counts', 'median')))
print('QC flag for sensitivity analysis: at least 500 UMIs and 200 detected genes; no rows dropped from the embedding.')
print(f'Isolated spots with self-only composition: {data.obs["isolated_spot"].sum():,}')

sc.pp.neighbors(data, use_rep='X_nicheformer', n_neighbors=15,
                metric='cosine', random_state=SEED, key_added='nicheformer_neighbors')
sc.tl.umap(data, neighbors_key='nicheformer_neighbors', random_state=SEED)
sc.pl.umap(data, color=['tissue_group', 'patient_id', 'log1p_total_counts',
                       'spot_state', 'spatial_niche', 'split'], ncols=3,
           wspace=0.35, size=3, show=False)
plt.savefig(OUTPUT / 'embedding_umap.png', dpi=150, bbox_inches='tight')
plt.show()

programs = {
    'Epithelial': ['EPCAM', 'KRT8', 'KRT18', 'KRT19'],
    'Alveolar': ['SFTPA1', 'SFTPB', 'SFTPC', 'SFTPD'],
    'T_cell': ['CD3D', 'CD3E', 'TRAC'],
    'B_cell': ['MS4A1', 'CD79A', 'CD79B'],
    'Cytotoxic': ['NKG7', 'GNLY', 'PRF1'],
    'Myeloid': ['LST1', 'TYROBP', 'FCER1G'],
    'Fibroblast': ['COL1A1', 'COL1A2', 'DCN'],
    'Endothelial': ['PECAM1', 'VWF', 'KDR'],
    'Hypoxia_like': ['CA9', 'SLC2A1', 'VEGFA', 'LDHA'],
}
marker_report = []
totals = data.obs['total_counts'].to_numpy()
for name, genes in programs.items():
    mask = symbols.isin(genes).to_numpy()
    present = symbols[mask].tolist()
    assert len(present) >= 2, (name, present)
    x = data.X[:, mask].astype(np.float32)
    x = sparse.diags(10000 / np.maximum(totals, 1)) @ x
    x.data = np.log1p(x.data)
    data.obs[f'score_{name}'] = np.asarray(x.mean(axis=1)).ravel()
    marker_report.append({'program': name, 'present': ', '.join(present),
                          'missing': ', '.join(sorted(set(genes) - set(present)))})
display(pd.DataFrame(marker_report))
score_cols = [f'score_{name}' for name in programs]
scores = data.obs[score_cols]
is_development = data.obs['patient_id'].isin(DEVELOPMENT)
z = (scores - scores.loc[is_development].mean()) / scores.loc[is_development].std().clip(lower=1e-6)
profile_frame = z.join(data.obs[['sample_id', 'patient_id', 'tissue_group', 'spot_state', 'spatial_niche', 'split']])
state_profiles = profile_frame.groupby(['sample_id', 'spot_state'], observed=True)[score_cols].mean()
state_profiles = state_profiles.groupby('spot_state', observed=True).mean()
fig, ax = plt.subplots(figsize=(10, 5))
sns.heatmap(state_profiles.rename(columns=lambda c: c.removeprefix('score_')),
            cmap='vlag', center=0, ax=ax)
ax.set_title('Spot states: marker profiles (equal weight per represented slide)')
plt.tight_layout(); plt.savefig(OUTPUT / 'state_markers.png', dpi=150); plt.show()

images = ad.read_h5ad(IMAGES)
assert set(images.uns['spatial']) == set(data.obs['sample_id'].astype(str)), 'Every slide needs its image.'

state_palette = dict(zip(data.obs['spot_state'].cat.categories, sns.color_palette('tab10', 8).as_hex()))
niche_palette = dict(zip(data.obs['spatial_niche'].cat.categories, sns.color_palette('Set2', 6).as_hex()))

def plot_slides(slides, colour, palette=None, vmin=None, vmax=None):
    from matplotlib.collections import PatchCollection
    from matplotlib.patches import Circle, Patch
    from matplotlib.colors import Normalize

    fig, axes = plt.subplots(1, len(slides), figsize=(5 * len(slides), 5), squeeze=False)
    for ax, slide_id in zip(axes.ravel(), slides):
        mask = data.obs['sample_id'].astype(str).eq(slide_id).to_numpy()
        bundle = images.uns['spatial'][slide_id]
        scale = bundle['scalefactors']['tissue_hires_scalef']
        radius = bundle['scalefactors']['spot_diameter_fullres'] * scale / 2
        xy = data.obsm['spatial'][mask] * scale
        patches = [Circle(point, radius) for point in xy]
        ax.imshow(bundle['images']['hires'], origin='upper')
        if palette is not None:
            colours = [palette[str(value)] for value in data.obs.loc[mask, colour]]
            collection = PatchCollection(patches, facecolors=colours, edgecolors='none', alpha=0.8)
        else:
            collection = PatchCollection(patches, cmap='viridis', norm=Normalize(vmin, vmax),
                                         edgecolors='none', alpha=0.8)
            collection.set_array(data.obs.loc[mask, colour].to_numpy(dtype=float))
            fig.colorbar(collection, ax=ax, fraction=0.035, pad=0.01, label=colour.removeprefix('score_'))
        ax.add_collection(collection)
        ax.set_title(f'{slide_id} | {data.obs.loc[mask, "split"].iloc[0]}')
        ax.set_aspect('equal'); ax.axis('off')
    if palette is not None:
        fig.legend(handles=[Patch(color=value, label=key) for key, value in palette.items()],
                   loc='lower center', ncol=len(palette), frameon=False, title=colour)
    fig.tight_layout(rect=(0, 0.12 if palette else 0, 1, 1))
    fig.savefig(OUTPUT / f'spatial_{colour}_{"_".join(slides)}.png', dpi=150, bbox_inches='tight')
    plt.show()

plot_slides(['P16_T1', 'P24_T1', 'P25_T1'], 'spot_state', state_palette)
plot_slides(['P16_T1', 'P24_T1', 'P25_T1'], 'spatial_niche', niche_palette)
plot_slides(['P24_B1', 'P24_T1', 'D2_1'], 'spatial_niche', niche_palette)
program = 'score_Myeloid'  # also try score_Fibroblast, score_T_cell or score_Epithelial
plot_slides(['P16_T1', 'P24_T1', 'P25_T1'], program,
            vmin=0, vmax=float(data.obs[program].quantile(0.98)))

slide_fractions = pd.crosstab(data.obs['sample_id'], data.obs['spatial_niche'], normalize='index')
patient_fractions = slide_fractions.join(slide_table[['patient', 'tissue', 'split']])
patient_fractions = patient_fractions.groupby(['patient', 'tissue', 'split'], observed=True).mean()
display(patient_fractions.round(3))
fig, ax = plt.subplots(figsize=(10, 8))
sns.heatmap(patient_fractions, cmap='Blues', vmin=0, vmax=1, ax=ax)
ax.set_title('Candidate niche fractions; equal weight per slide')
plt.tight_layout(); plt.savefig(OUTPUT / 'niche_patient_fractions.png', dpi=150); plt.show()

niche_profiles = profile_frame.groupby(['split', 'sample_id', 'spatial_niche'], observed=True)[score_cols].mean()
niche_profiles = niche_profiles.groupby(['split', 'spatial_niche'], observed=True).mean()
fig, axes = plt.subplots(1, 2, figsize=(16, 5))
for ax, split in zip(axes, ['Discovery', 'Held-out']):
    sns.heatmap(niche_profiles.loc[split].rename(columns=lambda c: c.removeprefix('score_')),
                cmap='vlag', center=0, vmin=-1, vmax=1, ax=ax)
    ax.set_title(f'{split}: niche marker profiles')
plt.tight_layout(); plt.savefig(OUTPUT / 'transferred_niche_markers.png', dpi=150); plt.show()

tumour_fraction = patient_fractions.xs('Tumour', level='tissue').droplevel('split')
background_fraction = patient_fractions.xs('Background', level='tissue').droplevel('split')
paired_enrichment = tumour_fraction - background_fraction
candidate_niches = paired_enrichment.columns[
    (paired_enrichment.loc[DEVELOPMENT].mean() > 0.05)
    & (paired_enrichment.loc[DEVELOPMENT].gt(0).sum() >= 4)].tolist()
print('Discovery-selected tumour-associated candidates:', candidate_niches or 'None meet the stated rule')
display(paired_enrichment.round(3))
fig, ax = plt.subplots(figsize=(9, 5))
sns.heatmap(paired_enrichment, cmap='vlag', center=0, vmin=-0.6, vmax=0.6, ax=ax)
ax.set_title('Paired tumour minus background niche fraction')
plt.tight_layout(); plt.savefig(OUTPUT / 'paired_niche_enrichment.png', dpi=150); plt.show()

# First average within slide, then within patient; compare held-out patients separately.
by_patient_profile = profile_frame.groupby(
    ['patient_id', 'sample_id', 'spatial_niche'], observed=True)[score_cols].mean()
by_patient_profile = by_patient_profile.groupby(['patient_id', 'spatial_niche'], observed=True).mean()
discovery_profile = by_patient_profile.loc[DEVELOPMENT].groupby('spatial_niche', observed=True).mean()
support = data.obs.groupby(['patient_id', 'spatial_niche'], observed=True).agg(
    n_spots=('sample_id', 'size'), n_slides=('sample_id', 'nunique'),
    state_coverage=('state_in_reference', 'mean'), niche_coverage=('niche_in_reference', 'mean'))
transfer_rows = []
for patient in HOLDOUT:
    for niche in data.obs['spatial_niche'].cat.categories:
        if (patient, niche) not in support.index:
            transfer_rows.append({'patient_id': patient, 'niche': niche, 'supported': False,
                                  'n_spots': 0, 'n_slides': 0, 'marker_correlation': np.nan})
            continue
        stats = support.loc[(patient, niche)]
        supported = stats['n_spots'] >= 50 and stats['n_slides'] >= 2
        reference = discovery_profile.loc[niche].to_numpy()
        query = by_patient_profile.loc[(patient, niche)].to_numpy()
        correlation = np.corrcoef(reference, query)[0, 1] if supported else np.nan
        transfer_rows.append({'patient_id': patient, 'niche': niche,
                              **stats.to_dict(), 'supported': supported,
                              'marker_correlation': correlation,
                              'tumour_minus_background': paired_enrichment.loc[patient, niche],
                              'discovery_candidate': niche in candidate_niches})
niche_transfer = pd.DataFrame(transfer_rows)
display(niche_transfer.round(3))

def spatial_state_enrichment(obs, adjacency, states, n_permutations, random_state):
    '''Share of neighbouring spots with the same state, against shuffles within each slide.'''
    rng = np.random.default_rng(random_state)
    states = np.asarray(states)
    edges = sparse.triu(adjacency, k=1).tocoo()
    slides = obs['sample_id'].astype(str).to_numpy()
    patients = obs['patient_id'].astype(str).to_numpy()
    rows = []
    for slide in np.unique(slides):
        nodes = np.flatnonzero(slides == slide)
        keep = slides[edges.row] == slide
        i, j = edges.row[keep], edges.col[keep]
        if len(i) == 0:
            continue
        observed = float(np.mean(states[i] == states[j]))
        shuffled = states.copy()
        null = []
        for _ in range(n_permutations):
            shuffled[nodes] = rng.permutation(states[nodes])
            null.append(np.mean(shuffled[i] == shuffled[j]))
        rows.append({'sample_id': slide, 'patient_id': patients[nodes[0]], 'n_edges': len(i),
                     'observed_same_state': observed, 'null_mean': float(np.mean(null)),
                     'null_sd': float(np.std(null)), 'excess': observed - float(np.mean(null))})
    return pd.DataFrame(rows)


spatial_results = spatial_state_enrichment(
    data.obs, discovery['adjacency'], discovery['states'],
    n_permutations=100, random_state=SEED,
)
display(spatial_results)
spatial_patient = spatial_results.groupby('patient_id').agg(
    observed=('observed_same_state', 'mean'),
    expected=('null_mean', 'mean'), excess=('excess', 'mean'))
display(spatial_patient)

from sklearn.linear_model import LogisticRegression
from sklearn.metrics import balanced_accuracy_score, roc_auc_score
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

QC_FEATURES = ['log1p_total_counts', 'n_genes_by_counts', 'pct_counts_mt']
TRAIN_PER_SLIDE = 500


def probe():
    return make_pipeline(StandardScaler(), LogisticRegression(max_iter=2000))


def centroid_scores(x_train, y_train, x_test):
    tumour, background = (normalize(x_train[y_train == k].mean(axis=0, keepdims=True)) for k in (1, 0))
    return (x_test @ (tumour - background).T).ravel()


def leave_one_patient_out(embedding, obs, n_permutations, random_state, include_logistic=True):
    '''Tumour versus background, holding out all slides of one cancer patient at a time.'''
    rng = np.random.default_rng(random_state)
    cancer = obs['tissue_group'].isin(['Tumour', 'Background']).to_numpy()
    x = normalize(embedding)
    qc = obs[QC_FEATURES].to_numpy(dtype=float)
    y = obs['tissue_group'].eq('Tumour').to_numpy().astype(int)
    patients = obs['patient_id'].astype(str).to_numpy()
    metrics, predictions, null = [], [], []
    for patient in sorted(np.unique(patients[cancer])):
        test = cancer & (patients == patient)
        train = balanced_sample(cancer & (patients != patient), obs, TRAIN_PER_SLIDE, rng)
        scores = {'embedding_centroid': centroid_scores(x[train], y[train], x[test]),
                  'qc_logistic': probe().fit(qc[train], y[train]).predict_proba(qc[test])[:, 1]}
        if include_logistic:
            scores['embedding_logistic'] = probe().fit(x[train], y[train]).predict_proba(x[test])[:, 1]
        for model, score in scores.items():
            threshold = 0 if model == 'embedding_centroid' else 0.5
            metrics.append({'patient_id': patient, 'model': model, 'n_spots': int(test.sum()),
                            'balanced_accuracy': balanced_accuracy_score(y[test], score > threshold),
                            'auroc': roc_auc_score(y[test], score)})
            predictions.append(pd.DataFrame({'spot_id': obs.index[test], 'patient_id': patient, 'model': model,
                                             'tumour': y[test], 'score': score}))
        # Control: swap tumour and background for a random half of the training patients.
        null_model = 'embedding_logistic' if include_logistic else 'embedding_centroid'
        for permutation in range(n_permutations):
            flipped = rng.choice(np.unique(patients[train]), size=len(np.unique(patients[train])) // 2, replace=False)
            y_null = np.where(np.isin(patients[train], flipped), 1 - y[train], y[train])
            if include_logistic:
                score = probe().fit(x[train], y_null).predict_proba(x[test])[:, 1]
                prediction = score > 0.5
            else:
                score = centroid_scores(x[train], y_null, x[test])
                prediction = score > 0
            null.append({'patient_id': patient, 'model': null_model, 'permutation': permutation,
                         'balanced_accuracy': balanced_accuracy_score(y[test], prediction),
                         'auroc': roc_auc_score(y[test], score)})
    null = pd.DataFrame(null)
    null_summary = null.groupby('model').agg(
        permutations=('permutation', 'nunique'),
        balanced_accuracy_mean=('balanced_accuracy', 'mean'), balanced_accuracy_sd=('balanced_accuracy', 'std'),
        auroc_mean=('auroc', 'mean'), auroc_sd=('auroc', 'std')).reset_index()
    return {'metrics': pd.DataFrame(metrics), 'predictions': pd.concat(predictions, ignore_index=True),
            'null_summary': null_summary}


evaluation = leave_one_patient_out(
    data.obsm['X_nicheformer'], data.obs,
    n_permutations=20, random_state=SEED,
)
patient_metrics = evaluation['metrics']
display(patient_metrics.round(3))
display(patient_metrics.groupby('model')[['balanced_accuracy', 'auroc']].agg(['mean', 'std']))
display(evaluation['null_summary'].round(3))
fig, axes = plt.subplots(1, 2, figsize=(12, 4))
for ax, metric in zip(axes, ['balanced_accuracy', 'auroc']):
    sns.pointplot(data=patient_metrics, x='patient_id', y=metric, hue='model',
                  dodge=0.3, errorbar=None, linestyles='none', ax=ax)
    ax.axhline(0.5, color='grey', linestyle='--')
    ax.set_ylim(0, 1.02)
    ax.set_title(f'Leave-one-patient-out {metric}')
plt.tight_layout(); plt.savefig(OUTPUT / 'heldout_patient_metrics.png', dpi=150); plt.show()

qc_mask = data.obs['qc_pass'].to_numpy()
qc_evaluation = leave_one_patient_out(
    data.obsm['X_nicheformer'][qc_mask], data.obs.loc[qc_mask],
    n_permutations=20, random_state=SEED, include_logistic=True,
)
qc_metrics = qc_evaluation['metrics']
display(qc_metrics.round(3))
display(qc_metrics.groupby('model')[['balanced_accuracy', 'auroc']].mean().round(3))

# This compact report is populated only from computed outputs.
means = patient_metrics.groupby('model')[['balanced_accuracy', 'auroc']].mean()
display(Markdown('**Measured evidence:**'))
display(means.round(3))
heldout_coverage = discovery['coverage'].loc[discovery['coverage']['patient_id'].isin(HOLDOUT)]
display(heldout_coverage)
transfer_check = niche_transfer.assign(passes=lambda f:
    f['supported'] & f['tumour_minus_background'].gt(0)
    & f['marker_correlation'].ge(0.5) & f['state_coverage'].ge(0.8))
recurrent_niches = [niche for niche in candidate_niches
                   if transfer_check.loc[transfer_check['niche'].eq(niche), 'passes'].sum() == len(HOLDOUT)]
centroid_ba = means.loc['embedding_centroid', 'balanced_accuracy']
probe_ba = means.loc['embedding_logistic', 'balanced_accuracy']
qc_ba = means.loc['qc_logistic', 'balanced_accuracy']
patient_range = patient_metrics.query("model == 'embedding_logistic'")['balanced_accuracy']
spatial_excess = float(spatial_patient['excess'].mean())
summary = (
    f'**Answer in this cohort.** Across eight patient-held-out folds, mean balanced accuracy '
    f'was **{centroid_ba:.3f}** for the Nicheformer centroid readout and **{probe_ba:.3f}** '
    f'for the supervised linear probe, versus **{qc_ba:.3f}** for QC alone. '
    f'Probe accuracy ranged from {patient_range.min():.3f} to {patient_range.max():.3f} across patients. '
    f'Across all ten subjects (including the two healthy donors), the mean excess of same-state physical neighbours over shuffled labels was '
    f'**{spatial_excess:.3f}**. Of {len(candidate_niches)} discovery-selected tumour-enriched '
    f'candidate niches, **{len(recurrent_niches)}** passed the stated exploratory transfer checks '
    f'in both held-out patients ({", ".join(recurrent_niches) or "none"}).'
)
interpretation = (
    'The combined results provide exploratory support for recurring tumour-associated spatial niches '
    'in frozen Nicheformer representations, with the patient-specific limitations shown above.'
    if recurrent_niches and probe_ba > qc_ba and spatial_excess > 0 else
    'The combined results do not establish reliable recovery of tumour-associated niches across '
    'both held-out patients under the stated checks; examine which component failed in the tables above.'
)
display(Markdown(summary + '\n\n' + interpretation +
    ' These spot mixtures are not validated malignant or immune-cell identities. A counts-derived '
    'baseline and independent spatial annotations are still needed to establish model-specific added value.'))
(OUTPUT / 'findings.md').write_text(summary + '\n\n' + interpretation + '\n')

import joblib

result = ad.AnnData(X=sparse.csr_matrix((data.n_obs, 0)), obs=data.obs.copy())
for key in ['X_nicheformer', 'X_umap', 'spatial', 'state_neighborhood_proportions']:
    result.obsm[key] = data.obsm[key]
result.obsp['physical_neighbors'] = data.obsp['physical_neighbors']
result.uns['source_h5ad'] = str(INPUT)
result.uns['development_patients'] = DEVELOPMENT
result.uns['heldout_patients'] = HOLDOUT
result.write_h5ad(OUTPUT / 'spatial_nicheformer_analysis.h5ad', compression='gzip')
patient_metrics.to_csv(OUTPUT / 'patient_metrics.csv', index=False)
evaluation['null_summary'].to_csv(OUTPUT / 'patient_label_shuffle_control.csv', index=False)
spatial_results.to_csv(OUTPUT / 'spatial_state_enrichment.csv', index=False)
patient_fractions.to_csv(OUTPUT / 'niche_patient_fractions.csv')
state_profiles.to_csv(OUTPUT / 'state_marker_profiles.csv')
niche_profiles.to_csv(OUTPUT / 'niche_marker_profiles.csv')
paired_enrichment.to_csv(OUTPUT / 'paired_niche_enrichment.csv')
niche_transfer.to_csv(OUTPUT / 'heldout_niche_transfer.csv', index=False)
qc_metrics.to_csv(OUTPUT / 'qc_sensitivity_patient_metrics.csv', index=False)
evaluation['predictions'].to_csv(OUTPUT / 'out_of_fold_predictions.csv', index=False)
discovery['coverage'].to_csv(OUTPUT / 'reference_coverage.csv', index=False)
joblib.dump(discovery['models'], OUTPUT / 'frozen_state_niche_models.joblib')
(OUTPUT / 'analysis_settings.json').write_text(json.dumps({
    'seed': SEED, 'development': DEVELOPMENT, 'holdout': HOLDOUT,
    'state_count': 8, 'niche_count': 6, 'programs': programs,
    'source': str(INPUT),
    'versions': {p: importlib.metadata.version(p) for p in
                 ['anndata', 'scanpy', 'numpy', 'pandas', 'scikit-learn']},
}, indent=2))
print(f'Results and figures saved to {OUTPUT}')