CRRT-UDE

Context-Residual Regulatory Transport via Universal Differential Equations

A biologically structured model that learns epithelial progression dynamics and quantifies how local tissue ecology alters progression trajectories beyond intrinsic cellular state.

Overview

CRRT-UDE addresses a core question in premalignant progression: which local tissue interactions alter the progression-aligned destination of an epithelial cell beyond what was already predictable from its intrinsic state, and through which regulatory programs?

The method learns:

An intrinsic regulatory progression field from epithelial state alone
A context-residual field capturing ecological contributions from the surrounding niche
The regulatory programs mediating these niche effects

This enables quantifying the per-receiver, spatially mappable context-residual regulatory transport that represents progression velocity explained by local ecology after accounting for intrinsic state.

Method

Design

CRRT-UDE operates on:

Regulatory state z in R^d (~773 dimensions on the LUAD cohort: 14 PROGENy signaling pathways + ~759 CollecTRI transcription factors), scored per spot via decoupleR from epithelial expression. The structured regulatory subvector r(z) that drives the interpretable terms B r(z) and the context gate C[r(z) ⊙ g(w)] is the 14 named PROGENy pathways; the full TF+pathway state is transported and read out per program (all ~773 dims appear in R_cond). The model is thus context-residual transport of coordinated TF + pathway activity; it is not (yet) transport over an explicit TF-network topology (no TF–TF adjacency/Laplacian enters the cost or field — that is a documented future extension).
Spatial context w: interpretable features from matched spatial transcriptomics (Visium) including cell composition (DestVI deconvolution), spatial organization (BANKSY neighborhood summaries), ECM markers, inflammation signatures, and spatial gradients
Pseudo-progression time tau in [0,1]: normalized pathological stage coordinate (not clock time)

The model decomposes epithelial progression velocity into:

dz/dtau = v_self(z, tau) + v_context(z, w, tau) + residuals

where:
  v_self = A*z + b + B*r(z) + rho_s              [intrinsic field]
  v_context = C * [r(z) ⊙ g(w)] + rho_c          [niche-gated field]
  g(w) = tanh(G*w + c)                            [context gate]

The structured vector field encodes biological mechanisms: A*z + b is linear drift; B*r(z) captures regulatory-program modulation of intrinsic progression; C * [r(z) ⊙ g(w)] is the context-gated regulatory interaction; rho_s, rho_c are neural residuals (small if the regulatory structure captures the dynamics).

Conditional-Stratified Coupling

A critical design decision emerged from prespecified identifiability benchmarking (Task 11A, see docs/manuscript/EXPERIMENT_LOG.md Experiments 5-6):

Initial approach: strict context-neutral optimal transport on epithelial state z alone, motivated by anti-circularity (don't embed the effect you're trying to measure into the pairing).
Empirical finding: strict-z coupling recovers a synthetic context effect in abstract latent space (Tier-1 cosine 0.82) but FAILS through a regulatory representation like pathway/TF scores (Tier-2 cosine 0.27).
Revision: CONDITIONAL-STRATIFIED coupling that constructs source-target transport WITHIN ecologically comparable populations, defined by a COARSE, prespecified niche representation c = a(w) (broad myeloid/stromal/lymphoid/endothelial fractions). This passes both Tier-1 (0.93) and Tier-2 regulatory recovery (0.84).

The coupling uses coarse niche classes to establish comparability; fine-grained mechanistic context is EXCLUDED from the coupling and reserved for the UDE interpretation. This follows the conditional-transport formulation of conditional Monge maps (Bunne et al., NeurIPS 2022).

Training

Three-stage optimal-transport conditional flow matching (OT-CFM):

Stage A: train v_self to source-target displacements from the conditional coupling
Stage B: FREEZE v_self; train v_context to the residual velocity v_self failed to explain
Stage C: brief anchored joint refinement with a penalty keeping intrinsic dynamics near the Stage A solution

