diff --git a/sdm/models/kumo/timeseries/forecast/model.py b/sdm/models/kumo/timeseries/forecast/model.py index 2c6fefd6e..86f8a3a9a 100644 --- a/sdm/models/kumo/timeseries/forecast/model.py +++ b/sdm/models/kumo/timeseries/forecast/model.py @@ -40,8 +40,7 @@ class KumoForecast(ICLModel): Optional cross-channel mixing lets forecasts use dependencies between variates after temporal encoding. The forecasting head predicts the native up to 72-step horizon directly from a 512-step history. - Forecast horizons beyond 72 steps are extended autoregressively, up to - 512 steps. + Forecast horizons beyond 72 steps are extended autoregressively. For longer observation histories, Domain-Aware Representation and Retrieval (DARR) can augment the direct forecast with patterns retrieved @@ -181,19 +180,7 @@ def _forward( f"Expected 'alpha' between zero and one (got {alpha})" ) - if y_context.size(-2) < self.model.context_size: - raise ValueError( - f"Expected at least {self.model.context_size} context rows " - f"(got {y_context.size(-2)})" - ) - horizon = x_query.size(-2) - if not 1 <= horizon <= 512: - raise ValueError( - "Number of query rows must be between 1 and 512 steps " - f"(got {horizon})" - ) - context_size = self.model.context_size prediction_size = self.model.prediction_size memory_horizon = max(prediction_size, horizon) @@ -218,6 +205,12 @@ def _forward( # Forecast: last_x = x[..., -self.model.context_size :] + if last_x.size(-1) < context_size: + last_x = torch.nn.functional.pad( + last_x, + pad=(context_size - last_x.size(-1), 0), + value=float("nan"), + ) encoded, mean, scale = self.model.encode(last_x) # [..., C, P, D] out = self.model.decode(encoded, mean, scale) diff --git a/test/models/kumo/timeseries/forecast/test_model.py b/test/models/kumo/timeseries/forecast/test_model.py index 79390666a..2eea41c86 100644 --- a/test/models/kumo/timeseries/forecast/test_model.py +++ b/test/models/kumo/timeseries/forecast/test_model.py @@ -14,10 +14,16 @@ @withCUDA +@pytest.mark.parametrize( + ("context_size", "horizon"), + [(1, 120), (37, 600), (1200, 600)], +) @pytest.mark.parametrize("cross_channel", [False, True]) @pytest.mark.parametrize("alpha", [0.0, 0.5, 1.0]) def test_forward( monkeypatch: pytest.MonkeyPatch, + context_size: int, + horizon: int, cross_channel: bool, alpha: float, device: torch.device, @@ -38,24 +44,24 @@ def test_forward( # Past and past-and-future covariates: x_context = TableTensor.from_tensor( - torch.randn(1024, 2, device=device), + torch.randn(context_size, 2, device=device), columns=["x1", "x2"], ) - x_query = TableTensor.from_tensor(torch.randn(120, 0, device=device)) + x_query = TableTensor.from_tensor(torch.randn(horizon, 0, device=device)) # Target variates: y_context = TableTensor.from_tensor( - torch.randn(1024, 3, device=device), + torch.randn(context_size, 3, device=device), columns=["y1", "y2", "y3"], ) with ( nullcontext() - if cross_channel or alpha < 1 + if cross_channel or (alpha < 1 and context_size >= 512 + horizon) else pytest.warns(RuntimeWarning, match="Historical covariates") ): out = model(x_context, y_context, x_query, alpha=alpha) - assert out.size() == (120, 3) + assert out.size() == (horizon, 3) assert out.device == device assert out.numerical.isfinite().all() assert "y1" in out.numerical_columns