diff --git a/src/decima/model/lightning.py b/src/decima/model/lightning.py index 1c20ac7..4725d64 100644 --- a/src/decima/model/lightning.py +++ b/src/decima/model/lightning.py @@ -549,23 +549,31 @@ def predict_on_dataset( compare_func: Optional[Union[str, Callable]] = None, float_precision: str = "32", ): - preds = super().predict_on_dataset( - dataset=dataset, - device=device, - num_workers=num_workers, - batch_size=batch_size, - augment_aggfunc=augment_aggfunc, - compare_func=compare_func, - float_precision=float_precision, - ) - expression = rearrange( - preds["expression"], - "(e b) t -> e b t", - e=len(self.models), - ) + # Run each constituent model independently so that allele ordering and + # variant ordering are handled correctly by LightningModel.predict_on_dataset. + # The previous approach (super().predict_on_dataset via a shared trainer) had + # two bugs: + # 1. With batch_size=1 the 4-model concat in predict_step scrambled the allele + # dimension, computing cross-model differences instead of alt-ref LFC. + # 2. With n_seqs > 1 the subsequent "(e b) t -> e b t" rearrange treated model + # index as the outer dimension, but the actual ordering coming out of the + # shared trainer is variant-outer, mixing LFCs across models and genes. + all_preds = [ + model.predict_on_dataset( + dataset=dataset, + device=device, + num_workers=num_workers, + batch_size=batch_size, + augment_aggfunc=augment_aggfunc, + compare_func=compare_func, + float_precision=float_precision, + ) + for model in self.models + ] + expression = np.stack([p["expression"] for p in all_preds]) # (e, b, T) return { - "expression": expression.mean(axis=0), - "warnings": preds["warnings"], + "expression": expression.mean(axis=0), # (b, T) + "warnings": all_preds[0]["warnings"], "ensemble_preds": expression, } diff --git a/tests/test_lightning.py b/tests/test_lightning.py index 40287f1..df92fae 100644 --- a/tests/test_lightning.py +++ b/tests/test_lightning.py @@ -1,8 +1,9 @@ import pytest import torch +import numpy as np from decima.constants import DECIMA_CONTEXT_SIZE, MODEL_METADATA, DEFAULT_ENSEMBLE from decima.data.dataset import VariantDataset -from decima.model.lightning import LightningModel, GeneMaskLightningModel +from decima.model.lightning import LightningModel, EnsembleLightningModel, GeneMaskLightningModel from decima.model.metrics import WarningType from conftest import device @@ -73,3 +74,48 @@ def test_GeneMaskLightningModel_forward(): ).to(device) preds = model(seq) assert preds.shape == (1, metadata["num_tasks"], 1) + + +def test_EnsembleLightningModel_predict_on_dataset_matches_individual_replicates(): + """Ensemble mean must equal the mean of per-model predictions for any n_seqs. + + Two bugs caused wrong results when n_seqs > 1: + - Bug 1 (batch_size=1): allele dimension scrambled by 4-model concat in predict_step. + - Bug 2 (n_seqs > 1): rearrange '(e b) -> e b' treated model as outer dim, but the + actual ordering from predict_step is variant-outer, so models and variants were mixed. + + We mock predict_on_dataset on each constituent model to isolate and test the ensemble + averaging logic without running expensive forward passes. + """ + from unittest.mock import patch, MagicMock + + n_variants, n_tasks = 5, 10 # n_variants > 1 exercises the n_seqs > 1 bug + warnings = {"allele_mismatch_with_reference_genome": 0, "unknown": 0} + + np.random.seed(0) + preds_m0 = np.random.randn(n_variants, n_tasks).astype(np.float32) + np.random.seed(1) + preds_m1 = np.random.randn(n_variants, n_tasks).astype(np.float32) + + m0 = LightningModel(model_params={"n_tasks": n_tasks, "init_borzoi": False}, name="v1_rep0") + m1 = LightningModel(model_params={"n_tasks": n_tasks, "init_borzoi": False}, name="v1_rep1") + ensemble = EnsembleLightningModel([m0, m1]) + + sentinel_dataset = MagicMock() + + with patch.object(m0, "predict_on_dataset", return_value={"expression": preds_m0, "warnings": warnings}) as mock0, \ + patch.object(m1, "predict_on_dataset", return_value={"expression": preds_m1, "warnings": warnings}) as mock1: + result = ensemble.predict_on_dataset(sentinel_dataset, device="cpu", batch_size=4) + + # Both constituent models must have been called with the forwarded arguments + mock0.assert_called_once_with(dataset=sentinel_dataset, device="cpu", num_workers=1, batch_size=4, + augment_aggfunc="mean", compare_func=None, float_precision="32") + mock1.assert_called_once_with(dataset=sentinel_dataset, device="cpu", num_workers=1, batch_size=4, + augment_aggfunc="mean", compare_func=None, float_precision="32") + + expected_mean = (preds_m0 + preds_m1) / 2 + np.testing.assert_allclose(result["expression"], expected_mean, rtol=1e-6) + + # Per-replicate outputs must match the corresponding individual model predictions + np.testing.assert_allclose(result["ensemble_preds"][0], preds_m0, rtol=1e-6) + np.testing.assert_allclose(result["ensemble_preds"][1], preds_m1, rtol=1e-6)