Integration via fixed-step RK4; endpoint calibration via an entropic Sinkhorn optimal-transport cost (squared-euclidean, (pi * cost).sum()) between predicted and real target-stage populations. Note: this is the plain entropic OT cost, not the Sinkhorn divergence with self-cost debiasing terms; the endpoint coupling is subsampled (default 2048/side, deterministic) so the cost is memory-bounded at cohort scale.

Primary Estimand

R_cond(z, w; c) = context-residual regulatory transport CONTRAST conditional on coarse niche stratum:

R_cond(z, w; c) = Φ_full(z, w | π_c*) − Φ_self(z | π_c*)

where Φ is the integrated UDE flow to tau=1 and π_c* is the conditional-stratified coupling. This is a per-receiver, model-based, spatially mappable quantity representing predicted regulatory displacement attributable to context AFTER restricting to ecologically comparable populations.

Honest framing: R_cond is conditional transport (identifiable given the coarse niche class), NOT an individual causal effect or historical cell-fate trajectory. The coupling identifies which cells are comparable; it does not recover exact pairings.

Key Results

Methods Development (Synthetic Benchmarks)

The prespecified identifiability benchmark (Task 11A) characterized the recovery boundary for context effects from unpaired populations:

Null controls pass: methods do not invent effects when ground-truth context coefficient = 0, even under state-context correlation or stage confounding
Recoverable regime: at realistic effect sizes (context/state-spread ratio ~0.5-0.7), conditional-stratified coupling recovers the true per-cell regulatory effect (correlation 0.63-0.91 across regimes, Tier-2 cosine 0.84)
Non-identifiable regime: when context effect dominates state spread (ratio >0.9), only oracle pairings recover the effect (guard holds)

This empirical characterization of identifiability — including the NEGATIVE finding that strict context-neutral coupling fails Tier-2 — is a methods contribution. The benchmark motivated the conditional-stratified revision and establishes the supported estimand before any real-data claim.

Real-Data Analysis

The patient-held-out LUAD run (v2lowrank_source model; AAH→AIS and AIS→invasive edges) and the PanIN cross-disease replication have completed. Full quantitative results — per-program R_cond, the falsification ladder, D_rho credibility, and the pathway-level findings — are reported in the thesis/manuscript, not restated here (the repository documents the method and how to reproduce it; the results paper carries the numbers).

The headline is a bounded positive result, reported honestly. Cell-intrinsic dynamics account for the majority of the modelled progression, and the model recovers a small, reproducible, program-specific context residual (a sign-consistent set of programs on each LUAD edge, stable across all five patient-held-out folds). But the falsification ladder bounds what that residual means at the endpoint: cell-resolved local context is only suggestively better than self-only transport (favored in 4/5 folds, paired difference not significant, P≈0.10) and is statistically equivalent to a stratum-matched shuffle of the same niche features (paired 95% CI tightly bounded around zero). A variance decomposition explains why — the ecological variation is overwhelmingly receiver-local and stage-invariant at Visium resolution — so the reproducible signal is specimen-scale rather than uniquely receiver-local in these cohorts. The degenerate single-branch fields collapse as required, confirming the grey-box decomposition, and the PanIN cross-disease replication returns a correctly calibrated null at four donors. The contribution is a prespecified estimand with a characterized recovery boundary, not an unqualified positive claim; full numbers are in the thesis/manuscript.

Installation

# Minimal dependencies (data manipulation, evaluation)
pip install -e .

# Full dependencies for feature extraction and UDE training (HPC)
pip install -e ".[hpc]"

# Development dependencies
pip install -e ".[dev]"

Note: Some spatial features require external packages:

BANKSY: pip install -e /path/to/Banksy (or vendor the neighbor-mean + gradient functions)

Reproduction

Snakemake Workflow

The repository uses Snakemake for reproducible pipeline execution.

Edge selection is explicit. LUAD targets require an edge-selector configfile so a command can never silently run the wrong thesis edge; a bare LUAD target stops with a message telling you to pick one. Prepend the selector to every LUAD command:

# AAH -> AIS (preinvasive; the PRIMARY thesis edge)
SEL="--configfile workflow/config.yaml workflow/config_aah_ais.yaml"
# AIS -> {MIA, LUAD} (pooled invasion edge)   -> use workflow/config_ais_invasive.yaml instead

