Model Fitting

Fit multidms models to simulated functional-score data across a grid of fusion-regularization values.

Outline

1. Load simulated functional scores from disk

2. Create multidms.Data objects (one per library x func_score_type)

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

4. Post-process and 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 _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-sim-vpl500-tol1e5"

config = load_config(config_path)
sim = config["simulation"]
fit_config = sim["fitting"]
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 simulated 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")
func_scores.head()

Loaded 182256 rows

	library	homolog	aa_substitutions	func_score_type	func_score	pre_sample	func_score_var	pre_count	post_count	pre_count_wt	post_count_wt	pseudocount	n_aa_substitutions	variant_class	latent_phenotype
0	lib_1	h1	L33P	observed_phenotype	0.008957	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	1 nonsynonymous	5.253866
1	lib_1	h1	N3V F16S	observed_phenotype	-0.031679	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	>1 nonsynonymous	4.413085
2	lib_1	h1	NaN	observed_phenotype	0.000000	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	wildtype	5.000000
3	lib_1	h1	L2Q H31*	observed_phenotype	-5.929515	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	stop	-5.282384
4	lib_1	h1	N3Y L17N D18I E27R	observed_phenotype	-5.039159	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	>1 nonsynonymous	-1.707815

Create Data objects

One multidms.Data object per (library, func_score_type) combination.

data_objects = []
for (lib, fst), group_df in func_scores.rename(
    columns={"homolog": "condition"}
).groupby(["library", "func_score_type"]):
    df = group_df.copy()
    df["aa_substitutions"] = df["aa_substitutions"].fillna("")
    data_objects.append(
        multidms.Data(
            df,
            reference="h1",
            alphabet=multidms.AAS_WITHSTOP_WITHGAP,
            verbose=False,
            name=f"{lib}_{fst}_func_score",
        )
    )

print(f"Created {len(data_objects)} Data objects:")
for d in data_objects:
    print(f"  {d.name}")

Created 6 Data objects:
  lib_1_observed_phenotype_func_score
  lib_1_loose_bottle_func_score
  lib_1_tight_bottle_func_score
  lib_2_observed_phenotype_func_score
  lib_2_loose_bottle_func_score
  lib_2_tight_bottle_func_score

Build fitting parameters and 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-05]
  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': 100, 'maxls': 40, 'jit': True, 'verbose': False}]
  cal_kwargs: [{'tol': 0.0001, 'maxiter': 100, 'maxls': 40, 'jit': True, 'verbose': False}]
  loss_kwargs: [{'δ': 1.0}]
  warmstart: [False]
  recompute_scale: [False]
  beta0_init: [{'h1': 5.0, 'h2': 0.0}]
  alpha_init: [6.0]
  share_alpha: [True]
  beta_clip_range: [(-10, 10)]

import os
from multidms.utils import explode_params_dict

# Determine n_processes: null in config = auto-detect
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)  # ensure at least 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 60 models with n_processes=30 (cpus=64)
Fit 60 models successfully, 0 failed

Post-process and save

Extract library and measurement_type from the dataset name.

fit_collection_df = fit_collection_df.assign(
    library=(
        fit_collection_df["dataset_name"]
        .str.split("_").str[0:2].str.join("_")
    ),
    measurement_type=(
        fit_collection_df["dataset_name"]
        .str.split("_").str[2:4].str.join("_")
    ),
)

fit_collection_df["measurement_type"] = pd.Categorical(
    fit_collection_df["measurement_type"],
    categories=["observed_phenotype", "loose_bottle", "tight_bottle"],
    ordered=True,
)

print(f"Post-processed {len(fit_collection_df)} models")

Post-processed 60 models

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-sim-vpl500-tol1e5/fit_collection.pkl (60 models)

Summary

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

	dataset_name	library	measurement_type	fusionreg	fit_time
