Injured PT stGP analysis: left kidney

Fit stGP on the left biological replicate (ident suffix L) of injured proximal tubule cells. Run this notebook with the stGP conda environment.

[1]:
import os, sys, warnings, pickle
from pathlib import Path

import numpy as np
import pandas as pd
import scanpy as sc
import matplotlib.pyplot as plt

warnings.filterwarnings("ignore", category=FutureWarning)
sys.path.insert(0, "..")

DATA_PROC = Path("data/processed")
RESULTS_DIR = Path("Results/stgp")
FIGURES_DIR = Path("Figures/Inj_PT_L")
FIGURES_DIR.mkdir(parents=True, exist_ok=True)

CELLTYPE = "Inj_PT"
SIDE = "L"
ANALYSIS_NAME = f"{CELLTYPE}_{SIDE}"

adata_inj_pt = sc.read_h5ad(DATA_PROC / f"{CELLTYPE}.h5ad")
adata_inj_pt.obs["side"] = adata_inj_pt.obs["ident"].astype(str).str[-1]
adata_inj_pt = adata_inj_pt[adata_inj_pt.obs["injury_time_days"] > 0].copy()
adata_inj_pt = adata_inj_pt[adata_inj_pt.obs["side"] == SIDE].copy()
adata_inj_pt.obs["age"] = adata_inj_pt.obs["injury_time_days"].copy()

print(adata_inj_pt)
print(adata_inj_pt.obs[["ident", "time", "injury_time_days", "side"]].drop_duplicates().sort_values("injury_time_days").to_string(index=False))
AnnData object with n_obs × n_vars = 69210 × 299
    obs: 'x_centroid', 'y_centroid', 'n_genes', 'n_counts', 'ident', 'region', 'celltype_plot', 'time', 'CN', 'injury_time_days', 'side', 'age'
    uns: 'CN_colors', 'celltype_plot_colors', 'ident_colors', 'neighbors', 'pca', 'umap'
    obsm: 'X_pca', 'X_pca_harmony', 'X_umap', 'spatial'
    varm: 'PCs'
    obsp: 'connectivities', 'distances'
  ident   time  injury_time_days side
 Hour4L  Hour4          0.166667    L
Hour12L Hour12          0.500000    L
  Day2L   Day2          2.000000    L
 Day14L  Day14         14.000000    L
 Week6L  Week6         42.000000    L
[2]:
import matplotlib as mpl

mpl.rcParams.update({
    "font.family": "sans-serif",
    "font.sans-serif": ["Arial", "Helvetica", "DejaVu Sans"],
    "font.size": 11,
    "pdf.fonttype": 42,
    "ps.fonttype": 42,
    "svg.fonttype": "none",
    "figure.dpi": 150,
    "savefig.dpi": 400,
    "figure.facecolor": "white",
    "axes.facecolor": "white",
    "axes.linewidth": 1.2,
    "axes.spines.top": False,
    "axes.spines.right": False,
    "axes.labelsize": 11,
    "axes.titlesize": 12,
    "xtick.direction": "out",
    "ytick.direction": "out",
    "legend.frameon": False,
    "legend.fontsize": 9,
    "lines.linewidth": 1.5,
})
[3]:
import time
from stgp.estimation import fit_pfactor_auto
from stgp.kernels import (
    bandwidth_select_spatial,
    bandwidth_select_temporal,
    build_K_age,
    build_K_spa_list_from_stacked,
)
from stgp.preprocessing import standardize_coords_list

OUT_DIR = RESULTS_DIR / ANALYSIS_NAME
OUT_DIR.mkdir(parents=True, exist_ok=True)
PKL_PATH = OUT_DIR / "stgp_result.pkl"

age_arr = pd.to_numeric(adata_inj_pt.obs["injury_time_days"], errors="coerce").to_numpy(float)
groups = adata_inj_pt.obs["ident"].astype(str).to_numpy()
uniq, inv = np.unique(groups, return_inverse=True)
idx_per_group = [np.sort(np.where(inv == t)[0]) for t in range(len(uniq))]

adata_prep = adata_inj_pt.copy()
sc.pp.scale(adata_prep)
Y_list = [adata_prep.X[ix] for ix in idx_per_group]
nlist = np.array([len(ix) for ix in idx_per_group])
ages = np.array([age_arr[ix[0]] for ix in idx_per_group])
sort_ord = np.argsort(ages)
ages = ages[sort_ord]
slices = uniq[sort_ord]
idx_sorted = [idx_per_group[i] for i in sort_ord]
nlist = nlist[sort_ord]
Y_list = [Y_list[i] for i in sort_ord]

coords_list = standardize_coords_list([adata_inj_pt.obsm["spatial"][ix] for ix in idx_per_group])
coords_list = [coords_list[i] for i in sort_ord]

