results

chronax.utils.results

Named tuple returned by optimize_theta_target_fn and used by ets_functions.

Fields: x, fn, nit, simplex

ensure_float

chronax.utils.ensure_float

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

ensure_float(y: jnp.ndarray) -> jnp.ndarray

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

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

calculate_sigma

chronax.utils.calculate_sigma

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

calculate_sigma(residuals: jnp.ndarray, n: int) -> jnp.ndarray

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

Returns: Scalar sigma value; returns 0.0 when n <= 0.

_quantiles

chronax.utils._quantiles

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

JAX equivalent of statsforecast.utils._quantiles().

_quantiles(level: List[Union[int, float]]) -> jnp.ndarray

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

Returns: Array of z-scores, one per level.

extract_demand

chronax.utils.extract_demand

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.

extract_demand(y: jnp.ndarray) -> jnp.ndarray

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

Returns: Array containing only positive values from y.

extract_probability

chronax.utils.extract_probability

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.

extract_probability(y: jnp.ndarray) -> jnp.ndarray

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

Returns: Binary array where 1 indicates demand occurred, 0 indicates no demand.

_repeat_val

chronax.utils._repeat_val

Repeat scalar value h times.

JAX equivalent of statsforecast.utils._repeat_val().

_repeat_val(val: float, h: int) -> jnp.ndarray

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

Returns: Array of length h filled with val.

_repeat_val_seas

chronax.utils._repeat_val_seas

Tile seasonal values to cover forecast horizon h.

JAX equivalent of statsforecast.utils._repeat_val_seas().

_repeat_val_seas(season_vals: jnp.ndarray, h: int) -> jnp.ndarray

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: Tiled pattern of length h.

_calculate_intervals

chronax.utils._calculate_intervals

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

_calculate_intervals(res: dict, level: List[int], h: int, sigmah: Union[jnp.ndarray, float]) -> dict

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 with 'lo-{lv}' and 'hi-{lv}' keys for each level.

_add_fitted_pi

chronax.utils._add_fitted_pi

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

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

_add_fitted_pi(res: dict, se: jnp.ndarray, level: Union[List[int], jnp.ndarray]) -> dict

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: Updated res dict with 'fitted-lo-{lv}' and 'fitted-hi-{lv}' keys.

_add_fitted_pi_1

chronax.utils._add_fitted_pi_1

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.

_add_fitted_pi_1(fitted: jnp.ndarray, sigmah: Union[jnp.ndarray, float], level: List[int]) -> dict

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 with 'fitted-lo-{lv}' and 'fitted-hi-{lv}' keys for each level.

_add_conformal_distribution_intervals

chronax.utils._add_conformal_distribution_intervals

Add symmetric conformal intervals using absolute residuals.

Takes the absolute value of signed conformity scores and constructs 2W forecast paths (mean +/- |scores|), producing intervals that are always symmetric around the mean.

_add_conformal_distribution_intervals(fcst: dict, cs: jnp.ndarray, level: Union[List[float], List[int]]) -> dict

Parameter	Type	Default	Description
fcst	dict	-	Forecast dict containing 'mean'.
cs	jnp.ndarray	-	Signed conformal scores of shape (W, h).
level	Union[List[float], List[int]]	-	Confidence levels (0-100).

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

_get_conformal_method

chronax.utils._get_conformal_method

Look up a conformal prediction interval method by name.

_get_conformal_method(method: str) -> Callable

Parameter	Type	Default	Description
method	str	-	Method name ('conformal_distribution' or 'conformal_signed').

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

_conformal_method

chronax.utils._conformal_method

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

_conformal_method(self) -> Callable

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

Returns: The conformal interval function.

_store_cs

chronax.utils._store_cs

Compute and store conformal scores on the model instance.

_store_cs(self, y: jnp.ndarray, X: Optional[jnp.ndarray]) -> None

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

Add conformal prediction intervals to a forecast dict.

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

_add_conformal_intervals(self, fcst: dict, y: Optional[jnp.ndarray], X: Optional[jnp.ndarray], level: Optional[List[int]]) -> dict

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: Updated forecast dict with interval keys.

_add_predict_conformal_intervals

chronax.utils._add_predict_conformal_intervals

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

_add_predict_conformal_intervals(self, fcst: dict, level: Optional[List[int]]) -> dict

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

Returns: Updated forecast dict with interval keys.

_intervals_c

chronax.utils._intervals_c

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

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

_intervals_c(x: jnp.ndarray) -> jnp.ndarray

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

Returns: Fixed-size array with intervals packed at start, rest NaN.

_intervals

chronax.utils._intervals

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.

_intervals(x: jnp.ndarray) -> jnp.ndarray

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

Returns: Float array of inter-arrival intervals.

_expand_fitted_demand

chronax.utils._expand_fitted_demand

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

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

_expand_fitted_demand(fitted: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray

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

Returns: Fitted values expanded to match y's length.

_expand_fitted_intervals

chronax.utils._expand_fitted_intervals

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.

_expand_fitted_intervals(fitted: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray

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

Returns: Fitted intervals expanded to match y's length.

_seasonal_exponential_smoothing

chronax.utils._seasonal_exponential_smoothing

Seasonal exponential smoothing forecast.

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

_seasonal_exponential_smoothing(y: jnp.ndarray, h: int, fitted: bool, season_length: int, alpha: float) -> Dict[str, jnp.ndarray]

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 with 'mean' and optionally 'fitted' keys.

_seasonal_naive

chronax.utils._seasonal_naive

JAX implementation of seasonal-naive forecast.

Repeats the last season_length observations as the forecast.

_seasonal_naive(y: jnp.ndarray, h: int, season_length: int, fitted: bool = False) -> Dict[str, jnp.ndarray]

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 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.

_window_average

chronax.utils._window_average

Window average forecast.

_window_average(y: jnp.ndarray, h: int, fitted: bool, window_size: int) -> Dict[str, jnp.ndarray]

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 with 'mean' key containing constant forecast of length h. Raises: NotImplementedError: If fitted=True.

_imapa

chronax.utils._imapa

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. vmap-compatible: all control flow uses JAX primitives (jnp.where, lax.cond).

_imapa(y: jnp.ndarray, h: int, fitted: bool) -> Dict[str, jnp.ndarray]

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 with 'mean' (shape (h,)) and optionally 'fitted' (shape (T,)).

is_constant

chronax.utils.is_constant

Check if all elements of an array are equal.

is_constant(x: jnp.ndarray) -> jnp.ndarray

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

Returns: Boolean scalar.

acf

chronax.utils.acf

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

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

acf(x: jnp.ndarray, nlags: int) -> jnp.ndarray

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

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

calculate_information_criteria

chronax.utils.calculate_information_criteria

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

calculate_information_criteria(residuals: jnp.ndarray, n_params: int, n: int) -> Dict[str, jnp.ndarray]

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

Returns: Dict with 'loglik', 'aic', 'bic', 'aicc' as JAX arrays.

seasonal_decompose

chronax.utils.seasonal_decompose

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.

seasonal_decompose(y: jnp.ndarray, model: str = 'additive', period: int = 1) -> Dict[str, jnp.ndarray]

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 with 'trend', 'seasonal', and 'resid' keys.