Fix HyperTreeSTL short-horizon forecasts

hypertrees/models/HyperTreeSTL.py

@@ -290,10 +290,13 @@ def get_params_loss(
 
         # Calculate losses for trend and add smoothing penalties
         loss_trend = self.loss_fn(trend, y_trend)
-        smooth_d1 = torch.nanmean(torch.diff(trend, dim=0, n=1) ** 2, dim=0)
-        smooth_d2 = torch.nanmean(torch.diff(trend, dim=0, n=2) ** 2, dim=0)
-        smooth_penalty = torch.nanmean(torch.cat([smooth_d1, smooth_d2], dim=0))
-        loss_trend += smooth_penalty
+        smooth_terms = []
+        if trend.shape[0] >= 2:
+            smooth_terms.append(torch.nanmean(torch.diff(trend, dim=0, n=1) ** 2, dim=0))
+        if trend.shape[0] >= 3:
+            smooth_terms.append(torch.nanmean(torch.diff(trend, dim=0, n=2) ** 2, dim=0))
+        if smooth_terms:
+            loss_trend += torch.nanmean(torch.cat(smooth_terms, dim=0))
 
         # Loss for seasonal component
         loss_seasonality = self.loss_fn(seasonality, y_seasonality)
@@ -429,21 +432,27 @@ def _forward_default(
         slope = params[:, :, 1]
         trend_raw = intercept + slope * time_idx
 
-        # Map logit -> odd window size W in [3, max_w]
-        max_w = min(2 * m + 1, 101)
-        w_logit = params[:, :, 2]
-        w_float = (max_w - 3) * torch.sigmoid(w_logit) + 3.0
-        W = int(torch.median(torch.round(w_float)).clamp(3, max_w).item())
-        if W % 2 == 0:
-            W += 1
-
-        # Grouped conv expects channels divisible by groups.
-        # Put series in the *channel* dimension: input (1, N, T), weight (N, 1, W), groups=N.
-        k = torch.ones((N, 1, W), dtype=dtype) / W  # (N,1,W)
-        xin = trend_raw.T.contiguous().unsqueeze(0)  # (1,N,T)
-        pad = W // 2
-        xpad = torch.nn.functional.pad(xin, (pad, pad), mode="reflect")  # (1,N,T+2p)
-        trend = torch.nn.functional.conv1d(xpad, k, groups=N).squeeze(0).T  # (T,N)
+        # Map logit -> odd window size W in [3, max_w]. Reflect padding requires
+        # pad < T, so cap the smoothing window by the current sequence length.
+        # This matters during short-horizon forecasting, e.g. fcst_h=3.
+        if T <= 1:
+            trend = trend_raw
+        else:
+            max_w = min(2 * m + 1, 101, 2 * T - 1)
+            min_w = 3
+            w_logit = params[:, :, 2]
+            w_float = (max_w - min_w) * torch.sigmoid(w_logit) + float(min_w)
+            W = int(torch.median(torch.round(w_float)).clamp(min_w, max_w).item())
+            if W % 2 == 0:
+                W += 1
+
+            # Grouped conv expects channels divisible by groups.
+            # Put series in the *channel* dimension: input (1, N, T), weight (N, 1, W), groups=N.
+            k = torch.ones((N, 1, W), dtype=dtype) / W  # (N,1,W)
+            xin = trend_raw.T.contiguous().unsqueeze(0)  # (1,N,T)
+            pad = W // 2
+            xpad = torch.nn.functional.pad(xin, (pad, pad), mode="reflect")  # (1,N,T+2p)
+            trend = torch.nn.functional.conv1d(xpad, k, groups=N).squeeze(0).T  # (T,N)
 
         # Seasonality: Fourier with per-cycle zero-mean centering
         H = (self.n_params - 3) // 2
@@ -463,8 +472,10 @@ def _forward_default(
         C = (T + m - 1) // m
         pad_T = C * m - T
         if pad_T > 0:
-            tail = torch.flip(seasonality[-min(T, pad_T):, :], dims=[0])
-            S_ext = torch.cat([seasonality, tail[:pad_T]], dim=0)  # (C*m,N)
+            tail = torch.flip(seasonality, dims=[0])
+            repeats = (pad_T + T - 1) // T
+            tail = tail.repeat((repeats, 1))[:pad_T]
+            S_ext = torch.cat([seasonality, tail], dim=0)  # (C*m,N)
         else:
             S_ext = seasonality
 

tests/test_hypertree_stl.py

@@ -498,6 +498,36 @@ def test_forward_default_method(self):
         assert trend.shape == (n_samples, n_series)
         assert seasonality.shape == (n_samples, n_series)
 
+    def test_forward_default_short_horizon(self):
+        """Default STL smoothing should support short forecast horizons."""
+        model = HyperTreeSTL(period=12, num_seasonal_components=2, fcst_h=3, type="default")
+        n_samples = 3
+        n_series = 1
+        n_params = model.n_params
+
+        params = torch.randn(n_samples, n_series, n_params)
+        time_idx = torch.arange(1, n_samples + 1, dtype=torch.float32).reshape(-1, n_series)
+
+        trend, seasonality = model._forward_default(params, time_idx)
+
+        assert trend.shape == (n_samples, n_series)
+        assert seasonality.shape == (n_samples, n_series)
+
+    def test_forward_default_single_step_horizon(self):
+        """Default STL smoothing should bypass padding for a one-step horizon."""
+        model = HyperTreeSTL(period=12, num_seasonal_components=2, fcst_h=1, type="default")
+        n_samples = 1
+        n_series = 1
+        n_params = model.n_params
+
+        params = torch.randn(n_samples, n_series, n_params)
+        time_idx = torch.ones(n_samples, n_series, dtype=torch.float32)
+
+        trend, seasonality = model._forward_default(params, time_idx)
+
+        assert trend.shape == (n_samples, n_series)
+        assert seasonality.shape == (n_samples, n_series)
+
     def test_forward_method_selection(self):
         """Test that the correct forward method is selected based on type."""
         model_paper = HyperTreeSTL(type="paper")
@@ -538,6 +568,21 @@ def test_get_params_loss_requires_grad(self):
         assert params.requires_grad is True
         assert loss.requires_grad is True
 
+    @pytest.mark.parametrize("n_samples", [1, 2])
+    def test_get_params_loss_short_horizon_smoothing_is_finite(self, n_samples):
+        model = HyperTreeSTL(period=12, num_seasonal_components=2, fcst_h=n_samples, type="default")
+        n_params = model.n_params
+        predt = np.random.randn(n_samples * n_params)
+        target = torch.randn(n_samples, 1)
+        model.n_series = 1
+        time_idx = torch.arange(1, n_samples + 1, dtype=torch.float32).reshape(-1, 1)
+
+        params, loss = model.get_params_loss(predt, target, time_idx, requires_grad=True)
+
+        assert params.requires_grad is True
+        assert loss.requires_grad is True
+        assert torch.isfinite(loss)
+
     def test_calculate_gradients_and_hessians(self):
         model = HyperTreeSTL(period=6, type="default")
         n_params = model.n_params  # Should be 5 for default type