interpretability refactor: attention/ → interpretability/ (publish diffae core, quarantine _internal)

pyproject.toml

@@ -84,6 +84,9 @@ source = "vcs"
 [tool.hatch.build.targets.wheel]
 only-include = ["src"]
 sources = ["src"]
+# not-for-release: internal analysis/paper tooling (atlas, shap, titration, embedding,
+# weighted_aggregation, viewer, kyle_pcs, gen_validation) lives under interpretability/toolkit
+exclude = ["src/ops_model/models/interpretability/toolkit/**"]
 
 [tool.hatch.metadata]
 allow-direct-references = true

src/ops_model/data/data_loader.py

@@ -39,8 +39,8 @@
 warnings.filterwarnings("ignore", category=zarr.errors.ZarrUserWarning)
 
 # Default name of the per-construct identifier column in adata.obs.
-# Override per experiment via OpsDataManager(guide_col=...) (e.g.
-# "minibinder_perturbation" for minibinder experiments).
+# Override per experiment via OpsDataManager(guide_col=...) (e.g. a custom
+# perturbation column for non-CRISPR libraries).
 DEFAULT_GUIDE_COL = "sgRNA"
 
 
@@ -592,18 +592,18 @@ def get_labels(self):
                 print(f"Reading link CSV from {csv_path}")
                 labels_tmp = pd.read_csv(csv_path)
 
-                # Minibinder back-compat: link CSVs from minibinder experiments
-                # don't have a "gene_name" column. Copy minibinder_perturbation
-                # into gene_name so downstream gene_name-aware code (e.g. the
-                # balanced-sampling and gene-label LUT helpers) keeps working.
-                # Follow-up: minibinder gene-level should ultimately use
-                # gene_target, not the construct id — tracked separately.
+                # Custom-perturbation back-compat: some link CSVs use a
+                # non-standard guide column (self.guide_col) and have no
+                # "gene_name" column. Copy the guide column into gene_name so
+                # downstream gene_name-aware code (balanced sampling, gene-label
+                # LUT helpers) keeps working.
                 if (
                     "gene_name" not in labels_tmp.columns
                     and "Gene name" not in labels_tmp.columns
-                    and "minibinder_perturbation" in labels_tmp.columns
+                    and self.guide_col != "gene_name"
+                    and self.guide_col in labels_tmp.columns
                 ):
