Update global-temperature example to use TimesFM 2.5 API

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)
-# NOTE: TimesFM 2.5 PyTorch checkpoint has a file format issue at time of writing.
-# 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
+# Load TimesFM 2.5 (PyTorch)
+print("\n🤖 Loading TimesFM 2.5 (200M) PyTorch...")
 
-hparams = timesfm.TimesFmHparams(horizon_len=12)
-checkpoint = timesfm.TimesFmCheckpoint(
-    huggingface_repo_id="google/timesfm-1.0-200m-pytorch"
+model = timesfm.TimesFM_2p5_200M_torch.from_pretrained(
+    "google/timesfm-2.5-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,
-    freq=frequency_input,
+    horizon=12,
+    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]
-# Index mapping: 0=10%, 1=20%, ..., 4=50% (median), ..., 9=99%
-quantile_labels = ["10%", "20%", "30%", "40%", "50%", "60%", "70%", "80%", "90%", "99%"]
+# TimesFM 2.5 columns: 0=mean, 1=10%, 2=20%, ..., 5=50% (median), ..., 9=90%
+quantile_labels = ["mean", "10%", "20%", "30%", "40%", "50%", "60%", "70%", "80%", "90%"]
 
 # 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],
-        "q20": quantiles[:, 1],
-        "q30": quantiles[:, 2],
-        "q40": quantiles[:, 3],
-        "q50": quantiles[:, 4],  # Median
-        "q60": quantiles[:, 5],
-        "q70": quantiles[:, 6],
-        "q80": quantiles[:, 7],
-        "q90": quantiles[:, 8],
-        "q99": quantiles[:, 9],
+        "mean": quantiles[:, 0],
+        "q10": quantiles[:, 1],
+        "q20": quantiles[:, 2],
+        "q30": quantiles[:, 3],
+        "q40": quantiles[:, 4],
+        "q50": quantiles[:, 5],  # Median
+        "q60": quantiles[:, 6],
+        "q70": quantiles[:, 7],
+        "q80": quantiles[:, 8],
+        "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(f"\n   {'Month':<10} {'Point':>8} {'80% CI':>15} {'90% CI':>15}")
+print("\n🌡️  Temperature Anomaly Forecast (°C above 1951-1980 baseline):")
+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[:, 7],  # 80%
-        quantiles[:, 0],  # 10%
-        quantiles[:, 8],  # 90%
+        quantiles[:, 2],  # 20%
+        quantiles[:, 8],  # 80%
+        quantiles[:, 1],  # 10%
+        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'}")