"""
BaseESN -- a single-reservoir Echo State Network.
The reservoir (fixed, random recurrent layer) transforms the input sequence into
a high-dimensional state trajectory; only a linear ridge readout is trained.
State updates use the BLAS-backed Cython kernel when available.
"""
from __future__ import annotations
from typing import List, Optional, Tuple, Union
import numpy as np
import scipy.sparse as _sp
from sklearn.linear_model import RidgeCV
from ._speedups import reservoir_update as _reservoir_update
from ._speedups import reservoir_update_sparse as _reservoir_update_sparse
from .weights import (
init_bias,
init_input_weights,
init_reservoir_weights,
validate_distribution,
)
class PersistenceMixin:
"""Adds ``save`` / ``load`` (joblib) to any model in the hierarchy."""
def save(self, filepath: str) -> None:
"""Serialize the fitted model to ``filepath`` with joblib."""
import joblib
joblib.dump(self, filepath)
@classmethod
def load(cls, filepath: str):
"""Load a model previously written with :meth:`save`."""
import joblib
return joblib.load(filepath)
[docs]
class BaseESN(PersistenceMixin):
"""Single-reservoir Echo State Network.
Parameters
----------
n_reservoir : int, default=100
Number of reservoir neurons.
spectral_radius : float, default=0.9
Target spectral radius of the reservoir weight matrix (controls memory
/ stability; keep < 1 for the echo-state property).
input_scaling : float, default=0.5
Scale of the random input weights.
input_init : str, default='uniform'
Distribution for the input weights: 'uniform', 'gaussian'/'normal',
'bernoulli'/'binary', or 'laplace'. See :mod:`memory_esn.weights`.
(Uniform is the paper's initialisation.)
reservoir_init : str, default='uniform'
Distribution for the reservoir weights (same options). The matrix is
still rescaled to ``spectral_radius`` afterwards.
bias_init : str, default='uniform'
Distribution for the neuron bias (same options; only used when
``bias_scaling > 0``).
leaky : float, default=1.0
Leaky-integration rate (1.0 = no leak).
activation : str, default='tanh'
Reservoir nonlinearity. One of: 'tanh', 'sigmoid', 'relu', 'leaky_relu',
'elu', 'selu', 'softplus', 'swish'/'silu', 'gelu', 'hard_tanh',
'softsign', 'mish', 'hard_sigmoid', 'relu6', 'tanhshrink', 'sin',
'identity'/'linear'.
bias_scaling : float, default=0.0
Scale of the random neuron bias (0.0 = no bias).
noise : float, default=0.0
Std-dev sigma of isotropic additive state noise eta(t) ~ N(0, sigma^2 I)
injected into the pre-activation each timestep (0.0 = no noise). Drawn
from the reservoir RNG, so results are reproducible for a fixed seed.
sparsity : float, default=0.9
Fraction of reservoir connections forced to zero (0.0 = dense, 0.9 = 90%
zeros). Matches the paper's sparsification proportion phi.
random_state : int or None, default=None
Seed for reproducibility.
alphas : tuple of float, default=(0.01, 0.1, 1.0, 10.0)
Candidate ridge penalties for :class:`~sklearn.linear_model.RidgeCV`.
ridge_cv_params : dict, optional
Extra keyword arguments forwarded to ``RidgeCV``.
concatenate_input : bool, default=True
If True, append the raw input to the reservoir states before the readout.
Examples
--------
>>> esn = BaseESN(n_reservoir=200, spectral_radius=0.95, random_state=0)
>>> esn.fit(X_train, y_train, washout=100)
>>> y_pred = esn.predict(X_test)
"""
def __init__(
self,
n_reservoir: int = 100,
spectral_radius: float = 0.9,
input_scaling: float = 0.5,
input_init: str = "uniform",
reservoir_init: str = "uniform",
bias_init: str = "uniform",
leaky: float = 1.0,
activation: str = "tanh",
bias_scaling: float = 0.0,
noise: float = 0.0,
sparsity: float = 0.9,
random_state: Optional[int] = None,
alphas: Tuple[float, ...] = (0.01, 0.1, 1.0, 10.0),
ridge_cv_params: Optional[dict] = None,
concatenate_input: bool = True,
):
validate_distribution(input_init, "input_init")
validate_distribution(reservoir_init, "reservoir_init")
validate_distribution(bias_init, "bias_init")
if noise < 0:
raise ValueError(f"noise (std-dev) must be >= 0, got {noise}")
self.n_reservoir = n_reservoir
self.spectral_radius = spectral_radius
self.input_scaling = input_scaling
self.input_init = input_init
self.reservoir_init = reservoir_init
self.bias_init = bias_init
self.leaky = leaky
self.activation = activation
self.bias_scaling = bias_scaling
self.noise = noise
self.sparsity = sparsity
self.random_state = random_state
self.alphas = alphas
self.ridge_cv_params = ridge_cv_params or {}
self.concatenate_input = concatenate_input
# Learned / initialised during fit
self.W_in_ = None
self.W_reservoir_ = None
self.bias_ = None
self.readout_ = None
self.n_inputs_ = None
self.n_outputs_ = None
self._is_fitted = False
# Last reservoir state, for prediction continuation
self.last_state_ = None
self.rng_ = np.random.RandomState(random_state)
# ------------------------------------------------------------------ setup
def _initialize_weights(self, n_inputs: int) -> None:
"""Draw input/reservoir/bias weights and rescale W to the target radius.
Draw order (input -> reservoir -> bias) is preserved, so the default
'gaussian' settings reproduce earlier results bit-for-bit given a seed.
"""
self.W_in_ = init_input_weights(
self.rng_, self.n_reservoir, n_inputs,
scaling=self.input_scaling, distribution=self.input_init,
)
self.W_reservoir_ = init_reservoir_weights(
self.rng_, self.n_reservoir,
spectral_radius=self.spectral_radius, distribution=self.reservoir_init,
sparsity=self.sparsity,
)
self.bias_ = init_bias(
self.rng_, self.n_reservoir,
scaling=self.bias_scaling, distribution=self.bias_init,
)
self.n_inputs_ = n_inputs
def _compute_states(
self,
X: np.ndarray,
x0: Optional[np.ndarray] = None,
continuation: bool = False,
) -> np.ndarray:
"""Run the reservoir over ``X`` and return the state trajectory.
The final state is cached in ``last_state_`` for continuation.
"""
if continuation and self.last_state_ is not None:
initial_state = self.last_state_.copy()
elif x0 is not None:
initial_state = x0
else:
initial_state = None
# Isotropic additive state noise eta(t) ~ N(0, noise^2 I), drawn from the
# reservoir RNG (reproducible). None when noise == 0 (default), so the
# RNG stream is untouched and default runs stay bit-identical.
if self.noise > 0:
noise_sequence = self.rng_.randn(X.shape[0], self.n_reservoir) * self.noise
else:
noise_sequence = None
# Large, highly-sparse reservoirs store W as a scipy CSR matrix; use the
# sparse matvec path for them and the dense BLAS dgemv kernel otherwise.
update = (
_reservoir_update_sparse
if _sp.issparse(self.W_reservoir_)
else _reservoir_update
)
states = update(
input_sequence=np.ascontiguousarray(X, dtype=np.float64),
W_in=self.W_in_,
W_reservoir=self.W_reservoir_,
activation=self.activation,
leaky=self.leaky,
bias=self.bias_,
x0=initial_state,
noise_sequence=noise_sequence,
)
self.last_state_ = states[-1].copy()
return states
@staticmethod
def _as_2d(a: np.ndarray) -> np.ndarray:
# A 1-D series of length N is a univariate time series -> (N, 1).
# (np.atleast_2d would wrongly make it (1, N): 1 timestep, N features.)
a = np.asarray(a)
if a.ndim == 1:
a = a.reshape(-1, 1)
return a
# ------------------------------------------------------------------- API
[docs]
def fit(
self,
X: np.ndarray,
y: np.ndarray,
washout: int = 0,
x0: Optional[np.ndarray] = None,
) -> "BaseESN":
"""Fit the reservoir readout.
Parameters
----------
X : ndarray, shape (n_timesteps, n_inputs)
y : ndarray, shape (n_timesteps, n_outputs)
washout : int, default=0
Number of initial timesteps to discard before training the readout.
x0 : ndarray, optional
Initial reservoir state.
"""
X = self._as_2d(X)
y = self._as_2d(y)
n_timesteps, n_inputs = X.shape
self.n_outputs_ = y.shape[1]
if self.W_in_ is None:
self._initialize_weights(n_inputs)
states = self._compute_states(X, x0=x0)
extended = np.hstack([states, X]) if self.concatenate_input else states
if washout > 0:
extended = extended[washout:]
y = y[washout:]
self.readout_ = RidgeCV(alphas=self.alphas, **self.ridge_cv_params)
self.readout_.fit(extended, y)
self._is_fitted = True
return self
[docs]
def predict(
self,
X: np.ndarray,
x0: Optional[np.ndarray] = None,
continuation: bool = False,
) -> np.ndarray:
"""Predict outputs for ``X`` (optionally continuing from the last state)."""
if not self._is_fitted:
raise RuntimeError("Model must be fitted before prediction. Call fit() first.")
X = self._as_2d(X)
states = self._compute_states(X, x0=x0, continuation=continuation)
extended = np.hstack([states, X]) if self.concatenate_input else states
pred = self.readout_.predict(extended)
if pred.ndim == 1:
pred = pred.reshape(-1, self.n_outputs_)
return pred
[docs]
def get_states(
self,
X: np.ndarray,
x0: Optional[np.ndarray] = None,
continuation: bool = False,
) -> np.ndarray:
"""Return reservoir states for ``X`` (weights must already be initialised)."""
if self.W_in_ is None:
raise RuntimeError("Weights not initialized. Call fit() first.")
return self._compute_states(self._as_2d(X), x0=x0, continuation=continuation)
[docs]
def reset_state(self) -> None:
"""Forget the cached reservoir state used for continuation."""
self.last_state_ = None
def __repr__(self) -> str:
return (
f"BaseESN(n_reservoir={self.n_reservoir}, "
f"spectral_radius={self.spectral_radius}, activation='{self.activation}')"
)