print(f"Prepared {ANALYSIS_NAME}: {adata_inj_pt.n_obs} cells across {len(slices)} slices")
print(pd.DataFrame({"slice": slices, "age": ages, "n_cells": nlist}).to_string(index=False))
/home/byual/.conda/envs/stGP/lib/python3.11/functools.py:909: UserWarning: zero-centering a sparse array/matrix densifies it.
  return dispatch(args[0].__class__)(*args, **kw)
Prepared Inj_PT_L: 69210 cells across 5 slices
  slice       age  n_cells
 Hour4L  0.166667    15139
Hour12L  0.500000    33544
  Day2L  2.000000    19449
 Day14L 14.000000      772
 Week6L 42.000000      306
[4]:
gamma_spa = bandwidth_select_spatial(coords_list, frac=0.01, rho=0.5)
gamma_age = bandwidth_select_temporal(ages, rho=np.exp(-1.5))
print(f"gamma_spa = {gamma_spa:.4f} | gamma_age = {gamma_age:.4f}")

K_age = build_K_age(ages, gamma_age, kernel="ar1", standardize=True)
K_spa_list = build_K_spa_list_from_stacked(
    np.vstack(coords_list), nlist, gamma_spa, standardize=False, jitter=1e-6
)
gamma_spa = 0.3181 | gamma_age = 0.3904
[5]:
t0 = time.perf_counter()
res = fit_pfactor_auto(
    Y_list=Y_list,
    Nlist=nlist,
    K_age=K_age,
    Kspa_list=K_spa_list,
    p_max=10,
    k=15,
    inner_rank1_tol=1e-4,
    rel_improve_total_tol=0.002,
    backfit_tol=1e-4,
    prune_energy_frac=0.005,
    random_state=0,
    verbose=1,
)
print(f"Runtime: {time.perf_counter() - t0:.1f}s | programs selected: {res['W'].shape[0]}")

res["gamma_age"] = gamma_age
res["gamma_spa"] = gamma_spa
with open(PKL_PATH, "wb") as f:
    pickle.dump(res, f)
