utils.py

Shared utilities for all Chronax forecasting models.

results

chronax.utils.results

Named tuple returned by optimization routines, containing the optimization result.

Attribute	Type	Description
x	-	The optimal parameter value.
fn	-	The objective function value at x.
nit	-	Number of iterations.
simplex	-	(undocumented)

ensure_float

chronax.utils.ensure_float(y)

Cast array to float32 if it is not already a floating-point dtype.

Parameter	Type	Default	Description
y	jnp.ndarray	-	Input JAX array of any dtype.

Returns: jnp.ndarray (The same array if already floating-point, otherwise cast to float32.)

calculate_sigma

chronax.utils.calculate_sigma(residuals, n)

Compute the root-mean-square of residuals (RMS sigma).

| Parameter | Type | Default | Description | |-----------|---------------|-------------| | residuals | jnp.ndarray | - | Residual values as a JAX array. | | n | int | - | Number of degrees of freedom (denominator). |

Returns: jnp.ndarray (Scalar sigma value; returns 0.0 when n <= 0.)

extract_demand

chronax.utils.extract_demand(y)

Extract positive (non-zero) demand values from a time series.

Used for intermittent demand models like TSB and Croston, where we need to separate demand occurrences from no-demand periods.

Parameter	Type	Default	Description
y	jnp.ndarray	-	Time series array that may contain zeros.

Returns: jnp.ndarray (Array containing only positive values from y.)

extract_probability

chronax.utils.extract_probability(y)

Convert time series to binary indicator (1=demand, 0=no demand).

Used for intermittent demand models like TSB to track the probability of demand occurrence at each time step.

Parameter	Type	Default	Description
y	jnp.ndarray	-	Time series array.

Returns: jnp.ndarray (Binary array where 1 indicates demand occurred, 0 indicates no demand.)

_repeat_val

chronax.utils._repeat_val(val, h)

Repeat scalar value h times.

JAX equivalent of statsforecast.utils._repeat_val().

Parameter	Type	Default	Description
val	float	-	Scalar value to repeat.
h	int	-	Number of repetitions (forecast horizon).

Returns: jnp.ndarray (Array of length h filled with val.)

_repeat_val_seas

chronax.utils._repeat_val_seas(season_vals, h)

Tile seasonal values to cover forecast horizon h.

JAX equivalent of statsforecast.utils._repeat_val_seas().

Parameter	Type	Default	Description
season_vals	jnp.ndarray	-	Seasonal pattern of shape (season_length,).
h	int	-	Forecast horizon (static — must be known at compile time).

Returns: jnp.ndarray (Tiled pattern of length h.)

_quantiles

chronax.utils._quantiles(level)

Convert confidence levels to z-scores using the normal inverse CDF.

JAX equivalent of statsforecast.utils._quantiles().

Parameter	Type	Default	Description
level	List[Union[int, float]]	-	List of confidence levels in [0, 100], e.g. [80, 95].

Returns: jnp.ndarray (Array of z-scores, one per level.)

_calculate_intervals

chronax.utils._calculate_intervals(res, level, h, sigmah)

Calculate native (non-conformal) prediction intervals using normal quantiles.

Parameter	Type	Default	Description
res	dict	-	Forecast result dict containing 'mean'.
level	List[int]	-	List of confidence levels (0-100).
h	int	-	Forecast horizon.
sigmah	Union[jnp.ndarray, float]	-	Standard error (scalar or array of length h).

Returns: dict (Dict with 'lo-{lv}' and 'hi-{lv}' keys for each level.)

_add_fitted_pi

chronax.utils._add_fitted_pi(res, se, level)

Add in-sample prediction intervals to a fitted result dict.

Used by theta/HW/ETS models. Works with scalar or vector se via reshaping.

Parameter	Type	Default	Description
res	dict	-	Result dict containing 'fitted'.
se	jnp.ndarray	-	Standard error (scalar or vector).
level	Union[List[int], jnp.ndarray]	-	Confidence levels (0-100).

Returns: dict (Updated res dict with 'fitted-lo-{lv}' and 'fitted-hi-{lv}' keys.)

_add_fitted_pi_1

chronax.utils._add_fitted_pi_1(fitted, sigmah, level)

Calculate native (non-conformal) fitted (in-sample) prediction intervals.

JAX equivalent of statsforecast.models._add_fitted_pi(). Used by historic_average and croston_classic models.

