Forecast a univariate series using the Pre-trained TimesFM Model

This guide shows you how to use the pre-trained TimesFM foundation model within Chronax to generate point forecasts for a univariate time series. TimesFM is used for inference directly, requiring no explicit training step on your data.

Prerequisites

Chronax installed (pip install chronax).
The input time series y must be a 1-D jnp.ndarray of type float32.
You must specify a frequency indicator (0, 1, or 2) corresponding to the data granularity.

Steps

1. Load the TimesFM foundation model

TimesFM is a large pre-trained model. We assume Chronax provides a utility to load the checkpoint directly into the JAX environment.

from chronax.models import TimesFM
import jax.numpy as jnp

# TimesFM requires complex configuration related to JAX backend and checkpoint IDs.
# We assume Chronax provides a simplified loading mechanism for the 500M checkpoint.
# TODO: confirm TimesFM loading API in Chronax
model = TimesFM.load_pretrained("timesfm-2.0-500m")

2. Prepare the input data and frequency indicator

Define your context series y (the data used for forecasting) and the forecast horizon h. TimesFM also requires a frequency indicator: 0 (high frequency, e.g., daily or finer), 1 (medium frequency, e.g., weekly/monthly), or 2 (low frequency, e.g., quarterly/yearly).

# Sample univariate time series (context)
y = jnp.array([10.0, 12.0, 15.0, 14.0, 16.0, 18.0, 20.0, 22.0, 25.0], dtype=jnp.float32)
h = 5  # Forecast 5 steps ahead

# Use 0 for high frequency data (default recommendation)
frequency_indicator = 0

3. Generate the forecast

Call the model.forecast() method, passing the context series y, the horizon h, and the required frequency indicator. Since TimesFM is a foundation model, no prior fit step is necessary.

# Generate the forecast
forecast_output = model.forecast(
    y=y,
    h=h,
    frequency=frequency_indicator
)

# Extract the mean forecast from the returned dictionary
mean_forecast = forecast_output["mean"]

print(f"Input context length: {len(y)}")
print(f"Forecast horizon (h): {h}")
print(f"Mean forecast: {mean_forecast}")

Full example

from chronax.models import TimesFM
import jax.numpy as jnp

# 1. Load the TimesFM foundation model
# TODO: confirm TimesFM loading API in Chronax
model = TimesFM.load_pretrained("timesfm-2.0-500m")

# 2. Prepare the input data and frequency indicator
y = jnp.array([10.0, 12.0, 15.0, 14.0, 16.0, 18.0, 20.0, 22.0, 25.0], dtype=jnp.float32)
h = 5
frequency_indicator = 0 # High frequency

# 3. Generate the forecast
forecast_output = model.forecast(
    y=y,
    h=h,
    frequency=frequency_indicator
)

# Extract the mean forecast
mean_forecast = forecast_output["mean"]

print(f"Input context: {y}")
print(f"Mean forecast (h={h}): {mean_forecast}")

Next steps

Explore how to use the frequency parameter to optimize results for different granularities (0, 1, 2).
Investigate the use of model.forecast_on_df if you are working with pandas.DataFrame inputs.
Review the TimesFM documentation for maximum context length limitations (e.g., 2048).