"""
wESN -- a :class:`~memory_esn.double.DoubleReservoirESN` whose second
reservoir sees a *wavelet-transformed then fractionally differenced* view of the input.
Pipeline for reservoir 2::
raw X -> MODWT decomposition -> select component(s) -> fractional differencing
Reservoir 1 still receives the raw series. By default the second reservoir sees
the level-``wavelet_level`` smooth (approximation) coefficients -- a multi-scale
low-frequency trend. Via ``wavelet_components`` you may instead (or additionally)
select individual detail levels; selecting several stacks them into a
**multivariate** input for the second reservoir.
"""
from __future__ import annotations
from typing import List, Optional, Tuple, Union
import numpy as np
from ._speedups import fractional_diff as _fractional_diff
from ._speedups import modwt as _modwt
from .double import DoubleReservoirESN, _resolve_d
Pair = Union[Tuple, list]
# PyWavelets short family names that are *continuous* wavelets (no decomposition
# filters -> unusable by the MODWT). Support for these is planned but not built.
_CONTINUOUS_FAMILIES = {"cgau", "cmor", "fbsp", "gaus", "mexh", "morl", "shan"}
# Aliases accepted for the smooth / approximation component.
_SMOOTH_ALIASES = {"smooth", "approx", "approximation", "a"}
def _family_of(name: str) -> str:
"""Leading alphabetic token of a wavelet name, e.g. 'gaus3' -> 'gaus'."""
i = 0
while i < len(name) and name[i].isalpha():
i += 1
return name[:i].lower()
def _validate_wavelet(name: str) -> None:
"""Raise a helpful error unless ``name`` is a discrete wavelet usable by MODWT.
Continuous wavelets raise :class:`NotImplementedError` (planned for a future
release); unknown names raise :class:`ValueError`.
"""
try:
import pywt
except ImportError as exc: # pragma: no cover
raise ImportError(
"wESN requires PyWavelets (`pip install pywavelets`)."
) from exc
try:
pywt.Wavelet(name) # succeeds only for discrete wavelets with dec filters
return
except (ValueError, Exception): # noqa: BLE001
pass
if _family_of(name) in _CONTINUOUS_FAMILIES or name in set(
pywt.wavelist(kind="continuous")
):
raise NotImplementedError(
f"Continuous wavelet '{name}' is not supported yet -- support for "
"continuous wavelets is planned for a future release. For now use a "
"discrete wavelet: haar, db*, sym*, coif*, bior*, rbio*, dmey "
"(see pywt.wavelist(kind='discrete'))."
)
raise ValueError(
f"Unknown wavelet '{name}'. Choose a discrete wavelet from the families "
"haar, db, sym, coif, bior, rbio, dmey "
"(full list: pywt.wavelist(kind='discrete'))."
)
[docs]
class wESN(DoubleReservoirESN):
"""Wavelet Echo State Network: dual-reservoir ESN with a MODWT +
fractional-differencing memory branch.
Parameters
----------
n_reservoir, spectral_radius, input_scaling, input_init, reservoir_init, bias_init, leaky, activation, bias_scaling, noise, sparsity :
Scalar or length-2 pair (reservoir 1 = raw, reservoir 2 = processed).
d : float or list of float, default=0.25
Fractional differencing order(s) applied to the selected wavelet
component(s). A single value (default) is univariate; a list applies
each order and stacks the results, so the second reservoir's input is the
**cross-product** of components and orders
(``n_features x n_components x n_d`` channels).
K : int, default=100
Number of lags in the fractional-difference filter.
wavelet : str, default='haar'
Discrete wavelet understood by PyWavelets (e.g. 'haar', 'db4', 'sym4').
Continuous wavelets ('mexh', 'morl', ...) are rejected -- see
:func:`_validate_wavelet`.
wavelet_level : int, default=2
MODWT decomposition depth ``J`` (number of detail levels produced).
wavelet_components : default='smooth'
Which MODWT component(s) feed reservoir 2. One of:
* ``'smooth'`` / ``'approx'`` -- the level-``J`` approximation (default,
i.e. the classic smooth trend).
* an int ``k`` in ``1..J`` -- detail coefficients at level ``k``.
* ``'details'`` -- all detail levels ``D_1..D_J``.
* ``'all'`` -- all details plus the smooth.
* a list mixing ints and ``'smooth'``, e.g. ``[1, 2, 'smooth']``.
Selecting more than one component makes the second reservoir's input
**multivariate** (one column per selected component, per input feature).
wavelet_norm : {'modwt', 'classic'}, default='modwt'
MODWT filter normalisation. 'modwt' uses the standard Percival-Walden
rescaling (Haar smooth filter [1/2, 1/2]); 'classic' uses the DWT
orthonormal filters ([1/sqrt2, 1/sqrt2]). The two differ only by a
per-level constant that the random memory-input weights absorb, so they
are model-equivalent; 'classic' reproduces the paper's example equation
literally.
skip_initial : bool, default=True
If True, drop the first ``K`` timesteps of the differenced branch.
random_state, alphas, ridge_cv_params, concatenate_inputs, verbose :
See :class:`~memory_esn.multi.MultiESN`.
Examples
--------
>>> # classic smooth trend (backward-compatible default)
>>> wesn = wESN(n_reservoir=(200, 150), wavelet='db4', wavelet_level=2)
>>> # multivariate: levels 1 & 2 detail plus the smooth
>>> wesn = wESN(wavelet='db4', wavelet_level=3,
... wavelet_components=[1, 2, 'smooth'])
"""
def __init__(
self,
n_reservoir: Union[int, Pair] = 100,
spectral_radius: Union[float, Pair] = 0.9,
input_scaling: Union[float, Pair] = 0.5,
input_init: Union[str, Pair] = "uniform",
reservoir_init: Union[str, Pair] = "uniform",
bias_init: Union[str, Pair] = "uniform",
leaky: Union[float, Pair] = 1.0,
activation: Union[str, Pair] = "tanh",
bias_scaling: Union[float, Pair] = 0.0,
noise: Union[float, Pair] = 0.0,
sparsity: Union[float, Pair] = 0.9,
d: float = 0.25,
K: int = 100,
wavelet: str = "haar",
wavelet_level: int = 2,
wavelet_components: Union[str, int, list, tuple] = "smooth",
wavelet_norm: str = "modwt",
skip_initial: bool = True,
random_state: Union[int, Pair, None] = None,
alphas: Tuple[float, ...] = (0.01, 0.1, 1.0, 10.0),
ridge_cv_params: Optional[dict] = None,
concatenate_inputs: bool = True,
verbose: bool = False,
):
if wavelet_level < 1:
raise ValueError(f"wavelet_level must be >= 1, got {wavelet_level}")
_validate_wavelet(wavelet)
if wavelet_norm not in ("modwt", "classic"):
raise ValueError(
f"wavelet_norm must be 'modwt' or 'classic', got '{wavelet_norm}'"
)
super().__init__(
n_reservoir=n_reservoir,
spectral_radius=spectral_radius,
input_scaling=input_scaling,
input_init=input_init,
reservoir_init=reservoir_init,
bias_init=bias_init,
leaky=leaky,
activation=activation,
bias_scaling=bias_scaling,
noise=noise,
sparsity=sparsity,
random_state=random_state,
alphas=alphas,
ridge_cv_params=ridge_cv_params,
concatenate_inputs=concatenate_inputs,
verbose=verbose,
)
self.d_ = d
self._d_list_ = _resolve_d(d) # normalized list of orders (cross with components)
self.K_ = K
self.wavelet_ = wavelet
self.wavelet_level_ = wavelet_level
self.wavelet_components = wavelet_components
self.wavelet_norm = wavelet_norm
# MODWT filter scale: 1/sqrt(2) = standard (Percival-Walden);
# 1.0 = classic DWT filters ([1/sqrt2, 1/sqrt2] Haar smooth).
self._norm_factor_ = 1.0 / np.sqrt(2.0) if wavelet_norm == "modwt" else 1.0
self.skip_initial_ = skip_initial
# Resolve the requested components to row indices in the (J+1, N) MODWT
# coefficient array: rows 0..J-1 are details D_1..D_J, row J is smooth A_J.
self._component_rows_, self._component_labels_ = self._resolve_components(
wavelet_components, wavelet_level
)
# History for boundary-free smoothing and fracdiff continuation
self.X_train_raw_: Optional[np.ndarray] = None
self._last_K_wav_inputs: Optional[np.ndarray] = None
# ---------------------------------------------------- component selection
@staticmethod
def _resolve_components(spec, J: int) -> Tuple[List[int], List[str]]:
"""Map a ``wavelet_components`` spec to (row indices, human labels)."""
def one(item) -> Tuple[int, str]:
if isinstance(item, str):
if item.lower() in _SMOOTH_ALIASES:
return J, f"A{J}" # smooth / approximation row
raise ValueError(
f"Unknown wavelet component '{item}'. Use an int level 1..{J} "
"or 'smooth'."
)
if isinstance(item, (int, np.integer)):
k = int(item)
if not (1 <= k <= J):
raise ValueError(
f"Detail level {k} is out of range 1..{J} "
f"(wavelet_level={J})."
)
return k - 1, f"D{k}" # detail level k -> row k-1
raise TypeError(
f"wavelet_components entries must be int or str, got {type(item).__name__}"
)
# String shortcuts that expand to several components
if isinstance(spec, str) and spec.lower() == "all":
rows = list(range(J + 1))
return rows, [f"D{k}" for k in range(1, J + 1)] + [f"A{J}"]
if isinstance(spec, str) and spec.lower() == "details":
return list(range(J)), [f"D{k}" for k in range(1, J + 1)]
# Single value or explicit list/tuple
items = spec if isinstance(spec, (list, tuple)) else [spec]
if len(items) == 0:
raise ValueError("wavelet_components must select at least one component.")
rows, labels = [], []
for it in items:
r, lbl = one(it)
rows.append(r)
labels.append(lbl)
return rows, labels
# ---------------------------------------------------------- transforms
@staticmethod
def _as_2d(X: np.ndarray) -> np.ndarray:
# 1-D univariate series -> (N, 1), not (1, N).
X = np.asarray(X)
if X.ndim == 1:
X = X.reshape(-1, 1)
return X
def _wavelet_features(self, X: np.ndarray) -> np.ndarray:
"""Return the selected MODWT component(s) as a (N, F * n_components) matrix.
For each input feature the MODWT is computed once and the requested rows
are extracted; multiple components stack as extra columns (multivariate).
"""
X = self._as_2d(X)
n_samples, n_features = X.shape
columns: List[np.ndarray] = []
for i in range(n_features):
coeffs = _modwt(
np.ascontiguousarray(X[:, i], dtype=np.float64),
self.wavelet_,
self.wavelet_level_,
self._norm_factor_,
) # shape (wavelet_level + 1, n_samples)
for row in self._component_rows_:
columns.append(coeffs[row, :])
return np.column_stack(columns) # (n_samples, n_features * n_components)
def _apply_fractional_diff(self, X: np.ndarray, continuation: bool) -> np.ndarray:
"""Fractionally difference the wavelet features at each order in ``d``.
With several orders the results are stacked, so the second reservoir sees
``n_wavelet_channels x n_d`` columns -- the cross-product of the selected
wavelet components and the differencing orders. History is prepended once
and shared by all orders so their output rows stay aligned.
"""
X = self._as_2d(X)
if continuation and self._last_K_wav_inputs is not None:
X_ext = np.vstack([self._last_K_wav_inputs, X])
else:
X_ext = X
blocks = []
for d in self._d_list_:
# Pure-memory filter ((1-B)^d - 1) applied to the smoothed signal s(t):
# drop the present (zero-lag) term -> f(t) = sum_{k>=1} omega_k s(t-k).
X_diff = _fractional_diff(
X_ext, d=d, K=self.K_, skip_initial=self.skip_initial_,
include_zero_lag=False,
)
if continuation and self._last_K_wav_inputs is not None and not self.skip_initial_:
X_diff = X_diff[self._last_K_wav_inputs.shape[0]:]
blocks.append(X_diff)
return blocks[0] if len(blocks) == 1 else np.hstack(blocks)
def _store_wav_history(self, X_wav: np.ndarray) -> None:
if X_wav.shape[0] >= self.K_:
self._last_K_wav_inputs = X_wav[-self.K_:, :].copy()
else:
pad = np.zeros((self.K_ - X_wav.shape[0], X_wav.shape[1]))
self._last_K_wav_inputs = np.vstack([pad, X_wav])
# ----------------------------------------------------------------- API
[docs]
def fit(
self,
X: np.ndarray,
y: np.ndarray,
washout: int = 0,
x0_list: Optional[list] = None,
) -> "wESN":
"""Fit on a single raw series ``X`` (wavelet -> fracdiff for reservoir 2)."""
X = self._as_2d(X)
y = self._as_2d(y)
self.X_train_raw_ = X.copy()
X_wav = self._wavelet_features(X)
X_frac = self._apply_fractional_diff(X_wav, continuation=False)
self._store_wav_history(X_wav)
if self.skip_initial_:
X_raw_aligned = X[self.K_:, :]
y_aligned = y[self.K_:, :]
else:
X_raw_aligned = X
y_aligned = y
super().fit([X_raw_aligned, X_frac], y_aligned, washout=washout, x0_list=x0_list)
return self
[docs]
def predict(
self,
X: np.ndarray,
continuation: bool = True,
x0_list: Optional[list] = None,
) -> np.ndarray:
"""Predict from a single raw series ``X``.
With ``continuation=True`` (default) the training history is prepended so
the wavelet transform is boundary-free and the output aligns with ``X``.
"""
if not self._is_fitted:
raise RuntimeError("Model must be fitted before prediction. Call fit() first.")
X = self._as_2d(X)
# 1. Wavelet features (use full history in continuation to avoid boundary effects)
if continuation and self.X_train_raw_ is not None:
X_full = np.vstack([self.X_train_raw_, X])
X_wav = self._wavelet_features(X_full)[-len(X):, :]
else:
X_wav = self._wavelet_features(X)
# 2. Fractional differencing (continuation prepends stored history)
X_frac = self._apply_fractional_diff(X_wav, continuation=continuation)
# 3. Align raw branch
if not continuation and self.skip_initial_:
X_raw_aligned = X[self.K_:, :]
else:
X_raw_aligned = X
return super().predict([X_raw_aligned, X_frac], x0_list=x0_list, continuation=continuation)
[docs]
def reset_states(self) -> None:
"""Reset reservoir states and the wavelet/fracdiff history buffers."""
super().reset_states()
self._last_K_wav_inputs = None
def __repr__(self) -> str:
d_repr = self._d_list_[0] if len(self._d_list_) == 1 else self._d_list_
return (
f"wESN(sizes={[r.n_reservoir for r in self.reservoirs_]}, "
f"d={d_repr}, K={self.K_}, wavelet='{self.wavelet_}', "
f"level={self.wavelet_level_}, components={self._component_labels_})"
)
# Backward-compatible alias (the class is now named ``wESN``)
WaveletESN = wESN
