Injured PT stGP analysis: right kidney

Fit stGP on the right biological replicate (ident suffix R) 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_R")
FIGURES_DIR.mkdir(parents=True, exist_ok=True)

CELLTYPE = "Inj_PT"
SIDE = "R"
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 = 76699 × 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
 Hour4R  Hour4          0.166667    R
Hour12R Hour12          0.500000    R
  Day2R   Day2          2.000000    R
 Day14R  Day14         14.000000    R
 Week6R  Week6         42.000000    R
[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_R: 76699 cells across 5 slices
  slice       age  n_cells
 Hour4R  0.166667    18287
Hour12R  0.500000    38001
  Day2R  2.000000    19408
 Day14R 14.000000      820
 Week6R 42.000000      183
[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.4106 | 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=1.656e-01 dTheta_rel=1.727e-01 time=5.957e+02
[sweep=002] dW_rel=1.568e-01 dTheta_rel=2.711e-02 time=4.321e+02
[sweep=003] dW_rel=8.584e-02 dTheta_rel=2.277e-02 time=4.079e+02
[sweep=004] dW_rel=6.787e-02 dTheta_rel=1.343e-02 time=3.357e+02
[sweep=005] dW_rel=7.023e-02 dTheta_rel=1.177e-02 time=3.423e+02
[sweep=006] dW_rel=1.876e-02 dTheta_rel=8.113e-03 time=2.782e+02
[sweep=007] dW_rel=6.750e-02 dTheta_rel=9.243e-03 time=2.930e+02
[sweep=008] dW_rel=1.176e-01 dTheta_rel=2.817e-02 time=4.225e+02
[sweep=009] dW_rel=1.425e-01 dTheta_rel=2.328e-02 time=4.059e+02
[sweep=010] dW_rel=2.780e-02 dTheta_rel=1.921e-02 time=4.274e+02
[sweep=011] dW_rel=7.550e-02 dTheta_rel=2.859e-02 time=3.203e+02
[sweep=012] dW_rel=1.575e-02 dTheta_rel=3.185e-03 time=1.747e+02
[sweep=013] dW_rel=7.294e-03 dTheta_rel=4.125e-03 time=1.259e+02
[sweep=014] dW_rel=4.750e-03 dTheta_rel=3.447e-03 time=1.283e+02
[sweep=015] dW_rel=3.394e-03 dTheta_rel=2.579e-03 time=9.984e+01
[sweep=016] dW_rel=2.509e-03 dTheta_rel=1.905e-03 time=1.208e+02
[sweep=017] dW_rel=1.885e-03 dTheta_rel=1.430e-03 time=1.382e+02
[sweep=018] dW_rel=1.429e-03 dTheta_rel=1.088e-03 time=1.590e+02
[sweep=019] dW_rel=1.088e-03 dTheta_rel=8.401e-04 time=7.327e+01
[sweep=020] dW_rel=8.303e-04 dTheta_rel=6.476e-04 time=7.566e+01
[sweep=021] dW_rel=6.372e-04 dTheta_rel=5.003e-04 time=2.745e+01
[sweep=022] dW_rel=4.894e-04 dTheta_rel=3.863e-04 time=2.772e+01
[sweep=023] dW_rel=3.768e-04 dTheta_rel=2.983e-04 time=4.946e+01
[sweep=024] dW_rel=2.905e-04 dTheta_rel=2.308e-04 time=5.718e+01
[sweep=025] dW_rel=2.237e-04 dTheta_rel=1.813e-04 time=3.391e+01
[sweep=026] dW_rel=1.735e-04 dTheta_rel=1.388e-04 time=1.904e+01
[sweep=027] dW_rel=1.335e-04 dTheta_rel=1.068e-04 time=1.577e+01
[sweep=028] dW_rel=1.033e-04 dTheta_rel=8.551e-05 time=1.583e+01
[auto_rank prune] p=8 -> 6 (dropped factors [6, 7])
[sweep=001] dW_rel=2.848e-01 dTheta_rel=3.871e-02 time=5.849e+01
[sweep=002] dW_rel=1.247e-01 dTheta_rel=3.304e-02 time=6.729e+01
[sweep=003] dW_rel=4.483e-02 dTheta_rel=2.465e-02 time=4.213e+01
[sweep=004] dW_rel=1.371e-02 dTheta_rel=9.353e-03 time=3.109e+01
[sweep=005] dW_rel=7.102e-03 dTheta_rel=4.718e-03 time=2.319e+01
[sweep=006] dW_rel=4.290e-03 dTheta_rel=2.881e-03 time=2.372e+01
[sweep=007] dW_rel=2.964e-03 dTheta_rel=1.976e-03 time=1.926e+01
[sweep=008] dW_rel=2.226e-03 dTheta_rel=1.447e-03 time=2.113e+01
[sweep=009] dW_rel=1.741e-03 dTheta_rel=1.096e-03 time=1.925e+01
[sweep=010] dW_rel=1.383e-03 dTheta_rel=8.476e-04 time=1.903e+01
[sweep=011] dW_rel=1.107e-03 dTheta_rel=6.660e-04 time=1.628e+01
[sweep=012] dW_rel=8.859e-04 dTheta_rel=5.235e-04 time=1.572e+01
[sweep=013] dW_rel=7.123e-04 dTheta_rel=4.176e-04 time=1.349e+01
[sweep=014] dW_rel=5.735e-04 dTheta_rel=3.310e-04 time=1.370e+01
[sweep=015] dW_rel=4.614e-04 dTheta_rel=2.665e-04 time=1.359e+01
[sweep=016] dW_rel=3.713e-04 dTheta_rel=2.142e-04 time=1.257e+01
[sweep=017] dW_rel=2.994e-04 dTheta_rel=1.692e-04 time=1.174e+01
[sweep=018] dW_rel=2.414e-04 dTheta_rel=1.377e-04 time=1.187e+01
[sweep=019] dW_rel=1.947e-04 dTheta_rel=1.112e-04 time=1.086e+01
[sweep=020] dW_rel=1.569e-04 dTheta_rel=8.966e-05 time=8.831e+00
[sweep=021] dW_rel=1.266e-04 dTheta_rel=7.218e-05 time=8.632e+00
[sweep=022] dW_rel=1.022e-04 dTheta_rel=5.805e-05 time=8.557e+00
Runtime: 8790.9s | programs selected: 6
Saved: Results/stgp/Inj_PT_R/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_R/adata_with_scores.h5ad
Saved: Results/stgp/Inj_PT_R/W.csv