Cross-Validation

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

Outline

1. Load simulated functional scores and pipeline config

2. Split data into train (80%) and test (20%) sets, stratifying WT variants

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 across regularization strengths

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, all_muts_known

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

# Parameters
config_path = "config/config.yaml"
output_dir = "results-prod-sim-vpl500-tol1e5"

config = load_config(config_path)
sim = config["simulation"]
fit_config = sim["fitting"]

seed = config["seed"]
train_frac = config["train_frac"]
if output_dir is None:
    output_dir = sim["output_dir"]

os.makedirs(output_dir, exist_ok=True)
print(f"Output directory: {output_dir}")

Output directory: results-prod-sim-vpl500-tol1e5

Load functional scores

func_scores = pd.read_csv(os.path.join(output_dir, "simulated_func_scores.csv"))
func_scores["func_score_type"] = pd.Categorical(
    func_scores["func_score_type"],
    categories=["observed_phenotype", "loose_bottle", "tight_bottle"],
    ordered=True,
)
print(f"Loaded {len(func_scores)} rows")

Loaded 182256 rows

Train/test split

Split functional scores into 80% train / 20% test. All wildtype variants are kept in the training set (stratified split).

train_datasets = []
test_data = {}

for (library, measurement), fs_df in func_scores.rename(
    columns={"homolog": "condition"}
).groupby(["library", "func_score_type"]):
    if "enrichment" in str(measurement):
        continue

    wt_mask = fs_df["aa_substitutions"].fillna("").str.strip() == ""
    wt_rows = fs_df[wt_mask]
    non_wt_rows = fs_df[~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"{library}_{measurement}"

    train_split = train_split.copy()
    train_split["aa_substitutions"] = train_split[
        "aa_substitutions"
    ].fillna("")

    train_datasets.append(
        multidms.Data(
            train_split,
            reference="h1",
            alphabet=multidms.AAS_WITHSTOP_WITHGAP,
            verbose=False,
            name=name,
        )
    )
    test_data[name] = test_split

print(
    f"Created {len(train_datasets)} train datasets, "
    f"{len(test_data)} test datasets"
)

Created 6 train datasets, 6 test datasets

Fit models on training data

cv_fitting_params = build_fit_params(fit_config, train_datasets)

# Determine n_processes for parallel fitting
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} (cpus={os.cpu_count()})")

n_fit_cv, n_failed_cv, fit_collection_cv = fit_models(
    cv_fitting_params, n_processes=n_processes
)

# Convert dict-valued columns to strings for groupby compatibility
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_cv} models, {n_failed_cv} failed")

Fitting 60 CV models with n_processes=30 (cpus=64)
CV: fit 60 models, 0 failed

Evaluate validation loss

Filter test variants to exclude those with mutations unseen during training, then compute validation loss via add_eval_loss.

mc_cv = ModelCollection(fit_collection_cv)

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)

    mask = test_df["aa_substitutions"].apply(
        lambda s: all_muts_known(s, known_muts)
    )
    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)

lib_1_observed_phenotype: dropped 1416/6087 test variants with unseen mutations
lib_1_loose_bottle: dropped 1417/6087 test variants with unseen mutations
lib_1_tight_bottle: dropped 1417/6087 test variants with unseen mutations
lib_2_observed_phenotype: dropped 1464/6063 test variants with unseen mutations
lib_2_loose_bottle: dropped 1367/6063 test variants with unseen mutations
lib_2_tight_bottle: dropped 1367/6063 test variants with unseen mutations

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, 0 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 can distinguish the two, and both require 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-sim-vpl500-tol1e5/cv_fit_collection.pkl (60 models)
CV convergence: 60/60 fits converged

Visualize cross-validation loss

Training vs validation loss (mean Huber loss per variant) across the fusion-regularization grid, faceted by measurement type.

NICE_TITLES = {
    "observed_phenotype": "Observed Phenotype",
    "loose_bottle": "Loose Bottleneck",
    "tight_bottle": "Tight Bottleneck",
}
MT_ORDER = ["observed_phenotype", "loose_bottle", "tight_bottle"]

cv_data = mc_cv.fit_models.copy()
cv_data["library"] = (
    cv_data["dataset_name"].str.split("_").str[0:2].str.join("_")
)
cv_data["measurement_type"] = (
    cv_data["dataset_name"].str.split("_").str[2:4].str.join("_")
)
cv_data["measurement_type"] = pd.Categorical(
    cv_data["measurement_type"],
    categories=["observed_phenotype", "loose_bottle", "tight_bottle"],
    ordered=True,
)

cv_data = (
    cv_data.melt(
        id_vars=[
            "fusionreg",
            "library",
            "measurement_type",
            "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_", ""),
        lib_dataset=lambda x: x["library"] + " " + x["dataset"],
    )
)

# 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"]

# Convert fusionreg to categorical string for x-axis
cv_data["fusionreg_cat"] = cv_data["fusionreg"].astype(str)

# Plot
mt_vals = [
    mt for mt in MT_ORDER
    if mt in cv_data["measurement_type"].dropna().values
]
fig, axes = plt.subplots(
    1, len(mt_vals), figsize=(12, 3.5), sharey=False, squeeze=False
)
for ax, mt in zip(axes[0], mt_vals):
    sub = cv_data[cv_data["measurement_type"] == mt]
    for (ld, ds, lib), grp in sub.groupby(
        ["lib_dataset", "dataset", "library"]
    ):
        grp = grp.sort_values("fusionreg")
        marker = "o" if "lib_1" in lib else "s"
        color = "tab:blue" if ds == "training" else "tab:orange"
        ax.plot(
            grp["fusionreg_cat"],
            grp["mean_loss"],
            marker=marker,
            color=color,
            label=ld,
            alpha=0.8,
        )
    ax.set_title(NICE_TITLES.get(mt, mt), fontsize=9)
    ax.set_xlabel("lasso penalty (\u03bb)")
    ax.set_ylabel("mean Huber loss per variant")
    ax.tick_params(axis="x", rotation=90)

# Consolidated legend outside figure
for ax in axes[0]:
    leg = ax.get_legend()
    if leg:
        leg.remove()
handles, labels = axes[0][0].get_legend_handles_labels()
fig.legend(
    handles, labels, loc="center left", bbox_to_anchor=(1.0, 0.5),
    frameon=False, fontsize=7,
)

fig.suptitle("Cross-validation loss", fontsize=11)
fig.tight_layout(w_pad=3)
fig.subplots_adjust(right=0.93)
plt.show()

Save results

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 (120 rows)