Key targets (LUAD; add $SEL):

# Full pipeline (snRNA preprocessing + spatial context + UDE + evaluation)
snakemake --cores 8 $SEL all

# Spatial preprocessing + context feature extraction only
snakemake --cores 4 $SEL all_spatial

# Generate all figures
snakemake --cores 4 $SEL all_figures

# Dry run to see execution plan
snakemake --cores 1 $SEL --dry-run all

PanIN targets read their own panin: block and need NO selector:

snakemake --cores 8 all_panin

For HPC (SLURM) with GPU:

# Via SLURM profile (see workflow/slurm/ for config); $SEL selects the LUAD edge
snakemake --profile workflow/slurm --jobs 20 $SEL all

# Individual GPU-dependent rules (LUAD rules still need the edge selector)
snakemake --profile workflow/slurm $SEL ude_train

Key rules (see workflow/Snakefile for full DAG):

ingest_snrna, spatial_ingest: raw count ingestion from AnnData
normalize, typing: normalization (scanpy/scclr) + cell-type annotation
self_features: regulatory programs (CollecTRI/PROGENy via decoupleR) from snRNA
spatial_context, spatial_regulatory: BANKSY + DestVI deconvolution + spatial regulatory scoring
assemble_context: construct per-spot context table w
ude_decode: decode epithelial receiver expression via per-sample DestVI decoders (fanned out one GPU job per sample; see below) + gather (gene-intersection + concat + metadata)
ude_gates: Phase-0 study-reference adequacy gate (real-data DestVI-decoder check; the synthetic identifiability benchmark that selected the coupling lives in crrt/ude/identifiability_benchmark.py / tests/)
ude_train: three-stage CRRT-UDE training (conditional-stratified coupling; fanned out one GPU job per CV fold)
ude_estimand: compute the context-residual estimand per receiver (code/output named R_e; R_cond is the design-doc name for the same conditional quantity) + per-program table (per-fold jobs + cross-fold stability gather)
ude_ladder: falsification ladder (2x2 coupling x context + shuffle controls; per-fold jobs + concat gather)
viability: patient-held-out self-only stage-prediction baseline + ROC

Opt-in downstream stages (off every default target; each fully isolated, graceful-gated, and non-blocking — nothing in all_ude/all_panin/all_figures depends on them):

all_comm / all_comm_panin: R_cond-gated cell–cell communication (crrt/comm/). Fuses LIANA ligand–receptor inference with Receptor→TF→program chains, gated by R_cond so only model-implicated context-driven channels survive; adds a comparative (source-vs-target) differential communication network, Ollivier–Ricci bottleneck detection, an optional pySCENIC/CollecTRI promoter (TF→target) layer, and a publication figure set (circos, chains, sankey). Retrain-free — reads existing all_ude outputs (typed h5ad + R_cond) + optional all_grn regulons.
all_grn / all_grn_full: pySCENIC + hdWGCNA GRN backends (own conda envs; docs/GRN_SETUP.md).
all_scenic_downstream: SCENIC AUCell + RSS panels + regulon-activity clustermap — the pySCENIC-protocol downstream figures (scenic_aucell → scenic_downstream_figures).

Parallel fan-out (SLURM)

The compute-heavy UDE steps are decomposed into independent per-unit SLURM jobs so the pipeline saturates the cluster instead of running one node at a time (the StageBridge wildcard pattern):

ude_decode fans out to one GPU job per compatible sample (ude_decode_sample, ~22 jobs); samples are discovered at parse time by audit_decoders (a filesystem scan, no scvi/GPU). A CPU gather then intersects genes, concatenates, and attaches metadata (preserving the sample:barcode spot index).
ude_train, ude_estimand, ude_ladder each fan out to one job per CV fold (*_fold rules), followed by a gather. ude_train_config writes the shared config.json once so parallel fold jobs never race on it.

