chronax.models.tcn

TCN: Temporal Convolution Network forecaster (flax.nnx port of neuralforecast.TCN).

TCN

chronax.models.tcn.tcn_model · inherits BaseForecaster

Temporal Convolution Network (TCN) Forecaster.

A deep learning model utilizing dilated causal convolutions for time series forecasting. This implementation is a Flax NNX port of the model found in neuralforecast.

Attributes: * uses_exog: bool (False) - Whether the model uses exogenous variables (X). Always False for TCN. * alias: str ("TCN") - The model alias. * model_: Optional[Module] - The fitted TCN model structure.

__init__(self, h, input_size=7, output_size=None, n_blocks=3, n_channels=(16, 16, 16), kernel_size=3, dropout_rate=0.1, activation=jnp.relu, loss='mae', optimizer='adam', learning_rate=0.001, max_steps=1000, early_stopping_patience=50, random_seed=42)

Initializes the TCN forecaster configuration.

Parameter	Type	Default	Description
h	int	-	Forecast horizon (required for initialization, unlike other models).
input_size	int	7	Length of the lookback window (context length).
output_size	Optional[int]	None	The number of output steps (usually equal to h).
n_blocks	int	3	Number of TCN residual blocks.
n_channels	Sequence[int]	(16, 16, 16)	Number of channels/filters in the convolutional layers.
kernel_size	int	3	Size of the convolutional kernel.
dropout_rate	float	0.1	Dropout rate applied after convolutions.
activation	Callable[[jnp.ndarray], jnp.ndarray]	jnp.relu	Activation function used within the blocks.
loss	str	'mae'	Loss function used for training (e.g., 'mae', 'mse').
optimizer	str	'adam'	Optax optimizer configuration (e.g., 'adam', 'sgd').
learning_rate	float	1e-3	Learning rate for the optimizer.
max_steps	int	1000	Maximum number of training steps.
early_stopping_patience	int	50	Patience for early stopping based on validation loss.
random_seed	int	42	Seed for reproducibility.

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

Fits the TCN model to the provided time series data.

The model expects y to be 2D (Time, Features) or 1D (Time,). It handles internal data preparation (windowing).

Parameters:

Parameter	Type	Default	Description
y	jnp.ndarray	-	Target time series data (T, F) or (T,).
X	Optional[jnp.ndarray]	None	Exogenous variables (ignored as uses_exog=False).

Returns: Self (the fitted forecaster instance (self); sets self.model_).

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

Generates forecasts based on the last input_size observations used during fitting.

Note: TCN is a point forecast model. Prediction intervals (level) are not supported and will raise a NotImplementedError if requested.

Parameters:

Parameter	Type	Default	Description
h	int	-	Forecast horizon. Must match self.h.
X	Optional[jnp.ndarray]	None	Exogenous variables (ignored).
level	Optional[Sequence[int]]	None	Prediction intervals (e.g., [80, 95]). Not supported.

Returns: dict (A dictionary containing the mean forecast: {"mean": jnp.ndarray}).

Raises: * NotImplementedError: If level is provided. * RuntimeError: If the model has not been fitted.

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

Fits the model and generates forecasts in a single call.

If fitted is True, it behaves like predict using the provided y to extract the context window.

Parameters:

Parameter	Type	Default	Description
y	jnp.ndarray	-	Target time series data (T, F) or (T,). Used for fitting or context extraction.
h	int	-	Forecast horizon.
X	Optional[jnp.ndarray]	None	Exogenous variables used during fitting (ignored).
X_future	Optional[jnp.ndarray]	None	Future exogenous variables (ignored).
level	Optional[Sequence[int]]	None	Prediction intervals (Not supported).
fitted	bool	False	If True, skips fitting and uses y only for context.

Returns: dict (A dictionary containing the mean forecast: {"mean": jnp.ndarray}).

Raises: * NotImplementedError: If level is provided.