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

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,
         }
 

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)