Parameter	Type	Default	Description
fitted	jnp.ndarray	-	Fitted values of shape (t,).
sigmah	Union[jnp.ndarray, float]	-	Standard error for predictions (scalar or array).
level	List[int]	-	Sorted list of confidence levels (0-100).

Returns: dict (Dict with 'fitted-lo-{lv}' and 'fitted-hi-{lv}' keys for each level.)

_add_conformal_distribution_intervals

chronax.utils._add_conformal_distribution_intervals(fcst, cs, level)

Add conformal intervals to forecast dict based on conformal scores.

Creates forecast paths from errors and calculates quantiles.

Parameter	Type	Default	Description
fcst	dict	-	Forecast dict containing 'mean'.
cs	jnp.ndarray	-	Conformal scores array.
level	Union[List[float], List[int]]	-	Sorted list of confidence levels (0-100).

Returns: dict (Updated fcst dict with 'lo-{lv}' and 'hi-{lv}' keys.)

_get_conformal_method

chronax.utils._get_conformal_method(method)

Look up a conformal prediction interval method by name.

Parameter	Type	Default	Description
method	str	-	Method name (currently only 'conformal_distribution').

Returns: Callable (The corresponding interval function.) Raises: ValueError (If method is not supported.)

_conformal_method

chronax.utils._conformal_method(self)

Retrieve the conformal method from a model's prediction_intervals config.

Parameter	Type	Default	Description
self	-	-	A forecaster instance with prediction_intervals attribute.

Returns: Callable (The conformal interval function.)

_store_cs

chronax.utils._store_cs(self, y, X)

Compute and store conformal scores on the model instance.

Parameter	Type	Default	Description
self	-	-	A forecaster instance with prediction_intervals and conformity_scores.
y	jnp.ndarray	-	Training time series.
X	Optional[jnp.ndarray]	-	Optional exogenous variables.

_add_conformal_intervals

chronax.utils._add_conformal_intervals(self, fcst, y, X, level)

Add conformal prediction intervals to a forecast dict.

If y is provided, computes fresh conformal scores; otherwise uses stored scores.

Parameter	Type	Default	Description
self	-	-	A forecaster instance.
fcst	dict	-	Forecast dict to augment.
y	Optional[jnp.ndarray]	-	Training series (None to use stored scores).
X	Optional[jnp.ndarray]	-	Optional exogenous variables.
level	Optional[List[int]]	-	Confidence levels (0-100).

Returns: dict (Updated forecast dict with interval keys.)

_add_predict_conformal_intervals

chronax.utils._add_predict_conformal_intervals(self, fcst, level)

Add conformal intervals for the predict() path (uses stored scores).

Parameter	Type	Default	Description
self	-	-	A fitted forecaster instance.
fcst	dict	-	Forecast dict to augment.
level	Optional[List[int]]	-	Confidence levels (0-100).

Returns: dict (Updated forecast dict with interval keys.)

_seasonal_naive

chronax.utils._seasonal_naive(y, h, season_length, fitted=False)

JAX implementation of seasonal-naive forecast.

Repeats the last season_length observations as the forecast.

Parameter	Type	Default	Description
y	jnp.ndarray	-	1-D array-like (length T). Converted to float32.
h	int	-	Forecast horizon (int >= 1).
season_length	int	-	Seasonal period m (int >= 1).
fitted	bool	False	If True, also return in-sample fitted values.

Returns: Dict[str, jnp.ndarray] (Dict with 'mean' (shape (h,)) and optionally 'fitted' (shape (T,)).) Raises: ValueError (If y is not 1-D, season_length <= 0, T < season_length, or h < 1.)

_seasonal_exponential_smoothing

chronax.utils._seasonal_exponential_smoothing(y, h, fitted, season_length, alpha)

Seasonal exponential smoothing forecast.

Applies SES independently to each seasonal sub-series, then tiles the forecasts to cover horizon h.

Parameter	Type	Default	Description
y	jnp.ndarray	-	Input time series.
h	int	-	Forecast horizon.
fitted	bool	-	Whether to return in-sample fitted values.
season_length	int	-	Seasonal period.
alpha	float	-	Smoothing parameter for SES.

Returns: Dict[str, jnp.ndarray] (Dict with 'mean' and optionally 'fitted' keys.)

_window_average

chronax.utils._window_average(y, h, fitted, window_size)

Window average forecast.

Parameter	Type	Default	Description
y	jnp.ndarray	-	Time series.
h	int	-	Forecasting horizon.
fitted	bool	-	Whether to return fitted values (not implemented).
window_size	int	-	Window size for averaging.

