Model Fitting

Fit multidms models to spike functional-score data across a grid of fusion-regularization values, independently for each replicate.

Outline

1. Load training functional scores

2. Aggregate per (condition, aa_substitutions) within each replicate

3. Create multidms.Data objects (one per replicate)

4. Fit models across the regularization grid via fit_models()

5. Save the fit collection

import warnings

warnings.filterwarnings("ignore")

import os
import pickle
import sys

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

import pandas as pd
import multidms
from multidms.model_collection import fit_models
from multidms.utils import explode_params_dict

from _common import load_config, build_fit_params

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

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

config = load_config(config_path)
spike = config["spike"]
fit_config = spike["fitting"]
reference = spike["reference"]

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

Create Data objects

Aggregate functional scores per (condition, aa_substitutions) within each replicate, then create one multidms.Data object per replicate.

data_objects = []
for rep_num, df_rep in func_score_df.groupby("replicate"):
    df_agg = (
        df_rep.groupby(["condition", "aa_substitutions"], dropna=False)
        .agg({"func_score": "mean"})
        .reset_index()
    )
    data = multidms.Data(
        df_agg,
        alphabet=multidms.AAS_WITHSTOP_WITHGAP,
        reference=reference,
        assert_site_integrity=False,
        name=f"rep_{rep_num}",
    )
    data_objects.append(data)
    print(f"rep_{rep_num}: {len(df_agg):,} variants, conditions={data.conditions}")

rep_1: 157,295 variants, conditions=('Delta', 'Omicron_BA1', 'Omicron_BA2')
rep_2: 145,712 variants, conditions=('Delta', 'Omicron_BA1', 'Omicron_BA2')

Fit models

fitting_params = build_fit_params(fit_config, data_objects)
print("Fitting parameters:")
for k, v in fitting_params.items():
    if k != "dataset":
        print(f"  {k}: {v}")

Fitting parameters:
  maxiter: [500]
  tol: [1e-06]
  recompute_scale: [False]
  fusionreg: [0.0, 5e-06, 1e-05, 2e-05, 4e-05, 8e-05, 0.00016, 0.00032, 0.00064, 0.00128]
  l2reg: [1e-06]
  beta0_ridge: [0.01]
  scale_fusion_by_n: [False]
  ge_type: ['Sigmoid']
  ge_kwargs: [{'tol': 0.0001, 'maxiter': 10, 'maxls': 40, 'jit': True, 'verbose': False}]
  cal_kwargs: [{'tol': 0.0001, 'maxiter': 10, 'maxls': 40, 'jit': True, 'verbose': False}]
  loss_kwargs: [{'δ': 1.0}]
  warmstart: [False]
  beta0_init: [{'Omicron_BA1': 0.0, 'Delta': 0.0, 'Omicron_BA2': 0.0}]
  alpha_init: [6.0]
  share_alpha: [True]
  beta_clip_range: [(-10, 10)]

n_models = len(explode_params_dict(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} models with n_processes={n_processes} (cpus={os.cpu_count()})")

n_fit, n_failed, fit_collection_df = fit_models(
    fitting_params, n_processes=n_processes
)

# Convert dict-valued columns to strings for groupby compatibility
for col in fit_collection_df.columns:
    if fit_collection_df[col].apply(lambda x: isinstance(x, dict)).any():
        fit_collection_df[col] = fit_collection_df[col].apply(str)

print(f"Fit {n_fit} models successfully, {n_failed} failed")

Fitting 20 models with n_processes=6 (cpus=64)
Fit 20 models successfully, 0 failed

Save

output_path = os.path.join(output_dir, "fit_collection.pkl")
with open(output_path, "wb") as f:
    pickle.dump(fit_collection_df, f)
print(f"Saved {output_path} ({len(fit_collection_df)} models)")

Saved results-prod-294-naive-baseline-arm/fit_collection.pkl (20 models)

display_cols = ["dataset_name", "fusionreg", "converged", "fit_time"]
display_cols = [c for c in display_cols if c in fit_collection_df.columns]
fit_collection_df[display_cols]

	dataset_name	fusionreg	fit_time
0	rep_1	0.0	604
1	rep_1	0.000005	629
2	rep_1	0.00001	656
3	rep_1	0.00002	552
4	rep_1	0.00004	294
5	rep_1	0.00008	569
6	rep_1	0.00016	226
7	rep_1	0.00032	497
8	rep_1	0.00064	429
9	rep_1	0.00128	1321
10	rep_2	0.0	829
11	rep_2	0.000005	815
12	rep_2	0.00001	788
13	rep_2	0.00002	707
14	rep_2	0.00004	308
15	rep_2	0.00008	187
16	rep_2	0.00016	186
17	rep_2	0.00032	520
18	rep_2	0.00064	263
19	rep_2	0.00128	278