50 lines
1.3 KiB
Python
Executable File
50 lines
1.3 KiB
Python
Executable File
import numpy as np
|
|
import pandas as pd
|
|
|
|
from timesfm_app.core import ForecastConfig, ForecastPipeline, TimesFMService
|
|
|
|
|
|
class FakeTimesFMService(TimesFMService):
|
|
"""
|
|
Replace the model call with deterministic values for tests.
|
|
Returns predictions: [1, 2, ..., max_h] for each series in the batch.
|
|
"""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
def forecast_batch(self, past_values, freq_code: int, max_h: int) -> np.ndarray:
|
|
base = np.arange(1, max_h + 1, dtype=float)
|
|
return np.tile(base, (len(past_values), 1))
|
|
|
|
|
|
def test_pipeline_long_output_shape():
|
|
df = pd.DataFrame(
|
|
{
|
|
"date": pd.date_range("2023-01-01", periods=5, freq="ME").strftime(
|
|
"%Y-%m-%d"
|
|
),
|
|
"value": [1, 2, 3, 4, 5],
|
|
"key": ["A"] * 5,
|
|
}
|
|
)
|
|
|
|
cfg = ForecastConfig(
|
|
model_name="dummy",
|
|
device="cpu",
|
|
batch_size=2,
|
|
timesfm_freq_code=2,
|
|
pandas_freq_fallback="M",
|
|
horizon_min=1,
|
|
horizon_max=3,
|
|
)
|
|
|
|
svc = FakeTimesFMService()
|
|
pipe = ForecastPipeline(svc=svc, config=cfg)
|
|
out = pipe.run(df)
|
|
|
|
assert len(out) == 3
|
|
assert out["key"].unique().tolist() == ["A"]
|
|
assert out["step"].tolist() == [1, 2, 3]
|
|
assert np.allclose(out["forecast"].values, [1.0, 2.0, 3.0])
|