mouse gastrulation 4-mod sweep: TSV + Briscoe DV markers + handoff

Adds 3-row Phase-1B-style slurm sweep for Argelaguet mouse gastrulation
multimodal training (RNA total + spliced + unspliced + ATAC):

- docs/notebooks/model_comparisons/mouse_gastrulation_slurm_jobs.tsv
  3 rows: _balencode, _balencode_zloc50x, _balencode_burst.
  Mirrors balanced_encoder_decoder_slurm_jobs.tsv schema (58 cols);
  queue=gpu-normal, mem=100GB, time=2-00:00:00. Template points to a
  new sibling of embryo_rna_atac_spliced_unspliced_slurm.ipynb in
  cell2state_embryo (committed separately).

- docs/notebooks/mouse_gastrulation/known_marker_genes_briscoe_dv_mouse.csv
  147 rows / 27 categories (14 radial-glia, 13 neuron). Briscoe-lab
  dorsoventral spinal-cord progenitor markers from
  cell2state_embryo/notebooks/benchmark/section_suspension_embryo_10x_kit/
  affinity_analysis_from_annotation_v2.ipynb Cell 151 (dv_progenitor_genes
  dict), title-cased per-hyphen-segment for mouse symbols (SOX2 -> Sox2,
  NKX2-2 -> Nkx2-2). Schema matches docs/notebooks/known_marker_genes.csv
  (gene, cell_type, lineage, category).

- .claude/plans/mouse_gastrulation_4mod_handoff.md
  Handoff for the ATAC re-aggregation step (blocking for full 11-sample
  multimodal). The current anndata_atac.h5ad (mtime 2021-04-21) only
  covers 6 of 11 samples; fragments for all 11 exist on disk but were
  never tiled. Next step: rebuild via cell2state SnapATAC2 loader and
  intersect peaks with the existing 203k peak set by genomic
  coordinate overlap (not exact name match).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: Vitalii Kleshchevnikov

.claude/plans/mouse_gastrulation_4mod_handoff.md

