Add tests for SMT estimator

Author: kashif

test/torch/model/test_smt.py

@@ -0,0 +1,115 @@
+# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# A copy of the License is located at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the "license" file accompanying this file. This file is distributed
+# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+# express or implied. See the License for the specific language governing
+# permissions and limitations under the License.
+
+from itertools import islice
+
+import numpy as np
+import pytest
+import torch
+
+from gluonts.dataset.common import ListDataset
+from gluonts.evaluation import make_evaluation_predictions, Evaluator
+from gluonts.torch.distributions import StudentTOutput, NormalOutput
+from gluonts.torch.model.smt import SMTEstimator
+
+
+PREDICTION_LENGTH = 6
+FREQ = "h"
+
+
+@pytest.fixture(autouse=True)
+def _allow_full_torch_load(monkeypatch):
+    # torch>=2.6 defaults torch.load(weights_only=True), which rejects the
+    # pickled hyperparameters in Lightning checkpoints reloaded by the
+    # estimator. Relax it for the duration of each test.
+    orig = torch.load
+    monkeypatch.setattr(
+        torch, "load", lambda *a, **k: orig(*a, **{**k, "weights_only": False})
+    )
+
+
+def _univariate_dataset(num_series=3, length=200, seed=0):
+    rng = np.random.default_rng(seed)
+    t = np.arange(length)
+    return ListDataset(
+        [
+            {
+                "start": "2021-01-01 00:00:00",
+                "target": (
+                    10.0
+                    + 5.0 * np.sin(2 * np.pi * t / 24)
+                    + rng.standard_normal(length)
+                ),
+            }
+            for _ in range(num_series)
+        ],
+        freq=FREQ,
+    )
+
+
+def _estimator(**kwargs):
+    defaults = dict(
+        freq=FREQ,
+        prediction_length=PREDICTION_LENGTH,
+        context_length=2 * PREDICTION_LENGTH,
+        d_model=16,
+        nhead=2,
+        num_encoder_layers=1,
+        num_decoder_layers=1,
+        num_rnn_layers=1,
+        mem_tokens=2,
+        batch_size=8,
+        num_batches_per_epoch=4,
+        num_parallel_samples=20,
+        trainer_kwargs=dict(
+            max_epochs=2, accelerator="cpu", enable_progress_bar=False
+        ),
+    )
+    defaults.update(kwargs)
+    return SMTEstimator(**defaults)
+
+
+def test_smt_univariate_train_predict():
+    """SMT trains and forecasts on univariate series with the right shapes."""
+    dataset = _univariate_dataset()
+    predictor = _estimator().train(dataset)
+
+    forecasts = list(predictor.predict(dataset, num_samples=20))
+
+    assert len(forecasts) == 3
+    for forecast in forecasts:
+        assert forecast.samples.shape == (20, PREDICTION_LENGTH)
+        assert np.isfinite(forecast.mean).all()
+
+
+def test_smt_evaluation_runs():
+    """The full make_evaluation_predictions + Evaluator path produces finite
+    metrics."""
+    dataset = _univariate_dataset()
+    predictor = _estimator().train(dataset)
+
+    forecast_it, ts_it = make_evaluation_predictions(
+        dataset=dataset, predictor=predictor, num_samples=20
+    )
+    agg, _ = Evaluator(quantiles=[0.5])(list(ts_it), list(forecast_it))
+
+    assert np.isfinite(agg["MASE"])
+    assert np.isfinite(agg["mean_wQuantileLoss"])
+
+
+@pytest.mark.parametrize("distr_output", [StudentTOutput(), NormalOutput()])
+def test_smt_distr_outputs(distr_output):
+    dataset = _univariate_dataset()
+    predictor = _estimator(distr_output=distr_output).train(dataset)
+    forecasts = list(islice(predictor.predict(dataset), 1))
+    assert forecasts[0].samples.shape == (20, PREDICTION_LENGTH)