Naive (independent per-condition) baseline

The joint model estimates a shared reference effect \(\beta_m\) per mutation plus a condition-specific shift \(\Delta_{d,m}\), fitting every condition together under a lasso that pulls shifts toward zero. The naive approach is the obvious alternative: fit each condition on its own, then call the difference of the two independently inferred effects a “shift”.

This notebook builds that baseline so the two can be compared on equal footing. Its output feeds manuscript Figure 3, whose panels ask a single question – do the shifts replicate across independent libraries? – of each approach in turn.

Two things make the comparison fair, and both are deliberate:

1. Equal convergence. The naive arm inherits the joint arm’s tol and maxiter. The original analysis gave the naive fits a tenth of the iteration budget, so part of its poor showing was under-convergence rather than method.

2. One shared mutation index. Both arms are filtered at times_seen_threshold=1 and compared on the same mutations. This knob moves the headline number more than anything else in the analysis.

Outline

1. Load the training functional scores

2. Fit one model per (replicate, condition)

3. Derive naive shifts by subtracting the reference condition’s betas

4. Report the per-condition scale parameters

5. Export

import warnings

warnings.filterwarnings("ignore")

import os
import sys

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

import pandas as pd

from _common import load_config
from _downstream import derive_naive_shifts, fit_naive_arm

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

reference = spike["reference"]
conditions = naive_config["conditions"]
times_seen_threshold = naive_config["times_seen_threshold"]

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

print(f"reference:  {reference}")
print(f"conditions: {conditions}")
print(f"times_seen_threshold: {times_seen_threshold}")
print(f"convergence: tol={fit_config['tol']:g} maxiter={fit_config['maxiter']}")

reference:  Omicron_BA1
conditions: ['Delta', 'Omicron_BA1', 'Omicron_BA2']
times_seen_threshold: 1
convergence: tol=1e-06 maxiter=500

Load the training data

Only the training CSV is read. fit_collection.pkl is 1.76 GB and needs roughly 7 GB of resident memory to unpickle; nothing here requires it, and loading it would make this rule far more expensive than the fits it performs.

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

replicates = sorted(func_score_df["replicate"].unique())
print(f"{len(func_score_df):,} variants across replicates {replicates}")
func_score_df.groupby(["replicate", "condition"]).size().rename("variants").to_frame()

303,007 variants across replicates [np.int64(1), np.int64(2)]

		variants
replicate	condition
1	Delta	28046
	Omicron_BA1	68699
	Omicron_BA2	60550
2	Delta	27734
	Omicron_BA1	60538
	Omicron_BA2	57440

Fit one model per (replicate, condition)

Each fit gets its own multidms.Data with reference set to that condition, so the model has a single condition and therefore no shift parameters at all. Two consequences follow:

• There is no lasso ladder. fusionreg and beta0_ridge only ever apply to non-reference conditions, so with one condition they are inert. The whole arm is six fits, not six times a ten-rung ladder, which is why it costs minutes rather than hours.

• The substitutions are already in BA.1 numbering, so each single-condition fit is expressed in the same mutation namespace as the others without any rewriting.

The convergence settings come from the joint arm’s own config. That is the fairness requirement: a baseline given a smaller budget than the method it is compared against tells you about the budget, not the method.

models, convergence_df, ge_params_df = fit_naive_arm(
    func_score_df, conditions, replicates, fit_config
)

n_converged = int(convergence_df["converged"].sum())
print(f"\n{n_converged}/{len(convergence_df)} fits converged")
convergence_df

  rep 1 / Delta: converged=True sweeps=59 alpha=11.72 (64.7s)
  rep 1 / Omicron_BA1: converged=True sweeps=63 alpha=13.06 (74.8s)
  rep 1 / Omicron_BA2: converged=True sweeps=126 alpha=12.47 (134.5s)
  rep 2 / Delta: converged=True sweeps=68 alpha=11.41 (78.0s)
  rep 2 / Omicron_BA1: converged=True sweeps=73 alpha=13.14 (97.3s)
  rep 2 / Omicron_BA2: converged=True sweeps=119 alpha=11.79 (155.0s)

6/6 fits converged

	replicate	condition	converged	n_outer_sweeps	final_objective_error	tol
0	1	Delta	True	59	9.439099e-07	0.000001
1	1	Omicron_BA1	True	63	8.843008e-07	0.000001
2	1	Omicron_BA2	True	126	1.544366e-07	0.000001
3	2	Delta	True	68	7.893149e-07	0.000001
4	2	Omicron_BA1	True	73	5.128082e-07	0.000001
5	2	Omicron_BA2	True	119	1.766770e-07	0.000001

Derive the naive shifts

For each replicate and condition,

\[\Delta^{\text{naive}}_{d,m} = \beta_{d,m} - \beta_{\text{ref},m}\]

computed on the inner join of the per-condition mutation indices.

The inner join is not a detail. Thirty-two spike sites carry different wildtype residues across these backgrounds, so at those sites the same physical substitution is written differently in each condition – T19I in one and I19T in another. A label can therefore only appear in every condition’s index when the wildtype residue agrees, which is what makes a plain intersection safe. A union would quietly pair labels that do not denote the same substitution.

naive_muts = derive_naive_shifts(
    models, reference=reference, times_seen_threshold=times_seen_threshold
)

