Downstream clustering and trajectory inference for stGP Microglia
This notebook uses the saved stGP outputs for the mouse brain MERFISH microglia analysis. It focuses on one mouse_id at a time, clusters cells in the stGP spatial embedding with a symmetric KNN graph and spectral clustering.
1. Setup and load stGP outputs
[1]:
import shutil, warnings
from pathlib import Path
import h5py
import numpy as np
import pandas as pd
import scanpy as sc
import scipy.sparse as sp
import matplotlib.pyplot as plt
from sklearn.neighbors import NearestNeighbors
from sklearn.cluster import SpectralClustering
from pyslingshot import Slingshot
warnings.filterwarnings('ignore', category=FutureWarning)
plt.rcParams['figure.dpi'] = 120
ROOT = Path.cwd()
CELLTYPE = 'Microglia'
STGP_DIR = ROOT / 'Results' / 'stgp' / CELLTYPE
ADATA_PATH = STGP_DIR / 'adata_with_scores.h5ad'
OUT_DIR = STGP_DIR / 'downstream_cluster'
OUT_DIR.mkdir(parents=True, exist_ok=True)
# Spatial coordinates are section-local, so run downstream trajectory on one mouse/section.
# Set MOUSE_ID = None to use all microglia, but spatial overlays across mice are not recommended.
MOUSE_ID = '97'
TRUTH_KEY = 'subregion'
ROOT_TRUTH_LABEL = 'CC/ACO'
START_CLUSTER = None # None selects the stGP domain most enriched for ROOT_TRUTH_LABEL.
N_EPOCHS = 10
RANDOM_STATE = 1234
def _prepare_nichescope_h5ad(path: Path, cache_dir: Path) -> Path:
"""Return a NicheScope-readable copy when the stGP h5ad contains null encodings."""
fixed = cache_dir / f'{path.stem}_niche_scope_compat.h5ad'
if fixed.exists() and fixed.stat().st_mtime >= path.stat().st_mtime:
return fixed
shutil.copy2(path, fixed)
with h5py.File(fixed, 'a') as h5:
# NicheScope's anndata cannot read the null-encoded Scanpy log1p base.
if 'uns/log1p/base' in h5:
del h5['uns/log1p/base']
return fixed
def read_h5ad_compat(path: Path, cache_dir: Path):
"""Read saved stGP AnnData from the NicheScope environment."""
fixed = _prepare_nichescope_h5ad(path, cache_dir)
print(f'Reading NicheScope compatibility copy: {fixed}')
return sc.read_h5ad(fixed)
if not ADATA_PATH.exists():
raise FileNotFoundError(f'Missing saved stGP AnnData: {ADATA_PATH}')
adata_full = read_h5ad_compat(ADATA_PATH, OUT_DIR)
if 'X_stgp_spatial' not in adata_full.obsm:
raise KeyError('Expected adata.obsm["X_stgp_spatial"] from the stGP Microglia output.')
if TRUTH_KEY not in adata_full.obs:
raise KeyError(f'Expected adata.obs[{TRUTH_KEY!r}] for clustering comparison.')
if MOUSE_ID is None:
adata = adata_full.copy()
label_suffix = 'all_microglia'
else:
mask = adata_full.obs['mouse_id'].astype(str).to_numpy() == str(MOUSE_ID)
if not mask.any():
available = sorted(adata_full.obs['mouse_id'].astype(str).unique())
raise ValueError(f'MOUSE_ID={MOUSE_ID!r} not found. Available mouse_id values: {available}')
adata = adata_full[mask].copy()
label_suffix = f'mouse_{MOUSE_ID}'
print(adata)
print(f'Analysis selection: {label_suffix}; n_obs={adata.n_obs}, n_vars={adata.n_vars}')
print(adata.obs[['mouse_id', 'slide_id', 'age', 'batch']].drop_duplicates().to_string(index=False))
display(adata_full.obs.groupby(['mouse_id', 'age', 'slide_id']).size().reset_index(name='n_cells').sort_values('n_cells', ascending=False).head(20))
/home/byual/anaconda3/envs/NicheScope/lib/python3.9/site-packages/pyslingshot/slingshot.py:11: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
from tqdm.autonotebook import tqdm
Reading NicheScope compatibility copy: /import/home4/byual/stGP-0529/RealData_MouseBrainMERFISH/Results/stgp/Microglia/downstream_cluster/adata_with_scores_niche_scope_compat.h5ad
AnnData object with n_obs × n_vars = 3626 × 220
obs: 'volume', 'center_x', 'center_y', 'min_x', 'min_y', 'max_x', 'max_y', 'transcript_count', 'num_detected_genes', 'barcodeCount', 'mouse_id', 'slide_id', 'cohort', 'age', 'batch', 'celltype', 'region', 'subregion', 'bbox_area', 'n_genes_by_counts', 'total_counts'
var: 'n_cells_by_counts', 'mean_counts', 'pct_dropout_by_counts', 'total_counts'
uns: 'neighbors', 'pca', 'preprocess_info', 'stgp', 'umap'
obsm: 'X_pca', 'X_stgp', 'X_stgp_spatial', 'X_umap', 'spatial'
varm: 'PCs'
obsp: 'connectivities', 'distances'
Analysis selection: mouse_97; n_obs=3626, n_vars=220
mouse_id slide_id age batch
97 B17 32.6 B
| mouse_id | age | slide_id | n_cells | |
|---|---|---|---|---|
| 3779 | 97 | 32.6 | B17 | 3626 |
| 3558 | 93 | 28.5 | B5 | 3113 |
| 2897 | 75 | 18.8 | B4 | 2870 |
| 3335 | 86 | 24.6 | B2 | 2826 |
| 1763 | 42 | 30.9 | A9 | 2796 |
| 664 | 14 | 12.9 | A10 | 2789 |
| 3107 | 80 | 19.8 | B4 | 2738 |
| 2229 | 57 | 4.3 | B17 | 2711 |
| 12 | 1 | 3.8 | A8 | 2543 |
| 1541 | 38 | 26.7 | A7 | 2543 |
| 1110 | 30 | 21.4 | A6 | 2538 |
| 3996 | 101 | 34.5 | B3 | 2518 |
| 2448 | 61 | 6.6 | B5 | 2498 |
| 233 | 7 | 5.4 | A9 | 2497 |
| 1982 | 46 | 33.2 | A8 | 2434 |
| 1324 | 33 | 23.5 | A10 | 2398 |
| 451 | 11 | 9.8 | A7 | 2390 |
| 870 | 19 | 15.5 | A6 | 2258 |
| 2006 | 53 | 3.4 | B3 | 2195 |
| 2685 | 70 | 15.8 | B2 | 2136 |
2. KNN spectral clustering in the stGP spatial embedding
[2]:
# Microglia has a single celltype label, so subregion is used as the comparison label.
b = np.asarray(adata.obsm['X_stgp_spatial'])
if b.shape[0] != adata.n_obs:
b = b.T
if b.shape[0] != adata.n_obs:
raise ValueError(f'X_stgp_spatial shape {b.shape} is not aligned with n_obs={adata.n_obs}')
truth = adata.obs[TRUTH_KEY].astype('category').cat.remove_unused_categories()
n_clusters = truth.cat.categories.size
n_clusters = 4
if n_clusters < 2:
raise ValueError(f'Need at least 2 observed {TRUTH_KEY} categories for spectral clustering.')
k_nn = min(max(2, int(np.round(np.sqrt(b.shape[0])))), b.shape[0] - 1)
print(f'truth={TRUTH_KEY}, n_clusters={n_clusters}, k_nn={k_nn}')
nn = NearestNeighbors(n_neighbors=k_nn + 1, metric='euclidean').fit(b).kneighbors(return_distance=False)[:, 1:]
rows = np.repeat(np.arange(nn.shape[0]), k_nn)
cols = nn.ravel()
knn_graph = sp.csr_matrix((np.ones(rows.size), (rows, cols)), shape=(b.shape[0], b.shape[0]))
knn_graph = knn_graph.maximum(knn_graph.T)
clusterlabel = SpectralClustering(
n_clusters=n_clusters,
affinity='precomputed',
assign_labels='kmeans',
random_state=RANDOM_STATE,
).fit_predict(knn_graph) + 1
adata.obs['stGP_domain'] = pd.Categorical(clusterlabel.astype(str))
ctab = pd.crosstab(adata.obs['stGP_domain'], truth, rownames=['stGP'], colnames=[TRUTH_KEY])
display(ctab)
ctab.to_csv(OUT_DIR / f'{label_suffix}_stgp_domain_{TRUTH_KEY}_crosstab.csv')
truth=subregion, n_clusters=4, k_nn=60
| subregion | CC/ACO | CTX_L1/MEN | CTX_L2/3 | CTX_L4/5/6 | STR_CP/ACB | STR_LS/NDB | VEN |
|---|---|---|---|---|---|---|---|
| stGP | |||||||
| 1 | 192 | 95 | 147 | 193 | 519 | 331 | 1 |
| 2 | 5 | 133 | 657 | 184 | 242 | 79 | 0 |
| 3 | 319 | 5 | 6 | 7 | 196 | 10 | 2 |
| 4 | 280 | 2 | 0 | 1 | 19 | 1 | 0 |
[3]:
adata_full
[3]:
AnnData object with n_obs × n_vars = 52417 × 220
obs: 'volume', 'center_x', 'center_y', 'min_x', 'min_y', 'max_x', 'max_y', 'transcript_count', 'num_detected_genes', 'barcodeCount', 'mouse_id', 'slide_id', 'cohort', 'age', 'batch', 'celltype', 'region', 'subregion', 'bbox_area', 'n_genes_by_counts', 'total_counts'
var: 'n_cells_by_counts', 'mean_counts', 'pct_dropout_by_counts', 'total_counts'
uns: 'neighbors', 'pca', 'preprocess_info', 'stgp', 'umap'
obsm: 'X_pca', 'X_stgp', 'X_stgp_spatial', 'X_umap', 'spatial'
varm: 'PCs'
obsp: 'connectivities', 'distances'
3. Slingshot trajectory inference
[4]:
DRM = np.asarray(adata.obsm['X_stgp_spatial'])
if DRM.shape[0] != adata.n_obs:
DRM = DRM.T
adata.obsm['X_DRM'] = np.asarray(DRM)
adata.obs['clusterlabel'] = pd.Categorical(adata.obs['stGP_domain'].astype(str))
cats = adata.obs['clusterlabel'].cat.categories.astype(str)
print('Cluster categories:', list(cats))
start_cluster = '4'
start_node = int(np.where(cats == start_cluster)[0][0])
print(f'Start cluster: {start_cluster} (start_node={start_node})')
sl = Slingshot(adata, celltype_key='clusterlabel', obsm_key='X_DRM', start_node=start_node)
sl.fit(num_epochs=N_EPOCHS)
Cluster categories: ['1', '2', '3', '4']
Start cluster: 4 (start_node=3)
0%| | 0/10 [00:00<?, ?it/s]/home/byual/anaconda3/envs/NicheScope/lib/python3.9/site-packages/pyslingshot/slingshot.py:333: RuntimeWarning: invalid value encountered in divide
cell_weights = z_prime / np.nanmax(z_prime, axis=1, keepdims=True) #rowMins(D) / D
10%|█ | 1/10 [00:02<00:22, 2.48s/it]/home/byual/anaconda3/envs/NicheScope/lib/python3.9/site-packages/pyslingshot/slingshot.py:333: RuntimeWarning: invalid value encountered in divide
cell_weights = z_prime / np.nanmax(z_prime, axis=1, keepdims=True) #rowMins(D) / D
20%|██ | 2/10 [00:04<00:19, 2.43s/it]
30%|███ | 3/10 [00:07<00:16, 2.35s/it]
40%|████ | 4/10 [00:09<00:13, 2.31s/it]
50%|█████ | 5/10 [00:11<00:11, 2.37s/it]
60%|██████ | 6/10 [00:14<00:09, 2.30s/it]
70%|███████ | 7/10 [00:16<00:06, 2.28s/it]
80%|████████ | 8/10 [00:18<00:04, 2.34s/it]
90%|█████████ | 9/10 [00:20<00:02, 2.26s/it]
100%|██████████| 10/10 [00:23<00:00, 2.28s/it]
100%|██████████| 10/10 [00:23<00:00, 2.31s/it]
[5]:
adata.obs['slingPseudotime_1'] = sl.unified_pseudotime
if sl.curves is not None and sl.cell_weights is not None:
for l_idx, curve in enumerate(sl.curves):
pt = curve.pseudotimes_interp.copy()
weight = sl.cell_weights[:, l_idx].copy()
pt[weight <= 0] = np.nan
adata.obs[f'slingPseudotime_{l_idx + 1}'] = pt
adata.obs[f'slingCurveWeight_{l_idx + 1}'] = weight
adata.obs[['clusterlabel', TRUTH_KEY, 'slingPseudotime_1']].head()
[5]:
| clusterlabel | subregion | slingPseudotime_1 | |
|---|---|---|---|
| 3615378900220100285-97 | 2 | CTX_L2/3 | 18.680730 |
| 3615378900220100290-97 | 2 | CTX_L2/3 | 21.746658 |
| 3615378900220200026-97 | 2 | CTX_L2/3 | 21.590094 |
| 3615378900230100262-97 | 2 | CTX_L1/MEN | 22.114490 |
| 3615378900230100309-97 | 2 | CTX_L2/3 | 22.111445 |
4. Visualise and save outputs
[6]:
xy = np.asarray(adata.obsm['spatial'])
pt = adata.obs['slingPseudotime_1'].to_numpy()
x = xy[:, 0]
y = -xy[:, 1]
fig = plt.figure(figsize=(6, 6), constrained_layout=True)
gs = fig.add_gridspec(
2, 2,
width_ratios=[1, 0.08],
height_ratios=[1, 1]
)
ax0 = fig.add_subplot(gs[0, 0])
ax1 = fig.add_subplot(gs[1, 0], sharex=ax0, sharey=ax0)
leg_ax = fig.add_subplot(gs[0, 1])
cax = fig.add_subplot(gs[1, 1])
for label in adata.obs['clusterlabel'].cat.categories:
mask = adata.obs['clusterlabel'].astype(str).to_numpy() == str(label)
ax0.scatter(x[mask], y[mask], s=5, linewidths=0, rasterized=True, label=str(label))
scat = ax1.scatter(x, y, c=pt, s=5, cmap='viridis', linewidths=0, rasterized=True)
for ax in [ax0, ax1]:
ax.set_aspect('equal')
ax.axis('off')
ax.set_xlim(x.min(), x.max())
ax.set_ylim(y.max(), y.min()) # 上下翻转两个共享空间子图
leg_ax.axis('off')
handles, labels = ax0.get_legend_handles_labels()
leg_ax.legend(handles, labels, title='clusterlabel', markerscale=2, loc='upper left')
plt.colorbar(scat, cax=cax, label='Slingshot pseudotime')
fig.savefig(OUT_DIR / f'{label_suffix}_slingshot_spatial.png', dpi=300, bbox_inches='tight')
plt.show()
[7]:
# Arrow overlay: local spatial pseudotime gradient, from lower to higher pseudotime.
GRIDNUM = 10
MIN_CELLS_PER_GRID = 10
ARROW_LENGTH_FRAC = 0.75
xy = np.asarray(adata.obsm['spatial'])
pt = adata.obs['slingPseudotime_1'].to_numpy(dtype=float)
valid = np.isfinite(pt) & np.all(np.isfinite(xy), axis=1)
xy_valid = xy[valid]
pt_valid = pt[valid]
x_edges = np.linspace(xy_valid[:, 0].min(), xy_valid[:, 0].max(), GRIDNUM + 1)
y_edges = np.linspace(xy_valid[:, 1].min(), xy_valid[:, 1].max(), GRIDNUM + 1)
x_bin = np.clip(np.digitize(xy_valid[:, 0], x_edges) - 1, 0, GRIDNUM - 1)
y_bin = np.clip(np.digitize(xy_valid[:, 1], y_edges) - 1, 0, GRIDNUM - 1)
mean_pt = np.full((GRIDNUM, GRIDNUM), np.nan)
mean_xy = np.full((GRIDNUM, GRIDNUM, 2), np.nan)
for i in range(GRIDNUM):
for j in range(GRIDNUM):
m = (x_bin == i) & (y_bin == j)
if int(m.sum()) >= MIN_CELLS_PER_GRID:
mean_pt[i, j] = float(np.nanmean(pt_valid[m]))
mean_xy[i, j] = np.nanmean(xy_valid[m], axis=0)
arrow_start, arrow_vec = [], []
for i in range(GRIDNUM):
for j in range(GRIDNUM):
if not np.isfinite(mean_pt[i, j]):
continue
grad = np.zeros(2, dtype=float)
center = mean_xy[i, j]
for di in (-1, 0, 1):
for dj in (-1, 0, 1):
if di == 0 and dj == 0:
continue
ni, nj = i + di, j + dj
if ni < 0 or ni >= GRIDNUM or nj < 0 or nj >= GRIDNUM or not np.isfinite(mean_pt[ni, nj]):
continue
direction = mean_xy[ni, nj] - center
dist = np.linalg.norm(direction)
if dist > 0:
grad += (mean_pt[ni, nj] - mean_pt[i, j]) * direction / dist
norm = np.linalg.norm(grad)
if norm > 0:
arrow_start.append(center)
arrow_vec.append(grad / norm)
arrow_start = np.asarray(arrow_start)
arrow_vec = np.asarray(arrow_vec)
cell_size = min(np.diff(x_edges).mean(), np.diff(y_edges).mean())
arrow_vec = arrow_vec * cell_size * ARROW_LENGTH_FRAC
fig, ax = plt.subplots(figsize=(7, 7), constrained_layout=True)
for label in adata.obs['clusterlabel'].cat.categories:
mask = adata.obs['clusterlabel'].astype(str).to_numpy() == str(label)
ax.scatter(xy[mask, 0], xy[mask, 1], s=5, linewidths=0, rasterized=True, label=str(label))
if len(arrow_start) > 0:
ax.quiver(arrow_start[:, 0], arrow_start[:, 1], arrow_vec[:, 0], arrow_vec[:, 1], angles='xy', scale_units='xy', scale=1, color='black', width=0.006, headwidth=4.5, headlength=6.0, headaxislength=5.0, label='pseudotime gradient')
ax.set_aspect('equal')
ax.invert_yaxis()
ax.axis('off')
ax.set_title('Spatial pseudotime direction overlay')
ax.legend(title='clusterlabel', markerscale=2, bbox_to_anchor=(1.02, 1), loc='upper left')
fig.savefig(OUT_DIR / f'{label_suffix}_slingshot_arrow_overlay.png', dpi=300, bbox_inches='tight')
plt.show()
[8]:
sc.pp.neighbors(adata, use_rep='X_DRM', n_neighbors=k_nn)
sc.tl.umap(adata, random_state=RANDOM_STATE)
sc.pl.umap(adata, color=['clusterlabel', 'slingPseudotime_1', TRUTH_KEY], size=10, wspace=0.4, show=False)
plt.savefig(OUT_DIR / f'{label_suffix}_slingshot_umap.png', dpi=300, bbox_inches='tight')
plt.show()
[9]:
OUT_H5AD = OUT_DIR / f'{label_suffix}_stgp_cluster_slingshot.h5ad'
OUT_OBS = OUT_DIR / f'{label_suffix}_stgp_cluster_slingshot_obs.csv'
adata.write_h5ad(OUT_H5AD, compression='gzip')
adata.obs.to_csv(OUT_OBS)
print(f'Saved trajectory AnnData: {OUT_H5AD.resolve()}')
print(f'Saved trajectory obs: {OUT_OBS.resolve()}')
Saved trajectory AnnData: /import/home4/byual/stGP-0529/RealData_MouseBrainMERFISH/Results/stgp/Microglia/downstream_cluster/mouse_97_stgp_cluster_slingshot.h5ad
Saved trajectory obs: /import/home4/byual/stGP-0529/RealData_MouseBrainMERFISH/Results/stgp/Microglia/downstream_cluster/mouse_97_stgp_cluster_slingshot_obs.csv