Forecast a univariate series using a Neural Model (GRU)

This guide shows you how to initialize, train, and forecast a single time series using a Chronax neural model, such as the GRU model. Use this approach when you need high performance and are working with large datasets where traditional statistical models are too slow.

Prerequisites

Install Chronax and JAX.
Import models from chronax.models.
The input time series y must be a 1-D jnp.ndarray of type float32.

import jax.numpy as jnp
from chronax.models import GRU # Assuming GRU is available

Steps

1. Prepare the input data

Define your historical time series y. Chronax models expect JAX arrays (jnp.ndarray) for all inputs.

# Example: 100 historical observations
y = jnp.arange(100, dtype=jnp.float32) + jnp.sin(jnp.linspace(0, 10, 100))

2. Initialize the GRU model

Initialize the model. Neural models often require parameters defining the architecture (e.g., number of layers, hidden size) and the lookback window (input_size). Since these parameters are not specified, they are marked with TODO.

# Initialize the model. Parameters must be confirmed against the GRU API.
model = GRU(
    # TODO: confirm input_size (lookback window)
    input_size=30,
    # TODO: confirm hidden_size
    hidden_size=64,
    # TODO: confirm learning rate or optimizer settings
)

3. Fit the model to the historical data

Use the fit method to train the model on the historical series y. Neural models require an explicit number of training epochs.

# Fit the model for a specified number of epochs
# TODO: confirm required arguments for fit (e.g., epochs, batch_size)
fit_result = model.fit(
    y,
    epochs=50,
    batch_size=32
)
print("Model fitted successfully.")

4. Predict the future horizon

Use the predict method, specifying the forecast horizon h. The result is a dictionary containing the forecast mean and potentially other metrics.

# Forecast 10 steps into the future
h = 10
forecast_output = model.predict(h=h)

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

print(f"Forecast shape: {mean_forecast.shape}")

Full example

This example combines the steps to train and forecast a simple synthetic series using the GRU model.

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

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

# 2. Initialize the GRU model
model = GRU(
    # TODO: confirm input_size (lookback window)
    input_size=30,
    # TODO: confirm hidden_size
    hidden_size=64,
    # TODO: confirm learning rate or optimizer settings
)

# 3. Fit the model to the historical data
fit_result = model.fit(
    y,
    epochs=50,
    batch_size=32
)

# 4. Predict the future horizon
forecast_output = model.predict(h=h)
mean_forecast = forecast_output["mean"]

print(f"Historical data length: {len(y)}")
print(f"Forecast horizon (h): {h}")
print(f"Mean forecast (first 5 values): {mean_forecast[:5]}")

Next steps

Consult the guide on "Adding prediction intervals" to quantify uncertainty in the forecast using the level parameter in predict.
Explore other neural models like PatchTST or KAN.
Learn how to incorporate exogenous features using the X argument in fit and predict.