data.py

Data pipeline for Chronax N-BEATS.

Three concerns are separated:

1. RobustScaler — per-window median/MAD normalisation (pure JAX).
2. Window building — build_windows / split_train_val_windows extract sliding windows from a 1-D series.
3. Batch construction — create_batch assembles JAX batches for the masked-loss training path (mirrors RNN data.py).

The window-based path (build_windows + RobustScaler) is the primary path used by NBEATSForecaster; it applies per-window scaling inside the train_step JIT kernel, so no global statistics need to be tracked.

No NumPy — every array is a jnp.ndarray, including the sliding-window extraction (done via a single vectorised gather).

RobustScaler

data.RobustScaler

Median + MAD scaler with std-based fallback when MAD == 0.

Stateless: stats computes (shift, scale) so the same statistics can be reused to inverse-transform predictions.

stats(self, x: jnp.ndarray, axis: int = 1) -> Tuple[jnp.ndarray, jnp.ndarray]

Computes the median (shift) and MAD-based scale for normalization.

Parameter	Type	Default	Description
x	jnp.ndarray	-	(undocumented)
axis	int	1	(undocumented)

Returns: Tuple[jnp.ndarray, jnp.ndarray] (shift, scale).

transform(self, x: jnp.ndarray, shift: jnp.ndarray, scale: jnp.ndarray) -> jnp.ndarray

Applies the robust scaling transformation: (x - shift) / scale.

Parameter	Type	Default	Description
x	jnp.ndarray	-	(undocumented)
shift	jnp.ndarray	-	(undocumented)
scale	jnp.ndarray	-	(undocumented)

Returns: jnp.ndarray

inverse(self, z: jnp.ndarray, shift: jnp.ndarray, scale: jnp.ndarray) -> jnp.ndarray

Applies the inverse robust scaling transformation: z * scale + shift.

Parameter	Type	Default	Description
z	jnp.ndarray	-	(undocumented)
shift	jnp.ndarray	-	(undocumented)
scale	jnp.ndarray	-	(undocumented)

Returns: jnp.ndarray

build_windows

data.build_windows

Return [n_windows, input_size + h] rolling windows (stride = 1).

Parameter	Type	Default	Description
y	jnp.ndarray	-	1-D float32 series of length T.
input_size	int	-	History window length L.
h	int	-	Forecast horizon.

Returns: jnp.ndarray

Raises: * ValueError: When T < input_size + h.

split_train_val_windows

data.split_train_val_windows

Chronological split: earliest windows train, latest windows validate.

When val_fraction <= 0 all windows go to the training set.

Parameter	Type	Default	Description
windows	jnp.ndarray	-	(undocumented)
val_fraction	float	0.1	(undocumented)

Returns: Tuple[jnp.ndarray, jnp.ndarray]

pad_sequence

data.pad_sequence

Left-pad / left-truncate a 1-D series to target_len.

Returns (padded, mask) where mask is 1 on real observations.

Parameter	Type	Default	Description
y	jnp.ndarray	-	(undocumented)
target_len	int	-	(undocumented)
pad_value	float	0.0	(undocumented)

Returns: Tuple[jnp.ndarray, jnp.ndarray]

create_batch

data.create_batch

Build a JAX batch from a list of complete time series.

Each series is split into (history, horizon). History is padded/truncated to input_size (yielding available_mask) and horizon is padded/truncated to h (yielding sample_mask).

Parameter	Type	Default	Description
y_series	List[jnp.ndarray]	-	(undocumented)
input_size	int	-	(undocumented)
h	int	-	(undocumented)
sample_mask_list	Optional[List[Optional[jnp.ndarray]]]	None	(undocumented)

Returns: Dict[str, Optional[jnp.ndarray]] Returns a dict with JAX arrays: * insample_y: [B, input_size] (2-D, no trailing feature dim) * outsample_y: [B, h] * available_mask: [B, input_size] (1 = real, 0 = padded) * sample_mask: [B, h] (1 = include in loss, 0 = ignore)

batch_generator

data.batch_generator

Iterate over y_series yielding fully-prepared JAX batches.

Parameter	Type	Default	Description
y_series	List[jnp.ndarray]	-	(undocumented)
input_size	int	-	(undocumented)
h	int	-	(undocumented)
batch_size	int	-	(undocumented)
sample_mask_list	Optional[List[Optional[jnp.ndarray]]]	None	(undocumented)
shuffle	bool	True	(undocumented)
seed	int	42	(undocumented)

Returns: Iterator[Dict[str, Optional[jnp.ndarray]]]