v0.4.0 ↔ main model equivalence check (spike analysis)

[Haddox]

Summary

Adds a self-contained validation experiment (experiments/v04-equivalence/) showing that the
manuscript multidms 0.4.0 global-epistasis-with-shifts model and the current main
(PR #164) model are one forward model in two coordinate systems, and that the change of
variables preserves the Huber loss.

The map from a fitted 0.4.0 model to main's parameters is:

• β_d = β + Δ_d — per-condition absolute effect vector (reference: Δ = 0 ⇒ β_ref = β)
• β0,d = β0 + α_d — fold the additive offset into the per-condition intercept
• α = θ_scale — the 0.4.0 sigmoid range becomes main's output scale

With these, both formulations evaluate the same latent phenotype φ_d(v); the two Huber residuals
differ only by a per-condition constant D_d = θ_bias + θ_scale·σ(φ_wt,d), which vanishes when the
fit places each homolog's wildtype at predicted score 0.

Dataset

The notebook evaluates the SARS-CoV-2 spike deep-mutational-scanning data from the manuscript
analysis (matsengrp/SARS-CoV-2_spike_multidms @ 6c98b7b), the same fit reported in the paper:

• Three homologs — Delta, Omicron_BA1 (the reference), Omicron_BA2 — jointly modeled with a
  shared effect vector β plus per-homolog shifts Δ_d.
• Per-mutation effects/shifts come from the committed mutations_df.csv (at the chosen shift-lasso
  4e-5); observed per-variant functional scores come from training_functional_scores.csv
  (replicate 2), collapsing identical variants by mean to match the fit's
  collapse_identical_variants="mean". These CSVs (~48 MB) are fetched at runtime from the pinned
  commit and are git-ignored.

Relationship to the two images

The committed CSVs give β and the shifts, but not the sigmoid global-epistasis shape
parameters (no fitted pickle is available). We recover those from two diagnostic figures from the
0.4.0 spike fit:

• ge_fits.png (global-epistasis curve: observed functional score vs. predicted latent φ,
  with the fitted sigmoid g(φ) = θ_scale·σ(φ) + θ_bias and each homolog's wildtype marked). From
  it we read the shape parameters used to reconstruct the 0.4.0 model: θ_scale ≈ 7.4 (range),
  θ_bias ≈ −3.4 (lower asymptote), and φ_wt ≈ −0.4 (wildtype latent). Crucially, the wildtype
  lines sit on the g(φ) = 0 crossing, i.e. D_d ≈ 0 — the condition that makes the loss identity
  exact.
• func_score_corr.png (observed vs. predicted functional score, per homolog: Delta r ≈ 0.80,
  BA1 ≈ 0.86–0.90, BA2 ≈ 0.86–0.88). This is the validation target: recomputing predicted
  functional scores from the reconstructed model should reproduce these correlations.

Both figures are included in the experiment folder for reference.

Key results

1. Huber loss preserved. With θ_bias calibrated so D_d = 0, the 0.4.0 and main
   per-condition Huber losses are identical to machine precision:

   condition	0.4.0	main	gap
   Delta	0.2892	0.2892	0.0e+00
   Omicron_BA1	0.2909	0.2909	0.0e+00
   Omicron_BA2	0.2368	0.2368	0.0e+00

   With the raw eyeballed θ_bias = −3.4 (D_d ≈ −0.43) the gap is still tiny (≤ 0.047 per
   condition) — confirming the losses coincide as D_d → 0.

2. Correlations match func_score_corr.png. Recomputed predicted-vs-observed Pearson r
   (identical under both forms): Delta 0.795, Omicron_BA1 0.891, Omicron_BA2 0.876 —
   matching the published ~0.80 / 0.86–0.90 / 0.86–0.88.

3. main code-path agrees. jaxmodels.functional_score_loss reproduces the main-form loss for
   all three homologs.

Initializing a main model from the 0.4.0 fit

Beyond validating the loss identity, convert_040_to_main_params.py materializes the mapped
parameters as CSVs, so a main jaxmodels.Model can be initialized at the exact point in
parameter space implied by the 0.4.0 fit — no pkl, no re-fit. It applies the map above
(β_d = β + Δ_d; β0,d = φ_wt,d − x_wt,d·β_d; α = θ_scale; logθ_d = 0; θ_bias calibrated so
D_d = 0), runs it on the spike replicate-2 parameter set, and writes three CSVs to
main_init_params/ that together fully specify a Model:

CSV	contents	initializes
main_params_beta.csv	long (condition, mutation, beta), 3 × 10,870 rows	Model.φ[c].β
main_params_latent.csv	per condition: beta0, logtheta, is_reference (+ phi_wt, D_d diagnostics)	Model.φ[c].β0, Model.logθ[c]
main_params_global.csv	key/value: alpha=7.4, reference_condition=Omicron_BA1, global_epistasis=Sigmoid, output_activation=IdentityOutput, huber_delta=1.0, calibrated theta_bias	Model.α, Model.reference_condition, structural choices

The script (and a mirrored notebook cell) then reads the CSVs back, rebuilds a Model, and
verifies it lands where intended: per-condition functional_score_loss matches an in-memory model
to machine precision (gaps ≤ 5.6e-17) and correlations reproduce func_score_corr.png
(0.795 / 0.891 / 0.876). The per-condition β0 differ (−0.452 / −0.400 / −0.561) precisely because
each folds in that homolog's bundle term −x_wt,d·β_d, placing all three wildtypes at the read
φ_wt = −0.40.

Contents

• check_equivalence_from_image.ipynb — executed notebook with the three checks, a worked
  single-variant walkthrough, and the converter/round-trip cell.
• convert_040_to_main_params.py — converts 0.4.0 parameters → the main-model init CSVs, with a
  built-in CSV round-trip verification.
• main_init_params/{main_params_beta,main_params_latent,main_params_global}.csv — the emitted
  main-model initialization parameters for spike replicate 2.
• compute_equivalence.py — script form of the loss-equivalence analysis.
• pick_example_variant.py — extracts the real-variant worked example.
• .gitignore — excludes the fetched data/ CSVs.
• ge_fits.png, func_score_corr.png — the reference figures.

Closes #281.

🤖 Generated with Claude Code