Folds and samples are independent, so fan-out does not change results — this is enforced by tests/test_ude_ladder.py::test_ladder_fanout_per_fold_equals_all_folds (per-fold concatenation equals the all-folds run to rtol=1e-6). Each script keeps its original monolithic path (no --fold/--sample/--gather flags) for single-job local runs.

GPU/CPU resources are declared inline in each rule (slurm_partition, gpu), not only in the profile, so the device request is visible in the rule and is not silently dropped by a command-line --set-resources/--exclude override. crrt/ude/receiver.py loads DestVI with accelerator="gpu" when CUDA is present (scvi otherwise defaults the decode forward pass to CPU).

Configuration

Local vs HPC paths live in workflow/config.yaml (paths.local / paths.hpc blocks), selected via --config env=local|hpc. The transport edge is NOT set in the base config; it is chosen by an edge-selector configfile (workflow/config_aah_ais.yaml for the primary AAH -> AIS edge, or workflow/config_ais_invasive.yaml for AIS -> {MIA, LUAD}), which a bare LUAD target requires. The corrected model version is set by model_tag (default v2lowrank_source: low-rank identifiable A, source-fixed context). Local smoke fixtures: configs/smoke.yaml, configs/smoke_spatial.yaml.

# local smoke test (AAH -> AIS edge)
snakemake -n all_ude --configfile workflow/config.yaml workflow/config_aah_ais.yaml --config env=local
# HPC (SLURM), full run (AAH -> AIS edge; swap the selector for the invasion edge)
snakemake all_ude --profile workflow/slurm \
  --configfile workflow/config.yaml workflow/config_aah_ais.yaml --config env=hpc --jobs 20
# PanIN cross-disease replication needs no edge selector:
snakemake all_panin --profile workflow/slurm --config env=hpc --jobs 20

Optional downstream stages

These run after all_ude (and, for the promoter layer, all_grn); each is fully isolated and does not affect the core pipeline. GRN backends need their pinned conda envs — rebuild from the frozen lock (see docs/GRN_SETUP.md), never from loose version ranges:

# GRN inference (pySCENIC + hdWGCNA), all edges in one command, each its own SLURM job
snakemake all_grn_full --profile workflow/slurm --configfile workflow/config.yaml \
  --config env=hpc scenic_env=/scratch/chaunzt1/envs/pyscenic hdwgcna_env=/scratch/chaunzt1/envs/hdwcgna --jobs 20

# SCENIC downstream: AUCell activity + RSS panels + regulon-activity clustermap
snakemake all_scenic_downstream --profile workflow/slurm --configfile workflow/config.yaml \
  --config env=hpc scenic_env=/scratch/chaunzt1/envs/pyscenic --jobs 4

# R_cond-gated cell-cell communication (per LUAD edge; needs $SEL)
snakemake all_comm --profile workflow/slurm $SEL --config env=hpc --jobs 4
snakemake all_comm_panin --profile workflow/slurm --config env=hpc --jobs 4   # PanIN (no selector)

Documentation

docs/DATA.md — data availability (GEO accessions, reference atlases, what is not redistributed)
docs/FIGURES.md — every figure mapped to its plotting function and Snakemake rule
docs/GRN_SETUP.md — pySCENIC + hdWGCNA environments (rebuild from the frozen lock)
docs/specs/ — locked design document · docs/manuscript/EXPERIMENT_LOG.md — identifiability findings

Repository Layout