0	lib_1_observed_phenotype_func_score	lib_1	observed_phenotype	0.0	317
1	lib_1_observed_phenotype_func_score	lib_1	observed_phenotype	0.000005	314
2	lib_1_observed_phenotype_func_score	lib_1	observed_phenotype	0.00001	314
3	lib_1_observed_phenotype_func_score	lib_1	observed_phenotype	0.00002	310
4	lib_1_observed_phenotype_func_score	lib_1	observed_phenotype	0.00004	306
5	lib_1_observed_phenotype_func_score	lib_1	observed_phenotype	0.00008	278
6	lib_1_observed_phenotype_func_score	lib_1	observed_phenotype	0.00016	242
7	lib_1_observed_phenotype_func_score	lib_1	observed_phenotype	0.00032	168
8	lib_1_observed_phenotype_func_score	lib_1	observed_phenotype	0.00064	115
9	lib_1_observed_phenotype_func_score	lib_1	observed_phenotype	0.00128	98
10	lib_1_loose_bottle_func_score	lib_1	loose_bottle	0.0	311
11	lib_1_loose_bottle_func_score	lib_1	loose_bottle	0.000005	316
12	lib_1_loose_bottle_func_score	lib_1	loose_bottle	0.00001	311
13	lib_1_loose_bottle_func_score	lib_1	loose_bottle	0.00002	303
14	lib_1_loose_bottle_func_score	lib_1	loose_bottle	0.00004	298
15	lib_1_loose_bottle_func_score	lib_1	loose_bottle	0.00008	274
16	lib_1_loose_bottle_func_score	lib_1	loose_bottle	0.00016	232
17	lib_1_loose_bottle_func_score	lib_1	loose_bottle	0.00032	154
18	lib_1_loose_bottle_func_score	lib_1	loose_bottle	0.00064	121
19	lib_1_loose_bottle_func_score	lib_1	loose_bottle	0.00128	92
20	lib_1_tight_bottle_func_score	lib_1	tight_bottle	0.0	198
21	lib_1_tight_bottle_func_score	lib_1	tight_bottle	0.000005	212
22	lib_1_tight_bottle_func_score	lib_1	tight_bottle	0.00001	201
23	lib_1_tight_bottle_func_score	lib_1	tight_bottle	0.00002	203
24	lib_1_tight_bottle_func_score	lib_1	tight_bottle	0.00004	191
25	lib_1_tight_bottle_func_score	lib_1	tight_bottle	0.00008	187
26	lib_1_tight_bottle_func_score	lib_1	tight_bottle	0.00016	154
27	lib_1_tight_bottle_func_score	lib_1	tight_bottle	0.00032	128
28	lib_1_tight_bottle_func_score	lib_1	tight_bottle	0.00064	92
29	lib_1_tight_bottle_func_score	lib_1	tight_bottle	0.00128	79
30	lib_2_observed_phenotype_func_score	lib_2	observed_phenotype	0.0	289
31	lib_2_observed_phenotype_func_score	lib_2	observed_phenotype	0.000005	280
32	lib_2_observed_phenotype_func_score	lib_2	observed_phenotype	0.00001	281
33	lib_2_observed_phenotype_func_score	lib_2	observed_phenotype	0.00002	274
34	lib_2_observed_phenotype_func_score	lib_2	observed_phenotype	0.00004	254
35	lib_2_observed_phenotype_func_score	lib_2	observed_phenotype	0.00008	241
36	lib_2_observed_phenotype_func_score	lib_2	observed_phenotype	0.00016	218
37	lib_2_observed_phenotype_func_score	lib_2	observed_phenotype	0.00032	160
38	lib_2_observed_phenotype_func_score	lib_2	observed_phenotype	0.00064	123
39	lib_2_observed_phenotype_func_score	lib_2	observed_phenotype	0.00128	77
40	lib_2_loose_bottle_func_score	lib_2	loose_bottle	0.0	225
41	lib_2_loose_bottle_func_score	lib_2	loose_bottle	0.000005	222
42	lib_2_loose_bottle_func_score	lib_2	loose_bottle	0.00001	216
43	lib_2_loose_bottle_func_score	lib_2	loose_bottle	0.00002	216
44	lib_2_loose_bottle_func_score	lib_2	loose_bottle	0.00004	205
45	lib_2_loose_bottle_func_score	lib_2	loose_bottle	0.00008	192
46	lib_2_loose_bottle_func_score	lib_2	loose_bottle	0.00016	159
47	lib_2_loose_bottle_func_score	lib_2	loose_bottle	0.00032	128
48	lib_2_loose_bottle_func_score	lib_2	loose_bottle	0.00064	100
49	lib_2_loose_bottle_func_score	lib_2	loose_bottle	0.00128	65
50	lib_2_tight_bottle_func_score	lib_2	tight_bottle	0.0	130
51	lib_2_tight_bottle_func_score	lib_2	tight_bottle	0.000005	132
52	lib_2_tight_bottle_func_score	lib_2	tight_bottle	0.00001	122
53	lib_2_tight_bottle_func_score	lib_2	tight_bottle	0.00002	121
54	lib_2_tight_bottle_func_score	lib_2	tight_bottle	0.00004	114
55	lib_2_tight_bottle_func_score	lib_2	tight_bottle	0.00008	113
56	lib_2_tight_bottle_func_score	lib_2	tight_bottle	0.00016	93
57	lib_2_tight_bottle_func_score	lib_2	tight_bottle	0.00032	78
58	lib_2_tight_bottle_func_score	lib_2	tight_bottle	0.00064	68
59	lib_2_tight_bottle_func_score	lib_2	tight_bottle	0.00128	62