Forecast a Univariate Series

This guide demonstrates the fundamental workflow for forecasting a single time series using Chronax. Since Chronax models are built on JAX, the first time you call fit, JAX performs Just-In-Time (JIT) compilation (the "Cold Start"), making all subsequent prediction calls extremely fast (the "Warm Run").

Prerequisites

• Chronax and JAX installed.
• The history array, y, must be a 1-D jnp.ndarray of type float32.
• We will use AutoARIMA as the example model.

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

Steps

1. Prepare the history data

Define your historical time series data (y) as a JAX array. Chronax requires float32 precision for optimal performance on accelerators.

# Example: 100 points of synthetic data
y = jnp.arange(100, dtype=jnp.float32) + jnp.sin(jnp.linspace(0, 10, 100))

2. Initialize the model

Instantiate the forecasting model you wish to use. Chronax models are stateless until they are fitted.

# Initialize the AutoARIMA model
model = AutoARIMA()

3. Fit the model (The Cold Start)

Call the fit method, passing the history array y. This step triggers the JAX JIT compilation (XLA), which may take a few seconds on the first run. This is the "Cold Start" phase.

# Fit the model to the historical data
fitted_model = model.fit(y)

4. Generate the forecast (The Warm Run)

Use the predict method on the fitted model to generate the forecast for the next h steps. Since the model is already compiled, this prediction step is highly optimized and represents the "Warm Run" performance.

We will forecast 10 steps ahead (h=10).

h = 10
forecast_output = fitted_model.predict(h=h)

5. Access the mean forecast

The predict method returns a dictionary containing the results. Access the primary forecast using the "mean" key.

mean_forecast = forecast_output["mean"]

print(f"Forecast shape: {mean_forecast.shape}")
print(f"First 5 forecast values: {mean_forecast[:5]}")

Full example

This complete example demonstrates the standard Chronax workflow, separating the compilation/training phase (fit) from the fast inference phase (predict).

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

# 1. Prepare the history data (must be float32)
y = jnp.arange(100, dtype=jnp.float32) + jnp.sin(jnp.linspace(0, 10, 100))
h = 10

# 2. Initialize the model
model = AutoARIMA()

# 3. Fit the model (Cold Start: JIT compilation occurs here)
print("Starting fit (JIT compilation)...")
fitted_model = model.fit(y)
print("Fit complete.")

# 4. Generate the forecast (Warm Run: Fast inference)
forecast_output = fitted_model.predict(h=h)

# 5. Access the mean forecast
mean_forecast = forecast_output["mean"]

print(f"\nForecast for {h} steps:")
print(mean_forecast)

Next steps

• To measure the overhead of generating prediction intervals, see the guide on adding level to the predict method.
• To use a single function call for both fitting and forecasting, explore the model.forecast() method.
• Learn how to incorporate external information using exogenous features (X).
• Explore other models available in chronax.models.