Skills › Scientific Agent Skills › timesfm-forecasting › forecast_csv.py
forecast_csv.py
232 lines · skills/timesfm-forecasting/scripts/forecast_csv.py at 5f2bc03
1#!/usr/bin/env python3
2"""Validated regular-grid CSV forecasting using the Apache-licensed TimesFM 2.5.
4Targets timesfm==3.0.2. Forecasts are nominal quantiles, not calibrated intervals.
5No model is downloaded until CSV validation and the resource preflight complete.
6"""
7from __future__ import annotations
9import argparse
10import csv
11import json
12import sys
13from importlib.metadata import version
14from pathlib import Path
16import numpy as np
17import pandas as pd
19MODEL_ID = "google/timesfm-2.5-200m-pytorch"
20MODEL_REVISION = "1d952420fba87f3c6dee4f240de0f1a0fbc790e3"
23def run_preflight() -> dict:
24 """Run the advisory resource check; stop on a failed resource threshold."""
25 from check_system import run_checks
26 report = run_checks("v2.5")
27 if not report.passed:
28 print(f"[ERROR] System check FAILED: {report.verdict_detail}")
29 sys.exit(1)
30 return report.to_dict()
33def forecast_config(horizon: int = 256, max_context: int = 1024,
34 batch_size: int = 1, nonnegative: bool = False,
35 return_backcast: bool = False):
36 """Validate and construct the release's supported quantile-head config."""
37 if not 1 <= horizon <= 1024:
38 raise ValueError("horizon must be 1..1024 with the 2.5 continuous quantile head")
39 if max_context < 1 or batch_size < 1:
40 raise ValueError("max_context and batch_size must be positive")
41 context = ((max_context + 31) // 32) * 32
42 max_horizon = ((horizon + 127) // 128) * 128
43 if context + max_horizon > 16384:
44 raise ValueError("patch-rounded context + horizon must be <= 16384")
45 import timesfm
46 return timesfm.ForecastConfig(
47 max_context=context, max_horizon=max_horizon,
48 normalize_inputs=True, per_core_batch_size=batch_size,
49 use_continuous_quantile_head=True, force_flip_invariance=True,
50 infer_is_positive=nonnegative, fix_quantile_crossing=True,
51 return_backcast=return_backcast,
52 )
55def load_model(batch_size: int = 1, horizon: int = 256, max_context: int = 1024,
56 nonnegative: bool = False, revision: str = MODEL_REVISION,
57 return_backcast: bool = False):
58 """Load pinned safetensors, without torch.compile startup overhead."""
59 import timesfm
60 config = forecast_config(horizon, max_context, batch_size, nonnegative, return_backcast)
61 print(f"Loading {MODEL_ID} at {revision}...")
62 model = timesfm.TimesFM_2p5_200M_torch.from_pretrained(
63 MODEL_ID, revision=revision, torch_compile=False,
64 )
65 model.compile(config)
66 return model
69def regular_frequency(dates: pd.Series, freq: str | None = None) -> str:
70 """Require a sorted unique complete grid; never invent a missing timestamp."""
71 idx = pd.DatetimeIndex(dates)
72 if idx.empty or idx.hasnans or not idx.is_unique or not idx.is_monotonic_increasing:
73 raise ValueError("dates must be nonempty, finite, unique, and increasing")
74 if freq is None:
75 freq = pd.infer_freq(idx) if len(idx) >= 3 else None
76 if freq is None:
77 raise ValueError("cannot infer a regular time grid; supply --freq and repair gaps")
78 expected = pd.date_range(idx[0], periods=len(idx), freq=freq)
79 if not idx.equals(expected):
80 raise ValueError("dates do not match the declared regular grid")
81 return pd.tseries.frequencies.to_offset(freq).freqstr
84def load_csv(path: str, date_col: str | None = None,
85 value_cols: list[str] | None = None, freq: str | None = None
86 ) -> tuple[pd.DataFrame, list[str], str | None]:
87 with open(path, newline="", encoding="utf-8-sig") as handle:
88 header = next(csv.reader(handle), [])
89 if len(header) != len(set(header)):
90 raise ValueError("duplicate CSV column names are not allowed")
91 df = pd.read_csv(path)
92 if df.empty:
93 raise ValueError("CSV has no observations")
94 if date_col:
95 if date_col not in df:
96 raise ValueError(f"date column {date_col!r} not found")
97 df[date_col] = pd.to_datetime(df[date_col], errors="raise")
98 df = df.sort_values(date_col, kind="stable").reset_index(drop=True)
99 df.attrs["frequency"] = regular_frequency(df[date_col], freq)
100 elif freq:
101 raise ValueError("--freq requires --date-col")
102 if value_cols is None:
103 value_cols = [c for c in df.select_dtypes(include=[np.number]) if c != date_col]
104 if not value_cols or len(value_cols) != len(set(value_cols)):
105 raise ValueError("select at least one unique numeric value column")
106 for col in value_cols:
107 if col not in df or col == date_col:
108 raise ValueError(f"value column {col!r} is missing or is the date column")
109 df[col] = pd.to_numeric(df[col], errors="raise")
110 return df, value_cols, date_col
113def prepare_inputs(df: pd.DataFrame, value_cols: list[str], missing: str = "error") -> list[np.ndarray]:
114 """Preserve grid positions. Optional interpolation stays inside each history."""
115 if missing not in {"error", "interpolate"}:
116 raise ValueError("missing must be error or interpolate")
117 inputs = []
118 for col in value_cols:
119 series = pd.to_numeric(df[col], errors="raise")
120 values = series.to_numpy(dtype=np.float32, copy=True)
121 if values.size == 0 or np.isinf(values).any():
122 raise ValueError(f"{col}: empty series or infinite values")
123 if np.isnan(values).any():
124 if missing == "error":
125 raise ValueError(f"{col}: missing values; choose an explicit missing-data policy")
126 if np.isnan(values[0]) or np.isnan(values[-1]):
127 raise ValueError(f"{col}: leading/trailing missing values change the forecast origin")
128 values = pd.Series(values).interpolate(limit_area="inside").to_numpy(dtype=np.float32)
129 if not np.isfinite(values).all():
130 raise ValueError(f"{col}: nonfinite input after preprocessing")
131 inputs.append(values)
132 return inputs
135def validate_forecast(point, quantiles, batch: int, horizon: int) -> tuple[np.ndarray, np.ndarray]:
136 point, quantiles = np.asarray(point), np.asarray(quantiles)
137 if point.shape != (batch, horizon) or quantiles.shape != (batch, horizon, 10):
138 raise ValueError("unexpected TimesFM 2.5 output shape; check return_backcast and model version")
139 if not np.isfinite(point).all() or not np.isfinite(quantiles).all():
140 raise ValueError("nonfinite model output")
141 if np.any(np.diff(quantiles[..., 1:], axis=-1) < 0):
142 raise ValueError("crossed quantiles")
143 if not np.allclose(point, quantiles[..., 5]):
144 raise ValueError("point forecast does not match q50")
145 return point, quantiles
148def forecast_series(model, df: pd.DataFrame, value_cols: list[str], horizon: int,
149 missing: str = "error") -> dict[str, dict]:
150 if horizon < 1:
151 raise ValueError("horizon must be positive")
152 inputs = prepare_inputs(df, value_cols, missing)
153 # The upstream 2.5 batching code appends padding series to the input list.
154 point, quantiles = model.forecast(horizon=horizon, inputs=list(inputs))
155 point, quantiles = validate_forecast(point, quantiles, len(value_cols), horizon)
156 return {
157 col: {"forecast": point[i].tolist(), "lower_80": quantiles[i, :, 1].tolist(),
158 "lower_60": quantiles[i, :, 2].tolist(), "median": quantiles[i, :, 5].tolist(),
159 "upper_60": quantiles[i, :, 8].tolist(), "upper_80": quantiles[i, :, 9].tolist()}
160 for i, col in enumerate(value_cols)
161 }
164def write_csv_output(results, output_path, df, date_col, horizon) -> None:
165 future_dates = None
166 if date_col:
167 freq = regular_frequency(df[date_col], df.attrs.get("frequency"))
168 future_dates = pd.date_range(df[date_col].iloc[-1], periods=horizon + 1, freq=freq)[1:]
169 rows = []
170 for col, data in results.items():
171 for h in range(horizon):
172 row = {"series": col, "step": h + 1, **{key: values[h] for key, values in data.items()}}
173 if future_dates is not None:
174 row["date"] = future_dates[h]
175 rows.append(row)
176 pd.DataFrame(rows).to_csv(output_path, index=False)
177 print(f"[OK] Wrote {len(rows)} forecast rows to {output_path}")
180def write_json_output(results, output_path) -> None:
181 with open(output_path, "w", encoding="utf-8") as handle:
182 json.dump(results, handle, indent=2, allow_nan=False)
183 print(f"[OK] Wrote forecasts for {len(results)} series to {output_path}")
186def main() -> None:
187 parser = argparse.ArgumentParser(description=__doc__)
188 parser.add_argument("input", help="Input CSV")
189 parser.add_argument("--horizon", type=int, required=True)
190 parser.add_argument("--date-col")
191 parser.add_argument("--freq", help="Pandas frequency, e.g. MS, D, h; required for <3 dates")
192 parser.add_argument("--value-cols", help="Comma-separated numeric columns")
193 parser.add_argument("--output", default="forecasts.csv")
194 parser.add_argument("--format", choices=["csv", "json"])
195 parser.add_argument("--batch-size", type=int, default=1)
196 parser.add_argument("--max-context", type=int, default=1024)
197 parser.add_argument("--missing", choices=["error", "interpolate"], default="error")
198 parser.add_argument("--nonnegative", action="store_true", help="Use only for a truly nonnegative target domain")
199 parser.add_argument("--revision", default=MODEL_REVISION, help="Immutable HF revision for reproducibility")
200 parser.add_argument("--skip-check", action="store_true")
201 args = parser.parse_args()
202 try:
203 cols = [c.strip() for c in args.value_cols.split(",")] if args.value_cols else None
204 df, cols, date_col = load_csv(args.input, args.date_col, cols, args.freq)
205 prepare_inputs(df, cols, args.missing)
206 config = forecast_config(args.horizon, args.max_context, args.batch_size, args.nonnegative)
207 if Path(args.input).resolve() == Path(args.output).resolve():
208 raise ValueError("output must differ from input")
209 if not args.skip_check:
210 run_preflight()
211 model = load_model(args.batch_size, args.horizon, args.max_context, args.nonnegative, args.revision)
212 results = forecast_series(model, df, cols, args.horizon, args.missing)
213 out_format = args.format or ("json" if args.output.lower().endswith(".json") else "csv")
214 if out_format == "json":
215 write_json_output(results, args.output)
216 else:
217 write_csv_output(results, args.output, df, date_col, args.horizon)
218 import dataclasses
219 metadata = {
220 "model": MODEL_ID, "revision": args.revision, "package": version("timesfm"),
221 "config": dataclasses.asdict(config), "horizon": args.horizon,
222 "frequency": df.attrs.get("frequency"), "missing_policy": args.missing,
223 "forecast_origin": str(df[date_col].iloc[-1]) if date_col else None,
224 "intervals": "nominal q10-q90 (80%) and q20-q80 (60%); calibration unverified",
225 }
226 Path(args.output + ".metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8")
227 except (ValueError, OSError) as exc:
228 parser.error(str(exc))
231if __name__ == "__main__":
232 main()