diff --git a/CLAUDE.md b/CLAUDE.md index c9cfe3a5..099eaeaa 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -97,10 +97,46 @@ GitHub Actions runs on push to main and PRs: ruff lint → black format check ## Remote Pipeline Execution -Before launching any remote pipeline: +### CRITICAL: set `n_processes` to at least the number of fits + +**JAX/XLA leaks executable JIT mappings across sequential fits in one process.** +A worker that runs more than ~5 fits dies with `LLVM ERROR: Unable to allocate +section memory` and `JaxRuntimeError: INTERNAL: Failed to materialize symbols` +— on an idle host with 1.4 TB free. It is not a memory shortage: the failing +request was 980 bytes. The XLA dylib counter climbs (`xla_jit_dylib_12` → +`_19` → `_31`) and mappings are never released. + +Failures land **by position in the work queue, not by hyperparameter.** A +serial run of the spike 20-fit grid passed indices 0–4 and failed 5–19, across +both replicates and every `fusionreg` including 0.0. Under a multi-worker pool +this surfaces as a misleading `ModelCollectionFitError: Failed fitting 1 of 20 +parameter sets` — which reads like one bad parameter set but is just the first +worker to exhaust its address space. + +**Rule: `n_processes >= n_fits`,** where `n_fits = len(fusionreg_values) × +n_datasets`. `fit_models()` calls `p.map` with no `chunksize`, so it resolves +to `ceil(n_fits / (n_processes × 4)) == 1` and every fit gets a fresh process. + +| pipeline | grid | `n_processes` | fits/worker | safe? | +|---|---|---|---|---| +| spike prod | 10 fusionreg × 2 reps = 20 | **20** | 1 | yes | +| spike prod (old) | 20 | 6 | 3–4 | **no — leaks** | +| simulation prod | 10 fusionreg × 6 datasets = 60 | 30 | 2 | yes, under the ~5 threshold | + +Never set `n_processes: null` — auto-selection picks workers by core count +alone, ignoring both memory and this leak. Raising `n_processes` is cheap here: +each worker holds one `Data` object (~2–3 GB), so 20 workers is ~60 GB on a +1511 GB host. + +If a fit ever needs more workers than the host has cores, fix the leak instead +of splitting the grid — the real repair is disposing of XLA state between fits +(or `maxtasksperchild=1` on the pool in `model_collection.py`). + +### Launch procedure 1. **Create a local worktree** (if not on main): `git worktree add ../multidms-wt- ` -2. **Scout first**: `bip scout` — pick a server with <20% CPU +2. **Scout first**: `bip scout` — pick a server with <20% CPU. Check available + RAM and cores too, not just the CPU percentage. 3. **Launch** (from the worktree): `pixi run remote-pipeline -- host=` - pipeline: `simulation` or `spike` - profile: `test`, `experimental`, or `prod` @@ -116,13 +152,91 @@ Convention: `output_dir` is always `results--` (e.g., `results- Never skip step 2. Never leave tmux sessions running after fetching results. +### Manual launch (when `remote-pipeline` cannot reach GitHub) + +`remote-pipeline` runs a remote `git fetch origin`, which fails when the +1Password SSH agent refuses to sign for **forwarded** sessions (`ssh-add -l` +lists the key, but every signature dies with `signing failed ... from agent`). +Retrying never fixes it. Push straight into the remote clone over the SSH +channel you already have: + +```bash +git push :/fh/fast/matsen_e/shared/multidms/multidms :refs/heads/ +ssh 'git -C worktree add ' +``` + +If the branch is already checked out in the remote worktree the push is +rejected; push to a temp ref and fast-forward instead: + +```bash +git push : :refs/heads/-incoming +ssh 'git -C merge --ff-only -incoming' +``` + +Then symlink the env into the new worktree (`.pixi/envs` **and** `pixi.lock` — +the lock is gitignored, so a fresh worktree lacks it and pixi re-solves every +platform from scratch) and drive snakemake from a here-doc'd script on the +remote, so `cd` sits on its own line where it cannot be dropped: + +```bash +ssh 'cat > ~/run-.sh <<"EOF" +#!/bin/bash -l +set -euo pipefail +export PATH="$HOME/.pixi/bin:$PATH" +cd +pixi run --frozen snakemake -s experiments/scv2-spike/Snakefile \ + --config output_dir=results-- -j1 "$@" +EOF' +ssh 'tmux new-session -d -s smk- "bash ~/run-.sh > ~/smk-.log 2>&1"' +``` + +`/fh/fast` is a **shared filesystem** — every orca host sees the same worktree, +env, and results. Moving hosts needs no re-sync, no re-download, and completed +rule outputs carry over. Always `--dry-run` first: a rule's declared outputs +are deleted at job start, and `config.yaml` is a rule `input:`, so any edit to +it (even a comment) invalidates the expensive fit. + +### Monitoring a running fit + +**Never judge liveness by `%CPU`** — `ps` reports a *lifetime average* that +decays slowly, so a deadlocked worker can sit at "60%" for hours. Sample +**cumulative CPU time** twice instead; if `TIME` is identical across a 20 s +gap, the process is doing nothing: + +```bash +ssh 'ps -p -o pid,stat,time --no-headers' # run twice +``` + +`stat` of `S` with `wchan` `futex_wait_queue` / `pipe_read` across the whole +tree means a **`multiprocessing.Pool` deadlock**: `_fit_fun` swallows worker +exceptions (`except Exception: return None`, no logging, `model_collection.py`) +and `p.map` blocks forever when a child dies. Count the workers — fewer alive +than `n_processes` confirms it. See also +`~/.claude/.../project_convergence_maxiter100_spawn_deadlock.md`. + +To find *which* fit fails, run the grid serially with per-fit capture +(`fit_one_model` in a `try/except` with `traceback.print_exc()`); the pipeline +itself only reports a count. + +Other recurring snags: +- **HTTP 429 from `raw.githubusercontent.com`** in `prepare_data` — GitHub + rate-limits the raw-data download. Seed the cache from a prior run instead: + `cp -a /raw_data /raw_data` (verify md5s; do not + `mkdir -p` the destination first or `cp -a` nests it one level deep). +- **Stale snakemake lock** after a killed run — verify no snakemake is running + on *any* host sharing the filesystem, then `snakemake --unlock`. +- **`IncompleteFilesException`** after a kill — `--rerun-incomplete` regenerates + only the interrupted output. +- **SSH `Connection closed by UNKNOWN port 65535`** — sshd `MaxStartups` + throttling from too many rapid connections. Back off several minutes and + batch probes into one session; it is not an auth failure. + ## Active Technologies - marimo (interactive dashboard for exploring ModelCollection results) - Python 3.9+ (matches existing CI matrix) + multidms (this package), snakemake, papermill, jupyter, matplotlib, seaborn, pandas, numpy, pyyaml (002-simulation-pipeline) - CSV intermediate files + pickle for fitted model collections; all in `experiments/simulation/results/` (002-simulation-pipeline) - Python 3.9+ (matches existing CI matrix) + multidms (this package), snakemake, papermill, jupyter, matplotlib, seaborn, pandas, numpy, pyyaml, requests (for data download) (003-spike-pipeline) - CSV intermediate files + pickle for fitted model collections; all in `experiments/scv2-spike/results/` (003-spike-pipeline) -- Snakemake experiment pipeline in `experiments/loss-normalization/` for validating `.mean()` loss normalization against V0.4.0 hyperparameter anchors (fusionreg × l2reg 2D grid) ## Recent Changes - 002-simulation-pipeline: Added Python 3.9+ (matches existing CI matrix) + multidms (this package), snakemake, papermill, jupyter, matplotlib, seaborn, pandas, numpy, pyyaml diff --git a/HANDOFF.md b/HANDOFF.md index 522f8d86..44b2c1ae 100644 --- a/HANDOFF.md +++ b/HANDOFF.md @@ -1,411 +1,252 @@ # Handoff -Practical guide for whoever maintains `multidms` next. It covers the four -things that are hard to reconstruct from the code alone: how to get results -off the remote server and look at them, how to build and deploy the docs, -how the repository is laid out, and what state the manuscript work is in. +**For whoever writes the manuscript revision next.** This assumes you know +the science — deep mutational scanning, global epistasis, the joint shift +model, what the paper claims. It assumes you have never read this codebase. -Everything below was verified against the repository on **2026-08-13**. +It answers, in order: where the work stands, **how to get the results and +look at them** (start there if you have just cloned), what is still open, +which figures exist, how to run the pipeline yourself, how to build the docs, +and which scientific findings must survive the handoff. + +Verified against the repository on **2026-08-18**. Reference material lives elsewhere and is not repeated here: | For | Read | |---|---| -| Package architecture, code style, dev commands | [`CLAUDE.md`](CLAUDE.md) | +| Package architecture, code style, dev commands, the XLA JIT leak | [`CLAUDE.md`](CLAUDE.md) | | Contribution workflow | [`CONTRIBUTING.rst`](CONTRIBUTING.rst) | | Pipeline internals, config tiers, run log | [`experiments/scv2-spike/README.md`](experiments/scv2-spike/README.md) | | Pipeline index and remote setup | [`experiments/README.md`](experiments/README.md) | +| **The legacy analysis notebook** — where the *published* manuscript figures were generated | [`SARS-CoV-2_spike_multidms` @ `6c98b7b`](https://github.com/matsengrp/SARS-CoV-2_spike_multidms/blob/6c98b7b607d7387b508cdaa192d659ee9fca7367/spike-analysis.ipynb) | ---- - -## 1. Prerequisites +> ⭐ **That pinned commit is the old-vs-new comparison baseline.** The figures +> in the published preprint came from it, under multidms v0.4.0. When you need +> to know whether a number moved because the model changed or because the code +> changed, compare against `6c98b7b` rather than against the current legacy +> `main`, which has drifted. -The project declares its environment tool: **use `pixi`**, never bare -`python`/`pip`. - -```bash -pixi install # one-command environment setup -pixi run test # pytest + doctests -pixi run fmt-check # black -``` - -> ⚠️ `pixi run lint` is **vacuous on `experiments/`** — `ruff` is configured -> to exclude that tree. To lint experiment code, pass paths explicitly: -> `pixi run ruff check experiments/scv2-spike/notebooks/_downstream.py` - -### Remote access +--- -Long fits run on the Matsen lab's shared servers, configured **outside the -repository** in `~/.config/multidms-experiments/remote.yaml`: +## 1. Where things stand -```yaml -host: ermine -remote_dir: /fh/fast/matsen_e/shared/multidms/multidms -``` +**What you can trust today:** -`worktree_base` defaults to `/multidms-worktrees`. -Any key can be overridden per-invocation with `host=orca03`. +- The **spike production fit** is final: 10 lasso rungs × 2 replicates, at + `tol 1e-6 / maxiter 500`, with the chosen λ = **8.0e-05**. Its + `fit_collection.pkl` is 1,761,133,462 bytes, md5 + `b2b4736073e475a6fd7b1b5260063d6c`. +- The **simulation fit** is final at 649,380,559 bytes, md5 + `f602baf4801bb9257a1f22281da99f49`. +- **10 of the manuscript's 22 figures** regenerate from the live pipeline + today; §4 names the other 12 and who owns each. +- The two scientific results in §7 (**λ moved**, **A419S**) are measured on + this fit and are ready to go into prose. -> Prefer hosts `orca01`–`orca05` (64 cores, 1.5 TB RAM). **Avoid `quokka`.** -> Check load before launching anything. +**What is not done:** six simulation figures (#316), the linear baseline arm +for S10 (#293), and the manuscript prose itself. --- -## 2. Fetching results and viewing them +## 2. Getting the results and looking at them -### The one-paragraph version +**Start here if you have just cloned the repo.** This is the full path from a +fresh clone to viewing the manuscript figures and exploring the fits. It +should take under an hour, most of it waiting on the ~4.5 GB download. -A production run writes to `experiments//results--/` -on the **remote** host. `run-pull` rsyncs that directory back. The `results/` -symlink then points at whichever run is canonical, and both the docs build and -the dashboard read through it. +### Step 1 — set up the environment -### Step by step +The project uses [pixi](https://pixi.sh). It manages its own Python; you do +not need conda, venv, or a system Python. **Never run bare `python` or `pip` +here** — everything goes through `pixi run`. ```bash -# 1. Launch (from the branch whose name sets the output directory) -pixi run remote-pipeline -- spike prod host=orca03 - -# 2. Monitor — poll sparsely, every 30s+, never in a tight loop -pixi run remote-status -- spike prod host=orca03 +pixi install # solves and creates .pixi/envs/ — a few minutes +pixi run test # sanity check: 134 tests should pass +``` -# 3. Fetch when finished -pixi run run-pull -- spike prod host=orca03 +The default environment (Python **3.11**) is the one you want; `py39` / +`py310` / `py311` / `py312` exist only for the CI matrix +(`pixi run -e py39 test`). Everything below assumes the default. -# 4. Point `results/` at what you just pulled -ln -sfn results-prod- experiments/scv2-spike/results -``` +> ⚠️ **`pixi.lock` is gitignored, so a fresh clone re-solves from scratch.** +> You do not get the exact environment that produced these results. That +> matters more than it looks: the `jax` / `jaxopt` pins in `pyproject.toml` +> carry comments explaining that resolver drift there previously caused a +> **SIGABRT inside jaxopt's proximal solver**. If you hit an inexplicable +> crash during fitting on a fresh machine, suspect the resolve before you +> suspect the model. -`pipeline` is `simulation` or `spike`; `profile` is `test`, `experimental`, -or `prod`. +### Step 2 — configure the remote -**The output directory is derived from the current branch name**, not chosen -by you: branch `292-spike-fit-tuning` + profile `prod` → -`results-prod-292-spike-fit-tuning`. The remote tmux session is named -`smk--`. Attach with: +Results are **not in git** — they are ~4.5 GB of pickles, CSVs and PDFs living +on `ermine`. Create this file once: ```bash -ssh orca03 -t "tmux attach -t smk-spike-" +mkdir -p ~/.config/multidms-experiments +cat > ~/.config/multidms-experiments/remote.yaml << 'EOF' +host: your-username@ermine +remote_dir: /fh/fast/matsen_e/shared/multidms/multidms +EOF ``` -> Never leave tmux sessions running after `run-pull`. Clean up the remote -> worktree and session when a run is done. - -### The `results/` symlink +`host` and `remote_dir` are both **required**; the scripts abort with a +template if the file is missing. -`experiments//results` is a **gitignored symlink** to a -`results-*` directory — it does not exist in a fresh clone. Create it by -hand, pointing at whichever run the docs should publish: +### Step 3 — fetch both payloads and link them ```bash -ln -sfn results-prod-292-spike-fit-tuning experiments/scv2-spike/results -ln -sfn results-prod-sim-vpl500-tol1e5 experiments/simulation/results +REMOTE=your-username@ermine +BASE=/fh/fast/matsen_e/shared/multidms/multidms/experiments + +# spike — 3.3 GB +rsync -a --info=progress2 \ + "$REMOTE:$BASE/scv2-spike/results-prod-294-naive-baseline-arm/" \ + experiments/scv2-spike/results-prod-294-naive-baseline-arm/ +ln -sfn results-prod-294-naive-baseline-arm experiments/scv2-spike/results + +# simulation — 1.2 GB +rsync -a --info=progress2 \ + "$REMOTE:$BASE/simulation/results-prod-sim-vpl500-tol1e5/" \ + experiments/simulation/results-prod-sim-vpl500-tol1e5/ +ln -sfn results-prod-sim-vpl500-tol1e5 experiments/simulation/results + +pixi run check-results # ✓ per pipeline, or an explanatory ✗ ``` -To see what is linked now and what runs are available: +> ⚠️ **The `ln -sfn` lines are not optional.** `results` is a gitignored +> symlink, so it never survives a clone and nothing recreates it for you. +> Without it `pixi run docs` dies with a bare "No such file or directory" and +> the dashboard finds nothing — the most common failure on a fresh machine. + +### Step 4 — view the figures + +Rendered figures are in `results/figures/`, as **both PDF and PNG**: ```bash -pixi run check-results # both pipelines -bash experiments/scripts/check-results.sh spike # just one +open experiments/scv2-spike/results/figures/ # macOS ``` -`check-results` **only reports** — it never creates or repoints a symlink. -Choosing a run decides which numbers the published docs show, so that call is -left to a human. When the link is missing or broken it lists every run on -disk and prints the `ln -sfn` command to fix it. +The manifest in §4 maps every manuscript figure number to its filename — read +it before hunting, because **two filenames are deliberately misleading traps** +(§4). The executed notebooks (`results/*.ipynb`) carry the same figures inline +with the code that made them. -> ⚠️ **Never point `results/` into `.worktrees/`.** A worktree is removed when -> its branch lands, leaving a dangling symlink that breaks the docs build in -> the main clone. This has already happened once. +### Step 5 — explore the fits interactively -### The dashboard +An interactive [marimo](https://marimo.io) dashboard explores fitted +`ModelCollection`s — convergence, GE landscape, parameter correlation, +replicate scatter, sparsity. ```bash pixi run dashboard # read-only pixi run dashboard-edit # editable ``` -An interactive [marimo](https://marimo.io) app for exploring fitted -`ModelCollection`s — convergence, GE landscape, parameter correlation, -replicate scatter, sparsity. +> ⚠️ **Launch it from the repository root.** It discovers `*.pkl` files below +> the directory you launch it from (`cwd`), not from a fixed path. Dot-hidden +> directories (`.git/`, `.pixi/`, `.worktrees/`) are pruned from the search. -> ⚠️ **It discovers `*.pkl` files below the directory you launch it from -> (`cwd`), not from a fixed path.** Launch from the repository root to see -> every run. Dot-hidden directories (`.git/`, `.pixi/`, `.worktrees/`) are -> pruned from the search. +**The long waits are not hangs.** This is the most common way to mistake +working software for broken software here: -Two gotchas recorded from experience: +- Loading the spike `fit_collection.pkl` needs **~7 GB RSS** and takes a + while. The file is 1.76 GB. +- **The Param Correlation tab is slow by design.** It is button-gated: select + the fits, set the threshold, press **Plot**, then wait — **minutes** on a + prod-sized collection. That is expected. Do not kill it. +- `marimo` is pinned `>=0.8,<0.23` in `pyproject.toml`; 0.23.x breaks + `mo.ui.table` selection. Do not raise it. -- **Pin `marimo<0.23`.** 0.23.x breaks `mo.ui.table` selection. -- Loading `fit_collection.pkl` for the spike prod run needs **~7 GB RSS** - (the file is 1.76 GB). Prefer the exported CSVs when you only need numbers. +### What is in a results directory ---- - -## 3. Building and deploying the docs - -```bash -pixi run docs # clean + build to docs/_build/html -pixi run docs-deploy # build, then push to the gh-pages branch -``` - -Published at from the `gh-pages` -branch, via `ghp-import`. - -### Why the docs need a completed run - -Ten `docs/*.nblink` files point at **executed** notebooks inside -`experiments//results/`: +| File | What it holds | +|---|---| +| `figures/` | The rendered PDFs and PNGs listed in §4. | +| `mutations_df.csv` | Per-mutation β and shifts — **the table most manuscript numbers come from.** | +| `cross_validation_loss.csv` | CV loss per λ rung; the evidence for the λ choice. | +| `fit_sparsity.csv`, `library_replicate_correlation.csv` | The other two λ selection criteria. | +| `fit_convergence.csv`, `convergence_trajectory.csv` | Per-fit convergence; check before trusting any fit. | +| `fit_collection.pkl` | The fitted `ModelCollection`. Large (1.76 GB for spike), ~7 GB RSS to load. | +| `*.ipynb` | The executed notebooks, with outputs, for every pipeline stage. | -``` -docs/spike_evaluate.nblink → ../experiments/scv2-spike/results/evaluate.ipynb -``` +> ⭐ **Prefer the CSVs.** Almost every number in the manuscript can be read +> from `mutations_df.csv` and the three selection-criterion CSVs in seconds, +> without ever loading the pickle. -Because `results/` is gitignored, a fresh clone has nothing to resolve, and -Sphinx fails with a bare, misleading error: +### Where this all lives on the remote ``` -InputError: [Errno 2] No such file or directory: - '../experiments/simulation/results/cross_validation.ipynb' +/fh/fast/matsen_e/shared/multidms/ +├── multidms/ ← canonical clone, on up-to-date main +│ └── experiments/ +│ ├── scv2-spike/ +│ │ ├── results-prod-294-naive-baseline-arm/ (3.3 GB) +│ │ └── results -> results-prod-294-naive-baseline-arm +│ └── simulation/ +│ ├── results-prod-sim-vpl500-tol1e5/ (1.2 GB) +│ └── results -> results-prod-sim-vpl500-tol1e5 +└── archive/ + ├── archive-2026-08-18/ ← 9 superseded payloads + loose dirs + └── ``` -**This is not a Sphinx problem.** `pixi run docs` and `docs-deploy` now depend -on `check-results`, which runs first and stops the build with a readable error -naming the pipeline, the runs available on disk, and the `ln -sfn` command to -fix it. Sphinx never starts, so there is no half-built output to clean up. - -> Diagnostic habit: when a docs build fails, run `readlink experiments/*/results` -> **first**. - -### Adding a docs page for a new analysis - -Every new pipeline analysis gets a page. Three steps: - -1. Create `docs/spike_.nblink`: - ```json - {"path": "../experiments/scv2-spike/results/.ipynb"} - ``` -2. Add `spike_` to the `Spike Analysis` toctree in `docs/index.rst`. -3. Give the notebook real narrative markdown. **These pages are the public - documentation of the method, not an execution log.** - -Verify with `pixi run docs` before opening a PR. +Nothing was deleted in the cleanup, only moved — if you need an older run, it +is under `archive/`. --- -## 4. Repository structure +## 3. Issue status -``` -multidms/ -├── multidms/ # the package -│ ├── jaxmodels.py # JAX-native core (equinox, BCOO sparse, jaxopt) -│ ├── data.py model.py # pandas/binarymap wrapper API over jaxmodels -│ ├── model_collection.py# parallel fitting over parameter grids, CV -│ ├── plot.py # ALL Altair rendering; classes delegate here -│ └── utils.py # mutation-string parsing, transforms -├── experiments/ # analysis pipelines (see below) -├── docs/ # Sphinx sources + .nblink stubs -├── tests/ -└── HANDOFF.md # this file -``` - -The package has **two API layers**: `jaxmodels` is the JAX-native core; -`data.py`/`model.py` are the friendlier pandas-facing wrappers. Convert -between them with `jaxmodels.Data.from_multidms()`. See `CLAUDE.md` for the -full architecture. +Re-queried 2026-08-18: **14 open issues** — 13 once this change lands and +closes #297 — and **#282 is the only open PR** besides it. -### `experiments/` +### Live — someone should act on these -| Directory | Status | What it is | +| # | What | Note | |---|---|---| -| `simulation/` | **Live pipeline** | Synthetic DMS with known ground truth. Manuscript Fig 2, S1–S5. | -| `scv2-spike/` | **Live pipeline** | SARS-CoV-2 spike DMS. Nine manuscript figures. | -| `scripts/` | **Infrastructure** | Remote execution + `check-results.sh`. | -| `dashboard.py`, `dashboard_helpers.py` | **Live tooling** | The marimo dashboard. | -| `convergence-lab/` | **Evidence record — keep** | A lab notebook, *not* a pipeline. Its README's "standing findings" are the cited justification for production hyperparameters in both live pipelines (e.g. `simulation/config/config.yaml` refers to it by name). Deleting it orphans those citations. | - -### How a pipeline is wired - -Both pipelines are Snakemake workflows executing parameterized notebooks via -papermill. Source notebooks live in `notebooks/`; executed copies land in -`results/`. - -The spike DAG: - -``` -prepare_data ──► training_functional_scores.csv - ├──────────────────────┐ - ▼ ▼ - fit_models cross_validation - │ │ - ▼ │ - evaluate ──► mutations_df.csv, collection_muts.csv, … - └──────────┬───────────┘ - ▼ - manuscript_figures ──► figures/*.pdf, *.png -``` - -### ⚠️ The config tier split — the thing most likely to be broken by accident - -The config is split so a **downstream-only edit cannot invalidate a ~2h20m -model fit**: - -| File | Holds | Invalidates the fit? | +| **#316** | Emit Fig 2 and S1–S5 from the simulation pipeline | The largest remaining figure gap. Mostly renames — see §4. | +| **#293** | Linear (Identity) baseline arm → SI S10 | Spec'd, unblocked, no compute dependency left. | +| **#282** | v0.4.0 ↔ main equivalence check | **The only open PR.** Open since 2026-07-15. Finished work; it is the evidence resolving #281. Land or close it. | +| **#318** | Write the three missing docs pages | The placeholders they replace were deleted. | +| **#313** | Spike prep diverges from legacy (codon deletions, replicate subset) | Affects data prep, not the fitted model. | +| **#312** | `mut_type()` mislabels in-frame codon deletions | Related to #313. | +| **#192** | Condition-specific mutation names in `get_mutations_df` | Fully spec'd. #302 was closed as a duplicate of it. | +| **#319** | `IndexError` in `mut_param_dataset_correlation` | Live on `main` at `model_collection.py:1471`, but **benign for the manuscript**: it needs a `(mut_param, x)` cell surviving in only one replicate, which arises under `strategy: continuation`. Every published figure comes from independent-strategy fits. Branch `fix/mut-param-correlation-1col` has a starting patch. | +| **#179** | Remove deprecated `phenotype_as_effect` | Small cleanup. | +| **#99** | Citation | Small. | + +### Parked and reference-only + +| # | What | Why | |---|---|---| -| `config.yaml` | `fitting:` block, data sourcing, filtering | **Yes** | -| `config_downstream.yaml` | `lasso_choice`, colors, `domain_dict`, `figures:` | **No** | - -Rules to preserve: - -1. Do **not** add `config_downstream.yaml` to the `input:` of `prepare_data`, - `cross_validation`, or `fit_models`. That silently restores the defect. -2. New downstream helpers go in `notebooks/_downstream.py`, **never** - `_common.py` — the latter is `input:` on all four fit-tier rules. -3. `manuscript_figures` reads **CSVs only**; it must never load - `fit_collection.pkl`. -4. Every config variant needs a matching `_downstream.yaml` sibling. - The path is derived by string substitution, so a missing sibling fails. - -> **`n_processes: 6` is pinned for spike.** Steady-state RSS is ~35–38 GB per -> worker, so six workers need ~230 GB. Restoring `null` auto-sizes by core -> count (~64 workers) and needs multiple TB. This has OOM'd a host before. - -Note `maxiter` is overloaded: top-level = outer sweeps; inside -`ge_kwargs`/`cal_kwargs` = inner solver steps. +| #290 | The epic this work was tracked under | Its phases are done or re-homed; close it when #316 and #293 land. | +| #243, #51 | Concurrent single-solve fitting; a ridge penalty that doesn't bias toward WT | Design questions, not manuscript blockers. | +| #281 | Re-express the 0.4.0 spike model in main's form | Already executed by PR #282. | -### Config variants that look like cruft but are not +> **Issues labelled `question` or `wontfix` were kept only for reference.** +> They record decisions and dead ends, not work. Unless you find one that +> matters to you, they can safely be deleted. -`config_recompute_false*.yaml` (three pairs) are unreachable from the -Snakefile by profile name and look like leftovers from a finished experiment. -They are **test fixtures**: `tests/test_config_tiers.py` iterates -`SPIKE_VARIANTS` and asserts on each. Deleting them fails four tests. +> ⚠️ **#240** is closed, but its `fit_models_path` truncation bug was never +> fixed in code. It is live and low-severity: +> prod sets no `strategy` key, so it defaults to `"independent"`. It becomes a +> real trap only if you switch to `"continuation"`. -Removing them is a deliberate two-step change — edit `SPIKE_VARIANTS` first, -then delete the YAMLs. +> **#295** was closed as delivered, and the log-x requirement for Figure 5 was +> **dropped, not deferred**. Its closed body remains the richest record of the +> Figure 4 zoom regions and the Figure 5 x-axis warning: that axis is +> `2 ** avg_predicted_func_score`, the predicted enrichment ratio — **not** β +> and **not** shift, despite the legacy name `predicted_beta`. --- -## 5. State of the manuscript work - -The active spine is **EPIC #290** — regenerating every manuscript figure from -the current model rather than the archived v0.4.0 notebook. - -``` -Phase 1 #291 ✅ landed (PR #303) simulation convergence -Phase 2 #292 ✅ landed (PR #309) spike refit + the nine-figure surface -Phase 3 #293 🔻 DESCOPED 2026-08-17 → standalone issue, spec'd and ready -Phase 4 #294 ✅ landed (PR #311) naive per-condition baseline → Fig 3 -Phase 5 #295 ✅ closed as delivered-by-Phase-2 (figures shipped in PR #309) -Phase 6 #296 ⬜ stub ▶ UNBLOCKED figure manifest + number-diff -Phase 7 #297 ⬜ stub written handoff for manuscript revision -``` - -**The epic's remaining work is #296 then #297 — both local, no compute.** All -four fit-bearing phases have landed and every remote run the epic needs is done. - -### Phase 3 was descoped — what that means - -**The linear (Identity) baseline arm, SI Figure S10, is not part of EPIC #290 -anymore.** It lives at **#293** as a standalone issue carrying its full spec, -and it is unblocked today: it needs only the `(tol, maxiter)` pair from #291 -and the cached spike fit from #292, both landed. - -Nothing in the repo implements it yet — there is no `linear_baseline.ipynb`, no -`rule linear_baseline`, no `spike.linear` config block, no `S10` entry in the -Snakefile's `FIGURE_NAMES`. A reader grepping for those and finding nothing is -seeing the correct state, not a broken checkout. - -Consequences to carry into any manuscript work: - -- **S10 is a carried-over-unchanged figure**, alongside S8 and S13–S15. It is - the one SI figure still showing v0.4.0 output while its neighbours were refit - under the new `(tol, maxiter)` and λ = 8.0e-05. -- **The linear-vs-sigmoid loss gap is unmeasured** — not "unchanged", and not - "moved". The paper's S10 claim is untested by this work. -- The paper's **central methodological claim** (joint R² ≈ 3.4× naive) is - unaffected: that is Figure 3, delivered by Phase 4. - -> ⚠️ Whoever picks up #293 should re-check its §1a analysis against the -> *current* fit rather than the state of the world when the spec was written. - -### Two live scientific results from the Phase 2 refit - -**λ moved: `4.0e-05` → `8.0e-05`.** All three selection criteria (CV loss, -replicate correlation, stop-codon sparsity) now agree on the chosen rung, but -the two rungs sit within **0.16%** of each other — corroboration, not a -decisive vote. The Methods paragraph in `main.tex` needs updating. - -**A419S retains its contrast — direction preserved.** At λ = 8.0e-05, -`2 ** avg(predicted_func_score)`: - -| | Delta | BA.1 | BA.2 | -|---|---|---|---| -| phenotypic effect | 0.854 | 0.132 | 0.137 | -| fold vs Delta | — | 6.4× | 6.2× | - -> ⚠️ The paper's **">1,000-fold"** figure describes **measured titers**, not -> the model's predicted enrichment ratio. The model reproduces the contrast -> *direction and ordering*. Stating it otherwise reads as a failed -> replication when it is not. - -### A suspected defect that turned out not to be one +## 4. Figure status -**The Figure S10 "erratum" was investigated and refuted. Do not report it.** - -An earlier version of this document — and the bodies of #290, #296 and #297 — -stated that legacy notebook cell 103 plots the *sigmoid* collection's CV loss -inside the linear-model figure, making the published S10 middle panel an -erratum against the preprint. **That is false.** The spec work on #293 recovered -the notebook at `fc89753:notebooks/spike-analysis.ipynb` (the previously cited -`6c98b7b` does not resolve in this repo) and read it by 0-based cell index: - -| idx | exec | what it actually does | -|---|---|---| -| 103 | 127 | the linear **fit call** — not a loss call | -| 104 | 128 | builds `linear_mc`, adds validation loss | -| 105 | 132 | `cross_validation_df = linear_mc.get_conditional_loss_df()` — **linear, and read** | -| 106 | 133 | renders `shrinkage_analysis_linear_models` — the S10 figure | - -Execution counts are monotone 128 → 132 → 133, so the last write to -`cross_validation_df` before S10 rendered was the linear one. Cells 101/102 -likewise rebind `sparsity_df` and `corr_df` from the linear collection, so -panels A and C are linear too. - -> ⚠️ **Do not tell Hugh the preprint contains an S10 erratum.** Reporting a -> defect that is not there is worse than reporting nothing, and with Phase 3 -> descoped there is no reproduction run left in the epic to catch the mistake -> before it reaches the manuscript. -> -> **The supportable sentence:** *"A suspected defect in S10 was investigated -> and refuted. S10 was not regenerated, so the check is not yet decisive — -> #293 carries the reproduction that would settle it."* - -The evidential limit is real and is why #293 still treats this as live: stored -execution counts record *an* execution order, not proof the saved PDF came from -it. But "unverified" is not "erratum". Full analysis: **#293 §1a**. - -The separate warning that cells 98–100 rebind module-level frames is still a -genuine *fragility* — the two arms share variable names, so a re-run in a -different order would silently mix them — and #293's spec keeps the arms in -separate namespaces for that reason. - -### Notation (paper ↔ code) - -| Paper | Meaning | Code | -|---|---|---| -| `β_m` | mutation effect in the reference experiment | `beta` | -| `Δ_{d,m}` | shift in experiment `d` vs reference | `shift` | -| `λ` | **"lasso regularization weight"** | `fusionreg` | -| `α_d` | experiment offset | `alpha` | -| `θ₀, θ₁` | sigmoid bias & scale | `theta` | - -> The paper never says "fusion regularization". Use **λ / "lasso -> regularization weight"** in prose and captions; `fusionreg` in code. - -### Figure manifest — what's regenerated, what's missing (#296) The manuscript includes **22 figures** (via `\includegraphics` in -`main.tex`/`si.tex` at `f79ac4a`, excluding two commented-out template +`main.tex`/`si.tex` at `f79ac4a`, excluding four commented-out template placeholders). The current spike pipeline (`experiments/scv2-spike/results/figures/`, symlinked to `results-prod-294-naive-baseline-arm/`) regenerates **10** of them. @@ -455,8 +296,9 @@ The other **12** are listed below with an owner for each. > work against the already-cached `fit_collection.pkl` (verified > 2026-08-18: `snakemake --touch` cleared an mtime-only cascade with > the 649,380,559-byte pickle staying byte-identical, and a subsequent -> forced dry run against a figure target reported `total: 1`, rule -> `manuscript_figures` alone). S3, S5, and S1's panels B/C need two new +> forced dry run against a figure target re-ran only +> `manuscript_figures` and the `all` aggregator — `total: 2`, no +> fit-tier rule). S3, S5, and S1's panels B/C need two new > simulation conditions and do require a simulate-and-fit pass. See > #316 for the full breakdown. > @@ -467,51 +309,134 @@ The other **12** are listed below with an owner for each. > `fig:shifts_3D_structure` but includes > `structure_and_neighbor_statistics_scatter.pdf` (a decoy > `shifts_3D_structure.pdf` also exists on disk and is not the included -> file). Both mismatches are deliberate and already documented in the -> Snakefile's `FIGURE_NAMES` comment — match on what's included, never -> on the label. +> file). Both mismatches are deliberate: match on what's included, never +> on the label. The Figure 3 trap is also recorded in the Snakefile's +> `FIGURE_NAMES` comment; **S13's is not, because S13 has no producer +> yet** — if you write one (#316-style), carry this warning into it. --- -## 6. Open loose ends - -**Unmerged PRs:** - -- **#282** — v0.4.0 ↔ main equivalence check. Finished work, open since - 2026-07-15. It is the evidence resolving #281 and half of #242. Land or - close it. -- **#239** — `IndexError` fix in `mut_param_dataset_correlation`, open since - 2026-05-07. - -**Backlog notes:** - -- **#240** (bug) is live and unfixed: `fit_models_path` truncates paths - silently. Lower severity today because prod's `config.yaml` sets no - `strategy` key at all, so it defaults to `"independent"` — but a real trap - if you switch to `"continuation"`. -- **#281** duplicates #242 and is already executed by PR #282. -- **#176 / #177 / #179** reference a **v2.0 release that never happened** - (tags go 0.4.2 → 1.0.0 → 1.3.0). #177's target module, `biophysical.py`, - no longer exists. Retire the "v2.0" vocabulary rather than trying to - satisfy it. -- **#192 and #302** overlap substantially — both re-express mutation effects - in a condition's own coordinates. Merge them when #302 is specced. -- **#295** was **closed as delivered-by-Phase-2**; the log-x requirement for - Figure 5 was explicitly dropped, not deferred. Its (closed) body remains the - richest record of the Figure 4 zoom regions and the Figure 5 x-axis warning — - the axis is `2 ** avg_predicted_func_score`, the predicted enrichment ratio, - **not** β and **not** shift, despite the legacy name `predicted_beta`. The - closing audit comment restates both. Note also that the manuscript's - ">1,000-fold" A419S claim refers to **measured titers**, not the model's - predicted ratio (predicted: Delta 0.854 / BA.1 0.132 / BA.2 0.137). - -**Housekeeping:** - -- `experiments/loss-normalization/` is dead: nothing references it, and it - still uses the pre-tier-split single-argument `load_config()`, so it could - not run today. Safe to delete along with its `CLAUDE.md` line. -- Branch `246-convergence-lab` is **ahead of its remote by one commit** - (`5ac7477`, "preserve uncommitted SWEEP_PLAN + sweep runner"). Push or - discard it before deleting the branch. -- Two `.claude/worktrees/agent-*` worktrees hold unreviewed experiments - (an `alpha_ridge` knob; a `BiasedSigmoid` GE with a fitted lower plateau). + +--- + +## 5. Running the pipeline + +Do this whenever the revision needs numbers that do not exist yet: a new +model arm, a different λ grid, an ablation a reviewer asked for, a rerun after +a code change. Both pipelines are Snakemake workflows driven by a profile. + +### Start with the test profile + +```bash +pixi run spike-test # ~10 min, 10% subsample +pixi run sim-test # ~5 min +``` + +**Always run `-test` before `-prod`.** It exercises the whole DAG end to end +in minutes, so a broken config, a bad path, or a notebook that raises fails +immediately rather than after hours of fitting. The prod profiles are +`pixi run spike-prod` and `pixi run sim-prod`; spike also has `experimental`. + +### Production runs go on a remote host + +The spike prod fit needs ~35–38 GB RSS **per worker** with 20 workers. That is +a server job, not a laptop job. + +```bash +# 1. pick an idle host first — a busy one will thrash at 20 workers +# (lab tooling: `bip scout`; otherwise check load however you normally do) +# 2. launch, poll, fetch: +pixi run remote-pipeline -- spike prod host=orca03 # launches in tmux +pixi run remote-status -- spike prod host=orca03 # poll progress +pixi run run-pull -- spike prod host=orca03 # fetch results back +``` + +All three take ` [key=value ...]`. The `--` matters: it +separates pixi's own arguments from the script's, and `host=` overrides +`~/.config/multidms-experiments/remote.yaml`, which you create once with +`host:` and `remote_dir:` keys (see `experiments/README.md`). + +**Commit before launching.** The launcher warns on a dirty tree because the +remote checks out your branch — it never sees uncommitted work. When the run +finishes, kill the tmux session; leaving it holds the host. + +> ⭐ **A new run cannot overwrite the manuscript payload.** `output_dir` is +> derived automatically as `results--`, and non-`main` +> branches get their own remote worktree. Work on a branch and your run lands +> in its own directory; the manuscript's `results-prod-294-naive-baseline-arm` +> is untouched. This is the property that makes experimenting safe. + +### Changing what gets fit + +Fitting knobs live in the `fitting:` block of each pipeline's `config.yaml` — +λ grid, `tol`, `maxiter`, the loss. Editing that file **intentionally** +invalidates the fit, which is exactly what you want when changing the model. +The **config tier split** — which edits cost a refit and which do not — is +documented in +[`experiments/scv2-spike/README.md`](experiments/scv2-spike/README.md); read +it before editing anything under `config/`. + +For a genuinely different configuration, prefer a **config variant** — +`config_.yaml` plus its required `config__downstream.yaml` +sibling — over editing the production config in place. That keeps the +manuscript's configuration reproducible alongside your new one. + +Pipeline internals, the DAG, and the run log of past production fits are in +[`experiments/scv2-spike/README.md`](experiments/scv2-spike/README.md); the +remote setup is in [`experiments/README.md`](experiments/README.md). + +--- + +## 6. Building the docs + +```bash +pixi run docs # build to docs/_build/html +pixi run docs-deploy # publish to gh-pages +``` + +> ⚠️ **The build needs the `results/` symlink.** The `.nblink` files in +> `docs/` point at `experiments/*/results/*.ipynb`. Those symlinks are not +> tracked in git, so a **fresh clone fails to build** until you fetch a +> results payload (§2) and recreate them. This is the usual cause of a +> mystifying docs failure on a machine that has never run a pipeline. + +The docs render the executed pipeline notebooks directly, so the published +site reflects whichever run `results` points at. + +--- + +## 7. Science to carry forward + +**λ moved: `4.0e-05` → `8.0e-05`.** All three selection criteria (CV loss, +replicate correlation, stop-codon sparsity) now agree on the chosen rung, but +the two rungs sit within **0.16%** of each other — corroboration, not a +decisive vote. The Methods paragraph in `main.tex` needs updating. + +**A419S retains its contrast — direction preserved.** At λ = 8.0e-05, +`2 ** avg(predicted_func_score)`: + +| | Delta | BA.1 | BA.2 | +|---|---|---|---| +| phenotypic effect | 0.854 | 0.132 | 0.137 | +| fold vs Delta | — | 6.4× | 6.2× | + +> ⚠️ The paper's **">1,000-fold"** figure describes **measured titers**, not +> the model's predicted enrichment ratio. The model reproduces the contrast +> *direction and ordering*. Stating it otherwise reads as a failed +> replication when it is not. + +### Notation (paper ↔ code) + +| Paper | Meaning | Code | +|---|---|---| +| `β_m` | mutation effect in the reference experiment | `beta` | +| `Δ_{d,m}` | shift in experiment `d` vs reference | `shift` | +| `λ` | **"lasso regularization weight"** | `fusionreg` | +| `α_d` | experiment offset | `alpha` | +| `θ₀, θ₁` | sigmoid bias & scale | `theta` | + +> The paper never says "fusion regularization". Use **λ / "lasso +> regularization weight"** in prose and captions; `fusionreg` in code. + + +*Verified against the repository on 2026-08-18.* diff --git a/docs/biophysical_model.ipynb b/docs/biophysical_model.ipynb deleted file mode 100644 index b42c54c2..00000000 --- a/docs/biophysical_model.ipynb +++ /dev/null @@ -1,27 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "placeholder-cell", - "metadata": {}, - "source": [ - "# Biophysical Model\n", - "\n", - "COMING SOON" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "name": "python", - "version": "3.11.0" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/docs/fit_delta_BA1_example.ipynb b/docs/fit_delta_BA1_example.ipynb deleted file mode 100644 index 2880f95d..00000000 --- a/docs/fit_delta_BA1_example.ipynb +++ /dev/null @@ -1,27 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "placeholder-cell", - "metadata": {}, - "source": [ - "# Fit Delta BA.1 Example\n", - "\n", - "COMING SOON" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "name": "python", - "version": "3.11.0" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/docs/index.rst b/docs/index.rst index b4554be3..4a15472b 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -20,10 +20,6 @@ and how much the effects differ between experiments. - The preprint is available on `bioRxiv `_. -- A concise description of the joint modeling approach is available in the `biophysical model `_ section. - -- A example fitting python with the python interface is available in the `usage examples documentation `_ page. - - For a more advanced example of the multidms interface, see our `manuscript SARS-CoV-2 spike analysis `_. - The source code is `on GitHub `_. @@ -40,9 +36,6 @@ and how much the effects differ between experiments. :caption: Contents installation - biophysical_model - simulation_validation - fit_delta_BA1_example multidms acknowledgments contributing diff --git a/docs/simulation_validation.ipynb b/docs/simulation_validation.ipynb deleted file mode 100644 index 33365ea4..00000000 --- a/docs/simulation_validation.ipynb +++ /dev/null @@ -1,27 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "placeholder-cell", - "metadata": {}, - "source": [ - "# Simulation Validation\n", - "\n", - "COMING SOON" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "name": "python", - "version": "3.11.0" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/experiments/convergence-lab/.gitignore b/experiments/convergence-lab/.gitignore deleted file mode 100644 index fbca2253..00000000 --- a/experiments/convergence-lab/.gitignore +++ /dev/null @@ -1 +0,0 @@ -results/ diff --git a/experiments/convergence-lab/CLAUDE.md b/experiments/convergence-lab/CLAUDE.md deleted file mode 100644 index a95a673f..00000000 --- a/experiments/convergence-lab/CLAUDE.md +++ /dev/null @@ -1,8 +0,0 @@ -# convergence-lab — agent instructions - -After EVERY harness run in this directory, append a Run log entry to README.md: -the config swept, what the fit collection showed (basin diagnostics, replicate -correlation — computed downstream from a ModelCollection), the conclusion, and -the next step. -Update Standing findings when a result confirms or overturns one. Never create -FINDINGS/SWEEP_PLAN files — README.md is the single source of truth for this lab. diff --git a/experiments/convergence-lab/README.md b/experiments/convergence-lab/README.md deleted file mode 100644 index 13c182ff..00000000 --- a/experiments/convergence-lab/README.md +++ /dev/null @@ -1,279 +0,0 @@ -# convergence-lab - -A fast, local harness for diagnosing the scv2-spike convergence and -reproducibility problem (issue #253; built on the `recompute_scale` fix, #246). -The full remote pipeline takes hours per turn; this fits tiny grids locally in -minutes so we can iterate on the two leading suspects (the α/β see-saw and the -L2 knee / fixed-scale interaction). - -## How to run - - pixi run python experiments/convergence-lab/harness.py \ - --config grids/smoke.yaml --cache smoke - # parallel worker count defaults to all-but-one core (capped at grid size); - # override per machine with --n-processes N - -Writes `results//fit_collection.pkl` — a *true* fit collection (the raw -`stack_fit_models` frame, the same schema the scv2-spike pipeline writes; -gitignored, regenerated on demand). A new experiment is a new `grids/*.yaml`; -the runner is generic and owns the constant scv2-spike data. - -The harness **only fits models** — it computes no derived metrics. Everything -downstream (replicate-shift correlation, basin diagnostics, plots) is computed -on the fly from a `ModelCollection` built over the frame, per experiment: - - import pickle - from multidms.model_collection import ModelCollection - mc = ModelCollection(pickle.load(open("results/smoke/fit_collection.pkl", "rb"))) - _, corr = mc.mut_param_dataset_correlation(x="fusionreg", return_data=True, r=1) - -**Parallelism:** the harness always fits via `fit_models(n_processes=N)` (the -`spawn` worker pool). This is safe because the fit call lives under the script's -`if __name__ == "__main__":` guard — see the warning on -`multidms.model_collection.fit_models`. Each worker rebuilds its own dataset -copy, so peak RAM grows with `n_processes` × dataset size; size `--n-processes` -to the host's free memory, not just its core count. - -**Visual review:** explore the fits in the multidms dashboard — from this -directory run `pixi run dashboard` (it discovers `fit_collection.pkl` below cwd). - -## Standing findings - -- **α/β see-saw degeneracy** (measured): `warmstart=True` → α 3–8 / Σβ² 350–1400 - (healthy); `warmstart=False` → α 1.4–3.3 / Σβ² 1700–75000 (β exploded so α·φ - stays bounded and Huber loss barely moves). Two near-degenerate basins fit - equally well, so noise decides which one each replicate lands in. -- **L2 knee near `3e-4`**; double-ended degeneracy: β explodes at `l2reg=0`, - β collapses + α explodes at `l2reg ≥ 1e-3`. **Below the knee, L2 is inert - against an active clip** (measured, #284, `cache=beta0-ridge-l2-scan`): - `l2reg` at 1e-8/1e-7/1e-6 — three-to-four orders under the knee — moves - nothing on top of `beta_clip_range=[-10,10]` across a 72-fit grid. Σβ² 0.96–1.01× - and α 1.00–1.11× vs the `l2reg=0` baseline, replicate-shift r within ±0.007, - all far under the pre-registered 2× / 0.05 effect thresholds. The clip is the - binding constraint on β; a whisper of L2 underneath it is a no-op. (The knee - is where L2 *starts* to bite — this pins the other end of that statement.) -- **A nonzero `beta0_ridge` (1e-4…1e-2) is also inert at this fixed block** - (measured, #284, same grid — the axis's first sweep in this lab). **What it - penalizes matters and is easy to get wrong:** `_beta_ridge_penalty` - (`jaxmodels.py:533-540`) is NOT an intercept ridge — it penalizes each - non-reference condition's β0 *difference from the reference*, - `beta0_ridge · Σ_{d≠ref}(β0_d − β0_ref)²`, never the reference's own β0. So it - is an **L2 shift-shrinkage on the intercepts**, structurally parallel to - `fusionreg`'s L1 on the β shift — cross the two and read its effect *within* - each fusionreg level, never marginalized. Its penalty is also **exactly 0 at - initialization** wherever `beta0_init` pins all conditions to the same value - (as every grid here does), and only bites as the β0s separate. Effect at - 1e-4…1e-2: α rises **monotonically** (16.1 → 16.5 → 17.5) and shift_Delta r at - fusionreg=0 lifts 0.594 → 0.617, so the penalty genuinely touches the optimum — - but every move is 1-2 orders below the effect thresholds. It is too weak to - matter here, not structurally dead; a stronger axis was not tested. -- **Caveat on both of the above — they were measured where convergence was - already 0/8.** #284's baseline is 0/8 converged at inner maxiter=10, so that - grid could not have detected a convergence *gain* from any regularizer. The - defensible claim is *"regularization cannot rescue convergence at an - inadequate inner cap"*, NOT *"these regularizers are inert"* in general. #285 - re-runs the same axes at inner maxiter=100 (the only 100%-converging cell - known) to separate the two. -- **A small `l2reg>0` tames the β-explosion across the whole prod fusion axis** - (measured, #256, `cache=l2-fusion`): at `l2reg=0` Σβ² sits at ~16–19k for - *every* `fusionreg` (0 → 6.4e-4); `l2reg=1e-4` collapses it ~25× into the - healthy 350–840 band and `3e-4` further to ~190–300, both holding at the - prod-max `fusionreg=6.4e-4`. `1e-4` is the working weight (α ~20–25); `3e-4` - over-regularizes toward the see-saw's collapse end (α ~30–40). **Strong - fusion is reproducibility-rescued only once β is penalized:** at `l2reg=0` - the `6.4e-4` column collapses `shift_Delta` replicate-r to 0.09, but at - `l2reg=1e-4` strong fusion *lifts* `shift_Omicron_BA2` r to 0.80 — the - data-poor-condition distortion the path-fitter targets is itself a symptom - of the unpenalized β-explosion. -- **β-bound choice — settled** (#263, Phase 1): head-to-head at cold-start / - recompute_scale=false / tol=1e-6, **`beta_clip_range=[-10,10]` wins** the - tri-criteria race against `l2reg=1e-4`. Clip is the *only* arm that converges - to tol=1e-6 (6/6, with 4/6 crossings in the epic's 20–50-iter band); l2 never - reaches tol (sentinel 101 on all 6). Clip keeps α~3 (healthy) while l2's α - blows up to ~55 (the see-saw's collapse end). Clip's strong-fusion - shift_Omicron_BA2 replicate-r reaches 0.81 vs l2's 0.66. Note the twist: clip - wins with a *large* Σβ² (~110k–181k) that is harmless because the hard φ-cap - keeps α·φ bounded — "large Σβ²" only signals an explosion when the fit *also* - fails to converge (as at l2reg=0 in #256). The epic's downstream phases inherit - the clip bound. -- **A free (per-condition) α deepens Delta's functional-score floor toward the - true −3.5** (measured, #273, `cache=273-free-alpha`): sweeping - `share_alpha` True/False across `fusionreg [0, 4e-5, 1.6e-4, 6.4e-4]` at PR - 270's exact fitting block (`warmstart=false`, `l2reg=0`, `recompute_scale=false`, - Sigmoid, inner blocks maxiter=10/maxls=10, outer tol=1e-5/maxiter=50). Delta's - predicted floor = `−α_Delta · sigmoid(φ_wt,Delta)` (the low asymptote of - `α·(g(φ)−g(φ_wt))` as φ→−∞). At the prod-strength `fusionreg=6.4e-4`, freeing α - lets Delta take its **own** α (5.2, decoupled from BA1's 17.8 / BA2's 14.4) - and drops Delta's floor from **−2.16 (shared) → −3.09 (free)** — most of the - way to the true −3.5, where the shared α is pinned by the other two conditions. - Counter-intuitively the deeper floor comes with a *smaller* Delta α: freeing α - lets the Delta latent push φ_wt more positive so `−α·sigmoid(φ_wt)` deepens. - The effect grows with fusion (floor gap free−shared: −0.15 at fusionreg=0, - −0.93 at 6.4e-4). BA1/BA2 floors barely move. All 16 fits hit maxiter=50 - (0/16 converged, same as PR 270's 0/18 at tol=1e-6 — the looser tol=1e-5 - didn't buy convergence), but obj_err is tiny (≤3e-4), so the α/floor numbers - are trustworthy while the binary flag is not. -- **Softplus floor × free-α is a complementary 2×2 — only BOTH together floor the - prod tail to the biological −3.5** (measured, #277, full in-harness 2×2: - `cache=277-softplus-floor{,-off}` share_alpha=true, `…-freealpha` / - `…-off-freealpha` share_alpha=false; `output_floor=-3.5` hinge 0.1 vs null; - 4 fusion × 2 reps each = 32 fits, 0 failed; Delta floor computed two ways — - analytic `t(−α·g(φ_wt))` and model composition at φ=−1e4 — agreed to 0.0e0 on - all 32). Delta floor (mean over 2 reps): - ``` - fusionreg free-OFF free-ON shared-OFF shared-ON - 0.0 -5.24 -3.50 -5.09 -3.50 - 4e-5 -4.82 -3.50 -4.69 -3.50 - 1.6e-4 -3.45 -3.50 -3.22 -3.50 - 6.4e-4(prod) -3.09 -3.50 -2.16 -2.95 - ``` - The **softplus is a blunt clip**: it hard-truncates the *pre-activation* floor to - −3.5 wherever that raw floor overshoots the hinge. At weak fusion (0…1.6e-4) both - α regimes drive Delta's pre-activation floor to −11…−16 (far past −3.5), so ON - pins every one of those cells to exactly −3.5. The floor is inside `predict_score` - (which the loss uses), so turning it on also perturbs the fit everywhere — fitted - α lands consistently higher ON than OFF in both regimes. **Free-α** does the - opposite thing: it *deepens the pre-activation floor selectively* by letting - Delta's own α decouple (Delta α at prod: shared 13.47 → free **5.23** OFF), which - by itself takes the OFF prod floor −2.16 → −3.09. **The prod cell is the whole - point:** softplus-alone gets −2.16 → −2.95 (α-driven, raw floor only −2.96 so the - hinge barely bites); free-α-alone gets −2.16 → −3.09; **both together** are the - *only* arm that reaches **−3.50** at prod, because free-α pushes the - pre-activation floor (−4.80) past the hinge, which then clamps it. The two levers - attack opposite regimes and compose at prod — neither alone floors the prod tail; - together they do. Ships **default-off** (`output_floor=None`); the shared-α OFF - control reproduced #274's separately-fit column to the last decimal, confirming - determinism and a fair baseline. -- **`recompute_scale=False`** (the fixed-scale objective normalizer) converges. -- **Inner optimizer maxiter is the dominant convergence lever; `recompute_scale=False` - + inner maxiter=100 is the only 100%-converging cell** (measured, 2026-07-10, - `cache=maxiter-scan`, 48 fits = ge/cal `maxiter` {1,10,100} × `recompute_scale` - {False,True} × fusionreg [0, 4e-5, 1.6e-4, 6.4e-4] × 2 reps, else #277's - softplus-off block verbatim: cold-start, l2reg=0, Sigmoid, α_init=6, - clip[-10,10], outer maxiter=50/tol=1e-5). Convergence rises monotonically with - inner maxiter and is 0 at maxiter=1: **maxiter=1 → 0/16, maxiter=10 → 3/16, - maxiter=100 → 12/16.** At maxiter=1 the optimizer barely moves (α≈init 6.1, - Σβ²≈0.4 — essentially unfit); at maxiter=100 the converged fits land in the - large-Σβ²/low-α clip basin the #263 finding calls healthy (α~4-6, Σβ²~60k–124k - harmless under the φ-clip). **`recompute_scale=False` dominates `True` at - maxiter=100 on both axes: 8/8 vs 4/8 converged AND 3.6× faster (mean fit_time - 788s vs 2808s).** So the fixed-scale objective isn't just convergent — it's the - cheap, reliably-convergent one, and the inner cap must be ≥~100 to actually - reach tol. NOTE: run serial or a *small* pool for heavy inner-maxiter grids — - the m100 level deadlocked under the 16-worker spawn pool (collapsed to one - process); `n_processes=4` ran clean (4 workers ~100% CPU each). -- **`fit_models` parallelism — settled** (`diagnostics/parallelism_probe.py`): - `n_processes=2` (the real spawn path) ran the full data-size × l2reg staircase - with **zero hangs**, at `l2reg=0` AND `l2reg=3e-4`. The l2reg-deadlock theory is - **refuted**. The original local hangs were the **execution context** — marimo - cells / `/tmp` scripts spawn workers without a clean `if __name__ == - "__main__"` guard, so each child re-runs module-level code on import. A - properly-guarded script runs the identical spawn path cleanly. -- **Dataset duplication is real but harmless** (measured): each spawn worker - carries its own dataset copy, so `n_processes=2` peaks **~0.9–1.9 GB above** - the `n_processes=1` in-process baseline, and the gap **grows with data size** - (+0.9 GB tiny → +1.9 GB full). But peak RSS at the full dataset is ~4.9 GB - against ~13 GB free — duplication never exhausts RAM. Memory was NOT the cause - of the hangs. The harness fits in parallel by default; size `--n-processes` to - the host's free RAM (per-worker dataset copy), not just its core count. - -## Run log - -(Append one entry per harness run: date, cache name, config swept, what the -fit collection showed — basin diagnostics and replicate correlation computed -downstream from a ModelCollection — the conclusion, the next step.) - -- 2026-07-15 | cache=beta0-ridge-l2-scan | **(#284)** sweep: `beta0_ridge` [1e-4, 1e-3, 1e-2] × `l2reg` [1e-8, 1e-7, 1e-6] × fusionreg [0.0, 4e-5, 1.6e-4, 6.4e-4], 2 reps = **72 fits**. Else `softplus-floor-off.yaml`'s fixed block VERBATIM (cold-start warmstart=false, recompute_scale=false, output_floor=null, share_alpha=true, Sigmoid, alpha_init=6, `beta_clip_range=[-10,10]` — the settled #263 bound, INHERITED not swept — `beta0_init` pins all three β0 to 0.0, scale_fusion_by_n=false, loss δ=1, outer maxiter=50/tol=1e-5, inner ge/cal maxiter=10/maxls=10/tol=1e-4), so the grid is a one-knob-at-a-time delta from its baseline (asserted in `diagnostics/test_beta0_ridge_l2_grid.py`, not eyeballed). Config: `grids/beta0-ridge-l2-scan.yaml`. Ran REMOTE on orca04 (64-core, 1.5 TB RAM, scouted idle at 0.3% CPU; GitHub reachable this time, so a real git clone at branch HEAD — no rsync fallback needed) at **n_processes=16**, wall **3320.1s (55m), 72 fit / 0 failed** (39,950s CPU-time in 3,320s wall = 12× speedup; the m100 16-worker deadlock does NOT reproduce at this grid's light inner maxiter=10 block). Downstream from `diagnostics/beta0_ridge_l2_report.py --cache beta0-ridge-l2-scan --baseline-cache 277-softplus-floor-off`. - **What `beta0_ridge` actually penalizes — NOT the intercepts.** `_beta_ridge_penalty` (jaxmodels.py:533-540) penalizes each non-reference condition's β0 *difference from the reference*, `beta0_ridge · Σ_{d≠ref}(β0_d − β0_ref)²`; the reference's own β0 is never touched. It is an **L2 shift-shrinkage on the intercepts**, structurally parallel to `fusionreg`'s **L1 shift penalty on the β** (`fusionreg · Σ|β_d − β_ref|`, jaxmodels.py:568-577) — hence crossing the two, and hence reading the effect WITHIN each fusionreg level rather than marginalized over it. Note the inherited `beta0_init` starts all three β0 identical at 0.0, so the penalty is **exactly 0 at initialization** and only bites as the β0s separate. - **Baseline computed here, not quoted — it existed nowhere.** The 277-softplus-floor-off entry says "pkl carries no converged/final_obj_err column", and the 0/16 it cites is #273's *different* free-alpha grid. Computed from the 8-fit pickle with the same report: **0/8 converged, median final_obj_err 2.29e-4, α 16.1, Σβ² 4191**. Baseline replicate-shift r — shift_Delta 0.594/0.571/0.604/0.498, shift_Omicron_BA2 0.382/0.482/0.619/0.767 across fusionreg 0/4e-5/1.6e-4/6.4e-4. - Convergence + basin per (beta0_ridge, l2reg), **n=8 per cell** (4 fusionreg × 2 reps); the baseline row is n=8 for its whole cache: - | beta0_ridge | l2reg | converged | median obj_err | α | Σβ² | - |---|---|---|---|---|---| - | *(baseline 0)* | *0* | *0/8* | *2.29e-4* | *16.1* | *4191* | - | 1e-4 | 1e-8 | 0/8 | 2.20e-4 | 16.1 | 4190 | - | 1e-4 | 1e-7 | 0/8 | 2.15e-4 | 16.1 | 4180 | - | 1e-4 | 1e-6 | 0/8 | 1.74e-4 | 16.3 | 4020 | - | 1e-3 | 1e-8 | 0/8 | 2.19e-4 | 16.5 | 4220 | - | 1e-3 | 1e-7 | 0/8 | 2.13e-4 | 16.5 | 4210 | - | 1e-3 | 1e-6 | 0/8 | 1.76e-4 | 16.7 | 4060 | - | 1e-2 | 1e-8 | 0/8 | 2.16e-4 | 17.5 | 4220 | - | 1e-2 | 1e-7 | 0/8 | 2.23e-4 | 17.6 | 4200 | - | 1e-2 | 1e-6 | 0/8 | 1.96e-4 | 17.8 | 4060 | - Replicate-shift Pearson r (rep_1 vs rep_2, `shift_*` rows only — there is no shift_Omicron_BA1, the reference carries no shift parameters), per cell across fusionreg. **Computed by slicing `mut_param_dataset_correlation` with `query=` per cell:** the bare call groups by (dataset_name, fusionreg) and mean-collapses everything else (model_collection.py:1444/787), which on this 72-fit frame would average the 9 (beta0_ridge, l2reg) fits per group — averaging away the two axes under test. Every cell reproduces the baseline column to ±0.02: - | beta0_ridge | mut_param | fr=0 | fr=4e-5 | fr=1.6e-4 | fr=6.4e-4 | - |---|---|---|---|---|---| - | 1e-4 | shift_Delta | 0.595 | 0.571 | 0.604 | 0.498 | - | 1e-3 | shift_Delta | 0.600 | 0.577 | 0.606 | 0.493 | - | 1e-2 | shift_Delta | 0.617 | 0.589 | 0.606 | 0.486 | - | 1e-4 | shift_Omicron_BA2 | 0.382 | 0.482 | 0.619 | 0.767 | - | 1e-3 | shift_Omicron_BA2 | 0.383 | 0.484 | 0.619 | 0.768 | - | 1e-2 | shift_Omicron_BA2 | 0.393 | 0.487 | 0.619 | 0.769 | - (l2reg collapsed above — it moves nothing; the three l2 levels agree to ±0.003 within every beta0_ridge row.) - **Adjudicated against #284's pre-registered thresholds** — an effect required ANY of: ≥1/8 converged where the baseline is 0/8; median final_obj_err ≥2× baseline; median shift-r moving ≥0.05; or α/Σβ² ≥2× at matched fusionreg. **Result: FLAT on all four, in all 9 cells.** Convergence 0/8 everywhere (primary DEGENERATE — 0.0 vs 0.0 discriminates nothing, so the pre-registered obj_err fallback adjudicates); obj_err ratios 0.76–0.97× (max deviation 24%, threshold 100%); α 1.00–1.11×, Σβ² 0.96–1.01×; shift-r deltas −0.006…+0.007 (threshold 0.05, so the largest is 7× under). - Conclusion: **NULL — neither a sub-knee `l2reg` (1e-8…1e-6) nor a nonzero `beta0_ridge` (1e-4…1e-2) buys convergence or reproducibility on top of the settled clip bound.** This CONFIRMS and sharpens the standing L2 finding: the ~3e-4 knee is real, and three-to-four orders below it the penalty is simply inert against an active `beta_clip_range=[-10,10]`, which is already the binding constraint on β. The axes are not *quite* dead, though, and the directions are informative: α rises **monotonically** with beta0_ridge (16.1 → 16.5 → 17.5 at fixed l2reg) and Σβ² dips ~4% at l2reg=1e-6, so both penalties do touch the optimum — they are 1-2 orders of magnitude too weak to move it. **The honest limit of this result:** convergence was 0/8 at the baseline BEFORE the scan began, so this grid could never have detected a convergence *gain* — it can only say the regularizers did not rescue a fit that the inner cap had already pinned. Per the 2026-07-10 maxiter-scan, the inner cap is the dominant convergence lever (0/16 → 3/16 → 12/16 over inner maxiter 1/10/100), and this grid inherits the middle-of-the-road maxiter=10. So the defensible claim is *"regularization cannot rescue convergence at an inadequate inner cap"* — NOT *"these regularizers are inert"* in general. - Next: **#285** (stub, blocked-by #284) re-runs these exact axes at inner maxiter=100 — the only 100%-converging cell known — which is what separates those two claims. Budget ~11h naive (m100 was 2h31m for 16 fits) and **do not reuse n_processes=16** (m100 deadlocked under it; use 4). A maxiter=100 `(0,0)` baseline does not exist yet and must be fit. Standing findings: the L2-knee finding is sharpened below; no finding is overturned. - -- 2026-07-10 | cache=maxiter-scan | sweep: ge/cal_kwargs `maxiter` {1, 10, 100} (COUPLED — ge and cal move together) × `recompute_scale` {False, True} × fusionreg [0.0, 4e-5, 1.6e-4, 6.4e-4], 2 reps = 48 fits. Else #277 softplus-floor-off.yaml fixed block VERBATIM (cold-start warmstart=False, output_floor=null, share_alpha=true, Sigmoid, l2reg=0, alpha_init=6, beta_clip_range=[-10,10], loss δ=1, outer maxiter=50/tol=1e-5, inner maxls=10/tol=1e-4). Configs: grids/maxiter-scan-m{1,10,100}.yaml (three configs because the harness Cartesian-products every sweep axis — ge_kwargs and cal_kwargs listed together would 3×3-cross to 9, not the 3 coupled cells; one config per level then `pd.concat` the three fit_collection.pkl). Ran REMOTE on orca04 (64-core, scouted idle) — source tree rsync'd to HEAD 9121748 (GitHub unreachable from orca04 via the 1Password-gated forwarded agent, so no remote git fetch; the one required data CSV, gitignored, rsync'd explicitly). m1/m10 fit at n_processes=16 (wall 154s / 639s, 16 fit 0 failed each); **m100 deadlocked under the 16-worker spawn pool** (collapsed to a single limping process, ~46 min no output — the documented spawn-under-JAX race, widened by the heavier 100-iter compile) → rerun at **n_processes=4**, which ran clean (4 workers ~100% CPU each), wall 9049s (2h31m), 16 fit 0 failed. Downstream basin/convergence computed inline from the combined frame. - Convergence rate + basin (α, Σβ² of ref-condition β, mean fit_time) per (ge/cal maxiter × recompute_scale): - | maxiter | recompute | converged | α | Σβ² | fit_time | - |---|---|---|---|---|---| - | 1 | False | 0/8 | 6.13 | 0.4 | 132s | - | 1 | True | 0/8 | 6.13 | 0.4 | 137s | - | 10 | False | 0/8 | 15.5 | 4,096 | 550s | - | 10 | True | 3/8 | 11.9 | 7,641 | 558s | - | 100 | False | **8/8** | 5.59 | 60,959 | 788s | - | 100 | True | 4/8 | 4.31 | 124,271 | 2,808s | - Convergence rises monotonically with inner maxiter (0/16 → 3/16 → 12/16 over 1/10/100). At maxiter=1 the optimizer barely moves off init (α≈6.1=alpha_init, Σβ²≈0.4 — essentially unfit, so its "0/8 not converged" is trivial, not pathological). At maxiter=100 the converged fits sit in the large-Σβ²/low-α clip basin the #263 finding calls HEALTHY (α~4-6, Σβ² 60k–124k held harmless by the [-10,10] φ-clip — "large Σβ²" only signals explosion when the fit ALSO fails to converge). - Conclusion: the inner optimizer cap is the DOMINANT convergence lever here — it must be ≥~100 to reach the outer tol=1e-5 at all; 1 and 10 are simply too few inner steps. And `recompute_scale=False` (fixed-scale objective) dominates `True` at maxiter=100 on BOTH axes: 8/8 vs 4/8 converged AND 3.6× faster (788s vs 2808s). This sharpens the terse standing "recompute_scale=False converges" into a quantified 2-D result and confirms the fixed-scale objective is the cheap, reliably-convergent one. Standing findings updated. Next: with 100% convergence now reachable (recompute=False, inner maxiter=100), re-run the reproducibility criterion (replicate-shift Pearson r via ModelCollection.mut_param_dataset_correlation) on this converged cell to check the α/β-basin choice is stable replicate-to-replicate now that fits actually reach tol; also record the n_processes=4 (not 16) requirement for heavy-maxiter grids in the harness docs. - -- 2026-07-07 | cache=273-free-alpha | sweep: share_alpha [True, False] × fusionreg [0.0, 4e-5, 1.6e-4, 6.4e-4], 2 reps (16 fits), PR 270 fitting block verbatim (warmstart=False, recompute_scale=False, l2reg=0, Sigmoid, alpha_init=6, beta_clip [-10,10], inner ge/cal_kwargs maxiter=10/maxls=10/tol=1e-4) + 3 deltas from PR 270 (outer tol 1e-6→1e-5, 4-point fusionreg, share_alpha swept). Independent fitting (#273). Wall 718s at n_processes=13 (local, Apple M4 Max). Numbers computed inline from a ModelCollection over the frame (dashboard-only exploration; no committed analysis code). Delta floor = −α_Delta·sigmoid(φ_wt,Delta), the low asymptote of α·(g(φ)−g(φ_wt)) as φ→−∞ (analytic and model-GE-at-φ=−1e4 evaluations agreed to 0.0e0). - Delta α + floor (mean over reps), shared vs free, across fusionreg 0 / 4e-5 / 1.6e-4 / 6.4e-4: SHARED α → α 16.8/17.0/14.8/13.5, floor −5.08/−4.69/−3.22/−2.16. FREE α → Delta's own α 16.8/18.1/12.9/5.2 (decoupled from BA1 17.8 / BA2 14.4 at 6.4e-4), floor −5.24/−4.82/−3.45/−3.09. At the prod-strength fusionreg=6.4e-4 freeing α drops Delta's floor from −2.16 → −3.09 (gap −0.93), most of the way to the true −3.5 that the shared α (pinned by BA1/BA2) cannot reach. Effect grows with fusion (floor gap free−shared −0.15 at fr=0 → −0.93 at fr=6.4e-4); BA1/BA2 floors barely move. Counter-intuitively the deeper floor comes with a *smaller* Delta α — freeing α lets the Delta latent push φ_wt more positive so −α·sigmoid(φ_wt) deepens. - Convergence: 0/16 converged (all hit maxiter=50), same as PR 270's 0/18 at tol=1e-6 — the looser tol=1e-5 did NOT buy convergence. But obj_err ≤3e-4 (range 5.9e-5–3.0e-4), so the α/floor numbers are trustworthy; the binary flag is not (same maxiter-truncation pattern as the l2-fusion / smoke runs). - Conclusion: CONFIRMS the PR 270 hypothesis — the single shared α is the reason Delta under-fits its low tail. A per-condition α lets Delta decouple and pull its floor toward the true −3.5, with the effect concentrated at strong fusion (exactly where PR 270 saw the −2.5 shortfall). Standing findings updated (new free-α floor finding). Next: confirm on the full scv2-spike pipeline (share_alpha=false prod run) that the deeper floor holds against held-out data and does not reintroduce the per-condition sigmoid degeneracy the shared-α default guards against; optionally pair with a small l2reg>0 (the β-explosion tamer) since this run sits in the l2reg=0 regime. - -- 2026-07-06 | cache=beta-control-clip + beta-control-l2 | Phase 1 (#263), EPIC #262. sweep: arm {clip[-10,10], l2reg=1e-4} × fusionreg [0.0, 4e-5, 6.4e-4], 2 reps (12 fits). Base regime: cold-start (warmstart=false), recompute_scale=false, share_alpha=true, Sigmoid, maxiter=100 (ceiling), tol=1e-6. Independent fitting. Local (M4 Max). Downstream from diagnostics/beta_control_report.py. - Basin diagnostics (Σβ², α per cell): clip arm → Σβ² 110,597–180,962, α 2.96–3.39, **all 6 converged=True** (final_obj_err 6e-8–9e-7). l2 arm → Σβ² 221–300, α 52.5–58.8, all 6 converged=False (final_obj_err 8e-5 at fusionreg=0 rising to ~2.8e-4 at 6.4e-4). This is the α/β see-saw in the open: clip caps φ so α stays low (~3) and the fit converges cleanly despite a large Σβ²; l2 collapses β and drives α to ~55, and it never reaches tol=1e-6. Large clip-Σβ² is NOT an explosion — α·φ is bounded by the clip and the objective converges (contrast #256's l2reg=0 Σβ²~16–19k which did NOT converge to tol). - maxiter each needs (outer iters to cross tol=1e-6 / 1e-4): clip → 1e-6 at {fusionreg=0: 60,26 · 4e-5: 30,39 · 6.4e-4: 55,13}, 1e-4 at 6–8 iters everywhere (4/6 of the 1e-6 crossings land in the epic's 20–50 band, the other two at 13 and 60). l2 → 1e-6 NEVER crossed (sentinel 101 on all 6); 1e-4 crossed only at fusionreg=0 (iter 88/89) and never at fusionreg≥4e-5 (101). The clip arm is the ONLY arm that reaches tol=1e-6, and it does so cheaply. - Replicate-shift Pearson r (shift_* rows, per arm, across fusionreg 0/4e-5/6.4e-4): clip → shift_Delta 0.49/0.46/0.59, shift_Omicron_BA2 0.38/0.56/0.81. l2 → shift_Delta 0.49/0.44/0.49, shift_Omicron_BA2 0.33/0.36/0.66. Strong fusion RESCUES the data-poor shift_Omicron_BA2 under both arms, but clip lifts it higher (0.81 vs 0.66) and also lifts shift_Delta at prod-max fusion (0.59 vs 0.49) — clip dominates the reproducibility criterion at every fusion strength. - Conclusion (tri-criteria — speed × basin health × rising reproducibility): **clip[-10,10] wins, decisively and on all three axes.** Speed: clip is the only arm that converges to tol=1e-6 (l2 never does), landing 4/6 fits in the 20–50 band. Basin health: clip keeps α~3 (healthy) while l2's α blows up to ~55 (the see-saw's collapse end) — the low Σβ² of l2 is the symptom, not the cure. Reproducibility: clip's strong-fusion shift_Omicron_BA2 r reaches 0.81 vs l2's 0.66. This overturns the gauge argument's tie-break expectation (it predicted clip would win, but on a *healthier Σβ²* — instead clip wins with a *large* Σβ² held harmless by the hard φ-cap). **The epic's downstream phases (#264 recompute_scale, #265 warmstart) inherit `beta_clip_range=[-10,10]` as the β-bound.** Next: Phase 2 / #264 — vary recompute_scale at the clip bound. - -- 2026-07-07 | cache=277-softplus-floor (ON) + 277-softplus-floor-off (OFF control) | MATCHED in-harness A/B: same 8-cell grid fit twice — sweep fusionreg [0.0, 4e-5, 1.6e-4, 6.4e-4] × 2 reps, share_alpha=true, warmstart=false, recompute_scale=false, Sigmoid, alpha_init=6.0, beta_clip_range=[-10,10], maxiter=50, tol=1e-5 (#274's free-alpha.yaml fixed block VERBATIM) — once with output_floor=-3.5 (softplus, hinge=0.1), once with output_floor=null, so the floor is the ONLY difference. 16 fits, 0 failed, ~307s + ~326s at n_processes=8 (local, Apple M4 Max). The OFF control reproduces #274's separately-fit shared-α column to the last decimal (floor −5.08/−4.69/−3.22/−2.16, α 16.81/17.01/14.76/13.47), validating both the fit determinism and #274's cross-PR baseline. - Delta floor + fitted shared α (mean over 2 reps); floor computed two ways (analytic `t(−α·sigmoid(φ_wt))` and model composition at φ=−1e4) — agreed to 0.0e0: - | fusionreg | OFF floor | ON floor | Δfloor | OFF α | ON α | Δα | ON pre-act | clamped? | - |---|---|---|---|---|---|---|---|---| - | 0.0 | −5.08 | −3.50 | +1.58 | 16.81 | 20.61 | +3.80 | −15.74 | yes | - | 4.0e-5 | −4.69 | −3.50 | +1.19 | 17.01 | 21.34 | +4.32 | −15.75 | yes | - | 1.6e-4 | −3.22 | −3.50 | −0.28 | 14.76 | 19.24 | +4.48 | −16.43 | yes | - | 6.4e-4 | −2.16 | −2.94 | −0.79 | 13.47 | 17.36 | +3.90 | −2.96 | no | - Conclusion: the softplus perturbs the FIT at every fusion strength — fitted shared α is a consistent +3.8…+4.5 higher ON than OFF (+29% at prod), and because the floor is inside predict_score (which the loss uses) this is a real optimization effect, not fit-to-fit noise. Two regimes: at weak fusion (0…1.6e-4) the shared-α sigmoid drives Delta's raw pre-activation floor to −15…−17 (far past the biological −3.5) and the softplus HARD-TRUNCATES it to exactly −3.5 (the assay detection floor); at prod fusion (6.4e-4) the raw floor only reaches −2.96 so the hinge itself clips almost nothing (−2.96→−2.94), but the OFF→ON floor still deepens −2.16→−2.94 driven by the genuinely higher fitted α the floor induced, not by the hinge biting. vs #274 free-α (which DEEPENS the under-shooting prod floor −2.16→−3.09 by selectively decoupling Delta's α): the softplus is a BLUNT clip that floors all three conditions and mainly bites the over-shooting weak-fusion tail — it does NOT pull the prod floor to −3.5 on its own. Complementary, opposite regimes. Convergence: pkl carries no converged/final_obj_err column (computed downstream); #274 saw 0/16 at these tols; all 16 fits completed, floor introduced no failures. Next: for a floored prod recommendation, pair the softplus with free-α or l2reg>0 (softplus alone does not fix the prod under-fit); counts-path floor deferred to a separate stub. [SUPERSEDED 2026-07-08 by the free-α arm below — the "pair with free-α" recommendation is now measured directly, not inferred cross-PR.] - -- 2026-07-08 | cache=277-softplus-floor-freealpha (ON) + 277-softplus-floor-off-freealpha (OFF control) | The MISSING free-α arm of the 2×2. Same 8-cell grid as the 2026-07-07 shared-α pair but share_alpha=FALSE (the #274 free-α regime, per-condition α) — fit twice, output_floor=-3.5 vs null. 16 fits, 0 failed; wall 1309s (ON) + 1912s (OFF) at n_processes=1. **Ran serial (n_processes=1) deliberately:** the free-α grids deadlocked under the spawn pool (`n_processes=8`) — all 8 workers went to 0% CPU with a shared multiprocessing pipe_handle after a worker re-spawned, a spawn-under-JAX race the heavier free-α compilation widened; serial fits the identical models with no pool. All 16 verified genuinely free-α (per-condition α dict on the fitted model, not shared scalar). Delta floor computed two ways (analytic + model φ=−1e4) — agreed to 0.0e0 on all evals. - Full 2×2 Delta floor + Delta's own α (mean over 2 reps): - | fusionreg | free OFF floor | free ON floor | shared OFF floor | shared ON floor | free OFF α_Δ | free ON α_Δ | shared α (OFF/ON) | - |---|---|---|---|---|---|---|---| - | 0.0 | −5.24 | −3.50 | −5.09 | −3.50 | 17.37 | 19.16 | 16.81/20.61 | - | 4.0e-5 | −4.82 | −3.50 | −4.69 | −3.50 | 18.07 | 21.20 | 17.01/21.34 | - | 1.6e-4 | −3.45 | −3.50 | −3.22 | −3.50 | 12.85 | 17.54 | 14.76/19.24 | - | 6.4e-4 | −3.09 | −3.50 | −2.16 | −2.95 | 5.23 | 11.00 | 13.47/17.36 | - Conclusion: the two levers are complementary and compose AT PROD. Free-α *deepens the pre-activation floor selectively* — Delta's α at prod drops shared-13.47 → free-5.23 (OFF), taking the OFF prod floor −2.16 → −3.09 with no clip. The softplus is a *blunt clip* to −3.5 wherever the raw pre-activation floor overshoots the hinge. At prod: softplus-alone −2.16→−2.95 (raw floor only −2.96, hinge barely bites, α-driven); free-α-alone −2.16→−3.09; **BOTH together −2.16→−3.50** — the only arm reaching the biological floor at prod, because free-α's deeper −4.80 pre-activation floor now exceeds the hinge and gets clamped. At weak fusion (0…1.6e-4) both α regimes overshoot to −11…−16 pre-activation, so ON pins every cell to exactly −3.5 regardless of share_alpha. This directly measures (not infers) the 2026-07-07 "pair softplus with free-α for a floored prod recommendation" next-step: confirmed — neither lever alone floors the prod tail, together they do. Standing finding rewritten to the 2×2. Next: none for this experiment; the softplus × free-α interaction is settled. Counts-path floor still deferred to stub #276. - -- 2026-06-30 | cache=l2-fusion | sweep: l2reg [0.0, 1e-4, 3e-4] × fusionreg [0.0, 4e-5, 6.4e-4], 2 reps (18 fits), warmstart=True, recompute_scale=False, share_alpha=True, Sigmoid, maxiter=25. Independent fitting (#256). Wall 1138s at n_processes=4 (local, Apple M4 Max). Downstream numbers from diagnostics/l2_fusion_report.py. - Basin diagnostics (Σβ², α per cell): at l2reg=0.0 → Σβ² 16,181–19,222, α 5.9–8.2 (β EXPLODED at EVERY fusionreg, incl. 6.4e-4). l2reg=1e-4 → Σβ² 554–841, α 19.5–24.6 (β tamed ~25× into the healthy 350–1400 band, holds across all fusion). l2reg=3e-4 → Σβ² 188–298, α 30.5–40.4 (β tamed further but α climbing toward the see-saw collapse end). All 18 converged=False, but final_obj_err is tiny (4e-6 at l2reg=0 to ~7e-4 at l2reg>0) — maxiter=25 truncation, not a pathology; β-magnitude and r are trustworthy, the binary flag is not. - Replicate-shift Pearson r (per l2reg slice, across fusionreg 0/4e-5/6.4e-4): l2reg=0 → shift_Delta 0.45/0.50/0.09, shift_Omicron_BA2 0.26/0.43/0.20 (strong fusion COLLAPSES shift_Delta at l2reg=0 — the unpenalized-β regime). l2reg=1e-4 → shift_Delta 0.47/0.45/0.33, shift_Omicron_BA2 0.37/0.43/0.80 (strong fusion now LIFTS shift_Omicron_BA2 to 0.80). l2reg=3e-4 → shift_Delta 0.40/0.37/0.41, shift_Omicron_BA2 0.43/0.49/0.76. - Conclusion: a small l2reg>0 SOLVES the β-explosion across the entire prod fusion axis; l2reg=1e-4 is the working weight (β healthy, α not yet over-driven). The data-poor-condition distortion under strong fusion is itself a symptom of the unpenalized explosion — once β is penalized, strong fusion rescues rather than wrecks replicate-r (shift_Omicron_BA2 0.20→0.80 at fusionreg=6.4e-4). Standing findings updated (L2-knee finding extended to 2-D). Next: maxiter sweep at l2reg=1e-4 to convert the tiny-but-nonzero final_obj_err into converged=True; optionally a continuation-path comparison along fusionreg at l2reg=1e-4 (separate experiment — independent fitting already suffices to tame β). - -- 2026-06-26 | cache=smoke | sweep: l2reg [0.0, 3e-4], warmstart=True, recompute_scale=False, fusionreg=0.0, Sigmoid, maxiter=25. - df_fits (4 fits, ~185s/fit, wall 740s): at l2reg=0.0 -> alpha 7.4-8.2, Sigma-beta^2 36,785-40,371 (beta EXPLODED); at l2reg=3e-4 -> alpha 36-40, Sigma-beta^2 438-626 (beta collapsed, alpha exploded). All 4 converged=False (final_obj_err ~4-6e-6 at l2reg=0, ~3e-4 at l2reg=3e-4). - df_corr (replicate-shift Pearson r): shift_Delta r=0.47 (l2reg=0) / 0.33 (3e-4); shift_Omicron_BA2 r=0.37 (l2reg=0) / 0.53 (3e-4). - Conclusion: harness reproduces the double-ended L2 degeneracy + sub-convergence locally in ~12 min. Even warmstart=True + recompute_scale=False shows beta-explosion at l2reg=0. Next: sweep warmstart True/False at fixed l2reg to isolate the see-saw; widen l2reg ladder around the 3e-4 knee. - -- 2026-06-27 | diagnostics/parallelism_probe.py --baseline | two-axis staircase (data-size tiny/small/medium/full × l2reg [0.0, 3e-4]), n_processes=2 spawn vs n_processes=1 in-process, cheap iters (block maxiter=6 / inner 8), per-step peak RSS sampled. - np=2 (spawn): 8/8 PASS, ZERO hangs. Wall 9-36s/step. Peak RSS climbs with data size: l2reg=0 -> 3690/3798/3956/4907 MB (tiny/small/medium/full); l2reg=3e-4 -> 4012/4225/4510/4687 MB (l2reg>0 costs ~+300-500 MB). - np=1 (in-process baseline, l2reg=0): tiny/small/medium/full -> 2754/2674/2662/3048 MB. (Baseline truncated after l2reg=0 by an operator kill; the l2reg=0 column is sufficient for the spawn-vs-in-process comparison.) - Spawn overhead = np2 - np1 at l2reg=0: +936/+1124/+1294/+1859 MB — grows with data size, the signature of per-worker dataset duplication. But 4.9 GB peak << ~13 GB free. - Conclusion: fit_models parallelism does NOT hang (l2reg-deadlock theory refuted); dataset duplication is real (~1 GB/worker) but never exhausts RAM. The original local hangs were the marimo/`/tmp` no-__main__-guard execution context. Next: none for parallelism — settled. - -- 2026-06-29 | cache=smoke | sweep: l2reg [0.0, 3e-4], warmstart=True, recompute_scale=False, fusionreg=0.0, Sigmoid, maxiter=25. Harness now PARALLEL (fit_models, n_processes=4 default). - Output: results/smoke/fit_collection.pkl — a TRUE fit collection (raw stack_fit_models frame, ModelCollection-loadable, dashboard-discoverable). 4/4 fit, wall 246s (vs ~740s sequential on 2026-06-26 — the 4 fits ran concurrently). No metrics stored in the pickle. - Downstream (computed on the fly from ModelCollection over the frame, NOT stored): replicate-shift Pearson r reproduces the 2026-06-26 numbers — shift_Delta r=0.47 (l2reg=0) / 0.33 (3e-4); shift_Omicron_BA2 r=0.37 (l2reg=0) / 0.53 (3e-4). - Conclusion: parallel harness produces an identical-schema fit collection the dashboard reads directly; correlation/basin analysis is now per-experiment downstream work, not baked into the harness output. Next: same science as before — sweep warmstart True/False at fixed l2reg to isolate the see-saw; widen the l2reg ladder around the 3e-4 knee. diff --git a/experiments/convergence-lab/diagnostics/beta0_ridge_l2_report.py b/experiments/convergence-lab/diagnostics/beta0_ridge_l2_report.py deleted file mode 100644 index 2cbff0c3..00000000 --- a/experiments/convergence-lab/diagnostics/beta0_ridge_l2_report.py +++ /dev/null @@ -1,308 +0,0 @@ -r"""Downstream report for the beta0_ridge × sub-knee l2reg scan (#284). - -Loads a convergence-lab ``fit_collection.pkl`` into a ``ModelCollection`` and -prints, per ``(beta0_ridge, l2reg)`` cell: the convergence rate, median -``final_obj_err``, basin diagnostics (Σβ², α), and replicate-shift Pearson r. -With ``--baseline-cache`` it prints the same tables for the ``(0,0)`` baseline -so the two are read side by side — the baseline's convergence rate has never -been computed and cannot be quoted from the README. - -Reuses ``l2_fusion_report.basin_row`` UNCHANGED (its contract is pinned by -``test_l2_fusion_report.py``) and adds the ``beta0_ridge`` column this grid -needs, mirroring how ``beta_control_report`` bolts on its ``arm`` column. - -Adjudication (pre-registered in #284) — an axis shows an effect if any of: - -* convergence: ≥1/8 converged in a cell where the baseline is 0/8; -* fallback: median ``final_obj_err`` differing ≥2× from baseline; -* repro: median replicate-r shifting ≥0.05 vs baseline; -* basin: Σβ² or α differing ≥2× at matched ``fusionreg``. - -The primary is expected to be DEGENERATE (0/72 vs 0/8) — when it is, the -``final_obj_err`` fallback IS the adjudicator, not a footnote. - -Run:: - - pixi run python experiments/convergence-lab/diagnostics/beta0_ridge_l2_report.py \\ - --cache beta0-ridge-l2-scan --baseline-cache 277-softplus-floor-off -""" - -from __future__ import annotations - -import argparse -import gc -import os -import pickle -import sys -import warnings -from pathlib import Path - -os.environ.setdefault("XLA_FLAGS", "--xla_cpu_multi_thread_eigen=false") -os.environ.setdefault("OMP_NUM_THREADS", "1") -os.environ.setdefault("JAX_PLATFORM_NAME", "cpu") - -warnings.filterwarnings("ignore") - -import pandas as pd # noqa: E402 - -# Reuse the harness's RESULTS_DIR and the l2-fusion report's basin_row. -sys.path.insert(0, str(Path(os.path.abspath(__file__)).parent.parent)) -import harness # noqa: E402 - -sys.path.insert(0, str(Path(os.path.abspath(__file__)).parent)) -import l2_fusion_report as l2rpt # noqa: E402 - -from multidms.model_collection import ModelCollection # noqa: E402 - - -def basin_row_with_ridge(fit) -> dict: - """``l2_fusion_report.basin_row`` plus the ``beta0_ridge`` cell value. - - ``basin_row`` predates this grid's ``beta0_ridge`` axis and does not carry - it. Rather than edit it (its contract is pinned by - ``test_l2_fusion_report.py``), bolt the column on here — the same move - ``beta_control_report.basin_table`` makes for its ``arm`` column. - - Args: - fit: A row (``pd.Series``) of the fit-collection frame. - - Returns: - ``basin_row``'s dict plus ``beta0_ridge``. - """ - row = l2rpt.basin_row(fit) - row["beta0_ridge"] = fit.beta0_ridge - return row - - -def basin_table(frame: pd.DataFrame) -> pd.DataFrame: - """Per-fit basin diagnostics, sorted with ``beta0_ridge`` outermost. - - ``l2_fusion_report.basin_table`` sorts by ``(l2reg, fusionreg, - dataset_name)`` and knows nothing of ``beta0_ridge``, so this grid needs its - own. - - Args: - frame: The raw fit-collection DataFrame. - - Returns: - One row per fit, sorted by - ``(beta0_ridge, l2reg, fusionreg, dataset_name)``. - """ - rows = [basin_row_with_ridge(frame.iloc[i]) for i in range(len(frame))] - return ( - pd.DataFrame(rows) - .sort_values(["beta0_ridge", "l2reg", "fusionreg", "dataset_name"]) - .reset_index(drop=True) - ) - - -def convergence_table(basin: pd.DataFrame) -> pd.DataFrame: - """Convergence rate + median diagnostics per ``(beta0_ridge, l2reg)`` cell. - - The denominator is explicit: each cell holds ``n`` fits (4 ``fusionreg`` × 2 - replicates = 8 for the full grid). ``median_obj_err`` is the pre-registered - fallback adjudicator for when ``conv_rate`` is degenerate (all-zero) — the - lab has repeatedly found the binary flag untrustworthy while - ``final_obj_err`` stays tiny and informative (#273: 0/16 converged, obj_err - ≤3e-4, parameters trustworthy). - - Args: - basin: Output of :func:`basin_table`. - - Returns: - One row per ``(beta0_ridge, l2reg)``: ``n``, ``n_converged``, - ``conv_rate``, ``median_obj_err``, ``median_alpha``, - ``median_sum_beta_sq``. - """ - out = ( - basin.groupby(["beta0_ridge", "l2reg"], dropna=False) - .agg( - n=("converged", "size"), - n_converged=("converged", "sum"), - median_obj_err=("final_obj_err", "median"), - median_alpha=("alpha", "median"), - median_sum_beta_sq=("sum_beta_sq", "median"), - ) - .reset_index() - ) - out["conv_rate"] = out["n_converged"] / out["n"] - return out[ - [ - "beta0_ridge", - "l2reg", - "n", - "n_converged", - "conv_rate", - "median_obj_err", - "median_alpha", - "median_sum_beta_sq", - ] - ] - - -def replicate_corr_table(frame: pd.DataFrame) -> pd.DataFrame: - """Replicate-shift Pearson r across fusionreg, per ``(beta0_ridge, l2reg)``. - - ``mut_param_dataset_correlation`` forwards to ``split_apply_combine_muts`` - with ``groupby=("dataset_name", x)`` (``model_collection.py:1444``) and - mean-collapses every other column (``aggregate_func="mean"``). On this - 72-fit frame each ``(rep, fusionreg)`` group holds 9 fits (3 ``beta0_ridge`` - × 3 ``l2reg``), so the BARE call would average away the two axes under test - and return a spuriously smooth r. Slicing with ``query=`` per cell is the - fix — the same move ``l2_fusion_report.replicate_corr_table`` makes for its - single ``l2reg`` loop, extended here to the nested pair. - - Each slice holds 2 datasets × 4 ``fusionreg`` = 8 fits, clearing the - ``<2 datasets`` guard at ``model_collection.py:1440``. The correlated - population is ``rep_1`` vs ``rep_2``; ``inner_merge_dataset_muts`` (default - True) restricts to mutations shared across both. - - Args: - frame: The raw fit-collection DataFrame. - - Returns: - Concatenated per-slice correlation frames (columns ``datasets``, - ``mut_param``, ``correlation``, ``fusionreg``, plus ``beta0_ridge`` and - ``l2reg`` tags), or an empty frame if no slice has ≥2 datasets. - """ - mc = ModelCollection(frame) - out = [] - for b in sorted(frame["beta0_ridge"].unique()): - for l2 in sorted(frame["l2reg"].unique()): - try: - _, df = mc.mut_param_dataset_correlation( - x="fusionreg", - return_data=True, - r=1, - query=f"beta0_ridge == {b} and l2reg == {l2}", - ) - except ValueError: - # Fewer than 2 datasets in this slice — skip it honestly. - continue - df = df.copy() - df["beta0_ridge"] = b - df["l2reg"] = l2 - out.append(df) - return pd.concat(out, ignore_index=True) if out else pd.DataFrame() - - -def shift_corr_pivot(corr: pd.DataFrame) -> pd.DataFrame: - """The ``shift_*`` rows of :func:`replicate_corr_table`, pivoted by fusionreg. - - ``mut_param_dataset_correlation`` returns a row per mutation-parameter type - (``beta_*``, ``shift_*``, ``predicted_func_score_*``). The #284 - reproducibility criterion is the **shift** rows specifically — the - condition-vs-reference parameters the fusion lasso penalizes. Note there is - no ``shift_Omicron_BA1``: it is the reference condition and carries no shift - parameters, so only ``shift_Delta`` and ``shift_Omicron_BA2`` exist. - - Args: - corr: Output of :func:`replicate_corr_table`. - - Returns: - A ``(beta0_ridge, l2reg, mut_param)`` × ``fusionreg`` pivot of Pearson - r, or an empty frame if ``corr`` carries no ``shift_*`` rows. - """ - if not len(corr): - return pd.DataFrame() - shifts = corr[corr["mut_param"].str.startswith("shift_")] - if not len(shifts): - return pd.DataFrame() - return shifts.pivot_table( - index=["beta0_ridge", "l2reg", "mut_param"], - columns="fusionreg", - values="correlation", - ).reset_index() - - -def _load(cache: str, results_dir: Path | None = None) -> pd.DataFrame: - """Load one cache's fit-collection frame. - - Args: - cache: Subdir under the results dir holding ``fit_collection.pkl``. - results_dir: Directory holding the caches. ``None`` → the harness's - own ``RESULTS_DIR``, which resolves relative to *this file*. That - default is wrong when the script runs from a worktree: ``results/`` - is gitignored and materializes only in the canonical clone, so pass - ``--results-dir`` there. - - Returns: - The raw fit-collection DataFrame. - """ - base = results_dir if results_dir is not None else harness.RESULTS_DIR - pkl = Path(base) / cache / "fit_collection.pkl" - frame = pickle.load(open(pkl, "rb")) - size_gb = pkl.stat().st_size / 1e9 - print(f"[report] loaded {len(frame)} fits from {pkl} ({size_gb:.1f} GB)") - return frame - - -def _report(cache: str, label: str, results_dir: Path | None = None) -> None: - """Load one cache, print its tables, and drop the frame before returning. - - Loads inside the call rather than taking a frame, so the caller never holds - two collections at once. These pickles carry the fitted models themselves: - ~88 MB/fit on disk but a measured **~213 MB/fit resident**, so the 72-fit - scan alone peaks near 15 GB. Holding the baseline alongside it would add - ~1.7 GB alongside for no reason, and on a 39 GB laptop that headroom is - worth keeping. - - Args: - cache: Subdir under the results dir holding ``fit_collection.pkl``. - label: Human-readable name for the section headers. - results_dir: Passed through to :func:`_load`. - """ - frame = _load(cache, results_dir) - basin = basin_table(frame) - with pd.option_context("display.width", 200, "display.max_rows", 100): - print(f"\n=== [{label}] convergence + basin per (beta0_ridge, l2reg) ===") - print(convergence_table(basin).to_string(index=False)) - print(f"\n=== [{label}] basin diagnostics (per fit) ===") - print(basin.to_string(index=False)) - corr = replicate_corr_table(frame) - print(f"\n=== [{label}] replicate-SHIFT Pearson r — the #284 criterion ===") - pivot = shift_corr_pivot(corr) - print(pivot.to_string(index=False) if len(pivot) else "(no shift_* rows)") - print(f"\n=== [{label}] all mut_param correlations (per cell slice) ===") - print(corr.to_string(index=False) if len(corr) else "(no ≥2-dataset slice)") - - # Drop the collection (and its fitted models) before the caller loads the - # next one — see the docstring's size note. - del frame - gc.collect() - - -def main() -> None: - """CLI: ``--cache`` (+ optional ``--baseline-cache``) → print all tables.""" - ap = argparse.ArgumentParser(description=__doc__) - ap.add_argument( - "--cache", - default="beta0-ridge-l2-scan", - help="results subdir holding fit_collection.pkl " - "(default: beta0-ridge-l2-scan)", - ) - ap.add_argument( - "--baseline-cache", - default=None, - help="optional (0,0) baseline cache to report alongside " - "(e.g. 277-softplus-floor-off)", - ) - ap.add_argument( - "--results-dir", - default=None, - type=Path, - help="directory holding the caches (default: the harness's results/, " - "resolved next to harness.py). results/ is gitignored and exists only " - "in the canonical clone, so pass this when running from a worktree.", - ) - args = ap.parse_args() - - _report(args.cache, args.cache, args.results_dir) - if args.baseline_cache: - _report( - args.baseline_cache, f"BASELINE {args.baseline_cache}", args.results_dir - ) - - -if __name__ == "__main__": - main() diff --git a/experiments/convergence-lab/diagnostics/beta_control_report.py b/experiments/convergence-lab/diagnostics/beta_control_report.py deleted file mode 100644 index 932aa5f5..00000000 --- a/experiments/convergence-lab/diagnostics/beta_control_report.py +++ /dev/null @@ -1,216 +0,0 @@ -r"""Downstream report for the Phase 1 β-control head-to-head (#263). - -Loads BOTH convergence-lab ``fit_collection.pkl``s (the clip arm and the -l2 arm), tags each fit with an ``arm`` label derived in Python, and prints -three tables: basin diagnostics (Σβ², α, converged, final_obj_err), a -maxiter-each-needs table (outer iteration where objective_error first fell -below 1e-6 and 1e-4), and replicate-shift Pearson r per arm. - -An "arm" is the ``(beta_clip_range, l2reg)`` pair. ``beta_clip_range`` is a -list/None column that pandas ``.query()`` cannot match cleanly, so arms are -sliced by a Python-derived ``arm`` column, not a query string. - -The convergence verdict is the maxiter-each-needs table, NOT the -``converged`` flag: at the strict ``tol=1e-6`` the flag is expected mostly -False (#256 logged all fits False even at looser truncation). - -Run:: - - pixi run python experiments/convergence-lab/diagnostics/beta_control_report.py \\ - --cache-clip beta-control-clip --cache-l2 beta-control-l2 -""" - -from __future__ import annotations - -import argparse -import os -import pickle -import sys -import warnings -from pathlib import Path - -os.environ.setdefault("XLA_FLAGS", "--xla_cpu_multi_thread_eigen=false") -os.environ.setdefault("OMP_NUM_THREADS", "1") -os.environ.setdefault("JAX_PLATFORM_NAME", "cpu") - -warnings.filterwarnings("ignore") - -import pandas as pd # noqa: E402 - -# Reuse the harness's RESULTS_DIR and the l2-fusion report's basin_row. -sys.path.insert(0, str(Path(os.path.abspath(__file__)).parent.parent)) -import harness # noqa: E402 - -sys.path.insert(0, str(Path(os.path.abspath(__file__)).parent)) -import l2_fusion_report as l2rpt # noqa: E402 - -from multidms.model_collection import ModelCollection # noqa: E402 - - -def arm_label(beta_clip_range) -> str: - """Return the arm label for one fit's ``beta_clip_range`` cell. - - The clip arm stores a Python list ``[-10, 10]``; the l2 arm stores the - absent bound, which round-trips through an object-dtype column as EITHER - ``None`` OR ``float('nan')``. Both must map to ``"l2"``. - - Args: - beta_clip_range: The cell value (a list, ``None``, or ``NaN``). - - Returns: - ``"clip"`` if the value is a list, else ``"l2"``. - """ - return "clip" if isinstance(beta_clip_range, (list, tuple)) else "l2" - - -def first_below(trajectory, threshold: float, sentinel: int = 101) -> int: - """First 1-based outer iteration where ``trajectory`` drops below ``threshold``. - - Searched from the FRONT (unlike ``basin_row``'s ``.iloc[-1]``), so the - 1e-4 crossing lands earlier than the last row. - - Args: - trajectory: Iterable of per-sweep ``objective_error`` values. - threshold: The floor to cross. - sentinel: Value returned when the threshold is never crossed - (default 101, one above the maxiter=100 ceiling). - - Returns: - The 1-based index of the first crossing, or ``sentinel``. - """ - for i, val in enumerate(trajectory, start=1): - if val < threshold: - return i - return sentinel - - -def tagged_frame(clip_frame: pd.DataFrame, l2_frame: pd.DataFrame) -> pd.DataFrame: - """Concatenate both arm frames, adding an ``arm`` column derived per row. - - Args: - clip_frame: Raw fit-collection frame from the clip-arm pickle. - l2_frame: Raw fit-collection frame from the l2-arm pickle. - - Returns: - The concatenation with an ``arm`` column (``"clip"`` / ``"l2"``). - """ - out = pd.concat([clip_frame, l2_frame], ignore_index=True) - out["arm"] = out["beta_clip_range"].map(arm_label) - return out - - -def basin_table(frame: pd.DataFrame) -> pd.DataFrame: - """Per-fit basin diagnostics for the tagged collection. - - Args: - frame: The tagged frame from :func:`tagged_frame` (carries ``arm``). - - Returns: - One row per fit (``arm`` + the ``basin_row`` fields), sorted by - ``(arm, fusionreg, dataset_name)``. - """ - rows = [] - for i in range(len(frame)): - fit = frame.iloc[i] - row = l2rpt.basin_row(fit) - row["arm"] = fit.arm - rows.append(row) - return ( - pd.DataFrame(rows) - .sort_values(["arm", "fusionreg", "dataset_name"]) - .reset_index(drop=True) - ) - - -def maxiter_table(frame: pd.DataFrame) -> pd.DataFrame: - """Per-fit outer-iteration counts to cross 1e-6 and 1e-4. - - Args: - frame: The tagged frame from :func:`tagged_frame`. - - Returns: - One row per fit with ``arm``, ``fusionreg``, ``dataset_name``, - ``iters_to_1e6``, ``iters_to_1e4``, sorted by - ``(arm, fusionreg, dataset_name)``. - """ - rows = [] - for i in range(len(frame)): - fit = frame.iloc[i] - traj = fit.model.convergence_trajectory_df["objective_error_trajectory"] - rows.append( - { - "arm": fit.arm, - "fusionreg": fit.fusionreg, - "dataset_name": fit.dataset_name, - "iters_to_1e6": first_below(traj, 1e-6), - "iters_to_1e4": first_below(traj, 1e-4), - } - ) - return ( - pd.DataFrame(rows) - .sort_values(["arm", "fusionreg", "dataset_name"]) - .reset_index(drop=True) - ) - - -def replicate_corr_table(frame: pd.DataFrame) -> pd.DataFrame: - """Replicate-shift Pearson r across fusionreg, per arm slice. - - ``mut_param_dataset_correlation`` groups by ``(dataset_name, fusionreg)``, - so it is run once per ``arm`` level to avoid mixing arms. The arm slice - is selected by an ``arm``-column query (a plain string column, unlike the - list-valued ``beta_clip_range``). - - Args: - frame: The tagged frame from :func:`tagged_frame`. - - Returns: - Concatenated per-arm correlation frames (columns ``datasets``, - ``mut_param``, ``correlation``, ``fusionreg`` plus an ``arm`` tag), - or an empty frame if no arm slice has ≥2 datasets. - """ - mc = ModelCollection(frame) - out = [] - for arm in sorted(frame["arm"].unique()): - try: - _, df = mc.mut_param_dataset_correlation( - x="fusionreg", - return_data=True, - r=1, - query=f"arm == '{arm}'", - ) - except ValueError: - # Fewer than 2 datasets in this slice — skip it honestly. - continue - df = df.copy() - df["arm"] = arm - out.append(df) - return pd.concat(out, ignore_index=True) if out else pd.DataFrame() - - -def main() -> None: - """CLI: ``--cache-clip/--cache-l2`` → print basin + maxiter + replicate-r.""" - ap = argparse.ArgumentParser(description=__doc__) - ap.add_argument("--cache-clip", default="beta-control-clip") - ap.add_argument("--cache-l2", default="beta-control-l2") - args = ap.parse_args() - - clip_pkl = harness.RESULTS_DIR / args.cache_clip / "fit_collection.pkl" - l2_pkl = harness.RESULTS_DIR / args.cache_l2 / "fit_collection.pkl" - clip_frame = pickle.load(open(clip_pkl, "rb")) - l2_frame = pickle.load(open(l2_pkl, "rb")) - frame = tagged_frame(clip_frame, l2_frame) - print(f"[report] {len(clip_frame)} clip + {len(l2_frame)} l2 = {len(frame)} fits\n") - - with pd.option_context("display.width", 200, "display.max_rows", 50): - print("=== basin diagnostics (per fit) ===") - print(basin_table(frame).to_string(index=False)) - print("\n=== maxiter each needs (iters to cross 1e-6 / 1e-4) ===") - print(maxiter_table(frame).to_string(index=False)) - print("\n=== replicate-shift Pearson r (per arm) ===") - corr = replicate_corr_table(frame) - print(corr.to_string(index=False) if len(corr) else "(no ≥2-dataset slice)") - - -if __name__ == "__main__": - main() diff --git a/experiments/convergence-lab/diagnostics/l2_fusion_report.py b/experiments/convergence-lab/diagnostics/l2_fusion_report.py deleted file mode 100644 index d2e88fc5..00000000 --- a/experiments/convergence-lab/diagnostics/l2_fusion_report.py +++ /dev/null @@ -1,155 +0,0 @@ -r"""Downstream report for the l2-fusion β-explosion sweep (#256). - -Loads a convergence-lab ``fit_collection.pkl`` into a ``ModelCollection`` -and prints a per-cell table of basin diagnostics (Σβ², α, converged, -final_obj_err) plus replicate-shift Pearson r per ``l2reg`` slice. The -harness only fits; this computes the derived metrics on the fly, matching -the convergence-lab design (nothing derived is stored in the pickle). - -The ``l2reg=0`` rows are the β-explosion control; success is at least one -``l2reg>0`` value showing bounded Σβ² + α across all three fusion -strengths, with replicate-r that does not collapse as fusionreg rises. - -Run:: - - pixi run python experiments/convergence-lab/diagnostics/l2_fusion_report.py \\ - --cache l2-fusion -""" - -from __future__ import annotations - -import argparse -import os -import pickle -import sys -import warnings -from pathlib import Path - -os.environ.setdefault("XLA_FLAGS", "--xla_cpu_multi_thread_eigen=false") -os.environ.setdefault("OMP_NUM_THREADS", "1") -os.environ.setdefault("JAX_PLATFORM_NAME", "cpu") - -warnings.filterwarnings("ignore") - -import pandas as pd # noqa: E402 - -# Reuse the harness's RESULTS_DIR / path constants. -sys.path.insert(0, str(Path(os.path.abspath(__file__)).parent.parent)) -import harness # noqa: E402 - -from multidms.model_collection import ModelCollection # noqa: E402 - - -def basin_row(fit) -> dict: - """Extract basin diagnostics from one fit-collection row. - - Args: - fit: A row (``pd.Series``) of the fit-collection frame; ``fit.model`` - is a fitted ``multidms.Model``. - - Returns: - Dict with ``l2reg``, ``fusionreg``, ``dataset_name``, ``sum_beta_sq`` - (Σβ² of the reference condition's mutation effects), ``alpha`` (shared - scalar), ``converged`` (bool), and ``final_obj_err``. - """ - model = fit.model - ref = model.data.reference - beta = model.params.φ[ref].β - sum_beta_sq = float((beta**2).sum()) - alpha = float(model.params.α) - try: - final_obj_err = float( - model.convergence_trajectory_df["objective_error_trajectory"].iloc[-1] - ) - except (TypeError, KeyError, IndexError): - final_obj_err = float("nan") - return { - "l2reg": fit.l2reg, - "fusionreg": fit.fusionreg, - "dataset_name": fit.dataset_name, - "sum_beta_sq": sum_beta_sq, - "alpha": alpha, - "converged": bool(model.converged), - "final_obj_err": final_obj_err, - } - - -def basin_table(frame: pd.DataFrame) -> pd.DataFrame: - """Per-fit basin diagnostics for the whole collection. - - Args: - frame: The raw fit-collection DataFrame. - - Returns: - One row per fit, sorted by ``(l2reg, fusionreg, dataset_name)``. - """ - rows = [basin_row(frame.iloc[i]) for i in range(len(frame))] - return ( - pd.DataFrame(rows) - .sort_values(["l2reg", "fusionreg", "dataset_name"]) - .reset_index(drop=True) - ) - - -def replicate_corr_table(frame: pd.DataFrame) -> pd.DataFrame: - """Replicate-shift Pearson r across fusionreg, per l2reg slice. - - ``mut_param_dataset_correlation`` groups correlation by ``(dataset_name, - fusionreg)``, so it must be run once per ``l2reg`` level to avoid mixing - l2 regimes. Returns the concatenated per-slice correlation frames, each - tagged with its ``l2reg``. The underlying frame carries a ``mut_param`` - column (the shift, e.g. ``shift_Delta``), so each ``(l2reg, fusionreg)`` - cell yields one row per shift. - - Args: - frame: The raw fit-collection DataFrame. - - Returns: - Concatenated correlation frames (columns ``datasets``, ``mut_param``, - ``correlation``, ``fusionreg`` plus an ``l2reg`` tag), or an empty - frame if no slice has ≥2 datasets. - """ - mc = ModelCollection(frame) - out = [] - for l2 in sorted(frame["l2reg"].unique()): - try: - _, df = mc.mut_param_dataset_correlation( - x="fusionreg", - return_data=True, - r=1, - query=f"l2reg == {l2}", - ) - except ValueError: - # Fewer than 2 datasets in this slice — skip it honestly. - continue - df = df.copy() - df["l2reg"] = l2 - out.append(df) - return pd.concat(out, ignore_index=True) if out else pd.DataFrame() - - -def main() -> None: - """CLI: ``--cache `` → print basin + replicate-r tables.""" - ap = argparse.ArgumentParser(description=__doc__) - ap.add_argument( - "--cache", - default="l2-fusion", - help="results subdir holding fit_collection.pkl (default: l2-fusion)", - ) - args = ap.parse_args() - - pkl = harness.RESULTS_DIR / args.cache / "fit_collection.pkl" - frame = pickle.load(open(pkl, "rb")) - print(f"[report] loaded {len(frame)} fits from {pkl}\n") - - basin = basin_table(frame) - with pd.option_context("display.width", 200, "display.max_rows", 50): - print("=== basin diagnostics (per fit) ===") - print(basin.to_string(index=False)) - print("\n=== replicate-shift Pearson r (per l2reg slice) ===") - corr = replicate_corr_table(frame) - print(corr.to_string(index=False) if len(corr) else "(no ≥2-dataset slice)") - - -if __name__ == "__main__": - main() diff --git a/experiments/convergence-lab/diagnostics/parallelism_probe.py b/experiments/convergence-lab/diagnostics/parallelism_probe.py deleted file mode 100644 index b02210db..00000000 --- a/experiments/convergence-lab/diagnostics/parallelism_probe.py +++ /dev/null @@ -1,267 +0,0 @@ -"""Two-axis parallelism + memory probe for ``fit_models`` (#253). - -Settles WHY ``fit_models`` hung locally, and whether memory is the cause. The -scv2-spike pipeline runs ``fit_models`` (spawn workers, ``n_processes>1``) from a -papermill notebook and works on servers — but every pipeline config uses -``l2reg: 0.0``, so the ``l2reg>0`` + spawn path was never exercised there. This -walks a data-size x l2reg staircase at ``n_processes=2`` (the real spawn path), -recording BOTH wall-clock and peak resident memory (RSS) per step. - -Why memory: the original hypothesis was that parallel fits duplicate the dataset -per worker and exhaust RAM. With ``spawn`` (not ``fork``) each worker re-imports -and rebuilds its own ``multidms.Data`` from the CSV, so datasets ARE duplicated. -The data-size axis (tiny -> full) is therefore also the memory-growth axis; this -probe reads the gauge a wall-clock-only probe could not. A worker OOM-killed by -the OS looks like a hang from the parent (child dies, queue stays empty), so we -distinguish DIED (child exited non-zero, e.g. SIGKILL) from HANG (child still -alive at timeout). - -Fits are deliberately cheap (low ``maxiter``): a deadlock or memory blowup -manifests regardless of iteration count, so we do not pay for convergence here. - -A proper ``if __name__ == "__main__":`` guard ensures spawn re-imports cleanly -(unlike the earlier marimo/``/tmp`` probes that confounded the diagnosis). - -Run:: - - PROBE=experiments/convergence-lab/diagnostics/parallelism_probe.py - pixi run python $PROBE # np=2 spawn staircase - pixi run python $PROBE --baseline # also the np=1 in-process twin - -``--baseline`` reruns the same staircase at ``n_processes=1`` (in-process) so a -failing step has a sequential twin for comparison (same fits, no spawn, one -process's memory) — isolating whether a problem is spawn-specific. -""" - -from __future__ import annotations - -import argparse -import multiprocessing as mp -import os -import sys -import threading -import time -from pathlib import Path - -os.environ.setdefault("XLA_FLAGS", "--xla_cpu_multi_thread_eigen=false") -os.environ.setdefault("OMP_NUM_THREADS", "1") -os.environ.setdefault("OPENBLAS_NUM_THREADS", "1") -os.environ.setdefault("MKL_NUM_THREADS", "1") -os.environ.setdefault("JAX_PLATFORM_NAME", "cpu") - -import psutil # noqa: E402 - -# Reuse the harness's constant data loader and path/REF constants. -sys.path.insert(0, str(Path(os.path.abspath(__file__)).parent.parent)) -import harness # noqa: E402 - -# Generous per-step ceiling: a healthy `full`-size step (2 parallel fits on the -# whole dataset) is minutes even with cheap iters; a true deadlock never returns -# and still trips this. Calibrated after the smoke run measured ~140-228s for a -# single full fit at maxiter=25 (this probe uses far fewer iters). -STEP_TIMEOUT_S = 420 -# All four sizes survive — the data-size axis IS the memory-growth axis. -SIZES = [("tiny", 200), ("small", 1000), ("medium", 5000), ("full", None)] -L2_LADDER = [0.0, 3e-4] -# Cheap fits: the probe measures parallel-execution mechanics + memory, not -# convergence. A deadlock/OOM is independent of iteration count. -_INNER = dict(tol=1e-4, maxiter=8, maxls=40, jit=True) -_BLOCK_MAXITER = 6 -_RSS_POLL_S = 0.4 - - -def _peak_rss_mb(pid: int, stop_evt: threading.Event, out: dict) -> None: - """Poll a process subtree's RSS until ``stop_evt`` is set; record peak MB. - - Sums the RSS of ``pid`` and all its descendants (the spawn workers are - children of the step process), sampling every ``_RSS_POLL_S`` seconds. - Stores ``out["peak_mb"]`` — the maximum total RSS observed across samples. - - Args: - pid: The step process PID (its children are the fit workers). - stop_evt: Set by the parent when the step finishes / is killed. - out: Mutable dict to receive ``peak_mb``. - """ - peak = 0.0 - while not stop_evt.is_set(): - try: - proc = psutil.Process(pid) - procs = [proc] + proc.children(recursive=True) - total = 0 - for p in procs: - try: - total += p.memory_info().rss - except (psutil.NoSuchProcess, psutil.AccessDenied): - continue - peak = max(peak, total) - except (psutil.NoSuchProcess, psutil.AccessDenied): - pass - time.sleep(_RSS_POLL_S) - out["peak_mb"] = round(peak / 1e6, 1) - - -def _fit_step(size_n, l2reg, n_processes, q): - """Child entry: fit 2 models in parallel on a data subset; report seconds.""" - import warnings - - warnings.filterwarnings("ignore") - import multidms - from multidms.model_collection import fit_models - - rep_data = harness.load_rep_data() - datasets = list(rep_data.values())[:2] - if size_n is not None: - # Subset each Data's variants by rebuilding from a truncated CSV view. - # Take head(size_n) PER CONDITION (not a blind head of the replicate): - # the CSV is condition-ordered, so a blind head() would drop the - # reference condition entirely and multidms.Data would reject it - # ("reference must be in condition factor levels"). Per-condition heads - # guarantee every condition — including the reference — is represented. - import pandas as pd - - raw = pd.read_csv(harness.DATA_CSV).fillna({"aa_substitutions": ""}) - datasets = [] - for rep in sorted(raw["replicate"].unique())[:2]: - df_rep = ( - raw[raw["replicate"] == rep] - .groupby("condition", dropna=False, group_keys=False) - .head(size_n) - ) - df_agg = ( - df_rep.groupby(["condition", "aa_substitutions"], dropna=False) - .agg({"func_score": "mean"}) - .reset_index() - ) - datasets.append( - multidms.Data( - df_agg, - alphabet=multidms.AAS_WITHSTOP_WITHGAP, - reference=harness.REF, - assert_site_integrity=False, - name=f"rep_{rep}", - verbose=False, - ) - ) - - params = { - "dataset": datasets, - "l2reg": [l2reg], - "warmstart": [True], - "recompute_scale": [False], - "share_alpha": [True], - "fusionreg": [0.0], - "ge_type": ["Sigmoid"], - "maxiter": [_BLOCK_MAXITER], - "tol": [1e-6], - "ge_kwargs": [dict(_INNER)], - "cal_kwargs": [dict(_INNER)], - } - t0 = time.time() - fit_models(params, n_processes=n_processes) - q.put(time.time() - t0) - - -def run_staircase(n_processes: int) -> None: - """Walk the size x l2reg staircase, recording time + peak RSS per step. - - Stops at the first step that does not return cleanly, distinguishing DIED - (the step process exited non-zero — e.g. an OOM SIGKILL) from HANG (still - alive at the timeout — a true deadlock). Peak RSS is sampled for every step, - including failing ones, so a memory blowup is visible even when the step - is killed. - - Args: - n_processes: Passed to ``fit_models`` (2 = real spawn path; 1 = in-process). - """ - ctx = mp.get_context("spawn") - vm = psutil.virtual_memory() - print(f"\n=== staircase (n_processes={n_processes}) ===", flush=True) - print( - f" system: total={vm.total / 1e9:.1f}GB available={vm.available / 1e9:.1f}GB", - flush=True, - ) - rows = [] - stop = False - for l2reg in L2_LADDER: - for size_label, size_n in SIZES: - q = ctx.Queue() - p = ctx.Process(target=_fit_step, args=(size_n, l2reg, n_processes, q)) - p.start() - - mem: dict = {"peak_mb": float("nan")} - stop_evt = threading.Event() - sampler = threading.Thread( - target=_peak_rss_mb, args=(p.pid, stop_evt, mem), daemon=True - ) - sampler.start() - - p.join(STEP_TIMEOUT_S) - - if p.is_alive(): - # Still running at the ceiling -> true deadlock / hang. - stop_evt.set() - sampler.join(timeout=2) - p.terminate() - p.join() - peak = mem.get("peak_mb", float("nan")) - rows.append( - (size_label, l2reg, n_processes, "HANG", STEP_TIMEOUT_S, peak) - ) - print( - f" {size_label:<7} l2reg={l2reg:<7} np={n_processes} -> HANG " - f"(>{STEP_TIMEOUT_S}s, peak {peak}MB) STOPPING", - flush=True, - ) - stop = True - break - - # Process finished — stop sampling and read the verdict. - stop_evt.set() - sampler.join(timeout=2) - peak = mem.get("peak_mb", float("nan")) - - if p.exitcode != 0 or q.empty(): - # Child exited non-zero (e.g. OOM SIGKILL) or produced no time. - rows.append((size_label, l2reg, n_processes, "DIED", p.exitcode, peak)) - print( - f" {size_label:<7} l2reg={l2reg:<7} np={n_processes} -> DIED " - f"(exitcode={p.exitcode}, peak {peak}MB) STOPPING", - flush=True, - ) - stop = True - break - - secs = round(q.get(), 1) - rows.append((size_label, l2reg, n_processes, "PASS", secs, peak)) - print( - f" {size_label:<7} l2reg={l2reg:<7} np={n_processes} -> PASS " - f"({secs}s, peak {peak}MB)", - flush=True, - ) - if stop: - break - - print("\n step data l2reg np result secs/code peak_MB") - for i, (size_label, l2reg, npr, result, val, peak) in enumerate(rows, 1): - print( - f" {i:>4} {size_label:<7} {l2reg:<7} {npr:>2} {result:<6} " - f"{str(val):<9} {peak}" - ) - - -def main() -> None: - """CLI: default np=2 staircase; ``--baseline`` adds the np=1 twin.""" - ap = argparse.ArgumentParser(description=__doc__) - ap.add_argument( - "--baseline", - action="store_true", - help="also run the staircase at n_processes=1 (in-process)", - ) - args = ap.parse_args() - run_staircase(n_processes=2) - if args.baseline: - run_staircase(n_processes=1) - - -if __name__ == "__main__": - main() diff --git a/experiments/convergence-lab/diagnostics/test_beta0_ridge_l2_grid.py b/experiments/convergence-lab/diagnostics/test_beta0_ridge_l2_grid.py deleted file mode 100644 index 60c80561..00000000 --- a/experiments/convergence-lab/diagnostics/test_beta0_ridge_l2_grid.py +++ /dev/null @@ -1,97 +0,0 @@ -"""Grid-config test: the beta0_ridge × l2reg scan expands to exactly 72 cells (#284). - -Pure config validation — no fitting, no pickle, so this always runs (unlike the -report tests, which need a fit collection). Guards the two things the #284 -acceptance criteria pin: the cell count and the inherited clip bound. -""" - -import math -import os -import sys -from pathlib import Path - -DIAG = Path(__file__).resolve().parent -sys.path.insert(0, str(DIAG.parent)) - -os.environ.setdefault("JAX_PLATFORM_NAME", "cpu") - -GRID = DIAG.parent / "grids" / "beta0-ridge-l2-scan.yaml" - - -def test_grid_expands_to_72_cells(): - """4 fusionreg × 3 beta0_ridge × 3 l2reg × 2 replicates = 72 fits.""" - import harness - - config = harness.load_config(GRID) - cells = harness.explode_grid(config) - assert len(cells) == 72, f"expected 72 cells, got {len(cells)}" - - -def test_grid_sweeps_the_three_axes(): - """The three swept axes carry exactly the specified values.""" - import harness - - config = harness.load_config(GRID) - assert config["sweep"]["fusionreg"] == [0.0, 4.0e-5, 1.6e-4, 6.4e-4] - assert config["sweep"]["beta0_ridge"] == [1.0e-4, 1.0e-3, 1.0e-2] - assert config["sweep"]["l2reg"] == [1.0e-8, 1.0e-7, 1.0e-6] - assert config["replicates"] == [1, 2] - - -def test_grid_inherits_the_clip_bound(): - """The #263 clip bound and the baseline's regime flags are inherited verbatim.""" - import harness - - config = harness.load_config(GRID) - assert config["fixed"]["beta_clip_range"] == [-10, 10] - assert config["fixed"]["output_floor"] is None - assert config["fixed"]["share_alpha"] is True - assert config["fixed"]["recompute_scale"] is False - - -def test_build_params_preserves_both_new_axes(): - """The swept axes survive build_params' distinct-value collapse. - - ``explode_grid`` proving 72 cells is not enough: ``build_params`` re-collects - each kwarg's distinct values into a list for ``fit_models`` to re-cross - (``harness.py:210-218``), and that collapse is what actually reaches the - fitter. A float axis that collapsed wrong would silently shrink the grid. - """ - import harness - - config = harness.load_config(GRID) - cells = harness.explode_grid(config) - params = harness.build_params(cells, {"rep_1": "DATA1", "rep_2": "DATA2"}) - - assert params["beta0_ridge"] == [1.0e-4, 1.0e-3, 1.0e-2] - assert params["l2reg"] == [1.0e-8, 1.0e-7, 1.0e-6] - assert params["fusionreg"] == [0.0, 4.0e-5, 1.6e-4, 6.4e-4] - assert len(params["dataset"]) == 2 - - n_fits = len(params["dataset"]) * math.prod( - len(v) for k, v in params.items() if k != "dataset" - ) - assert n_fits == 72, f"fit_models would fit {n_fits}, not 72" - - -def test_grid_matches_the_baseline_fixed_block(): - """Every fixed key equals softplus-floor-off's, minus the two swept axes. - - This is the #284 AC1 check in executable form: the grid must be a - one-knob-at-a-time delta from its baseline, so any drift in the inherited - block is a bug, not a choice. - """ - import harness - - baseline = harness.load_config(DIAG.parent / "grids" / "softplus-floor-off.yaml") - scan = harness.load_config(GRID) - - moved = {"l2reg", "beta0_ridge"} - assert moved <= set(baseline["fixed"]), "baseline should fix the axes we sweep" - assert moved <= set(scan["sweep"]), "scan should sweep them" - assert not (moved & set(scan["fixed"])), "scan must not also fix them" - - expected_fixed = {k: v for k, v in baseline["fixed"].items() if k not in moved} - assert scan["fixed"] == expected_fixed, "inherited fixed block drifted" - assert scan["sweep"]["fusionreg"] == baseline["sweep"]["fusionreg"] - assert scan["replicates"] == baseline["replicates"] diff --git a/experiments/convergence-lab/diagnostics/test_beta0_ridge_l2_report.py b/experiments/convergence-lab/diagnostics/test_beta0_ridge_l2_report.py deleted file mode 100644 index 0e313a00..00000000 --- a/experiments/convergence-lab/diagnostics/test_beta0_ridge_l2_report.py +++ /dev/null @@ -1,228 +0,0 @@ -"""Tests for the beta0_ridge × l2reg report's extraction and grouping (#284).""" - -import os -import pickle -import sys -from pathlib import Path - -import pandas as pd -import pytest - -DIAG = Path(__file__).resolve().parent -sys.path.insert(0, str(DIAG)) - - -def _sample_pickle(): - """Locate a sample fit_collection.pkl to extract from. - - ``results/`` is gitignored and never materialized in a worktree, so honor a - ``BETA0_RIDGE_TEST_PKL`` override (point it at the canonical clone's copy); - otherwise fall back to the in-tree baseline cache. Returns ``None`` when no - fixture is reachable, so the test skips rather than fails. - """ - override = os.environ.get("BETA0_RIDGE_TEST_PKL") - if override and Path(override).exists(): - return Path(override) - in_tree = DIAG.parent / "results" / "277-softplus-floor-off" / "fit_collection.pkl" - return in_tree if in_tree.exists() else None - - -def _synthetic_basin() -> pd.DataFrame: - """A basin frame with a known answer: 4 cells × 4 fits, half converged. - - 2 beta0_ridge × 2 l2reg × 2 fusionreg × 2 replicates = 16 rows. Exactly one - replicate per cell converges, so every cell's rate must be 0.5. - """ - return pd.DataFrame( - [ - { - "beta0_ridge": b, - "l2reg": l2, - "fusionreg": f, - "dataset_name": d, - "sum_beta_sq": 100.0, - "alpha": 5.0, - "converged": conv, - "final_obj_err": 1e-4, - } - for b in (1e-4, 1e-3) - for l2 in (1e-8, 1e-7) - for f in (0.0, 6.4e-4) - for d, conv in (("rep_1", True), ("rep_2", False)) - ] - ) - - -def test_convergence_table_groups_by_both_axes(): - """convergence_table yields one row per (beta0_ridge, l2reg) with a rate. - - Uses a synthetic basin frame — no pickle needed, so this always runs. - """ - import beta0_ridge_l2_report as rpt - - out = rpt.convergence_table(_synthetic_basin()) - assert len(out) == 4, "2 beta0_ridge × 2 l2reg = 4 cells" - assert set(out.columns) >= { - "beta0_ridge", - "l2reg", - "n", - "n_converged", - "conv_rate", - "median_obj_err", - } - assert (out["n"] == 4).all(), "each cell holds 2 fusionreg × 2 reps = 4 fits" - assert (out["conv_rate"] == 0.5).all(), "1 of 2 reps converged per cell" - - -def test_convergence_table_reports_a_degenerate_rate_honestly(): - """An all-unconverged grid yields conv_rate 0.0 with a finite median obj_err. - - This is the expected real-world case (#284 pre-registers a likely 0/72), and - it is exactly when the median_obj_err fallback becomes the adjudicator — so - the fallback column must still carry a real number when the rate is - degenerate. - """ - import beta0_ridge_l2_report as rpt - - basin = _synthetic_basin() - basin["converged"] = False - out = rpt.convergence_table(basin) - assert (out["conv_rate"] == 0.0).all(), "degenerate rate reported as 0.0" - assert out["median_obj_err"].notna().all(), "fallback must survive a 0/n rate" - - -def test_basin_table_sort_order_without_a_pickle(monkeypatch): - """basin_table sorts beta0_ridge-outermost — checked without a fit collection. - - The pickle-backed sort test skips on a fresh checkout (results/ is - gitignored), so the sort contract would go unverified exactly where it is - most likely to regress. Stub the per-row extraction to cover the ordering on - synthetic rows, which is all the sort itself depends on. - """ - import beta0_ridge_l2_report as rpt - - rows = [ - { - "beta0_ridge": b, - "l2reg": l2, - "fusionreg": f, - "dataset_name": d, - "sum_beta_sq": 1.0, - "alpha": 1.0, - "converged": False, - "final_obj_err": 1e-4, - } - # Deliberately shuffled relative to the expected sort order. - for b in (1e-2, 1e-4) - for l2 in (1e-6, 1e-8) - for f in (6.4e-4, 0.0) - for d in ("rep_2", "rep_1") - ] - frame = pd.DataFrame(rows) - monkeypatch.setattr( - rpt, "basin_row_with_ridge", lambda fit: dict(fit), raising=True - ) - - table = rpt.basin_table(frame) - keys = ["beta0_ridge", "l2reg", "fusionreg", "dataset_name"] - assert table[keys].values.tolist() == ( - frame.sort_values(keys).reset_index(drop=True)[keys].values.tolist() - ) - # The outermost key really is beta0_ridge, not l2reg (the prior art's order). - assert table["beta0_ridge"].is_monotonic_increasing - - -def test_shift_corr_pivot_keeps_only_shift_rows(): - """shift_corr_pivot drops beta_*/predicted_* and pivots shift_* by fusionreg. - - The real correlation frame carries a row per mutation-parameter type; #284's - reproducibility criterion is the shift_* rows only. Synthetic — always runs. - """ - import beta0_ridge_l2_report as rpt - - corr = pd.DataFrame( - [ - { - "datasets": "rep_1,rep_2", - "mut_param": mp, - "correlation": 0.5, - "fusionreg": f, - "beta0_ridge": 1e-4, - "l2reg": 1e-8, - } - for mp in ( - "beta_Delta", - "shift_Delta", - "predicted_func_score_Delta", - "shift_Omicron_BA2", - ) - for f in (0.0, 6.4e-4) - ] - ) - out = rpt.shift_corr_pivot(corr) - assert len(out) == 2, "only shift_Delta and shift_Omicron_BA2 survive" - assert set(out["mut_param"]) == {"shift_Delta", "shift_Omicron_BA2"} - for col in (0.0, 6.4e-4): - assert col in out.columns, f"fusionreg {col} should be a pivoted column" - - -def test_shift_corr_pivot_handles_empty_input(): - """An empty or shift-less correlation frame yields an empty pivot, not a raise.""" - import beta0_ridge_l2_report as rpt - - assert not len(rpt.shift_corr_pivot(pd.DataFrame())) - beta_only = pd.DataFrame( - [ - { - "datasets": "rep_1,rep_2", - "mut_param": "beta_Delta", - "correlation": 0.5, - "fusionreg": 0.0, - "beta0_ridge": 1e-4, - "l2reg": 1e-8, - } - ] - ) - assert not len(rpt.shift_corr_pivot(beta_only)) - - -@pytest.mark.skipif( - _sample_pickle() is None, - reason="needs a sample fit_collection.pkl (set BETA0_RIDGE_TEST_PKL)", -) -def test_basin_row_with_ridge_adds_beta0_ridge(): - """basin_row_with_ridge returns basin_row's keys plus beta0_ridge.""" - import beta0_ridge_l2_report as rpt - - frame = pickle.load(open(_sample_pickle(), "rb")) - out = rpt.basin_row_with_ridge(frame.iloc[0]) - assert "beta0_ridge" in out - for key in ( - "l2reg", - "fusionreg", - "dataset_name", - "sum_beta_sq", - "alpha", - "converged", - "final_obj_err", - ): - assert key in out, f"lost {key} from basin_row's contract" - assert out["sum_beta_sq"] >= 0.0 - assert isinstance(out["converged"], bool) - - -@pytest.mark.skipif( - _sample_pickle() is None, - reason="needs a sample fit_collection.pkl (set BETA0_RIDGE_TEST_PKL)", -) -def test_basin_table_sorts_by_beta0_ridge_first(): - """basin_table sorts by (beta0_ridge, l2reg, fusionreg, dataset_name).""" - import beta0_ridge_l2_report as rpt - - frame = pickle.load(open(_sample_pickle(), "rb")) - table = rpt.basin_table(frame) - assert len(table) == len(frame) - expected = table.sort_values( - ["beta0_ridge", "l2reg", "fusionreg", "dataset_name"] - ).reset_index(drop=True) - pd.testing.assert_frame_equal(table, expected) diff --git a/experiments/convergence-lab/diagnostics/test_beta_control_report.py b/experiments/convergence-lab/diagnostics/test_beta_control_report.py deleted file mode 100644 index 0758f999..00000000 --- a/experiments/convergence-lab/diagnostics/test_beta_control_report.py +++ /dev/null @@ -1,67 +0,0 @@ -"""Unit tests for the Phase 1 β-control report (#263).""" - -import os -import pickle -import sys -from pathlib import Path - -import pandas as pd -import pytest - -DIAG = Path(__file__).resolve().parent -sys.path.insert(0, str(DIAG)) - - -def _sample_pickle() -> Path | None: - """Locate a sample fit_collection.pkl for the end-to-end test. - - The convergence-lab pickles are gitignored (absent in a fresh worktree). - Honor a ``BETA_CONTROL_TEST_PKL`` env override (point it at a canonical - clone's copy); otherwise fall back to any beta-control pickle in-tree. - Returns ``None`` when none is reachable, so the end-to-end test skips. - """ - override = os.environ.get("BETA_CONTROL_TEST_PKL") - if override and Path(override).exists(): - return Path(override) - for cache in ("beta-control-clip", "beta-control-l2"): - p = DIAG.parent / "results" / cache / "fit_collection.pkl" - if p.exists(): - return p - return None - - -def test_arm_label_treats_none_and_nan_as_l2(): - """[-10,10]→clip; both None and NaN→l2 (object-column coercion case).""" - import beta_control_report as rpt - - frame = pd.DataFrame({"beta_clip_range": [[-10, 10], None, float("nan")]}) - labels = frame["beta_clip_range"].map(rpt.arm_label).tolist() - assert labels == ["clip", "l2", "l2"] - - -def test_first_below_crossing_and_sentinel(): - """first_below finds the 1-based crossing index, else the 101 sentinel.""" - import beta_control_report as rpt - - traj = [1.0, 1e-3, 1e-5, 1e-7] - assert rpt.first_below(traj, 1e-4) == 3 # index 2, 1-based - assert rpt.first_below(traj, 1e-6) == 4 # index 3, 1-based - never = [1.0, 1e-2, 1e-3] - assert rpt.first_below(never, 1e-6) == 101 - - -@pytest.mark.skipif( - _sample_pickle() is None, - reason="needs a beta-control fit_collection.pkl (set BETA_CONTROL_TEST_PKL)", -) -def test_end_to_end_tables_nonempty(): - """tagged_frame → basin/maxiter tables are non-empty on a real pickle.""" - import beta_control_report as rpt - - frame = pickle.load(open(_sample_pickle(), "rb")) - # Reuse the same real frame for both arms; tagged_frame just concatenates - # and labels, so the test exercises the table builders end-to-end. - tagged = rpt.tagged_frame(frame, frame) - assert "arm" in tagged.columns - assert len(rpt.basin_table(tagged)) == len(tagged) - assert len(rpt.maxiter_table(tagged)) == len(tagged) diff --git a/experiments/convergence-lab/diagnostics/test_l2_fusion_report.py b/experiments/convergence-lab/diagnostics/test_l2_fusion_report.py deleted file mode 100644 index ac6abb2b..00000000 --- a/experiments/convergence-lab/diagnostics/test_l2_fusion_report.py +++ /dev/null @@ -1,47 +0,0 @@ -"""Smoke test for the l2-fusion report's per-fit metric extraction.""" - -import os -import pickle -import sys -from pathlib import Path - -import pytest - -DIAG = Path(__file__).resolve().parent -sys.path.insert(0, str(DIAG)) - - -def _sample_pickle() -> Path | None: - """Locate a sample fit_collection.pkl to extract from. - - The simulation collection is gitignored, so it is absent from a fresh - checkout (and from a worktree, where ``results/`` is never materialized). - Honor a ``L2_FUSION_TEST_PKL`` env override (point it at the canonical - clone's copy when running from a worktree); otherwise fall back to the - in-tree path. Returns ``None`` when no fixture is reachable, so the test - skips rather than fails. - """ - override = os.environ.get("L2_FUSION_TEST_PKL") - if override and Path(override).exists(): - return Path(override) - in_tree = ( - Path(__file__).resolve().parents[3] - / "experiments/simulation/results/fit_collection.pkl" - ) - return in_tree if in_tree.exists() else None - - -@pytest.mark.skipif( - _sample_pickle() is None, - reason="needs a sample fit_collection.pkl (set L2_FUSION_TEST_PKL)", -) -def test_basin_row_extracts_finite_numbers(): - """basin_row returns finite Σβ², α, and a bool converged for a real fit.""" - import l2_fusion_report as rpt - - frame = pickle.load(open(_sample_pickle(), "rb")) - row = frame.iloc[0] - out = rpt.basin_row(row) - assert out["sum_beta_sq"] >= 0.0 - assert isinstance(out["converged"], bool) - assert out["alpha"] == out["alpha"] # not NaN diff --git a/experiments/convergence-lab/grids/beta-control-clip.yaml b/experiments/convergence-lab/grids/beta-control-clip.yaml deleted file mode 100644 index 844dcbbd..00000000 --- a/experiments/convergence-lab/grids/beta-control-clip.yaml +++ /dev/null @@ -1,16 +0,0 @@ -# Phase 1 (#263) — β-control head-to-head, CLIP arm. -# 3 fusionreg × 2 reps = 6 fits. beta_clip_range=[-10,10], no l2reg. -sweep: - fusionreg: [0.0, 4.0e-5, 6.4e-4] # no-fusion / prod lasso_choice / prod-max -fixed: - beta_clip_range: [-10, 10] - l2reg: 0.0 - warmstart: false - recompute_scale: false - share_alpha: true - ge_type: Sigmoid - maxiter: 100 - tol: 1.0e-6 - ge_kwargs: {tol: 1.0e-4, maxiter: 20, maxls: 40, jit: true} - cal_kwargs: {tol: 1.0e-4, maxiter: 20, maxls: 40, jit: true} -replicates: [1, 2] diff --git a/experiments/convergence-lab/grids/beta-control-l2.yaml b/experiments/convergence-lab/grids/beta-control-l2.yaml deleted file mode 100644 index d1a65af1..00000000 --- a/experiments/convergence-lab/grids/beta-control-l2.yaml +++ /dev/null @@ -1,16 +0,0 @@ -# Phase 1 (#263) — β-control head-to-head, L2 arm. -# 3 fusionreg × 2 reps = 6 fits. l2reg=1e-4 (#256 working weight), no clip. -sweep: - fusionreg: [0.0, 4.0e-5, 6.4e-4] -fixed: - beta_clip_range: null - l2reg: 1.0e-4 - warmstart: false - recompute_scale: false - share_alpha: true - ge_type: Sigmoid - maxiter: 100 - tol: 1.0e-6 - ge_kwargs: {tol: 1.0e-4, maxiter: 20, maxls: 40, jit: true} - cal_kwargs: {tol: 1.0e-4, maxiter: 20, maxls: 40, jit: true} -replicates: [1, 2] diff --git a/experiments/convergence-lab/grids/beta0-ridge-l2-scan.yaml b/experiments/convergence-lab/grids/beta0-ridge-l2-scan.yaml deleted file mode 100644 index 02ff23ce..00000000 --- a/experiments/convergence-lab/grids/beta0-ridge-l2-scan.yaml +++ /dev/null @@ -1,36 +0,0 @@ -# beta0_ridge × sub-knee l2reg scan (#284) — does a whisper of L2 (three-to-four -# orders below the measured ~3e-4 knee) or a nonzero beta0_ridge buy convergence -# or replicate reproducibility ON TOP OF the settled beta_clip_range=[-10,10] -# bound (#263)? Baseline = the separately-fit (l2reg=0, beta0_ridge=0) cache -# `277-softplus-floor-off`, whose fixed block this inherits VERBATIM except for -# the two axes below moving fixed: -> sweep:. -# -# NOTE beta0_ridge does NOT penalize the intercepts — it penalizes each -# non-reference condition's β0 DIFFERENCE from the reference, (β0_d − β0_ref)² -# (jaxmodels.py:533-540). It is an L2 shift-shrinkage, structurally parallel to -# fusionreg's L1 on the β shift — hence crossing both axes here. Its penalty is -# also exactly 0 at initialization (beta0_init pins all three conditions to 0.0) -# and only bites as the β0s separate during fitting. -# -# 4 fusionreg × 3 beta0_ridge × 3 l2reg × 2 reps = 72 fits. -sweep: - fusionreg: [0.0, 4.0e-5, 1.6e-4, 6.4e-4] # same 4-point axis as #274/#277 - beta0_ridge: [1.0e-4, 1.0e-3, 1.0e-2] # NEW axis — never swept in this lab - l2reg: [1.0e-8, 1.0e-7, 1.0e-6] # NEW axis — all sub-knee (knee ~3e-4) -fixed: - output_floor: null # OFF — matches the baseline cache - share_alpha: true # #274 baseline (shared-α) condition - # --- softplus-floor-off.yaml fixed block, VERBATIM minus the two swept axes --- - warmstart: false - recompute_scale: false - ge_type: Sigmoid - scale_fusion_by_n: false - alpha_init: 6.0 - beta0_init: {Omicron_BA1: 0.0, Delta: 0.0, Omicron_BA2: 0.0} - beta_clip_range: [-10, 10] # settled β-bound (#263) — inherited - loss_kwargs: {δ: 1.0} - maxiter: 50 # outer block cap - tol: 1.0e-5 # outer tol - ge_kwargs: {tol: 1.0e-4, maxiter: 10, maxls: 10, jit: true, verbose: false} - cal_kwargs: {tol: 1.0e-4, maxiter: 10, maxls: 10, jit: true, verbose: false} -replicates: [1, 2] diff --git a/experiments/convergence-lab/grids/free-alpha.yaml b/experiments/convergence-lab/grids/free-alpha.yaml deleted file mode 100644 index 8f398f12..00000000 --- a/experiments/convergence-lab/grids/free-alpha.yaml +++ /dev/null @@ -1,25 +0,0 @@ -# free-alpha grid (#273) — does a per-condition α (share_alpha=false) lower -# Delta's functional-score floor toward the true −3.5 that a shared α can't -# reach? A smaller, local version of the PR 270 recompute-scale run. -# 2 share_alpha × 4 fusionreg × 2 reps = 16 fits. Diagnostic, NOT tuning. -sweep: - share_alpha: [true, false] # the toggle under test - fusionreg: [0.0, 4.0e-5, 1.6e-4, 6.4e-4] # reduced from PR 270's 9-point grid -fixed: - # --- PR 270 config_recompute_false.yaml fitting block, verbatim --- - warmstart: false - recompute_scale: false - ge_type: Sigmoid - l2reg: 0.0 - beta0_ridge: 0.0 - scale_fusion_by_n: false - alpha_init: 6.0 - beta0_init: {Omicron_BA1: 0.0, Delta: 0.0, Omicron_BA2: 0.0} - beta_clip_range: [-10, 10] - loss_kwargs: {δ: 1.0} - maxiter: 50 # outer block cap (PR 270) - # --- deltas from PR 270 --- - tol: 1.0e-5 # outer tol LOOSENED from PR 270's 1e-6 - ge_kwargs: {tol: 1.0e-4, maxiter: 10, maxls: 10, jit: true, verbose: false} - cal_kwargs: {tol: 1.0e-4, maxiter: 10, maxls: 10, jit: true, verbose: false} -replicates: [1, 2] diff --git a/experiments/convergence-lab/grids/l2-fusion.yaml b/experiments/convergence-lab/grids/l2-fusion.yaml deleted file mode 100644 index ee67ffad..00000000 --- a/experiments/convergence-lab/grids/l2-fusion.yaml +++ /dev/null @@ -1,16 +0,0 @@ -# l2-fusion grid (#256) — does an l2reg penalty tame the β-explosion across -# the prod fusionreg axis? 3 l2reg × 3 fusionreg × 2 reps = 18 fits. -# Diagnostic confirmation, NOT hyperparameter tuning. Independent fitting only. -sweep: - l2reg: [0.0, 1.0e-4, 3.0e-4] # control (explosion) + 2 candidate fixes, below the 3e-4 knee - fusionreg: [0.0, 4.0e-5, 6.4e-4] # prod sweep: no-fusion, positional midpoint (lasso_choice), prod max -fixed: - warmstart: true - recompute_scale: false # scale held fixed (#246) — never recomputed - share_alpha: true - ge_type: Sigmoid - maxiter: 25 # diagnostic, capped for speed (same as smoke) - tol: 1.0e-6 - ge_kwargs: {tol: 1.0e-4, maxiter: 20, maxls: 40, jit: true} - cal_kwargs: {tol: 1.0e-4, maxiter: 20, maxls: 40, jit: true} -replicates: [1, 2] diff --git a/experiments/convergence-lab/grids/maxiter-scan-m1.yaml b/experiments/convergence-lab/grids/maxiter-scan-m1.yaml deleted file mode 100644 index 23cfb0bc..00000000 --- a/experiments/convergence-lab/grids/maxiter-scan-m1.yaml +++ /dev/null @@ -1,28 +0,0 @@ -# maxiter-scan (inner ge/cal maxiter × recompute_scale) — LEVEL m1 (ge=cal=1). -# Aligned with softplus-floor-off.yaml: same fixed fitting block VERBATIM, -# but recompute_scale is now SWEPT [true, false] and the inner ge/cal_kwargs -# maxiter is pinned to 1 (coupled: ge and cal move together). Outer `maxiter` -# (the block cap) stays at 50 — this scan is the INNER optimizer cap only. -# Grid: recompute_scale[2] × fusionreg[4] × reps[2] = 16 fits. -# Companion levels: maxiter-scan-m10.yaml (ge=cal=10), maxiter-scan-m100.yaml -# (ge=cal=100). Combine the three fit_collection.pkl for the full 48-fit study. -sweep: - fusionreg: [0.0, 4.0e-5, 1.6e-4, 6.4e-4] # same 4-point axis as #277 - recompute_scale: [true, false] # SWEPT (was fixed=false) -fixed: - output_floor: null # OFF control — matched in-harness baseline - share_alpha: true # #274 baseline (shared-α) condition - warmstart: false - ge_type: Sigmoid - l2reg: 0.0 - beta0_ridge: 0.0 - scale_fusion_by_n: false - alpha_init: 6.0 - beta0_init: {Omicron_BA1: 0.0, Delta: 0.0, Omicron_BA2: 0.0} - beta_clip_range: [-10, 10] - loss_kwargs: {δ: 1.0} - maxiter: 50 # outer block cap (unchanged) - tol: 1.0e-5 # outer tol - ge_kwargs: {tol: 1.0e-4, maxiter: 1, maxls: 10, jit: true, verbose: false} - cal_kwargs: {tol: 1.0e-4, maxiter: 1, maxls: 10, jit: true, verbose: false} -replicates: [1, 2] diff --git a/experiments/convergence-lab/grids/maxiter-scan-m10.yaml b/experiments/convergence-lab/grids/maxiter-scan-m10.yaml deleted file mode 100644 index a4865f2d..00000000 --- a/experiments/convergence-lab/grids/maxiter-scan-m10.yaml +++ /dev/null @@ -1,29 +0,0 @@ -# maxiter-scan (inner ge/cal maxiter × recompute_scale) — LEVEL m10 (ge=cal=10). -# Aligned with softplus-floor-off.yaml: same fixed fitting block VERBATIM, -# but recompute_scale is now SWEPT [true, false] and the inner ge/cal_kwargs -# maxiter is pinned to 10 (coupled: ge and cal move together). This level's -# ge/cal maxiter (10) matches softplus-floor-off.yaml exactly. Outer `maxiter` -# (the block cap) stays at 50 — this scan is the INNER optimizer cap only. -# Grid: recompute_scale[2] × fusionreg[4] × reps[2] = 16 fits. -# Companion levels: maxiter-scan-m1.yaml (ge=cal=1), maxiter-scan-m100.yaml -# (ge=cal=100). Combine the three fit_collection.pkl for the full 48-fit study. -sweep: - fusionreg: [0.0, 4.0e-5, 1.6e-4, 6.4e-4] # same 4-point axis as #277 - recompute_scale: [true, false] # SWEPT (was fixed=false) -fixed: - output_floor: null # OFF control — matched in-harness baseline - share_alpha: true # #274 baseline (shared-α) condition - warmstart: false - ge_type: Sigmoid - l2reg: 0.0 - beta0_ridge: 0.0 - scale_fusion_by_n: false - alpha_init: 6.0 - beta0_init: {Omicron_BA1: 0.0, Delta: 0.0, Omicron_BA2: 0.0} - beta_clip_range: [-10, 10] - loss_kwargs: {δ: 1.0} - maxiter: 50 # outer block cap (unchanged) - tol: 1.0e-5 # outer tol - ge_kwargs: {tol: 1.0e-4, maxiter: 10, maxls: 10, jit: true, verbose: false} - cal_kwargs: {tol: 1.0e-4, maxiter: 10, maxls: 10, jit: true, verbose: false} -replicates: [1, 2] diff --git a/experiments/convergence-lab/grids/maxiter-scan-m100.yaml b/experiments/convergence-lab/grids/maxiter-scan-m100.yaml deleted file mode 100644 index 2f85f60b..00000000 --- a/experiments/convergence-lab/grids/maxiter-scan-m100.yaml +++ /dev/null @@ -1,28 +0,0 @@ -# maxiter-scan (inner ge/cal maxiter × recompute_scale) — LEVEL m100 (ge=cal=100). -# Aligned with softplus-floor-off.yaml: same fixed fitting block VERBATIM, -# but recompute_scale is now SWEPT [true, false] and the inner ge/cal_kwargs -# maxiter is pinned to 100 (coupled: ge and cal move together). Outer `maxiter` -# (the block cap) stays at 50 — this scan is the INNER optimizer cap only. -# Grid: recompute_scale[2] × fusionreg[4] × reps[2] = 16 fits. -# Companion levels: maxiter-scan-m1.yaml (ge=cal=1), maxiter-scan-m10.yaml -# (ge=cal=10). Combine the three fit_collection.pkl for the full 48-fit study. -sweep: - fusionreg: [0.0, 4.0e-5, 1.6e-4, 6.4e-4] # same 4-point axis as #277 - recompute_scale: [true, false] # SWEPT (was fixed=false) -fixed: - output_floor: null # OFF control — matched in-harness baseline - share_alpha: true # #274 baseline (shared-α) condition - warmstart: false - ge_type: Sigmoid - l2reg: 0.0 - beta0_ridge: 0.0 - scale_fusion_by_n: false - alpha_init: 6.0 - beta0_init: {Omicron_BA1: 0.0, Delta: 0.0, Omicron_BA2: 0.0} - beta_clip_range: [-10, 10] - loss_kwargs: {δ: 1.0} - maxiter: 50 # outer block cap (unchanged) - tol: 1.0e-5 # outer tol - ge_kwargs: {tol: 1.0e-4, maxiter: 100, maxls: 10, jit: true, verbose: false} - cal_kwargs: {tol: 1.0e-4, maxiter: 100, maxls: 10, jit: true, verbose: false} -replicates: [1, 2] diff --git a/experiments/convergence-lab/grids/smoke.yaml b/experiments/convergence-lab/grids/smoke.yaml deleted file mode 100644 index 3a3aaed4..00000000 --- a/experiments/convergence-lab/grids/smoke.yaml +++ /dev/null @@ -1,15 +0,0 @@ -# Smoke grid — the harness's worked example and "did I break it" check. -# 2 l2reg × 2 replicates = 4 fits, fit in parallel (capped iters). -sweep: - l2reg: [0.0, 3.0e-4] # below and at the L2 knee -fixed: - warmstart: true - recompute_scale: false # the fixed-scale convergence fix (#246) - share_alpha: true - fusionreg: 0.0 - ge_type: Sigmoid - maxiter: 25 # block coordinate-descent iters (capped for speed) - tol: 1.0e-6 - ge_kwargs: {tol: 1.0e-4, maxiter: 20, maxls: 40, jit: true} - cal_kwargs: {tol: 1.0e-4, maxiter: 20, maxls: 40, jit: true} -replicates: [1, 2] diff --git a/experiments/convergence-lab/grids/softplus-floor-freealpha.yaml b/experiments/convergence-lab/grids/softplus-floor-freealpha.yaml deleted file mode 100644 index 542219d3..00000000 --- a/experiments/convergence-lab/grids/softplus-floor-freealpha.yaml +++ /dev/null @@ -1,26 +0,0 @@ -# softplus-floor × FREE-α, floor ON (#277) — the missing free-α arm of the 2×2. -# #274 (Free-Alpha) is defined by share_alpha=false; this grid runs the softplus -# floor (lower_bound=-3.5) ON in that regime so the softplus × share_alpha -# interaction is measured in ONE harness rather than quoted cross-PR. -# fusionreg [0, 4e-5, 1.6e-4, 6.4e-4] × 2 reps = 8 fits. -sweep: - fusionreg: [0.0, 4.0e-5, 1.6e-4, 6.4e-4] # same 4-point axis as #274 -fixed: - output_floor: -3.5 # THE new lever — on for every cell - share_alpha: false # #274 free-α condition (per-condition α) - # --- #274 free-alpha.yaml fixed block, VERBATIM (= PR 270 recompute-false) --- - warmstart: false - recompute_scale: false - ge_type: Sigmoid - l2reg: 0.0 - beta0_ridge: 0.0 - scale_fusion_by_n: false - alpha_init: 6.0 - beta0_init: {Omicron_BA1: 0.0, Delta: 0.0, Omicron_BA2: 0.0} - beta_clip_range: [-10, 10] - loss_kwargs: {δ: 1.0} - maxiter: 50 # outer block cap - tol: 1.0e-5 # outer tol (loosened from PR 270's 1e-6) - ge_kwargs: {tol: 1.0e-4, maxiter: 10, maxls: 10, jit: true, verbose: false} - cal_kwargs: {tol: 1.0e-4, maxiter: 10, maxls: 10, jit: true, verbose: false} -replicates: [1, 2] diff --git a/experiments/convergence-lab/grids/softplus-floor-off-freealpha.yaml b/experiments/convergence-lab/grids/softplus-floor-off-freealpha.yaml deleted file mode 100644 index 72f6d103..00000000 --- a/experiments/convergence-lab/grids/softplus-floor-off-freealpha.yaml +++ /dev/null @@ -1,26 +0,0 @@ -# softplus-floor × FREE-α, floor OFF (#277) — the matched OFF control for the -# free-α arm of the 2×2. share_alpha=false (the #274 regime) with output_floor -# null, so the softplus toggle is a controlled comparison against -# softplus-floor-freealpha.yaml within ONE harness. -# fusionreg [0, 4e-5, 1.6e-4, 6.4e-4] × 2 reps = 8 fits. -sweep: - fusionreg: [0.0, 4.0e-5, 1.6e-4, 6.4e-4] # same 4-point axis as #274 -fixed: - output_floor: null # OFF control — matched in-harness baseline - share_alpha: false # #274 free-α condition (per-condition α) - # --- #274 free-alpha.yaml fixed block, VERBATIM (= PR 270 recompute-false) --- - warmstart: false - recompute_scale: false - ge_type: Sigmoid - l2reg: 0.0 - beta0_ridge: 0.0 - scale_fusion_by_n: false - alpha_init: 6.0 - beta0_init: {Omicron_BA1: 0.0, Delta: 0.0, Omicron_BA2: 0.0} - beta_clip_range: [-10, 10] - loss_kwargs: {δ: 1.0} - maxiter: 50 # outer block cap - tol: 1.0e-5 # outer tol (loosened from PR 270's 1e-6) - ge_kwargs: {tol: 1.0e-4, maxiter: 10, maxls: 10, jit: true, verbose: false} - cal_kwargs: {tol: 1.0e-4, maxiter: 10, maxls: 10, jit: true, verbose: false} -replicates: [1, 2] diff --git a/experiments/convergence-lab/grids/softplus-floor-off.yaml b/experiments/convergence-lab/grids/softplus-floor-off.yaml deleted file mode 100644 index e0dbe94b..00000000 --- a/experiments/convergence-lab/grids/softplus-floor-off.yaml +++ /dev/null @@ -1,25 +0,0 @@ -# softplus-floor OFF control (#277) — does an explicit softplus output floor -# (lower_bound=-3.5) deepen Delta's under-fit low tail the way freeing α did in -# #274? Floor OFF (output_floor=null) — the MATCHED in-harness baseline for the ON run, -# so the A/B is a controlled toggle in ONE harness. fusionreg [0, 4e-5, 1.6e-4, 6.4e-4] × 2 reps = 8 fits. -sweep: - fusionreg: [0.0, 4.0e-5, 1.6e-4, 6.4e-4] # same 4-point axis as #274 -fixed: - output_floor: null # OFF control — matched in-harness baseline - share_alpha: true # #274 baseline (shared-α) condition - # --- #274 free-alpha.yaml fixed block, VERBATIM (= PR 270 recompute-false) --- - warmstart: false - recompute_scale: false - ge_type: Sigmoid - l2reg: 0.0 - beta0_ridge: 0.0 - scale_fusion_by_n: false - alpha_init: 6.0 - beta0_init: {Omicron_BA1: 0.0, Delta: 0.0, Omicron_BA2: 0.0} - beta_clip_range: [-10, 10] - loss_kwargs: {δ: 1.0} - maxiter: 50 # outer block cap - tol: 1.0e-5 # outer tol (loosened from PR 270's 1e-6) - ge_kwargs: {tol: 1.0e-4, maxiter: 10, maxls: 10, jit: true, verbose: false} - cal_kwargs: {tol: 1.0e-4, maxiter: 10, maxls: 10, jit: true, verbose: false} -replicates: [1, 2] diff --git a/experiments/convergence-lab/grids/softplus-floor.yaml b/experiments/convergence-lab/grids/softplus-floor.yaml deleted file mode 100644 index 987e3c11..00000000 --- a/experiments/convergence-lab/grids/softplus-floor.yaml +++ /dev/null @@ -1,25 +0,0 @@ -# softplus-floor grid (#277) — does an explicit softplus output floor -# (lower_bound=-3.5) deepen Delta's under-fit low tail the way freeing α did in -# #274? Floor fixed ON for all cells; A/B is cross-PR vs #274's shared-α column. -# output_floor [-3.5] × fusionreg [0, 4e-5, 1.6e-4, 6.4e-4] × 2 reps = 8 fits. -sweep: - fusionreg: [0.0, 4.0e-5, 1.6e-4, 6.4e-4] # same 4-point axis as #274 -fixed: - output_floor: -3.5 # THE new lever — on for every cell - share_alpha: true # #274 baseline (shared-α) condition - # --- #274 free-alpha.yaml fixed block, VERBATIM (= PR 270 recompute-false) --- - warmstart: false - recompute_scale: false - ge_type: Sigmoid - l2reg: 0.0 - beta0_ridge: 0.0 - scale_fusion_by_n: false - alpha_init: 6.0 - beta0_init: {Omicron_BA1: 0.0, Delta: 0.0, Omicron_BA2: 0.0} - beta_clip_range: [-10, 10] - loss_kwargs: {δ: 1.0} - maxiter: 50 # outer block cap - tol: 1.0e-5 # outer tol (loosened from PR 270's 1e-6) - ge_kwargs: {tol: 1.0e-4, maxiter: 10, maxls: 10, jit: true, verbose: false} - cal_kwargs: {tol: 1.0e-4, maxiter: 10, maxls: 10, jit: true, verbose: false} -replicates: [1, 2] diff --git a/experiments/convergence-lab/harness.py b/experiments/convergence-lab/harness.py deleted file mode 100644 index 6d77ae72..00000000 --- a/experiments/convergence-lab/harness.py +++ /dev/null @@ -1,336 +0,0 @@ -r"""Convergence-lab harness (#253): config-driven parallel model fitting. - -Reads a grid YAML (a ``sweep:`` Cartesian product + optional ``fixed:`` -overrides + optional ``replicates:``), fits every cell in parallel via -``multidms.model_collection.fit_models`` (``n_processes`` spawn workers), -and pickles the raw fit-collection DataFrame to -``results//fit_collection.pkl``. - -The output is a *true* ``fit_collection.pkl`` — the same -``stack_fit_models`` frame the scv2-spike pipeline writes — so the marimo -dashboard discovers it directly (it ``rglob``s ``fit_collection.pkl`` below -cwd) and all downstream analysis, including replicate-shift correlation, is -done on the fly from a ``ModelCollection`` built over the frame. The harness -only fits; it computes no derived metrics. - -CPU parallelism uses ``multiprocessing`` ``spawn``, which re-imports this -module in every worker. That is safe here because the fitting is reachable -only under the ``if __name__ == "__main__":`` guard at the bottom — see the -warning on ``multidms.model_collection.fit_models``. ``--n-processes`` -selects worker count (default: all but one core, capped at the grid size). - -The runner owns the constant scv2-spike data; configs carry only swept -knobs and fixed overrides. See ``README.md`` in this directory. - -Usage:: - - pixi run python experiments/convergence-lab/harness.py \\ - --config grids/smoke.yaml --cache smoke - pixi run python experiments/convergence-lab/harness.py \\ - --config grids/smoke.yaml --cache smoke --n-processes 4 -""" - -from __future__ import annotations - -import argparse # noqa: F401 — used in later tasks (CLI) -import os - -# Pin per-process CPU threading BEFORE JAX/XLA imports (multidms imports JAX at -# module load; these env vars are only read at XLA init). Harmless for the -# sequential harness; required if a future config restores parallelism. -os.environ.setdefault("XLA_FLAGS", "--xla_cpu_multi_thread_eigen=false") -os.environ.setdefault("OMP_NUM_THREADS", "1") -os.environ.setdefault("OPENBLAS_NUM_THREADS", "1") -os.environ.setdefault("MKL_NUM_THREADS", "1") -os.environ.setdefault("JAX_PLATFORM_NAME", "cpu") - -import pickle # noqa: E402 -import time # noqa: E402 -import warnings # noqa: E402 -from pathlib import Path # noqa: E402 - -import pandas as pd # noqa: E402 -import yaml # noqa: E402 - -warnings.filterwarnings("ignore") - -import multidms # noqa: E402 -from multidms.model_collection import fit_models # noqa: E402 -from multidms.utils import explode_params_dict # noqa: E402 - -# --- Constant data the runner owns (resolved relative to this file) ---------- -HARNESS_DIR = Path(os.path.abspath(__file__)).parent -DATA_CSV = ( - HARNESS_DIR.parent - / "scv2-spike" - / "results-prod-235-times-seen-threshold" - / "training_functional_scores.csv" -) -RESULTS_DIR = HARNESS_DIR / "results" -REF = "Omicron_BA1" - -# Exact fit_one_model signature (model_collection.py:61) minus `dataset` -# (runner-supplied) and `verbose`. The allowlist for config validation. -VALID_KWARGS = frozenset( - { - "ge_type", - "l2reg", - "fusionreg", - "beta0_ridge", - "scale_fusion_by_n", - "output_floor", - "output_floor_hinge", - "loss_type", - "maxiter", - "tol", - "warmstart", - "recompute_scale", - "beta0_init", - "beta_init", - "alpha_init", - "share_alpha", - "beta_clip_range", - "ge_kwargs", - "cal_kwargs", - "loss_kwargs", - } -) - - -def load_config(path) -> dict: - """Load and validate a grid YAML. - - The config has two kwargs sections — ``sweep`` (each value a list, crossed - in a Cartesian product) and ``fixed`` (each value a scalar, applied to every - fit) — plus an optional harness-level ``replicates`` list (default - ``[1, 2]``). Every key in ``sweep``/``fixed`` must be a valid - ``fit_one_model`` kwarg (``VALID_KWARGS``); ``fit_one_model``'s ``**kwargs`` - would otherwise swallow typos silently, so we reject them here. - - Args: - path: Path to the YAML config. - - Returns: - ``{"sweep": dict, "fixed": dict, "replicates": list[int]}``. - - Raises: - ValueError: If ``sweep`` is missing/empty, or any ``sweep``/``fixed`` - key is not in ``VALID_KWARGS`` (the message names the bad key). - """ - with open(path) as fh: - raw = yaml.safe_load(fh) or {} - - sweep = raw.get("sweep") or {} - fixed = raw.get("fixed") or {} - replicates = raw.get("replicates", [1, 2]) - - if not sweep: - raise ValueError(f"config {path}: must define a non-empty 'sweep:' section") - - for section_name, section in (("sweep", sweep), ("fixed", fixed)): - for key in section: - if key not in VALID_KWARGS: - raise ValueError( - f"config {path}: '{key}' in '{section_name}:' is not a valid " - f"fit_one_model kwarg. Valid keys: {sorted(VALID_KWARGS)}" - ) - - return {"sweep": dict(sweep), "fixed": dict(fixed), "replicates": list(replicates)} - - -def explode_grid(config: dict) -> list[dict]: - """Cartesian-product the sweep × replicates, merging fixed into each cell. - - Args: - config: The dict returned by :func:`load_config`. - - Returns: - A list of kwarg dicts, one per fit. Each carries every ``fixed`` key and - an integer ``replicate`` key (the runner maps it to the rep's Data). - """ - sweep = dict(config["sweep"]) - sweep["replicate"] = list(config["replicates"]) - exploded = explode_params_dict(sweep) - for cell in exploded: - cell.update(config["fixed"]) - return exploded - - -def load_rep_data() -> dict[str, multidms.Data]: - """Build one ``multidms.Data`` per replicate from the prod spike CSV. - - Mirrors the pipeline ``fit_models.ipynb``: aggregate ``func_score`` by - ``(condition, aa_substitutions)`` with ``.mean()`` within each replicate, - then construct a ``multidms.Data`` with the gap-inclusive alphabet. - - Returns: - Mapping ``"rep_" -> multidms.Data`` (reference = ``Omicron_BA1``). - """ - raw = pd.read_csv(DATA_CSV).fillna({"aa_substitutions": ""}) - rep_data: dict[str, multidms.Data] = {} - for rep in sorted(raw["replicate"].unique()): - df_rep = raw[raw["replicate"] == rep] - df_agg = ( - df_rep.groupby(["condition", "aa_substitutions"], dropna=False) - .agg({"func_score": "mean"}) - .reset_index() - ) - rep_data[f"rep_{rep}"] = multidms.Data( - df_agg, - alphabet=multidms.AAS_WITHSTOP_WITHGAP, - reference=REF, - assert_site_integrity=False, - name=f"rep_{rep}", - verbose=False, - ) - return rep_data - - -def build_params(exploded: list[dict], rep_data: dict) -> dict: - """Turn exploded grid cells into a ``fit_models`` ``params`` dict. - - ``fit_models`` takes a ``params`` dict whose every value is a list and - crosses them internally. Each exploded cell already carries an integer - ``replicate`` (which selects a Data) plus its fit kwargs; this maps the - replicate to the rep's Data object and collects each kwarg's distinct - values into a list, so ``fit_models`` reproduces exactly the cells in - ``exploded``. - - Args: - exploded: Output of :func:`explode_grid` (sweep × replicates × fixed). - rep_data: Output of :func:`load_rep_data` (``"rep_" -> Data``). - - Returns: - A ``params`` dict for :func:`multidms.model_collection.fit_models`: - ``dataset`` is the list of Data objects, every other key the sorted - distinct values seen across cells (singletons stay singleton lists). - """ - reps = sorted({c["replicate"] for c in exploded}) - datasets = [rep_data[f"rep_{rep}"] for rep in reps] - params: dict = {"dataset": datasets} - for cell in exploded: - for key, val in cell.items(): - if key == "replicate": - continue - params.setdefault(key, []) - if val not in params[key]: - params[key].append(val) - return params - - -def run_fits(params: dict, n_processes: int) -> pd.DataFrame: - """Fit the whole grid in parallel via ``fit_models`` and return the frame. - - Delegates to :func:`multidms.model_collection.fit_models`, which explodes - ``params`` (``dataset`` × swept kwargs), fits each combination in a - ``spawn`` worker pool (``n_processes`` workers), and stacks the results. - Failures are tolerated (``failures="tolerate"``): a failed fit is dropped - and the grid continues, so one bad cell does not poison the run; if every - fit fails, ``fit_models`` raises ``ModelCollectionFitError``. - - This MUST be reached only under the module's ``if __name__ == - "__main__":`` guard — ``spawn`` re-imports this module in every worker. - - Args: - params: Output of :func:`build_params`. - n_processes: Worker count passed to ``fit_models`` (>= 1; 1 runs - in-process with no pool). - - Returns: - The raw ``stack_fit_models`` fit-collection DataFrame (one row per - fit, the ``model`` column plus the fit kwargs and ``dataset_name``). - This is a *true* ``fit_collection.pkl`` — no derived metrics. - """ - from math import prod - - n = len(params["dataset"]) * prod( - len(v) for k, v in params.items() if k != "dataset" - ) - print(f"[harness] fitting {n} models with n_processes={n_processes} …", flush=True) - t0 = time.time() - n_fit, n_failed, fit_df = fit_models( - params, n_processes=n_processes, failures="tolerate" - ) - print( - f"[harness] wall {time.time() - t0:.1f}s — {n_fit} fit, {n_failed} failed", - flush=True, - ) - return fit_df - - -def default_n_processes(grid_size: int) -> int: - """Worker count to use when ``--n-processes`` is not given. - - All but one CPU core (leave one for the OS / parent), never more than the - number of fits (extra workers would idle), and at least 1. - - Args: - grid_size: Number of fits in the exploded grid. - - Returns: - The default worker count. - """ - cores = os.cpu_count() or 1 - return max(1, min(grid_size, cores - 1)) - - -def run(config_path, cache: str, n_processes: int | None = None) -> Path: - """Load config, fit the grid in parallel, and pickle the fit collection. - - Writes ``results//fit_collection.pkl`` — the raw fit-collection - DataFrame, the same schema the scv2-spike pipeline writes. The marimo - dashboard discovers it by name; build a ``ModelCollection`` over it for - any downstream analysis (correlation, basin diagnostics, plots). - - Args: - config_path: Path to the grid YAML. - cache: Subdirectory name under ``results/`` to hold the pickle. - n_processes: Worker count; ``None`` → :func:`default_n_processes`. - - Returns: - Path to the written ``fit_collection.pkl``. - """ - config = load_config(config_path) - rep_data = load_rep_data() - exploded = explode_grid(config) - if n_processes is None: - n_processes = default_n_processes(len(exploded)) - params = build_params(exploded, rep_data) - fit_df = run_fits(params, n_processes) - - out_dir = RESULTS_DIR / cache - out_dir.mkdir(parents=True, exist_ok=True) - out_path = out_dir / "fit_collection.pkl" - with open(out_path, "wb") as fh: - pickle.dump(fit_df, fh) - print(f"[harness] wrote {out_path} ({len(fit_df)} fits)", flush=True) - return out_path - - -def main() -> None: - """CLI: ``--config --cache [--n-processes N]``.""" - ap = argparse.ArgumentParser(description=__doc__) - ap.add_argument("--config", required=True, help="path to grid YAML") - ap.add_argument( - "--cache", - required=True, - help="output subdir → results//fit_collection.pkl", - ) - ap.add_argument( - "--n-processes", - type=int, - default=None, - help="parallel worker count (default: all but one core, capped at " - "the grid size)", - ) - args = ap.parse_args() - # Resolve a relative --config against the harness dir (so `grids/smoke.yaml` - # works regardless of cwd), matching how DATA_CSV is resolved. - cfg_path = Path(args.config) - if not cfg_path.is_absolute() and not cfg_path.exists(): - cfg_path = HARNESS_DIR / args.config - run(cfg_path, args.cache, args.n_processes) - - -if __name__ == "__main__": - main() diff --git a/experiments/loss-normalization/.gitignore b/experiments/loss-normalization/.gitignore deleted file mode 100644 index 1ab70ba5..00000000 --- a/experiments/loss-normalization/.gitignore +++ /dev/null @@ -1,2 +0,0 @@ -results/ -results_test/ diff --git a/experiments/loss-normalization/Snakefile b/experiments/loss-normalization/Snakefile deleted file mode 100644 index f18bd945..00000000 --- a/experiments/loss-normalization/Snakefile +++ /dev/null @@ -1,84 +0,0 @@ -"""Snakemake workflow for the loss normalization experiment. - -Two notebooks executed via papermill: - fit_models → evaluate - -Run with: - snakemake -s experiments/loss-normalization/Snakefile -j4 - snakemake -s experiments/loss-normalization/Snakefile --config profile=test -j4 -""" - -import os -import yaml - -# Resolve paths relative to the Snakefile location -EXP_DIR = os.path.dirname(workflow.snakefile) -NOTEBOOKS = os.path.join(EXP_DIR, "notebooks") - -# Config: use test profile if --config profile=test is passed -profile = config.get("profile", None) -if profile == "test": - CONFIG_PATH = os.path.join(EXP_DIR, "config", "config_test.yaml") -else: - CONFIG_PATH = os.path.join(EXP_DIR, "config", "config.yaml") - -# Read output_dir from the YAML config -with open(CONFIG_PATH) as _f: - _cfg = yaml.safe_load(_f) -_yaml_output_dir = _cfg.get("experiment", {}).get("output_dir", "results") -output_dir = config.get("output_dir", _yaml_output_dir) -RESULTS = os.path.join(EXP_DIR, output_dir) - -# Relative config path for papermill parameter injection -CONFIG_PATH_REL = os.path.relpath(CONFIG_PATH, EXP_DIR) - - -rule all: - input: - os.path.join(RESULTS, "evaluate.ipynb"), - os.path.join(RESULTS, "fit_summary.csv"), - os.path.join(RESULTS, "sparsity_vs_fusionreg.pdf"), - os.path.join(RESULTS, "correlation_vs_fusionreg.pdf"), - - -rule fit_models: - input: - config=CONFIG_PATH, - notebook=os.path.join(NOTEBOOKS, "fit_models.ipynb"), - output: - fit_collection=os.path.join(RESULTS, "fit_collection.pkl"), - executed_notebook=os.path.join(RESULTS, "fit_models.ipynb"), - params: - config_path=CONFIG_PATH_REL, - output_dir=output_dir, - shell: - """ - papermill {input.notebook} {output.executed_notebook} \ - -p config_path {params.config_path} \ - -p output_dir {params.output_dir} \ - --cwd {EXP_DIR} \ - && jupyter trust {output.executed_notebook} - """ - - -rule evaluate: - input: - config=CONFIG_PATH, - fit_collection=os.path.join(RESULTS, "fit_collection.pkl"), - notebook=os.path.join(NOTEBOOKS, "evaluate.ipynb"), - output: - fit_summary=os.path.join(RESULTS, "fit_summary.csv"), - sparsity_plot=os.path.join(RESULTS, "sparsity_vs_fusionreg.pdf"), - corr_plot=os.path.join(RESULTS, "correlation_vs_fusionreg.pdf"), - executed_notebook=os.path.join(RESULTS, "evaluate.ipynb"), - params: - config_path=CONFIG_PATH_REL, - output_dir=output_dir, - shell: - """ - papermill {input.notebook} {output.executed_notebook} \ - -p config_path {params.config_path} \ - -p output_dir {params.output_dir} \ - --cwd {EXP_DIR} \ - && jupyter trust {output.executed_notebook} - """ diff --git a/experiments/loss-normalization/config/config.yaml b/experiments/loss-normalization/config/config.yaml deleted file mode 100644 index 9516b6a4..00000000 --- a/experiments/loss-normalization/config/config.yaml +++ /dev/null @@ -1,33 +0,0 @@ -# Configuration for the loss normalization experiment. -# Sweeps fusionreg × l2reg to validate .mean() loss normalization -# against V0.4.0 hyperparameter anchors. - -seed: 2 -train_frac: 0.8 - -experiment: - data: - func_scores: "data/simulated_func_scores.csv" - true_effects: "data/simulated_muteffects.csv" - - fitting: - maxiter: 15 - tol: 1.0e-6 - ge_type: "Sigmoid" - warmstart: false - beta0_ridge: 0.0 - delta: 1.0 - share_alpha: true - beta0_init: {h1: 5.0, h2: 0.0} - alpha_init: 6.0 - beta_clip_range: [-10, 10] - loss_kwargs: {"δ": 1.0} - ge_kwargs: {tol: 1.0e-5, maxiter: 1000, maxls: 40, jit: true, verbose: false} - cal_kwargs: {tol: 1.0e-4, maxiter: 1000, maxls: 40, jit: true, verbose: false} - n_processes: null # null = auto - - # Sweep dimensions (V0.4.0-anchored) - fusionreg_values: [0.0, 5.0e-6, 1.0e-5, 2.0e-5, 4.0e-5, 8.0e-5, 1.6e-4, 3.2e-4, 6.4e-4] - l2reg_values: [0.0, 1.0e-8, 1.0e-7, 1.0e-6] - - output_dir: results diff --git a/experiments/loss-normalization/config/config_test.yaml b/experiments/loss-normalization/config/config_test.yaml deleted file mode 100644 index aac0364b..00000000 --- a/experiments/loss-normalization/config/config_test.yaml +++ /dev/null @@ -1,33 +0,0 @@ -# Test configuration for the loss normalization experiment. -# Reduced grid for fast iteration. - -seed: 2 -train_frac: 0.8 - -experiment: - data: - func_scores: "data/simulated_func_scores.csv" - true_effects: "data/simulated_muteffects.csv" - subsample_frac: 0.05 # subsample to ~5% for fast test runs - - fitting: - maxiter: 2 - tol: 1.0e-6 - ge_type: "Sigmoid" - warmstart: false - beta0_ridge: 0.0 - delta: 1.0 - share_alpha: true - beta0_init: {h1: 5.0, h2: 0.0} - alpha_init: 6.0 - beta_clip_range: [-10, 10] - loss_kwargs: {"δ": 1.0} - ge_kwargs: {tol: 1.0e-4, maxiter: 2, maxls: 40, jit: true, verbose: false} - cal_kwargs: {tol: 1.0e-4, maxiter: 2, maxls: 40, jit: true, verbose: false} - n_processes: null - - # Minimal grid for testing - fusionreg_values: [0.0, 4.0e-5] - l2reg_values: [0.0, 1.0e-7] - - output_dir: results_test diff --git a/experiments/loss-normalization/data/.gitkeep b/experiments/loss-normalization/data/.gitkeep deleted file mode 100644 index e69de29b..00000000 diff --git a/experiments/loss-normalization/notebooks/_common.py b/experiments/loss-normalization/notebooks/_common.py deleted file mode 100644 index 9f887ca1..00000000 --- a/experiments/loss-normalization/notebooks/_common.py +++ /dev/null @@ -1,59 +0,0 @@ -"""Shared utilities for the loss normalization experiment pipeline.""" - -import yaml - - -def load_config(config_path): - """Load a pipeline configuration YAML file. - - Parameters - ---------- - config_path : str - Path to the YAML config file. - - Returns - ------- - dict - Configuration dictionary. - """ - with open(config_path) as f: - config = yaml.safe_load(f) - return config - - -def build_fit_params(fit_config, datasets): - """Build a fitting parameter dict that sweeps fusionreg AND l2reg. - - Unlike the standard simulation pipeline's ``build_fit_params`` which - fixes ``l2reg`` to a single value, this version maps - ``l2reg_values`` to a list for the ``l2reg`` sweep dimension. - - Parameters - ---------- - fit_config : dict - The ``fitting`` subsection of the experiment config. - datasets : list - List of ``multidms.Data`` objects to fit. - - Returns - ------- - dict - Ready to pass to ``multidms.model_collection.fit_models()``. - """ - return { - "maxiter": [fit_config["maxiter"]], - "tol": [fit_config["tol"]], - "fusionreg": fit_config["fusionreg_values"], - "l2reg": fit_config["l2reg_values"], - "beta0_ridge": [fit_config["beta0_ridge"]], - "ge_type": [fit_config["ge_type"]], - "ge_kwargs": [fit_config["ge_kwargs"]], - "cal_kwargs": [fit_config["cal_kwargs"]], - "loss_kwargs": [fit_config["loss_kwargs"]], - "warmstart": [fit_config["warmstart"]], - "beta0_init": [fit_config["beta0_init"]], - "alpha_init": [fit_config["alpha_init"]], - "share_alpha": [fit_config.get("share_alpha", True)], - "beta_clip_range": [tuple(fit_config["beta_clip_range"])], - "dataset": datasets, - } diff --git a/experiments/loss-normalization/notebooks/evaluate.ipynb b/experiments/loss-normalization/notebooks/evaluate.ipynb deleted file mode 100644 index 490be946..00000000 --- a/experiments/loss-normalization/notebooks/evaluate.ipynb +++ /dev/null @@ -1,121 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "01c71fc3", - "source": "# Loss Normalization Experiment: Evaluation\n\nEvaluate the fitted models from the `fusionreg × l2reg` grid against\nground-truth mutation effects.\n\n**Key outputs:**\n1. Sparsity vs fusionreg (one curve per l2reg)\n2. β correlation with ground truth vs fusionreg (one curve per l2reg)\n3. `fit_summary.csv` — per-model summary table", - "metadata": {} - }, - { - "cell_type": "code", - "id": "d1f2f4d6", - "source": "import warnings\n\nwarnings.filterwarnings(\"ignore\")\n\nimport os\nimport pickle\nimport sys\n\nsys.path.insert(0, \"notebooks\")\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport multidms\n\nfrom _common import load_config", - "metadata": {}, - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "code", - "id": "e27e4e9a", - "source": "config_path = \"config/config.yaml\"\noutput_dir = None", - "metadata": { - "tags": [ - "parameters" - ] - }, - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "code", - "id": "daa2f78e", - "source": "config = load_config(config_path)\nexp = config[\"experiment\"]\nif output_dir is None:\n output_dir = exp[\"output_dir\"]\n\ntrue_effects_path = exp[\"data\"][\"true_effects\"]\nprint(f\"Output directory: {output_dir}\")", - "metadata": {}, - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "markdown", - "id": "060e7434", - "source": "## Load data", - "metadata": {} - }, - { - "cell_type": "code", - "id": "ee35107b", - "source": "fit_collection_df = pickle.load(\n open(os.path.join(output_dir, \"fit_collection.pkl\"), \"rb\")\n)\ntrue_effects = pd.read_csv(true_effects_path)\nprint(f\"Loaded {len(fit_collection_df)} fitted models\")\nprint(f\"Ground truth: {len(true_effects)} mutations\")", - "metadata": {}, - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "markdown", - "id": "a67e24c0", - "source": "## Build per-model summary\n\nFor each fitted model, compute:\n- Shift sparsity (fraction of zero shifts)\n- Pearson correlation of β with ground truth\n- Pearson correlation of shifts with ground truth", - "metadata": {} - }, - { - "cell_type": "code", - "id": "3853d485", - "source": "def compute_model_summary(row, true_effects):\n \"\"\"Compute accuracy metrics for a single fitted model.\"\"\"\n model = row[\"model\"]\n muts_df = model.get_mutations_df().reset_index()\n\n # Merge with ground truth — shared columns (beta_h1, beta_h2) get _x/_y suffixes\n merged = muts_df.merge(true_effects, on=\"mutation\", how=\"inner\", suffixes=(\"_fit\", \"_true\"))\n\n result = {\n \"fusionreg\": row[\"fusionreg\"],\n \"l2reg\": row[\"l2reg\"],\n \"dataset_name\": row[\"dataset_name\"],\n \"library\": row.get(\"library\", \"\"),\n \"measurement_type\": row.get(\"measurement_type\", \"\"),\n \"fit_time\": row.get(\"fit_time\", np.nan),\n }\n\n # β correlation (reference condition): fitted beta_h1 vs true beta_h1\n fit_beta = \"beta_h1_fit\" if \"beta_h1_fit\" in merged.columns else \"beta_h1\"\n true_beta = \"beta_h1_true\" if \"beta_h1_true\" in merged.columns else \"beta_h1\"\n if fit_beta in merged.columns and true_beta in merged.columns:\n mask = merged[fit_beta].notna() & merged[true_beta].notna()\n if mask.sum() > 2:\n result[\"beta_corr\"] = merged.loc[mask, fit_beta].corr(\n merged.loc[mask, true_beta]\n )\n\n # Shift correlation: fitted shift_h2 vs true shift\n fit_shift = \"shift_h2\" if \"shift_h2\" in merged.columns else None\n true_shift = \"shift\" if \"shift\" in merged.columns else None\n if fit_shift and true_shift and fit_shift in merged.columns and true_shift in merged.columns:\n mask = merged[fit_shift].notna() & merged[true_shift].notna()\n if mask.sum() > 2:\n result[\"shift_corr\"] = merged.loc[mask, fit_shift].corr(\n merged.loc[mask, true_shift]\n )\n # Sparsity: fraction of exactly zero shifts\n result[\"shift_sparsity\"] = (merged[fit_shift] == 0).mean()\n\n # True sparsity\n if \"shifted_site\" in merged.columns:\n result[\"true_sparsity\"] = 1.0 - merged[\"shifted_site\"].mean()\n elif \"bundle_mut\" in merged.columns:\n result[\"true_sparsity\"] = 1.0 - merged[\"bundle_mut\"].mean()\n\n return result\n\n\nsummaries = []\nfor _, row in fit_collection_df.iterrows():\n summaries.append(compute_model_summary(row, true_effects))\n\nsummary_df = pd.DataFrame(summaries)\nsummary_df.to_csv(os.path.join(output_dir, \"fit_summary.csv\"), index=False)\nprint(f\"Summary: {len(summary_df)} rows\")\nprint(f\"Columns: {list(summary_df.columns)}\")\nsummary_df.head(10)", - "metadata": {}, - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "markdown", - "id": "33522de3", - "source": "## 1. Sparsity vs fusionreg\n\nOne curve per `l2reg` value. Shows fraction of zero-shift parameters as a\nfunction of fusion regularization strength, compared against true sparsity.", - "metadata": {} - }, - { - "cell_type": "code", - "id": "d1ca9811", - "source": "# Average across datasets (libraries × measurement types) for each (fusionreg, l2reg)\n# Only aggregate columns that exist in the summary\nagg_dict = {}\nif \"shift_sparsity\" in summary_df.columns:\n agg_dict[\"shift_sparsity\"] = (\"shift_sparsity\", \"mean\")\nif \"true_sparsity\" in summary_df.columns:\n agg_dict[\"true_sparsity\"] = (\"true_sparsity\", \"mean\")\n\nif agg_dict:\n agg = summary_df.groupby([\"fusionreg\", \"l2reg\"]).agg(**agg_dict).reset_index()\n\n fig, ax = plt.subplots(figsize=(8, 5))\n if \"shift_sparsity\" in agg.columns:\n for l2val, grp in agg.groupby(\"l2reg\"):\n grp_sorted = grp.sort_values(\"fusionreg\")\n ax.plot(\n grp_sorted[\"fusionreg\"],\n grp_sorted[\"shift_sparsity\"],\n \"o-\",\n label=f\"l2reg={l2val:.1e}\",\n )\n\n # True sparsity reference line\n if \"true_sparsity\" in agg.columns and agg[\"true_sparsity\"].notna().any():\n true_sp = agg[\"true_sparsity\"].mean()\n ax.axhline(true_sp, color=\"k\", linestyle=\"--\", alpha=0.5,\n label=f\"true sparsity={true_sp:.2f}\")\n\n ax.set_xlabel(\"fusionreg\")\n ax.set_ylabel(\"Shift sparsity (fraction zero)\")\n ax.set_title(\"Shift sparsity vs fusion regularization\")\n ax.legend()\n ax.set_xscale(\"symlog\", linthresh=1e-6)\n plt.tight_layout()\n plt.savefig(os.path.join(output_dir, \"sparsity_vs_fusionreg.pdf\"))\n plt.show()\nelse:\n print(\"No sparsity data available — creating empty plot\")\n fig, ax = plt.subplots(figsize=(8, 5))\n ax.set_title(\"Shift sparsity vs fusion regularization (no data)\")\n plt.savefig(os.path.join(output_dir, \"sparsity_vs_fusionreg.pdf\"))\n plt.show()", - "metadata": {}, - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "markdown", - "id": "f547acd2", - "source": "## 2. β correlation with ground truth vs fusionreg\n\nPearson correlation between inferred and true mutation effects, broken out by `l2reg`.", - "metadata": {} - }, - { - "cell_type": "code", - "id": "2f555451", - "source": "corr_agg_dict = {}\nif \"beta_corr\" in summary_df.columns:\n corr_agg_dict[\"beta_corr\"] = (\"beta_corr\", \"mean\")\nif \"shift_corr\" in summary_df.columns:\n corr_agg_dict[\"shift_corr\"] = (\"shift_corr\", \"mean\")\n\nif corr_agg_dict:\n agg_corr = summary_df.groupby([\"fusionreg\", \"l2reg\"]).agg(**corr_agg_dict).reset_index()\n\n fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n\n if \"beta_corr\" in agg_corr.columns:\n for l2val, grp in agg_corr.groupby(\"l2reg\"):\n grp_sorted = grp.sort_values(\"fusionreg\")\n axes[0].plot(\n grp_sorted[\"fusionreg\"],\n grp_sorted[\"beta_corr\"],\n \"o-\",\n label=f\"l2reg={l2val:.1e}\",\n )\n axes[0].set_xlabel(\"fusionreg\")\n axes[0].set_ylabel(\"Pearson r (β vs truth)\")\n axes[0].set_title(\"β accuracy vs fusion regularization\")\n axes[0].legend()\n axes[0].set_xscale(\"symlog\", linthresh=1e-6)\n\n if \"shift_corr\" in agg_corr.columns:\n for l2val, grp in agg_corr.groupby(\"l2reg\"):\n grp_sorted = grp.sort_values(\"fusionreg\")\n axes[1].plot(\n grp_sorted[\"fusionreg\"],\n grp_sorted[\"shift_corr\"],\n \"s-\",\n label=f\"l2reg={l2val:.1e}\",\n )\n axes[1].set_xlabel(\"fusionreg\")\n axes[1].set_ylabel(\"Pearson r (shift vs truth)\")\n axes[1].set_title(\"Shift accuracy vs fusion regularization\")\n axes[1].legend()\n axes[1].set_xscale(\"symlog\", linthresh=1e-6)\n\n plt.tight_layout()\n plt.savefig(os.path.join(output_dir, \"correlation_vs_fusionreg.pdf\"))\n plt.show()\nelse:\n print(\"No correlation data available — creating empty plot\")\n agg_corr = pd.DataFrame(columns=[\"fusionreg\", \"l2reg\"])\n fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n axes[0].set_title(\"β accuracy (no data)\")\n axes[1].set_title(\"Shift accuracy (no data)\")\n plt.savefig(os.path.join(output_dir, \"correlation_vs_fusionreg.pdf\"))\n plt.show()", - "metadata": {}, - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "markdown", - "id": "8f087751", - "source": "## 3. Best model per l2reg\n\nIdentify the best (fusionreg, l2reg) combination for each l2reg value.", - "metadata": {} - }, - { - "cell_type": "code", - "id": "572650d9", - "source": "if \"beta_corr\" in agg_corr.columns:\n best = agg_corr.loc[agg_corr.groupby(\"l2reg\")[\"beta_corr\"].idxmax()]\n print(\"Best β correlation per l2reg:\")\n print(best[[\"l2reg\", \"fusionreg\", \"beta_corr\"]].to_string(index=False))\n print()\n\n overall_best = agg_corr.loc[agg_corr[\"beta_corr\"].idxmax()]\n print(f\"Overall best: fusionreg={overall_best['fusionreg']:.2e}, \"\n f\"l2reg={overall_best['l2reg']:.2e}, \"\n f\"beta_corr={overall_best['beta_corr']:.4f}\")\nelse:\n print(\"No beta_corr data available\")", - "metadata": {}, - "execution_count": null, - "outputs": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "name": "python", - "version": "3.9.0" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} \ No newline at end of file diff --git a/experiments/loss-normalization/notebooks/fit_models.ipynb b/experiments/loss-normalization/notebooks/fit_models.ipynb deleted file mode 100644 index a397cbea..00000000 --- a/experiments/loss-normalization/notebooks/fit_models.ipynb +++ /dev/null @@ -1,129 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "08ff80f2", - "source": "# Loss Normalization Experiment: Model Fitting\n\nFit `multidms` models across a 2D grid of `fusionreg × l2reg` to validate\nthe `.mean()` loss normalization against V0.4.0 hyperparameter anchors.\n\n**Outline**\n1. Load pre-generated simulation data\n2. Create `multidms.Data` objects\n3. Fit models across the 2D hyperparameter grid\n4. Save the fit collection", - "metadata": {} - }, - { - "cell_type": "code", - "id": "fb03faa9", - "source": "import warnings\n\nwarnings.filterwarnings(\"ignore\")\n\nimport os\nimport pickle\nimport sys\n\nsys.path.insert(0, \"notebooks\")\n\nimport pandas as pd\nimport multidms\nfrom multidms.model_collection import fit_models\nfrom multidms.utils import explode_params_dict\n\nfrom _common import load_config, build_fit_params", - "metadata": {}, - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "code", - "id": "0aa9eef4", - "source": "config_path = \"config/config.yaml\"\noutput_dir = None", - "metadata": { - "tags": [ - "parameters" - ] - }, - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "code", - "id": "8346a08b", - "source": "config = load_config(config_path)\nexp = config[\"experiment\"]\nfit_config = exp[\"fitting\"]\nif output_dir is None:\n output_dir = exp[\"output_dir\"]\n\nos.makedirs(output_dir, exist_ok=True)\nprint(f\"Output directory: {output_dir}\")\nprint(f\"fusionreg grid: {fit_config['fusionreg_values']}\")\nprint(f\"l2reg grid: {fit_config['l2reg_values']}\")", - "metadata": {}, - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "markdown", - "id": "a0d87129", - "source": "## Load simulated functional scores", - "metadata": {} - }, - { - "cell_type": "code", - "id": "07789c0f", - "source": "func_scores = pd.read_csv(exp[\"data\"][\"func_scores\"])\nfunc_scores[\"func_score_type\"] = pd.Categorical(\n func_scores[\"func_score_type\"],\n categories=[\"observed_phenotype\", \"loose_bottle\", \"tight_bottle\"],\n ordered=True,\n)\n\n# Optional subsampling for test runs (preserves wildtype rows)\nsubsample_frac = exp[\"data\"].get(\"subsample_frac\")\nif subsample_frac is not None and subsample_frac < 1.0:\n is_wt = func_scores[\"aa_substitutions\"].isna() | (\n func_scores[\"aa_substitutions\"].str.strip() == \"\"\n )\n wt_rows = func_scores[is_wt]\n mut_rows = func_scores[~is_wt].groupby(\n [\"homolog\", \"library\", \"func_score_type\"], observed=True\n ).sample(frac=subsample_frac, random_state=config.get(\"seed\", 42))\n func_scores = pd.concat([wt_rows, mut_rows]).sort_index()\n print(f\"Subsampled to {subsample_frac:.0%}: {len(func_scores)} rows \"\n f\"({len(wt_rows)} wt kept)\")\nelse:\n print(f\"Loaded {len(func_scores)} rows\")\n\nfunc_scores.head()", - "metadata": {}, - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "markdown", - "id": "87b5015d", - "source": "## Create Data objects\n\nOne `multidms.Data` per (library, func_score_type) combination.", - "metadata": {} - }, - { - "cell_type": "code", - "id": "38ba51ae", - "source": "data_objects = []\nfor (lib, fst), group_df in func_scores.rename(\n columns={\"homolog\": \"condition\"}\n).groupby([\"library\", \"func_score_type\"]):\n df = group_df.copy()\n df[\"aa_substitutions\"] = df[\"aa_substitutions\"].fillna(\"\")\n data_objects.append(\n multidms.Data(\n df,\n reference=\"h1\",\n alphabet=multidms.AAS_WITHSTOP_WITHGAP,\n verbose=False,\n name=f\"{lib}_{fst}_func_score\",\n )\n )\n\nprint(f\"Created {len(data_objects)} Data objects:\")\nfor d in data_objects:\n print(f\" {d.name}\")", - "metadata": {}, - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "markdown", - "id": "b04a3ae1", - "source": "## Build fitting parameters and fit models\n\nThis sweeps both `fusionreg` and `l2reg` (2D grid).", - "metadata": {} - }, - { - "cell_type": "code", - "id": "fafe9cdc", - "source": "fitting_params = build_fit_params(fit_config, data_objects)\nprint(\"Fitting parameters:\")\nfor k, v in fitting_params.items():\n if k != \"dataset\":\n print(f\" {k}: {v}\")", - "metadata": {}, - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "code", - "id": "9208c311", - "source": "n_models = len(explode_params_dict(fitting_params))\ncfg_n_processes = fit_config.get(\"n_processes\")\n\nif cfg_n_processes is None:\n n_processes = min(os.cpu_count() // 2, n_models)\nelse:\n n_processes = min(int(cfg_n_processes), n_models)\n\nn_processes = max(n_processes, 1)\nprint(f\"Fitting {n_models} models with n_processes={n_processes} (cpus={os.cpu_count()})\")\n\nn_fit, n_failed, fit_collection_df = fit_models(\n fitting_params, n_processes=n_processes\n)\n\n# Convert dict-valued columns to strings for groupby compatibility\nfor col in fit_collection_df.columns:\n if fit_collection_df[col].apply(lambda x: isinstance(x, dict)).any():\n fit_collection_df[col] = fit_collection_df[col].apply(str)\n\nprint(f\"Fit {n_fit} models successfully, {n_failed} failed\")", - "metadata": {}, - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "markdown", - "id": "69a22d0d", - "source": "## Post-process and save", - "metadata": {} - }, - { - "cell_type": "code", - "id": "8091ab80", - "source": "fit_collection_df = fit_collection_df.assign(\n library=(\n fit_collection_df[\"dataset_name\"]\n .str.split(\"_\").str[0:2].str.join(\"_\")\n ),\n measurement_type=(\n fit_collection_df[\"dataset_name\"]\n .str.split(\"_\").str[2:4].str.join(\"_\")\n ),\n)\n\nfit_collection_df[\"measurement_type\"] = pd.Categorical(\n fit_collection_df[\"measurement_type\"],\n categories=[\"observed_phenotype\", \"loose_bottle\", \"tight_bottle\"],\n ordered=True,\n)\n\noutput_path = os.path.join(output_dir, \"fit_collection.pkl\")\nwith open(output_path, \"wb\") as f:\n pickle.dump(fit_collection_df, f)\nprint(f\"Saved {output_path} ({len(fit_collection_df)} models)\")", - "metadata": {}, - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "markdown", - "id": "8bb201c6", - "source": "## Summary", - "metadata": {} - }, - { - "cell_type": "code", - "id": "436b0065", - "source": "summary_cols = [\n \"dataset_name\", \"library\", \"measurement_type\",\n \"fusionreg\", \"l2reg\", \"fit_time\",\n]\ndisplay_cols = [c for c in summary_cols if c in fit_collection_df.columns]\nfit_collection_df[display_cols]", - "metadata": {}, - "execution_count": null, - "outputs": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "name": "python", - "version": "3.9.0" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} \ No newline at end of file diff --git a/experiments/scv2-spike/README.md b/experiments/scv2-spike/README.md index 3e824507..a4584470 100644 --- a/experiments/scv2-spike/README.md +++ b/experiments/scv2-spike/README.md @@ -113,6 +113,13 @@ silently vanishing into notebook output. Every config variant has a matching `_downstream.yaml` sibling. See **Configuration tiers** below for which keys live where and why it matters. +The `config_recompute_false*.yaml` variants (three pairs) are unreachable from +the Snakefile by profile name and look like leftovers from a finished +experiment. They are **test fixtures**: `tests/test_config_tiers.py` iterates +`SPIKE_VARIANTS` and asserts on each, so deleting them fails four tests. +Removing them is a deliberate two-step change — edit `SPIKE_VARIANTS` first, +then delete the YAMLs. + ## Configuration tiers The config is split by dependency tier so that a downstream-only edit cannot @@ -140,6 +147,16 @@ the fit. > `prepare_data`, `cross_validation`, or `fit_models`. That would silently > restore the defect. See issue #287. +> ⚠️ **`config.yaml` is itself a rule `input:`, so Snakemake hashes the file, +> not its meaning.** Even a comment-only edit marks the fit out of date. When +> you know the change cannot affect results, `snakemake --touch` re-stamps the +> outputs instead of refitting — but verify `fit_collection.pkl` is +> byte-identical afterward, and never point `--touch` at a hand-picked subset. + +Note that `maxiter` is overloaded: the top-level value counts **outer +sweeps**, while the one inside `ge_kwargs` / `cal_kwargs` counts **inner +solver steps**. + Retuning the chosen lasso weight, changing a plot color, or adding a downstream analysis therefore reuses the cached `fit_collection.pkl`. @@ -279,10 +296,10 @@ corroboration rather than a decisive independent vote. `beta0_ridge=0.01`, `l2reg=1e-6`, `maxiter=200` (outer) / `10` (inner `ge_kwargs`, `cal_kwargs`), fusionreg = the 9-value manuscript ladder. -> These are **convergence-lab-derived** values, not the manuscript's ridge -> weights (manuscript: β ridge 1e-7, α ridge 1e-3). Note also that -> `beta0_ridge` is *shift shrinkage* — it penalizes `(β0_d − β0_ref)²`, -> the intercept **differences**, not the intercept magnitudes. +> These are not the manuscript's ridge weights (manuscript: β ridge 1e-7, +> α ridge 1e-3). Note `beta0_ridge` is **shift shrinkage** — it penalizes +> `(β0_d − β0_ref)²`, the intercept **differences**, not the intercept +> magnitudes. **Shift replicate correlation vs the pre-change baseline** (delta = new − baseline): diff --git a/experiments/scv2-spike/config/config.yaml b/experiments/scv2-spike/config/config.yaml index 70da6d01..36e481a1 100644 --- a/experiments/scv2-spike/config/config.yaml +++ b/experiments/scv2-spike/config/config.yaml @@ -29,38 +29,17 @@ spike: # Outer block-coordinate-descent ceiling. 500 is headroom, not a target; # a fit that reaches it has NOT converged regardless of any success count. maxiter: 500 - # 1e-6, tightened from the 1e-5 that produced 20/20 in the first Stage 1 - # run. The tolerance is a *between-sweep relative change*, not a gradient - # norm, so the question was whether the inner solvers (tol 1e-4) floor it. - # They do not: every one of those 20 fits was still decaying geometrically - # at its stopping point (median tail rate 0.941/sweep), so the tail had not - # flattened. Extrapolating each fit's own decay rate to 1e-6 projects a - # median of 127 sweeps and a worst case of 193 — well inside maxiter 500. + # Convergence tolerance on the between-sweep relative change in the + # objective — not a gradient norm. tol: 1.0e-6 - # false => the objective normalizer is computed once after warmstart and - # held constant, so the proximal lasso threshold `fusionreg / scale` is - # stationary across sweeps. Under `true` the threshold drifts between - # sweeps, which makes a lasso-selection study measure a moving target, and - # pins obj_old to 1.0 so the convergence test never sees a between-sweep - # change. Matches the simulation pipeline after PR #303. + # Hold the objective normalizer `scale` fixed for the whole fit so the + # convergence test and the proximal lasso threshold (fusionreg / scale) + # stay stationary across sweeps. recompute_scale: false - # Extended by one rung (1.28e-3) above the manuscript's 6.4e-4 ceiling, to - # test whether stop-codon sparsity is still climbing past the top rung. - # It is not: sparsity saturates at 1.000 by 3.2e-4 and is flat thereafter. - # - # The 1.28e-3 rung is RETAINED here but EXCLUDED from analysis (gate G5). - # Under tol 1e-6 it is the only rung that fails to converge, and the - # trajectory shows why: rep_1 reaches objective_error 2.7e-06 at sweep 100, - # then diverges -- the error climbs back to ~1.8e-04 and flattens (tail rate - # ~1.000), and the objective rises 10% above its sweep-109 minimum. That is - # instability at this lasso strength, not slow convergence, so raising - # maxiter would not fix it. The tol 1e-5 run stopped at sweep 94, before the - # divergence began, and so reported this rung as converged. - # - # It stays in the ladder because removing it would invalidate this config - # and force a full refit purely to *stop* computing one rung, discarding the - # evidence above. Nothing selects models by ladder position -- evaluate - # queries `fusionreg == lasso_choice` -- so an excluded rung is inert. + # Lasso ladder swept by the pipeline. The top rung (1.28e-3) is retained + # but excluded from analysis by gate G5: it is unstable at this lasso + # strength. An excluded rung is inert — evaluate selects by + # `fusionreg == lasso_choice`, never by ladder position. fusionreg_values: [0.0, 5.0e-6, 1.0e-5, 2.0e-5, 4.0e-5, 8.0e-5, 1.6e-4, 3.2e-4, 6.4e-4, 1.28e-3] l2reg: 1.0e-6 beta0_ridge: 0.01 @@ -72,12 +51,15 @@ spike: share_alpha: true beta_clip_range: [-10, 10] loss_kwargs: {"δ": 1.0} - # Inner solver maxiter deliberately stays at 10 while the outer loop moves - # to tol 1e-5 / maxiter 500. Simulation runs 100, but changing two - # convergence knobs in one run makes a bad result unattributable, and - # inner=10 has never been shown to be spike's problem. If fits converge - # poorly, raising this to 100 is the first lever — on evidence. + # Inner block-solver settings. NOTE `maxiter` here is inner solver steps, + # distinct from the top-level `maxiter` above, which is outer sweeps. ge_kwargs: {tol: 1.0e-4, maxiter: 10, maxls: 40, jit: true, verbose: false} cal_kwargs: {tol: 1.0e-4, maxiter: 10, maxls: 40, jit: true, verbose: false} - n_processes: 6 # pinned: auto (null) picks workers by core count alone, - # ignoring memory, which exhausted a 36 GB host mid-run. Never restore null. + # Worker count. Pinned to one worker per fit (10 fusionreg x 2 replicates + # = 20). JAX/XLA leaks executable JIT mappings across sequential fits in a + # single process, so a worker handling more than ~5 fits dies with + # "Unable to allocate section memory" even on a host with 1.4 TB free. + # One fit per worker means no process compiles twice. Never set to null: + # auto-sizing picks workers by core count alone, ignoring both memory and + # this leak, and has exhausted a host mid-run. ~35-38 GB RSS per worker. + n_processes: 20 diff --git a/experiments/simulation/README.md b/experiments/simulation/README.md index dfb2aa1c..1d62805f 100644 --- a/experiments/simulation/README.md +++ b/experiments/simulation/README.md @@ -129,10 +129,10 @@ pipeline has no `condition_colors`, `condition_titles`, or `domain_dict`. `beta0_ridge=0.01`, `l2reg=1e-6`, `maxiter=200` (outer) / `10` (inner `ge_kwargs`, `cal_kwargs`), fusionreg = the 9-value manuscript ladder. -> These are **convergence-lab-derived** values, not the manuscript's ridge -> weights (manuscript: β ridge 1e-7, α ridge 1e-3). `beta0_ridge` is -> *shift shrinkage* — it penalizes `(β0_d − β0_ref)²`, the intercept -> **differences**, not the intercept magnitudes. +> These are not the manuscript's ridge weights (manuscript: β ridge 1e-7, +> α ridge 1e-3). Note `beta0_ridge` is **shift shrinkage** — it penalizes +> `(β0_d − β0_ref)²`, the intercept **differences**, not the intercept +> magnitudes. #### AC7 — ground-truth shift recovery at `fusionreg = 8e-5` diff --git a/experiments/simulation/config/config.yaml b/experiments/simulation/config/config.yaml index 8e3f15f3..c838dd2c 100644 --- a/experiments/simulation/config/config.yaml +++ b/experiments/simulation/config/config.yaml @@ -16,15 +16,9 @@ simulation: n_shifted_identical_sites: 4 shift_gauss_variance: 0.666 sigmoid_phenotype_scale: 6 - # Halved 1000 -> 500 to match the manuscript simulation - # (SARS-CoV-2_spike_multidms @ 6c98b7b, simulation_validation.ipynb), whose - # cross-validation panel shows an interior validation-loss minimum that the - # pipeline at 1000 does not reproduce. At 1000 the fit sees ~29 variants per - # mutation, so every beta is pinned by direct evidence and shrinkage has no - # variance left to trade against bias -- validation loss just tracks - # training loss and rises monotonically in lambda. This is the single-knob - # A/B against results-prod-291-sim-trimmed-ladder-inner100; every other - # fit-tier key is held at its pipeline value. + # Variants simulated per library, as a multiple of genelength. Sets how much + # direct evidence each mutation gets; too high and shrinkage has no variance + # left to trade against bias, flattening the cross-validation minimum. variants_per_lib_genelength_scalar: 500 avgmuts: 2.0 bclen: 16 @@ -38,134 +32,32 @@ simulation: output_dir: results fitting: # Outer block-coordinate-descent sweep cap and convergence tolerance. - # Raised 200 -> 500 and tightened 1.0e-4 -> 1.0e-6 (see the loosening to - # 1.0e-5 noted below, which supersedes the 1e-6 value) alongside - # recompute_scale: false (below), to give every fit a fair chance to - # reach tol rather than running out of sweeps. Under the previous - # recompute_scale=true the tolerance was not comparable across sweeps - # anyway (objective_error was a within-sweep change against an obj_old - # pinned to 1.0), so tightening it only becomes meaningful now that the - # scale is held fixed. - # - # LOOSENED 1.0e-6 -> 1.0e-5 for this run. At 1e-6 the 500-variant run - # converged 60/60 on the full-data fit but only 59/60 in - # cross-validation: lib_2_observed_phenotype at the top rung - # (fusionreg 1.28e-3) exhausted all 500 outer sweeps and stopped at - # objective_error 2.32e-6, ~2.3x short of tol. This is a single-knob A/B - # against sim-vpl500 to see whether a 10x looser tolerance clears that - # holdout and how much it shifts the fits that already converged. + # A fit that reaches maxiter has NOT converged, whatever else reports. maxiter: 500 tol: 1.0e-5 - # Ladder extended past 6.4e-4, now topping out at 1.28e-3. - # - # The first recompute_scale=false run, before this extension, had not - # turned over at the old top rung (6.4e-4): sparsity was still climbing - # (0.59 vs ground truth 0.81) and - # recovery was still 0.748. Extending by two doublings bracketed both - # optima -- sparsity peaks at 1.28e-3 (0.848, ~= truth) and recovery peaks - # at 1.6e-4 (r=0.874), both turning over inside the plotted range. - # - # The 2.56e-3 rung was then DROPPED. It is where the optimizer breaks down, - # not where the model does: all 7 fits whose objective minimum was away - # from the final sweep sat there (drift up to +256% past the minimum), one - # burned 470 of 500 outer sweeps, and sparsity DECREASED rather than - # saturating. Dropping it took per-dataset sparsity monotonicity from 1/6 - # to 5/6. - # - # The remaining holdout is lib_1_tight_bottle, which drops 0.7284 -> 0.4695 - # at the top rung (1.28e-3) while the other five sit at 0.92-0.94, and is - # also the single remaining fit whose objective minimum is away from its - # final sweep. That is the same breakdown signature that justified dropping - # 2.56e-3, one rung lower and on the noisiest dataset only. 1.28e-3 is kept - # because it is the only rung that reaches the sparsity peak (0.848 vs - # ground truth 0.81), but whether it belongs in the PUBLISHED ladder is an - # open question -- see the follow-up noted in PR #303. + # Lasso ladder swept by the pipeline; `lasso_choice` in the downstream + # config selects which rung the figures use. fusionreg_values: [0.0, 5.0e-6, 1.0e-5, 2.0e-5, 4.0e-5, 8.0e-5, 1.6e-4, 3.2e-4, 6.4e-4, 1.28e-3] l2reg: 1.0e-6 beta0_ridge: 0.01 ge_type: "Sigmoid" warmstart: false - # Hold the objective normalizer `scale` fixed for the whole fit instead of - # recomputing it every outer sweep (the library default, used by every - # simulation run before this one). - # - # With recompute_scale: true, each sweep sets scale = |objective| and then - # obj_old = raw_obj/scale = 1.0 BY CONSTRUCTION, so the convergence test - # `|obj_old - obj| < tol` measures the WITHIN-sweep relative decrease - # against a moving normalizer -- not a true between-sweep change. The - # proximal lasso threshold (fusionreg / scale) also drifts with it, so the - # problem being minimized is not stationary across sweeps. - # - # convergence-lab measured false as strictly dominant at inner maxiter=100: - # 8/8 vs 4/8 fits converged AND 3.6x faster (788s vs 2808s). - # See experiments/convergence-lab/README.md, entry 2026-07-10, - # cache=maxiter-scan. + # Hold the objective normalizer `scale` fixed for the whole fit rather + # than recomputing it each outer sweep, so the convergence test and the + # proximal lasso threshold (fusionreg / scale) stay stationary across + # sweeps. recompute_scale: false beta0_init: {h1: 5.0, h2: 0.0} alpha_init: 6.0 share_alpha: true beta_clip_range: [-10, 10] loss_kwargs: {"δ": 1.0} - # Inner block-solver cap. Raised 10 -> 100 in #291 (Phase 1 of EPIC #290). - # PR #288 lowered it 100 -> 10 as a side effect of the config tier split; - # that is the measured cause of the non-monotonic shift-sparsity curve. - # Single-knob arms at inner=100 on lib_1_observed_phenotype_func_score: - # sparsity +0.341 at fusionreg=3.2e-4 (0.5189 -> 0.8600) and +0.383 at - # 6.4e-4 (0.2600 -> 0.6432), with alpha returning from ~9.3 to ~6.2-7.0 - # (alpha_init 6.0) and objective_total down 58%. Zeroing l2reg (+0.000) or - # beta0_ridge (+0.067) does not reproduce this. - # - # 50 was then tried and REVERTED to 100 -- measured A/B at outer 500 / - # tol 1e-6, recompute_scale=false. The two arms ran on different ladders - # (inner=100 on the 11-rung ladder still carrying 2.56e-3, inner=50 on the - # 10-rung trimmed ladder), hence the different denominators: - # - # inner=100 inner=50 - # converged 66/66 59/60 <- lost the guarantee - # max outer sweeps 470 500 <- hit the cap - # drift fits 7 (@2.56e-3) 3 (@1.28e-3) - # mean sparsity 0.848 0.765 <- further from GT 0.81 - # @1.28e-3 - # median s/fit 436 419 <- only ~4% cheaper - # - # Halving inner effort bought ~4% wall-clock and cost the 100% - # convergence guarantee, pushed a fit into the 500-sweep cap, and - # regressed peak sparsity. Inner iterations are not wasted work: they are - # what keeps the OUTER loop short, so weakening them lengthens the outer - # tail. Worse, the drift region MIGRATED DOWN the ladder from the rung we - # dropped (2.56e-3) into a rung we keep (1.28e-3). - # - # Convergence triple for later EPIC #290 phases to inherit -- - # SIMULATION-PIPELINE SCOPE ONLY: - # (inner maxiter, outer maxiter, tol) = (100, 500, 1.0e-6) - # experiments/scv2-spike/ deliberately stays at inner maxiter = 10: the - # evidence above is simulation-only (noise-free data, known ground-truth - # sparsity 0.81), the cost on spike is unmeasured, and changing spike's - # fit-tier config.yaml would invalidate its cached fits. Any spike change - # must be measured on spike data first. + # Inner block-solver settings. NOTE `maxiter` here is inner solver steps, + # distinct from the top-level `maxiter` above, which is outer sweeps. ge_kwargs: {tol: 1.0e-4, maxiter: 100, maxls: 40, jit: true, verbose: false} cal_kwargs: {tol: 1.0e-4, maxiter: 100, maxls: 40, jit: true, verbose: false} - # Worker count for the 60-fit grid (6 datasets x 10 fusionreg). NEVER set - # this to null: auto-selection picks workers by core count alone, ignoring - # memory, which exhausted a 36 GB host mid-run. - # - # This is HOST-DEPENDENT. Size it to cores first, then sanity-check RAM. - # - # CORRECTION: an earlier version of this comment claimed ~42 GB resident - # per worker and budgeted 45 GB/worker. That number was never measured -- - # it was back-inferred from a 36 GB laptop dying at 7 workers. Direct - # measurement of a worker running a fit from this grid gives a peak RSS of - # ~2.9 GB, corroborated independently by the repo's own instrumented probe - # (experiments/convergence-lab/diagnostics/parallelism_probe.py: ~1 GB - # per-worker duplication, ~4.9 GB total peak). The real figure is ~15x - # smaller, so MEMORY IS NOT THE BINDING CONSTRAINT ON THIS GRID -- cores - # are. Each task carries exactly one Data object (build_fit_params puts - # "dataset" on the product axis), so per-worker memory is O(1 dataset), - # not O(all datasets). - # - # Set to 30: the grid is 6 datasets x 10 fusionreg = 60 fits, which divides - # into exactly 2 waves on a 64-core host, leaving cores free for the JAX - # thread pools each worker spins up. At ~3 GB/worker that is ~90 GB, ~6% - # of a 1511 GB host. Lower this on a smaller host -- and never set it to - # null, since auto-selection picks workers by core count alone. + # Worker count for the 60-fit grid (6 datasets x 10 fusionreg). Never set + # to null: auto-sizing picks workers by core count alone, ignoring memory, + # which has exhausted a host mid-run. Host-dependent -- size to cores, then + # sanity-check RAM. Cores, not memory, bind this grid (~3 GB/worker). n_processes: 30 diff --git a/experiments/simulation/notebooks/_common.py b/experiments/simulation/notebooks/_common.py index 51bcb035..5c13345a 100644 --- a/experiments/simulation/notebooks/_common.py +++ b/experiments/simulation/notebooks/_common.py @@ -283,17 +283,8 @@ def build_fit_params(fit_config, datasets): "cal_kwargs": [fit_config["cal_kwargs"]], "loss_kwargs": [fit_config["loss_kwargs"]], "warmstart": [fit_config["warmstart"]], - # Objective normalizer. False holds `scale` fixed for the whole fit, so - # the optimization problem, the lasso threshold (fusionreg / scale), and - # `objective_error` are all stationary across outer sweeps. True (the - # library default, and what every simulation run before this one used) - # recomputes it every sweep, which makes `objective_error` a - # within-sweep change measured against an obj_old that is 1.0 by - # construction -- not a true relative change between sweeps. - # convergence-lab measured False as strictly dominant at inner - # maxiter=100: 8/8 vs 4/8 fits converged AND 3.6x faster (788s vs - # 2808s). See experiments/convergence-lab/README.md, entry - # 2026-07-10, cache=maxiter-scan. + # Hold the objective normalizer fixed for the whole fit so the + # convergence test and lasso threshold stay stationary across sweeps. # # NOTE: this dict is an explicit whitelist -- a key added to the YAML # but missing here is silently dropped. Optional via .get() so older diff --git a/tests/test_convergence_lab_harness.py b/tests/test_convergence_lab_harness.py deleted file mode 100644 index de5e72e5..00000000 --- a/tests/test_convergence_lab_harness.py +++ /dev/null @@ -1,145 +0,0 @@ -"""Unit tests for the convergence-lab harness (issue #253).""" - -import sys -from pathlib import Path - -import pytest -import yaml - -# The harness lives under experiments/, not on the package path; import by path. -HARNESS_DIR = Path(__file__).parent.parent / "experiments" / "convergence-lab" -sys.path.insert(0, str(HARNESS_DIR)) - -import harness # noqa: E402 - - -def _write_yaml(tmp_path, body): - """Write ``body`` as a YAML grid file under ``tmp_path`` and return its path.""" - p = tmp_path / "grid.yaml" - p.write_text(yaml.safe_dump(body)) - return p - - -def test_load_config_defaults_replicates(tmp_path): - """load_config defaults replicates to [1, 2] and leaves sweep/fixed intact.""" - cfg = harness.load_config(_write_yaml(tmp_path, {"sweep": {"l2reg": [0.0, 3e-4]}})) - assert cfg["replicates"] == [1, 2] - assert cfg["sweep"] == {"l2reg": [0.0, 3e-4]} - assert cfg["fixed"] == {} - - -def test_load_config_rejects_unknown_sweep_key(tmp_path): - """An unknown key in the sweep section raises ValueError naming the key.""" - with pytest.raises(ValueError, match="not_a_kwarg"): - harness.load_config(_write_yaml(tmp_path, {"sweep": {"not_a_kwarg": [1, 2]}})) - - -def test_load_config_rejects_unknown_fixed_key(tmp_path): - """An unknown key in the fixed section raises ValueError naming the key.""" - with pytest.raises(ValueError, match="bogus"): - harness.load_config( - _write_yaml( - tmp_path, - {"sweep": {"l2reg": [0.0]}, "fixed": {"bogus": True}}, - ) - ) - - -def test_load_config_rejects_dataset_key(tmp_path): - """`dataset` is runner-supplied, not a config key, so it is rejected.""" - with pytest.raises(ValueError, match="dataset"): - harness.load_config(_write_yaml(tmp_path, {"sweep": {"dataset": [1]}})) - - -def test_load_config_requires_sweep(tmp_path): - """A config with no sweep section raises ValueError.""" - with pytest.raises(ValueError, match="sweep"): - harness.load_config(_write_yaml(tmp_path, {"fixed": {"warmstart": True}})) - - -def test_explode_grid_cartesian_with_replicates(tmp_path): - """explode_grid crosses sweep x replicates and merges fixed into each cell.""" - cfg = harness.load_config( - _write_yaml( - tmp_path, - { - "sweep": {"l2reg": [0.0, 3e-4]}, - "fixed": {"warmstart": True}, - "replicates": [1, 2], - }, - ) - ) - exploded = harness.explode_grid(cfg) - assert len(exploded) == 4 # 2 l2reg × 2 replicates - assert all(d["warmstart"] is True for d in exploded) - assert {d["replicate"] for d in exploded} == {1, 2} - assert {d["l2reg"] for d in exploded} == {0.0, 3e-4} - - -@pytest.mark.needs_local_data -def test_load_rep_data_builds_two_reps(): - """load_rep_data builds rep_1/rep_2 Data with the expected conditions. - - Reads the gitignored prod spike CSV, so this runs only where the pipeline - has been run locally. - """ - rep_data = harness.load_rep_data() - assert set(rep_data) == {"rep_1", "rep_2"} - for name, data in rep_data.items(): - assert data.name == name - assert "Omicron_BA1" in data.conditions - assert "Delta" in data.conditions - - -def test_build_params_maps_replicates_and_collects_sweep(): - """build_params maps replicate→Data and collects each kwarg's distinct values. - - The result is a ``fit_models`` ``params`` dict: ``dataset`` is the list of - rep Data objects, every other key the distinct values seen across cells. - Crossing dataset × those lists must reproduce exactly the exploded cells. - """ - rep_data = {"rep_1": "DATA1", "rep_2": "DATA2"} - exploded = [ - {"l2reg": 0.0, "warmstart": True, "replicate": 1}, - {"l2reg": 3e-4, "warmstart": True, "replicate": 1}, - {"l2reg": 0.0, "warmstart": True, "replicate": 2}, - {"l2reg": 3e-4, "warmstart": True, "replicate": 2}, - ] - params = harness.build_params(exploded, rep_data) - assert params["dataset"] == ["DATA1", "DATA2"] - assert params["l2reg"] == [0.0, 3e-4] - assert params["warmstart"] == [True] - assert "replicate" not in params - - -def test_run_fits_delegates_to_fit_models(monkeypatch): - """run_fits calls fit_models with the params + n_processes and returns its frame.""" - import pandas as pd - - captured = {} - - def fake_fit_models(params, n_processes, failures): - captured["params"] = params - captured["n_processes"] = n_processes - captured["failures"] = failures - frame = pd.DataFrame({"l2reg": [0.0, 3e-4], "model": [object(), object()]}) - return (2, 0, frame) - - monkeypatch.setattr(harness, "fit_models", fake_fit_models) - params = {"dataset": ["D1", "D2"], "l2reg": [0.0, 3e-4]} - df = harness.run_fits(params, n_processes=3) - assert captured["n_processes"] == 3 - assert captured["failures"] == "tolerate" - assert len(df) == 2 - - -def test_default_n_processes_caps_at_grid_size(monkeypatch): - """default_n_processes never exceeds the grid size and is at least 1.""" - monkeypatch.setattr(harness.os, "cpu_count", lambda: 8) - # Grid smaller than cores-1 → capped at grid size. - assert harness.default_n_processes(3) == 3 - # Grid larger than cores-1 → capped at cores-1. - assert harness.default_n_processes(100) == 7 - # Always at least 1. - monkeypatch.setattr(harness.os, "cpu_count", lambda: 1) - assert harness.default_n_processes(10) == 1