Fix EnsembleLightningModel VEP: run models independently to avoid allele/variant dimension scrambling

[avantikalal]

Problem

predict_variant_effect with model='ensemble' produced wrong log-fold-change predictions. The ensemble mean did not match the arithmetic mean of the four constituent replicate models run individually.

Confirmed by running rs2158799 (chr7:28237488 C>G) against JAZF1 in blood cell types:

Approach	Classical monocyte median LFC
model='ensemble' (broken)	−0.043
Each replicate individually, averaged	−0.268 (correct)

Root cause: two bugs in EnsembleLightningModel.predict_on_dataset

Bug 1 — batch_size=1 allele scrambling

EnsembleLightningModel.predict_step concatenates the outputs of all 4 models along dim=0 within each batch. With the default batch_size=1, ref and alt alleles land in separate batches. After model-concatenation the flattened tensor becomes:

[m0_ref, m1_ref, m2_ref, m3_ref,  m0_alt, m1_alt, m2_alt, m3_alt]

The subsequent rearrange("(b n a) t -> b n a t", a=2) in LightningModel.predict_on_dataset then groups m0_ref with m1_ref as the two "alleles" for a single variant, and the LFC expression[:,:,1,:] - expression[:,:,0,:] computes cross-model differences instead of alt−ref.

Bug 2 — n_seqs > 1 model/variant dimension swap

With batch_size=n_alleles=2 (which does pair alleles correctly), the LFC tensor after all batches are collected is ordered variant-outer:

[JAZF1_m0, JAZF1_m1, JAZF1_m2, JAZF1_m3,
 JAZF1-AS1_m0, JAZF1-AS1_m1, ...,
 CREB5_m0, ...]

EnsembleLightningModel.predict_on_dataset then applied rearrange("(e b) t -> e b t", e=4), which assumed model-outer ordering. With n_seqs=3 genes this sliced off the wrong blocks, mixing LFCs across genes and models. The ensemble mean and the save_replicates per-replicate columns were both wrong.

Bug 2 only manifests when a variant overlaps more than one gene (n_seqs > 1). Gene expression prediction (n_alleles=1) was not affected by either bug.

Fix

Replace the super().predict_on_dataset() + broken rearrange with a loop that runs each constituent model's predict_on_dataset independently, then stacks and averages:

all_preds = [model.predict_on_dataset(dataset, ...) for model in self.models]
expression = np.stack([p["expression"] for p in all_preds])  # (e, b, T)
return {"expression": expression.mean(axis=0), "ensemble_preds": expression, ...}

Each LightningModel.predict_on_dataset correctly handles allele ordering for any batch_size and any n_seqs. The explicit np.stack replaces the incorrect (e b) t -> e b t rearrange.

Test

Added test_EnsembleLightningModel_predict_on_dataset_matches_individual_replicates to tests/test_lightning.py. The test mocks predict_on_dataset on each constituent model (so no GPU forward passes needed) and asserts:

1. The ensemble mean equals the arithmetic mean of the mocked per-model predictions.
2. ensemble_preds[i] matches model i's predictions exactly.
3. Each constituent model was called with the correct forwarded arguments.

The test runs in ~5 seconds and does not require a long-running mark.

Verification

45 passed, 0 failed (full test suite)

🤖 Generated with Claude Code