Skip to content

Commit 67128b8

Browse files
committed
update
1 parent 47d0301 commit 67128b8

2 files changed

Lines changed: 43 additions & 9 deletions

File tree

‎sdm/models/kumo/timeseries/forecast/model.py‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,8 @@ class KumoForecast(ICLModel):
4040
Optional cross-channel mixing lets forecasts use dependencies between
4141
variates after temporal encoding. The forecasting head predicts the native
4242
up to 72-step horizon directly from a 512-step history.
43+
Shorter histories are left-padded with missing values, which are excluded
44+
from normalization statistics and replaced with zeros after normalization.
4345
Forecast horizons beyond 72 steps are extended autoregressively, up to
4446
512 steps.
4547
@@ -181,12 +183,6 @@ def _forward(
181183
f"Expected 'alpha' between zero and one (got {alpha})"
182184
)
183185

184-
if y_context.size(-2) < self.model.context_size:
185-
raise ValueError(
186-
f"Expected at least {self.model.context_size} context rows "
187-
f"(got {y_context.size(-2)})"
188-
)
189-
190186
horizon = x_query.size(-2)
191187
if not 1 <= horizon <= 512:
192188
raise ValueError(
@@ -218,6 +214,12 @@ def _forward(
218214

219215
# Forecast:
220216
last_x = x[..., -self.model.context_size :]
217+
if last_x.size(-1) < context_size:
218+
last_x = torch.nn.functional.pad(
219+
last_x,
220+
pad=(context_size - last_x.size(-1), 0),
221+
value=float("nan"),
222+
)
221223
encoded, mean, scale = self.model.encode(last_x) # [..., C, P, D]
222224
out = self.model.decode(encoded, mean, scale)
223225

‎test/models/kumo/timeseries/forecast/test_model.py‎

Lines changed: 35 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,12 @@
1616
@withCUDA
1717
@pytest.mark.parametrize("cross_channel", [False, True])
1818
@pytest.mark.parametrize("alpha", [0.0, 0.5, 1.0])
19+
@pytest.mark.parametrize("context_size", [1, 37, 1024])
1920
def test_forward(
2021
monkeypatch: pytest.MonkeyPatch,
2122
cross_channel: bool,
2223
alpha: float,
24+
context_size: int,
2325
device: torch.device,
2426
) -> None:
2527
monkeypatch.setattr(
@@ -38,20 +40,20 @@ def test_forward(
3840

3941
# Past and past-and-future covariates:
4042
x_context = TableTensor.from_tensor(
41-
torch.randn(1024, 2, device=device),
43+
torch.randn(context_size, 2, device=device),
4244
columns=["x1", "x2"],
4345
)
4446
x_query = TableTensor.from_tensor(torch.randn(120, 0, device=device))
4547

4648
# Target variates:
4749
y_context = TableTensor.from_tensor(
48-
torch.randn(1024, 3, device=device),
50+
torch.randn(context_size, 3, device=device),
4951
columns=["y1", "y2", "y3"],
5052
)
5153

5254
with (
5355
nullcontext()
54-
if cross_channel or alpha < 1
56+
if cross_channel or (alpha < 1 and context_size >= 632)
5557
else pytest.warns(RuntimeWarning, match="Historical covariates")
5658
):
5759
out = model(x_context, y_context, x_query, alpha=alpha)
@@ -61,3 +63,33 @@ def test_forward(
6163
assert "y1" in out.numerical_columns
6264
assert "y2" in out.numerical_columns
6365
assert "y3" in out.numerical_columns
66+
67+
if context_size < 512:
68+
padded_x_context = TableTensor.from_tensor(
69+
torch.nn.functional.pad(
70+
x_context.numerical,
71+
pad=(0, 0, 512 - context_size, 0),
72+
value=float("nan"),
73+
),
74+
columns=x_context.numerical_columns,
75+
)
76+
padded_y_context = TableTensor.from_tensor(
77+
torch.nn.functional.pad(
78+
y_context.numerical,
79+
pad=(0, 0, 512 - context_size, 0),
80+
value=float("nan"),
81+
),
82+
columns=y_context.numerical_columns,
83+
)
84+
with (
85+
nullcontext()
86+
if cross_channel
87+
else pytest.warns(RuntimeWarning, match="Historical covariates")
88+
):
89+
expected = model(
90+
padded_x_context,
91+
padded_y_context,
92+
x_query,
93+
alpha=alpha,
94+
)
95+
torch.testing.assert_close(out.numerical, expected.numerical)

0 commit comments

Comments
 (0)