shared = naive_muts.groupby("replicate")["mutation"].nunique()
print("Shared mutation index per replicate:")
print(shared.to_string())
print("\nThe manuscript reports 5,934 mutations at times_seen_threshold=1.")
naive_muts.head()

Shared mutation index per replicate:
replicate
1    5926
2    5969

The manuscript reports 5,934 mutations at times_seen_threshold=1.

	mutation	wts	sites	muts	replicate	condition	beta	naive_shift	times_seen
0	M1I	M	1	I	1	Delta	-0.978232	0.807336	3
1	F2L	F	2	L	1	Delta	0.077378	-0.075700	8
2	F2S	F	2	S	1	Delta	0.125780	0.005258	5
3	F2V	F	2	V	1	Delta	-0.156814	-0.386161	1
4	V3A	V	3	A	1	Delta	-0.081212	-0.121458	11

The per-condition scale parameters

With one condition, share_alpha has nothing to share with, so each fit finds its own \(\alpha\) – the scale carrying latent phenotype to functional score. The naive shift \(\beta_d - \beta_{\text{ref}}\) therefore mixes a genuine difference in effect with a difference in scale.

This is not corrected here, and the direction of its effect is measured below rather than argued, because the plausible argument gets it backwards.

That argument runs: the mismatch is stable across replicates (the same condition draws a similar \(\alpha\) in library 1 and library 2), so it adds a correlated component to both replicates’ shifts, which can only push the naive replicate correlation up – making the joint-versus-naive gap reported downstream a lower bound on the true gap.

The premise holds: \(\alpha\) really is replicate-stable. The conclusion does not. Rescaling each condition’s betas by its fitted \(\alpha\) before subtracting raises the naive replicate \(R^2\) rather than lowering it, so the scale mismatch is suppressing the naive correlation, not inflating it, and the gap is not a lower bound. The check is run below so the claim travels with its evidence.

print(ge_params_df.to_string(index=False))

alpha_spread = ge_params_df.groupby("condition")["alpha"].agg(["min", "max"])
print("\nAlpha by condition (replicate-stable if min and max are close):")
print(alpha_spread.to_string())

# Does the alpha mismatch inflate or suppress the naive replicate
# correlation? Rescale each condition's betas by its own fitted alpha before
# subtracting, and compare the resulting replicate R^2 against the
# uncorrected one. Rescaling removes the scale difference, so if the
# uncorrected number is the inflated one, this should come out LOWER.
from scipy.stats import pearsonr

alpha = ge_params_df.set_index(["replicate", "condition"])["alpha"]
wide = naive_muts.pivot_table(
    index="mutation", columns=["replicate", "condition"], values="beta"
)

print("\nNaive shift replicate R^2, uncorrected vs alpha-rescaled:")
for cond in [c for c in conditions if c != reference]:
    raw, scaled = {}, {}
    for rep in replicates:
        b_cond, b_ref = wide[(rep, cond)], wide[(rep, reference)]
        raw[rep] = b_cond - b_ref
        scaled[rep] = alpha[(rep, cond)] * b_cond - alpha[(rep, reference)] * b_ref

    raw_df = pd.DataFrame(raw).dropna()
    scaled_df = pd.DataFrame(scaled).dropna()
    raw_r2 = pearsonr(raw_df.iloc[:, 0], raw_df.iloc[:, 1])[0] ** 2
    scaled_r2 = pearsonr(scaled_df.iloc[:, 0], scaled_df.iloc[:, 1])[0] ** 2
    verdict = "suppressing" if scaled_r2 > raw_r2 else "inflating"
    print(
        f"  {cond:14s} uncorrected {raw_r2:.3f}   rescaled {scaled_r2:.3f}"
        f"   -> alpha mismatch is {verdict}"
    )

 replicate   condition     alpha     beta0  n_variants
         1       Delta 11.717550 -0.515767       28046
         1 Omicron_BA1 13.056016 -0.800739       68699
         1 Omicron_BA2 12.470536 -0.825371       60550
         2       Delta 11.405512 -0.653749       27734
         2 Omicron_BA1 13.142683 -0.821267       60538
         2 Omicron_BA2 11.792196 -0.700683       57440

Alpha by condition (replicate-stable if min and max are close):
                   min        max
condition
Delta        11.405512  11.717550
Omicron_BA1  13.056016  13.142683
Omicron_BA2  11.792196  12.470536

Naive shift replicate R^2, uncorrected vs alpha-rescaled:
  Delta          uncorrected 0.340   rescaled 0.416   -> alpha mismatch is suppressing
  Omicron_BA2    uncorrected 0.161   rescaled 0.179   -> alpha mismatch is suppressing

Export

Three tables, all long-form. The legacy analysis carried these as wide columns named {replicate}-{condition}_beta, which is easy to mis-parse; one row per (mutation, replicate, condition) says the same thing without encoding structure in column names.

outputs = {
    "naive_muts.csv": naive_muts,
    "naive_convergence.csv": convergence_df,
    "naive_ge_params.csv": ge_params_df,
}

for name, frame in outputs.items():
    path = os.path.join(output_dir, name)
    frame.to_csv(path, index=False)
    print(f"  wrote {path}  ({len(frame):,} rows)")

  wrote results-prod-294-naive-baseline-arm/naive_muts.csv  (35,685 rows)
  wrote results-prod-294-naive-baseline-arm/naive_convergence.csv  (6 rows)
  wrote results-prod-294-naive-baseline-arm/naive_ge_params.csv  (6 rows)