@@ -0,0 +1,176 @@
+# Handoff: 4-modality RegularizedMultimodalVI on Argelaguet mouse gastrulation
+
+**Status**: blocked on ATAC re-aggregation. Notebook + marker CSV + TSV all built, but cannot submit yet — the existing ATAC h5ad covers only 6/11 samples (the older 2021 build pre-dates additions to the dataset). User has decided to rebuild ATAC from fragments for all 11 samples, then intersect with the *existing* peak set by **genomic coordinate overlap** (not exact name match).
+
+---
+
+## Goal
+
+Train **4-modality `RegularizedMultimodalVI`** (RNA total + spliced + unspliced + ATAC) on the **Argelaguet mouse gastrulation multiome** (Pijuan-Sala / Marioni 2022). Submit a Phase-1B-style slurm sweep with **3 conditions**: `_balencode`, `_balencode_zloc50x`, `_balencode_burst`.
+
+Downstream use: per-sample UMAPs across the gastrulation stages, with cell-type overlays from the Briscoe-lab dorsoventral spinal-cord progenitor marker list (title-cased to mouse symbols).
+
+---
+
+## Existing plan & approved approach
+
+- **Plan file**: [/nemo/lab/briscoej/home/users/kleshcv/.claude/plans/now-that-we-found-precious-kettle.md](/nemo/lab/briscoej/home/users/kleshcv/.claude/plans/now-that-we-found-precious-kettle.md) — approved by user 2026-05-13, then revised after the ATAC re-aggregation decision (this handoff supersedes the appendix A1 entry, which now becomes the critical path).
+- Template lineage: `embryo_rna_atac_spliced_unspliced_slurm.ipynb` (4-mod slurm/papermill, autogenerated 2026-04-19 from `embryo_rna_atac_spliced_unspliced.ipynb`, already Crick-pathed).
+- Settings reference: Phase 1B sweep [balanced_encoder_decoder_slurm_jobs.tsv](docs/notebooks/model_comparisons/balanced_encoder_decoder_slurm_jobs.tsv), 58-column schema, described in [phase-1-new-init-golden-wind.md](/nemo/lab/briscoej/home/users/kleshcv/.claude/plans/phase-1-new-init-golden-wind.md).
+
+---
+
+## Current progress (built and verified)
+
+### Files created on 2026-05-13
+
+| File | Purpose | Status |
+|---|---|---|
+| [cell2state_embryo/notebooks/benchmark/regularizedvi/mouse_gastrulation_rna_atac_spliced_unspliced_slurm.ipynb](/nemo/lab/briscoej/home/users/kleshcv/cell2state_embryo/notebooks/benchmark/regularizedvi/mouse_gastrulation_rna_atac_spliced_unspliced_slurm.ipynb) | 34-cell 4-mod training notebook | syntax-clean, all 17 imports resolve, **but data-loading cell still points at the OLD `anndata_atac.h5ad` (6 samples) — needs swap to the rebuilt h5ad + coordinate-overlap var intersection** |
+| [regularizedvi/docs/notebooks/mouse_gastrulation/known_marker_genes_briscoe_dv_mouse.csv](docs/notebooks/mouse_gastrulation/known_marker_genes_briscoe_dv_mouse.csv) | 147 rows, 27 categories (14 radial-glia, 13 neuron). Mouse-symbol-cased Briscoe DV progenitor markers (Sox2, Nkx2-2, …) | ready, no changes needed |
+| [regularizedvi/docs/notebooks/model_comparisons/mouse_gastrulation_slurm_jobs.tsv](docs/notebooks/model_comparisons/mouse_gastrulation_slurm_jobs.tsv) | 3 sweep rows; 58 cols matching Phase 1B schema; queue=gpu-normal, mem=100GB, time=2 days | ready (will need re-validation after the ATAC rebuild because mem may need bumping) |
+
+### Notebook edits applied (via `/tmp/edit_mouse_gastrulation_nb.py`)
+- **Edit 1**: Crick-paths header + title → mouse gastrulation
+- **Edit 2**: papermill `parameters` cell — full param block with mouse gastrulation paths; `wandb_project=regularizedvi_mouse_gastrulation`; `stratify_validation_key="celltype+sample"`; `markers_csv` + `c2m_results_folder` added; CRE parquet params removed
+- **Edit 3**: data-loading cell replaced — reads 3 h5ads, intersects on `obs_names`, splits velocyto layers into spliced/unspliced AnnDatas, builds MuData. **Currently broken on the var intersection point** — see "Blocking issue" below.
+- **Edit 4**: SEACell join (per-sample CSVs → `obs['SEACell']`, `obs['celltype_SEACell']`) — plotting only
+- **Edit 5**: `setup_mudata` call — `ambient_covariate_keys=["sample"]`, `dispersion_key="sample"`, `library_size_key="sample"`; `nn_conditioning_covariate_keys=["all"]` and `feature_scaling_covariate_keys=["all"]` use the constant `obs['all']="all"` column for single-category soft-disable (signature-checked: this signature accepts `list[str]`)
+- **Edit 6**: marker overlay cell — reads `known_marker_genes_briscoe_dv_mouse.csv` and plots per-domain
+- **Edit 7**: per-sample UMAP loop (mirrors cell2module per-sample style)
+- **Edit 8–13**: sweep of remaining embryo-specific obs column references (`Embryo`, `10x_kit`, `sample_id`, `cell_type_lvl*`, `Experiment`, `Section` → `sample`, `stage`, `celltype`, `celltype_SEACell`, etc.) in Cells 8 (data summary), 21 (UMAP color_vars), 22 (standard UMAP), 23 (small annotation set), 29 (LISI), 31 (UMAP grid)
+
+### Verification passes
+- Marker CSV: 147 rows, 27 categories, 61 unique genes — correct
+- `setup_mudata` signature accepts `list[str]` for nn_conditioning/feature_scaling (`["all"]` is valid)
+- All 17 imports resolve in the regularizedvi env
+- Notebook syntax checks clean (`/syntax-check`)
+- TSV column count consistent (58 cols × 3 rows + header)
+
+---
+
+## Blocking issue — the ATAC h5ad coverage gap
+
+**Root cause discovered during data-loading dry run**: `/nemo/lab/briscoej/home/users/kleshcv/large_data/gastrulation_multiome_anndata/anndata_atac.h5ad` has mtime 2021-04-21 and contains only **6 samples** (E7.5/E8.0/E8.5 × rep1/rep2). The 5 missing samples (`E7.75_rep1`, `E8.75_rep1`, `E8.75_rep2`, `E8.5_CRISPR_T_KO`, `E8.5_CRISPR_T_WT`) are **NOT QC-failed** — they have ATAC fragment files (1.4–3.3 GB each, mtime 2022-08-18) but were never aggregated into the h5ad because the h5ad pre-dates them by ~16 months.
+
+Result: a naive multimodal intersection currently yields **32,420 cells × 6 samples × 3 stages (E7.5, E8.0, E8.5)** — missing the later E8.75 timepoint the user cares about.
+
+**The fix**: rebuild ATAC for all 11 samples from fragments via the cell2state SnapATAC2 tile loader, then intersect var names with the existing 203k peaks by **genomic coordinate overlap** (peak `chr/start/end` ∩ tile `chr/start/end`) — not exact name match, because tile names are like `chr1:1000-2000` while peak names are like `chr1-1234-2345` (Argelaguet's custom format).
+
+---
+
+## What worked
+
+- **Editing large notebooks via a Python helper script** rather than NotebookEdit. The notebook (67 KB JSON / ~26k tokens) is over the Read tool's 25k token limit, which blocks NotebookEdit (which requires a prior Read). Workaround: write a small `nbformat`-style helper at `/tmp/edit_mouse_gastrulation_nb.py` that loads JSON, edits cells by source-string match, and writes back. Uses `bash scripts/helper_scripts/run_python_cmd.sh` to invoke with the correct env. Successful pattern for all 13 edits applied.
+- **Copying the 4-mod slurm template as base** (rather than the 2-mod `embryo_rna_atac.ipynb` or the model_comparisons bone marrow tutorial). The slurm template is already Crick-pathed and is the exact notebook the Phase 1B TSV invokes — minimal structural drift.
+- **Constant "all" obs column for soft-disabled covariate keys** — passes signature validation, avoids per-dataset code-path branches, keeps the API call structurally identical to the embryo template.
+- **Building the TSV via a script that copies reference Phase 1B rows + applies overrides** — guarantees column-count consistency.
+- **Title-casing gene symbols hyphen-segment-wise**: `lambda g: "-".join(seg[:1].upper() + seg[1:].lower() for seg in g.split("-"))` handles both `SOX2 → Sox2` and `NKX2-2 → Nkx2-2`.
+- **All Phase 1B settings already wired** in the embryo template — `use_modality_balanced_encoder_init`, `z_loc_init_scale`, `decoder_init_variance_target_fraction`, `decoder_burst_size_init_variance_target_fraction`, papermill type-conversion plumbing — no model-side code changes needed.
+
+## What didn't work / pitfalls to avoid
+
+- **Assuming the 5 missing ATAC samples failed QC**. They did NOT — they were simply never aggregated into the 2021 h5ad. Always verify against `original_with_atac/{sample}/atac_fragments.tsv.gz` existence and mtimes before claiming "QC fail".
+- **Read tool on full notebook**: fails on notebooks larger than ~25k tokens (~67 KB JSON). Use `inspect-notebook` skill for cell-level reads and the JSON-editing helper script pattern for writes.
+- **Wholesale notebook re-creation**: not needed. The 4-mod embryo template's model-side cells (training, attribution, LISI, parameter inspection) all reuse cleanly — only data-loading and plotting columns need swap.
+- **Exact-name var intersection between ATAC h5ads**: WILL FAIL. The existing `anndata_atac.h5ad` peaks use Argelaguet's custom naming (`chr-start-end` hyphenated) while SnapATAC2 tiles use `chr:start-end`. Always intersect by coordinates (`chr` + `start` + `end` tuples from `var` columns), never by `var_names`.
+- **`embryo_rna_atac.ipynb` (2-mod) as base**: has Sanger `/nfs/team205` paths and no slurm metadata cell — more rewriting than starting from the 4-mod slurm sibling.
+- **Bone marrow tutorial (`bone_marrow_multimodal_tutorial_early_stopping.ipynb`) as base**: lacks the papermill `parameters` block needed for TSV-driven slurm submission.
+
+## What's open / unresolved
+
+- **Validation for the burst row in mouse gastrulation**: bursting uses `dispersion_init=variance_burst_size`, which requires `dispersion_init_bio_frac` to be set (per Phase 1B plan). The current TSV row inherits `-` for `dispersion_init_bio_frac` from the embryo 4-mod row. Verify this is OK for a 2-modality MuData-like (RNA+ATAC+spliced+unspliced); embryo had `bio_frac='{"rna":0.99,"spliced":0.99,"unspliced":0.99,"atac":0.9}'` set in the notebook params cell, not the TSV.
+- **GPU memory**: TSV requests 100 GB, but this is a guess. Embryo 4-mod (424k cells) hit 99% of 210 GB. Mouse gastrulation ~50–80k paired cells (after the ATAC rebuild) should be smaller — but 4 modalities × wider feature space (~28k RNA + 32k spliced + 32k unspliced + ~200k ATAC peaks) may still push it. Monitor first run; bump to 150 GB if needed.
+- **SEACell coverage**: 6/11 SEACell.csv files were verified to exist; the other 5 samples (E7.75, E8.75 × 2, CRISPR × 2) — check whether matching SEACell.csv files exist before relying on them in the join. They should, since cell2module ran on all 11 RNA samples, but verify.
+
+---
+
+## Next steps for fresh agent (in plan mode)
+
+### Step 1 — Rebuild ATAC h5ad for all 11 samples
+
+**Template**: [docs/notebooks/immune_integration/bm_pbmc_atac_loading.ipynb](docs/notebooks/immune_integration/bm_pbmc_atac_loading.ipynb) (env: `cell2state_v2026_cuda124_torch25`).
+
+**New notebook** (suggested path): `docs/notebooks/mouse_gastrulation/mouse_gastrulation_atac_loading.ipynb`.
+
+Steps inside the new notebook:
+1. Read the **RNA h5ad** (`/nemo/.../latest/data/processed/rna/anndata.h5ad`, 65,420 cells × 11 samples) — this drives which cells / barcodes will be counted.
+2. Add `fragment_file_path` column mapping `obs['sample']` → `/nemo/.../latest/data/original_with_atac/{sample}/atac_fragments.tsv.gz`. Verify all 11 paths exist (sizes 1.4–3.3 GB, mtime 2022-08-18; see this handoff's "Blocking issue" section for the table).
+3. Call `cell2state.utils.aggregation_v2.concatenate_h5ad(adata, variable_type="atac_tiles", batch_key="sample", loading_kwargs={path_to_reference=".../refdata-cellranger-arc-mm10-2020-A-2.0.0", path_to_fragment_file_key="fragment_file_path", max_frag_size_split=120, bin_size=1000, counting_strategy="paired-insertion", use_complete_path=True})` — exactly mirrors `bm_pbmc_atac_loading.ipynb` Cell 10. Output: ~65k cells × ~2.7M genome-wide 1kb tiles.
+4. Compute ATAC QC (`total_fragments`, `n_tiles_accessible`) per cell. Apply the QC threshold the user specifies (or use the embryo / immune defaults — confirm with user).
+5. Save to `/nemo/.../large_data/gastrulation_multiome_anndata/anndata_atac_tiles_1000bp_split120_11samples.h5ad`.
+
+**Submission**: this is a one-shot slurm job, not a TSV sweep. May need `cell2state_v2026_cuda124_torch25` env (not `regularizedvi`) — check `bm_pbmc_atac_loading.ipynb` Cell 0 for the env declaration. Memory budget: ~150–200 GB depending on chunking strategy. Genome reference `refdata-cellranger-arc-mm10-2020-A-2.0.0` must be on Crick — verify `/nemo/.../refdata-cellranger-arc-mm10-2020-A-2.0.0/genes/genes.gtf.gz` and `star/chrNameLength.txt` exist.
+
+### Step 2 — Coordinate-based var intersection helper
+
+User's exact phrasing: *"intersect them with vars from the already created object - exact names won't match, but we are looking for coordinate overlaps."*
+
+Need: a function that takes (a) the new tile h5ad (~2.7M genome-wide 1kb tiles, `var.index = "chr:start-end"`) and (b) the existing peak h5ad (~203k Argelaguet peaks, `var.index = "chr-start-end"`, columns `chr/start/end`), and returns a tile-subset whose genomic interval overlaps any peak.
+
+**Implementation approaches** (pick one in the plan):
+- **pyranges / bedtools-style** join via `pyranges` (already in regularizedvi env? check) or `pybedtools`. Convert both to BED-style intervals, do interval-join, take the set of tile IDs that hit ≥1 peak.
+- **Native pandas** with sorted `(chr, start)` and `np.searchsorted` for overlap detection — slower but no dependency.
+- **bedtools intersect** via subprocess if pyranges/pybedtools not available.
+
+Decide based on what's in the regularizedvi conda env — check `pyranges` import first.
+
+**Output**: a filtered ATAC tile h5ad (~200k–500k tiles, the union of all tiles overlapping any of the 203k peaks). This becomes the ATAC modality fed to RegularizedMultimodalVI.
+
+### Step 3 — Update the training notebook's data-loading cell
+
+Edit Cell 6 of [mouse_gastrulation_rna_atac_spliced_unspliced_slurm.ipynb](/nemo/lab/briscoej/home/users/kleshcv/cell2state_embryo/notebooks/benchmark/regularizedvi/mouse_gastrulation_rna_atac_spliced_unspliced_slurm.ipynb):
+- Replace the `atac_h5ad` param target with the **new tile h5ad** path (from Step 1).
+- Add the **coordinate-overlap var intersection** call (helper from Step 2) immediately after reading the new ATAC h5ad — load the existing `anndata_atac.h5ad` only to extract its `var[['chr','start','end']]` for the overlap set, then drop the existing h5ad.
+- Update the intersection-on-`obs_names` step: should yield ~55–75k cells × 11 samples × 5 stages (E7.5, E7.75, E8.0, E8.5, E8.75) plus the 2 CRISPR samples — verify this count in a dry run.
+
+Same JSON-editing helper script pattern as before: write `/tmp/edit_mouse_gastr_nb_v2.py` with one `stage` that only touches Cell 6 of the data loader.
+
+### Step 4 — Re-verify and submit
+
+- `/check-signature` on any new helper functions used.
+- `/syntax-check` on the updated notebook.
+- `/check-imports` for any new imports (`pyranges` etc.).
+- Data-loading dry run: confirm cell count (expect 55–75k), confirm 11 samples / 5 stages.
+- `/verify-implementation` — mandatory pre-submission gate per CLAUDE.md.
+- Submit via `bash scripts/helper_scripts/submit_papermill_slurm.sh docs/notebooks/model_comparisons/mouse_gastrulation_slurm_jobs.tsv`.
+
+---
+
+## Key paths quick-reference
+
+```
+# Plan (read this first)
+/nemo/lab/briscoej/home/users/kleshcv/.claude/plans/now-that-we-found-precious-kettle.md
+
+# Built files (ready)
+/nemo/lab/briscoej/home/users/kleshcv/cell2state_embryo/notebooks/benchmark/regularizedvi/mouse_gastrulation_rna_atac_spliced_unspliced_slurm.ipynb
+/nemo/lab/briscoej/home/users/kleshcv/my_packages/regularizedvi/docs/notebooks/mouse_gastrulation/known_marker_genes_briscoe_dv_mouse.csv
+/nemo/lab/briscoej/home/users/kleshcv/my_packages/regularizedvi/docs/notebooks/model_comparisons/mouse_gastrulation_slurm_jobs.tsv
+
+# Templates to copy
+/nemo/lab/briscoej/home/users/kleshcv/cell2state_embryo/notebooks/benchmark/regularizedvi/embryo_rna_atac_spliced_unspliced_slurm.ipynb  # 4-mod training (already copied)
+/nemo/lab/briscoej/home/users/kleshcv/my_packages/regularizedvi/docs/notebooks/immune_integration/bm_pbmc_atac_loading.ipynb  # ATAC tile loader (next step)
+
+# Data inputs
+/nemo/lab/briscoej/home/users/kleshcv/large_data/gastrulation_multiome_anndata/latest/data/processed/rna/anndata.h5ad                      # RNA total (65k × 32k, 11 samples)
+/nemo/lab/briscoej/home/users/kleshcv/large_data/gastrulation_multiome_anndata/latest/data/processed/rna/velocyto/anndata_velocyto.h5ad    # spliced+unspliced (65k × 32k, 11 samples)
+/nemo/lab/briscoej/home/users/kleshcv/large_data/gastrulation_multiome_anndata/anndata_atac.h5ad                                          # OLD ATAC peaks (38k × 203k, 6 samples) — keep ONLY for peak var coordinates
+/nemo/lab/briscoej/home/users/kleshcv/large_data/gastrulation_multiome_anndata/latest/data/original_with_atac/{sample}/atac_fragments.tsv.gz  # ATAC fragments for ALL 11 samples
+
+# SEACell partitions (for plotting only, joined post-training)
+/nemo/lab/briscoej/home/users/kleshcv/my_packages/cell2module/docs/notebooks/results/mouse_gastrulation/per_sample_analysis_{sample}__cell2module/SEACell.csv
+
+# Briscoe DV marker source notebook
+/nemo/lab/briscoej/home/users/kleshcv/cell2state_embryo/notebooks/benchmark/section_suspension_embryo_10x_kit/affinity_analysis_from_annotation_v2.ipynb  # Cell 151 = dv_progenitor_genes dict; Cell 152 = dv_domain_label_v2 derivation
+```
+
+## Project conventions to follow
+
+- All Python via `bash scripts/helper_scripts/run_python_cmd.sh` (NOT bare `python3` or `conda run`).
+- Default float dtype: `float32`. Integer dtypes are nuanced — ask before assuming.
+- Helper scripts go in `scripts/claude_helper_scripts/`. One-off scripts in `/tmp/` are fine but should also follow the `run_python_cmd.sh` convention.
+- `gh` CLI invocation on Crick: `~/.conda/envs/cell2state_v2026_cuda124_torch25/bin/gh <command>`.
+- GPU specs (per CLAUDE.md): bsub on Sanger, sbatch on Crick. Mouse gastrulation 4-mod target queue = `gpu-normal`, partition `ga100` (non-burst) / `gh100` (burst).
+- Always invoke `/verify-implementation` before any irreversible action (submission, push).
+- Notebook editing: use `inspect-notebook` skill to read cells; use the JSON-editing helper script pattern to write — the harness Read tool fails on large notebooks (>25k tokens).