Model Evaluation

Evaluate fitted spike models: convergence diagnostics, sparsity analysis, replicate parameter correlations, global epistasis plots, and mutation parameter export.

Outline

1. Load fitted models and training data

2. Convergence diagnostics

3. Shift sparsity analysis

4. Replicate parameter correlations

5. Global epistasis plots

6. Export mutations DataFrame and intermediate CSVs

import warnings

warnings.filterwarnings("ignore")

import os
import pickle
import sys

sys.path.insert(0, "notebooks")

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import multidms.plot
from multidms.model_collection import ModelCollection

from _common import load_config, combine_replicate_muts

config_path = "config/config.yaml"
downstream_config_path = "config/config_downstream.yaml"
output_dir = None

# Parameters
config_path = "config/config.yaml"
downstream_config_path = "config/config_downstream.yaml"
output_dir = "results-prod-294-naive-baseline-arm"

config = load_config(config_path, downstream_config_path)
spike = config["spike"]
fit_config = spike["fitting"]
lasso_choice = spike["lasso_choice"]
condition_titles = spike["condition_titles"]
condition_colors = spike["condition_colors"]
experiment_conditions = spike["experiment_conditions"]

if output_dir is None:
    output_dir = spike.get("output_dir", "results")

Load data

with open(os.path.join(output_dir, "fit_collection.pkl"), "rb") as f:
    fit_collection_df = pickle.load(f)

func_score_df = pd.read_csv(
    os.path.join(output_dir, "training_functional_scores.csv")
).fillna({"aa_substitutions": ""})

model_collection = ModelCollection(fit_collection_df)
print(f"Loaded {len(fit_collection_df)} fitted models")
print(f"Loaded {len(func_score_df):,} training variants")

Loaded 20 fitted models
Loaded 303,007 training variants

Convergence diagnostics

# Build summary table with fitted alpha and beta0 parameters
summary_rows = []
for _, row in model_collection.fit_models.iterrows():
    jm = row.model._jax_model
    r = {
        "dataset": row.dataset_name,
        "fusionreg": row.fusionreg,
        "converged": row.converged,
        "fit_time": row.fit_time,
    }
    # Alpha: shared scalar or per-condition dict
    if hasattr(jm.α, "items"):
        for cond in experiment_conditions:
            title = condition_titles.get(cond, cond)
            r[f"alpha_{title}"] = float(jm.α[cond])
    else:
        r["alpha"] = float(jm.α)
    # Per-condition loss
    for cond in experiment_conditions:
        loss_col = f"{cond}_loss_training"
        if loss_col in model_collection.fit_models.columns:
            r[f"loss_{condition_titles.get(cond, cond)}"] = row[loss_col]
    # Fitted beta0 per condition
    for cond in experiment_conditions:
        title = condition_titles.get(cond, cond)
        r[f"beta0_{title}"] = float(jm.φ[cond].β0)
    summary_rows.append(r)

summary_df = pd.DataFrame(summary_rows)
print(f"{summary_df['converged'].sum()}/{len(summary_df)} models converged\n")
summary_df.round(3)

19/20 models converged

	dataset	fusionreg	converged	fit_time	alpha	loss_Delta	loss_BA.1	loss_BA.2	beta0_Delta	beta0_BA.1	beta0_BA.2
