Collagen distance model¶
Supplementary Notes
Overview and notation We analyzed spatial-omics data summarized at the level of (condition, cell type, batch, region). For each observation $i$: ● : condition index ● : cell type index ● : batch index ● : region label ● : number of “positive” cells ● : total number of cells ● : standardized logit-area covariate ● : region indicator coded as (core) and (border) Likelihood: Beta–Binomial regression Counts are bounded by and arise from heterogeneous spatial sampling. We used:
We parameterized the Beta distribution using the mean and a cell-type–specific precision :
The mean probability is linked to $\eta_i$ via:
Linear predictor and parameter meanings We decomposed technical and biological variation using:
Algebra of region coding Region was effect-coded as:
Define the region-invariant part of the linear predictor:
Then the contrast is derived as:
Exponentiating gives the odds ratio:
Prior specification and weakly informative interpretation We specified weakly informative priors to guide estimation in a stable, data-driven way. These priors offer enough regularization to prevent overfitting—particularly in smaller groups—while maintaining flexibility to reflect biologically plausible effects. Global effects:
Baseline condition cell type effects:
Hierarchical region contrasts : We decomposed δ_c,k into global, condition, cell type, and interaction components: Supplementary Notes
Overview and notation We analyzed spatial-omics data summarized at the level of (condition, cell type, batch, region). For each observation $i$: ● : condition index ● : cell type index ● : batch index ● : region label ● : number of “positive” cells ● : total number of cells ● : standardized logit-area covariate ● : region indicator coded as (core) and (border) Likelihood: Beta–Binomial regression Counts are bounded by and arise from heterogeneous spatial sampling. We used:
We parameterized the Beta distribution using the mean and a cell-type–specific precision :
The mean probability is linked to $\eta_i$ via:
Linear predictor and parameter meanings We decomposed technical and biological variation using:
Algebra of region coding Region was effect-coded as:
Define the region-invariant part of the linear predictor:
Then the contrast is derived as:
Exponentiating gives the odds ratio:
Prior specification and weakly informative interpretation We specified weakly informative priors to guide estimation in a stable, data-driven way. These priors offer enough regularization to prevent overfitting—particularly in smaller groups—while maintaining flexibility to reflect biologically plausible effects. Global effects:
Baseline condition cell type effects:
Hierarchical region contrasts : We decomposed δ_c,k into global, condition, cell type, and interaction components:
Centering δ_c and δ_k enforces sum-to-zero identifiability and improves interpretability: δ₀ captures the overall average border-vs-core effect, while δ_c and δ_k capture deviations by condition and cell type. Batch random intercept (non-centered):
Cell-type–specific precision: Precision parameters were modeled on the log scale with hierarchical shrinkage:
Anchoring μ_logφ around log(mean(n)) ties dispersion to the typical sample-size scale while allowing cell types to deviate.
Prior predictive simulation
We performed prior predictive simulation to verify that the chosen weakly informative priors generate plausible log-odds, probabilities, and count fractions before conditioning on the observed data. For each prior draw, model parameters were sampled from their priors, including the global intercept α\alphaα, offset effect βoff\beta_{\text{off}}βoff, batch random effects usu_sus, condition×cell-type baseline effects μc,k\mu_{c,k}μc,k, region effects δc,k\delta_{c,k}δc,k (hierarchically parameterized), and cell-type-specific dispersion ϕk\phi_kϕk. Using the observed design indices (ci,ki,si)(c_i,k_i,s_i)(ci,ki,si), standardized offset offseti\text{offset}_ioffseti, and region coding zi=−0.5z_i=-0.5zi=−0.5 (core) or zi=+0.5z_i=+0.5zi=+0.5 (border), we constructed the linear predictor:
and transformed to probabilities pi=logit−1(ηi)p_i=\text{logit}^{-1}(\eta_i)pi=logit−1(ηi). Given the total count nin_ini, we simulated prior predictive counts via the Beta–Binomial parameterization:
As a compact diagnostic, we summarized (i) the pooled distribution of η\etaη over all observations and prior draws (Fig. Sx, panel a), and (ii) the fraction of observations per prior draw with near-boundary probabilities p<ϵp<\epsilonp<ϵ or p>1−ϵp>1-\epsilonp>1−ϵ for ϵ=10−3\epsilon=10^{-3}ϵ=10−3 (Fig. Sx, panels b–c). These “extreme-mass” summaries are sensitive to priors that inadvertently imply near-deterministic outcomes for large portions of the dataset. In our settings, the bulk of prior mass concentrated away from the boundaries, indicating that the priors provide regularization while retaining flexibility to accommodate both low and high positivity rates across strata.
Centering δ_c and δ_k enforces sum-to-zero identifiability and improves interpretability: δ₀ captures the overall average border-vs-core effect, while δ_c and δ_k capture deviations by condition and cell type. Batch random intercept (non-centered):
Cell-type–specific precision: Precision parameters were modeled on the log scale with hierarchical shrinkage:
Anchoring μ_logφ around log(mean(n)) ties dispersion to the typical sample-size scale while allowing cell types to deviate.
Prior predictive simulation
We performed prior predictive simulation to verify that the chosen weakly informative priors generate plausible log-odds, probabilities, and count fractions before conditioning on the observed data. For each prior draw, model parameters were sampled from their priors, including the global intercept α\alphaα, offset effect βoff\beta_{\text{off}}βoff, batch random effects usu_sus, condition×cell-type baseline effects μc,k\mu_{c,k}μc,k, region effects δc,k\delta_{c,k}δc,k (hierarchically parameterized), and cell-type-specific dispersion ϕk\phi_kϕk. Using the observed design indices (ci,ki,si)(c_i,k_i,s_i)(ci,ki,si), standardized offset offseti\text{offset}_ioffseti, and region coding zi=−0.5z_i=-0.5zi=−0.5 (core) or zi=+0.5z_i=+0.5zi=+0.5 (border), we constructed the linear predictor:
and transformed to probabilities pi=logit−1(ηi)p_i=\text{logit}^{-1}(\eta_i)pi=logit−1(ηi). Given the total count nin_ini, we simulated prior predictive counts via the Beta–Binomial parameterization:
As a compact diagnostic, we summarized (i) the pooled distribution of η\etaη over all observations and prior draws (Fig. Sx, panel a), and (ii) the fraction of observations per prior draw with near-boundary probabilities p<ϵp<\epsilonp<ϵ or p>1−ϵp>1-\epsilonp>1−ϵ for ϵ=10−3\epsilon=10^{-3}ϵ=10−3 (Fig. Sx, panels b–c). These “extreme-mass” summaries are sensitive to priors that inadvertently imply near-deterministic outcomes for large portions of the dataset. In our settings, the bulk of prior mass concentrated away from the boundaries, indicating that the priors provide regularization while retaining flexibility to accommodate both low and high positivity rates across strata.
In [2]:
Copied!
# mamba activate pymc_latest
import scanpy as sc
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
from matplotlib.colors import ListedColormap
from matplotlib.colors import Normalize, LinearSegmentedColormap
from matplotlib.lines import Line2D
import seaborn as sns
import os
import gzip
import numpy as np
import statsmodels.api as sm
from scipy.special import softmax
import random
import re
import arviz as az
import xarray as xr
import pymc as pm
import math
RANDOM_SEED = 8927
rng = np.random.default_rng(RANDOM_SEED)
import pickle
import anndata as ad
# mamba activate pymc_latest
import scanpy as sc
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
from matplotlib.colors import ListedColormap
from matplotlib.colors import Normalize, LinearSegmentedColormap
from matplotlib.lines import Line2D
import seaborn as sns
import os
import gzip
import numpy as np
import statsmodels.api as sm
from scipy.special import softmax
import random
import re
import arviz as az
import xarray as xr
import pymc as pm
import math
RANDOM_SEED = 8927
rng = np.random.default_rng(RANDOM_SEED)
import pickle
import anndata as ad
In [ ]:
Copied!
# Large input files for this notebook are hosted on Zenodo; download and unpack the Zenodo bundle locally and set the base data path below to your local folder.
DATA_ROOT = "data_external/nguyenlab-il15act-dataset-v1"
path_tables = os.path.join(DATA_ROOT, "02_tables_inputs")
table_model = pd.read_csv(os.path.join(path_tables, "colga_distance_input_table.csv"))
# Large input files for this notebook are hosted on Zenodo; download and unpack the Zenodo bundle locally and set the base data path below to your local folder.
DATA_ROOT = "data_external/nguyenlab-il15act-dataset-v1"
path_tables = os.path.join(DATA_ROOT, "02_tables_inputs")
table_model = pd.read_csv(os.path.join(path_tables, "colga_distance_input_table.csv"))
Practical interpretation. A prior predictive distribution of η that is excessively wide (e.g., placing substantial mass at ∣η∣≫10) typically yields probabilities indistinguishable from 0 or 1 and can indicate overly diffuse priors. Likewise, extreme-mass histograms shifted toward large fractions suggest priors that frequently generate near-deterministic pip_ipi across many observations
In [4]:
Copied!
z_map = {"core": -0.5, "border": +0.5}
z = table_model["roi"].map(z_map).values.astype(float)
# --- factorize indices (you already have these)
ind_cells, cells_levels = table_model['cluster'].factorize()
ind_cond, cond_levels = table_model['condition'].factorize()
ind_batch, batch_levels = table_model['batch'].factorize()
ind_roi, roi_levels = table_model['roi'].factorize() # core/border only (2)
# data
logit_offset = table_model['area_odr'].values
offset_in = (logit_offset - np.mean(logit_offset)) / np.std(logit_offset)
y = table_model['positive'].values
N = table_model['total'].values
mean_N = np.mean(N)
coords = {
"obs_id": np.arange(len(y)),
"Cells": cells_levels,
"Condition": cond_levels,
"S": batch_levels,
"r": roi_levels,
}
with pm.Model(coords=coords) as model8:
# indices
c = pm.Data("c", ind_cond, dims="obs_id")
k = pm.Data("k", ind_cells, dims="obs_id")
s = pm.Data("s", ind_batch, dims="obs_id")
r = pm.Data("r", ind_roi, dims="obs_id")
z_d = pm.Data("z", z, dims="obs_id")
# data
offset = pm.Data("offset", offset_in, dims="obs_id")
n = pm.Data("n", N, dims="obs_id")
Y = pm.Data("y", y, dims="obs_id")
# intercept
alpha = pm.Normal("alpha", 0, 2)
# condition × cell baseline (shared across regions)
mu_ck = pm.Normal("mu_ck", 0, 2, dims=("Condition", "Cells"))
delta0 = pm.Normal("delta0", 0, 2)
delta_c_raw = pm.Normal("delta_c_raw", 0, 2, dims=("Condition",))
delta_k_raw = pm.Normal("delta_k_raw", 0, 2, dims=("Cells",))
delta_c = pm.Deterministic("delta_c", delta_c_raw - pm.math.mean(delta_c_raw), dims=("Condition",))
delta_k = pm.Deterministic("delta_k", delta_k_raw - pm.math.mean(delta_k_raw), dims=("Cells",))
sigma_delta = pm.HalfNormal("sigma_delta", 1)
delta_raw = pm.Normal("delta_raw", 0, 2, dims=("Condition", "Cells"))
delta_ck = pm.Deterministic(
"delta_ck",
delta0 + delta_c[:, None] + delta_k[None, :] + sigma_delta * delta_raw,
dims=("Condition", "Cells"),
)
# batch random intercept (non-centered)
sigma_batch = pm.HalfNormal("sigma_batch", 1)
u_batch_raw = pm.Normal("u_batch_raw", 0, 2, dims=("S",))
u_batch = pm.Deterministic("u_batch", u_batch_raw * sigma_batch, dims=("S",))
# dispersion
mu_logphi = pm.Normal("mu_logphi", mu=np.log(mean_N), sigma=2)
sigma_logphi = pm.HalfNormal("sigma_logphi", 2)
logphi_k_raw = pm.Normal("logphi_k_raw", mu=0, sigma=2, dims=("Cells",))
logphi_k = pm.Deterministic("logphi_k",mu_logphi + sigma_logphi * logphi_k_raw, dims=("Cells",))
phi_k = pm.Deterministic("phi_k", 1 + pm.math.exp(logphi_k), dims=("Cells",))
# linear predictor
beta_off = pm.Normal("beta_off", 0.0, 2)
eta = pm.Deterministic("eta", beta_off*offset + alpha + u_batch[s] + mu_ck[c, k] + delta_ck[c, k] * z_d, dims="obs_id")
p = pm.Deterministic("p", pm.math.sigmoid(eta), dims="obs_id")
phi_obs = phi_k[k]
a = p * phi_obs
b = (1 - p) * phi_obs
# Use BetaBinomial correctly for prior predictive
y_pred = pm.BetaBinomial("y_pred", alpha=a, beta=b, n=n, observed=Y, dims="obs_id")
# Run Prior Predictive Check inside the model context
print("Running Prior Predictive Simulation...")
idata_prior = pm.sample_prior_predictive(samples=500, random_seed=RANDOM_SEED)
print("Done.")
z_map = {"core": -0.5, "border": +0.5}
z = table_model["roi"].map(z_map).values.astype(float)
# --- factorize indices (you already have these)
ind_cells, cells_levels = table_model['cluster'].factorize()
ind_cond, cond_levels = table_model['condition'].factorize()
ind_batch, batch_levels = table_model['batch'].factorize()
ind_roi, roi_levels = table_model['roi'].factorize() # core/border only (2)
# data
logit_offset = table_model['area_odr'].values
offset_in = (logit_offset - np.mean(logit_offset)) / np.std(logit_offset)
y = table_model['positive'].values
N = table_model['total'].values
mean_N = np.mean(N)
coords = {
"obs_id": np.arange(len(y)),
"Cells": cells_levels,
"Condition": cond_levels,
"S": batch_levels,
"r": roi_levels,
}
with pm.Model(coords=coords) as model8:
# indices
c = pm.Data("c", ind_cond, dims="obs_id")
k = pm.Data("k", ind_cells, dims="obs_id")
s = pm.Data("s", ind_batch, dims="obs_id")
r = pm.Data("r", ind_roi, dims="obs_id")
z_d = pm.Data("z", z, dims="obs_id")
# data
offset = pm.Data("offset", offset_in, dims="obs_id")
n = pm.Data("n", N, dims="obs_id")
Y = pm.Data("y", y, dims="obs_id")
# intercept
alpha = pm.Normal("alpha", 0, 2)
# condition × cell baseline (shared across regions)
mu_ck = pm.Normal("mu_ck", 0, 2, dims=("Condition", "Cells"))
delta0 = pm.Normal("delta0", 0, 2)
delta_c_raw = pm.Normal("delta_c_raw", 0, 2, dims=("Condition",))
delta_k_raw = pm.Normal("delta_k_raw", 0, 2, dims=("Cells",))
delta_c = pm.Deterministic("delta_c", delta_c_raw - pm.math.mean(delta_c_raw), dims=("Condition",))
delta_k = pm.Deterministic("delta_k", delta_k_raw - pm.math.mean(delta_k_raw), dims=("Cells",))
sigma_delta = pm.HalfNormal("sigma_delta", 1)
delta_raw = pm.Normal("delta_raw", 0, 2, dims=("Condition", "Cells"))
delta_ck = pm.Deterministic(
"delta_ck",
delta0 + delta_c[:, None] + delta_k[None, :] + sigma_delta * delta_raw,
dims=("Condition", "Cells"),
)
# batch random intercept (non-centered)
sigma_batch = pm.HalfNormal("sigma_batch", 1)
u_batch_raw = pm.Normal("u_batch_raw", 0, 2, dims=("S",))
u_batch = pm.Deterministic("u_batch", u_batch_raw * sigma_batch, dims=("S",))
# dispersion
mu_logphi = pm.Normal("mu_logphi", mu=np.log(mean_N), sigma=2)
sigma_logphi = pm.HalfNormal("sigma_logphi", 2)
logphi_k_raw = pm.Normal("logphi_k_raw", mu=0, sigma=2, dims=("Cells",))
logphi_k = pm.Deterministic("logphi_k",mu_logphi + sigma_logphi * logphi_k_raw, dims=("Cells",))
phi_k = pm.Deterministic("phi_k", 1 + pm.math.exp(logphi_k), dims=("Cells",))
# linear predictor
beta_off = pm.Normal("beta_off", 0.0, 2)
eta = pm.Deterministic("eta", beta_off*offset + alpha + u_batch[s] + mu_ck[c, k] + delta_ck[c, k] * z_d, dims="obs_id")
p = pm.Deterministic("p", pm.math.sigmoid(eta), dims="obs_id")
phi_obs = phi_k[k]
a = p * phi_obs
b = (1 - p) * phi_obs
# Use BetaBinomial correctly for prior predictive
y_pred = pm.BetaBinomial("y_pred", alpha=a, beta=b, n=n, observed=Y, dims="obs_id")
# Run Prior Predictive Check inside the model context
print("Running Prior Predictive Simulation...")
idata_prior = pm.sample_prior_predictive(samples=500, random_seed=RANDOM_SEED)
print("Done.")
Running Prior Predictive Simulation...
Sampling: [alpha, beta_off, delta0, delta_c_raw, delta_k_raw, delta_raw, logphi_k_raw, mu_ck, mu_logphi, sigma_batch, sigma_delta, sigma_logphi, u_batch_raw, y_pred]
Done.
In [5]:
Copied!
# --- Chunk C: Plotting Prior Predictive Results from PyMC Model ---
def plot_prior_check(idata):
prior = idata.prior
# prior_pred = idata.prior_predictive # Not needed if we don't plot y/n
def plot_hist_ax(ax, x, title, xlabel, color='gray'):
# Flatten and filter nan if any
x_flat = x.ravel()
ax.hist(x_flat, bins=40, density=True, color=color, alpha=0.7, ec='white', lw=0.3)
ax.set_title(title, fontsize=8, fontweight='bold')
ax.set_xlabel(xlabel, fontsize=7)
ax.set_ylabel("Density", fontsize=7)
ax.tick_params(axis='both', labelsize=6)
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['left'].set_linewidth(0.5)
ax.spines['bottom'].set_linewidth(0.5)
# Extract variables
eta_vals = prior["eta"].values
p_vals = prior["p"].values
# Calculate Extreme Mass Diagnostics
extreme_eps = 1e-3
p_flat = p_vals.reshape(-1, p_vals.shape[-1]) # (samples, obs)
frac_low = (p_flat < extreme_eps).mean(axis=1)
frac_high = (p_flat > 1 - extreme_eps).mean(axis=1)
# Combined Figure: 3 Panels side-by-side in one row
# Full page width ~7.2 inches
fig, axes = plt.subplots(1, 3, figsize=(7.2, 2.4), dpi=300)
# 1. Eta (Log-odds)
# Title changed as requested (removed "Prior Predictive")
plot_hist_ax(axes[0], eta_vals, "a. Linear Predictor: η (pooled)", "Log-odds (η)")
# 2. Mass < epsilon
plot_hist_ax(axes[1], frac_low, f"b. Mass < {extreme_eps:g}", "Fraction of Obs per Draw", color='tab:blue')
# 3. Mass > 1-epsilon
plot_hist_ax(axes[2], frac_high, f"c. Mass > 1-{extreme_eps:g}", "Fraction of Obs per Draw", color='tab:red')
plt.tight_layout()
plt.show()
if 'idata_prior' in locals():
print("Generating journal-ready prior predictive plots (3-panel row)...")
plot_prior_check(idata_prior)
else:
print("idata_prior not found. Did the model run completely?")
# --- Chunk C: Plotting Prior Predictive Results from PyMC Model ---
def plot_prior_check(idata):
prior = idata.prior
# prior_pred = idata.prior_predictive # Not needed if we don't plot y/n
def plot_hist_ax(ax, x, title, xlabel, color='gray'):
# Flatten and filter nan if any
x_flat = x.ravel()
ax.hist(x_flat, bins=40, density=True, color=color, alpha=0.7, ec='white', lw=0.3)
ax.set_title(title, fontsize=8, fontweight='bold')
ax.set_xlabel(xlabel, fontsize=7)
ax.set_ylabel("Density", fontsize=7)
ax.tick_params(axis='both', labelsize=6)
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['left'].set_linewidth(0.5)
ax.spines['bottom'].set_linewidth(0.5)
# Extract variables
eta_vals = prior["eta"].values
p_vals = prior["p"].values
# Calculate Extreme Mass Diagnostics
extreme_eps = 1e-3
p_flat = p_vals.reshape(-1, p_vals.shape[-1]) # (samples, obs)
frac_low = (p_flat < extreme_eps).mean(axis=1)
frac_high = (p_flat > 1 - extreme_eps).mean(axis=1)
# Combined Figure: 3 Panels side-by-side in one row
# Full page width ~7.2 inches
fig, axes = plt.subplots(1, 3, figsize=(7.2, 2.4), dpi=300)
# 1. Eta (Log-odds)
# Title changed as requested (removed "Prior Predictive")
plot_hist_ax(axes[0], eta_vals, "a. Linear Predictor: η (pooled)", "Log-odds (η)")
# 2. Mass < epsilon
plot_hist_ax(axes[1], frac_low, f"b. Mass < {extreme_eps:g}", "Fraction of Obs per Draw", color='tab:blue')
# 3. Mass > 1-epsilon
plot_hist_ax(axes[2], frac_high, f"c. Mass > 1-{extreme_eps:g}", "Fraction of Obs per Draw", color='tab:red')
plt.tight_layout()
plt.show()
if 'idata_prior' in locals():
print("Generating journal-ready prior predictive plots (3-panel row)...")
plot_prior_check(idata_prior)
else:
print("idata_prior not found. Did the model run completely?")
Generating journal-ready prior predictive plots (3-panel row)...
MCMC sampling and diagnostics¶
We sampled the posterior using NUTS in PyMC with 4 chains, 3 000 tuning iterations, and 1 500 retained draws per chain (total posterior draws = 6 000), with target_accept = 0.999 and init = "adapt_diag".
Convergence diagnostics:
| Diagnostic | Value |
|---|---|
| Divergences | 0 |
| Max $\widehat{R}$ | 1.0028 |
Effective sample size (ESS) — minimum across all elements of each parameter:
| Parameter | Bulk ESS (min) | Tail ESS (min) |
|---|---|---|
| $\delta_{c,k}$ | 4 682 | 4 615 |
| $\mu_{c,k}$ | 1 280 | 1 969 |
| $\sigma_\delta$ | 2 717 | 2 567 |
| $\sigma_{\text{batch}}$ | 2 566 | 3 826 |
| $\phi_k$ | 4 530 | 4 515 |
| $\alpha$ | 1 319 | 1 938 |
| $\beta_{\text{off}}$ | 4 609 | 4 119 |
| $\delta_0$ | 4 317 | 4 422 |
Zero divergences indicate stable Hamiltonian trajectories under the chosen parameterization and target acceptance rate. All $\widehat{R}$ values are below 1.01, consistent with convergence. The lowest bulk ESS (1 280 for $\mu_{c,k}$) still provides >200 effective samples per chain, adequate for reliable posterior summaries and 94% credible intervals.
In [6]:
Copied!
z_map = {"core": -0.5, "border": +0.5}
z = table_model["roi"].map(z_map).values.astype(float)
# --- factorize indices (you already have these)
ind_cells, cells_levels = table_model['cluster'].factorize()
ind_cond, cond_levels = table_model['condition'].factorize()
ind_batch, batch_levels = table_model['batch'].factorize()
ind_roi, roi_levels = table_model['roi'].factorize() # core/border only (2)
# NOTE: if you also have ROI id (tile/animal/slide), factorize that instead (see optional section below)
# data
logit_offset = table_model['area_odr'].values
offset_in = (logit_offset - np.mean(logit_offset)) / np.std(logit_offset)
y = table_model['positive'].values
N = table_model['total'].values
mean_N = np.mean(N)
coords = {
"obs_id": np.arange(len(y)),
"Cells": cells_levels,
"Condition": cond_levels,
"S": batch_levels,
"r": roi_levels, # here: r=core/border (2)
}
with pm.Model(coords=coords) as model6:
# indices
c = pm.Data("c", ind_cond, dims="obs_id")
k = pm.Data("k", ind_cells, dims="obs_id")
s = pm.Data("s", ind_batch, dims="obs_id")
r = pm.Data("r", ind_roi, dims="obs_id") # (core/border)
z_d = pm.Data("z", z, dims="obs_id") # -0.5/+0.5
# data
offset = pm.Data("offset", offset_in, dims="obs_id")
n = pm.Data("n", N, dims="obs_id")
Y = pm.Data("y", y, dims="obs_id")
# intercept
alpha = pm.Normal("alpha", 0, 1.5)
# condition × cell baseline (shared across regions)
mu_ck = pm.Normal("mu_ck", 0, 1.5, dims=("Condition", "Cells"))
delta0 = pm.Normal("delta0", 0, 1.0)
delta_c_raw = pm.Normal("delta_c_raw", 0, 0.7, dims=("Condition",))
delta_k_raw = pm.Normal("delta_k_raw", 0, 0.7, dims=("Cells",))
delta_c = pm.Deterministic("delta_c", delta_c_raw - pm.math.mean(delta_c_raw), dims=("Condition",))
delta_k = pm.Deterministic("delta_k", delta_k_raw - pm.math.mean(delta_k_raw), dims=("Cells",))
sigma_delta = pm.HalfNormal("sigma_delta", 0.25) # or even 0.2
delta_raw = pm.Normal("delta_raw", 0, 1, dims=("Condition", "Cells"))
delta_ck = pm.Deterministic(
"delta_ck",
delta0 + delta_c[:, None] + delta_k[None, :] + sigma_delta * delta_raw,
dims=("Condition", "Cells"),
)
# batch random intercept (non-centered)
sigma_batch = pm.HalfNormal("sigma_batch", 0.3)
u_batch_raw = pm.Normal("u_batch_raw", 0, 1, dims=("S",))
u_batch = pm.Deterministic("u_batch", u_batch_raw * sigma_batch, dims=("S",))
# optional small region baseline (core/border) if needed
# (usually not required if offset already captures geometry, but harmless)
#alpha_r = pm.Normal("alpha_r", 0, 0.5, dims=("r",))
# dispersion per cell type (your idea)
mu_logphi = pm.Normal("mu_logphi", mu=np.log(mean_N), sigma=0.5)
sigma_logphi = pm.HalfNormal("sigma_logphi", 0.3)
logphi_k_raw = pm.Normal("logphi_k_raw", mu=0, sigma=0.5, dims=("Cells",))
logphi_k = pm.Deterministic("logphi_k",mu_logphi + sigma_logphi * logphi_k_raw, dims=("Cells",))
phi_k = pm.Deterministic("phi_k", 1 + pm.math.exp(logphi_k), dims=("Cells",))
# linear predictor
beta_off = pm.Normal("beta_off", 0.0, 0.5)
eta = beta_off*offset + alpha + u_batch[s] + mu_ck[c, k] + delta_ck[c, k] * z_d
p = pm.Deterministic("p", pm.math.sigmoid(eta), dims="obs_id")
phi_obs = phi_k[k]
a = p * phi_obs
b = (1 - p) * phi_obs
count = pm.BetaBinomial("count", alpha=a, beta=b, n=n, observed=Y, dims="obs_id")
trace6 = pm.sample(
draws=1500,
tune=3000, # more adaptation
chains=4,
target_accept=0.999, # higher = fewer divergences
init="adapt_diag", # safer than jitter+adapt_diag sometimes
random_seed=RANDOM_SEED,
progressbar=True,
idata_kwargs={"log_likelihood": True},
)
z_map = {"core": -0.5, "border": +0.5}
z = table_model["roi"].map(z_map).values.astype(float)
# --- factorize indices (you already have these)
ind_cells, cells_levels = table_model['cluster'].factorize()
ind_cond, cond_levels = table_model['condition'].factorize()
ind_batch, batch_levels = table_model['batch'].factorize()
ind_roi, roi_levels = table_model['roi'].factorize() # core/border only (2)
# NOTE: if you also have ROI id (tile/animal/slide), factorize that instead (see optional section below)
# data
logit_offset = table_model['area_odr'].values
offset_in = (logit_offset - np.mean(logit_offset)) / np.std(logit_offset)
y = table_model['positive'].values
N = table_model['total'].values
mean_N = np.mean(N)
coords = {
"obs_id": np.arange(len(y)),
"Cells": cells_levels,
"Condition": cond_levels,
"S": batch_levels,
"r": roi_levels, # here: r=core/border (2)
}
with pm.Model(coords=coords) as model6:
# indices
c = pm.Data("c", ind_cond, dims="obs_id")
k = pm.Data("k", ind_cells, dims="obs_id")
s = pm.Data("s", ind_batch, dims="obs_id")
r = pm.Data("r", ind_roi, dims="obs_id") # (core/border)
z_d = pm.Data("z", z, dims="obs_id") # -0.5/+0.5
# data
offset = pm.Data("offset", offset_in, dims="obs_id")
n = pm.Data("n", N, dims="obs_id")
Y = pm.Data("y", y, dims="obs_id")
# intercept
alpha = pm.Normal("alpha", 0, 1.5)
# condition × cell baseline (shared across regions)
mu_ck = pm.Normal("mu_ck", 0, 1.5, dims=("Condition", "Cells"))
delta0 = pm.Normal("delta0", 0, 1.0)
delta_c_raw = pm.Normal("delta_c_raw", 0, 0.7, dims=("Condition",))
delta_k_raw = pm.Normal("delta_k_raw", 0, 0.7, dims=("Cells",))
delta_c = pm.Deterministic("delta_c", delta_c_raw - pm.math.mean(delta_c_raw), dims=("Condition",))
delta_k = pm.Deterministic("delta_k", delta_k_raw - pm.math.mean(delta_k_raw), dims=("Cells",))
sigma_delta = pm.HalfNormal("sigma_delta", 0.25) # or even 0.2
delta_raw = pm.Normal("delta_raw", 0, 1, dims=("Condition", "Cells"))
delta_ck = pm.Deterministic(
"delta_ck",
delta0 + delta_c[:, None] + delta_k[None, :] + sigma_delta * delta_raw,
dims=("Condition", "Cells"),
)
# batch random intercept (non-centered)
sigma_batch = pm.HalfNormal("sigma_batch", 0.3)
u_batch_raw = pm.Normal("u_batch_raw", 0, 1, dims=("S",))
u_batch = pm.Deterministic("u_batch", u_batch_raw * sigma_batch, dims=("S",))
# optional small region baseline (core/border) if needed
# (usually not required if offset already captures geometry, but harmless)
#alpha_r = pm.Normal("alpha_r", 0, 0.5, dims=("r",))
# dispersion per cell type (your idea)
mu_logphi = pm.Normal("mu_logphi", mu=np.log(mean_N), sigma=0.5)
sigma_logphi = pm.HalfNormal("sigma_logphi", 0.3)
logphi_k_raw = pm.Normal("logphi_k_raw", mu=0, sigma=0.5, dims=("Cells",))
logphi_k = pm.Deterministic("logphi_k",mu_logphi + sigma_logphi * logphi_k_raw, dims=("Cells",))
phi_k = pm.Deterministic("phi_k", 1 + pm.math.exp(logphi_k), dims=("Cells",))
# linear predictor
beta_off = pm.Normal("beta_off", 0.0, 0.5)
eta = beta_off*offset + alpha + u_batch[s] + mu_ck[c, k] + delta_ck[c, k] * z_d
p = pm.Deterministic("p", pm.math.sigmoid(eta), dims="obs_id")
phi_obs = phi_k[k]
a = p * phi_obs
b = (1 - p) * phi_obs
count = pm.BetaBinomial("count", alpha=a, beta=b, n=n, observed=Y, dims="obs_id")
trace6 = pm.sample(
draws=1500,
tune=3000, # more adaptation
chains=4,
target_accept=0.999, # higher = fewer divergences
init="adapt_diag", # safer than jitter+adapt_diag sometimes
random_seed=RANDOM_SEED,
progressbar=True,
idata_kwargs={"log_likelihood": True},
)
Initializing NUTS using adapt_diag... Multiprocess sampling (4 chains in 4 jobs) NUTS: [alpha, mu_ck, delta0, delta_c_raw, delta_k_raw, sigma_delta, delta_raw, sigma_batch, u_batch_raw, mu_logphi, sigma_logphi, logphi_k_raw, beta_off]
Output()
Sampling 4 chains for 3_000 tune and 1_500 draw iterations (12_000 + 6_000 draws total) took 657 seconds.
In [17]:
Copied!
# ── MCMC diagnostic plots ──────────────────────────────────────────────────
# Requires: trace6 from the sampling cell above
key_params = ["delta_ck", "mu_ck", "sigma_delta", "sigma_batch",
"phi_k", "alpha", "beta_off", "delta0"]
# ── 1. Trace plots for selected scalar parameters ─────────────────────────
scalar_params = ["alpha", "beta_off", "delta0", "sigma_delta", "sigma_batch"]
axes_trace = az.plot_trace(trace6, var_names=scalar_params, compact=True,
figsize=(10, 1.6 * len(scalar_params)))
plt.suptitle("Trace plots — key scalar parameters", fontsize=11, fontweight="bold", y=1.01)
plt.tight_layout()
plt.savefig("fig_trace_diagnostics.pdf", bbox_inches="tight")
plt.show()
# ── 2. Energy plot (NUTS diagnostic) ──────────────────────────────────────
fig, ax = plt.subplots(figsize=(5, 3), dpi=300)
az.plot_energy(trace6, ax=ax)
ax.set_title("NUTS energy diagnostic", fontsize=10, fontweight="bold")
plt.tight_layout()
plt.savefig("fig_energy_diagnostic.pdf", bbox_inches="tight")
plt.show()
# ── 3. Forest plot with R-hat ─────────────────────────────────────────────
az.plot_forest(trace6, var_names=scalar_params, kind="forestplot",
r_hat=True, ess=True, combined=True, figsize=(10, 4))
plt.suptitle("Forest plot with $\\widehat{R}$ and ESS", fontsize=11, fontweight="bold", y=1.01)
plt.tight_layout()
plt.savefig("fig_rhat_ess_diagnostics.pdf", bbox_inches="tight")
plt.show()
# ── 4. ESS evolution plot ─────────────────────────────────────────────────
axes_ess = az.plot_ess(trace6, var_names=scalar_params, kind="evolution",
figsize=(12, 3))
plt.suptitle("Bulk ESS evolution across draws", fontsize=11, fontweight="bold", y=1.02)
plt.tight_layout()
plt.savefig("fig_ess_evolution.pdf", bbox_inches="tight")
plt.show()
print("Saved: fig_trace_diagnostics.pdf, fig_energy_diagnostic.pdf,")
print(" fig_rhat_ess_diagnostics.pdf, fig_ess_evolution.pdf")
# ── MCMC diagnostic plots ──────────────────────────────────────────────────
# Requires: trace6 from the sampling cell above
key_params = ["delta_ck", "mu_ck", "sigma_delta", "sigma_batch",
"phi_k", "alpha", "beta_off", "delta0"]
# ── 1. Trace plots for selected scalar parameters ─────────────────────────
scalar_params = ["alpha", "beta_off", "delta0", "sigma_delta", "sigma_batch"]
axes_trace = az.plot_trace(trace6, var_names=scalar_params, compact=True,
figsize=(10, 1.6 * len(scalar_params)))
plt.suptitle("Trace plots — key scalar parameters", fontsize=11, fontweight="bold", y=1.01)
plt.tight_layout()
plt.savefig("fig_trace_diagnostics.pdf", bbox_inches="tight")
plt.show()
# ── 2. Energy plot (NUTS diagnostic) ──────────────────────────────────────
fig, ax = plt.subplots(figsize=(5, 3), dpi=300)
az.plot_energy(trace6, ax=ax)
ax.set_title("NUTS energy diagnostic", fontsize=10, fontweight="bold")
plt.tight_layout()
plt.savefig("fig_energy_diagnostic.pdf", bbox_inches="tight")
plt.show()
# ── 3. Forest plot with R-hat ─────────────────────────────────────────────
az.plot_forest(trace6, var_names=scalar_params, kind="forestplot",
r_hat=True, ess=True, combined=True, figsize=(10, 4))
plt.suptitle("Forest plot with $\\widehat{R}$ and ESS", fontsize=11, fontweight="bold", y=1.01)
plt.tight_layout()
plt.savefig("fig_rhat_ess_diagnostics.pdf", bbox_inches="tight")
plt.show()
# ── 4. ESS evolution plot ─────────────────────────────────────────────────
axes_ess = az.plot_ess(trace6, var_names=scalar_params, kind="evolution",
figsize=(12, 3))
plt.suptitle("Bulk ESS evolution across draws", fontsize=11, fontweight="bold", y=1.02)
plt.tight_layout()
plt.savefig("fig_ess_evolution.pdf", bbox_inches="tight")
plt.show()
print("Saved: fig_trace_diagnostics.pdf, fig_energy_diagnostic.pdf,")
print(" fig_rhat_ess_diagnostics.pdf, fig_ess_evolution.pdf")
Saved: fig_trace_diagnostics.pdf, fig_energy_diagnostic.pdf,
fig_rhat_ess_diagnostics.pdf, fig_ess_evolution.pdf
In [ ]:
Copied!
# ── Model fit summary and interpretation guidance ──────────────────────────
# Requires: trace6, model6, ind_cond, ind_cells, cond_levels, cells_levels
from scipy.special import expit # logistic / inverse-logit
# ── 1. MCMC diagnostics ───────────────────────────────────────────────────
print("=" * 72)
print("MCMC DIAGNOSTICS (model6 — narrow-prior Beta–Binomial)")
print("=" * 72)
# Divergences
n_divs = int(trace6.sample_stats["diverging"].sum())
print(f"\nDivergences : {n_divs}")
# Rhat
rhat_all = az.rhat(trace6)
rhat_vals = np.concatenate([v.values.ravel() for v in rhat_all.data_vars.values()])
rhat_max = float(np.nanmax(rhat_vals))
print(f"Max R-hat : {rhat_max:.4f}")
# ESS (bulk & tail) for key parameters
key_params = ["delta_ck", "mu_ck", "sigma_delta", "sigma_batch", "phi_k",
"alpha", "beta_off", "delta0"]
ess_bulk = az.ess(trace6, var_names=key_params, method="bulk")
ess_tail = az.ess(trace6, var_names=key_params, method="tail")
print("\n── Effective sample size (bulk / tail) for key parameters ──")
for var in key_params:
b = ess_bulk[var].values
t = ess_tail[var].values
print(f" {var:>16s} bulk min={float(np.nanmin(b)):7.0f} "
f"tail min={float(np.nanmin(t)):7.0f}")
# ── 2. δ_{c,k} summary with odds ratios ───────────────────────────────────
print("\n" + "=" * 72)
print("δ_{c,k} POSTERIOR SUMMARY (border-minus-core log-odds)")
print("=" * 72)
delta_ck_draws = trace6.posterior["delta_ck"].values # (chain, draw, C, K)
delta_flat = delta_ck_draws.reshape(-1, *delta_ck_draws.shape[2:]) # (samples, C, K)
rows = []
for ci, cond in enumerate(cond_levels):
for ki, cell in enumerate(cells_levels):
d = delta_flat[:, ci, ki]
m = float(np.mean(d))
hdi = az.hdi(d, hdi_prob=0.94)
odr = float(np.exp(m))
odr_lo = float(np.exp(hdi[0]))
odr_hi = float(np.exp(hdi[1]))
rows.append({
"Condition": cond, "Cell type": cell,
"δ_mean": round(m, 4),
"δ_HDI_lo": round(float(hdi[0]), 4),
"δ_HDI_hi": round(float(hdi[1]), 4),
"OR (border/core)": round(odr, 3),
"OR_lo": round(odr_lo, 3),
"OR_hi": round(odr_hi, 3),
})
df_delta = pd.DataFrame(rows)
print(df_delta.to_string(index=False))
# ── 3. Predicted border / core probabilities ──────────────────────────────
print("\n" + "=" * 72)
print("PREDICTED PROBABILITIES p_border, p_core (posterior mean ± 94 % HDI)")
print("=" * 72)
# A_{c,k} = alpha + mean(u_batch) + mu_ck (region-invariant baseline)
alpha_draws = trace6.posterior["alpha"].values.ravel() # (samples,)
mu_ck_draws = trace6.posterior["mu_ck"].values # (chain, draw, C, K)
mu_ck_flat = mu_ck_draws.reshape(-1, *mu_ck_draws.shape[2:]) # (samples, C, K)
u_batch_draws = trace6.posterior["u_batch"].values # (chain, draw, S)
u_batch_mean = u_batch_draws.reshape(-1, u_batch_draws.shape[-1]).mean(axis=1) # (samples,)
rows_prob = []
for ci, cond in enumerate(cond_levels):
for ki, cell in enumerate(cells_levels):
A = alpha_draws + u_batch_mean + mu_ck_flat[:, ci, ki]
d = delta_flat[:, ci, ki]
eta_border = A + 0.5 * d
eta_core = A - 0.5 * d
p_border = expit(eta_border)
p_core = expit(eta_core)
pb_m, pb_hdi = float(np.mean(p_border)), az.hdi(p_border, hdi_prob=0.94)
pc_m, pc_hdi = float(np.mean(p_core)), az.hdi(p_core, hdi_prob=0.94)
rows_prob.append({
"Condition": cond, "Cell type": cell,
"p_border_mean": round(pb_m, 4),
"p_border_lo": round(float(pb_hdi[0]), 4),
"p_border_hi": round(float(pb_hdi[1]), 4),
"p_core_mean": round(pc_m, 4),
"p_core_lo": round(float(pc_hdi[0]), 4),
"p_core_hi": round(float(pc_hdi[1]), 4),
})
df_prob = pd.DataFrame(rows_prob)
print(df_prob.to_string(index=False))
# ── 4. Publication figure: forest plot of δ_{c,k} with OR axis ────────────
n_cond = len(cond_levels)
n_cell = len(cells_levels)
n_rows_plot = n_cond * n_cell
fig, ax = plt.subplots(figsize=(6, max(4, 0.35 * n_rows_plot)), dpi=300)
plt.rcParams.update({"font.size": 8})
yticks, ylabels = [], []
cmap = plt.cm.Set2
for idx, (ci, cond) in enumerate(enumerate(cond_levels)):
for ki, cell in enumerate(cells_levels):
row_idx = ci * n_cell + ki
d = delta_flat[:, ci, ki]
m = float(np.mean(d))
hdi = az.hdi(d, hdi_prob=0.94)
color = cmap(ci / max(n_cond - 1, 1))
ax.plot([hdi[0], hdi[1]], [row_idx, row_idx], lw=2, color=color, solid_capstyle="round")
ax.plot(m, row_idx, "o", ms=4, color=color, zorder=5)
yticks.append(row_idx)
ylabels.append(f"{cond} | {cell}")
ax.axvline(0, ls="--", lw=0.8, color="k", zorder=0)
ax.set_yticks(yticks)
ax.set_yticklabels(ylabels, fontsize=7)
ax.set_xlabel("$\\delta_{c,k}$ (border − core log-odds)", fontsize=9)
ax.set_title("Posterior $\\delta_{c,k}$ with 94 % HDI", fontsize=10, fontweight="bold")
ax.invert_yaxis()
# Secondary x-axis for odds ratio
ax2 = ax.twiny()
lo, hi = ax.get_xlim()
ax2.set_xlim(np.exp(lo), np.exp(hi))
ax2.set_xscale("log")
ax2.set_xlabel("Odds ratio (border / core)", fontsize=9)
for sp in ("top", "right"):
ax.spines[sp].set_visible(False)
ax.spines["left"].set_linewidth(0.5)
ax.spines["bottom"].set_linewidth(0.5)
plt.tight_layout()
plt.savefig("fig_delta_ck_forest_model6.pdf", bbox_inches="tight")
plt.show()
# Save tables
df_delta.to_csv("model6_delta_ck_summary.csv", index=False)
df_prob.to_csv("model6_predicted_probabilities.csv", index=False)
print("\nSaved: model6_delta_ck_summary.csv, model6_predicted_probabilities.csv")
# ── Model fit summary and interpretation guidance ──────────────────────────
# Requires: trace6, model6, ind_cond, ind_cells, cond_levels, cells_levels
from scipy.special import expit # logistic / inverse-logit
# ── 1. MCMC diagnostics ───────────────────────────────────────────────────
print("=" * 72)
print("MCMC DIAGNOSTICS (model6 — narrow-prior Beta–Binomial)")
print("=" * 72)
# Divergences
n_divs = int(trace6.sample_stats["diverging"].sum())
print(f"\nDivergences : {n_divs}")
# Rhat
rhat_all = az.rhat(trace6)
rhat_vals = np.concatenate([v.values.ravel() for v in rhat_all.data_vars.values()])
rhat_max = float(np.nanmax(rhat_vals))
print(f"Max R-hat : {rhat_max:.4f}")
# ESS (bulk & tail) for key parameters
key_params = ["delta_ck", "mu_ck", "sigma_delta", "sigma_batch", "phi_k",
"alpha", "beta_off", "delta0"]
ess_bulk = az.ess(trace6, var_names=key_params, method="bulk")
ess_tail = az.ess(trace6, var_names=key_params, method="tail")
print("\n── Effective sample size (bulk / tail) for key parameters ──")
for var in key_params:
b = ess_bulk[var].values
t = ess_tail[var].values
print(f" {var:>16s} bulk min={float(np.nanmin(b)):7.0f} "
f"tail min={float(np.nanmin(t)):7.0f}")
# ── 2. δ_{c,k} summary with odds ratios ───────────────────────────────────
print("\n" + "=" * 72)
print("δ_{c,k} POSTERIOR SUMMARY (border-minus-core log-odds)")
print("=" * 72)
delta_ck_draws = trace6.posterior["delta_ck"].values # (chain, draw, C, K)
delta_flat = delta_ck_draws.reshape(-1, *delta_ck_draws.shape[2:]) # (samples, C, K)
rows = []
for ci, cond in enumerate(cond_levels):
for ki, cell in enumerate(cells_levels):
d = delta_flat[:, ci, ki]
m = float(np.mean(d))
hdi = az.hdi(d, hdi_prob=0.94)
odr = float(np.exp(m))
odr_lo = float(np.exp(hdi[0]))
odr_hi = float(np.exp(hdi[1]))
rows.append({
"Condition": cond, "Cell type": cell,
"δ_mean": round(m, 4),
"δ_HDI_lo": round(float(hdi[0]), 4),
"δ_HDI_hi": round(float(hdi[1]), 4),
"OR (border/core)": round(odr, 3),
"OR_lo": round(odr_lo, 3),
"OR_hi": round(odr_hi, 3),
})
df_delta = pd.DataFrame(rows)
print(df_delta.to_string(index=False))
# ── 3. Predicted border / core probabilities ──────────────────────────────
print("\n" + "=" * 72)
print("PREDICTED PROBABILITIES p_border, p_core (posterior mean ± 94 % HDI)")
print("=" * 72)
# A_{c,k} = alpha + mean(u_batch) + mu_ck (region-invariant baseline)
alpha_draws = trace6.posterior["alpha"].values.ravel() # (samples,)
mu_ck_draws = trace6.posterior["mu_ck"].values # (chain, draw, C, K)
mu_ck_flat = mu_ck_draws.reshape(-1, *mu_ck_draws.shape[2:]) # (samples, C, K)
u_batch_draws = trace6.posterior["u_batch"].values # (chain, draw, S)
u_batch_mean = u_batch_draws.reshape(-1, u_batch_draws.shape[-1]).mean(axis=1) # (samples,)
rows_prob = []
for ci, cond in enumerate(cond_levels):
for ki, cell in enumerate(cells_levels):
A = alpha_draws + u_batch_mean + mu_ck_flat[:, ci, ki]
d = delta_flat[:, ci, ki]
eta_border = A + 0.5 * d
eta_core = A - 0.5 * d
p_border = expit(eta_border)
p_core = expit(eta_core)
pb_m, pb_hdi = float(np.mean(p_border)), az.hdi(p_border, hdi_prob=0.94)
pc_m, pc_hdi = float(np.mean(p_core)), az.hdi(p_core, hdi_prob=0.94)
rows_prob.append({
"Condition": cond, "Cell type": cell,
"p_border_mean": round(pb_m, 4),
"p_border_lo": round(float(pb_hdi[0]), 4),
"p_border_hi": round(float(pb_hdi[1]), 4),
"p_core_mean": round(pc_m, 4),
"p_core_lo": round(float(pc_hdi[0]), 4),
"p_core_hi": round(float(pc_hdi[1]), 4),
})
df_prob = pd.DataFrame(rows_prob)
print(df_prob.to_string(index=False))
# ── 4. Publication figure: forest plot of δ_{c,k} with OR axis ────────────
n_cond = len(cond_levels)
n_cell = len(cells_levels)
n_rows_plot = n_cond * n_cell
fig, ax = plt.subplots(figsize=(6, max(4, 0.35 * n_rows_plot)), dpi=300)
plt.rcParams.update({"font.size": 8})
yticks, ylabels = [], []
cmap = plt.cm.Set2
for idx, (ci, cond) in enumerate(enumerate(cond_levels)):
for ki, cell in enumerate(cells_levels):
row_idx = ci * n_cell + ki
d = delta_flat[:, ci, ki]
m = float(np.mean(d))
hdi = az.hdi(d, hdi_prob=0.94)
color = cmap(ci / max(n_cond - 1, 1))
ax.plot([hdi[0], hdi[1]], [row_idx, row_idx], lw=2, color=color, solid_capstyle="round")
ax.plot(m, row_idx, "o", ms=4, color=color, zorder=5)
yticks.append(row_idx)
ylabels.append(f"{cond} | {cell}")
ax.axvline(0, ls="--", lw=0.8, color="k", zorder=0)
ax.set_yticks(yticks)
ax.set_yticklabels(ylabels, fontsize=7)
ax.set_xlabel("$\\delta_{c,k}$ (border − core log-odds)", fontsize=9)
ax.set_title("Posterior $\\delta_{c,k}$ with 94 % HDI", fontsize=10, fontweight="bold")
ax.invert_yaxis()
# Secondary x-axis for odds ratio
ax2 = ax.twiny()
lo, hi = ax.get_xlim()
ax2.set_xlim(np.exp(lo), np.exp(hi))
ax2.set_xscale("log")
ax2.set_xlabel("Odds ratio (border / core)", fontsize=9)
for sp in ("top", "right"):
ax.spines[sp].set_visible(False)
ax.spines["left"].set_linewidth(0.5)
ax.spines["bottom"].set_linewidth(0.5)
plt.tight_layout()
plt.savefig("fig_delta_ck_forest_model6.pdf", bbox_inches="tight")
plt.show()
# Save tables
df_delta.to_csv("model6_delta_ck_summary.csv", index=False)
df_prob.to_csv("model6_predicted_probabilities.csv", index=False)
print("\nSaved: model6_delta_ck_summary.csv, model6_predicted_probabilities.csv")
========================================================================
MCMC DIAGNOSTICS (model6 — narrow-prior Beta–Binomial)
========================================================================
Divergences : 0
Max R-hat : 1.0028
── Effective sample size (bulk / tail) for key parameters ──
delta_ck bulk min= 4682 tail min= 4615
mu_ck bulk min= 1280 tail min= 1969
sigma_delta bulk min= 2717 tail min= 2567
sigma_batch bulk min= 2566 tail min= 3826
phi_k bulk min= 4530 tail min= 4515
alpha bulk min= 1319 tail min= 1938
beta_off bulk min= 4609 tail min= 4119
delta0 bulk min= 4317 tail min= 4422
========================================================================
δ_{c,k} POSTERIOR SUMMARY (border-minus-core log-odds)
========================================================================
Condition Cell type δ_mean δ_HDI_lo δ_HDI_hi OR (border/core) OR_lo OR_hi
IL-15 DC 0.6631 0.3207 1.0040 1.941 1.378 2.729
IL-15 M1 TAM 0.2520 -0.1408 0.6598 1.287 0.869 1.934
IL-15 M2 TAM 0.8056 0.4941 1.1329 2.238 1.639 3.105
IL-15 Cytotoxic T-cells 0.7031 0.3812 0.9796 2.020 1.464 2.663
IL-15 Proliferating T-cells 0.4988 0.1890 0.8187 1.647 1.208 2.268
Both DC -0.1652 -0.5837 0.2300 0.848 0.558 1.259
Both M1 TAM -0.3983 -0.8143 -0.0006 0.671 0.443 0.999
Both M2 TAM 0.0776 -0.2656 0.4657 1.081 0.767 1.593
Both Cytotoxic T-cells -0.0956 -0.4545 0.2832 0.909 0.635 1.327
Both Proliferating T-cells -0.2500 -0.6524 0.1317 0.779 0.521 1.141
TCR DC -0.4747 -0.8844 -0.0850 0.622 0.413 0.918
TCR M1 TAM -0.6940 -1.1384 -0.3223 0.500 0.320 0.724
TCR M2 TAM 0.0199 -0.3667 0.3828 1.020 0.693 1.466
TCR Cytotoxic T-cells -0.3157 -0.6544 0.0225 0.729 0.520 1.023
TCR Proliferating T-cells -0.4818 -0.8946 -0.0430 0.618 0.409 0.958
========================================================================
PREDICTED PROBABILITIES p_border, p_core (posterior mean ± 94 % HDI)
========================================================================
Condition Cell type p_border_mean p_border_lo p_border_hi p_core_mean p_core_lo p_core_hi
IL-15 DC 0.6100 0.5037 0.7230 0.4488 0.3354 0.5652
IL-15 M1 TAM 0.4457 0.3197 0.5714 0.3856 0.2697 0.5052
IL-15 M2 TAM 0.4837 0.3690 0.5957 0.2975 0.2016 0.3957
IL-15 Cytotoxic T-cells 0.5093 0.3983 0.6228 0.3417 0.2401 0.4477
IL-15 Proliferating T-cells 0.5239 0.4133 0.6366 0.4023 0.2982 0.5163
Both DC 0.5829 0.4247 0.7242 0.6214 0.4825 0.7760
Both M1 TAM 0.3637 0.2270 0.5120 0.4572 0.3022 0.6029
Both M2 TAM 0.3652 0.2321 0.5077 0.3479 0.2226 0.4965
Both Cytotoxic T-cells 0.3975 0.2587 0.5445 0.4201 0.2737 0.5691
Both Proliferating T-cells 0.4398 0.2872 0.5931 0.5003 0.3504 0.6574
TCR DC 0.5203 0.3992 0.6311 0.6335 0.5064 0.7411
TCR M1 TAM 0.2728 0.1764 0.3685 0.4263 0.2975 0.5451
TCR M2 TAM 0.3544 0.2506 0.4572 0.3502 0.2460 0.4677
TCR Cytotoxic T-cells 0.4024 0.2968 0.5110 0.4790 0.3683 0.6033
TCR Proliferating T-cells 0.4358 0.3214 0.5664 0.5535 0.4215 0.6837
Saved: model6_delta_ck_summary.csv, model6_predicted_probabilities.csv
Posterior predictive checks (PPC) and calibration-style checks¶
Posterior predictive checks assessed whether the fitted model reproduces observed fractions and dispersion patterns. For each posterior draw, we generated replicated counts (y_i^{\text{rep}}) from the Beta–Binomial likelihood and summarized:
- Observed vs predicted mean fraction: scatter of (y_i/n_i) versus (\mathbb{E}[y_i^{\text{rep}}/n_i\mid \text{data}]), with the identity line. Systematic curvature indicates mis-specification of the link/linear predictor.
- Residuals vs (n_i): ( \mathbb{E}[y_i^{\text{rep}}/n_i]-y_i/n_i ) versus (n_i) to detect size-dependent bias (e.g., underfitting at small strata sizes).
- Predictive sharpness / uncertainty: distribution of 94% HDI widths for (y_i^{\text{rep}}/n_i) (narrower indicates more informative predictions; excessively narrow intervals can signal underdispersion).
(Optional calibration-style checks)
- Coverage of posterior predictive intervals for (y_i/n_i) stratified by (n_i) bins (e.g., small/medium/large (n)).
- Posterior predictive distribution of counts (not only fractions) for low-(n) strata, where discretization is strongest.
(Insert your PPC figure references here: [PPC figure placeholder].)
In [7]:
Copied!
# ── Posterior predictive checks (PPC) and calibration-style checks ──────────
# Requires: model6, trace6, y, N (from the sampling cell above)
# 1. Sample posterior predictive
with model6:
ppc6 = pm.sample_posterior_predictive(trace6, random_seed=RANDOM_SEED)
# Observed fraction
obs_frac = y / N
# Posterior-predictive replicated fractions (draws × obs)
y_rep = ppc6.posterior_predictive["count"].values # (chain, draw, obs)
y_rep_flat = y_rep.reshape(-1, y_rep.shape[-1]) # (samples, obs)
rep_frac = y_rep_flat / N[None, :]
# Posterior mean predicted fraction per observation
pred_mean = rep_frac.mean(axis=0)
# Residuals
resid = pred_mean - obs_frac
# 94 % HDI width per observation
hdi_lo = np.percentile(rep_frac, 3, axis=0)
hdi_hi = np.percentile(rep_frac, 97, axis=0)
hdi_width = hdi_hi - hdi_lo
# ── Publication-quality figure (3-panel row) ───────────────────────────────
fig, axes = plt.subplots(1, 3, figsize=(7.2, 2.8), dpi=300)
plt.rcParams.update({"font.size": 7, "axes.labelsize": 8, "axes.titlesize": 9})
# Panel a: Observed vs predicted mean fraction
ax = axes[0]
ax.scatter(obs_frac, pred_mean, s=8, alpha=0.45, edgecolors="none", c="steelblue")
lims = [0, max(obs_frac.max(), pred_mean.max()) * 1.05]
ax.plot(lims, lims, ls="--", lw=0.8, c="k", zorder=0)
ax.set_xlabel("Observed $y_i / n_i$", fontsize=8)
ax.set_ylabel("Posterior mean $\\hat{y}_i / n_i$", fontsize=8)
ax.set_title("a. Observed vs. predicted fraction", fontsize=9, fontweight="bold")
ax.set_xlim(lims); ax.set_ylim(lims)
for sp in ("top", "right"): ax.spines[sp].set_visible(False)
# Panel b: Residuals vs n_i
ax = axes[1]
ax.scatter(N, resid, s=8, alpha=0.45, edgecolors="none", c="darkorange")
ax.axhline(0, ls="--", lw=0.8, c="k")
ax.set_xlabel("Total count $n_i$", fontsize=8)
ax.set_ylabel("Residual ($\\hat{y}/n - y/n$)", fontsize=8)
ax.set_title("b. Residuals vs. sample size", fontsize=9, fontweight="bold")
for sp in ("top", "right"): ax.spines[sp].set_visible(False)
# Panel c: HDI width histogram
ax = axes[2]
ax.hist(hdi_width, bins=40, density=True, color="mediumpurple", alpha=0.75, ec="white", lw=0.3)
ax.set_xlabel("94 % HDI width ($y^{\\mathrm{rep}}/n$)", fontsize=8)
ax.set_ylabel("Density", fontsize=8)
ax.set_title("c. Predictive sharpness", fontsize=9, fontweight="bold")
for sp in ("top", "right"): ax.spines[sp].set_visible(False)
plt.tight_layout()
plt.savefig("fig_ppc_model6.pdf", bbox_inches="tight")
plt.show()
# ── Coverage summary (optional calibration check) ─────────────────────────
coverage_94 = np.mean((obs_frac >= hdi_lo) & (obs_frac <= hdi_hi))
print(f"Empirical coverage of 94 % posterior-predictive interval: {coverage_94:.3f}")
# Stratified by n bins
n_bins = pd.qcut(N, q=3, labels=["small n", "medium n", "large n"])
for lab in ["small n", "medium n", "large n"]:
mask = n_bins == lab
cov = np.mean((obs_frac[mask] >= hdi_lo[mask]) & (obs_frac[mask] <= hdi_hi[mask]))
print(f" {lab:>10s} (n obs = {mask.sum():4d}): coverage = {cov:.3f}")
# ── Posterior predictive checks (PPC) and calibration-style checks ──────────
# Requires: model6, trace6, y, N (from the sampling cell above)
# 1. Sample posterior predictive
with model6:
ppc6 = pm.sample_posterior_predictive(trace6, random_seed=RANDOM_SEED)
# Observed fraction
obs_frac = y / N
# Posterior-predictive replicated fractions (draws × obs)
y_rep = ppc6.posterior_predictive["count"].values # (chain, draw, obs)
y_rep_flat = y_rep.reshape(-1, y_rep.shape[-1]) # (samples, obs)
rep_frac = y_rep_flat / N[None, :]
# Posterior mean predicted fraction per observation
pred_mean = rep_frac.mean(axis=0)
# Residuals
resid = pred_mean - obs_frac
# 94 % HDI width per observation
hdi_lo = np.percentile(rep_frac, 3, axis=0)
hdi_hi = np.percentile(rep_frac, 97, axis=0)
hdi_width = hdi_hi - hdi_lo
# ── Publication-quality figure (3-panel row) ───────────────────────────────
fig, axes = plt.subplots(1, 3, figsize=(7.2, 2.8), dpi=300)
plt.rcParams.update({"font.size": 7, "axes.labelsize": 8, "axes.titlesize": 9})
# Panel a: Observed vs predicted mean fraction
ax = axes[0]
ax.scatter(obs_frac, pred_mean, s=8, alpha=0.45, edgecolors="none", c="steelblue")
lims = [0, max(obs_frac.max(), pred_mean.max()) * 1.05]
ax.plot(lims, lims, ls="--", lw=0.8, c="k", zorder=0)
ax.set_xlabel("Observed $y_i / n_i$", fontsize=8)
ax.set_ylabel("Posterior mean $\\hat{y}_i / n_i$", fontsize=8)
ax.set_title("a. Observed vs. predicted fraction", fontsize=9, fontweight="bold")
ax.set_xlim(lims); ax.set_ylim(lims)
for sp in ("top", "right"): ax.spines[sp].set_visible(False)
# Panel b: Residuals vs n_i
ax = axes[1]
ax.scatter(N, resid, s=8, alpha=0.45, edgecolors="none", c="darkorange")
ax.axhline(0, ls="--", lw=0.8, c="k")
ax.set_xlabel("Total count $n_i$", fontsize=8)
ax.set_ylabel("Residual ($\\hat{y}/n - y/n$)", fontsize=8)
ax.set_title("b. Residuals vs. sample size", fontsize=9, fontweight="bold")
for sp in ("top", "right"): ax.spines[sp].set_visible(False)
# Panel c: HDI width histogram
ax = axes[2]
ax.hist(hdi_width, bins=40, density=True, color="mediumpurple", alpha=0.75, ec="white", lw=0.3)
ax.set_xlabel("94 % HDI width ($y^{\\mathrm{rep}}/n$)", fontsize=8)
ax.set_ylabel("Density", fontsize=8)
ax.set_title("c. Predictive sharpness", fontsize=9, fontweight="bold")
for sp in ("top", "right"): ax.spines[sp].set_visible(False)
plt.tight_layout()
plt.savefig("fig_ppc_model6.pdf", bbox_inches="tight")
plt.show()
# ── Coverage summary (optional calibration check) ─────────────────────────
coverage_94 = np.mean((obs_frac >= hdi_lo) & (obs_frac <= hdi_hi))
print(f"Empirical coverage of 94 % posterior-predictive interval: {coverage_94:.3f}")
# Stratified by n bins
n_bins = pd.qcut(N, q=3, labels=["small n", "medium n", "large n"])
for lab in ["small n", "medium n", "large n"]:
mask = n_bins == lab
cov = np.mean((obs_frac[mask] >= hdi_lo[mask]) & (obs_frac[mask] <= hdi_hi[mask]))
print(f" {lab:>10s} (n obs = {mask.sum():4d}): coverage = {cov:.3f}")
Sampling: [count]
Output()
Empirical coverage of 94 % posterior-predictive interval: 0.976
small n (n obs = 43): coverage = 0.953
medium n (n obs = 42): coverage = 0.976
large n (n obs = 42): coverage = 1.000
Model fit summary and interpretation guidance¶
Primary estimand: (\delta_{c,k}). (\delta_{c,k}) measures the border-minus-core shift in collagen proximity on the log-odds scale, conditional on geometry and batch effects:
- (\delta_{c,k}>0): higher odds of collagen proximity in Border than Core.
- (\delta_{c,k}<0): higher odds in Core.
- (\exp(\delta_{c,k})): border/core odds ratio.
Baseline term: (\mu_{c,k}). (\mu_{c,k}) captures the region-averaged baseline log-odds for condition (c) and cell type (k) (in combination with (\alpha), batch, and offset). It is useful for interpreting absolute levels of proximity rather than border-core differences.
Predicted probabilities. For any observation (i), the model yields (p_i=\text{logit}^{-1}(\eta_i)), enabling reporting of posterior mean and credible intervals for predicted proximity fractions, and derived border/core probabilities: [ p_{\text{border}}=\text{logit}^{-1}!\left(A_i+\frac{\delta_{c,k}}{2}\right),\qquad p_{\text{core}}=\text{logit}^{-1}!\left(A_i-\frac{\delta_{c,k}}{2}\right). ]
In [8]:
Copied!
# ── Model fit summary and interpretation guidance ──────────────────────────
# Requires: trace6, model6, ind_cond, ind_cells, cond_levels, cells_levels
from scipy.special import expit # logistic / inverse-logit
# ── 1. MCMC diagnostics ───────────────────────────────────────────────────
print("=" * 72)
print("MCMC DIAGNOSTICS (model6 — narrow-prior Beta–Binomial)")
print("=" * 72)
# Divergences
n_divs = int(trace6.sample_stats["diverging"].sum())
print(f"\nDivergences : {n_divs}")
# Rhat
rhat_all = az.rhat(trace6)
rhat_vals = np.concatenate([v.values.ravel() for v in rhat_all.data_vars.values()])
rhat_max = float(np.nanmax(rhat_vals))
print(f"Max R-hat : {rhat_max:.4f}")
# ESS (bulk & tail) for key parameters
key_params = ["delta_ck", "mu_ck", "sigma_delta", "sigma_batch", "phi_k",
"alpha", "beta_off", "delta0"]
ess_bulk = az.ess(trace6, var_names=key_params, method="bulk")
ess_tail = az.ess(trace6, var_names=key_params, method="tail")
print("\n── Effective sample size (bulk / tail) for key parameters ──")
for var in key_params:
b = ess_bulk[var].values
t = ess_tail[var].values
print(f" {var:>16s} bulk min={float(np.nanmin(b)):7.0f} "
f"tail min={float(np.nanmin(t)):7.0f}")
# ── 2. δ_{c,k} summary with odds ratios ───────────────────────────────────
print("\n" + "=" * 72)
print("δ_{c,k} POSTERIOR SUMMARY (border-minus-core log-odds)")
print("=" * 72)
delta_ck_draws = trace6.posterior["delta_ck"].values # (chain, draw, C, K)
delta_flat = delta_ck_draws.reshape(-1, *delta_ck_draws.shape[2:]) # (samples, C, K)
rows = []
for ci, cond in enumerate(cond_levels):
for ki, cell in enumerate(cells_levels):
d = delta_flat[:, ci, ki]
m = float(np.mean(d))
hdi = az.hdi(d, hdi_prob=0.94)
odr = float(np.exp(m))
odr_lo = float(np.exp(hdi[0]))
odr_hi = float(np.exp(hdi[1]))
rows.append({
"Condition": cond, "Cell type": cell,
"δ_mean": round(m, 4),
"δ_HDI_lo": round(float(hdi[0]), 4),
"δ_HDI_hi": round(float(hdi[1]), 4),
"OR (border/core)": round(odr, 3),
"OR_lo": round(odr_lo, 3),
"OR_hi": round(odr_hi, 3),
})
df_delta = pd.DataFrame(rows)
print(df_delta.to_string(index=False))
# ── 3. Predicted border / core probabilities ──────────────────────────────
print("\n" + "=" * 72)
print("PREDICTED PROBABILITIES p_border, p_core (posterior mean ± 94 % HDI)")
print("=" * 72)
# A_{c,k} = alpha + mean(u_batch) + mu_ck (region-invariant baseline)
alpha_draws = trace6.posterior["alpha"].values.ravel() # (samples,)
mu_ck_draws = trace6.posterior["mu_ck"].values # (chain, draw, C, K)
mu_ck_flat = mu_ck_draws.reshape(-1, *mu_ck_draws.shape[2:]) # (samples, C, K)
u_batch_draws = trace6.posterior["u_batch"].values # (chain, draw, S)
u_batch_mean = u_batch_draws.reshape(-1, u_batch_draws.shape[-1]).mean(axis=1) # (samples,)
rows_prob = []
for ci, cond in enumerate(cond_levels):
for ki, cell in enumerate(cells_levels):
A = alpha_draws + u_batch_mean + mu_ck_flat[:, ci, ki]
d = delta_flat[:, ci, ki]
eta_border = A + 0.5 * d
eta_core = A - 0.5 * d
p_border = expit(eta_border)
p_core = expit(eta_core)
pb_m, pb_hdi = float(np.mean(p_border)), az.hdi(p_border, hdi_prob=0.94)
pc_m, pc_hdi = float(np.mean(p_core)), az.hdi(p_core, hdi_prob=0.94)
rows_prob.append({
"Condition": cond, "Cell type": cell,
"p_border_mean": round(pb_m, 4),
"p_border_lo": round(float(pb_hdi[0]), 4),
"p_border_hi": round(float(pb_hdi[1]), 4),
"p_core_mean": round(pc_m, 4),
"p_core_lo": round(float(pc_hdi[0]), 4),
"p_core_hi": round(float(pc_hdi[1]), 4),
})
df_prob = pd.DataFrame(rows_prob)
print(df_prob.to_string(index=False))
# ── 4. Publication figure: forest plot of δ_{c,k} with OR axis ────────────
n_cond = len(cond_levels)
n_cell = len(cells_levels)
n_rows_plot = n_cond * n_cell
fig, ax = plt.subplots(figsize=(6, max(4, 0.35 * n_rows_plot)), dpi=300)
plt.rcParams.update({"font.size": 8})
yticks, ylabels = [], []
cmap = plt.cm.Set2
for idx, (ci, cond) in enumerate(enumerate(cond_levels)):
for ki, cell in enumerate(cells_levels):
row_idx = ci * n_cell + ki
d = delta_flat[:, ci, ki]
m = float(np.mean(d))
hdi = az.hdi(d, hdi_prob=0.94)
color = cmap(ci / max(n_cond - 1, 1))
ax.plot([hdi[0], hdi[1]], [row_idx, row_idx], lw=2, color=color, solid_capstyle="round")
ax.plot(m, row_idx, "o", ms=4, color=color, zorder=5)
yticks.append(row_idx)
ylabels.append(f"{cond} | {cell}")
ax.axvline(0, ls="--", lw=0.8, color="k", zorder=0)
ax.set_yticks(yticks)
ax.set_yticklabels(ylabels, fontsize=7)
ax.set_xlabel("$\\delta_{c,k}$ (border − core log-odds)", fontsize=9)
ax.set_title("Posterior $\\delta_{c,k}$ with 94 % HDI", fontsize=10, fontweight="bold")
ax.invert_yaxis()
# Secondary x-axis for odds ratio
ax2 = ax.twiny()
lo, hi = ax.get_xlim()
ax2.set_xlim(np.exp(lo), np.exp(hi))
ax2.set_xscale("log")
ax2.set_xlabel("Odds ratio (border / core)", fontsize=9)
for sp in ("top", "right"):
ax.spines[sp].set_visible(False)
ax.spines["left"].set_linewidth(0.5)
ax.spines["bottom"].set_linewidth(0.5)
plt.tight_layout()
plt.savefig("fig_delta_ck_forest_model6.pdf", bbox_inches="tight")
plt.show()
# Save tables
df_delta.to_csv("model6_delta_ck_summary.csv", index=False)
df_prob.to_csv("model6_predicted_probabilities.csv", index=False)
print("\nSaved: model6_delta_ck_summary.csv, model6_predicted_probabilities.csv")
# ── Model fit summary and interpretation guidance ──────────────────────────
# Requires: trace6, model6, ind_cond, ind_cells, cond_levels, cells_levels
from scipy.special import expit # logistic / inverse-logit
# ── 1. MCMC diagnostics ───────────────────────────────────────────────────
print("=" * 72)
print("MCMC DIAGNOSTICS (model6 — narrow-prior Beta–Binomial)")
print("=" * 72)
# Divergences
n_divs = int(trace6.sample_stats["diverging"].sum())
print(f"\nDivergences : {n_divs}")
# Rhat
rhat_all = az.rhat(trace6)
rhat_vals = np.concatenate([v.values.ravel() for v in rhat_all.data_vars.values()])
rhat_max = float(np.nanmax(rhat_vals))
print(f"Max R-hat : {rhat_max:.4f}")
# ESS (bulk & tail) for key parameters
key_params = ["delta_ck", "mu_ck", "sigma_delta", "sigma_batch", "phi_k",
"alpha", "beta_off", "delta0"]
ess_bulk = az.ess(trace6, var_names=key_params, method="bulk")
ess_tail = az.ess(trace6, var_names=key_params, method="tail")
print("\n── Effective sample size (bulk / tail) for key parameters ──")
for var in key_params:
b = ess_bulk[var].values
t = ess_tail[var].values
print(f" {var:>16s} bulk min={float(np.nanmin(b)):7.0f} "
f"tail min={float(np.nanmin(t)):7.0f}")
# ── 2. δ_{c,k} summary with odds ratios ───────────────────────────────────
print("\n" + "=" * 72)
print("δ_{c,k} POSTERIOR SUMMARY (border-minus-core log-odds)")
print("=" * 72)
delta_ck_draws = trace6.posterior["delta_ck"].values # (chain, draw, C, K)
delta_flat = delta_ck_draws.reshape(-1, *delta_ck_draws.shape[2:]) # (samples, C, K)
rows = []
for ci, cond in enumerate(cond_levels):
for ki, cell in enumerate(cells_levels):
d = delta_flat[:, ci, ki]
m = float(np.mean(d))
hdi = az.hdi(d, hdi_prob=0.94)
odr = float(np.exp(m))
odr_lo = float(np.exp(hdi[0]))
odr_hi = float(np.exp(hdi[1]))
rows.append({
"Condition": cond, "Cell type": cell,
"δ_mean": round(m, 4),
"δ_HDI_lo": round(float(hdi[0]), 4),
"δ_HDI_hi": round(float(hdi[1]), 4),
"OR (border/core)": round(odr, 3),
"OR_lo": round(odr_lo, 3),
"OR_hi": round(odr_hi, 3),
})
df_delta = pd.DataFrame(rows)
print(df_delta.to_string(index=False))
# ── 3. Predicted border / core probabilities ──────────────────────────────
print("\n" + "=" * 72)
print("PREDICTED PROBABILITIES p_border, p_core (posterior mean ± 94 % HDI)")
print("=" * 72)
# A_{c,k} = alpha + mean(u_batch) + mu_ck (region-invariant baseline)
alpha_draws = trace6.posterior["alpha"].values.ravel() # (samples,)
mu_ck_draws = trace6.posterior["mu_ck"].values # (chain, draw, C, K)
mu_ck_flat = mu_ck_draws.reshape(-1, *mu_ck_draws.shape[2:]) # (samples, C, K)
u_batch_draws = trace6.posterior["u_batch"].values # (chain, draw, S)
u_batch_mean = u_batch_draws.reshape(-1, u_batch_draws.shape[-1]).mean(axis=1) # (samples,)
rows_prob = []
for ci, cond in enumerate(cond_levels):
for ki, cell in enumerate(cells_levels):
A = alpha_draws + u_batch_mean + mu_ck_flat[:, ci, ki]
d = delta_flat[:, ci, ki]
eta_border = A + 0.5 * d
eta_core = A - 0.5 * d
p_border = expit(eta_border)
p_core = expit(eta_core)
pb_m, pb_hdi = float(np.mean(p_border)), az.hdi(p_border, hdi_prob=0.94)
pc_m, pc_hdi = float(np.mean(p_core)), az.hdi(p_core, hdi_prob=0.94)
rows_prob.append({
"Condition": cond, "Cell type": cell,
"p_border_mean": round(pb_m, 4),
"p_border_lo": round(float(pb_hdi[0]), 4),
"p_border_hi": round(float(pb_hdi[1]), 4),
"p_core_mean": round(pc_m, 4),
"p_core_lo": round(float(pc_hdi[0]), 4),
"p_core_hi": round(float(pc_hdi[1]), 4),
})
df_prob = pd.DataFrame(rows_prob)
print(df_prob.to_string(index=False))
# ── 4. Publication figure: forest plot of δ_{c,k} with OR axis ────────────
n_cond = len(cond_levels)
n_cell = len(cells_levels)
n_rows_plot = n_cond * n_cell
fig, ax = plt.subplots(figsize=(6, max(4, 0.35 * n_rows_plot)), dpi=300)
plt.rcParams.update({"font.size": 8})
yticks, ylabels = [], []
cmap = plt.cm.Set2
for idx, (ci, cond) in enumerate(enumerate(cond_levels)):
for ki, cell in enumerate(cells_levels):
row_idx = ci * n_cell + ki
d = delta_flat[:, ci, ki]
m = float(np.mean(d))
hdi = az.hdi(d, hdi_prob=0.94)
color = cmap(ci / max(n_cond - 1, 1))
ax.plot([hdi[0], hdi[1]], [row_idx, row_idx], lw=2, color=color, solid_capstyle="round")
ax.plot(m, row_idx, "o", ms=4, color=color, zorder=5)
yticks.append(row_idx)
ylabels.append(f"{cond} | {cell}")
ax.axvline(0, ls="--", lw=0.8, color="k", zorder=0)
ax.set_yticks(yticks)
ax.set_yticklabels(ylabels, fontsize=7)
ax.set_xlabel("$\\delta_{c,k}$ (border − core log-odds)", fontsize=9)
ax.set_title("Posterior $\\delta_{c,k}$ with 94 % HDI", fontsize=10, fontweight="bold")
ax.invert_yaxis()
# Secondary x-axis for odds ratio
ax2 = ax.twiny()
lo, hi = ax.get_xlim()
ax2.set_xlim(np.exp(lo), np.exp(hi))
ax2.set_xscale("log")
ax2.set_xlabel("Odds ratio (border / core)", fontsize=9)
for sp in ("top", "right"):
ax.spines[sp].set_visible(False)
ax.spines["left"].set_linewidth(0.5)
ax.spines["bottom"].set_linewidth(0.5)
plt.tight_layout()
plt.savefig("fig_delta_ck_forest_model6.pdf", bbox_inches="tight")
plt.show()
# Save tables
df_delta.to_csv("model6_delta_ck_summary.csv", index=False)
df_prob.to_csv("model6_predicted_probabilities.csv", index=False)
print("\nSaved: model6_delta_ck_summary.csv, model6_predicted_probabilities.csv")
========================================================================
MCMC DIAGNOSTICS (model6 — narrow-prior Beta–Binomial)
========================================================================
Divergences : 0
Max R-hat : 1.0028
── Effective sample size (bulk / tail) for key parameters ──
delta_ck bulk min= 4682 tail min= 4615
mu_ck bulk min= 1280 tail min= 1969
sigma_delta bulk min= 2717 tail min= 2567
sigma_batch bulk min= 2566 tail min= 3826
phi_k bulk min= 4530 tail min= 4515
alpha bulk min= 1319 tail min= 1938
beta_off bulk min= 4609 tail min= 4119
delta0 bulk min= 4317 tail min= 4422
========================================================================
δ_{c,k} POSTERIOR SUMMARY (border-minus-core log-odds)
========================================================================
Condition Cell type δ_mean δ_HDI_lo δ_HDI_hi OR (border/core) OR_lo OR_hi
IL-15 DC 0.6631 0.3207 1.0040 1.941 1.378 2.729
IL-15 M1 TAM 0.2520 -0.1408 0.6598 1.287 0.869 1.934
IL-15 M2 TAM 0.8056 0.4941 1.1329 2.238 1.639 3.105
IL-15 Cytotoxic T-cells 0.7031 0.3812 0.9796 2.020 1.464 2.663
IL-15 Proliferating T-cells 0.4988 0.1890 0.8187 1.647 1.208 2.268
Both DC -0.1652 -0.5837 0.2300 0.848 0.558 1.259
Both M1 TAM -0.3983 -0.8143 -0.0006 0.671 0.443 0.999
Both M2 TAM 0.0776 -0.2656 0.4657 1.081 0.767 1.593
Both Cytotoxic T-cells -0.0956 -0.4545 0.2832 0.909 0.635 1.327
Both Proliferating T-cells -0.2500 -0.6524 0.1317 0.779 0.521 1.141
TCR DC -0.4747 -0.8844 -0.0850 0.622 0.413 0.918
TCR M1 TAM -0.6940 -1.1384 -0.3223 0.500 0.320 0.724
TCR M2 TAM 0.0199 -0.3667 0.3828 1.020 0.693 1.466
TCR Cytotoxic T-cells -0.3157 -0.6544 0.0225 0.729 0.520 1.023
TCR Proliferating T-cells -0.4818 -0.8946 -0.0430 0.618 0.409 0.958
========================================================================
PREDICTED PROBABILITIES p_border, p_core (posterior mean ± 94 % HDI)
========================================================================
Condition Cell type p_border_mean p_border_lo p_border_hi p_core_mean p_core_lo p_core_hi
IL-15 DC 0.6100 0.5037 0.7230 0.4488 0.3354 0.5652
IL-15 M1 TAM 0.4457 0.3197 0.5714 0.3856 0.2697 0.5052
IL-15 M2 TAM 0.4837 0.3690 0.5957 0.2975 0.2016 0.3957
IL-15 Cytotoxic T-cells 0.5093 0.3983 0.6228 0.3417 0.2401 0.4477
IL-15 Proliferating T-cells 0.5239 0.4133 0.6366 0.4023 0.2982 0.5163
Both DC 0.5829 0.4247 0.7242 0.6214 0.4825 0.7760
Both M1 TAM 0.3637 0.2270 0.5120 0.4572 0.3022 0.6029
Both M2 TAM 0.3652 0.2321 0.5077 0.3479 0.2226 0.4965
Both Cytotoxic T-cells 0.3975 0.2587 0.5445 0.4201 0.2737 0.5691
Both Proliferating T-cells 0.4398 0.2872 0.5931 0.5003 0.3504 0.6574
TCR DC 0.5203 0.3992 0.6311 0.6335 0.5064 0.7411
TCR M1 TAM 0.2728 0.1764 0.3685 0.4263 0.2975 0.5451
TCR M2 TAM 0.3544 0.2506 0.4572 0.3502 0.2460 0.4677
TCR Cytotoxic T-cells 0.4024 0.2968 0.5110 0.4790 0.3683 0.6033
TCR Proliferating T-cells 0.4358 0.3214 0.5664 0.5535 0.4215 0.6837
Saved: model6_delta_ck_summary.csv, model6_predicted_probabilities.csv
Hypothesis contrasts: M2 TAM border–core effect across conditions¶
We focus on M2 TAM as the primary cell type of interest and visualize the posterior $\delta_{c,k}$ (border-minus-core log-odds) across all three conditions (IL-15, Both, TCR). Significance annotations indicate the strength of evidence that $\delta_{c,k} \neq 0$:
- ** : Bayesian $p_B < 0.01$
- * : Bayesian $p_B < 0.05$
where $p_B = 2 \cdot \min\!\left(\Pr(\delta_{c,k}>0),\; 1-\Pr(\delta_{c,k}>0)\right)$.
In [19]:
Copied!
# ── Hypothesis contrasts: M2 TAM forest plot with significance stars ─────
# Requires: trace6, cond_levels, cells_levels, delta_flat (from cell 13)
target_cell = "M2 TAM"
HDI_PROB = 0.94
# Find M2 TAM index
ki_m2 = list(cells_levels).index(target_cell)
# Collect posterior draws for M2 TAM across all conditions
fig, ax = plt.subplots(figsize=(6, 2.5), dpi=300)
plt.rcParams.update({"font.size": 8})
cmap = plt.cm.Set2
yticks, ylabels = [], []
for ci, cond in enumerate(cond_levels):
d = delta_flat[:, ci, ki_m2]
m = float(np.mean(d))
hdi = az.hdi(d, hdi_prob=HDI_PROB)
# Bayesian p-value: p_B = 2 * min(Pr(δ>0), 1-Pr(δ>0))
prob_pos = np.mean(d > 0)
p_bayes = 2 * min(prob_pos, 1.0 - prob_pos)
# Significance star
if p_bayes < 0.01:
star = "**"
elif p_bayes < 0.05:
star = "*"
else:
star = ""
color = cmap(ci / max(len(cond_levels) - 1, 1))
ax.plot([hdi[0], hdi[1]], [ci, ci], lw=3, color=color, solid_capstyle="round")
ax.plot(m, ci, "o", ms=6, color=color, zorder=5)
# Annotate significance star to the right of the HDI bar
if star:
ax.text(hdi[1] + 0.03, ci, star, fontsize=14, fontweight="bold",
va="center", ha="left", color="k")
yticks.append(ci)
ylabels.append(f"{cond}")
# Print summary
or_val = np.exp(m)
print(f"{cond:>8s} | {target_cell}: δ = {m:.3f} [{hdi[0]:.3f}, {hdi[1]:.3f}] "
f"OR = {or_val:.2f} Pr(δ>0) = {prob_pos:.3f} p_B = {p_bayes:.4f} {star}")
ax.axvline(0, ls="--", lw=0.8, color="k", zorder=0)
ax.set_yticks(yticks)
ax.set_yticklabels(ylabels, fontsize=9)
ax.set_xlabel("$\\delta_{c,k}$ (border − core log-odds)", fontsize=10)
ax.set_title(f"Posterior $\\delta_{{c,k}}$ for {target_cell} with 94 % HDI",
fontsize=11, fontweight="bold")
ax.invert_yaxis()
# Secondary x-axis for odds ratio
# ax2 = ax.twiny()
# lo, hi = ax.get_xlim()
# ax2.set_xlim(np.exp(lo), np.exp(hi))
# ax2.set_xscale("log")
# ax2.set_xlabel("Odds ratio (border / core)", fontsize=10)
for sp in ("top", "right"):
ax.spines[sp].set_visible(False)
ax.spines["left"].set_linewidth(0.5)
ax.spines["bottom"].set_linewidth(0.5)
plt.tight_layout()
plt.savefig("fig_delta_M2TAM_significance.pdf", bbox_inches="tight")
plt.show()
print(f"\nSaved: fig_delta_M2TAM_significance.pdf")
# ── Hypothesis contrasts: M2 TAM forest plot with significance stars ─────
# Requires: trace6, cond_levels, cells_levels, delta_flat (from cell 13)
target_cell = "M2 TAM"
HDI_PROB = 0.94
# Find M2 TAM index
ki_m2 = list(cells_levels).index(target_cell)
# Collect posterior draws for M2 TAM across all conditions
fig, ax = plt.subplots(figsize=(6, 2.5), dpi=300)
plt.rcParams.update({"font.size": 8})
cmap = plt.cm.Set2
yticks, ylabels = [], []
for ci, cond in enumerate(cond_levels):
d = delta_flat[:, ci, ki_m2]
m = float(np.mean(d))
hdi = az.hdi(d, hdi_prob=HDI_PROB)
# Bayesian p-value: p_B = 2 * min(Pr(δ>0), 1-Pr(δ>0))
prob_pos = np.mean(d > 0)
p_bayes = 2 * min(prob_pos, 1.0 - prob_pos)
# Significance star
if p_bayes < 0.01:
star = "**"
elif p_bayes < 0.05:
star = "*"
else:
star = ""
color = cmap(ci / max(len(cond_levels) - 1, 1))
ax.plot([hdi[0], hdi[1]], [ci, ci], lw=3, color=color, solid_capstyle="round")
ax.plot(m, ci, "o", ms=6, color=color, zorder=5)
# Annotate significance star to the right of the HDI bar
if star:
ax.text(hdi[1] + 0.03, ci, star, fontsize=14, fontweight="bold",
va="center", ha="left", color="k")
yticks.append(ci)
ylabels.append(f"{cond}")
# Print summary
or_val = np.exp(m)
print(f"{cond:>8s} | {target_cell}: δ = {m:.3f} [{hdi[0]:.3f}, {hdi[1]:.3f}] "
f"OR = {or_val:.2f} Pr(δ>0) = {prob_pos:.3f} p_B = {p_bayes:.4f} {star}")
ax.axvline(0, ls="--", lw=0.8, color="k", zorder=0)
ax.set_yticks(yticks)
ax.set_yticklabels(ylabels, fontsize=9)
ax.set_xlabel("$\\delta_{c,k}$ (border − core log-odds)", fontsize=10)
ax.set_title(f"Posterior $\\delta_{{c,k}}$ for {target_cell} with 94 % HDI",
fontsize=11, fontweight="bold")
ax.invert_yaxis()
# Secondary x-axis for odds ratio
# ax2 = ax.twiny()
# lo, hi = ax.get_xlim()
# ax2.set_xlim(np.exp(lo), np.exp(hi))
# ax2.set_xscale("log")
# ax2.set_xlabel("Odds ratio (border / core)", fontsize=10)
for sp in ("top", "right"):
ax.spines[sp].set_visible(False)
ax.spines["left"].set_linewidth(0.5)
ax.spines["bottom"].set_linewidth(0.5)
plt.tight_layout()
plt.savefig("fig_delta_M2TAM_significance.pdf", bbox_inches="tight")
plt.show()
print(f"\nSaved: fig_delta_M2TAM_significance.pdf")
IL-15 | M2 TAM: δ = 0.806 [0.494, 1.133] OR = 2.24 Pr(δ>0) = 1.000 p_B = 0.0000 **
Both | M2 TAM: δ = 0.078 [-0.266, 0.466] OR = 1.08 Pr(δ>0) = 0.661 p_B = 0.6783
TCR | M2 TAM: δ = 0.020 [-0.367, 0.383] OR = 1.02 Pr(δ>0) = 0.534 p_B = 0.9323
Saved: fig_delta_M2TAM_significance.pdf
Model comparison: hierarchical vs flat $\delta_{c,k}$ (LOO)¶
To assess whether the structured hierarchical decomposition of the region effect $\delta_{c,k}$ improves out-of-sample predictive accuracy, we compared two models using Pareto-smoothed importance-sampling leave-one-out cross-validation (PSIS-LOO; Vehtari et al., 2017):
Model 6 (hierarchical) — the primary model used throughout this analysis, in which $\delta_{c,k}$ is decomposed into global, condition, cell-type, and interaction components: $$\delta_{c,k} = \delta_0 + \delta_c + \delta_k + \sigma_\delta \cdot \varepsilon_{c,k}, \quad \varepsilon_{c,k} \sim \mathcal{N}(0,1)$$ with centering constraints on $\delta_c$ and $\delta_k$ (sum-to-zero) and partial pooling through $\sigma_\delta \sim \text{HalfNormal}(0.25)$.
Model flat (non-hierarchical) — an otherwise identical model in which each $\delta_{c,k}$ receives an independent prior without structural decomposition: $$\delta_{c,k} \sim \mathcal{N}(0, 1)$$ This model has the same number of effective $\delta_{c,k}$ parameters (condition × cell type = 15), but does not share information across conditions or cell types.
The hierarchical model encourages borrowing of strength across condition–cell-type combinations through the shared $\sigma_\delta$, $\delta_c$, and $\delta_k$ components. If the data support such structure, the hierarchical model should produce better (or comparable) out-of-sample predictions; if the structure is too restrictive, the flat model will be preferred.
We report $\widehat{\text{elpd}}_{\text{LOO}}$ (expected log pointwise predictive density), $\Delta\widehat{\text{elpd}}$ relative to the best model with its standard error, and Pareto-$k$ diagnostics. When $|\Delta\text{elpd}| \lesssim 2 \times \text{SE}$, the models are considered practically equivalent in predictive performance.
In [ ]:
Copied!
# ── Model comparison: hierarchical vs flat δ_{c,k} (LOO) ─────────────────
# Compare model6 (hierarchical δ decomposition) vs model_flat (independent δ_{c,k})
# Requires: model6, trace6 already in memory from earlier cells.
# ═══════════════════════════════════════════════════════════════════════════
# 1. Define and sample the FLAT (non-hierarchical) model
# ═══════════════════════════════════════════════════════════════════════════
# Identical to model6 except δ_{c,k} ~ Normal(0, 1) directly,
# i.e. NO decomposition into δ₀ + δ_c + δ_k + σ_δ·ε_{c,k}.
print("═" * 72)
print("Defining and sampling model_flat (non-hierarchical δ_{c,k}) …")
print("═" * 72)
with pm.Model(coords=coords) as model_flat:
# indices (reuse same data)
c_f = pm.Data("c", ind_cond, dims="obs_id")
k_f = pm.Data("k", ind_cells, dims="obs_id")
s_f = pm.Data("s", ind_batch, dims="obs_id")
r_f = pm.Data("r", ind_roi, dims="obs_id")
z_f = pm.Data("z", z, dims="obs_id")
offset_f = pm.Data("offset", offset_in, dims="obs_id")
n_f = pm.Data("n", N, dims="obs_id")
Y_f = pm.Data("y", y, dims="obs_id")
# ── Priors identical to model6 (except δ structure) ───────────────────
alpha_f = pm.Normal("alpha", 0, 1.5)
mu_ck_f = pm.Normal("mu_ck", 0, 1.5, dims=("Condition", "Cells"))
# ── FLAT δ_{c,k}: independent Normal, no hierarchical decomposition ──
delta_ck_f = pm.Normal("delta_ck", 0, 1.0, dims=("Condition", "Cells"))
# batch random intercept (non-centered, same as model6)
sigma_batch_f = pm.HalfNormal("sigma_batch", 0.3)
u_batch_raw_f = pm.Normal("u_batch_raw", 0, 1, dims=("S",))
u_batch_f = pm.Deterministic("u_batch", u_batch_raw_f * sigma_batch_f, dims=("S",))
# dispersion (same as model6)
mu_logphi_f = pm.Normal("mu_logphi", mu=np.log(mean_N), sigma=0.5)
sigma_logphi_f = pm.HalfNormal("sigma_logphi", 0.3)
logphi_k_raw_f = pm.Normal("logphi_k_raw", mu=0, sigma=0.5, dims=("Cells",))
logphi_k_f = pm.Deterministic("logphi_k",
mu_logphi_f + sigma_logphi_f * logphi_k_raw_f,
dims=("Cells",))
phi_k_f = pm.Deterministic("phi_k", 1 + pm.math.exp(logphi_k_f), dims=("Cells",))
# linear predictor
beta_off_f = pm.Normal("beta_off", 0.0, 0.5)
eta_f = (beta_off_f * offset_f + alpha_f + u_batch_f[s_f]
+ mu_ck_f[c_f, k_f] + delta_ck_f[c_f, k_f] * z_f)
p_f = pm.Deterministic("p", pm.math.sigmoid(eta_f), dims="obs_id")
phi_obs_f = phi_k_f[k_f]
a_f = p_f * phi_obs_f
b_f = (1 - p_f) * phi_obs_f
count_f = pm.BetaBinomial("count", alpha=a_f, beta=b_f, n=n_f,
observed=Y_f, dims="obs_id")
trace_flat = pm.sample(
draws=1500,
tune=3000,
chains=4,
target_accept=0.999,
init="adapt_diag",
random_seed=RANDOM_SEED,
progressbar=True,
idata_kwargs={"log_likelihood": True},
)
# Quick diagnostics for the flat model
n_divs_flat = int(trace_flat.sample_stats["diverging"].sum())
rhat_flat = az.rhat(trace_flat)
rhat_flat_max = float(np.nanmax(
np.concatenate([v.values.ravel() for v in rhat_flat.data_vars.values()])
))
print(f"\nmodel_flat diagnostics: divergences = {n_divs_flat}, max R-hat = {rhat_flat_max:.4f}")
# ── Model comparison: hierarchical vs flat δ_{c,k} (LOO) ─────────────────
# Compare model6 (hierarchical δ decomposition) vs model_flat (independent δ_{c,k})
# Requires: model6, trace6 already in memory from earlier cells.
# ═══════════════════════════════════════════════════════════════════════════
# 1. Define and sample the FLAT (non-hierarchical) model
# ═══════════════════════════════════════════════════════════════════════════
# Identical to model6 except δ_{c,k} ~ Normal(0, 1) directly,
# i.e. NO decomposition into δ₀ + δ_c + δ_k + σ_δ·ε_{c,k}.
print("═" * 72)
print("Defining and sampling model_flat (non-hierarchical δ_{c,k}) …")
print("═" * 72)
with pm.Model(coords=coords) as model_flat:
# indices (reuse same data)
c_f = pm.Data("c", ind_cond, dims="obs_id")
k_f = pm.Data("k", ind_cells, dims="obs_id")
s_f = pm.Data("s", ind_batch, dims="obs_id")
r_f = pm.Data("r", ind_roi, dims="obs_id")
z_f = pm.Data("z", z, dims="obs_id")
offset_f = pm.Data("offset", offset_in, dims="obs_id")
n_f = pm.Data("n", N, dims="obs_id")
Y_f = pm.Data("y", y, dims="obs_id")
# ── Priors identical to model6 (except δ structure) ───────────────────
alpha_f = pm.Normal("alpha", 0, 1.5)
mu_ck_f = pm.Normal("mu_ck", 0, 1.5, dims=("Condition", "Cells"))
# ── FLAT δ_{c,k}: independent Normal, no hierarchical decomposition ──
delta_ck_f = pm.Normal("delta_ck", 0, 1.0, dims=("Condition", "Cells"))
# batch random intercept (non-centered, same as model6)
sigma_batch_f = pm.HalfNormal("sigma_batch", 0.3)
u_batch_raw_f = pm.Normal("u_batch_raw", 0, 1, dims=("S",))
u_batch_f = pm.Deterministic("u_batch", u_batch_raw_f * sigma_batch_f, dims=("S",))
# dispersion (same as model6)
mu_logphi_f = pm.Normal("mu_logphi", mu=np.log(mean_N), sigma=0.5)
sigma_logphi_f = pm.HalfNormal("sigma_logphi", 0.3)
logphi_k_raw_f = pm.Normal("logphi_k_raw", mu=0, sigma=0.5, dims=("Cells",))
logphi_k_f = pm.Deterministic("logphi_k",
mu_logphi_f + sigma_logphi_f * logphi_k_raw_f,
dims=("Cells",))
phi_k_f = pm.Deterministic("phi_k", 1 + pm.math.exp(logphi_k_f), dims=("Cells",))
# linear predictor
beta_off_f = pm.Normal("beta_off", 0.0, 0.5)
eta_f = (beta_off_f * offset_f + alpha_f + u_batch_f[s_f]
+ mu_ck_f[c_f, k_f] + delta_ck_f[c_f, k_f] * z_f)
p_f = pm.Deterministic("p", pm.math.sigmoid(eta_f), dims="obs_id")
phi_obs_f = phi_k_f[k_f]
a_f = p_f * phi_obs_f
b_f = (1 - p_f) * phi_obs_f
count_f = pm.BetaBinomial("count", alpha=a_f, beta=b_f, n=n_f,
observed=Y_f, dims="obs_id")
trace_flat = pm.sample(
draws=1500,
tune=3000,
chains=4,
target_accept=0.999,
init="adapt_diag",
random_seed=RANDOM_SEED,
progressbar=True,
idata_kwargs={"log_likelihood": True},
)
# Quick diagnostics for the flat model
n_divs_flat = int(trace_flat.sample_stats["diverging"].sum())
rhat_flat = az.rhat(trace_flat)
rhat_flat_max = float(np.nanmax(
np.concatenate([v.values.ravel() for v in rhat_flat.data_vars.values()])
))
print(f"\nmodel_flat diagnostics: divergences = {n_divs_flat}, max R-hat = {rhat_flat_max:.4f}")
In [12]:
Copied!
# ═══════════════════════════════════════════════════════════════════════════
# 2. Compute PSIS-LOO for both models
# ═══════════════════════════════════════════════════════════════════════════
print("\nComputing PSIS-LOO …")
def _ll_var(idata):
"""Return the name of the first log-likelihood variable."""
return list(idata.log_likelihood.data_vars)[0]
loo_hier = az.loo(trace6, var_name=_ll_var(trace6), pointwise=True)
loo_flat = az.loo(trace_flat, var_name=_ll_var(trace_flat), pointwise=True)
print(f"\nmodel6 (hierarchical): elpd_loo = {loo_hier.elpd_loo:.1f} (SE = {loo_hier.se:.1f})")
print(f"model_flat (flat) : elpd_loo = {loo_flat.elpd_loo:.1f} (SE = {loo_flat.se:.1f})")
# ═══════════════════════════════════════════════════════════════════════════
# 3. az.compare — publication-ready table
# ═══════════════════════════════════════════════════════════════════════════
comp = az.compare(
{"model6 (hierarchical δ)": trace6,
"model_flat (flat δ)": trace_flat},
ic="loo",
)
print("\n" + "=" * 80)
print("LOO MODEL COMPARISON — hierarchical vs flat δ_{c,k}")
print("=" * 80)
print(comp.to_string())
# ═══════════════════════════════════════════════════════════════════════════
# 4. Pareto-k diagnostic summary
# ═══════════════════════════════════════════════════════════════════════════
for label, loo_obj in [("hierarchical", loo_hier), ("flat", loo_flat)]:
pk = loo_obj.pareto_k.values if hasattr(loo_obj, "pareto_k") else loo_obj.pareto_k
n_bad = int(np.sum(pk > 0.7))
n_warn = int(np.sum((pk > 0.5) & (pk <= 0.7)))
print(f"\n {label} Pareto k: >0.7 (bad) = {n_bad}, "
f"0.5–0.7 (marginal) = {n_warn}, "
f"max k = {float(np.max(pk)):.3f}")
# ═══════════════════════════════════════════════════════════════════════════
# 5. Publication figure: ELPD comparison bar chart
# ═══════════════════════════════════════════════════════════════════════════
fig, ax = plt.subplots(figsize=(4.5, 3), dpi=300)
plt.rcParams.update({"font.size": 8})
model_names = comp.index.tolist()
elpd_vals = comp["elpd_loo"].values
se_vals = comp["se"].values
# --- FIX: Check for correct column name for diff ---
if "elpd_diff" in comp.columns:
d_elpd = comp["elpd_diff"].values
else:
d_elpd = comp["diff_elpd_loo"].values # older arviz versions
dse_vals = comp["dse"].values
colors = ["#4C72B0", "#DD8452"]
bars = ax.barh(range(len(model_names)), elpd_vals, xerr=se_vals,
height=0.5, color=colors[:len(model_names)], alpha=0.85,
edgecolor="white", capsize=3, error_kw={"lw": 1})
ax.set_yticks(range(len(model_names)))
ax.set_yticklabels(model_names, fontsize=8)
ax.set_xlabel("$\\widehat{\\mathrm{elpd}}_{\\mathrm{LOO}}$", fontsize=9)
ax.set_title("LOO model comparison: hierarchical vs flat $\\delta_{c,k}$",
fontsize=10, fontweight="bold")
ax.invert_yaxis()
# Annotate Δelpd ± SE on bars
for i, (name, de, ds) in enumerate(zip(model_names, d_elpd, dse_vals)):
if de != 0:
ax.text(elpd_vals[i] + se_vals[i] + 0.5, i,
f"Δ = {de:.1f} ± {ds:.1f}", va="center", fontsize=7, color="gray")
for sp in ("top", "right"):
ax.spines[sp].set_visible(False)
plt.tight_layout()
plt.savefig("fig_loo_comparison.pdf", bbox_inches="tight")
plt.show()
# Save comparison table
comp.to_csv("model_comparison_loo.csv")
print("\nSaved: model_comparison_loo.csv, fig_loo_comparison.pdf")
# ═══════════════════════════════════════════════════════════════════════════
# 2. Compute PSIS-LOO for both models
# ═══════════════════════════════════════════════════════════════════════════
print("\nComputing PSIS-LOO …")
def _ll_var(idata):
"""Return the name of the first log-likelihood variable."""
return list(idata.log_likelihood.data_vars)[0]
loo_hier = az.loo(trace6, var_name=_ll_var(trace6), pointwise=True)
loo_flat = az.loo(trace_flat, var_name=_ll_var(trace_flat), pointwise=True)
print(f"\nmodel6 (hierarchical): elpd_loo = {loo_hier.elpd_loo:.1f} (SE = {loo_hier.se:.1f})")
print(f"model_flat (flat) : elpd_loo = {loo_flat.elpd_loo:.1f} (SE = {loo_flat.se:.1f})")
# ═══════════════════════════════════════════════════════════════════════════
# 3. az.compare — publication-ready table
# ═══════════════════════════════════════════════════════════════════════════
comp = az.compare(
{"model6 (hierarchical δ)": trace6,
"model_flat (flat δ)": trace_flat},
ic="loo",
)
print("\n" + "=" * 80)
print("LOO MODEL COMPARISON — hierarchical vs flat δ_{c,k}")
print("=" * 80)
print(comp.to_string())
# ═══════════════════════════════════════════════════════════════════════════
# 4. Pareto-k diagnostic summary
# ═══════════════════════════════════════════════════════════════════════════
for label, loo_obj in [("hierarchical", loo_hier), ("flat", loo_flat)]:
pk = loo_obj.pareto_k.values if hasattr(loo_obj, "pareto_k") else loo_obj.pareto_k
n_bad = int(np.sum(pk > 0.7))
n_warn = int(np.sum((pk > 0.5) & (pk <= 0.7)))
print(f"\n {label} Pareto k: >0.7 (bad) = {n_bad}, "
f"0.5–0.7 (marginal) = {n_warn}, "
f"max k = {float(np.max(pk)):.3f}")
# ═══════════════════════════════════════════════════════════════════════════
# 5. Publication figure: ELPD comparison bar chart
# ═══════════════════════════════════════════════════════════════════════════
fig, ax = plt.subplots(figsize=(4.5, 3), dpi=300)
plt.rcParams.update({"font.size": 8})
model_names = comp.index.tolist()
elpd_vals = comp["elpd_loo"].values
se_vals = comp["se"].values
# --- FIX: Check for correct column name for diff ---
if "elpd_diff" in comp.columns:
d_elpd = comp["elpd_diff"].values
else:
d_elpd = comp["diff_elpd_loo"].values # older arviz versions
dse_vals = comp["dse"].values
colors = ["#4C72B0", "#DD8452"]
bars = ax.barh(range(len(model_names)), elpd_vals, xerr=se_vals,
height=0.5, color=colors[:len(model_names)], alpha=0.85,
edgecolor="white", capsize=3, error_kw={"lw": 1})
ax.set_yticks(range(len(model_names)))
ax.set_yticklabels(model_names, fontsize=8)
ax.set_xlabel("$\\widehat{\\mathrm{elpd}}_{\\mathrm{LOO}}$", fontsize=9)
ax.set_title("LOO model comparison: hierarchical vs flat $\\delta_{c,k}$",
fontsize=10, fontweight="bold")
ax.invert_yaxis()
# Annotate Δelpd ± SE on bars
for i, (name, de, ds) in enumerate(zip(model_names, d_elpd, dse_vals)):
if de != 0:
ax.text(elpd_vals[i] + se_vals[i] + 0.5, i,
f"Δ = {de:.1f} ± {ds:.1f}", va="center", fontsize=7, color="gray")
for sp in ("top", "right"):
ax.spines[sp].set_visible(False)
plt.tight_layout()
plt.savefig("fig_loo_comparison.pdf", bbox_inches="tight")
plt.show()
# Save comparison table
comp.to_csv("model_comparison_loo.csv")
print("\nSaved: model_comparison_loo.csv, fig_loo_comparison.pdf")
Computing PSIS-LOO …
/data/kanferg/conda/envs/pymc_latest/lib/python3.12/site-packages/arviz/stats/stats.py:795: UserWarning: Estimated shape parameter of Pareto distribution is greater than 0.70 for one or more samples. You should consider using a more robust model, this is because importance sampling is less likely to work well if the marginal posterior and LOO posterior are very different. This is more likely to happen with a non-robust model and highly influential observations. warnings.warn( /data/kanferg/conda/envs/pymc_latest/lib/python3.12/site-packages/arviz/stats/stats.py:795: UserWarning: Estimated shape parameter of Pareto distribution is greater than 0.70 for one or more samples. You should consider using a more robust model, this is because importance sampling is less likely to work well if the marginal posterior and LOO posterior are very different. This is more likely to happen with a non-robust model and highly influential observations. warnings.warn(
model6 (hierarchical): elpd_loo = -576.2 (SE = 17.6) model_flat (flat) : elpd_loo = -580.5 (SE = 17.3)
/data/kanferg/conda/envs/pymc_latest/lib/python3.12/site-packages/arviz/stats/stats.py:795: UserWarning: Estimated shape parameter of Pareto distribution is greater than 0.70 for one or more samples. You should consider using a more robust model, this is because importance sampling is less likely to work well if the marginal posterior and LOO posterior are very different. This is more likely to happen with a non-robust model and highly influential observations. warnings.warn( /data/kanferg/conda/envs/pymc_latest/lib/python3.12/site-packages/arviz/stats/stats.py:795: UserWarning: Estimated shape parameter of Pareto distribution is greater than 0.70 for one or more samples. You should consider using a more robust model, this is because importance sampling is less likely to work well if the marginal posterior and LOO posterior are very different. This is more likely to happen with a non-robust model and highly influential observations. warnings.warn(
================================================================================
LOO MODEL COMPARISON — hierarchical vs flat δ_{c,k}
================================================================================
rank elpd_loo p_loo elpd_diff weight se dse warning scale
model6 (hierarchical δ) 0 -576.154889 42.668621 0.000000 0.941327 17.550501 0.000000 True log
model_flat (flat δ) 1 -580.526204 48.097985 4.371315 0.058673 17.317732 3.087436 True log
hierarchical Pareto k: >0.7 (bad) = 7, 0.5–0.7 (marginal) = 19, max k = 0.975
flat Pareto k: >0.7 (bad) = 6, 0.5–0.7 (marginal) = 37, max k = 0.954
Saved: model_comparison_loo.csv, fig_loo_comparison.pdf
Implementation notes (PyMC)¶
The model was implemented in PyMC with indexed data containers for (c_i,k_i,s_i,z_i,\text{offset}i,n_i,y_i). Batch effects used a non-centered parameterization, and (\delta{c,k}) used a partially pooled hierarchical decomposition (global + condition + cell type + interaction scaled by (\sigma_\delta)). The cell-type precision (\phi_k) was modeled via a log-normal hierarchy with a positive transform (\phi_k = 1+\exp(\log\phi_k)) to ensure (\phi_k>1).