Fix unit tests

Author: robin-janssen

codes/benchmark/bench_fcts.py

@@ -375,8 +375,10 @@ def evaluate_iterative_predictions(
     else:
         batch_size = conf["batch_size"]
 
-    # container for the piecewise predictions
+    # container for the piecewise predictions; seed t=0 with ground truth so errors
+    # are computed only on actual predictions for t>=1 while keeping shape intact
     iterative_preds = np.zeros_like(targets)
+    iterative_preds[:, 0, :] = targets[:, 0, :]
 
     # number of chunks
     n_chunks = (n_timesteps + iter_interval - 1) // iter_interval
@@ -429,7 +431,9 @@ def evaluate_iterative_predictions(
         preds_chunk, _ = model.predict(
             data_loader=train_loader, leave_log=True, leave_norm=True
         )
-        iterative_preds[:, start:end, :] = (
+        # We predict steps 1..(chunk_len-1) relative to the provided init state (index 0).
+        # Map these to global indices [start+1 .. end] inclusively.
+        iterative_preds[:, start + 1 : end + 1, :] = (
             preds_chunk[:, 1 : model.n_timesteps, :].detach().cpu().numpy()
         )
 

test/test_bench_compare.py

@@ -13,7 +13,10 @@ def make_metrics_for_main(surr_names, n_timesteps=4, n_quantities=3):
     for i, name in enumerate(surr_names):
         metrics[name] = {
             "timesteps": np.zeros(n_timesteps),
-            "accuracy": {"absolute_errors": np.zeros((1, n_timesteps, n_quantities))},
+            "accuracy": {
+                "absolute_errors": np.zeros((1, n_timesteps, n_quantities)),
+                "absolute_errors_log": np.zeros((1, n_timesteps, n_quantities)),
+            },
             "n_params": 123 + i,
         }
     return metrics
@@ -47,7 +50,10 @@ def make_metrics_for_dynamic(surr_names, n_timesteps=4, n_quantities=2):
     metrics = {}
     for name in surr_names:
         metrics[name] = {
-            "accuracy": {"absolute_errors": np.zeros((1, n_timesteps, n_quantities))},
+            "accuracy": {
+                "absolute_errors": np.zeros((1, n_timesteps, n_quantities)),
+                "absolute_errors_log": np.zeros((1, n_timesteps, n_quantities)),
+            },
             "gradients": {
                 "gradients": np.ones((1, n_timesteps, n_quantities)),
                 "avg_correlation": 0.5,
@@ -80,6 +86,75 @@ def make_metrics_for_generalization(surr_names):
     return base
 
 
+def make_metrics_for_interpolation(surr_names):
+    base = {}
+    for name in surr_names:
+        base[name] = {
+            "interpolation": {
+                "intervals": np.array([1, 2, 4]),
+                "model_errors": np.array([0.1, 0.2, 0.25]),
+            }
+        }
+    return base
+
+
+def make_metrics_for_extrapolation(surr_names, timesteps_len=5):
+    base = {}
+    for name in surr_names:
+        base[name] = {
+            "extrapolation": {
+                "cutoffs": np.array([2, timesteps_len]),
+                "model_errors": np.array([0.3, 0.22]),
+            }
+        }
+    return base
+
+
+def make_metrics_for_sparse(surr_names):
+    base = {}
+    for name in surr_names:
+        base[name] = {
+            "sparse": {
+                "n_train_samples": np.array([100, 50, 25]),
+                "model_errors": np.array([0.15, 0.2, 0.28]),
+            }
+        }
+    return base
+
+
+def make_metrics_for_batchsize(surr_names):
+    base = {}
+    for name in surr_names:
+        base[name] = {
+            "batch_size": {
+                "batch_elements": np.array([32, 64, 128]),
+                "model_errors": np.array([0.12, 0.11, 0.13]),
+            }
+        }
+    return base
+
+
+def make_metrics_for_UQ(surr_names, timesteps):
+    base = {}
+    T = len(timesteps)
+    for name in surr_names:
+        uq_std = np.full((1, T, 1), 0.25)
+        uq_err = np.full((1, T, 1), 0.2)
+        base[name] = {
+            "timesteps": np.array(timesteps),
+            "accuracy": {"absolute_errors_log": np.full((1, T, 1), 0.18)},
+            "UQ": {
+                "pred_uncertainty_log": uq_std,
+                "absolute_errors_log": uq_err,
+                "axis_max": 1,
+                "max_counts": 1,
+                "correlation_metrics_log": 0.4,
+                "targets_log": np.zeros((1, T, 1)),
+            },
+        }
+    return base
+
+
 @pytest.fixture(autouse=True)
 def stub_plots_and_io(monkeypatch):
     calls = []
@@ -90,12 +165,13 @@ def stub_plots_and_io(monkeypatch):
         "plot_generalization_error_comparison",
         "plot_uncertainty_over_time_comparison",
         "plot_comparative_error_correlation_heatmaps",
+        "plot_mean_deltadex_over_time_main_vs_ensemble",
+        "plot_catastrophic_detection_curves",
+        "plot_errors_over_time",
         "plot_error_distribution_comparative",
-        "plot_uncertainty_confidence",
         "plot_loss_comparison",
         "plot_loss_comparison_equal",
         "plot_loss_comparison_train_duration",
-        "plot_relative_errors",
         "plot_error_distribution_comparative",
     ]:
         monkeypatch.setattr(bf, fn, lambda *a, _n=fn, **k: calls.append((_n, a, k)))
@@ -146,22 +222,35 @@ def load(self, *args, **kw):
     ]
 
 
-def test_compare_relative_errors(stub_plots_and_io, cfg):
+def test_compare_errors(stub_plots_and_io, cfg):
     timesteps = [0.0, 1.0, 2.0]
     metrics = make_metrics_for_relative(["M1"], timesteps)
-    bf.compare_relative_errors(metrics, cfg)
+    # also include deltadex branch to verify both paths
+    metrics["M1"]["accuracy"]["absolute_errors_log"] = np.arange(
+        len(timesteps)
+    ).reshape(1, len(timesteps), 1)
+
+    bf.compare_errors(metrics, cfg)
     # mean and median come from np.mean/median over rel errors
     mean_err = np.mean(metrics["M1"]["accuracy"]["relative_errors"], axis=(0, 2))
     median_err = np.median(metrics["M1"]["accuracy"]["relative_errors"], axis=(0, 2))
-    # first call to plot_relative_errors
+    # first call to plot_errors_over_time (relative)
     _n, args, kw = stub_plots_and_io[0]
-    assert _n == "plot_relative_errors"
+    assert _n == "plot_errors_over_time"
     # args = ( mean_dict, median_dict, timesteps, cfg )
     assert pytest.approx(list(args[0].values())[0]) == mean_err
     assert pytest.approx(list(args[1].values())[0]) == median_err
     assert np.all(args[2] == timesteps)
+    assert kw.get("mode") == "relative"
     # second call
     assert stub_plots_and_io[1][0] == "plot_error_distribution_comparative"
+    assert stub_plots_and_io[1][2].get("mode") == "relative"
+
+    # third and fourth calls should be for Δdex branch
+    assert stub_plots_and_io[2][0] == "plot_errors_over_time"
+    assert stub_plots_and_io[2][2].get("mode") == "deltadex"
+    assert stub_plots_and_io[3][0] == "plot_error_distribution_comparative"
+    assert stub_plots_and_io[3][2].get("mode") == "deltadex"
 
 
 def test_compare_inference_time(stub_plots_and_io, cfg):
@@ -197,38 +286,112 @@ def test_compare_gradients(stub_plots_and_io, cfg):
     assert kw.get("show_title", False) is True
 
 
-def test_compare_UQ_and_confidence(stub_plots_and_io, cfg, monkeypatch):
-    base = make_metrics_for_generalization(["M1"])
-    # ADD a dummy timesteps array
-    base["M1"]["timesteps"] = np.array([0.0, 1.0])
+def test_compare_interpolation(stub_plots_and_io, cfg):
+    m = make_metrics_for_interpolation(["M1", "M2"])
+    bf.compare_interpolation(m, cfg)
+    name, args, kw = stub_plots_and_io[0]
+    assert name == "plot_generalization_error_comparison"
+    surrogates, intervals, model_errors, xlabel, filename, conf = args
+    assert surrogates == ["M1", "M2"]
+    assert xlabel == "Interpolation Interval"
+    assert filename == "errors_interpolation.png"
+    assert all(isinstance(arr, np.ndarray) for arr in intervals)
+    assert all(isinstance(arr, np.ndarray) for arr in model_errors)
+    assert conf is cfg
+    assert kw.get("show_title", False) is True
 
-    # stub out plot_uncertainty_confidence to return known scores
-    monkeypatch.setattr(bf, "plot_uncertainty_confidence", lambda *a, **k: {"M1": 0.42})
 
-    bf.compare_UQ(base, cfg)
+def test_compare_extrapolation(stub_plots_and_io, cfg):
+    m = make_metrics_for_extrapolation(["M1"])
+    bf.compare_extrapolation(m, cfg)
+    name, args, kw = stub_plots_and_io[0]
+    assert name == "plot_generalization_error_comparison"
+    surrogates, cutoffs, model_errors, xlabel, filename, conf = args
+    assert surrogates == ["M1"]
+    assert xlabel == "Extrapolation Cutoff"
+    assert filename == "errors_extrapolation.png"
+    assert isinstance(cutoffs[0], np.ndarray)
+    assert isinstance(model_errors[0], np.ndarray)
+    assert conf is cfg
+
 
-    # after compare_UQ, confidence_scores should exist in metrics
-    assert base["M1"]["UQ"]["confidence_scores"] == 0.42
+def test_compare_sparse(stub_plots_and_io, cfg):
+    m = make_metrics_for_sparse(["M1"])
+    bf.compare_sparse(m, cfg)
+    name, args, kw = stub_plots_and_io[0]
+    assert name == "plot_generalization_error_comparison"
+    surrogates, n_train_samples, model_errors, xlabel, filename, conf = args
+    assert surrogates == ["M1"]
+    assert xlabel == "Number of Training Samples"
+    assert filename == "errors_sparse.png"
+    assert isinstance(n_train_samples[0], np.ndarray)
+    assert isinstance(model_errors[0], np.ndarray)
+    assert conf is cfg
+
+
+def test_compare_batchsize(stub_plots_and_io, cfg):
+    m = make_metrics_for_batchsize(["M1"])
+    bf.compare_batchsize(m, cfg)
+    name, args, kw = stub_plots_and_io[0]
+    assert name == "plot_generalization_error_comparison"
+    surrogates, batch_elements, model_errors, xlabel, filename, conf = args
+    assert surrogates == ["M1"]
+    assert xlabel == "Batch Size"
+    assert filename == "errors_batch_size.png"
+    assert isinstance(batch_elements[0], np.ndarray)
+    assert isinstance(model_errors[0], np.ndarray)
+    assert conf is cfg
+
+
+def test_compare_UQ(stub_plots_and_io, cfg):
+    timesteps = [0.0, 1.0, 2.0]
+    m = make_metrics_for_UQ(["M1"], timesteps)
+    bf.compare_UQ(m, cfg)
+    names = [c[0] for c in stub_plots_and_io]
+    assert names[:4] == [
+        "plot_mean_deltadex_over_time_main_vs_ensemble",
+        "plot_uncertainty_over_time_comparison",
+        "plot_comparative_error_correlation_heatmaps",
+        "plot_catastrophic_detection_curves",
+    ]
 
 
 def test_tabular_comparison_creates_files(tmp_path, stub_plots_and_io, monkeypatch):
     metrics = {
         "M1": {
             "accuracy": {
-                "mean_squared_error": 0.1,
-                "mean_absolute_error": 0.2,
+                "root_mean_squared_error_real": 0.1,
+                "mean_absolute_error_real": 0.2,
+                "median_absolute_error_real": 0.15,
+                "percentile_absolute_error_real": 0.25,
+                "root_mean_squared_error_log": 0.05,
+                "mean_absolute_error_log": 0.04,
+                "median_absolute_error_log": 0.035,
+                "percentile_absolute_error_log": 0.06,
                 "mean_relative_error": 0.3,
+                "median_relative_error": 0.25,
+                "percentile_relative_error": 0.35,
                 "main_model_epochs": 4,
                 "main_model_training_time": 7.0,
+                "error_percentile": 99,
             }
         },
         "M2": {
             "accuracy": {
-                "mean_squared_error": 0.01,
-                "mean_absolute_error": 0.02,
-                "mean_relative_error": 0.03,
+                "root_mean_squared_error_real": 0.08,
+                "mean_absolute_error_real": 0.18,
+                "median_absolute_error_real": 0.14,
+                "percentile_absolute_error_real": 0.22,
+                "root_mean_squared_error_log": 0.04,
+                "mean_absolute_error_log": 0.03,
+                "median_absolute_error_log": 0.028,
+                "percentile_absolute_error_log": 0.05,
+                "mean_relative_error": 0.25,
+                "median_relative_error": 0.2,
+                "percentile_relative_error": 0.3,
                 "main_model_epochs": 5,
                 "main_model_training_time": 3.0,
+                "error_percentile": 99,
             }
         },
     }
@@ -244,6 +407,7 @@ def test_tabular_comparison_creates_files(tmp_path, stub_plots_and_io, monkeypat
         "sparse": {"enabled": False},
         "batch_scaling": {"enabled": False},
         "verbose": False,
+        "iterative": False,
     }
 
     # run inside tmp_path

test/test_bench_main.py

@@ -14,7 +14,7 @@ def __init__(self, *, train_duration=2.5, n_quantities=2):
     def load(self, training_id, surr_name, model_identifier):
         self.load_calls.append((training_id, surr_name, model_identifier))
 
-    def predict(self, *, data_loader, leave_log=None, leave_norm=None):
+    def predict(self, data_loader, leave_log=None, leave_norm=None):
         """
         Return preds and targets of shape [1, 3, n_quantities] where
         preds are always 2x targets.
@@ -193,7 +193,7 @@ def __init__(self, n_timesteps, n_quantities):
         def load(self, training_id, surr_name, model_identifier):
             self.load_calls.append((training_id, surr_name, model_identifier))
 
-        def predict(self, *, data_loader, leave_log=None, leave_norm=None):
+        def predict(self, data_loader, leave_log: bool = False, leave_norm=None):
             # preds == targets == ones
             shape = (1, self.n_timesteps, self.n_quantities)
             ones = torch.ones(shape, dtype=torch.float32)
@@ -219,14 +219,15 @@ def denormalize(self, arr, leave_log=None, leave_norm=None):
         surr_name=surr,
         timesteps=timesteps,
         val_loader="dummy_val_loader",
+        val_params=None,
         conf=simple_conf,
         labels=["q1", "q2"],
     )
 
     # ensure we loaded the main model
     assert model.load_calls == [("TID", surr, f"{surr.lower()}_main")]
 
-    # since preds == targets, all errors should be zero
+    # ensure that each m
     for key in [
         "root_mean_squared_error_log",
         "mean_absolute_error_log",
@@ -235,6 +236,7 @@ def denormalize(self, arr, leave_log=None, leave_norm=None):
         "mean_absolute_error",
         "absolute_errors",
     ]:
+        print(metrics[key])
         assert metrics[key] == pytest.approx(0.0)
 
     # array shapes should be (1, T, Q)

test/test_bench_modalities.py

@@ -17,16 +17,21 @@
 class DummyModel:
     def __init__(self, device, n_quantities, n_timesteps, n_parameters, config):
         self._loads = []
+        self.n_timesteps = n_timesteps
 
     def load(self, training_id, surr_name, model_identifier):
         self._loads.append(model_identifier)
 
-    def predict(self, data_loader):
+    def predict(self, data_loader, leave_log=None, leave_norm=None):
         # targets always zero.  Shape (batch=2, timesteps=4, quantities=1).
         preds = torch.rand(2, 4, 1)
         targets = torch.rand(2, 4, 1)
         return preds, targets
 
+    def __call__(self, inputs):
+        # Return a dummy output tensor with 3 dims (B, T, Q)
+        return torch.zeros(1, self.n_timesteps, 1), None
+
 
 # Two standalone fakes: one for heatmap (returns tuple), one for all others (returns None)
 def _fake_heatmap(*args, **kwargs):
@@ -82,7 +87,8 @@ def test_modality_variations(raw_vals, cfg_key, func, main_bs, expected_nums):
         cfg["batch_size"] = [main_bs]
 
     timesteps = np.arange(4)
-    loader = object()
+    # minimal iterable loader; evaluate_batchsize will call next(iter(loader))
+    loader = [np.zeros((1,))]
     labels = ["q"] if func is evaluate_interpolation else None
 
     model = DummyModel(None, 1, len(timesteps), 0, {})
@@ -133,4 +139,4 @@ def test_modality_variations(raw_vals, cfg_key, func, main_bs, expected_nums):
         for num in expected_nums:
             assert f"{prefix} {num}" in metrics
     else:
-        assert "average_uncertainty" in metrics
+        assert "average_uncertainty_log" in metrics