crrt/                       Core package
├── data/                   ingest, access, audit, spatial_ingest, patient splits
├── features/               normalize, cell_typing, regulatory, spatial_context,
│                           destvi_compose, assemble_context, embedding, preprocess
├── ude/                    Model + identifiability core
│   ├── field.py                 Structured v_self + tanh-gated v_context (CRRTField)
│   ├── coupling.py              Context-neutral + conditional-stratified OT couplings
│   ├── _ot.py / _semot.py       Vendored log-domain Sinkhorn coupling + entropic OT-cost primitives
│   ├── train.py                 Three-stage OT-CFM trainer (conditional coupling)
│   ├── integrate.py             Fixed-step RK4 / Euler integrators
│   ├── estimand.py              R_e context-residual + endpoint metrics + D_N diagnostics
│   ├── ladder.py                Falsification ladder (2x2 + shuffle controls)
│   ├── gates.py                 Phase-0 study-reference adequacy gate
│   ├── receiver.py              DestVI gamma -> decoded expression -> regulatory state
│   ├── dataset.py               Join receivers + context, patient folds
│   ├── synthetic_identifiability.py   Task 11A benchmark DGP (regimes A-E)
│   └── identifiability_benchmark.py   Coupling comparison + model selection
├── comm/                   R_cond-gated cell–cell communication (opt-in `all_comm`)
│   ├── infer_lr.py              LIANA ligand–receptor inference (per stage)
│   ├── intracellular.py         Receptor→TF→program chains + R_cond gate
│   ├── network.py               differential CCI network + topology + Ollivier–Ricci curvature
│   ├── promoter.py              TF→target layer (pySCENIC regulons > cisTarget > CollecTRI)
│   └── figures.py               circos / chain networks / sankey / heatmap
├── evaluation/             baseline (self-only viability), viability_plot
└── viz/                    figure modules: ude_result / ude_dynamics / cnv_convergence /
                            regulatory / composition / niche / spatial / spatialstats /
                            dynamics_grid / trajectory_animation / scenic_figures

workflow/                   Snakemake pipeline
├── Snakefile               Main DAG (rules inline; model + spatial + comm + GRN + scenic-downstream)
├── config.yaml             local + hpc path blocks
├── config_aah_ais.yaml     edge selector: AAH → AIS (primary edge)
├── config_ais_invasive.yaml  edge selector: AIS → {MIA, LUAD} (pooled invasion)
├── envs/                   conda env specs
│   ├── crrt.yaml, hdwgcna.yaml, scclr.yaml
│   └── locks/              FROZEN pySCENIC lock (conda-explicit + pip-freeze) + restore_scenic.sh
├── slurm/                  SLURM profile (account chaunzt1)
└── scripts/                Per-rule driver scripts (ude_*, comm_analysis, scenic_*, etc.)

docs/
├── DATA.md                 Data availability (GEO accessions, references, what is not redistributed)
├── FIGURES.md              Figure → plotting-function → Snakemake-rule catalog
├── GRN_SETUP.md            pySCENIC + hdWGCNA env build (rebuild from the frozen lock)
├── specs/                  Locked design document
├── manuscript/             EXPERIMENT_LOG.md (all identifiability results) + deliverables plan
├── figures/                Model schematic (TikZ + matplotlib fallback)
└── plans/                  Task breakdowns, gate criteria

tests/                      Pytest suite (unit + synthetic benchmarks + paired discriminators)
configs/                    smoke.yaml, smoke_spatial.yaml (local smoke fixtures)

Caveats and Limitations

Stated Upfront

tau is pseudo-progression, not clock time: the model's time coordinate is normalized pathological stage, not days/months. It captures the direction and magnitude of state change aligned with progression, not growth rates or real-time dynamics.
Conditional transport, not individual causality: R_cond is a model-based conditional transport contrast. The conditional coupling identifies which source-target pairs are ecologically comparable; it does NOT recover exact historical cell lineages or prove that a specific niche caused a specific cell's endpoint. The estimand is justified by the conditional Monge-map formulation (Bunne et al. 2022) but is not an individual causal effect.
Inferred receiver states: spatial receiver states are DestVI-decoded epithelial expression from Visium spots (mixed-cell assay), projected into regulatory space. They are NOT dissociated single nuclei. Adequacy gates (decoder audit, anchor reconstruction, cross-patient consistency) validate that decoded states recapitulate study-defined premalignant biology before use.
Intrinsic field fitted on paired data: v_self is trained on source-target pairs from the conditional coupling, which contain ecological structure. The context-residual R_cond is therefore a conservative estimate under the fitted decomposition, not a formal lower bound on the true ecological effect (that would require unproven independence assumptions).
Contingent on HPC adequacy gate: real-data conclusions are contingent on the DestVI decoder adequacy audit (anchor reconstruction Spearman >= 0.90, decoder-identity R^2 < 0.10, program directionality agreement). If existing decoders fail, the pipeline retrains one shared CondSCVI on matched study snRNA before proceeding.