Returns: Dict[str, jnp.ndarray] (Dict with 'mean' key containing constant forecast of length h.) Raises: NotImplementedError (If fitted=True.)

_intervals

chronax.utils._intervals(x)

Intervals between nonzero elements (IMAPA variant).

Unlike _intervals_c, returns a compact array of diffs (no NaN padding) and prepends the position of the first nonzero element.

Parameter	Type	Default	Description
x	jnp.ndarray	-	Input array.

Returns: jnp.ndarray (Float array of inter-arrival intervals.)

_intervals_c

chronax.utils._intervals_c(x)

Compute intervals between non-zero elements (Croston variant, JIT-compiled).

Returns fixed-size NaN-padded array for JIT compatibility. Used by Croston-family models.

Parameter	Type	Default	Description
x	jnp.ndarray	-	Input array.

Returns: jnp.ndarray (Fixed-size array with intervals packed at start, rest NaN.)

_expand_fitted_intervals

chronax.utils._expand_fitted_intervals(fitted, y)

Expand interval fitted values back to original series length (JIT-compiled).

Used by Croston-family models. Uses lax.fori_loop for JIT compatibility. Avoids division by zero by replacing zero fitted values with 1.

Parameter	Type	Default	Description
fitted	jnp.ndarray	-	SES fitted values for intervals (length = num_nonzero + 1).
y	jnp.ndarray	-	Original time series.

Returns: jnp.ndarray (Fitted intervals expanded to match y's length.)

_expand_fitted_demand

chronax.utils._expand_fitted_demand(fitted, y)

Expand demand fitted values back to original series length (JIT-compiled).

Used by Croston-family models. Uses lax.fori_loop for JIT compatibility.

Parameter	Type	Default	Description
fitted	jnp.ndarray	-	SES fitted values for demand (length = num_nonzero + 1).
y	jnp.ndarray	-	Original time series.

Returns: jnp.ndarray (Fitted values expanded to match y's length.)

_imapa

chronax.utils._imapa(y, h, fitted)

IMAPA forecaster in pure JAX (intermittent demand).

Detects inter-arrival spacing, computes mean interval as max aggregation level K, then for each k = 1..K: chunks, sums, fits SES with golden-section alpha optimization, and scales back by 1/k. Averages per-k forecasts for the final constant-mean forecast.

Parameter	Type	Default	Description
y	jnp.ndarray	-	Input time series.
h	int	-	Forecast horizon.
fitted	bool	-	Whether to compute in-sample fitted values (O(T^2), expensive).

Returns: Dict[str, jnp.ndarray] (Dict with 'mean' (shape (h,)) and optionally 'fitted' (shape (T,)).)

calculate_information_criteria

chronax.utils.calculate_information_criteria(residuals, n_params, n)

Calculate AIC, BIC, and AICc from residuals (JIT-compiled).

Parameter	Type	Default	Description
residuals	jnp.ndarray	-	Model residuals.
n_params	int	-	Number of estimated parameters.
n	int	-	Number of observations.

Returns: Dict[str, jnp.ndarray] (Dict with 'loglik', 'aic', 'bic', 'aicc' as JAX arrays.)

is_constant

chronax.utils.is_constant(x)

Check if all elements of an array are equal.

Parameter	Type	Default	Description
x	jnp.ndarray	-	Input array.

Returns: jnp.ndarray (Boolean scalar.)

acf

chronax.utils.acf(x, nlags)

Compute autocorrelation function up to nlags for a 1-D array.

Equivalent to statsmodels.tsa.stattools.acf(x, nlags=nlags).

Parameter	Type	Default	Description
x	jnp.ndarray	-	Input 1-D array.
nlags	int	-	Number of lags to compute.

Returns: jnp.ndarray (Array of ACF values from lag 0 to nlags (length nlags+1).)

seasonal_decompose

chronax.utils.seasonal_decompose(y, model='additive', period=1)

Classical seasonal decomposition using centered moving average.

Uses mode='valid' convolution with NaN-padding and half-weights for even periods (proper centered MA), NaN-aware seasonal averaging, and correct normalization.

Parameter	Type	Default	Description
y	jnp.ndarray	-	Input time series array.
model	str	'additive'	Decomposition type, 'additive' or 'multiplicative'.
period	int	1	Seasonal period length.

Returns: Dict[str, jnp.ndarray] (Dict with 'trend', 'seasonal', and 'resid' keys.)