"""
Weight-initialisation distributions for :mod:`memory_esn`.
Only the *distribution* the weights are drawn from is configurable here -- the
reservoir is still rescaled to the target spectral radius and the optional
sparsity mask is still applied, exactly as before. This keeps the knob simple:
pick a distribution for the input weights, the reservoir weights, and the bias.
Supported distributions (ESN literature draws input weights from a uniform range
and reservoir weights from a normal law; bimodal/Bernoulli weights are also
common):
* ``'gaussian'`` (aliases ``'normal'``, ``'randn'``) -- N(0, scale^2) [default]
* ``'uniform'`` -- U(-scale, scale)
* ``'bernoulli'`` (aliases ``'binary'``, ``'bimodal'``, ``'sign'``) -- {-scale, +scale}
* ``'laplace'`` -- Laplace(0, scale)
``scale`` is ``input_scaling`` for the input weights and ``bias_scaling`` for the
bias; for the reservoir it is 1.0 (absolute scale is set afterwards by the
spectral-radius rescaling, so only the *shape* of the distribution matters there).
"""
from __future__ import annotations
from typing import Optional
import numpy as np
# ---------------------------------------------------------------------------
# Sparse-reservoir switch-over thresholds.
#
# For small/mid or dense reservoirs the dense BLAS ``dgemv`` update is fastest,
# so ``W`` is stored as a dense NumPy array (unchanged behaviour). For large,
# highly-sparse reservoirs a dense matvec wastes work on the ~90% zeros, so we
# store ``W`` as a scipy CSR matrix and use a sparse matvec instead. The
# crossover (see ``benchmarks/RESULTS.md``) is around 1000 units at >=80%
# sparsity. Both conditions must hold for the sparse path to switch on; the
# defaults (n_reservoir=100) stay on the dense path, bit-for-bit unchanged.
# ---------------------------------------------------------------------------
SPARSE_RESERVOIR_MIN_UNITS = 1000
SPARSE_RESERVOIR_MIN_SPARSITY = 0.8
# Canonical names -> the set of accepted aliases.
_ALIASES = {
"gaussian": {"gaussian", "normal", "randn"},
"uniform": {"uniform"},
"bernoulli": {"bernoulli", "binary", "bimodal", "sign"},
"laplace": {"laplace"},
}
_ALL_NAMES = {alias for names in _ALIASES.values() for alias in names}
DISTRIBUTIONS = tuple(_ALIASES.keys()) # public, canonical list
def _canonical(distribution: str) -> str:
d = distribution.lower()
for canon, aliases in _ALIASES.items():
if d in aliases:
return canon
raise ValueError(
f"Unknown weight distribution '{distribution}'. "
f"Choose from {sorted(_ALL_NAMES)}."
)
[docs]
def validate_distribution(distribution: str, arg_name: str = "init") -> None:
"""Raise ``ValueError`` unless ``distribution`` is a known name/alias."""
try:
_canonical(distribution)
except ValueError as exc:
raise ValueError(f"{arg_name}: {exc}") from None
def _draw(rng: np.random.RandomState, shape, distribution: str, scale: float) -> np.ndarray:
"""Draw an array of ``shape`` from ``distribution`` at the given ``scale``."""
canon = _canonical(distribution)
if canon == "gaussian":
return rng.randn(*shape) * scale
if canon == "uniform":
return rng.uniform(-scale, scale, size=shape)
if canon == "bernoulli":
return np.where(rng.randint(0, 2, size=shape) == 0, -scale, scale).astype(np.float64)
if canon == "laplace":
return rng.laplace(0.0, scale, size=shape)
raise ValueError(f"Unhandled distribution '{distribution}'") # pragma: no cover
[docs]
def use_sparse_reservoir(n_reservoir: int, sparsity: float) -> bool:
"""Whether to store ``W`` as a scipy sparse matrix for this configuration.
True only for large, highly-sparse reservoirs (see the module-level
thresholds), where a sparse matvec beats a dense BLAS ``dgemv``. Small,
mid-sized, or dense reservoirs keep the dense NumPy path.
"""
return (
sparsity >= SPARSE_RESERVOIR_MIN_SPARSITY
and n_reservoir >= SPARSE_RESERVOIR_MIN_UNITS
)
def _sparse_spectral_radius(W_sp) -> float:
"""Largest-magnitude eigenvalue of a sparse matrix via ARPACK.
Uses :func:`scipy.sparse.linalg.eigs` instead of a dense
``np.linalg.eigvals`` -- the latter would densify the matrix and cost
O(n^3), which is exactly what the sparse path exists to avoid. Falls back to
the dense computation on the rare occasion ARPACK fails to converge.
ARPACK's Arnoldi iteration is seeded with a *random* start vector and, for a
non-symmetric reservoir, ``k=1`` can return a not-quite-converged Ritz value
whose magnitude wobbles from run to run (and from one SciPy build to the
next). That makes the rescaling non-deterministic: the radius used to scale
``W`` disagrees with the radius measured afterwards. To pin it down we use a
fixed start vector, iterate to machine precision (``tol=0``), and request a
few eigenvalues so a dominant *complex-conjugate pair* is captured reliably.
"""
from scipy.sparse.linalg import ArpackError, ArpackNoConvergence, eigs
n = W_sp.shape[0]
k = min(6, n - 2) # eigs requires k < n - 1
eigs_kwargs = dict(
k=k,
which="LM",
return_eigenvectors=False,
v0=np.ones(n), # deterministic start vector
maxiter=n * 100,
tol=0, # converge to machine precision
)
try:
vals = eigs(W_sp, **eigs_kwargs)
return float(np.max(np.abs(vals)))
except ArpackNoConvergence as exc: # partial result may still be usable
if exc.eigenvalues.size:
return float(np.max(np.abs(exc.eigenvalues)))
return float(np.max(np.abs(np.linalg.eigvals(W_sp.toarray()))))
except ArpackError: # pragma: no cover - defensive
return float(np.max(np.abs(np.linalg.eigvals(W_sp.toarray()))))
[docs]
def init_reservoir_weights(
rng: np.random.RandomState,
n_reservoir: int,
spectral_radius: float = 0.9,
distribution: str = "gaussian",
sparsity: float = 0.0,
sparse: Optional[bool] = None,
):
"""Reservoir matrix ``W`` drawn from ``distribution``, sparsified, then
rescaled so its spectral radius equals ``spectral_radius``.
Returns a dense ``np.ndarray`` (default) or, for large highly-sparse
reservoirs, a scipy CSR matrix. ``sparse`` forces the storage choice; when
``None`` it is decided automatically by :func:`use_sparse_reservoir`. The
random draw and sparsity mask are identical in both cases, so the two paths
hold the same weights -- only the storage format (and the matvec used at
runtime) differs.
"""
W = _draw(rng, (n_reservoir, n_reservoir), distribution, 1.0)
if sparsity > 0:
mask = rng.rand(n_reservoir, n_reservoir) > sparsity
W = W * mask
if sparse is None:
sparse = use_sparse_reservoir(n_reservoir, sparsity)
if sparse:
from scipy.sparse import csr_matrix
W_sp = csr_matrix(W)
current_radius = _sparse_spectral_radius(W_sp)
if current_radius > 0:
W_sp = W_sp * (spectral_radius / current_radius)
return W_sp.tocsr()
eigenvalues = np.linalg.eigvals(W)
current_radius = np.max(np.abs(eigenvalues))
if current_radius > 0:
W = W * (spectral_radius / current_radius)
return W
[docs]
def init_bias(
rng: np.random.RandomState,
n_reservoir: int,
scaling: float = 0.0,
distribution: str = "gaussian",
) -> Optional[np.ndarray]:
"""Bias vector of shape (n_reservoir,), or ``None`` when ``scaling <= 0``."""
if scaling <= 0:
return None
return _draw(rng, (n_reservoir,), distribution, scaling)