0	rep_1	0.000	True	604	20.917	0.226	0.233	0.234	-1.383	-1.384	-1.436
1	rep_1	0.000	True	629	20.721	0.226	0.232	0.233	-1.373	-1.380	-1.421
2	rep_1	0.000	True	656	20.673	0.226	0.231	0.234	-1.379	-1.386	-1.425
3	rep_1	0.000	True	552	20.450	0.228	0.230	0.238	-1.374	-1.396	-1.424
4	rep_1	0.000	True	294	19.980	0.235	0.234	0.253	-1.299	-1.323	-1.422
5	rep_1	0.000	True	569	17.786	0.271	0.232	0.281	-1.174	-1.191	-1.228
6	rep_1	0.000	True	226	15.177	0.359	0.235	0.325	-0.909	-0.957	-1.036
7	rep_1	0.000	True	497	14.535	0.456	0.232	0.366	-0.917	-0.924	-0.974
8	rep_1	0.001	True	429	14.165	0.471	0.232	0.399	-0.884	-0.889	-0.960
9	rep_1	0.001	False	1321	4.082	0.511	0.242	0.434	1.615	1.578	1.638
10	rep_2	0.000	True	829	21.582	0.266	0.190	0.190	-1.431	-1.431	-1.481
11	rep_2	0.000	True	815	21.508	0.266	0.190	0.190	-1.429	-1.430	-1.482
12	rep_2	0.000	True	788	20.868	0.266	0.189	0.191	-1.396	-1.399	-1.448
13	rep_2	0.000	True	707	20.988	0.267	0.187	0.193	-1.417	-1.426	-1.468
14	rep_2	0.000	True	308	20.783	0.277	0.190	0.208	-1.356	-1.388	-1.532
15	rep_2	0.000	True	187	19.274	0.307	0.194	0.236	-1.214	-1.260	-1.464
16	rep_2	0.000	True	186	17.151	0.386	0.195	0.280	-1.029	-1.103	-1.252
17	rep_2	0.000	True	520	15.154	0.513	0.189	0.337	-0.959	-0.982	-1.031
18	rep_2	0.001	True	263	14.312	0.559	0.191	0.382	-0.873	-0.895	-0.957
19	rep_2	0.001	True	278	13.869	0.564	0.191	0.409	-0.813	-0.850	-0.907

conv_data = model_collection.convergence_trajectory_df(
    id_vars=("dataset_name", "fusionreg")
)
multidms.plot.convergence_trajectory(
    conv_data,
    id_cols=["dataset_name", "fusionreg"],
    title="Convergence trajectories",
)

Shift sparsity

Fraction of shift parameters that are exactly zero, across the regularization grid. Uses the ModelCollection.shift_sparsity method which returns an interactive Altair chart faceted by dataset and shift parameter.

sparsity_chart, sparsity_data = model_collection.shift_sparsity(return_data=True)
sparsity_chart

cache miss - this could take a moment

Replicate parameter correlations

Correlation of mutation parameters (beta, shift) between replicates across the regularization grid. Uses the ModelCollection.mut_param_dataset_correlation method which returns an interactive Altair chart.

corr_chart, corr_data = model_collection.mut_param_dataset_correlation(
    return_data=True,
    times_seen_threshold=1,
)
corr_chart

cache miss - this could take a moment

Global epistasis plots

Global epistasis (GE) landscape at the chosen lasso strength, showing the fitted sigmoid mapping from latent to observed phenotype.

from IPython.display import display, Image
import tempfile

for ds_name in fit_collection_df["dataset_name"].unique():
    representative = (
        model_collection.fit_models
        .query(f"fusionreg == {lasso_choice} and dataset_name == '{ds_name}'")
    )
    if len(representative) == 0:
        print(f"No model at lasso={lasso_choice} for {ds_name}")
        continue
    model = representative.iloc[0].model
    print(f"{ds_name} (fusionreg={lasso_choice}):")
    chart = model.plot_ge_landscape()
    with tempfile.NamedTemporaryFile(suffix=".png") as tmp:
        chart.save(tmp.name, format="png", scale_factor=2)
        display(Image(filename=tmp.name))

rep_1 (fusionreg=8e-05):

rep_2 (fusionreg=8e-05):

Export mutations DataFrame

Merge mutation parameters from both replicates at the chosen lasso strength.

fit_dict = {}
for _, row in model_collection.fit_models.query(
    f"fusionreg == {lasso_choice}"
).iterrows():
    fit_dict[row.dataset_name] = row.model

mutations_df = combine_replicate_muts(fit_dict, times_seen_threshold=1)

mutations_df["sense"] = np.where(
    mutations_df["muts"].str.contains("*", regex=False),
    "stop",
    "nonsynonymous",
)

print(f"mutations_df: {len(mutations_df):,} mutations")
print(f"  nonsynonymous: {(mutations_df['sense'] == 'nonsynonymous').sum():,}")
print(f"  stop: {(mutations_df['sense'] == 'stop').sum():,}")
mutations_df.head()

mutations_df: 5,809 mutations
  nonsynonymous: 5,654
  stop: 155

	mutation	wts	sites	muts	rep_1_beta_Delta	rep_2_beta_Delta	avg_beta_Delta	rep_1_beta_Omicron_BA1	rep_2_beta_Omicron_BA1	avg_beta_Omicron_BA1	rep_1_beta_Omicron_BA2	rep_2_beta_Omicron_BA2	avg_beta_Omicron_BA2	rep_1_shift_Delta	rep_2_shift_Delta	avg_shift_Delta	rep_1_shift_Omicron_BA2	rep_2_shift_Omicron_BA2	avg_shift_Omicron_BA2	sense
0	M1I	M	1	I	-1.469682	-1.460722	-1.465202	-1.469682	-1.460722	-1.465202	-1.469682	-1.460722	-1.465202	0.000000	0.000000	0.000000	0.000000	0.0	0.000000	nonsynonymous
1	F2L	F	2	L	0.137920	0.112844	0.125382	0.137920	0.112720	0.125320	-0.058976	0.112720	0.026872	0.000000	0.000125	0.000062	-0.196895	0.0	-0.098448	nonsynonymous
2	F2S	F	2	S	0.115862	0.216455	0.166158	0.115862	-0.043240	0.036311	0.115862	-0.043240	0.036311	0.000000	0.259695	0.129847	0.000000	0.0	0.000000	nonsynonymous
3	F2V	F	2	V	0.190730	-0.023614	0.083558	0.190730	-0.023614	0.083558	0.190730	-0.023614	0.083558	0.000000	0.000000	0.000000	0.000000	0.0	0.000000	nonsynonymous
4	V3A	V	3	A	-0.058728	-0.095125	-0.076926	0.018800	-0.025574	-0.003387	0.018800	-0.025574	-0.003387	-0.077528	-0.069551	-0.073539	0.000000	0.0	0.000000	nonsynonymous

Save outputs

groupby = ("dataset_name", "fusionreg")
collection_muts_df = model_collection.split_apply_combine_muts(
    groupby=groupby,
    times_seen_threshold=1,
)

mutations_df.to_csv(os.path.join(output_dir, "mutations_df.csv"), index=False)
print(f"Saved mutations_df.csv ({len(mutations_df):,} rows)")

# `split_apply_combine_muts` returns a MultiIndexed frame (mutation +
# groupby keys). Writing it with index=False DISCARDS dataset_name and
# fusionreg rather than relocating them, leaving no way to tell which
# replicate or which lasso weight a row belongs to. reset_index() first.
collection_muts_df = collection_muts_df.reset_index()
assert {"dataset_name", "fusionreg"}.issubset(collection_muts_df.columns), (
    "collection_muts.csv lost its grouping keys: "
    f"{sorted(collection_muts_df.columns)}"
)
collection_muts_df.to_csv(os.path.join(output_dir, "collection_muts.csv"), index=False)
print(f"Saved collection_muts.csv ({len(collection_muts_df):,} rows)")

sparsity_data.to_csv(os.path.join(output_dir, "fit_sparsity.csv"), index=False)
print(f"Saved fit_sparsity.csv ({len(sparsity_data)} rows)")

corr_data.to_csv(os.path.join(output_dir, "library_replicate_correlation.csv"), index=False)
print(f"Saved library_replicate_correlation.csv ({len(corr_data)} rows)")

Saved mutations_df.csv (5,809 rows)
Saved collection_muts.csv (125,150 rows)
Saved fit_sparsity.csv (80 rows)
Saved library_replicate_correlation.csv (80 rows)

Export convergence diagnostics

Spike could not previously answer “did the fits converge?” from artifacts: summary_df and the trajectory frame were computed here but never persisted, and fit_models.ipynb reports crashes, not convergence. A fit that exhausts its outer iteration budget without reaching tol is counted as a success there.

Two CSVs close that gap:

• ``fit_convergence.csv`` — one row per (dataset_name, fusionreg) carrying converged, the final objective error, the sweep count, and the drift columns below.

• ``convergence_trajectory.csv`` — the tidy per-sweep trace, which is also the data behind manuscript figure S16.

argmin_sweep is the load-bearing column. A fit whose objective minimum sits away from its final sweep has drifted past that minimum — an optimizer artifact, not a scientific signal. Reading such a fit’s λ as a real turnover is a mistake, so every lasso rung is checked this way rather than by final objective value alone.

conv_traj = model_collection.convergence_trajectory_df(
    id_vars=("dataset_name", "fusionreg")
)
conv_traj.to_csv(
    os.path.join(output_dir, "convergence_trajectory.csv"), index=False
)
print(f"Saved convergence_trajectory.csv ({len(conv_traj):,} rows)")

# Per-fit convergence + drift. drift_frac measures how far the objective climbs
# above its own minimum by the final sweep; Phase 1 dropped a lasso rung on
# exactly this signal.
conv_rows = []
for _, row in model_collection.fit_models.iterrows():
    traj = row.model.convergence_trajectory_df
    obj = traj["objective_total_trajectory"]
    argmin_sweep = int(obj.idxmin())
    objective_at_argmin = float(obj.min())
    objective_final = float(obj.iloc[-1])
    conv_rows.append(
        {
            "dataset_name": row.dataset_name,
            "fusionreg": row.fusionreg,
            "converged": bool(row.model.converged),
            "final_objective_error": float(
                traj["objective_error_trajectory"].iloc[-1]
            ),
            "n_outer_sweeps": len(traj),
            "tol": row.model._fit_tol,
            "argmin_sweep": argmin_sweep,
            "objective_at_argmin": objective_at_argmin,
            "objective_final": objective_final,
            "drift_frac": (objective_final - objective_at_argmin)
            / abs(objective_at_argmin),
        }
    )

fit_convergence = pd.DataFrame(conv_rows).sort_values(
    ["dataset_name", "fusionreg"]
)
fit_convergence.to_csv(
    os.path.join(output_dir, "fit_convergence.csv"), index=False
)

n_conv = int(fit_convergence["converged"].sum())
print(
    f"Saved fit_convergence.csv — {n_conv}/{len(fit_convergence)} fits converged"
)

# A fit at the outer cap has not converged, whatever any success counter says.
at_cap = fit_convergence[fit_convergence["n_outer_sweeps"] >= fit_config["maxiter"]]
if len(at_cap):
    print(f"\nHIT OUTER maxiter={fit_config['maxiter']}:")
    print(at_cap.to_string(index=False))

drifting = fit_convergence[fit_convergence["drift_frac"] > 0.05]
if len(drifting):
    print("\nDRIFTED past objective minimum (drift_frac > 0.05):")
    print(
        drifting[
            ["dataset_name", "fusionreg", "argmin_sweep", "n_outer_sweeps", "drift_frac"]
        ].to_string(index=False)
    )

Saved convergence_trajectory.csv (4,028 rows)
Saved fit_convergence.csv — 19/20 fits converged

HIT OUTER maxiter=500:
dataset_name  fusionreg  converged  final_objective_error  n_outer_sweeps      tol  argmin_sweep  objective_at_argmin  objective_final  drift_frac
       rep_1    0.00128      False               0.000174             500 0.000001           105             1.149385         1.266373    0.101783

DRIFTED past objective minimum (drift_frac > 0.05):
dataset_name  fusionreg  argmin_sweep  n_outer_sweeps  drift_frac
       rep_1    0.00128           105             500    0.101783

Export the global epistasis landscape

Figure S17 needs per-variant latent phenotypes, and no existing CSV carries one — plot_ge_landscape() derives its scatter straight from the pickle. The figures rule must not load fit_collection.pkl (1.76 GB, ~7 GB resident), so the export happens here, in the rule that already has it open.

Exported at lasso_choice only, for both replicates: the per-variant scatter is large, and no figure needs it at the other rungs.

Both global-epistasis spaces are exported, because S17 plots both:

• space="fitness" — the shared sigmoid \(g(\phi)\), one curve for all conditions.

• space="func_score" — the per-condition curve \(\alpha \cdot (g(\phi) - g(\phi_{wt}))\), which is what the measured functional scores are on the scale of.

The site_map is exported alongside them. It records the wildtype residue at every site for each condition, and Figure 4 needs it for three distinct marks: the black x at the reference wildtype, the filled circle where a homolog’s wildtype differs from the reference, and the triangles above the per-site scatter at non-identical sites. It is a property of Data, not of the fit, so it is identical across fits — written once, from the first model.

ge_variants, ge_curves, ge_fs_curves, ge_params = [], [], [], []

for _, row in model_collection.fit_models.query(
    f"fusionreg == {lasso_choice}"
).iterrows():
    variants_df, curve_df = row.model.get_ge_landscape_df(space="fitness")
    # Same variants frame both times; only the curve differs by space.
    _, fs_curve_df = row.model.get_ge_landscape_df(space="func_score")
    params_df = row.model.get_ge_params_df()
    for frame, sink in (
        (variants_df, ge_variants),
        (curve_df, ge_curves),
        (fs_curve_df, ge_fs_curves),
        (params_df, ge_params),
    ):
        sink.append(frame.assign(dataset_name=row.dataset_name))

ge_variants_df = pd.concat(ge_variants, ignore_index=True)
ge_curve_df = pd.concat(ge_curves, ignore_index=True)
ge_fs_curve_df = pd.concat(ge_fs_curves, ignore_index=True)
ge_params_df = pd.concat(ge_params, ignore_index=True)

# wildtype_latent == beta0 + bundle_sum is an exact algebraic identity, so any
# failure here is a real bug (mis-specified reference or corrupted parameter
# extraction) rather than a tolerance judgement. Those failure modes produce
# perfectly converged and completely wrong figures.
residual = (
    ge_params_df["wildtype_latent"]
    - (ge_params_df["beta0"] + ge_params_df["bundle_sum"])
).abs()
assert (residual < 1e-6).all(), f"wildtype_latent identity broken: max {residual.max()}"

ref_rows = ge_params_df[ge_params_df["condition"] == spike["reference"]]
assert (ref_rows["bundle_sum"] == 0).all(), "reference condition has a non-empty bundle"
assert (ref_rows["n_bundle_mutations"] == 0).all(), (
    "reference condition has bundle mutations"
)

# The site map records each condition's wildtype residue per site. It is a
# property of the sequences, not of the fit -- but the replicates do NOT cover
# an identical site set (rep_1 has 1237 sites, rep_2 has 1238), because each
# observes a slightly different mutation set. What must hold, and what Figure 4
# actually depends on, is that they agree on the wildtype residue wherever they
# overlap. Assert that, then take the union so no site is silently dropped.
_site_maps = {
    row.dataset_name: row.model.data.site_map
    for _, row in model_collection.fit_models.query(
        f"fusionreg == {lasso_choice}"
    ).iterrows()
}
_names = list(_site_maps)
_first = _site_maps[_names[0]]
for _name in _names[1:]:
    _other = _site_maps[_name]
    _shared = _first.index.intersection(_other.index)
    _clash = (_first.loc[_shared] != _other.loc[_shared]).any(axis=1)
    assert not _clash.any(), (
        f"{_names[0]} and {_name} disagree on the wildtype residue at sites "
        f"{list(_shared[_clash])}"
    )

site_map_df = (
    pd.concat(_site_maps.values())
    .groupby(level=0)
    .first()
    .rename_axis("sites")
    .reset_index()
)
print(
    f"site_map: {len(site_map_df):,} sites "
    f"(union over {', '.join(f'{k}={len(v)}' for k, v in _site_maps.items())})"
)

for name, frame in (
    ("ge_landscape_variants", ge_variants_df),
    ("ge_landscape_curve", ge_curve_df),
    ("ge_landscape_curve_func_score", ge_fs_curve_df),
    ("ge_params", ge_params_df),
    ("site_map", site_map_df),
):
    frame.to_csv(os.path.join(output_dir, f"{name}.csv"), index=False)
    print(f"Saved {name}.csv ({len(frame):,} rows)")

site_map: 1,240 sites (union over rep_1=1237, rep_2=1238)
Saved ge_landscape_variants.csv (279,017 rows)
Saved ge_landscape_curve.csv (400 rows)
Saved ge_landscape_curve_func_score.csv (1,200 rows)
Saved ge_params.csv (6 rows)
Saved site_map.csv (1,240 rows)