chronax.models.softs_model

SOFTS forecaster: BaseForecaster wrapper around the JAX/Flax-NNX backbone.

SOFTS

chronax.models.softs_model.SOFTS · inherits BaseForecaster

Univariate SOFTS forecaster (JAX/Flax-NNX port of neuralforecast.SOFTS). SOFTS (Han et al. 2024, "SOFTS: Efficient Multivariate Time Series Forecasting with Series-Core Fusion") shares iTransformer's inverted encoder — each variate's lookback window is embedded into a token (Linear(input_size -> hidden_size)) — but replaces self-attention with the STAD (STar Aggregate-Dispatch) module: the series are fused into a single global core (d_core) and the core is dispatched back onto each series. This makes the encoder O(C) in the number of series instead of attention's O(C^2). RevIN-style per-window mean/std normalization (use_norm) is applied inside the network and inverted on the output. Trained in original scale with Optax adam and a pluggable point loss. float32 throughout. The series is univariate (n_series = 1), so the encoder operates on a single token; the backbone is written N-generically to allow a future multivariate path.

Maintenance status: Active univariate forecaster. Integrates with the BaseForecaster interface, including conformal prediction intervals via predict(level=...), pickle round-trip, and forecast(fitted=True).

Attributes

Name	Type	Description
uses_exog	bool	False

__init__(self, h, input_size=3, hidden_size=512, d_core=512, e_layers=2, d_ff=2048, dropout=0.1, use_norm=True, use_boxcox=False, activation='gelu', max_steps=1000, learning_rate=0.001, windows_batch_size=32, random_seed=1, alias='SOFTS', loss='mae')

Initialize a SOFTS forecaster.

Stores hyperparameters; the network is built lazily at fit time so construction is cheap and side-effect free. Defaults match neuralforecast.SOFTS (hidden_size=512, d_core=512, e_layers=2, d_ff=2048, dropout=0.1, use_norm=True, max_steps=1000, learning_rate=1e-3, windows_batch_size=32). input_size=-1 resolves to 3 * h. loss is a registry name ("mae"/"mse"/"huber") or a callable; learning_rate is a scalar or an optax.ScalarOrSchedule; activation is "gelu" or "relu". If the fitted estimator will be pickled, any callable passed for loss/learning_rate must itself be picklable (a class-based callable or module-level function). use_boxcox (default False) applies a variance-stabilizing Box-Cox transform to the series before modelling and inverts it on the forecast, mirroring the use_boxcox option of :class:chronax.models.TBATS. The lambda is selected once at fit time by maximizing the Box-Cox profile log-likelihood. It helps multiplicative / strongly-trending series (e.g. airline passengers) and requires strictly positive values. Left off, the model is a faithful port of neuralforecast.SOFTS.

Parameter	Type	Default	Description
h	int	-	(undocumented)
input_size	int	-1	(undocumented)
hidden_size	int	512	(undocumented)
d_core	int	512	(undocumented)
e_layers	int	2	(undocumented)
d_ff	int	2048	(undocumented)
dropout	float	0.1	(undocumented)
use_norm	bool	True	(undocumented)
use_boxcox	bool	False	Applies a variance-stabilizing Box-Cox transform to the series before modelling and inverts it on the forecast. Requires strictly positive values.
activation	str	"gelu"	(undocumented)
max_steps	int	1000	(undocumented)
learning_rate	Union[float, Callable[[int], float]]	1e-3	(undocumented)
windows_batch_size	int	32	(undocumented)
random_seed	int	1	(undocumented)
alias	str	"SOFTS"	(undocumented)
loss	Union[str, LossFn]	"mae"	(undocumented)

fit(self, y, X=None) -> Self

Fit the network on a 1-D series.

Builds the network and runs max_steps Adam steps over rolling windows of length input_size + h, sampled per step the way neuralforecast does.

Parameters:

Parameter	Type	Default	Description
y	jnp.ndarray	-	1-D series of length >= input_size + h.
X	jnp.ndarray \| None	None	Reserved for exogenous regressors; must be None.

Returns: Self (the fitted forecaster; sets self.model_). Raises: * NotImplementedError: If X is provided. * ValueError: If y is not 1-D or is shorter than input_size + h. * RuntimeError: If a non-finite training loss is observed (divergence).

predict(self, h, X=None, level=None) -> dict

Forecast h steps from the fitted context.

Parameters:

Parameter	Type	Default	Description
h	int	-	Forecast horizon; must satisfy 1 <= h <= self.h (the model is direct-decoded for self.h steps and sliced).
X	jnp.ndarray \| None	None	Reserved for exogenous regressors; ignored.
level	list[int \| float] \| None	None	Optional confidence levels (e.g. [80, 95]). When set, returns conformal lo-XX/hi-XX keys via the inherited BaseForecaster path and requires self.conformal_params. Each call re-fits the model per CV window under vmap — expect minutes.

Returns: dict ({"mean": jnp.ndarray of shape (h,)} plus interval keys when level is provided.) Raises: * RuntimeError: If called before fit. * ValueError: If h < 1 or h > self.h, or if level is given without self.conformal_params set.

forecast(self, y, h, X=None, X_future=None, level=None, fitted=False) -> dict

Stateless fit-then-predict on y.

Equivalent to self.fit(y).predict(h=h, level=level), optionally adding a "fitted" key with one-step-ahead in-sample predictions.

Parameters:

Parameter	Type	Default	Description
y	jnp.ndarray	-	1-D training series.
h	int	-	Forecast horizon (<= self.h).
X	jnp.ndarray \| None	None	Reserved for exogenous regressors; must be None.
X_future	jnp.ndarray \| None	None	Reserved for exogenous regressors; must be None.
level	list[int \| float] \| None	None	Optional confidence levels; see predict.
fitted	bool	False	If True, include "fitted" — one-step-ahead values over the training series, NaN for the first input_size entries.

Returns: dict ({"mean": ..., optional "fitted": ...}.) Raises: * NotImplementedError: If X or X_future is provided. * ValueError / RuntimeError: Forwarded from fit / predict.