Citation

See CITATION.cff for structured citation metadata (author, title, repository URL, license).

License: MIT (see LICENSE or pyproject.toml).

Author: AJ Book (SecondBook5)

References

This work builds on and cites (see thesis.bib in the companion thesis repository for full entries):

Optimal transport and the Sinkhorn iteration: Cuturi (NeurIPS 2013)
Conditional Monge maps: Bunne et al. (NeurIPS 2022)
Schrödinger bridge (Gaussian closed form): Bunne et al. (2022)
Flow matching: Lipman et al. (2022)
Regulatory scoring: PROGENy (Schubert et al. 2018), CollecTRI (Müller-Dott et al. 2023), decoupleR (Badia-i-Mompel et al. 2022)
Spatial deconvolution and context: DestVI (Lopez et al. 2022), BANKSY (Singhal et al. 2024)
Regulatory-network inference: SCENIC (Aibar et al. 2017), i-cisTarget (Imrichová et al. 2015), hdWGCNA (Morabito et al. 2023), fgsea (Korotkevich et al. 2016)
Cell--cell communication: LIANA (Dimitrov et al. 2024), COMMOT (Cang et al. 2023), OmniPath (Türei et al. 2026)
Population-level optimal transport: moscot (Klein et al. 2025)

Positioning: CRRT-UDE extends spatially-conditioned progression models by introducing an explicit context-residual decomposition with receiver-conditioned regulatory mediation and a prespecified identifiability benchmark. The key delta is the structured regulatory UDE + the R_cond estimand with characterized recovery boundary.

Status

Design locked (2026-07-17). Model built, identifiability benchmark complete (conditional-stratified coupling selected). Patient-held-out LUAD analysis (AAH→AIS, AIS→invasive) and PanIN cross-disease replication complete; optional GRN, SCENIC-downstream, and R_cond-gated communication stages implemented. Quantitative results are reported in the thesis/manuscript.

See docs/specs/2026-07-17-crrt-ude-design.md for the complete locked specification and docs/manuscript/EXPERIMENT_LOG.md for all identifiability findings (Experiments 1-6, including negative results that motivated the conditional-stratified revision).

Name		Name	Last commit message	Last commit date
Latest commit History 247 Commits
.github/workflows		.github/workflows
configs		configs
crrt		crrt
docs		docs
experiments		experiments
outputs/identifiability		outputs/identifiability
scripts		scripts
tests		tests
workflow		workflow
.gitignore		.gitignore
.pre-commit-config.yaml		.pre-commit-config.yaml
CHANGELOG.md		CHANGELOG.md
CITATION.cff		CITATION.cff
CONTRIBUTING.md		CONTRIBUTING.md
LICENSE		LICENSE
README.md		README.md
pyproject.toml		pyproject.toml

Navigation Menu

Repository files navigation

CRRT-UDE

Overview

Method

Design

Conditional-Stratified Coupling

Training

Primary Estimand

Key Results

Methods Development (Synthetic Benchmarks)

Real-Data Analysis

Installation

Reproduction

Snakemake Workflow

Parallel fan-out (SLURM)

Configuration

Optional downstream stages

Documentation

Repository Layout

Caveats and Limitations

Stated Upfront

Citation

References

Status

About

Releases

Packages

Contributors

Languages

Folders and files

Latest commit

History

Repository files navigation

CRRT-UDE

Overview

Method

Design

Conditional-Stratified Coupling

Training

Primary Estimand

Key Results

Methods Development (Synthetic Benchmarks)

Real-Data Analysis

Installation

Reproduction

Snakemake Workflow

Parallel fan-out (SLURM)

Configuration

Optional downstream stages

Documentation

Repository Layout

Caveats and Limitations

Stated Upfront

Citation

References

Status

About

Resources

Contributing

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages