Cross-Validation

Evaluate model generalization via 80/20 train/test cross-validation across the fusion-regularization grid.

Outline

1. Load training functional scores and pipeline config

2. Split data into train (80%) and test (20%) sets per replicate

3. Fit models on training data across the fusionreg grid

4. Evaluate validation loss on held-out test data

5. Visualize training vs validation loss

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
from multidms.model_collection import ModelCollection, fit_models
from multidms.utils import explode_params_dict

from _common import load_config, build_fit_params

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"]
reference = spike["reference"]

seed = config["seed"]
train_frac = config["train_frac"]
condition_titles = spike["condition_titles"]

if output_dir is None:
    output_dir = spike.get("output_dir", "results")
os.makedirs(output_dir, exist_ok=True)

Load training functional scores

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

Loaded 303,007 variants

Train/test split

Split functional scores 80/20 per replicate. Wildtype variants are kept in the training set.

train_datasets = []
test_data = {}

for rep_num, df_rep in func_score_df.groupby("replicate"):
    # Aggregate per (condition, aa_substitutions)
    df_agg = (
        df_rep.groupby(["condition", "aa_substitutions"], dropna=False)
        .agg({"func_score": "mean"})
        .reset_index()
    )

    wt_mask = df_agg["aa_substitutions"].str.strip() == ""
    wt_rows = df_agg[wt_mask]
    non_wt_rows = df_agg[~wt_mask].sample(frac=1, random_state=seed)

    n_train = int(len(non_wt_rows) * train_frac)
    train_split = pd.concat([wt_rows, non_wt_rows.iloc[:n_train]])
    test_split = non_wt_rows.iloc[n_train:]

    name = f"rep_{rep_num}"
    train_datasets.append(
        multidms.Data(
            train_split,
            reference=reference,
            alphabet=multidms.AAS_WITHSTOP_WITHGAP,
            assert_site_integrity=False,
            verbose=False,
            name=name,
        )
    )
    test_data[name] = test_split
    print(f"{name}: {len(train_split):,} train, {len(test_split):,} test")

rep_1: 125,836 train, 31,459 test
rep_2: 116,570 train, 29,142 test

Fit CV models

cv_fitting_params = build_fit_params(fit_config, train_datasets)

n_models = len(explode_params_dict(cv_fitting_params))
cfg_n_processes = fit_config.get("n_processes")

