chronax

A high-performance, JAX-accelerated time-series forecasting library. chronax provides a comprehensive suite of classical and modern forecasting models — including AutoARIMA, AutoETS, AutoTheta, TBATS, MFLES, GARCH, and more — with a unified fit / predict interface and hardware-accelerated execution via JAX.

Features

• ⚡ JAX-accelerated — JIT-compiled model fitting and forecasting on CPU, GPU, or TPU
• 📈 20+ forecasting models including AutoARIMA, AutoETS, AutoTheta, TBATS, MFLES, GARCH, STL, and more
• 🔁 Unified API — every model follows the same fit() → predict() pattern
• 📊 Prediction intervals — built-in conformal and native interval support
• ✅ NumPy compatible — accepts and returns standard array types
• 🧪 Benchmarked against established libraries for correctness and speed

Installation

From PyPI (recommended)

Requires Python ≥ 3.11

pip install chronax

From TestPyPI (pre-release testing)

pip install -i https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ chronax

From GitHub source

Install directly from the latest commit:

pip install git+https://github.com/Smlcrm/ml-library-chronax.git

For local development:

git clone https://github.com/Smlcrm/ml-library-chronax.git
cd ml-library-chronax
python -m venv .venv
source .venv/bin/activate  # On Windows use: .venv\Scripts\activate
pip install -e .

Usage Overview

Fitting a model and forecasting

All chronax models follow the same interface: instantiate, fit(), then predict().

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

# Sample time series
y = jnp.array([112, 118, 132, 129, 121, 135, 148, 148, 136, 119, 104, 118,
               115, 126, 141, 135, 125, 149, 170, 170, 158, 133, 114, 140])

# Fit
model = AutoARIMA(season_length=12)
model = model.fit(y)

# Forecast 6 steps ahead
forecast = model.predict(h=6)
print("Forecast:", forecast['mean'])

Prediction intervals

Request probabilistic forecasts by passing level:

forecast = model.predict(h=6, level=[80, 95])

print("Point forecast:", forecast['mean'])
print("95% lower:", forecast['lo-95'])
print("95% upper:", forecast['hi-95'])

In-sample fitted values

Retrieve the model's in-sample predictions after fitting:

insample = model.predict_in_sample()
print("Fitted values:", insample['fitted'])

Memory-efficient forecasting

Use forecast() to fit and predict in a single call without storing model state:

from chronax.models import AutoETS

model = AutoETS(season_length=12)
result = model.forecast(y, h=6, level=[90])
print("Forecast:", result['mean'])

Comparing multiple models

from chronax.models import AutoARIMA, AutoETS, AutoTheta

models = {
    "AutoARIMA": AutoARIMA(season_length=12),
    "AutoETS": AutoETS(season_length=12),
    "AutoTheta": AutoTheta(season_length=12),
}

for name, m in models.items():
    m = m.fit(y)
    pred = m.predict(h=6)
    print(f"{name}: {pred['mean']}")

Available Models

Automatic Forecasting

Automatic model-selection wrappers that search over candidate configurations.

ARIMA Family

Autoregressive integrated moving-average models for autocorrelated series.

Theta Family

Theta-method forecasters for trend and seasonality decomposition.

Multiple Seasonalities & Decomposition

Models designed for multiple seasonal patterns or explicit trend-seasonal decomposition.

Volatility Models

Models specialized for time-varying variance and heteroskedastic dynamics.

Baseline Models

Simple reference forecasters used as strong, interpretable baselines.

Exponential Smoothing

Level, trend, and seasonal smoothing models with recursive state updates.

Sparse / Intermittent Demand

Forecasters tailored to sparse series with many zeros or irregular demand arrivals.

All models are importable from chronax.models.

Tutorial: Forecast a Time Series in Five Steps

1. Install the package

pip install chronax

2. Create a project

mkdir my-forecast && cd my-forecast
python -m venv .venv
source .venv/bin/activate  # On Windows: .venv\Scripts\activate
pip install chronax

3. Write a forecast script (forecast.py)

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

def main() -> None:
    # Monthly airline passengers (subset)
    y = jnp.array([112, 118, 132, 129, 121, 135, 148, 148, 136, 119, 104, 118,
                    115, 126, 141, 135, 125, 149, 170, 170, 158, 133, 114, 140])

    model = AutoETS(season_length=12)
    model = model.fit(y)

    forecast = model.predict(h=6, level=[80, 95])
    print("Point forecast:", forecast['mean'].tolist())
    print("80% interval:", list(zip(forecast['lo-80'].tolist(), forecast['hi-80'].tolist())))
    print("95% interval:", list(zip(forecast['lo-95'].tolist(), forecast['hi-95'].tolist())))

if __name__ == "__main__":
    main()

4. Run the script

python forecast.py

5. Explore further

Try swapping AutoETS for AutoARIMA or AutoTheta and compare results — the API is identical across all models.

Evaluation Benchmarks

We provide a benchmarking suite to evaluate chronax against other time-series libraries.

Note: Benchmark dependencies (such as statsforecast, pandas, etc.) are not included in the core package to keep the installation lightweight.

# Install benchmark dependencies
pip install statsforecast pandas matplotlib

# Run the benchmark suite
python benchmarks/benchmark_suite.py

Documentation

Explore inline docstrings for detailed parameter and return-type information.

License

MIT © Simulacrum, Inc.