-                    labels_tmp["gene_name"] = labels_tmp["minibinder_perturbation"]
+                    labels_tmp["gene_name"] = labels_tmp[self.guide_col]
 
                 if self.guide_col not in labels_tmp.columns:
                     raise ValueError(

src/ops_model/data/labels.py

@@ -10,8 +10,9 @@
 
 import numpy as np
 import pandas as pd
+from ops_model.paths import BASE_PATH
 
-_DEFAULT_BASE_PATH = "/hpc/projects/intracellular_dashboard/fast_ops"
+_DEFAULT_BASE_PATH = f"{BASE_PATH}"
 
 # Backward-compatible filename templates for legacy csv_source values
 SOURCE_FILENAME_TEMPLATES = {
@@ -71,7 +72,7 @@ def load_immunostaining_labels(
         filename_template: Filename pattern with {well} placeholder,
             e.g. "cell_painting_linked_{well}.csv" or "four_i_linked_{well}.csv"
         base_path: Base directory containing per-experiment subdirectories.
-            Defaults to /hpc/projects/intracellular_dashboard/fast_ops.
+            Defaults to /hpc/projects/icd.fast.ops.
 
     Returns:
         labels_df ready to pass to OpsDataManager.construct_dataloaders()

src/ops_model/data/paths.py

@@ -14,7 +14,7 @@ def _resolve_base() -> Path:
         return Path(
             os.environ.get(
                 "OPS_OUTPUT_BASE_DIR",
-                "/hpc/projects/icd.fast.ops",
+                f"{BASE_PATH}",
             )
         )
 
@@ -85,7 +85,7 @@ def __init__(self, experiment: str, well: str = None):
         }
 
         self.other = {
-            "gene_library": "/hpc/projects/intracellular_dashboard/ops/configs/annotated_guide_library_123-UpdateJuly28_2025.csv",
+            "gene_library": f"{BASE_PATH}/configs/annotated_guide_library_123-UpdateJuly28_2025.csv",
         }
 
     def reformat_well_name(self, well: str) -> str:

src/ops_model/features/anndata_utils.py

@@ -26,11 +26,12 @@
 import matplotlib.pyplot as plt
 
 from ops_utils.data.feature_metadata import FeatureMetadata
+from ops_model.paths import BASE_PATH
 
 
 DEFAULT_SEARCH_DIRS = [
-    Path("/hpc/projects/icd.fast.ops"),
-    Path("/hpc/projects/icd.ops"),
+    Path(f"{BASE_PATH}"),
+    Path(f"{BASE_PATH}"),
 ]
 
 DEFAULT_GUIDE_COL = "sgRNA"
@@ -1441,7 +1442,7 @@ def load_multiple_experiments(
         List of paths to .h5ad files
 
     Example:
-        >>> base_dir = "/hpc/projects/intracellular_dashboard/ops"
+        >>> base_dir = f"{BASE_PATH}"
         >>> experiments = ["ops0089_20251119", "ops0084_20250101"]
         >>> paths = load_multiple_experiments(base_dir, experiments)
         >>> adata_combined = concatenate_anndata_objects(paths)

src/ops_model/features/batch_process_embeddings.py

@@ -29,10 +29,11 @@
 
 
 from ops_model.features.processing_common import process_features_csv
+from ops_model.paths import BASE_PATH
 
 
 # Base directory for OPS experiments
-BASE_DIR = Path("/hpc/projects/intracellular_dashboard/ops")
+BASE_DIR = Path(f"{BASE_PATH}")
 
 
 def check_csv_exists(

src/ops_model/models/cellprofiler/cp_extraction.py

@@ -243,8 +243,8 @@ def create_subset(
         bounds: [start, end] index range
         out_channels: List of channel names (default: ["Phase2D", "mCherry"])
         guide_col: Name of the per-construct identifier column in the link CSV
-            (default: "sgRNA"; e.g. "minibinder_perturbation" for minibinder
-            experiments)
+            (default: "sgRNA"; e.g. a custom perturbation column for
+            non-CRISPR libraries)
 
     Returns:
         Tuple of (Subset dataset, label lookup table)

src/ops_model/models/interpretability/diffae/CROPSEQ_TO_MORPHOLOGY.md

@@ -0,0 +1,79 @@
+# Idea: transcriptome-controlled morphology generation (CROP-seq → DiffAE)
+
+**Goal:** use paired CROP-seq on the same geneKO library to let the DiffAE generate *how a cell's phenotype
+changes as its transcriptome moves toward a KO state* — i.e. drive the traversal by a **transcriptional
+direction** (CROP-seq) instead of (or in addition to) the CellDINO morphological direction.
+
+## Reframing vs today
+Current DiffEx traversal = real cell → DDIM-inverted `xT` (identity/nuisance) + **CellDINO gene-direction**
+(morphology conditioning), morph α NTC→KO, decode image. This idea keeps the entire image decoder (DDIM +
+inversion + guidance `w`) and swaps the *driver* to a transcriptional signature.
+
+## Hard constraint (shapes everything): no single-cell pairing
+CROP-seq is destructive scRNA-seq; OPS is imaging — **no cell is in both modalities**. So supervision is only
+**per-perturbation** (gene KO → mean transcriptional shift Δt_g AND a morphological distribution), never
+cell-level `(transcriptome → image)`.
+- Model learns **transcriptional-signature → morphological-distribution**; within-gene image variation comes
+  from the stochastic `xT` (same as today).
+- Conditioning vectors are **per-gene pseudobulk** (or per-guide if guide calls are clean).
+
+## Two paths
+
+### Path B — reuse the trained DiffAE via a transcriptome→CellDINO map (POC first)
+Fit a perturbation-level regressor `Δt_g → ΔCellDINO_g` (linear → small MLP) over the shared KOs. A
+transcriptional vector → predicted CellDINO shift → **existing morpho DiffAE renders it**.
+- Pros: reuses the whole trained pipeline + viewer; days not weeks. **The map's R² is itself a headline
+  result** ("fraction of KO morphology predictable from KO transcriptome").
+- Cons: bottlenecked through CellDINO.
+
+### Path A — condition the DiffAE directly on transcriptome (full, CPA-flavored)
+Project t through `cond_proj` into the FiLM/cross-attn slot the CellDINO emb uses now; train on
+`(image_i, t_{gene(i)})`. Cleanest is a **CPA-style shared perturbation embedding** `e_g`: a transcriptome
+decoder reconstructs CROP-seq (`NTC + e_g`), the DiffAE decodes the image (`anchor xT + e_g`), `e_g` shared →
+ties the modalities through one latent. Any transcriptional state → `e_g` → image.
+- Pros: end-to-end; supports unseen/combined signatures + continuous "dial a pathway, watch morphology".
+- Cons: real training effort; guard against collapse to gene-means (xT + guidance mitigate, as today).
+
+**Plan:** B as a weekend POC (also tests whether transcriptome predicts morphology at all) → A if promising.
+
+## Transcriptional vector options (cheapest first)
+1. pseudobulk logFC vs NTC (per gene); 2. learned scRNA latent (scVI/PCA) mean per gene; 3. pathway/program
+module scores (most interpretable "dials"). Start with (1)/(2) for the direction, expose (3) as the control.
+
+## Concrete CROP-seq source — Duo's sVAE+ gene-program embeddings (June 2025)
+Duo Peng built the CROP-seq embeddings we should use as option (2)/(3). This is a **sparse VAE (sVAE+)** — Lopez
+et al. 2023, *Learning Causal Representations of Single Cells via Sparse Mechanism Shift Modeling* — so the latent
+axes are interpretable **gene programs**, which is exactly the "dial a pathway" control we wanted.
+- Confluence: [sVAE approach to gene programs v3](https://czbiohub.atlassian.net/wiki/spaces/dashboard/pages/5199986706/sVAE+approach+to+gene+programs+v3)
+  — the **purple "sVAE embeddings" section** has the embeddings file.
+- **Run the encoder to embed new expression profiles** — point setup at the parent results folder:
+  `/hpc/projects/data.science/duo.peng/sVAEplus/sVAEplus/6000HVG/svaeplus_results_2_256_1_200_0.5/`
+  - trained encoder: `best_model/model.pt`; params `best_params.json` (n_layers=2, n_hidden=256,
+    sparse_mask_penalty=1.0, kl_warmup=200, dropout=0.05); code root `.../sVAEplus/sVAEplus/sVAE-main` + `ops_utils`, `install.sh`.
+  - expression values (normalized, sVAE+-compat, filtered, 6000 HVG):
+    `.../svaeplus_results_2_256_1_200_0.5/CropSeq_June2025_filtered_normalized_compat_forsvaeplus_filtered.h5ad`
+- ⚠️ **`gene_loadings.csv` (gene → gene-program activity) is a post-hoc *linear* summary — do NOT use it as the
+  mapping.** The real expression → program mapping is **non-linear**; get it by running the encoder on expression
+  values, not by the loadings matrix.
+- Fit for the plan: these program embeddings are the transcriptional vector for **Path B** (`Δprogram_g →
+  ΔCellDINO_g`) and the shared-latent seed / conditioning signal for **Path A**. Per-gene means over the encoder
+  output give Δt_g; NTC cells in the same h5ad define the control baseline.
+
+## The novel payoff: transcriptome↔morphology divergence map
+Plot every gene by (transcriptional effect size, morphological effect size). The DiffAE then lets you *see*:
+- **transcriptionally loud, morphologically silent** → counterfactual "what it would look like if it manifested"
+- **morphologically loud, transcriptionally quiet** → morphology carrying signal transcriptome misses
+- cross-modal interpolation between two genes' transcriptomes; agreement w/ CellDINO morph = validation,
+  divergence = the interesting biology.
+
+## Viewer tab concept
+"Transcriptome → Morphology" tab: α-slider drives the **transcriptional** traversal; side-by-side vs the
+existing CellDINO-driven morph (agreement = validation, divergence = biology). Reuses traversal/montage render.
+
+## To scope
+1. ~~CROP-seq path/format~~ → **resolved**: Duo's sVAE+ h5ad + trained encoder (see section above). Still TBD:
+   gene-overlap of the CROP-seq library with the imaging 1000-lib; whether to embed with the encoder or use
+   Duo's precomputed embeddings from the purple Confluence section.
+2. per-gene vs per-guide signatures.
+3. matched NTC/control in CROP-seq to define Δ.
+4. payoff emphasis: generator vs divergence-map.

src/ops_model/models/interpretability/diffae/README.md

@@ -0,0 +1,40 @@
+# DiffEx — counterfactual interpretability for the attention atlas
+
+Explain geneKO / protein-complex phenotypes **in image space**: generate
+counterfactual single-cell morphs ("if this control cell were a KD, what would it
+look like?") and the per-pixel change map, instead of relying on OP/CP features.
+Adapted from DiffEx (arXiv:2502.09663) and Alex Lin's EvolutionaryScale pipeline,
+working in the **CellDINO embedding space** that the SetTransformer already uses.
+
+See [PLAN.md](PLAN.md) for the design rationale and the full running log.
+
+## Pipeline (three stages, each a subpackage)
+
+| stage | package | what it does |
+|---|---|---|
+| 1 | [`classifier/`](classifier/) | per-class single-cell classifier on **top-attention cells** — the model whose decision DiffEx explains / that ranks directions. B = ResNet on phase crops; **C = MLP on CellDINO features** (chosen). |
+| 2 | [`generator/`](generator/) | **conditional diffusion** generator (the DiffAE): UNet that generates a cell image conditioned on its CellDINO embedding (conditioning dropout + EMA + CFG). |
+| 3 | [`directions/`](directions/) | **contrastive direction discovery** (InfoNCE + decorrelation, unsupervised) → rank directions by a control-vs-target classifier → **CFG traversal** α∈[−,+] → DDIM-sample a counterfactual strip + Δ-pixel heatmap, verified by re-encoded score. |
+
+## Run order (each stage has `run.py` for local + `submit.py` for SLURM)
+
+```bash
+# Stage 1 — classifier (per gene/complex, or sweep --all-classes)
+python -m ops_model.models.interpretability.diffae.classifier.submit --grain complex --all-classes --models C
+python -m ops_model.models.interpretability.diffae.classifier.aggregate --grain complex --model C
+
+# Stage 2 — train the conditional DiffAE (resume-able; gate = embedding/noise ratio)
+python -m ops_model.models.interpretability.diffae.generator.submit --epochs 120 --batch-size 48
+python -m ops_model.models.interpretability.diffae.generator.diagnose_conditioning   # conditioning-strength check
+
+# Stage 3 — directions + counterfactual traversal for a target
+python -m ops_model.models.interpretability.diffae.directions.submit --grain geneKO --target HSPA5
+```
+
+Outputs: `/hpc/projects/icd.fast.ops/models/diffex/{<grain>,diffae,directions}/…`.
+
+## Status
+Stages 1 & 3 built and validated end-to-end; Stage-2 DiffAE conditioning was the
+hard part — see PLAN.md (the v1 generator ignored the embedding; the rebuild with
+conditioning dropout + EMA fixes it). Current focus: training the DiffAE to a
+conditioning ratio high enough for visible morphs, then scaling across targets.

src/ops_model/models/interpretability/diffae/classifier/README.md

@@ -0,0 +1,61 @@
+# DiffEx single-cell classifier PoC (options B & C)
+
+The classifier DiffEx will explain (see [../PLAN.md](../PLAN.md), classifier §1).
+PoC = **binary HSPA5-vs-rest** on **phase** single-cell crops.
+
+- **B** — ResNet18 on the 160×160 phase crops (pixels → class logits).
+- **C** — MLP on **CellDINO** embeddings of the *same* crops (`ops_model.models.cell_dino`).
+
+Both score **top-attention cells** (settled), differing only in the feature space.
+
+## Locked design (Gav, 2026-06-16)
+- Positives: top-1000 attention cells of HSPA5 (`pma_phase_cells_v2_all.parquet`).
+- Negatives: 1000 cells sampled from the **top-5** attention cells of **other genes**
+  (the "distinct" contrast — strong-vs-strong).
+- Crop 160×160, phase-only (`Phase2D`), no cell mask (full crop context).
+- Split: **3-way train/val/test, grouped by experiment** (confound guard — val & test
+  cells come from experiments never trained on). val = model selection; **test = the
+  clean reported number** (scored once, never used for selection). Stratified-random
+  fallback if a class is missing from a side.
+- Success: held-out **test AUROC ≫ 0.5** (generalizing across experiments ⇒ biology, not batch).
+
+## Decision: C reuses the local CellDINO encoder on the same crops
+Rather than join Alex's per-gene dumps, option C runs the local encoder
+(`CellDinoModel`: channel-adaptive DINO ViT-L/16, resize 224 + per-image z-score,
+`in_channels=1`) on the identical crops B uses, and caches the embeddings. One crop
+pipeline; B and C see identical cells.
+
+## Run (GPU)
+```bash
+# single run, interactively on a GPU node
+python -m ops_model.models.attention.diffex.classifier.run --model B --gene HSPA5
+python -m ops_model.models.attention.diffex.classifier.run --model C --gene HSPA5
+
+# or submit both to SLURM (one GPU job each)
+python -m ops_model.models.attention.diffex.classifier.submit --gene HSPA5
+
+# sweep: all 98 EBI complexes + NTC control, model C
+python -m ops_model.models.attention.diffex.classifier.submit --grain complex --all-classes --models C
+# then rank them
+python -m ops_model.models.attention.diffex.classifier.aggregate --grain complex --model C
+```
+`--grain {geneKO,complex}` selects the parquet + class column (`gene` vs `predicted_class`).
+NTC is included as a negative-control bin (its AUROC should be near chance).
+Outputs land under `<out-dir>/<gene>/` (default out-dir
+`/hpc/projects/icd.fast.ops/models/diffex`): `model_{B,C}.pt`, `metrics_{B,C}.json`,
+and a shared `cache/` (crops + CellDINO features). SLURM logs →
+`ops_mono/slurm_logs/diffex_clf/`.
+
+## Layout
+- `config.py` — all params (the locked defaults above).
+- `data.py` — cell-table query, crop materialization (`BaseDataset`), split.
+- `models.py` — ResNet (B) + MLP head (C).
+- `celldino_features.py` — embed crops with the local CellDINO encoder (cached).
+- `train.py` — shared train/eval loop (AUROC).
+- `run.py` — orchestrator + `run_poc()` entry point.
+- `submit.py` — SLURM submission (`submit_parallel_jobs`).
+
+## Status
+Pipeline verified end-to-end on CPU for **B** (tiny config): cell table → crops
+(non-degenerate, masked) → train → AUROC → artifacts. **C** needs a GPU (CellDINO).
+Next: run B & C on HSPA5 at full scale (GPU), compare AUROC, pick the DiffEx target.

src/ops_model/models/interpretability/diffae/classifier/__init__.py

@@ -0,0 +1,4 @@
+"""DiffEx single-cell classifier PoC (options B and C).
+
+The classifier DiffEx will explain. See ../PLAN.md (classifier §1) and README.md.
+"""