if cfg_n_processes is None:
    n_processes = min(os.cpu_count() // 2, n_models)
else:
    n_processes = min(int(cfg_n_processes), n_models)
n_processes = max(n_processes, 1)

print(f"Fitting {n_models} CV models with n_processes={n_processes}")

n_fit, n_failed, fit_collection_cv = fit_models(
    cv_fitting_params, n_processes=n_processes
)

for col in fit_collection_cv.columns:
    if fit_collection_cv[col].apply(lambda x: isinstance(x, dict)).any():
        fit_collection_cv[col] = fit_collection_cv[col].apply(str)

print(f"CV: fit {n_fit} models, {n_failed} failed")

Fitting 20 CV models with n_processes=6
CV: fit 20 models, 0 failed

Evaluate validation loss

mc_cv = ModelCollection(fit_collection_cv)

# Filter test variants to those with known mutations only
filtered_test = {}
for name, test_df in test_data.items():
    model = (
        mc_cv.fit_models.query(f"dataset_name == '{name}'")
        .iloc[0]
        .model
    )
    known_muts = set(model.data.mutations)

    def all_muts_known(subs):
        if pd.isna(subs) or subs.strip() == "":
            return True
        return all(m in known_muts for m in subs.split())

    mask = test_df["aa_substitutions"].apply(all_muts_known)
    filtered_test[name] = test_df[mask]
    n_dropped = (~mask).sum()
    if n_dropped > 0:
        print(f"{name}: dropped {n_dropped}/{len(test_df)} test variants with unseen mutations")

mc_cv.add_eval_loss(filtered_test, overwrite=True)

rep_1: dropped 6062/31459 test variants with unseen mutations
rep_2: dropped 5413/29142 test variants with unseen mutations

Visualize CV loss

cv_data = mc_cv.fit_models.copy()
cv_data = (
    cv_data.melt(
        id_vars=["fusionreg", "dataset_name"],
        value_vars=["total_loss_training", "total_loss_validation"],
        var_name="dataset",
        value_name="loss",
    )
    .assign(dataset=lambda x: x["dataset"].str.replace("total_loss_", ""))
)

# total_loss_{training,validation} are already per-variant averages
# (functional_score_loss uses .mean() as of PR #230).
# Keep `mean_loss` as an alias for downstream CSV consumers.
cv_data["mean_loss"] = cv_data["loss"]
cv_data["fusionreg_cat"] = cv_data["fusionreg"].astype(str)

# Plot
fig, ax = plt.subplots(figsize=(6, 4))
for (rep, ds), grp in cv_data.groupby(["dataset_name", "dataset"]):
    grp = grp.sort_values("fusionreg")
    marker = "o" if "rep_1" in rep else "s"
    color = "tab:blue" if ds == "training" else "tab:orange"
    ax.plot(
        grp["fusionreg_cat"], grp["mean_loss"],
        marker=marker, color=color, alpha=0.8,
        label=f"{rep} {ds}",
    )

ax.set_xlabel("Lasso penalty (\u03bb)")
ax.set_ylabel("Mean Huber loss per variant")
ax.set_title("Cross-validation loss")
ax.legend(fontsize=8, frameon=False)
ax.spines[["top", "right"]].set_visible(False)
plt.tight_layout()
plt.show()

Save the CV fit collection

Persist the fitted CV models the same way fit_models persists the full-data collection. Written after add_eval_loss so the pickle carries validation losses alongside training losses.

Without this the CV models exist only in this kernel and are lost on exit, leaving no way to audit them after the fact. The fit N models, M failed line above counts crashes, not convergence — a fit that exhausts its iteration budget without reaching tol is reported as a success there. Only model.converged and model.convergence_trajectory_df distinguish the two, and both need the model objects.

cv_collection_path = os.path.join(output_dir, "cv_fit_collection.pkl")
with open(cv_collection_path, "wb") as f:
    pickle.dump(mc_cv.fit_models, f)
print(f"Saved {cv_collection_path} ({len(mc_cv.fit_models)} models)")

# Convergence per fit. `fit_models` reports failures, not convergence, so
# report it here where the model objects are still in hand.
cv_convergence = pd.DataFrame(
    [
        {
            "dataset_name": row.dataset_name,
            "fusionreg": row.fusionreg,
            "converged": bool(row.model.converged),
            "final_objective_error": (
                float(
                    row.model.convergence_trajectory_df[
                        "objective_error_trajectory"
                    ].iloc[-1]
                )
                if row.model.convergence_trajectory_df is not None
                and len(row.model.convergence_trajectory_df)
                else np.nan
            ),
            "n_outer_sweeps": (
                len(row.model.convergence_trajectory_df)
                if row.model.convergence_trajectory_df is not None
                else 0
            ),
            "tol": row.model._fit_tol,
        }
        for _, row in mc_cv.fit_models.iterrows()
    ]
).sort_values(["dataset_name", "fusionreg"])

cv_convergence.to_csv(os.path.join(output_dir, "cv_convergence.csv"), index=False)

n_converged = int(cv_convergence["converged"].sum())
print(f"CV convergence: {n_converged}/{len(cv_convergence)} fits converged")

if n_converged < len(cv_convergence):
    print("\nNOT CONVERGED:")
    print(cv_convergence[~cv_convergence["converged"]].to_string(index=False))

at_cap = cv_convergence[cv_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))

Saved results-prod-294-naive-baseline-arm/cv_fit_collection.pkl (20 models)
CV convergence: 20/20 fits converged

Save

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

Saved cross_validation_loss.csv (40 rows)