StemGNN
chronax.models.StemGNN · inherits BaseForecaster
StemGNN: Spectral Temporal Graph Neural Network (flax.nnx port of neuralforecast's StemGNN). Cao et al., 2020 — https://arxiv.org/abs/2103.07719. A latent correlation layer (GRU over the series axis + additive attention) learns a graph over series; two residual StockBlockLayer stacks apply a 4-term Chebyshev "GFT" of its normalized Laplacian, a 4-point DFT over the Chebyshev-order axis with GLU filtering (the Spe-Seq cell), a per-order graph-conv kernel, and sigmoid-gated forecast/backcast heads; a final MLP maps to the horizon. Trained with a StepLR schedule (num_lr_decays). This wrapper is univariate (n_series = 1 internally — see the module docstring for the N=1 degeneracy and the chebyshev_first_term escape hatch). No exogenous support. Point or multi-quantile losses; conformal or native quantile intervals. float32 throughout.
Attributes:
* uses_exog: bool = False
* alias: str
* conformal_params: ConformalIntervals | None
* model_: StemGNNNet | None
__init__(self, h, input_size=-1, n_stacks=2, multi_layer=5, dropout_rate=0.5, leaky_rate=0.2, max_steps=1000, learning_rate=1e-3, num_lr_decays=3, windows_batch_size=32, scaler_type='robust', loss='mae', quantile_sort=True, chebyshev_first_term='nf_zero', random_seed=1, alias='StemGNN')
| Parameter | Type | Default | Description |
|---|---|---|---|
h |
- | - | (undocumented) |
input_size |
int |
-1 |
(undocumented) |
n_stacks |
int |
2 |
(undocumented) |
multi_layer |
int |
5 |
(undocumented) |
dropout_rate |
float |
0.5 |
(undocumented) |
leaky_rate |
float |
0.2 |
(undocumented) |
max_steps |
int |
1000 |
(undocumented) |
learning_rate |
float |
1e-3 |
(undocumented) |
num_lr_decays |
int |
3 |
(undocumented) |
windows_batch_size |
int |
32 |
(undocumented) |
scaler_type |
str |
"robust" |
(undocumented) |
loss |
str |
"mae" |
(undocumented) |
quantile_sort |
bool |
True |
(undocumented) |
chebyshev_first_term |
str |
"nf_zero" |
(undocumented) |
random_seed |
int |
1 |
(undocumented) |
alias |
str |
"StemGNN" |
(undocumented) |
fit(self, y, X=None) -> Self
| Parameter | Type | Default | Description |
|---|---|---|---|
y |
- | - | (undocumented) |
X |
- | None |
(undocumented) |
Returns: Self (the fitted forecaster; sets self.model_).
Raises: NotImplementedError if X is not None. ValueError if y is not 1-D or too short.
predict(self, h, X=None, level=None) -> dict
| Parameter | Type | Default | Description |
|---|---|---|---|
h |
- | - | (undocumented) |
X |
- | None |
(undocumented) |
level |
- | None |
(undocumented) |
Returns: dict containing forecast arrays.
* If trained with point loss: {"mean": jnp.ndarray}. If level is provided, includes conformal intervals: {"lo-LV": jnp.ndarray, "hi-LV": jnp.ndarray}.
* If trained with quantile loss: {"mean": jnp.ndarray} (median quantile). If level is provided, includes native quantile intervals: {"lo-LV": jnp.ndarray, "hi-LV": jnp.ndarray}.
forecast(self, y, h, X=None, X_future=None, level=None, fitted=False) -> dict
Stateless fit-then-predict. Exogenous inputs are unsupported.
| Parameter | Type | Default | Description |
|---|---|---|---|
y |
- | - | (undocumented) |
h |
- | - | (undocumented) |
X |
- | None |
(undocumented) |
X_future |
- | None |
(undocumented) |
level |
- | None |
(undocumented) |
fitted |
bool |
False |
(undocumented) |
Returns: dict. Same keys as predict. If fitted=True, also includes "fitted": jnp.ndarray (one-step fitted values).
Raises: NotImplementedError if X_future is not None.