print(f"Saved: {PKL_PATH}")
[sweep=001] dW_rel=2.171e-01 dTheta_rel=8.388e-02 time=2.631e+02
[sweep=002] dW_rel=1.062e-01 dTheta_rel=1.977e-02 time=1.853e+02
[sweep=003] dW_rel=1.269e-01 dTheta_rel=1.709e-02 time=1.433e+02
[sweep=004] dW_rel=3.793e-02 dTheta_rel=1.203e-02 time=1.390e+02
[sweep=005] dW_rel=2.382e-02 dTheta_rel=9.420e-03 time=1.097e+02
[sweep=006] dW_rel=9.642e-02 dTheta_rel=3.352e-02 time=1.372e+02
[sweep=007] dW_rel=1.482e-02 dTheta_rel=1.190e-02 time=1.205e+02
[sweep=008] dW_rel=1.086e-02 dTheta_rel=7.258e-03 time=1.082e+02
[sweep=009] dW_rel=6.806e-02 dTheta_rel=6.524e-03 time=1.229e+02
[sweep=010] dW_rel=2.376e-02 dTheta_rel=3.832e-03 time=1.020e+02
[sweep=011] dW_rel=1.126e-02 dTheta_rel=2.049e-03 time=8.972e+01
[sweep=012] dW_rel=7.652e-03 dTheta_rel=1.144e-03 time=7.338e+01
[sweep=013] dW_rel=5.741e-03 dTheta_rel=7.454e-04 time=6.404e+01
[sweep=014] dW_rel=8.669e-02 dTheta_rel=2.594e-03 time=1.062e+02
[sweep=015] dW_rel=1.170e-02 dTheta_rel=3.616e-03 time=1.264e+02
[sweep=016] dW_rel=7.876e-03 dTheta_rel=1.845e-03 time=9.233e+01
[sweep=017] dW_rel=5.649e-03 dTheta_rel=8.460e-04 time=5.447e+01
[sweep=018] dW_rel=4.251e-03 dTheta_rel=4.966e-04 time=6.370e+01
[sweep=019] dW_rel=3.291e-03 dTheta_rel=4.087e-04 time=5.610e+01
[sweep=020] dW_rel=7.504e-02 dTheta_rel=9.172e-04 time=8.994e+01
[sweep=021] dW_rel=4.454e-02 dTheta_rel=1.640e-03 time=8.880e+01
[sweep=022] dW_rel=4.498e-02 dTheta_rel=1.260e-02 time=1.613e+02
[sweep=023] dW_rel=4.910e-02 dTheta_rel=1.025e-02 time=1.177e+02
[sweep=024] dW_rel=1.543e-02 dTheta_rel=5.719e-03 time=9.460e+01
[sweep=025] dW_rel=8.107e-02 dTheta_rel=8.919e-03 time=1.229e+02
[sweep=026] dW_rel=4.479e-03 dTheta_rel=2.150e-03 time=1.119e+02
[sweep=027] dW_rel=2.500e-03 dTheta_rel=9.286e-04 time=8.642e+01
[sweep=028] dW_rel=1.569e-03 dTheta_rel=5.092e-04 time=8.423e+01
[sweep=029] dW_rel=1.085e-03 dTheta_rel=3.084e-04 time=5.609e+01
[sweep=030] dW_rel=7.869e-04 dTheta_rel=2.051e-04 time=4.658e+01
[sweep=031] dW_rel=5.875e-04 dTheta_rel=1.431e-04 time=5.048e+01
[sweep=032] dW_rel=4.474e-04 dTheta_rel=1.042e-04 time=3.142e+01
[sweep=033] dW_rel=3.450e-04 dTheta_rel=7.981e-05 time=3.013e+01
[sweep=034] dW_rel=2.682e-04 dTheta_rel=6.406e-05 time=3.516e+01
[sweep=035] dW_rel=2.092e-04 dTheta_rel=5.099e-05 time=2.963e+01
[sweep=036] dW_rel=1.639e-04 dTheta_rel=4.025e-05 time=2.040e+01
[sweep=037] dW_rel=1.281e-04 dTheta_rel=3.236e-05 time=2.414e+01
[sweep=038] dW_rel=1.011e-04 dTheta_rel=2.774e-05 time=3.239e+01
[auto_rank prune] p=8 -> 7 (dropped factors [7])
[sweep=001] dW_rel=2.877e-01 dTheta_rel=9.598e-02 time=3.005e+02
[sweep=002] dW_rel=1.178e-01 dTheta_rel=2.802e-02 time=1.665e+02
[sweep=003] dW_rel=2.279e-02 dTheta_rel=6.407e-03 time=1.130e+02
[sweep=004] dW_rel=1.265e-02 dTheta_rel=2.698e-03 time=7.828e+01
[sweep=005] dW_rel=7.551e-03 dTheta_rel=1.545e-03 time=6.819e+01
[sweep=006] dW_rel=4.718e-03 dTheta_rel=9.438e-04 time=5.467e+01
[sweep=007] dW_rel=3.048e-03 dTheta_rel=5.661e-04 time=4.435e+01
[sweep=008] dW_rel=2.039e-03 dTheta_rel=3.307e-04 time=4.407e+01
[sweep=009] dW_rel=1.408e-03 dTheta_rel=1.945e-04 time=3.429e+01
[sweep=010] dW_rel=1.006e-03 dTheta_rel=1.306e-04 time=2.840e+01
[sweep=011] dW_rel=7.344e-04 dTheta_rel=1.044e-04 time=2.690e+01
[sweep=012] dW_rel=5.496e-04 dTheta_rel=9.075e-05 time=3.719e+01
[sweep=013] dW_rel=4.155e-04 dTheta_rel=8.234e-05 time=2.347e+01
[sweep=014] dW_rel=3.161e-04 dTheta_rel=7.395e-05 time=1.910e+01
[sweep=015] dW_rel=2.413e-04 dTheta_rel=6.326e-05 time=1.871e+01
[sweep=016] dW_rel=1.839e-04 dTheta_rel=5.513e-05 time=1.953e+01
[sweep=017] dW_rel=1.415e-04 dTheta_rel=4.816e-05 time=1.874e+01
[sweep=018] dW_rel=1.065e-04 dTheta_rel=3.972e-05 time=1.615e+01
Runtime: 6613.8s | programs selected: 7
Saved: Results/stgp/Inj_PT_L/stgp_result.pkl
[6]:
ADATA_PATH = OUT_DIR / "adata_with_scores.h5ad"

adata = adata_inj_pt.copy()
all_idx = np.concatenate(idx_sorted)
H_arr = np.empty_like(res["H"])
b_arr = np.empty_like(res["b"])
H_arr[all_idx] = res["H"]
b_arr[all_idx] = res["b"]

adata.obsm["X_stgp"] = H_arr.astype(np.float32)
adata.obsm["X_stgp_spatial"] = b_arr.astype(np.float32)
adata.uns["stgp"] = dict(
    groups=slices.tolist(),
    ages=ages.tolist(),
    gamma_age=float(res["gamma_age"]),
    gamma_spa=float(res["gamma_spa"]),
    p_selected=res["W"].shape[0],
    alpha=np.asarray(res["alpha"]).tolist(),
    alpha_lower=np.asarray(res["alpha_lower"]).tolist(),
    alpha_upper=np.asarray(res["alpha_upper"]).tolist(),
    theta=np.asarray(res["theta"]).tolist(),
    sigma2e=float(res.get("sigma2e", np.nan)),
)
adata.write_h5ad(str(ADATA_PATH), compression="gzip")
print(f"Saved: {ADATA_PATH}")

p_sel = res["W"].shape[0]
W_df = pd.DataFrame(
    res["W"],
    index=[f"stGP{j + 1}" for j in range(p_sel)],
    columns=adata.var_names.astype(str),
)
W_df.to_csv(OUT_DIR / "W.csv")
print(f"Saved: {OUT_DIR / 'W.csv'}")
Saved: Results/stgp/Inj_PT_L/adata_with_scores.h5ad
Saved: Results/stgp/Inj_PT_L/W.csv