diff --git a/experiments/v04-equivalence/.gitignore b/experiments/v04-equivalence/.gitignore
new file mode 100644
index 00000000..a136352f
--- /dev/null
+++ b/experiments/v04-equivalence/.gitignore
@@ -0,0 +1,3 @@
+# Large committed spike-analysis CSVs are fetched at runtime from the pinned
+# SARS-CoV-2_spike_multidms commit; do not commit them here.
+data/
diff --git a/experiments/v04-equivalence/check_equivalence_from_image.ipynb b/experiments/v04-equivalence/check_equivalence_from_image.ipynb
new file mode 100644
index 00000000..2c6f0564
--- /dev/null
+++ b/experiments/v04-equivalence/check_equivalence_from_image.ipynb
@@ -0,0 +1,675 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "id": "206d3404",
+ "metadata": {},
+ "source": [
+ "# Expressing the 0.4.0 spike model in the PR #164 (`main`) functional form\n",
+ "\n",
+ "This notebook shows that the manuscript's **multidms 0.4.0** global-epistasis-with-shifts model\n",
+ "and the current **`main` (PR #164)** model are one forward model in two coordinate systems, and\n",
+ "that the transformation between them **preserves the Huber loss**.\n",
+ "\n",
+ "We do *not* use the fitted pickle. Instead we read the sigmoid global-epistasis shape parameters\n",
+ "approximately from the committed diagnostic plot `ge_fits.png`, and take the per-mutation effects /\n",
+ "shifts and the observed functional scores from the committed spike-analysis CSVs\n",
+ "(`matsengrp/SARS-CoV-2_spike_multidms` @ `6c98b7b`).\n",
+ "\n",
+ "**Three checks:**\n",
+ "1. the 0.4.0 and `main` Huber losses coincide (exactly as `D_d → 0`);\n",
+ "2. recomputed predicted-vs-observed correlations match `func_score_corr.png` (Pearson r ≈ 0.80–0.90);\n",
+ "3. `main`'s own `functional_score_loss` reproduces the main-form loss."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "5dc11a1b",
+ "metadata": {},
+ "source": [
+ "## The transformation\n",
+ "\n",
+ "Let $\\phi_d(v)$ be the latent phenotype of variant $v$ in condition $d$ and $\\sigma$ the logistic sigmoid.\n",
+ "\n",
+ "**0.4.0** (identity output, $\\gamma=0$): latent $\\phi_d(v)=\\beta_0+\\alpha_d+X_v\\,(\\beta+\\Delta_d)$,\n",
+ "GE $g(z)=\\theta_\\text{scale}\\,\\sigma(z)+\\theta_\\text{bias}$, and the Huber loss is fit to the **raw**\n",
+ "prediction $\\hat y_d(v)=\\theta_\\text{scale}\\,\\sigma(\\phi_d(v))+\\theta_\\text{bias}$ (WT subtracted only when reporting).\n",
+ "\n",
+ "**`main`**: latent $\\phi_d(v)=\\beta_{0,d}+X_v\\,\\beta_d$, parameter-free GE $g(z)=\\sigma(z)$, and the loss\n",
+ "is fit to the **wildtype-subtracted** prediction $\\hat y_d(v)=\\alpha\\,(\\sigma(\\phi_d(v))-\\sigma(\\phi_{\\mathrm{wt},d}))$.\n",
+ "\n",
+ "**Map (0.4.0 → main):**\n",
+ "\n",
+ "| `main` parameter | from 0.4.0 |\n",
+ "|---|---|\n",
+ "| $\\beta_d = \\beta+\\Delta_d$ | per-condition absolute effect vector (reference: $\\Delta=0\\Rightarrow\\beta_\\text{ref}=\\beta$) |\n",
+ "| $\\beta_{0,d} = \\beta_0+\\alpha_d$ | fold the additive offset into the per-condition intercept |\n",
+ "| $\\alpha = \\theta_\\text{scale}$ | the sigmoid *range* becomes `main`'s output scale |\n",
+ "\n",
+ "With these, $\\phi_d(v)$ is identical variant-for-variant, so the two Huber residuals differ only by a\n",
+ "per-condition constant\n",
+ "\n",
+ "$$\\text{residual}_{0.4.0}-\\text{residual}_\\text{main}=\\theta_\\text{bias}+\\theta_\\text{scale}\\,\\sigma(\\phi_{\\mathrm{wt},d})\\equiv D_d .$$\n",
+ "\n",
+ "`main` has no free per-condition output constant (only the structural $-\\alpha\\,\\sigma(\\phi_{\\mathrm{wt},d})$),\n",
+ "so the losses coincide exactly when $D_d=0$ — i.e. when the fit puts each condition's WT at raw score $0$.\n",
+ "In `ge_fits.png` the WT vertical lines sit on the $g(\\phi)=0$ crossing, i.e. $D_d\\approx0$."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "id": "60aaff55",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-15T15:36:19.133690Z",
+ "iopub.status.busy": "2026-07-15T15:36:19.133342Z",
+ "iopub.status.idle": "2026-07-15T15:36:32.420625Z",
+ "shell.execute_reply": "2026-07-15T15:36:32.418817Z"
+ }
+ },
+ "outputs": [],
+ "source": [
+ "import os\n",
+ "import numpy as np\n",
+ "import pandas as pd\n",
+ "import scipy.sparse\n",
+ "import multidms\n",
+ "import multidms.jaxmodels as jm\n",
+ "import jax.numpy as jnp\n",
+ "\n",
+ "HERE = os.getcwd() # run this notebook from experiments/v04-equivalence/\n",
+ "DATA = os.path.join(HERE, \"data\")\n",
+ "\n",
+ "REFERENCE = \"Omicron_BA1\"\n",
+ "CHOSEN_LASSO = 4e-5\n",
+ "REPLICATE = \"rep-2\" # committed label in mutations_df.csv\n",
+ "REPLICATE_NUM = 2 # matching value in training_functional_scores.csv\n",
+ "DELTA = 1.0 # Huber delta (0.4.0 and main default)\n",
+ "\n",
+ "# ---- sigmoid GE-shape parameters read approximately from ge_fits.png ----\n",
+ "THETA_SCALE = 7.4 # sigmoid range (midpoint read; upper asymptote off-screen)\n",
+ "THETA_BIAS_RAW = -3.4 # lower asymptote\n",
+ "# WT vertical lines cluster near phi = -0.4 for all conditions (replicate-2 panel)\n",
+ "PHI_WT_READ = {\"Delta\": -0.4, \"Omicron_BA1\": -0.4, \"Omicron_BA2\": -0.4}\n",
+ "\n",
+ "sigmoid = lambda z: 1.0 / (1.0 + np.exp(-z))\n",
+ "def huber(r, delta=DELTA):\n",
+ " a = np.abs(r)\n",
+ " return np.where(a <= delta, 0.5 * a * a, delta * (a - 0.5 * delta))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "56edc195",
+ "metadata": {},
+ "source": [
+ "## Load committed data and build the reference-frame encoding"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "id": "2db13341",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-15T15:36:32.424205Z",
+ "iopub.status.busy": "2026-07-15T15:36:32.423695Z",
+ "iopub.status.idle": "2026-07-15T15:44:06.390726Z",
+ "shell.execute_reply": "2026-07-15T15:44:06.389574Z"
+ }
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "variants: raw=333910 collapsed=149006\n"
+ ]
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "n_mutations=10870 conditions=['Delta', 'Omicron_BA1', 'Omicron_BA2'] reference=Omicron_BA1\n"
+ ]
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "theta_bias: raw=-3.4 calibrated(D_ref=0)=-2.9697\n"
+ ]
+ }
+ ],
+ "source": [
+ "# Fetch the committed spike-analysis CSVs (pinned commit) if not already present locally.\n",
+ "import urllib.request\n",
+ "os.makedirs(DATA, exist_ok=True)\n",
+ "BASE = (\"https://raw.githubusercontent.com/matsengrp/SARS-CoV-2_spike_multidms/\"\n",
+ " \"6c98b7b607d7387b508cdaa192d659ee9fca7367/results/spike_analysis\")\n",
+ "for _f in [\"mutations_df.csv\", \"training_functional_scores.csv\"]:\n",
+ " _p = os.path.join(DATA, _f)\n",
+ " if not os.path.exists(_p):\n",
+ " print(\"downloading\", _f)\n",
+ " urllib.request.urlretrieve(f\"{BASE}/{_f}\", _p)\n",
+ "\n",
+ "# observed functional scores (one replicate), collapsing identical variants by mean\n",
+ "# (matches the fit's collapse_identical_variants=\"mean\")\n",
+ "fs = pd.read_csv(os.path.join(DATA, \"training_functional_scores.csv\"))\n",
+ "fs = fs[fs[\"replicate\"] == REPLICATE_NUM].copy()\n",
+ "fs[\"aa_substitutions\"] = fs[\"aa_substitutions\"].fillna(\"\")\n",
+ "variants = fs.groupby([\"condition\", \"aa_substitutions\"], as_index=False)[\"func_score\"].mean()\n",
+ "print(f\"variants: raw={len(fs)} collapsed={len(variants)}\")\n",
+ "\n",
+ "mdata = multidms.Data(variants, reference=REFERENCE,\n",
+ " alphabet=multidms.AAS_WITHSTOP_WITHGAP, assert_site_integrity=False)\n",
+ "muts = list(mdata.mutations)\n",
+ "idx = {m: i for i, m in enumerate(muts)}\n",
+ "conditions = list(mdata.conditions)\n",
+ "print(f\"n_mutations={len(muts)} conditions={conditions} reference={mdata.reference}\")\n",
+ "\n",
+ "# per-mutation beta / shifts (chosen lasso + replicate)\n",
+ "md_ = pd.read_csv(os.path.join(DATA, \"mutations_df.csv\"))\n",
+ "md_ = md_[(md_[\"dataset_name\"] == REPLICATE)\n",
+ " & np.isclose(md_[\"scale_coeff_lasso_shift\"].astype(float), CHOSEN_LASSO)].set_index(\"mutation\")\n",
+ "beta = np.zeros(len(muts)); shift = {c: np.zeros(len(muts)) for c in conditions}\n",
+ "for m in muts:\n",
+ " if m in md_.index:\n",
+ " beta[idx[m]] = md_.at[m, \"beta\"]\n",
+ " shift[\"Delta\"][idx[m]] = md_.at[m, \"shift_Delta\"]\n",
+ " shift[\"Omicron_BA2\"][idx[m]] = md_.at[m, \"shift_Omicron_BA2\"]\n",
+ "beta_d = {c: beta + shift[c] for c in conditions}\n",
+ "\n",
+ "def cond_arrays(c):\n",
+ " X = mdata.arrays[\"X\"][c] # BCOO, WT at row 0\n",
+ " Xs = scipy.sparse.csr_array((np.asarray(X.data),\n",
+ " (np.asarray(X.indices[:, 0]), np.asarray(X.indices[:, 1]))), shape=X.shape)\n",
+ " x_wt = np.asarray(Xs[[0], :].todense()).ravel()\n",
+ " return Xs[1:], x_wt, np.asarray(mdata.arrays[\"y\"][c])[1:]\n",
+ "\n",
+ "# intercept from the read WT latent: phi_wt = beta0 + x_wt . beta_d\n",
+ "beta0 = {}\n",
+ "for c in conditions:\n",
+ " _, x_wt, _ = cond_arrays(c)\n",
+ " beta0[c] = PHI_WT_READ[c] - float(x_wt @ beta_d[c])\n",
+ "\n",
+ "THETA_BIAS_CAL = -THETA_SCALE * sigmoid(PHI_WT_READ[REFERENCE]) # enforce D_ref = 0\n",
+ "print(f\"theta_bias: raw={THETA_BIAS_RAW} calibrated(D_ref=0)={THETA_BIAS_CAL:.4f}\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "a4189107",
+ "metadata": {},
+ "source": [
+ "## Check 1 — the Huber loss is preserved by the transformation"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "id": "221e58f6",
+ "metadata": {
+ "execution": {
+ "iopub.execute_input": "2026-07-15T15:44:06.397711Z",
+ "iopub.status.busy": "2026-07-15T15:44:06.397393Z",
+ "iopub.status.idle": "2026-07-15T15:44:06.672694Z",
+ "shell.execute_reply": "2026-07-15T15:44:06.671102Z"
+ }
+ },
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Raw eyeballed theta_bias = -3.400\n"
+ ]
+ },
+ {
+ "data": {
+ "text/html": [
+ "\n",
+ "\n",
+ "
\n",
+ " \n",
+ " \n",
+ " | \n",
+ " condition | \n",
+ " D_d | \n",
+ " loss_0_4_0 | \n",
+ " loss_main | \n",
+ " abs_gap | \n",
+ "
\n",
+ " \n",
+ " \n",
+ " \n",
+ " | 0 | \n",
+ " Delta | \n",
+ " -0.430289 | \n",
+ " 0.324947 | \n",
+ " 0.289200 | \n",
+ " 0.035747 | \n",
+ "
\n",
+ " \n",
+ " | 1 | \n",
+ " Omicron_BA1 | \n",
+ " -0.430289 | \n",
+ " 0.243844 | \n",
+ " 0.290895 | \n",
+ " 0.047050 | \n",
+ "
\n",
+ " \n",
+ " | 2 | \n",
+ " Omicron_BA2 | \n",
+ " -0.430289 | \n",
+ " 0.238204 | \n",
+ " 0.236814 | \n",
+ " 0.001390 | \n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
\n",
+ "\n",
+ "
\n",
+ " \n",
+ " \n",
+ " | \n",
+ " condition | \n",
+ " D_d | \n",
+ " loss_0_4_0 | \n",
+ " loss_main | \n",
+ " abs_gap | \n",
+ "
\n",
+ " \n",
+ " \n",
+ " \n",
+ " | 0 | \n",
+ " Delta | \n",
+ " 0.0 | \n",
+ " 0.289200 | \n",
+ " 0.289200 | \n",
+ " 0.0 | \n",
+ "
\n",
+ " \n",
+ " | 1 | \n",
+ " Omicron_BA1 | \n",
+ " 0.0 | \n",
+ " 0.290895 | \n",
+ " 0.290895 | \n",
+ " 0.0 | \n",
+ "
\n",
+ " \n",
+ " | 2 | \n",
+ " Omicron_BA2 | \n",
+ " 0.0 | \n",
+ " 0.236814 | \n",
+ " 0.236814 | \n",
+ " 0.0 | \n",
+ "
\n",
+ " \n",
+ "
\n",
+ "