Remove generated output files from git tracking, add to .gitignore

- stock_forecast.py: Basic AAPL price forecast
- stock_backtest.py: V1 original price prediction backtest
- stock_backtest_optimized.py: V2 log-return + ensemble optimization
- stock_backtest_xreg.py: V3 XReg covariates (volume/RSI/SPY)
- stock_backtest_trick.py: V4 SPY trend guidance + volatility adjustment
- Visualization charts for all versions

.gitignore

@@ -8,3 +8,6 @@ datasets/
 results/
 uv.lock
 development_setup.md
+debug.log
+# Generated output files
+timesfm-forecasting/examples/*/output/

README.md

@@ -9,8 +9,10 @@ model developed by Google Research for time-series forecasting.
 *   All checkpoints:
     [TimesFM Hugging Face Collection](https://huggingface.co/collections/google/timesfm-release-66e4be5fdb56e960c1e482a6).
 *   [Google Research blog](https://research.google/blog/a-decoder-only-foundation-model-for-time-series-forecasting/).
-*   [TimesFM in BigQuery](https://cloud.google.com/bigquery/docs/timesfm-model):
+*   TimesFM in Google 1P Products:
-    an official Google product.
+    *   [BigQuery ML](https://cloud.google.com/bigquery/docs/timesfm-model): Enterprise level SQL queries for scalability and reliability.
+    *   [Google Sheets](https://workspaceupdates.googleblog.com/2026/02/forecast-data-in-connected-sheets-BigQueryML-TimesFM.html): For your daily spreadsheet. 
+    *   [Vertex Model Garden](https://pantheon.corp.google.com/vertex-ai/publishers/google/model-garden/timesfm): Dockerized endpoint for agentic calling.
 
 This open version is not an officially supported Google product.
 
@@ -21,17 +23,21 @@ This open version is not an officially supported Google product.
 -   1.0 and 2.0: relevant code archived in the sub directory `v1`. You can `pip
     install timesfm==1.3.0` to install an older version of this package to load
     them.
+## Update - June 5, 2026
+
+Updated PyPI to `timesfm=2.0.0`. See [Install](https://github.com/google-research/timesfm#from-pypi).
 
 ## Update - Apr. 9, 2026
 
 Added fine-tuning example using HuggingFace Transformers + PEFT (LoRA) — see
 [`timesfm-forecasting/examples/finetuning/`](timesfm-forecasting/examples/finetuning/).
-Also added unit tests (`tests/`), fixed per-input ridge regression in XReg to
+Also added unit tests (`tests/`) and incorporated several community fixes.
-prevent data leakage, and incorporated several community fixes.
+
+Shoutout to [@kashif](https://github.com/kashif) and [@darkpowerxo](https://github.com/darkpowerxo). 
 
 ## Update - Mar. 19, 2026
 
-Huge shoutout to [@borealBytes](https://github.com/borealBytes) for adding the support for [AGENTS](https://github.com/google-research/timesfm/blob/master/AGENTS.md)! TimesFM [SKILL.md](https://github.com/google-research/timesfm/blob/master/timesfm-forecasting/SKILL.md) is out.
+Huge shoutout to [@borealBytes](https://github.com/borealBytes) for adding the support for [AGENTS](https://github.com/google-research/timesfm/blob/master/AGENTS.md)! TimesFM [SKILL.md](https://github.com/google-research/timesfm/tree/master/timesfm-forecasting) is out.
 
 ## Update - Oct. 29, 2025
 
@@ -61,6 +67,19 @@ Since the Sept. 2025 launch, the following improvements have been completed:
 
 ### Install
 
+#### From `PyPI`
+
+```shell
+# Install the package with torch
+pip install timesfm[torch]
+# Or with Flax
+pip install timesfm[flax]
+# And when XReg is needed
+pip install timesfm[xreg]
+```
+
+#### Local Install
+
 1.  Clone the repository:
     ```shell
     git clone https://github.com/google-research/timesfm.git
@@ -79,7 +98,7 @@ Since the Sept. 2025 launch, the following improvements have been completed:
     uv pip install -e .[torch]
     # Or with flax
     uv pip install -e .[flax]
-    # Or XReg is needed
+    # And when XReg is needed
     uv pip install -e .[xreg]
     ```
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "timesfm"
-version = "2.0.0"
+version = "2.0.1"
 description = "A time series foundation model."
 authors = [
     {name = "Rajat Sen", email = "senrajat@google.com"},

src/timesfm/timesfm_2p5/timesfm_2p5_flax.py

@@ -447,6 +447,21 @@ class TimesFM_2p5_200M_flax(timesfm_2p5_base.TimesFM_2p5):
 
   model: nnx.Module = TimesFM_2p5_200M_flax_module()
 
+  def __init__(self, **kwargs):
+    self.model = TimesFM_2p5_200M_flax_module()
+
+  def load_checkpoint(self, path: str):
+    """Loads a TimesFM model from a checkpoint."""
+    if os.path.isdir(path):
+      model_file_path = path
+    else:
+      model_file_path = os.path.dirname(path)
+
+    checkpointer = ocp.StandardCheckpointer()
+    graph, state = nnx.split(self.model)
+    state = checkpointer.restore(model_file_path, state)
+    self.model = nnx.merge(graph, state)
+
   @classmethod
   def from_pretrained(
     cls,
@@ -485,10 +500,7 @@ class TimesFM_2p5_200M_flax(timesfm_2p5_base.TimesFM_2p5):
       )
       logging.info("Loading checkpoint from: %s", model_file_path)
 
-    checkpointer = ocp.StandardCheckpointer()
+    instance.load_checkpoint(model_file_path)
-    graph, state = nnx.split(instance.model)
-    state = checkpointer.restore(model_file_path, state)
-    instance.model = nnx.merge(graph, state)
     return instance
 
   def compile(

src/timesfm/timesfm_2p5/timesfm_2p5_torch.py

@@ -257,7 +257,7 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
       to_concat = [t_pf[:, -1, ...]]
       if t_ar is not None:
         to_concat.append(t_ar.reshape(1, -1, self.q))
-      torch_forecast = torch.cat(to_concat, dim=1)[..., :horizon]
+      torch_forecast = torch.cat(to_concat, dim=1)[:, :horizon, :]
       torch_forecast = torch_forecast.squeeze(0)
       outputs.append(torch_forecast.detach().cpu().numpy())
     return outputs
@@ -283,12 +283,26 @@ class TimesFM_2p5_200M_torch(
     self,
     torch_compile: bool = True,
     config: Optional[dict] = None,
+    **kwargs,
   ):
     self.model = TimesFM_2p5_200M_torch_module()
     self.torch_compile = torch_compile
     if config is not None:
       self._hub_mixin_config = config
 
+  def load_checkpoint(self, path: str, **kwargs):
+    """Loads a TimesFM model from a checkpoint directory or file."""
+    if os.path.isdir(path):
+      model_file_path = os.path.join(path, self.WEIGHTS_FILENAME)
+      if not os.path.exists(model_file_path):
+        raise FileNotFoundError(
+          f"{self.WEIGHTS_FILENAME} not found in directory {path}"
+        )
+    else:
+      model_file_path = path
+
+    self.model.load_checkpoint(model_file_path, **kwargs)
+
   @classmethod
   def _from_pretrained(
     cls,
@@ -333,7 +347,7 @@ class TimesFM_2p5_200M_torch(
 
     logging.info("Loading checkpoint from: %s", model_file_path)
     # Load the weights into the model.
-    instance.model.load_checkpoint(
+    instance.load_checkpoint(
       model_file_path, torch_compile=instance.torch_compile
     )
     return instance

src/timesfm/utils/xreg_lib.py

@@ -370,20 +370,11 @@ class BatchedInContextXRegBase:
       x_train = np.concatenate(x_train, axis=1)
       x_test = np.concatenate(x_test, axis=1)
 
-      # Normalize per-input for robustness (batch-wide normalization
+      # Normalize for robustness.
-      # would make each input's result depend on batch composition).
+      x_mean = np.mean(x_train, axis=0, keepdims=True)
-      train_splits = np.cumsum(self.train_lens)[:-1]
+      x_std = np.where((w := np.std(x_train, axis=0, keepdims=True)) > _TOL, w, 1.0)
-      test_splits = np.cumsum(self.test_lens)[:-1]
+      x_train = [(x_train - x_mean) / x_std]
-      train_parts = np.split(x_train, train_splits, axis=0)
+      x_test = [(x_test - x_mean) / x_std]
-      test_parts = np.split(x_test, test_splits, axis=0)
-      norm_train, norm_test = [], []
-      for tr, te in zip(train_parts, test_parts):
-        m = np.mean(tr, axis=0, keepdims=True)
-        s = np.where((w := np.std(tr, axis=0, keepdims=True)) > _TOL, w, 1.0)
-        norm_train.append((tr - m) / s)
-        norm_test.append((te - m) / s)
-      x_train = [np.concatenate(norm_train, axis=0)]
-      x_test = [np.concatenate(norm_test, axis=0)]
 
     # Categorical features. Encode one by one.
     one_hot_encoder = preprocessing.OneHotEncoder(
@@ -472,57 +463,58 @@ class BatchedInContextXRegLinear(BatchedInContextXRegBase):
       assert_covariate_shapes=assert_covariate_shapes,
     )
 
+    x_train = x_train_raw.copy()
+    if max_rows_per_col:
+      nrows, ncols = x_train.shape
+      if nrows > (w := ncols * max_rows_per_col):
+        subsample = jax.random.choice(
+          jax.random.PRNGKey(max_rows_per_col_sample_seed),
+          nrows,
+          (w,),
+          replace=False,
+        )
+        x_train = x_train[subsample]
+        flat_targets = flat_targets[subsample]
+
     device = jax.devices("cpu")[0] if force_on_cpu else None
+    # Runs jitted version of the solvers which are quicker at the cost of
+    # running jitting during the first time calling. Re-jitting happens whenever
+    # new (padded) shapes are encountered.
+    # Ocassionally it helps with the speed and the accuracy if we force single
+    # thread execution on cpu for accelerator machines:
+    # 1. Avoid moving data to accelarator memory.
+    # 2. Avoid precision loss if any.
+    with jax.default_device(device):
+      x_train_raw = _to_padded_jax_array(x_train_raw)
+      x_train = _to_padded_jax_array(x_train)
+      flat_targets = _to_padded_jax_array(flat_targets)
+      x_test = _to_padded_jax_array(x_test)
+      beta_hat = (
+        jnp.linalg.pinv(
+          x_train.T @ x_train + ridge * jnp.eye(x_train.shape[1]),
+          hermitian=True,
+        )
+        @ x_train.T
+        @ flat_targets
+      )
+      y_hat = x_test @ beta_hat
+      y_hat_context = x_train_raw @ beta_hat if debug_info else None
+
     outputs = []
     outputs_context = []
-    train_idx, test_idx = 0, 0
 
-    with jax.default_device(device):
+    # Reconstruct the ragged 2-dim batched forecasts from flattened linear fits.
-      for trl, tel in zip(self.train_lens, self.test_lens):
+    train_index, test_index = 0, 0
-        x_tr = x_train_raw[train_idx : train_idx + trl]
+    for train_index_delta, test_index_delta in zip(self.train_lens, self.test_lens):
-        x_te = x_test[test_idx : test_idx + tel]
+      outputs.append(np.array(y_hat[test_index : (test_index + test_index_delta)]))
-        y_tr = flat_targets[train_idx : train_idx + trl]
+      if debug_info:
-
+        outputs_context.append(
-        x_tr_fit = x_tr.copy()
+          np.array(y_hat_context[train_index : (train_index + train_index_delta)])
-        if max_rows_per_col:
-          nrows, ncols = x_tr_fit.shape
-          if nrows > (w := ncols * max_rows_per_col):
-            subsample = jax.random.choice(
-              jax.random.PRNGKey(max_rows_per_col_sample_seed),
-              nrows,
-              (w,),
-              replace=False,
-            )
-            x_tr_fit = x_tr_fit[subsample]
-            y_tr = y_tr[subsample]
-
-        x_tr_raw_j = _to_padded_jax_array(x_tr)
-        x_tr_j = _to_padded_jax_array(x_tr_fit)
-        y_tr_j = _to_padded_jax_array(y_tr)
-        x_te_j = _to_padded_jax_array(x_te)
-
-        beta_hat = (
-          jnp.linalg.pinv(
-            x_tr_j.T @ x_tr_j + ridge * jnp.eye(x_tr_j.shape[1]),
-            hermitian=True,
-          )
-          @ x_tr_j.T
-          @ y_tr_j
         )
-        outputs.append(np.array(x_te_j @ beta_hat)[:tel])
+      train_index += train_index_delta
-        if debug_info:
+      test_index += test_index_delta
-          outputs_context.append(np.array(x_tr_raw_j @ beta_hat)[:trl])
-
-        train_idx += trl
-        test_idx += tel
 
     if debug_info:
-      return (
+      return outputs, outputs_context, flat_targets, x_train, x_test
-        outputs,
-        outputs_context,
-        _to_padded_jax_array(flat_targets),
-        _to_padded_jax_array(x_train_raw),
-        _to_padded_jax_array(x_test),
-      )
     else:
       return outputs

stock_backtest.py

@@ -0,0 +1,115 @@
+import yfinance as yf
+import numpy as np
+import torch
+import timesfm
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from datetime import datetime, timedelta
+
+# 1. 获取 AAPL 数据（到今天为止）
+print("下载 AAPL 股票数据...", flush=True)
+end = datetime.now()
+start = end - timedelta(days=365)
+df = yf.download("AAPL", start=start.strftime("%Y-%m-%d"), end=end.strftime("%Y-%m-%d"), progress=False)
+close = df["Close"].values.flatten().astype(np.float32)
+dates = df.index
+
+# 2. 分割：6/15 之前为训练集，6/16 之后为实际值
+split_date = "2026-06-15"
+split_idx = None
+for i, d in enumerate(dates):
+    if str(d.date()) <= split_date:
+        split_idx = i
+# split_idx 是 <= 6/15 的最后一个索引
+train_data = close[: split_idx + 1]
+train_dates = dates[: split_idx + 1]
+actual_data = close[split_idx + 1 :]
+actual_dates = dates[split_idx + 1 :]
+
+HORIZON = len(actual_data)
+print(f"训练数据: {len(train_data)} 天（截至 {train_dates[-1].date()}）", flush=True)
+print(f"实际数据: {HORIZON} 天（{actual_dates[0].date()} ~ {actual_dates[-1].date()}）", flush=True)
+print(f"分割日收盘价: {train_data[-1]:.2f}", flush=True)
+
+# 3. 加载模型
+print()
+print("加载 TimesFM 模型...", flush=True)
+torch.set_float32_matmul_precision("high")
+model = timesfm.TimesFM_2p5_200M_torch.from_pretrained(
+    "google/timesfm-2.5-200m-pytorch", torch_compile=False
+)
+print("模型加载完成", flush=True)
+
+# 4. 编译并预测
+model.compile(
+    timesfm.ForecastConfig(
+        max_context=512,
+        max_horizon=128,
+        normalize_inputs=True,
+        use_continuous_quantile_head=True,
+        force_flip_invariance=False,
+        infer_is_positive=True,
+        fix_quantile_crossing=True,
+    )
+)
+print(f"预测未来 {HORIZON} 个交易日...", flush=True)
+pf, qf = model.forecast(horizon=HORIZON, inputs=[train_data])
+print("预测完成！", flush=True)
+
+# 5. 计算误差指标
+actual = actual_data
+pred = pf[0][:HORIZON]
+mae = np.mean(np.abs(actual - pred))
+rmse = np.sqrt(np.mean((actual - pred) ** 2))
+mape = np.mean(np.abs((actual - pred) / actual)) * 100
+
+# 方向准确率
+actual_dir = np.diff(actual)
+pred_dir = np.diff(pred)
+dir_acc = np.mean(actual_dir * pred_dir > 0) * 100
+
+# 置信区间覆盖率
+in_80 = np.mean((actual >= qf[0, :HORIZON, 1]) & (actual <= qf[0, :HORIZON, 9])) * 100
+in_40 = np.mean((actual >= qf[0, :HORIZON, 3]) & (actual <= qf[0, :HORIZON, 7])) * 100
+
+print()
+print("=== 回测结果：6/16 ~ 6/30 预测 vs 实际 ===", flush=True)
+print()
+print("  日期          实际价    预测价    误差    误差%", flush=True)
+for i in range(HORIZON):
+    err = pred[i] - actual[i]
+    err_pct = err / actual[i] * 100
+    print(
+        f"  {actual_dates[i].date()}  {actual[i]:7.2f}  {pred[i]:7.2f}  {err:+7.2f}  {err_pct:+6.2f}%",
+        flush=True,
+    )
+
+print()
+print("=== 误差指标 ===", flush=True)
+print(f"  MAE  (平均绝对误差):     ${mae:.2f}", flush=True)
+print(f"  RMSE (均方根误差):       ${rmse:.2f}", flush=True)
+print(f"  MAPE (平均绝对百分比误差): {mape:.2f}%", flush=True)
+print(f"  方向准确率:               {dir_acc:.1f}%", flush=True)
+print(f"  80%置信区间覆盖率:        {in_80:.1f}%", flush=True)
+print(f"  40%置信区间覆盖率:        {in_40:.1f}%", flush=True)
+
+# 6. 画图
+fig, ax = plt.subplots(figsize=(14, 6))
+show_n = min(60, len(train_data))
+ax.plot(range(show_n), train_data[-show_n:], label="Historical Close", color="steelblue", linewidth=1.5)
+x_actual = range(show_n, show_n + HORIZON)
+ax.plot(x_actual, actual, label="Actual", color="forestgreen", linewidth=2, marker="o", markersize=3)
+ax.plot(x_actual, pred, label="Forecast", color="tomato", linewidth=2, linestyle="--")
+ax.fill_between(x_actual, qf[0, :HORIZON, 1], qf[0, :HORIZON, 9], alpha=0.15, color="tomato", label="80% CI")
+ax.fill_between(x_actual, qf[0, :HORIZON, 3], qf[0, :HORIZON, 7], alpha=0.3, color="tomato", label="40% CI")
+ax.axvline(x=show_n - 1, color="gray", linestyle="--", alpha=0.5, label="Forecast Start")
+ax.set_title(f"AAPL Backtest: Forecast vs Actual (MAPE={mape:.2f}%, Dir.Acc={dir_acc:.0f}%)", fontsize=14)
+ax.set_xlabel("Trading Days")
+ax.set_ylabel("Price (USD)")
+ax.legend(loc="upper left")
+plt.tight_layout()
+plt.savefig("aapl_backtest.png", dpi=150)
+print()
+print("Chart saved: aapl_backtest.png", flush=True)
+print("Done!", flush=True)

stock_backtest_optimized.py

@@ -0,0 +1,184 @@
+import yfinance as yf
+import numpy as np
+import torch
+import timesfm
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from datetime import datetime, timedelta
+
+# ---------- 技术指标计算 ----------
+def calc_rsi(prices, period=14):
+    deltas = np.diff(prices)
+    gains = np.where(deltas > 0, deltas, 0.0)
+    losses = np.where(deltas < 0, -deltas, 0.0)
+    avg_gain = np.convolve(gains, np.ones(period) / period, mode="valid")
+    avg_loss = np.convolve(losses, np.ones(period) / period, mode="valid")
+    avg_loss = np.where(avg_loss == 0, 1e-10, avg_loss)
+    rs = avg_gain / avg_loss
+    rsi = 100.0 - (100.0 / (1.0 + rs))
+    # pad front to match length
+    return np.concatenate([np.full(period, 50.0), rsi])
+
+
+def calc_sma(prices, period):
+    sma = np.convolve(prices, np.ones(period) / period, mode="valid")
+    return np.concatenate([np.full(period - 1, sma[0] if len(sma) > 0 else 0.0), sma])
+
+
+# ---------- 1. 获取数据 ----------
+print("下载 AAPL 股票数据...", flush=True)
+end = datetime.now()
+start = end - timedelta(days=365)
+df = yf.download("AAPL", start=start.strftime("%Y-%m-%d"), end=end.strftime("%Y-%m-%d"), progress=False)
+close = df["Close"].values.flatten().astype(np.float32)
+volume = df["Volume"].values.flatten().astype(np.float32)
+dates = df.index
+
+# 分割
+split_date = "2026-06-15"
+split_idx = None
+for i, d in enumerate(dates):
+    if str(d.date()) <= split_date:
+        split_idx = i
+
+train_close = close[: split_idx + 1]
+train_vol = volume[: split_idx + 1]
+train_dates = dates[: split_idx + 1]
+actual_close = close[split_idx + 1 :]
+actual_dates = dates[split_idx + 1 :]
+HORIZON = len(actual_close)
+
+print(f"训练数据: {len(train_close)} 天（截至 {train_dates[-1].date()}）", flush=True)
+print(f"实际数据: {HORIZON} 天（{actual_dates[0].date()} ~ {actual_dates[-1].date()}）", flush=True)
+
+# ---------- 2. 计算对数收益率 ----------
+# log returns = ln(P_t / P_{t-1})，更平稳
+train_logret = np.diff(np.log(train_close)).astype(np.float32)  # length = N-1
+actual_logret = np.diff(np.log(np.concatenate([train_close[-1:], actual_close]))).astype(np.float32)
+
+print(f"对数收益率: 均值={train_logret.mean():.6f}, 标准差={train_logret.std():.6f}", flush=True)
+
+# ---------- 3. 加载模型 ----------
+print()
+print("加载 TimesFM 模型...", flush=True)
+torch.set_float32_matmul_precision("high")
+model = timesfm.TimesFM_2p5_200M_torch.from_pretrained(
+    "google/timesfm-2.5-200m-pytorch", torch_compile=False
+)
+print("模型加载完成", flush=True)
+
+# ---------- 4. 多窗口集成预测（对数收益率） ----------
+CONTEXTS = [128, 256, 512]  # 不同 context 长度
+all_preds = []
+all_quantiles = []
+
+for ctx_len in CONTEXTS:
+    ctx_data = train_logret[-ctx_len:] if len(train_logret) >= ctx_len else train_logret
+    actual_ctx = min(ctx_len, len(ctx_data))
+    # 对收益率: infer_is_positive=False (可负), normalize_inputs=True
+    model.compile(
+        timesfm.ForecastConfig(
+            max_context=actual_ctx,
+            max_horizon=128,
+            normalize_inputs=True,
+            use_continuous_quantile_head=True,
+            force_flip_invariance=False,
+            infer_is_positive=False,  # 收益率可正可负
+            fix_quantile_crossing=True,
+        )
+    )
+    pf, qf = model.forecast(horizon=HORIZON, inputs=[ctx_data])
+    all_preds.append(pf[0][:HORIZON])
+    all_quantiles.append(qf[0][:HORIZON])
+    print(f"  context={actual_ctx} 预测完成", flush=True)
+
+# 集成: 取平均
+ensemble_pred_logret = np.mean(all_preds, axis=0)
+ensemble_q_logret = np.mean(all_quantiles, axis=0)
+
+# ---------- 5. 转换回价格 ----------
+# P_t = P_{t-1} * exp(r_t)
+last_price = train_close[-1]
+pred_prices = []
+for i in range(HORIZON):
+    last_price = last_price * np.exp(ensemble_pred_logret[i])
+    pred_prices.append(last_price)
+pred_prices = np.array(pred_prices)
+
+# 分位数价格
+q_prices = np.zeros((HORIZON, 10))
+for qi in range(10):
+    p = train_close[-1]
+    for i in range(HORIZON):
+        p = p * np.exp(ensemble_q_logret[i, qi])
+        q_prices[i, qi] = p
+
+# ---------- 6. 计算误差 ----------
+actual = actual_close
+pred = pred_prices
+mae = np.mean(np.abs(actual - pred))
+rmse = np.sqrt(np.mean((actual - pred) ** 2))
+mape = np.mean(np.abs((actual - pred) / actual)) * 100
+
+actual_dir = np.diff(actual)
+pred_dir = np.diff(pred)
+dir_acc = np.mean(actual_dir * pred_dir > 0) * 100
+
+in_80 = np.mean((actual >= q_prices[:, 1]) & (actual <= q_prices[:, 9])) * 100
+in_40 = np.mean((actual >= q_prices[:, 3]) & (actual <= q_prices[:, 7])) * 100
+
+print()
+print("=== 优化版回测结果：6/16 ~ 6/30 ===", flush=True)
+print()
+print("  日期          实际价    预测价    误差    误差%", flush=True)
+for i in range(HORIZON):
+    err = pred[i] - actual[i]
+    err_pct = err / actual[i] * 100
+    print(
+        f"  {actual_dates[i].date()}  {actual[i]:7.2f}  {pred[i]:7.2f}  {err:+7.2f}  {err_pct:+6.2f}%",
+        flush=True,
+    )
+
+print()
+print("=== 误差指标（优化版 vs 原始版）===", flush=True)
+print(f"  MAE:    ${mae:.2f}        (原始: $8.46)", flush=True)
+print(f"  RMSE:   ${rmse:.2f}       (原始: $11.29)", flush=True)
+print(f"  MAPE:   {mape:.2f}%       (原始: 2.98%)", flush=True)
+print(f"  方向准确率: {dir_acc:.1f}%       (原始: 55.6%)", flush=True)
+print(f"  80%CI覆盖率: {in_80:.1f}%        (原始: 90.0%)", flush=True)
+print(f"  40%CI覆盖率: {in_40:.1f}%        (原始: 60.0%)", flush=True)
+
+# ---------- 7. 画对比图 ----------
+fig, axes = plt.subplots(2, 1, figsize=(14, 10), sharex=False)
+
+# 上图: 价格对比
+ax = axes[0]
+show_n = min(60, len(train_close))
+ax.plot(range(show_n), train_close[-show_n:], label="Historical", color="steelblue", linewidth=1.5)
+x_actual = range(show_n, show_n + HORIZON)
+ax.plot(x_actual, actual, label="Actual", color="forestgreen", linewidth=2, marker="o", markersize=4)
+ax.plot(x_actual, pred, label="Optimized Forecast", color="tomato", linewidth=2, linestyle="--")
+ax.fill_between(x_actual, q_prices[:, 1], q_prices[:, 9], alpha=0.15, color="tomato", label="80% CI")
+ax.fill_between(x_actual, q_prices[:, 3], q_prices[:, 7], alpha=0.3, color="tomato", label="40% CI")
+ax.axvline(x=show_n - 1, color="gray", linestyle="--", alpha=0.5)
+ax.set_title(f"Optimized: Log-Return + Ensemble (MAPE={mape:.2f}%, Dir={dir_acc:.0f}%)", fontsize=13)
+ax.set_ylabel("Price (USD)")
+ax.legend(loc="upper left")
+
+# 下图: 预测误差对比
+ax2 = axes[1]
+errors = pred - actual
+ax2.bar(range(HORIZON), errors, color=["tomato" if e > 0 else "steelblue" for e in errors], alpha=0.7)
+ax2.axhline(y=0, color="black", linewidth=0.8)
+ax2.set_title("Forecast Error per Day (Optimized)", fontsize=13)
+ax2.set_xlabel("Trading Days After Split")
+ax2.set_ylabel("Error (USD)")
+ax2.set_xticks(range(HORIZON))
+ax2.set_xticklabels([str(d.date()) for d in actual_dates], rotation=45, fontsize=8)
+
+plt.tight_layout()
+plt.savefig("aapl_backtest_optimized.png", dpi=150)
+print()
+print("Chart saved: aapl_backtest_optimized.png", flush=True)
+print("Done!", flush=True)

stock_backtest_trick.py

@@ -0,0 +1,217 @@
+"""TimesFM 纯技巧优化：
+1. 对数收益率预测（更平稳）
+2. 多窗口集成（128/256/512）
+3. SPY 大盘走势引导：先预测 SPY，用 SPY 预测的趋势辅助判断 AAPL 方向
+4. 波动率调整：用近期波动率缩放置信区间
+5. infer_is_positive=False（收益率可负）
+"""
+import yfinance as yf
+import numpy as np
+import torch
+import timesfm
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from datetime import datetime, timedelta
+
+
+# 1. 获取数据
+print("下载 AAPL + SPY 数据...", flush=True)
+end = datetime.now()
+start = end - timedelta(days=365)
+df_aapl = yf.download("AAPL", start=start.strftime("%Y-%m-%d"), end=end.strftime("%Y-%m-%d"), progress=False)
+df_spy = yf.download("SPY", start=start.strftime("%Y-%m-%d"), end=end.strftime("%Y-%m-%d"), progress=False)
+
+close = df_aapl["Close"].values.flatten().astype(np.float32)
+spy_close = df_spy["Close"].values.flatten().astype(np.float32)
+dates = df_aapl.index
+
+min_len = min(len(close), len(spy_close))
+close = close[-min_len:]
+spy_close = spy_close[-min_len:]
+dates = dates[-min_len:]
+
+# 分割
+split_date = "2026-06-15"
+split_idx = None
+for i, d in enumerate(dates):
+    if str(d.date()) <= split_date:
+        split_idx = i
+
+train_close = close[: split_idx + 1]
+train_spy = spy_close[: split_idx + 1]
+train_dates = dates[: split_idx + 1]
+actual_close = close[split_idx + 1 :]
+actual_spy = spy_close[split_idx + 1 :]
+actual_dates = dates[split_idx + 1 :]
+HORIZON = len(actual_close)
+
+print(f"训练数据: {len(train_close)} 天 | 预测: {HORIZON} 天", flush=True)
+
+# 2. 对数收益率
+train_logret = np.diff(np.log(train_close)).astype(np.float32)
+train_spy_logret = np.diff(np.log(train_spy)).astype(np.float32)
+actual_logret = np.diff(np.log(np.concatenate([train_close[-1:], actual_close]))).astype(np.float32)
+
+# 3. 加载模型
+print()
+print("加载 TimesFM 模型...", flush=True)
+torch.set_float32_matmul_precision("high")
+model = timesfm.TimesFM_2p5_200M_torch.from_pretrained(
+    "google/timesfm-2.5-200m-pytorch", torch_compile=False
+)
+print("模型加载完成", flush=True)
+
+# 4. 多窗口集成预测 AAPL 收益率
+CONTEXTS = [128, 256, 512]
+all_preds = []
+all_quantiles = []
+
+for ctx_len in CONTEXTS:
+    ctx_data = train_logret[-ctx_len:] if len(train_logret) >= ctx_len else train_logret
+    actual_ctx = min(ctx_len, len(ctx_data))
+    # 向上取整到 32 的倍数
+    actual_ctx = ((actual_ctx + 31) // 32) * 32
+    ctx_data = train_logret[-actual_ctx:]
+
+    model.compile(
+        timesfm.ForecastConfig(
+            max_context=actual_ctx,
+            max_horizon=128,
+            normalize_inputs=True,
+            use_continuous_quantile_head=True,
+            force_flip_invariance=False,
+            infer_is_positive=False,
+            fix_quantile_crossing=True,
+        )
+    )
+    pf, qf = model.forecast(horizon=HORIZON, inputs=[ctx_data])
+    all_preds.append(pf[0][:HORIZON])
+    all_quantiles.append(qf[0][:HORIZON])
+    print(f"  AAPL context={actual_ctx} done", flush=True)
+
+# 5. 预测 SPY 收益率（大盘趋势引导）
+spy_ctx = train_spy_logret[-256:]
+spy_ctx_len = ((len(spy_ctx) + 31) // 32) * 32
+spy_ctx = train_spy_logret[-spy_ctx_len:]
+model.compile(
+    timesfm.ForecastConfig(
+        max_context=spy_ctx_len,
+        max_horizon=128,
+        normalize_inputs=True,
+        use_continuous_quantile_head=True,
+        force_flip_invariance=False,
+        infer_is_positive=False,
+        fix_quantile_crossing=True,
+    )
+)
+spy_pf, _ = model.forecast(horizon=HORIZON, inputs=[spy_ctx])
+print(f"  SPY context={spy_ctx_len} done", flush=True)
+
+# 6. 集成 + SPY 趋势调整
+ensemble_pred_logret = np.mean(all_preds, axis=0)
+ensemble_q_logret = np.mean(all_quantiles, axis=0)
+spy_pred_logret = spy_pf[0][:HORIZON]
+
+# SPY 趋势调整：如果 SPY 预测下跌，对 AAPL 预测施加向下的调整
+# 计算 AAPL 对 SPY 的 beta（敏感度）
+beta = np.corrcoef(train_logret[-60:], train_spy_logret[-60:])[0, 1]
+print(f"  AAPL-SPY 60日相关系数: {beta:.3f}", flush=True)
+
+# 调整：将 SPY 预测的偏离均值部分 * beta 加到 AAPL 预测上
+spy_mean = np.mean(train_spy_logret[-60:])
+spy_deviation = spy_pred_logret - spy_mean  # SPY 偏离其均值的部分
+adjustment = beta * spy_deviation * 0.3  # 0.3 是调整强度，避免过度修正
+adjusted_pred_logret = ensemble_pred_logret + adjustment
+
+# 7. 波动率调整置信区间
+recent_vol = np.std(train_logret[-20:])
+long_vol = np.std(train_logret[-60:])
+print(f"  近20日波动率: {recent_vol:.5f} | 近60日波动率: {long_vol:.5f}", flush=True)
+
+# 如果近期波动率高于长期，扩大置信区间
+vol_ratio = recent_vol / max(long_vol, 1e-8)
+vol_scale = max(vol_ratio, 1.0)  # 只扩大不缩小
+adjusted_q_logret = ensemble_q_logret.copy()
+median_idx = 5
+for qi in range(10):
+    if qi != median_idx:
+        adjusted_q_logret[:, qi] = ensemble_q_logret[:, median_idx] + (
+            ensemble_q_logret[:, qi] - ensemble_q_logret[:, median_idx]
+        ) * vol_scale
+
+# 8. 转换回价格
+last_price = train_close[-1]
+pred_prices = []
+for i in range(HORIZON):
+    last_price = last_price * np.exp(adjusted_pred_logret[i])
+    pred_prices.append(last_price)
+pred_prices = np.array(pred_prices)
+
+q_prices = np.zeros((HORIZON, 10))
+for qi in range(10):
+    p = train_close[-1]
+    for i in range(HORIZON):
+        p = p * np.exp(adjusted_q_logret[i, qi])
+        q_prices[i, qi] = p
+
+# 9. 误差指标
+actual = actual_close
+pred = pred_prices
+mae = np.mean(np.abs(actual - pred))
+rmse = np.sqrt(np.mean((actual - pred) ** 2))
+mape = np.mean(np.abs((actual - pred) / actual)) * 100
+actual_dir = np.diff(actual)
+pred_dir = np.diff(pred)
+dir_acc = np.mean(actual_dir * pred_dir > 0) * 100
+in_80 = np.mean((actual >= q_prices[:, 1]) & (actual <= q_prices[:, 9])) * 100
+in_40 = np.mean((actual >= q_prices[:, 3]) & (actual <= q_prices[:, 7])) * 100
+
+print()
+print("=== 纯技巧优化版回测结果 ===", flush=True)
+print()
+print("  日期          实际价    预测价    误差    误差%", flush=True)
+for i in range(HORIZON):
+    err = pred[i] - actual[i]
+    err_pct = err / actual[i] * 100
+    print(f"  {actual_dates[i].date()}  {actual[i]:7.2f}  {pred[i]:7.2f}  {err:+7.2f}  {err_pct:+6.2f}%", flush=True)
+
+print()
+print("=== 误差指标（四版对比）===", flush=True)
+print(f"  MAE:    ${mae:.2f}        (原始: $8.46, 对数收益率: $8.00, XReg: $9.21)", flush=True)
+print(f"  RMSE:   ${rmse:.2f}       (原始: $11.29, 对数收益率: $11.09, XReg: $12.47)", flush=True)
+print(f"  MAPE:   {mape:.2f}%       (原始: 2.98%, 对数收益率: 2.82%, XReg: 3.25%)", flush=True)
+print(f"  方向准确率: {dir_acc:.1f}%       (原始: 55.6%, 对数收益率: 44.4%, XReg: 44.4%)", flush=True)
+print(f"  80%CI覆盖率: {in_80:.1f}%        (原始: 90.0%, 对数收益率: 100.0%, XReg: 70.0%)", flush=True)
+print(f"  40%CI覆盖率: {in_40:.1f}%        (原始: 60.0%, 对数收益率: 80.0%, XReg: 40.0%)", flush=True)
+
+# 10. 画图
+fig, axes = plt.subplots(2, 1, figsize=(14, 10))
+show_n = min(60, len(train_close))
+ax = axes[0]
+ax.plot(range(show_n), train_close[-show_n:], label="Historical", color="steelblue", linewidth=1.5)
+x_actual = range(show_n, show_n + HORIZON)
+ax.plot(x_actual, actual, label="Actual", color="forestgreen", linewidth=2, marker="o", markersize=4)
+ax.plot(x_actual, pred, label="Trick Forecast", color="purple", linewidth=2, linestyle="--")
+ax.fill_between(x_actual, q_prices[:, 1], q_prices[:, 9], alpha=0.15, color="purple", label="80% CI")
+ax.fill_between(x_actual, q_prices[:, 3], q_prices[:, 7], alpha=0.3, color="purple", label="40% CI")
+ax.axvline(x=show_n - 1, color="gray", linestyle="--", alpha=0.5)
+ax.set_title(f"Trick: LogRet + Ensemble + SPY + VolAdj (MAPE={mape:.2f}%, Dir={dir_acc:.0f}%)", fontsize=13)
+ax.set_ylabel("Price (USD)")
+ax.legend(loc="upper left")
+
+ax2 = axes[1]
+errors = pred - actual
+ax2.bar(range(HORIZON), errors, color=["purple" if e > 0 else "steelblue" for e in errors], alpha=0.7)
+ax2.axhline(y=0, color="black", linewidth=0.8)
+ax2.set_title("Trick Forecast Error per Day", fontsize=13)
+ax2.set_xlabel("Trading Days After Split")
+ax2.set_ylabel("Error (USD)")
+ax2.set_xticks(range(HORIZON))
+ax2.set_xticklabels([str(d.date()) for d in actual_dates], rotation=45, fontsize=8)
+
+plt.tight_layout()
+plt.savefig("aapl_backtest_trick.png", dpi=150)
+print()
+print("Chart saved: aapl_backtest_trick.png", flush=True)
+print("Done!", flush=True)

stock_backtest_xreg.py

@@ -0,0 +1,203 @@
+"""TimesFM + XReg 协变量预测：加入成交量、RSI、SPY 大盘指数"""
+import yfinance as yf
+import numpy as np
+import torch
+import timesfm
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from datetime import datetime, timedelta
+
+
+def calc_rsi(prices, period=14):
+    deltas = np.diff(prices)
+    gains = np.where(deltas > 0, deltas, 0.0)
+    losses = np.where(deltas < 0, -deltas, 0.0)
+    avg_gain = np.convolve(gains, np.ones(period) / period, mode="valid")
+    avg_loss = np.convolve(losses, np.ones(period) / period, mode="valid")
+    avg_loss = np.where(avg_loss == 0, 1e-10, avg_loss)
+    rs = avg_gain / avg_loss
+    rsi = 100.0 - (100.0 / (1.0 + rs))
+    return np.concatenate([np.full(period, 50.0), rsi])
+
+
+# 1. 获取数据
+print("下载 AAPL + SPY 数据...", flush=True)
+end = datetime.now()
+start = end - timedelta(days=365)
+df_aapl = yf.download("AAPL", start=start.strftime("%Y-%m-%d"), end=end.strftime("%Y-%m-%d"), progress=False)
+df_spy = yf.download("SPY", start=start.strftime("%Y-%m-%d"), end=end.strftime("%Y-%m-%d"), progress=False)
+
+close = df_aapl["Close"].values.flatten().astype(np.float64)
+volume = df_aapl["Volume"].values.flatten().astype(np.float64)
+spy_close = df_spy["Close"].values.flatten().astype(np.float64)
+dates = df_aapl.index
+
+# 对齐长度
+min_len = min(len(close), len(spy_close))
+close = close[-min_len:]
+volume = volume[-min_len:]
+spy_close = spy_close[-min_len:]
+dates = dates[-min_len:]
+
+# 计算 RSI
+rsi = calc_rsi(close)
+
+# 分割
+split_date = "2026-06-15"
+split_idx = None
+for i, d in enumerate(dates):
+    if str(d.date()) <= split_date:
+        split_idx = i
+
+train_close = close[: split_idx + 1]
+train_vol = volume[: split_idx + 1]
+train_spy = spy_close[: split_idx + 1]
+train_rsi = rsi[: split_idx + 1]
+train_dates = dates[: split_idx + 1]
+
+actual_close = close[split_idx + 1 :]
+actual_dates = dates[split_idx + 1 :]
+HORIZON = len(actual_close)
+
+print(f"训练数据: {len(train_close)} 天", flush=True)
+print(f"预测目标: {HORIZON} 天", flush=True)
+
+# 2. 准备协变量
+# 动态数值协变量需要覆盖 context + horizon 的完整长度
+# 对每个序列: train 部分用实际值，test 部分需要"未来值"
+# 对于 RSI 和 Volume，我们没有未来值，用最后一个值填充
+# 对于 SPY，用训练集最后一个值填充（因为我们无法预知未来 SPY）
+
+# 协变量需要: 每个协变量是一个 list，每个元素对应一个输入序列的完整长度(context+horizon)
+full_len = len(train_close) + HORIZON
+
+# 成交量: train 用实际值，future 用最近 5 日均值
+vol_future = np.mean(train_vol[-5:])
+vol_full = np.concatenate([train_vol, np.full(HORIZON, vol_future)])
+
+# RSI: train 用实际值，future 用 50（中性）
+rsi_full = np.concatenate([train_rsi, np.full(HORIZON, 50.0)])
+
+# SPY: train 用实际值，future 用最后一个值
+spy_future = train_spy[-1]
+spy_full = np.concatenate([train_spy, np.full(HORIZON, spy_future)])
+
+# 3. 加载模型
+print()
+print("加载 TimesFM 模型...", flush=True)
+torch.set_float32_matmul_precision("high")
+model = timesfm.TimesFM_2p5_200M_torch.from_pretrained(
+    "google/timesfm-2.5-200m-pytorch", torch_compile=False
+)
+print("模型加载完成", flush=True)
+
+# 4. 编译（return_backcast=True 是 XReg 必需的）
+model.compile(
+    timesfm.ForecastConfig(
+        max_context=512,
+        max_horizon=128,
+        normalize_inputs=True,
+        use_continuous_quantile_head=True,
+        force_flip_invariance=False,
+        infer_is_positive=True,
+        fix_quantile_crossing=True,
+        return_backcast=True,  # XReg 需要
+    )
+)
+print("编译完成", flush=True)
+
+# 5. 用 XReg 预测
+print("运行 XReg 协变量预测...", flush=True)
+
+# 动态数值协变量: dict[str, list[list[float]]]
+# 每个协变量是一个 list，其中每个元素是一个序列（对应一个输入）
+# 这里只有一个输入序列
+dynamic_num_covs = {
+    "volume": [vol_full],
+    "rsi": [rsi_full],
+    "spy_close": [spy_full],
+}
+
+# 静态数值协变量
+static_num_covs = {
+    "avg_volume": [np.mean(train_vol)],
+}
+
+point_outputs, quantile_outputs = model.forecast_with_covariates(
+    inputs=[train_close],
+    dynamic_numerical_covariates=dynamic_num_covs,
+    static_numerical_covariates=static_num_covs,
+    xreg_mode="xreg + timesfm",  # 先回归再预测残差
+    normalize_xreg_target_per_input=True,
+    ridge=1.0,
+)
+print("XReg 预测完成！", flush=True)
+
+# 6. 提取结果
+pred = np.array(point_outputs[0][:HORIZON])
+q = np.array(quantile_outputs[0])  # (horizon, 10) or (full, 10)
+# quantile_outputs 可能包含 backcast，取最后 HORIZON 个
+if q.shape[0] > HORIZON:
+    q = q[-HORIZON:]
+
+# 7. 计算误差
+actual = actual_close
+mae = np.mean(np.abs(actual - pred))
+rmse = np.sqrt(np.mean((actual - pred) ** 2))
+mape = np.mean(np.abs((actual - pred) / actual)) * 100
+actual_dir = np.diff(actual)
+pred_dir = np.diff(pred)
+dir_acc = np.mean(actual_dir * pred_dir > 0) * 100
+in_80 = np.mean((actual >= q[:, 1]) & (actual <= q[:, 9])) * 100
+in_40 = np.mean((actual >= q[:, 3]) & (actual <= q[:, 7])) * 100
+
+print()
+print("=== XReg 协变量版回测结果 ===", flush=True)
+print()
+print("  日期          实际价    预测价    误差    误差%", flush=True)
+for i in range(HORIZON):
+    err = pred[i] - actual[i]
+    err_pct = err / actual[i] * 100
+    print(f"  {actual_dates[i].date()}  {actual[i]:7.2f}  {pred[i]:7.2f}  {err:+7.2f}  {err_pct:+6.2f}%", flush=True)
+
+print()
+print("=== 误差指标（三版对比）===", flush=True)
+print(f"  MAE:    ${mae:.2f}        (原始: $8.46, 优化: $8.00)", flush=True)
+print(f"  RMSE:   ${rmse:.2f}       (原始: $11.29, 优化: $11.09)", flush=True)
+print(f"  MAPE:   {mape:.2f}%       (原始: 2.98%, 优化: 2.82%)", flush=True)
+print(f"  方向准确率: {dir_acc:.1f}%       (原始: 55.6%, 优化: 44.4%)", flush=True)
+print(f"  80%CI覆盖率: {in_80:.1f}%        (原始: 90.0%, 优化: 100.0%)", flush=True)
+print(f"  40%CI覆盖率: {in_40:.1f}%        (原始: 60.0%, 优化: 80.0%)", flush=True)
+
+# 8. 画图
+fig, axes = plt.subplots(2, 1, figsize=(14, 10))
+show_n = min(60, len(train_close))
+ax = axes[0]
+ax.plot(range(show_n), train_close[-show_n:], label="Historical", color="steelblue", linewidth=1.5)
+x_actual = range(show_n, show_n + HORIZON)
+ax.plot(x_actual, actual, label="Actual", color="forestgreen", linewidth=2, marker="o", markersize=4)
+ax.plot(x_actual, pred, label="XReg Forecast", color="darkorange", linewidth=2, linestyle="--")
+ax.fill_between(x_actual, q[:, 1], q[:, 9], alpha=0.15, color="darkorange", label="80% CI")
+ax.fill_between(x_actual, q[:, 3], q[:, 7], alpha=0.3, color="darkorange", label="40% CI")
+ax.axvline(x=show_n - 1, color="gray", linestyle="--", alpha=0.5)
+ax.set_title(f"XReg: AAPL + Volume/RSI/SPY (MAPE={mape:.2f}%, Dir={dir_acc:.0f}%)", fontsize=13)
+ax.set_ylabel("Price (USD)")
+ax.legend(loc="upper left")
+
+# 误差柱状图
+ax2 = axes[1]
+errors = pred - actual
+ax2.bar(range(HORIZON), errors, color=["darkorange" if e > 0 else "steelblue" for e in errors], alpha=0.7)
+ax2.axhline(y=0, color="black", linewidth=0.8)
+ax2.set_title("XReg Forecast Error per Day", fontsize=13)
+ax2.set_xlabel("Trading Days After Split")
+ax2.set_ylabel("Error (USD)")
+ax2.set_xticks(range(HORIZON))
+ax2.set_xticklabels([str(d.date()) for d in actual_dates], rotation=45, fontsize=8)
+
+plt.tight_layout()
+plt.savefig("aapl_backtest_xreg.png", dpi=150)
+print()
+print("Chart saved: aapl_backtest_xreg.png", flush=True)
+print("Done!", flush=True)

stock_forecast.py

@@ -0,0 +1,84 @@
+import yfinance as yf
+import numpy as np
+import torch
+import timesfm
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from datetime import datetime, timedelta
+
+# 1. 获取 AAPL 近 1 年收盘价
+print("下载 AAPL 股票数据...", flush=True)
+end = datetime.now()
+start = end - timedelta(days=365)
+df = yf.download("AAPL", start=start.strftime("%Y-%m-%d"), end=end.strftime("%Y-%m-%d"), progress=False)
+close = df["Close"].values.flatten().astype(np.float32)
+dates = df.index
+print(f"获取到 {len(close)} 个交易日", flush=True)
+print(f"价格范围: {close.min():.2f} ~ {close.max():.2f}", flush=True)
+print(f"最近 5 日收盘价: {close[-5:]}", flush=True)
+
+# 2. 加载 TimesFM 模型
+print()
+print("加载 TimesFM 模型...", flush=True)
+torch.set_float32_matmul_precision("high")
+model = timesfm.TimesFM_2p5_200M_torch.from_pretrained(
+    "google/timesfm-2.5-200m-pytorch", torch_compile=False
+)
+print("模型加载完成", flush=True)
+
+# 3. 编译并预测未来 20 个交易日
+HORIZON = 20
+model.compile(
+    timesfm.ForecastConfig(
+        max_context=512,
+        max_horizon=128,
+        normalize_inputs=True,
+        use_continuous_quantile_head=True,
+        force_flip_invariance=False,
+        infer_is_positive=True,
+        fix_quantile_crossing=True,
+    )
+)
+print("编译完成，开始预测...", flush=True)
+
+pf, qf = model.forecast(horizon=HORIZON, inputs=[close])
+print("预测完成！", flush=True)
+
+# 4. 输出结果
+print()
+print("=== AAPL 未来 20 个交易日预测 ===", flush=True)
+print(f"当前价格: {close[-1]:.2f}", flush=True)
+print()
+print("  日期(估计)    点预测    q10(低)   q90(高)", flush=True)
+last_date = dates[-1]
+for i in range(HORIZON):
+    est_date = last_date + timedelta(days=i + 1)
+    print(
+        f"  {est_date.strftime('%Y-%m-%d')}  {pf[0][i]:7.2f}  {qf[0][i,1]:7.2f}  {qf[0][i,9]:7.2f}",
+        flush=True,
+    )
+
+print()
+print(f"预测均价: {pf[0].mean():.2f}", flush=True)
+print(f"预测涨跌: {(pf[0][-1] - close[-1]) / close[-1] * 100:+.2f}%", flush=True)
+print(f"80%置信区间: {qf[0,-1,1]:.2f} ~ {qf[0,-1,9]:.2f}", flush=True)
+
+# 5. 画图
+fig, ax = plt.subplots(figsize=(14, 6))
+show_n = min(60, len(close))
+ax.plot(range(show_n), close[-show_n:], label="历史收盘价", color="steelblue", linewidth=1.5)
+x_fc = range(show_n, show_n + HORIZON)
+ax.plot(x_fc, pf[0], label="点预测(中位数)", color="tomato", linewidth=2)
+ax.fill_between(x_fc, qf[0, :, 1], qf[0, :, 9], alpha=0.2, color="tomato", label="80% 置信区间")
+ax.fill_between(x_fc, qf[0, :, 3], qf[0, :, 7], alpha=0.35, color="tomato", label="40% 置信区间")
+ax.axvline(x=show_n - 1, color="gray", linestyle="--", alpha=0.5, label="预测起点")
+ax.set_title("AAPL 收盘价预测 (TimesFM 2.5)", fontsize=14)
+ax.set_xlabel("交易日")
+ax.set_ylabel("价格 (USD)")
+ax.legend(loc="upper left")
+plt.tight_layout()
+plt.savefig("aapl_forecast.png", dpi=150)
+print()
+print("图表已保存: aapl_forecast.png", flush=True)
+print("✅ 完成！", flush=True)

tests/test_model_loading.py

@@ -0,0 +1,68 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for loading TimesFM 2.5 models."""
+
+import os
+import tempfile
+
+from timesfm.timesfm_2p5.timesfm_2p5_torch import TimesFM_2p5_200M_torch
+from timesfm.timesfm_2p5.timesfm_2p5_flax import TimesFM_2p5_200M_flax
+
+
+class TestModelLoading:
+  """Tests to verify model instantiation, loading, and compatibility."""
+
+  def test_torch_load_checkpoint_and_from_pretrained_local(self):
+    """Verifies that PyTorch load_checkpoint and from_pretrained work locally."""
+    # 1. Instantiate the model wrapper with compilation disabled
+    tfm = TimesFM_2p5_200M_torch(torch_compile=False)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+      # 2. Save the model's randomly-initialized weights
+      tfm._save_pretrained(tmpdir)
+
+      # Verify weights file is written
+      weights_path = os.path.join(tmpdir, "model.safetensors")
+      assert os.path.exists(weights_path)
+
+      # 3. Verify that load_checkpoint works from the temp directory path
+      tfm2 = TimesFM_2p5_200M_torch(torch_compile=False)
+      tfm2.load_checkpoint(tmpdir, torch_compile=False)
+
+      # 4. Verify that from_pretrained works with a local directory path
+      # and accepts/ignores extra kwargs (like proxies) without raising TypeError
+      tfm3 = TimesFM_2p5_200M_torch.from_pretrained(
+          tmpdir,
+          torch_compile=False,
+          proxies={"http": "http://dummy.proxy"},
+          custom_kwarg="dummy_value",
+      )
+      assert tfm3 is not None
+      assert not tfm3.torch_compile
+
+      # 5. Run a simple prediction step to verify the loaded model performs forward pass
+      import numpy as np
+      inputs = [np.random.randn(32)]
+      forecasts = tfm3.model.forecast_naive(horizon=10, inputs=inputs)
+      assert len(forecasts) == 1
+      assert forecasts[0].shape == (10, 10)
+
+  def test_flax_model_init_kwargs(self):
+    """Verifies that Flax model wrapper constructor accepts arbitrary kwargs."""
+    tfm = TimesFM_2p5_200M_flax(
+        proxies={"http": "http://dummy.proxy"},
+        custom_kwarg="dummy_value",
+    )
+    assert tfm is not None

timesfm-forecasting/examples/anomaly-detection/output/anomaly_detection.json

@@ -1,448 +0,0 @@
-{
-  "method": "two_phase",
-  "context_method": "linear_detrend_zscore",
-  "forecast_method": "quantile_prediction_intervals",
-  "thresholds": {
-    "critical_z": 3.0,
-    "warning_z": 2.0,
-    "pi_critical_pct": 80,
-    "pi_warning_pct": 60
-  },
-  "context_summary": {
-    "total": 36,
-    "critical": 1,
-    "warning": 0,
-    "normal": 35,
-    "res_std": 0.11362
-  },
-  "forecast_summary": {
-    "total": 12,
-    "critical": 4,
-    "warning": 1,
-    "normal": 7
-  },
-  "context_detections": [
-    {
-      "date": "2022-01",
-      "value": 0.89,
-      "trend": 0.837,
-      "residual": 0.053,
-      "z_score": 0.467,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2022-02",
-      "value": 0.89,
-      "trend": 0.8514,
-      "residual": 0.0386,
-      "z_score": 0.34,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2022-03",
-      "value": 1.02,
-      "trend": 0.8658,
-      "residual": 0.1542,
-      "z_score": 1.357,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2022-04",
-      "value": 0.88,
-      "trend": 0.8803,
-      "residual": -0.0003,
-      "z_score": -0.002,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2022-05",
-      "value": 0.85,
-      "trend": 0.8947,
-      "residual": -0.0447,
-      "z_score": -0.394,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2022-06",
-      "value": 0.88,
-      "trend": 0.9092,
-      "residual": -0.0292,
-      "z_score": -0.257,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2022-07",
-      "value": 0.88,
-      "trend": 0.9236,
-      "residual": -0.0436,
-      "z_score": -0.384,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2022-08",
-      "value": 0.9,
-      "trend": 0.9381,
-      "residual": -0.0381,
-      "z_score": -0.335,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2022-09",
-      "value": 0.88,
-      "trend": 0.9525,
-      "residual": -0.0725,
-      "z_score": -0.638,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2022-10",
-      "value": 0.95,
-      "trend": 0.9669,
-      "residual": -0.0169,
-      "z_score": -0.149,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2022-11",
-      "value": 0.77,
-      "trend": 0.9814,
-      "residual": -0.2114,
-      "z_score": -1.86,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2022-12",
-      "value": 0.78,
-      "trend": 0.9958,
-      "residual": -0.2158,
-      "z_score": -1.9,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2023-01",
-      "value": 0.87,
-      "trend": 1.0103,
-      "residual": -0.1403,
-      "z_score": -1.235,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2023-02",
-      "value": 0.98,
-      "trend": 1.0247,
-      "residual": -0.0447,
-      "z_score": -0.394,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2023-03",
-      "value": 1.21,
-      "trend": 1.0392,
-      "residual": 0.1708,
-      "z_score": 1.503,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2023-04",
-      "value": 1.0,
-      "trend": 1.0536,
-      "residual": -0.0536,
-      "z_score": -0.472,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2023-05",
-      "value": 0.94,
-      "trend": 1.0681,
-      "residual": -0.1281,
-      "z_score": -1.127,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2023-06",
-      "value": 1.08,
-      "trend": 1.0825,
-      "residual": -0.0025,
-      "z_score": -0.022,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2023-07",
-      "value": 1.18,
-      "trend": 1.0969,
-      "residual": 0.0831,
-      "z_score": 0.731,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2023-08",
-      "value": 1.24,
-      "trend": 1.1114,
-      "residual": 0.1286,
-      "z_score": 1.132,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2023-09",
-      "value": 1.47,
-      "trend": 1.1258,
-      "residual": 0.3442,
-      "z_score": 3.029,
-      "severity": "CRITICAL"
-    },
-    {
-      "date": "2023-10",
-      "value": 1.32,
-      "trend": 1.1403,
-      "residual": 0.1797,
-      "z_score": 1.582,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2023-11",
-      "value": 1.18,
-      "trend": 1.1547,
-      "residual": 0.0253,
-      "z_score": 0.222,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2023-12",
-      "value": 1.16,
-      "trend": 1.1692,
-      "residual": -0.0092,
-      "z_score": -0.081,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2024-01",
-      "value": 1.22,
-      "trend": 1.1836,
-      "residual": 0.0364,
-      "z_score": 0.32,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2024-02",
-      "value": 1.35,
-      "trend": 1.1981,
-      "residual": 0.1519,
-      "z_score": 1.337,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2024-03",
-      "value": 1.34,
-      "trend": 1.2125,
-      "residual": 0.1275,
-      "z_score": 1.122,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2024-04",
-      "value": 1.26,
-      "trend": 1.2269,
-      "residual": 0.0331,
-      "z_score": 0.291,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2024-05",
-      "value": 1.15,
-      "trend": 1.2414,
-      "residual": -0.0914,
-      "z_score": -0.804,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2024-06",
-      "value": 1.2,
-      "trend": 1.2558,
-      "residual": -0.0558,
-      "z_score": -0.491,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2024-07",
-      "value": 1.24,
-      "trend": 1.2703,
-      "residual": -0.0303,
-      "z_score": -0.266,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2024-08",
-      "value": 1.3,
-      "trend": 1.2847,
-      "residual": 0.0153,
-      "z_score": 0.135,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2024-09",
-      "value": 1.28,
-      "trend": 1.2992,
-      "residual": -0.0192,
-      "z_score": -0.169,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2024-10",
-      "value": 1.27,
-      "trend": 1.3136,
-      "residual": -0.0436,
-      "z_score": -0.384,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2024-11",
-      "value": 1.22,
-      "trend": 1.328,
-      "residual": -0.108,
-      "z_score": -0.951,
-      "severity": "NORMAL"
-    },
-    {
-      "date": "2024-12",
-      "value": 1.2,
-      "trend": 1.3425,
-      "residual": -0.1425,
-      "z_score": -1.254,
-      "severity": "NORMAL"
-    }
-  ],
-  "forecast_detections": [
-    {
-      "date": "2025-01",
-      "actual": 1.2821,
-      "forecast": 1.2593,
-      "q10": 1.1407,
-      "q20": 1.1881,
-      "q80": 1.324,
-      "q90": 1.3679,
-      "severity": "NORMAL",
-      "was_injected": false
-    },
-    {
-      "date": "2025-02",
-      "actual": 1.1522,
-      "forecast": 1.2857,
-      "q10": 1.1406,
-      "q20": 1.1961,
-      "q80": 1.3751,
-      "q90": 1.4254,
-      "severity": "WARNING",
-      "was_injected": false
-    },
-    {
-      "date": "2025-03",
-      "actual": 1.3358,
-      "forecast": 1.295,
-      "q10": 1.1269,
-      "q20": 1.1876,
-      "q80": 1.4035,
-      "q90": 1.4643,
-      "severity": "NORMAL",
-      "was_injected": false
-    },
-    {
-      "date": "2025-04",
-      "actual": 2.0594,
-      "forecast": 1.2208,
-      "q10": 1.0353,
-      "q20": 1.1042,
-      "q80": 1.331,
-      "q90": 1.4017,
-      "severity": "CRITICAL",
-      "was_injected": true
-    },
-    {
-      "date": "2025-05",
-      "actual": 1.0747,
-      "forecast": 1.1703,
-      "q10": 0.9691,
-      "q20": 1.0431,
-      "q80": 1.2892,
-      "q90": 1.3632,
-      "severity": "NORMAL",
-      "was_injected": false
-    },
-    {
-      "date": "2025-06",
-      "actual": 1.1442,
-      "forecast": 1.1456,
-      "q10": 0.942,
-      "q20": 1.0111,
-      "q80": 1.2703,
-      "q90": 1.3454,
-      "severity": "NORMAL",
-      "was_injected": false
-    },
-    {
-      "date": "2025-07",
-      "actual": 1.2917,
-      "forecast": 1.1702,
-      "q10": 0.9504,
-      "q20": 1.0348,
-      "q80": 1.2998,
-      "q90": 1.3807,
-      "severity": "NORMAL",
-      "was_injected": false
-    },
-    {
-      "date": "2025-08",
-      "actual": 1.2519,
-      "forecast": 1.2027,
-      "q10": 0.9709,
-      "q20": 1.0594,
-      "q80": 1.3408,
-      "q90": 1.4195,
-      "severity": "NORMAL",
-      "was_injected": false
-    },
-    {
-      "date": "2025-09",
-      "actual": 0.6364,
-      "forecast": 1.191,
-      "q10": 0.9594,
-      "q20": 1.0404,
-      "q80": 1.3355,
-      "q90": 1.417,
-      "severity": "CRITICAL",
-      "was_injected": true
-    },
-    {
-      "date": "2025-10",
-      "actual": 1.2073,
-      "forecast": 1.1491,
-      "q10": 0.9079,
-      "q20": 0.9953,
-      "q80": 1.2869,
-      "q90": 1.3775,
-      "severity": "NORMAL",
-      "was_injected": false
-    },
-    {
-      "date": "2025-11",
-      "actual": 1.3851,
-      "forecast": 1.0805,
-      "q10": 0.8361,
-      "q20": 0.926,
-      "q80": 1.2284,
-      "q90": 1.3122,
-      "severity": "CRITICAL",
-      "was_injected": false
-    },
-    {
-      "date": "2025-12",
-      "actual": 1.8294,
-      "forecast": 1.0613,
-      "q10": 0.8022,
-      "q20": 0.8952,
-      "q80": 1.2169,
-      "q90": 1.296,
-      "severity": "CRITICAL",
-      "was_injected": true
-    }
-  ]
-}

timesfm-forecasting/examples/covariates-forecasting/output/covariates_metadata.json

@@ -1,59 +0,0 @@
-{
-  "description": "Synthetic retail sales data with covariates for TimesFM XReg demo",
-  "note_on_real_data": "For real datasets (e.g., Kaggle Rossmann Store Sales), download to tempfile.mkdtemp() -- do NOT commit to this repo.",
-  "stores": {
-    "store_A": {
-      "type": "premium",
-      "region": "urban",
-      "base_sales": 1000,
-      "mean_sales_context": 1148.7
-    },
-    "store_B": {
-      "type": "standard",
-      "region": "suburban",
-      "base_sales": 750,
-      "mean_sales_context": 907.0
-    },
-    "store_C": {
-      "type": "discount",
-      "region": "rural",
-      "base_sales": 500,
-      "mean_sales_context": 645.3
-    }
-  },
-  "dimensions": {
-    "context_length": 24,
-    "horizon_length": 12,
-    "total_length": 36,
-    "num_stores": 3,
-    "csv_rows": 108
-  },
-  "covariates": {
-    "dynamic_numerical": [
-      "price"
-    ],
-    "dynamic_categorical": [
-      "promotion",
-      "holiday",
-      "day_of_week"
-    ],
-    "static_categorical": [
-      "store_type",
-      "region"
-    ]
-  },
-  "effect_magnitudes": {
-    "holiday": "+200 units per holiday week",
-    "promotion": "+150 units per promotion week",
-    "price": "-20 units per $1 above base price"
-  },
-  "xreg_modes": {
-    "xreg + timesfm": "Regression on TimesFM residuals (default)",
-    "timesfm + xreg": "TimesFM on regression residuals"
-  },
-  "bug_fixes_history": [
-    "v1: Variable-shadowing -- all stores had identical covariates",
-    "v2: Fixed shadowing; CONTEXT_LEN 48->24",
-    "v3: Added component decomposition (base, price/promo/holiday effects); 2x2 sharex viz"
-  ]
-}

timesfm-forecasting/examples/covariates-forecasting/output/sales_with_covariates.csv

@@ -1,109 +0,0 @@
-store_id,week,split,sales,base_sales,price,price_effect,promotion,holiday,day_of_week,store_type,region
-store_A,0,context,1369.59,1012.19,11.6299,7.4,1,1,0,premium,urban
-store_A,1,context,973.53,973.04,11.9757,0.49,0,0,1,premium,urban
-store_A,2,context,1064.63,1059.16,11.7269,5.46,0,0,2,premium,urban
-store_A,3,context,1077.59,1080.99,12.1698,-3.4,0,0,3,premium,urban
-store_A,4,context,980.39,979.14,11.9372,1.26,0,0,4,premium,urban
-store_A,5,context,1011.7,1018.36,12.3327,-6.65,0,0,5,premium,urban
-store_A,6,context,1084.16,1088.16,12.2003,-4.01,0,0,6,premium,urban
-store_A,7,context,1085.98,1082.23,11.8124,3.75,0,0,0,premium,urban
-store_A,8,context,1098.52,1105.17,12.3323,-6.65,0,0,1,premium,urban
-store_A,9,context,1075.62,1081.71,12.3048,-6.1,0,0,2,premium,urban
-store_A,10,context,1312.23,1159.98,11.8875,2.25,1,0,3,premium,urban
-store_A,11,context,1368.02,1163.79,11.7883,4.23,0,1,4,premium,urban
-store_A,12,context,1138.41,1142.06,12.1825,-3.65,0,0,5,premium,urban
-store_A,13,context,1197.29,1190.09,11.6398,7.2,0,0,6,premium,urban
-store_A,14,context,1174.12,1168.12,11.6999,6.0,0,0,0,premium,urban
-store_A,15,context,1128.16,1118.3,11.5074,9.85,0,0,1,premium,urban
-store_A,16,context,1163.81,1169.55,12.2869,-5.74,0,0,2,premium,urban
-store_A,17,context,1114.18,1117.48,12.1649,-3.3,0,0,3,premium,urban
-store_A,18,context,1186.87,1190.98,12.2052,-4.1,0,0,4,premium,urban
-store_A,19,context,1147.27,1152.88,12.2807,-5.61,0,0,5,premium,urban
-store_A,20,context,1146.48,1145.66,11.9589,0.82,0,0,6,premium,urban
-store_A,21,context,1121.83,1123.21,12.0687,-1.37,0,0,0,premium,urban
-store_A,22,context,1203.28,1196.08,11.6398,7.2,0,0,1,premium,urban
-store_A,23,context,1344.9,1137.19,11.6145,7.71,0,1,2,premium,urban
-store_A,24,horizon,1118.64,1122.01,12.1684,-3.37,0,0,3,premium,urban
-store_A,25,horizon,1121.14,1120.56,11.9711,0.58,0,0,4,premium,urban
-store_A,26,horizon,1149.99,1151.29,12.0652,-1.3,0,0,5,premium,urban
-store_A,27,horizon,1284.67,1139.97,12.265,-5.3,1,0,6,premium,urban
-store_A,28,horizon,1284.67,1137.36,12.1347,-2.69,1,0,0,premium,urban
-store_A,29,horizon,1132.79,1133.86,12.0536,-1.07,0,0,1,premium,urban
-store_A,30,horizon,1197.3,1198.49,12.0592,-1.18,0,0,2,premium,urban
-store_A,31,horizon,1247.22,1093.3,11.804,3.92,1,0,3,premium,urban
-store_A,32,horizon,1095.84,1086.46,11.5308,9.38,0,0,4,premium,urban
-store_A,33,horizon,1073.83,1072.57,11.9367,1.27,0,0,5,premium,urban
-store_A,34,horizon,1134.51,1128.8,11.7146,5.71,0,0,6,premium,urban
-store_A,35,horizon,1351.15,1149.32,11.9085,1.83,0,1,0,premium,urban
-store_B,0,context,1062.53,712.0,9.9735,0.53,1,1,0,standard,suburban
-store_B,1,context,904.49,749.83,9.767,4.66,1,0,1,standard,suburban
-store_B,2,context,813.63,810.26,9.8316,3.37,0,0,2,standard,suburban
-store_B,3,context,720.11,720.53,10.0207,-0.41,0,0,3,standard,suburban
-store_B,4,context,820.78,819.55,9.9389,1.22,0,0,4,standard,suburban
-store_B,5,context,833.27,823.7,9.5216,9.57,0,0,5,standard,suburban
-store_B,6,context,795.26,801.78,10.3263,-6.53,0,0,6,standard,suburban
-store_B,7,context,770.37,778.29,10.3962,-7.92,0,0,0,standard,suburban
-store_B,8,context,855.92,848.72,9.6402,7.2,0,0,1,standard,suburban
-store_B,9,context,832.33,833.41,10.054,-1.08,0,0,2,standard,suburban
-store_B,10,context,1029.44,871.61,9.6086,7.83,1,0,3,standard,suburban
-store_B,11,context,1066.35,869.8,10.1722,-3.44,0,1,4,standard,suburban
-store_B,12,context,942.86,938.49,9.7812,4.38,0,0,5,standard,suburban
-store_B,13,context,1015.99,869.18,10.1594,-3.19,1,0,6,standard,suburban
-store_B,14,context,836.44,840.98,10.227,-4.54,0,0,0,standard,suburban
-store_B,15,context,885.72,891.1,10.2686,-5.37,0,0,1,standard,suburban
-store_B,16,context,901.45,893.6,9.6077,7.85,0,0,2,standard,suburban
-store_B,17,context,1080.63,938.95,10.416,-8.32,1,0,3,standard,suburban
-store_B,18,context,922.14,916.74,9.7302,5.4,0,0,4,standard,suburban
-store_B,19,context,904.66,895.41,9.5374,9.25,0,0,5,standard,suburban
-store_B,20,context,935.48,936.58,10.0549,-1.1,0,0,6,standard,suburban
-store_B,21,context,979.23,826.64,9.8709,2.58,1,0,0,standard,suburban
-store_B,22,context,837.49,844.09,10.3298,-6.6,0,0,1,standard,suburban
-store_B,23,context,1021.39,827.56,10.3083,-6.17,0,1,2,standard,suburban
-store_B,24,horizon,847.21,843.55,9.8171,3.66,0,0,3,standard,suburban
-store_B,25,horizon,789.27,798.33,10.4529,-9.06,0,0,4,standard,suburban
-store_B,26,horizon,877.09,872.91,9.7909,4.18,0,0,5,standard,suburban
-store_B,27,horizon,832.42,832.72,10.0151,-0.3,0,0,6,standard,suburban
-store_B,28,horizon,781.9,777.02,9.756,4.88,0,0,0,standard,suburban
-store_B,29,horizon,781.04,789.76,10.436,-8.72,0,0,1,standard,suburban
-store_B,30,horizon,844.57,837.86,9.6646,6.71,0,0,2,standard,suburban
-store_B,31,horizon,863.43,854.33,9.5449,9.1,0,0,3,standard,suburban
-store_B,32,horizon,898.12,896.82,9.9351,1.3,0,0,4,standard,suburban
-store_B,33,horizon,1070.58,930.42,10.4924,-9.85,1,0,5,standard,suburban
-store_B,34,horizon,820.4,828.24,10.3917,-7.83,0,0,6,standard,suburban
-store_B,35,horizon,965.86,770.83,10.2486,-4.97,0,1,0,standard,suburban
-store_C,0,context,709.12,501.23,7.1053,7.89,0,1,0,discount,rural
-store_C,1,context,651.44,492.78,7.0666,8.67,1,0,1,discount,rural
-store_C,2,context,659.15,511.04,7.5944,-1.89,1,0,2,discount,rural
-store_C,3,context,733.06,575.98,7.1462,7.08,1,0,3,discount,rural
-store_C,4,context,712.21,568.7,7.8247,-6.49,1,0,4,discount,rural
-store_C,5,context,615.23,611.44,7.3103,3.79,0,0,5,discount,rural
-store_C,6,context,568.99,561.87,7.1439,7.12,0,0,6,discount,rural
-store_C,7,context,541.12,549.54,7.921,-8.42,0,0,0,discount,rural
-store_C,8,context,583.57,576.88,7.1655,6.69,0,0,1,discount,rural
-store_C,9,context,607.34,603.04,7.2847,4.31,0,0,2,discount,rural
-store_C,10,context,613.79,606.86,7.1536,6.93,0,0,3,discount,rural
-store_C,11,context,919.49,561.8,7.1155,7.69,1,1,4,discount,rural
-store_C,12,context,622.61,613.04,7.0211,9.58,0,0,5,discount,rural
-store_C,13,context,630.52,621.63,7.0554,8.89,0,0,6,discount,rural
-store_C,14,context,721.62,715.12,7.1746,6.51,0,0,0,discount,rural
-store_C,15,context,699.18,690.25,7.0534,8.93,0,0,1,discount,rural
-store_C,16,context,578.85,580.67,7.5911,-1.82,0,0,2,discount,rural
-store_C,17,context,598.23,601.84,7.6807,-3.61,0,0,3,discount,rural
-store_C,18,context,554.43,552.3,7.3936,2.13,0,0,4,discount,rural
-store_C,19,context,587.39,583.75,7.318,3.64,0,0,5,discount,rural
-store_C,20,context,615.58,615.67,7.5045,-0.09,0,0,6,discount,rural
-store_C,21,context,638.68,646.18,7.875,-7.5,0,0,0,discount,rural
-store_C,22,context,555.99,563.01,7.8511,-7.02,0,0,1,discount,rural
-store_C,23,context,768.83,559.7,7.0435,9.13,0,1,2,discount,rural
-store_C,24,horizon,499.62,493.25,7.1815,6.37,0,0,3,discount,rural
-store_C,25,horizon,570.9,565.64,7.2367,5.27,0,0,4,discount,rural
-store_C,26,horizon,677.52,522.5,7.2494,5.01,1,0,5,discount,rural
-store_C,27,horizon,685.25,536.68,7.5712,-1.42,1,0,6,discount,rural
-store_C,28,horizon,517.46,515.78,7.4163,1.67,0,0,0,discount,rural
-store_C,29,horizon,549.38,540.36,7.0493,9.01,0,0,1,discount,rural
-store_C,30,horizon,470.04,467.51,7.3736,2.53,0,0,2,discount,rural
-store_C,31,horizon,622.9,473.37,7.5238,-0.48,1,0,3,discount,rural
-store_C,32,horizon,620.09,612.12,7.1017,7.97,0,0,4,discount,rural
-store_C,33,horizon,614.45,471.12,7.8335,-6.67,1,0,5,discount,rural
-store_C,34,horizon,484.25,475.29,7.052,8.96,0,0,6,discount,rural
-store_C,35,horizon,781.64,590.14,7.9248,-8.5,0,1,0,discount,rural

timesfm-forecasting/examples/global-temperature/output/forecast_output.csv

@@ -1,13 +0,0 @@
-date,point_forecast,q10,q20,q30,q40,q50,q60,q70,q80,q90,q99
-2025-01-01,1.2593384,1.248188,1.140702,1.1880752,1.2137158,1.2394564,1.2593384,1.2767732,1.297132,1.32396,1.367888
-2025-02-01,1.2856668,1.2773758,1.1406044,1.1960833,1.2322671,1.2593892,1.2856668,1.3110137,1.3400218,1.3751202,1.4253658
-2025-03-01,1.2950127,1.2869918,1.126852,1.1876173,1.234988,1.2675052,1.2950127,1.328448,1.354729,1.4035482,1.4642649
-2025-04-01,1.2207624,1.2084007,1.0352504,1.1041918,1.151865,1.1853008,1.2207624,1.256663,1.2898555,1.3310349,1.4016538
-2025-05-01,1.1702554,1.153313,0.9691495,1.0431063,1.0932612,1.1276176,1.1702554,1.201966,1.2390311,1.2891905,1.3632389
-2025-06-01,1.1455553,1.1275499,0.94203794,1.0110554,1.0658777,1.1061188,1.1455553,1.1806211,1.2180579,1.2702757,1.345366
-2025-07-01,1.1702348,1.1510556,0.9503718,1.0347577,1.0847733,1.1287677,1.1702348,1.2114835,1.2482276,1.2997853,1.3807325
-2025-08-01,1.2026825,1.1859496,0.9709255,1.0594383,1.1106675,1.1579902,1.2026825,1.2399211,1.2842004,1.3408126,1.419526
-2025-09-01,1.1909748,1.1784849,0.95943713,1.0403702,1.103606,1.1511956,1.1909748,1.2390201,1.2832941,1.3354731,1.416972
-2025-10-01,1.1490841,1.1264795,0.9079477,0.99529266,1.0548235,1.1052223,1.1490841,1.1897774,1.240414,1.2868769,1.3775467
-2025-11-01,1.0804785,1.0624356,0.8361266,0.9259792,0.9882403,1.0386353,1.0804785,1.1281581,1.1759715,1.228377,1.3122478
-2025-12-01,1.0613453,1.0366092,0.80220693,0.89521873,0.9593707,1.0152239,1.0613453,1.1032857,1.15315,1.216908,1.2959521

timesfm-forecasting/examples/global-temperature/output/forecast_output.json

@@ -1,188 +0,0 @@
-{
-  "model": "TimesFM 1.0 (200M) PyTorch",
-  "input": {
-    "source": "NOAA GISTEMP Global Temperature Anomaly",
-    "n_observations": 36,
-    "date_range": "2022-01 to 2024-12",
-    "mean_anomaly_c": 1.09
-  },
-  "forecast": {
-    "horizon": 12,
-    "dates": [
-      "2025-01",
-      "2025-02",
-      "2025-03",
-      "2025-04",
-      "2025-05",
-      "2025-06",
-      "2025-07",
-      "2025-08",
-      "2025-09",
-      "2025-10",
-      "2025-11",
-      "2025-12"
-    ],
-    "point": [
-      1.25933837890625,
-      1.285666823387146,
-      1.2950127124786377,
-      1.2207623720169067,
-      1.170255422592163,
-      1.1455552577972412,
-      1.1702347993850708,
-      1.2026824951171875,
-      1.1909748315811157,
-      1.1490840911865234,
-      1.080478549003601,
-      1.0613453388214111
-    ],
-    "quantiles": {
-      "10%": [
-        1.2481880187988281,
-        1.2773758172988892,
-        1.286991834640503,
-        1.2084007263183594,
-        1.1533130407333374,
-        1.1275498867034912,
-        1.1510555744171143,
-        1.1859495639801025,
-        1.1784849166870117,
-        1.1264795064926147,
-        1.0624356269836426,
-        1.036609172821045
-      ],
-      "20%": [
-        1.1407020092010498,
-        1.1406043767929077,
-        1.126852035522461,
-        1.0352504253387451,
-        0.9691494703292847,
-        0.9420379400253296,
-        0.9503718018531799,
-        0.970925509929657,
-        0.9594371318817139,
-        0.9079477190971375,
-        0.8361266255378723,
-        0.8022069334983826
-      ],
-      "30%": [
-        1.1880751848220825,
-        1.1960833072662354,
-        1.187617301940918,
-        1.104191780090332,
-        1.0431063175201416,
-        1.01105535030365,
-        1.0347577333450317,
-        1.0594383478164673,
-        1.040370225906372,
-        0.9952926635742188,
-        0.9259791970252991,
-        0.8952187299728394
-      ],
-      "40%": [
-        1.2137157917022705,
-        1.232267141342163,
-        1.2349879741668701,
-        1.151865005493164,
-        1.0932612419128418,
-        1.0658776760101318,
-        1.084773302078247,
-        1.1106674671173096,
-        1.1036059856414795,
-        1.0548235177993774,
-        0.9882403016090393,
-        0.9593706727027893
-      ],
-      "50%": [
-        1.2394564151763916,
-        1.2593891620635986,
-        1.267505168914795,
-        1.1853008270263672,
-        1.127617597579956,
-        1.1061187982559204,
-        1.128767728805542,
-        1.1579902172088623,
-        1.1511956453323364,
-        1.1052223443984985,
-        1.03863525390625,
-        1.0152238607406616
-      ],
-      "60%": [
-        1.25933837890625,
-        1.285666823387146,
-        1.2950127124786377,
-        1.2207623720169067,
-        1.170255422592163,
-        1.1455552577972412,
-        1.1702347993850708,
-        1.2026824951171875,
-        1.1909748315811157,
-        1.1490840911865234,
-        1.080478549003601,
-        1.0613453388214111
-      ],
-      "70%": [
-        1.27677321434021,
-        1.3110136985778809,
-        1.3284480571746826,
-        1.2566629648208618,
-        1.2019660472869873,
-        1.1806211471557617,
-        1.2114834785461426,
-        1.2399210929870605,
-        1.2390201091766357,
-        1.1897773742675781,
-        1.1281580924987793,
-        1.1032856702804565
-      ],
-      "80%": [
-        1.2971320152282715,
-        1.3400218486785889,
-        1.3547290563583374,
-        1.2898554801940918,
-        1.2390310764312744,
-        1.2180578708648682,
-        1.248227596282959,
-        1.2842004299163818,
-        1.2832940816879272,
-        1.240414023399353,
-        1.175971508026123,
-        1.153149962425232
-      ],
-      "90%": [
-        1.3239599466323853,
-        1.3751201629638672,
-        1.403548240661621,
-        1.3310348987579346,
-        1.2891905307769775,
-        1.2702757120132446,
-        1.2997852563858032,
-        1.3408125638961792,
-        1.3354730606079102,
-        1.286876916885376,
-        1.2283769845962524,
-        1.2169079780578613
-      ],
-      "99%": [
-        1.3678879737854004,
-        1.4253658056259155,
-        1.4642648696899414,
-        1.40165376663208,
-        1.3632389307022095,
-        1.3453660011291504,
-        1.380732536315918,
-        1.4195259809494019,
-        1.416972041130066,
-        1.3775466680526733,
-        1.3122477531433105,
-        1.2959520816802979
-      ]
-    }
-  },
-  "summary": {
-    "forecast_mean_c": 1.186,
-    "forecast_max_c": 1.295,
-    "forecast_min_c": 1.061,
-    "vs_last_year_mean": -0.067
-  }
-}

timesfm-forecasting/examples/global-temperature/run_forecast.py

@@ -11,6 +11,7 @@ from pathlib import Path
 
 import numpy as np
 import pandas as pd
+import timesfm
 
 # Preflight check
 print("=" * 60)
@@ -35,27 +36,28 @@ print(
 # TimesFM expects a list of 1D numpy arrays
 input_series = df["anomaly_c"].values.astype(np.float32)
 
-# Load TimesFM 1.0 (PyTorch)
+# Load TimesFM 2.5 (PyTorch)
-# NOTE: TimesFM 2.5 PyTorch checkpoint has a file format issue at time of writing.
+print("\n🤖 Loading TimesFM 2.5 (200M) PyTorch...")
-# The model.safetensors file is not loadable via torch.load().
-# Using TimesFM 1.0 PyTorch which works correctly.
-print("\n🤖 Loading TimesFM 1.0 (200M) PyTorch...")
-import timesfm
 
-hparams = timesfm.TimesFmHparams(horizon_len=12)
+model = timesfm.TimesFM_2p5_200M_torch.from_pretrained(
-checkpoint = timesfm.TimesFmCheckpoint(
+    "google/timesfm-2.5-200m-pytorch",
-    huggingface_repo_id="google/timesfm-1.0-200m-pytorch"
+    torch_compile=False,
 )
-model = timesfm.TimesFm(hparams=hparams, checkpoint=checkpoint)
+model.compile(timesfm.ForecastConfig(
+    max_context=512,
+    max_horizon=12,
+    normalize_inputs=True,
+    use_continuous_quantile_head=True,
+    fix_quantile_crossing=True,
+))
 
 # Forecast
 print("\n📈 Running forecast (12 months ahead)...")
 forecast_input = [input_series]
-frequency_input = [0]  # Monthly data
 
 point_forecast, experimental_quantile_forecast = model.forecast(
-    forecast_input,
+    horizon=12,
-    freq=frequency_input,
+    inputs=forecast_input,
 )
 
 print(f"   Point forecast shape: {point_forecast.shape}")
@@ -65,9 +67,8 @@ print(f"   Quantile forecast shape: {experimental_quantile_forecast.shape}")
 point = point_forecast[0]  # Shape: (horizon,)
 quantiles = experimental_quantile_forecast[0]  # Shape: (horizon, num_quantiles)
 
-# TimesFM quantiles: [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.99]
+# TimesFM 2.5 columns: 0=mean, 1=10%, 2=20%, ..., 5=50% (median), ..., 9=90%
-# Index mapping: 0=10%, 1=20%, ..., 4=50% (median), ..., 9=99%
+quantile_labels = ["mean", "10%", "20%", "30%", "40%", "50%", "60%", "70%", "80%", "90%"]
-quantile_labels = ["10%", "20%", "30%", "40%", "50%", "60%", "70%", "80%", "90%", "99%"]
 
 # Create forecast dates (2025 monthly)
 last_date = df["date"].max()
@@ -80,16 +81,16 @@ output_df = pd.DataFrame(
     {
         "date": forecast_dates.strftime("%Y-%m-%d"),
         "point_forecast": point,
-        "q10": quantiles[:, 0],
+        "mean": quantiles[:, 0],
-        "q20": quantiles[:, 1],
+        "q10": quantiles[:, 1],
-        "q30": quantiles[:, 2],
+        "q20": quantiles[:, 2],
-        "q40": quantiles[:, 3],
+        "q30": quantiles[:, 3],
-        "q50": quantiles[:, 4],  # Median
+        "q40": quantiles[:, 4],
-        "q60": quantiles[:, 5],
+        "q50": quantiles[:, 5],  # Median
-        "q70": quantiles[:, 6],
+        "q60": quantiles[:, 6],
-        "q80": quantiles[:, 7],
+        "q70": quantiles[:, 7],
-        "q90": quantiles[:, 8],
+        "q80": quantiles[:, 8],
-        "q99": quantiles[:, 9],
+        "q90": quantiles[:, 9],
     }
 )
 
@@ -100,7 +101,7 @@ output_df.to_csv(output_dir / "forecast_output.csv", index=False)
 
 # JSON output for the report
 output_json = {
-    "model": "TimesFM 1.0 (200M) PyTorch",
+    "model": "TimesFM 2.5 (200M) PyTorch",
     "input": {
         "source": "NOAA GISTEMP Global Temperature Anomaly",
         "n_observations": len(df),
@@ -135,24 +136,24 @@ print("=" * 60)
 print(
     f"\n📅 Forecast period: {forecast_dates[0].strftime('%Y-%m')} to {forecast_dates[-1].strftime('%Y-%m')}"
 )
-print(f"\n🌡️  Temperature Anomaly Forecast (°C above 1951-1980 baseline):")
+print("\n🌡️  Temperature Anomaly Forecast (°C above 1951-1980 baseline):")
-print(f"\n   {'Month':<10} {'Point':>8} {'80% CI':>15} {'90% CI':>15}")
+print(f"\n   {'Month':<10} {'Point':>8} {'60% CI':>15} {'80% CI':>15}")
 print(f"   {'-' * 10} {'-' * 8} {'-' * 15} {'-' * 15}")
-for i, (date, pt, q10, q90, q05, q95) in enumerate(
+for i, (date, pt, q20, q80, q10, q90) in enumerate(
     zip(
         forecast_dates.strftime("%Y-%m"),
         point,
-        quantiles[:, 1],  # 20%
+        quantiles[:, 2],  # 20%
-        quantiles[:, 7],  # 80%
+        quantiles[:, 8],  # 80%
-        quantiles[:, 0],  # 10%
+        quantiles[:, 1],  # 10%
-        quantiles[:, 8],  # 90%
+        quantiles[:, 9],  # 90%
     )
 ):
     print(
-        f"   {date:<10} {pt:>8.3f} [{q10:>6.3f}, {q90:>6.3f}] [{q05:>6.3f}, {q95:>6.3f}]"
+        f"   {date:<10} {pt:>8.3f} [{q20:>6.3f}, {q80:>6.3f}] [{q10:>6.3f}, {q90:>6.3f}]"
     )
 
-print(f"\n📊 Summary Statistics:")
+print("\n📊 Summary Statistics:")
 print(f"   Mean forecast:  {point.mean():.3f}°C")
 print(
     f"   Max forecast:   {point.max():.3f}°C (Month: {forecast_dates[point.argmax()].strftime('%Y-%m')})"
@@ -162,6 +163,6 @@ print(
 )
 print(f"   vs 2024 mean:   {point.mean() - df['anomaly_c'].iloc[-12:].mean():+.3f}°C")
 
-print(f"\n✅ Output saved to:")
+print("\n✅ Output saved to:")
 print(f"   {output_dir / 'forecast_output.csv'}")
 print(f"   {output_dir / 'forecast_output.json'}")

timesfm-forecasting/examples/global-temperature/visualize_forecast.py

@@ -16,6 +16,8 @@ from __future__ import annotations
 import json
 from pathlib import Path
 
+import matplotlib
+matplotlib.use("Agg")
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
@@ -57,11 +59,11 @@ def main() -> None:
         label="Historical (NOAA GISTEMP)",
     )
 
-    # Plot 90% CI (outer band)
+    # Plot 80% CI (outer band)
-    ax.fill_between(dates, q10, q90, alpha=0.2, color="#dc2626", label="90% CI")
+    ax.fill_between(dates, q10, q90, alpha=0.2, color="#dc2626", label="80% CI")
 
-    # Plot 80% CI (inner band)
+    # Plot 60% CI (inner band)
-    ax.fill_between(dates, q20, q80, alpha=0.3, color="#dc2626", label="80% CI")
+    ax.fill_between(dates, q20, q80, alpha=0.3, color="#dc2626", label="60% CI")
 
     # Plot point forecast
     ax.plot(