API reference¶

def ridge_statistics(Z: np.ndarray, Y: np.ndarray):
    """Return sufficient statistics ``A = Z^T Z`` and ``B = Z^T Y``.

    The Gram matrix is accumulated in float64 even when the states are float32:
    the harvest keeps the float32 speed/memory benefit, while the (small,
    ill-conditioned) ridge solve stays numerically stable.
    """
    Z = np.asarray(Z, dtype=np.float64)
    Y = np.asarray(Y, dtype=np.float64)
    return Z.T @ Z, Z.T @ Y

def solve_readout(A: np.ndarray, B: np.ndarray, ridge: float) -> np.ndarray:
    """Solve ``(A + ridge * I) W_out = B`` for the readout weights.

    ``A`` is a Gram matrix, hence symmetric positive semi-definite, so the system
    is solved via a Cholesky factorisation of the regularised matrix. With
    ``ridge = 0`` (or a rank-deficient ``A``) that system can be singular; rather
    than propagating a bare ``LinAlgError`` we fall back to the minimum-norm
    least-squares solution, which is the natural limit of ridge regression as the
    regularisation vanishes.
    """
    A = np.asarray(A, dtype=np.float64)
    B = np.asarray(B, dtype=np.float64)
    if A.ndim != 2 or A.shape[0] != A.shape[1]:
        raise ValueError(f"A must be a square matrix, got shape {A.shape}")
    if B.shape[0] != A.shape[0]:
        raise ValueError(
            f"A and B disagree: A is {A.shape}, B is {B.shape}"
        )
    if ridge < 0:
        raise ValueError(f"ridge must be >= 0, got {ridge}")
    G = A + ridge * np.eye(A.shape[0])
    try:
        # Symmetric-positive-definite solve: ~2x faster than a general LU and it
        # fails cleanly (rather than silently amplifying error) when G is singular.
        chol = np.linalg.cholesky(G)
        return np.linalg.solve(chol.T, np.linalg.solve(chol, B))
    except np.linalg.LinAlgError:
        pass
    try:
        return np.linalg.solve(G, B)
    except np.linalg.LinAlgError:
        warnings.warn(
            "the regularised Gram matrix is singular; falling back to the "
            "minimum-norm least-squares readout. Increase `ridge` for a "
            "well-posed solve.",
            stacklevel=2,
        )
        return np.linalg.lstsq(G, B, rcond=None)[0]

def random_reservoir(n: int, density: float = 0.1, rng=None) -> np.ndarray:
    """Erdos-Renyi reservoir: each directed edge present with prob. ``density``."""
    n = _check_n(n)
    if not 0.0 <= density <= 1.0:
        raise ValueError(f"density must be in [0, 1], got {density}")
    rng = _as_rng(rng)
    mask = rng.uniform(size=(n, n)) < density
    np.fill_diagonal(mask, False)
    W = np.zeros((n, n))
    W[mask] = rng.uniform(-1.0, 1.0, size=int(mask.sum()))
    return W

def small_world_reservoir(
    n: int, k: int = 6, p: float = 0.1, rng=None
) -> np.ndarray:
    """Watts-Strogatz small-world reservoir.

    A ring lattice where each node connects to its ``k`` nearest neighbours, with
    each edge rewired with probability ``p``. Small-world reservoirs combine high
    clustering with short path lengths.
    """
    n = _check_n(n)
    if not 0.0 <= p <= 1.0:
        raise ValueError(f"rewiring probability p must be in [0, 1], got {p}")
    if not 1 <= k < n:
        raise ValueError(
            f"the ring lattice needs 1 <= k < n neighbours, got k={k} for n={n}"
        )
    rng = _as_rng(rng)
    seed = int(rng.integers(0, 2**31 - 1))
    g = nx.watts_strogatz_graph(n, k, p, seed=seed)
    adj = nx.to_numpy_array(g)
    return _weight_edges(adj, rng)

def scale_free_reservoir(n: int, m: int = 3, rng=None) -> np.ndarray:
    """Barabasi-Albert scale-free reservoir.

    Growth with preferential attachment; each new node attaches to ``m`` existing
    nodes. Produces a power-law degree distribution with a few high-degree hubs.
    """
    n = _check_n(n)
    if not 1 <= m < n:
        raise ValueError(
            f"preferential attachment needs 1 <= m < n, got m={m} for n={n}"
        )
    rng = _as_rng(rng)
    seed = int(rng.integers(0, 2**31 - 1))
    g = nx.barabasi_albert_graph(n, m, seed=seed)
    adj = nx.to_numpy_array(g)
    return _weight_edges(adj, rng)

def ring_reservoir(n: int, weight: float = 1.0, rng=None) -> np.ndarray:
    """Deterministic uni-directional ring (a.k.a. simple cycle reservoir).

    Each node feeds the next; node ``n-1`` feeds node ``0``. Despite its
    simplicity this minimal-complexity reservoir is competitive on many tasks
    (Rodan & Tino, 2011).
    """
    n = _check_n(n)
    rng = _as_rng(rng)
    W = np.zeros((n, n))
    signs = rng.choice([-1.0, 1.0], size=n)
    for i in range(n):
        W[(i + 1) % n, i] = weight * signs[i]
    return W

def make_reservoir(kind: str, n: int, rng=None, **kwargs) -> np.ndarray:
    """Dispatch to a named generator from :data:`GENERATORS`."""
    if kind not in GENERATORS:
        raise KeyError(f"unknown topology {kind!r}; choices: {list(GENERATORS)}")
    return GENERATORS[kind](n, rng=rng, **kwargs)

def graph_metrics(W: np.ndarray) -> dict:
    """Return basic graph descriptors of a reservoir weight matrix.

    Useful to characterise structural heterogeneity across federated nodes.
    ``n_edges`` and ``density`` count *directed* edges; ``mean_degree`` and
    ``clustering`` are computed on the undirected projection, so a reciprocal
    pair contributes to the degree of both endpoints once.

    An empty matrix yields all-zero / ``nan`` descriptors rather than raising.
    """
    W = np.asarray(W)
    if W.ndim != 2 or W.shape[0] != W.shape[1]:
        raise ValueError(f"W must be a square 2-D matrix, got shape {W.shape}")
    adj = (W != 0).astype(int)
    n = adj.shape[0]
    if n == 0:
        # nx.average_clustering divides by the node count, so guard the empty case.
        return {"n_nodes": 0, "n_edges": 0, "density": 0.0, "mean_degree": 0.0,
                "clustering": float("nan"), "avg_path_length": float("nan")}
    g = nx.from_numpy_array(adj, create_using=nx.DiGraph)
    ug = g.to_undirected()
    n_edges = int(adj.sum())
    degrees = np.array([d for _, d in ug.degree()])
    try:
        avg_path = (nx.average_shortest_path_length(ug)
                    if nx.is_connected(ug) else float("nan"))
    except (nx.NetworkXError, nx.NetworkXPointlessConcept):
        avg_path = float("nan")
    return {
        "n_nodes": n,
        "n_edges": n_edges,
        "density": n_edges / (n * (n - 1)) if n > 1 else 0.0,
        "mean_degree": float(degrees.mean()) if degrees.size else 0.0,
        "clustering": float(nx.average_clustering(ug)),
        "avg_path_length": float(avg_path),
    }

def leaking_rates(n, kind="uniform", low=0.1, high=1.0, n_layers=3, rng=None):
    """Per-node leaking rates ``a`` for a *heterogeneous / multi-scale* reservoir.

    Different neurons integrate at different speeds, which helps tasks with mixed
    time-scales (finance, chaotic systems such as Mackey-Glass).

    Parameters
    ----------
    n
        Number of reservoir nodes.
    kind
        ``"uniform"`` (i.i.d. in ``[low, high]``), ``"log_uniform"`` (log-spaced,
        emphasising fast nodes), ``"layered"`` (``n_layers`` contiguous blocks
        with decreasing rates, high → low, i.e. fast → slow groups), or
        ``"constant"`` (all ``high``).
    low, high
        Range of leaking rates (each in ``(0, 1]``).
    n_layers
        Number of decreasing blocks for ``kind="layered"``.
    rng
        Seed or ``numpy`` generator.

    Returns
    -------
    a : ndarray of shape (n,)
        Per-node leaking rates, ready to pass as ``EchoStateNetwork(leaking_rate=a)``.
    """
    n = _check_n(n)
    # A leaking rate of 0 freezes a neuron for ever, and low > high silently
    # produced an empty/NaN range, so both are rejected up front.
    if not np.isfinite(low) or low <= 0.0:
        raise ValueError(f"low must be finite and > 0, got {low}")
    if not np.isfinite(high) or high <= 0.0:
        raise ValueError(f"high must be finite and > 0, got {high}")
    if low > high:
        raise ValueError(f"need low <= high, got low={low}, high={high}")
    if n_layers < 1:
        raise ValueError(f"n_layers must be >= 1, got {n_layers}")
    rng = _as_rng(rng)
    if kind == "uniform":
        return rng.uniform(low, high, size=n)
    if kind == "log_uniform":
        return np.exp(rng.uniform(np.log(low), np.log(high), size=n))
    if kind == "constant":
        return np.full(n, float(high))
    if kind == "layered":
        levels = np.linspace(high, low, max(1, n_layers))
        return np.concatenate([
            np.full(len(idx), levels[k])
            for k, idx in enumerate(np.array_split(np.arange(n), max(1, n_layers)))
        ])
    raise ValueError(
        f"unknown kind {kind!r}; choices: "
        "'uniform', 'log_uniform', 'layered', 'constant'"
    )

def mixed_activations(n, types=("tanh", "sigmoid", "sin"), weights=None, rng=None):
    """Assign a node nonlinearity to each reservoir node (*multi-type* reservoir).

    Mixing activation functions within one reservoir mimics biological neuronal
    diversity and broadens the basis of dynamics. Returns an array of activation
    names ready to pass as ``EchoStateNetwork(activation=...)``.

    Parameters
    ----------
    n
        Number of reservoir nodes.
    types
        Activation names to mix (see :data:`esnfed.esn.ACTIVATIONS`).
    weights
        Optional mixing proportions (defaults to uniform).
    rng
        Seed or ``numpy`` generator.
    """
    from .esn import ACTIVATIONS

    n = _check_n(n)
    rng = _as_rng(rng)
    types = list(types)
    if not types:
        raise ValueError("types must name at least one activation")
    unknown = [t for t in types if t not in ACTIVATIONS]
    if unknown:
        raise ValueError(
            f"unknown activation(s) {unknown}; choices: {sorted(ACTIVATIONS)}"
        )
    p = None
    if weights is not None:
        p = np.asarray(weights, dtype=float)
        if p.shape != (len(types),):
            raise ValueError(
                f"weights must have one entry per type ({len(types)}), got {p.shape}"
            )
        if np.any(p < 0):
            raise ValueError("weights must be non-negative")
        total = p.sum()
        if not total > 0:
            raise ValueError("weights must sum to a positive value")
        p = p / total
    return rng.choice(types, size=n, p=p)

def narma10(n: int, rng=None, warmup: int = 50, *,
            max_retries: int = 20) -> tuple[np.ndarray, np.ndarray]:
    """Generate a NARMA-10 input/target sequence of length ``n``.

    The NARMA-10 recursion is only *conditionally* stable: for an unlucky input
    stream the cubic feedback term diverges and the target overflows to ``inf``
    (this happens for roughly one seed in a hundred). Such a sequence is useless
    -- it silently poisons any model trained on it, and every error metric becomes
    ``nan`` -- so a diverging draw is rejected and re-drawn from the same
    generator, which is the standard practice for this benchmark. A warning names
    how many draws were discarded, and ``ValueError`` is raised if no stable
    sequence is found within ``max_retries``.
    """
    if n < 1:
        raise ValueError(f"n must be >= 1, got {n}")
    if warmup < 0:
        raise ValueError(f"warmup must be >= 0, got {warmup}")
    rng = np.random.default_rng(rng) if not isinstance(rng, np.random.Generator) else rng
    total = n + warmup
    for attempt in range(max_retries + 1):
        u = rng.uniform(0.0, 0.5, size=total)
        y = np.zeros(total)
        with np.errstate(over="ignore", invalid="ignore"):
            for t in range(10, total):
                y[t] = (
                    0.3 * y[t - 1]
                    + 0.05 * y[t - 1] * np.sum(y[t - 10 : t])
                    + 1.5 * u[t - 10] * u[t - 1]
                    + 0.1
                )
        # A stable NARMA-10 orbit stays in roughly [0, 1]; anything that leaves a
        # generous bound has entered the divergent regime.
        if np.all(np.isfinite(y)) and np.max(np.abs(y)) < _NARMA10_STABLE_BOUND:
            if attempt:
                warnings.warn(
                    f"discarded {attempt} diverging NARMA-10 draw(s) before "
                    "obtaining a stable sequence",
                    stacklevel=2,
                )
            return u[warmup:].reshape(-1, 1), y[warmup:].reshape(-1, 1)
    raise ValueError(
        f"could not generate a stable NARMA-10 sequence of length {n} in "
        f"{max_retries + 1} attempts; the recursion kept diverging"
    )

def mackey_glass(
    n: int,
    tau: int = 17,
    beta: float = 0.2,
    gamma: float = 0.1,
    power: int = 10,
    dt: float = 1.0,
    seed: int | None = None,
    discard: int = 250,
) -> tuple[np.ndarray, np.ndarray]:
    """Generate a Mackey-Glass series; return ``(x_t, x_{t+1})`` for prediction."""
    rng = np.random.default_rng(seed)
    history_len = tau + 1
    total = n + discard + 1
    x = np.empty(total + history_len)
    x[:history_len] = 1.2 + 0.2 * (rng.random(history_len) - 0.5)
    for t in range(history_len, total + history_len):
        x_tau = x[t - tau]
        x[t] = x[t - 1] + dt * (
            beta * x_tau / (1.0 + x_tau**power) - gamma * x[t - 1]
        )
    series = x[history_len + discard :]
    u = series[:-1].reshape(-1, 1)
    y = series[1:].reshape(-1, 1)
    return u, y

def lorenz(
    n: int,
    dt: float = 0.02,
    sigma: float = 10.0,
    rho: float = 28.0,
    beta: float = 8.0 / 3.0,
    seed: int | None = None,
    discard: int = 500,
) -> tuple[np.ndarray, np.ndarray]:
    """Generate the Lorenz attractor; return ``(x_t, x_{t+1})`` for prediction."""
    rng = np.random.default_rng(seed)
    state = np.array([1.0, 1.0, 1.0]) + 0.01 * rng.standard_normal(3)
    total = n + discard + 1
    xs = np.empty(total)
    s = state
    for t in range(total):
        xs[t] = s[0]
        dx = sigma * (s[1] - s[0])
        dy = s[0] * (rho - s[2]) - s[1]
        dz = s[0] * s[1] - beta * s[2]
        s = s + dt * np.array([dx, dy, dz])
    series = xs[discard:]
    # Normalise to a sensible range for tanh reservoirs.
    series = (series - series.mean()) / series.std()
    return series[:-1].reshape(-1, 1), series[1:].reshape(-1, 1)

def split(
    u: np.ndarray, y: np.ndarray, train_frac: float = 0.7
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Chronological train/test split (no shuffling, to preserve dynamics)."""
    if len(u) != len(y):
        raise ValueError(f"u and y disagree on length: {len(u)} vs {len(y)}")
    if not 0.0 < train_frac < 1.0:
        raise ValueError(f"train_frac must be in (0, 1), got {train_frac}")
    cut = int(len(u) * train_frac)
    if cut == 0 or cut == len(u):
        raise ValueError(
            f"train_frac={train_frac} leaves one side of the split empty for a "
            f"series of {len(u)} steps"
        )
    return u[:cut], y[:cut], u[cut:], y[cut:]

def partition_iid(
    u: np.ndarray, y: np.ndarray, n_clients: int, rng=None
) -> list[tuple[np.ndarray, np.ndarray]]:
    """Split a sequence into ``n_clients`` contiguous blocks (federated clients).

    Contiguous blocks (rather than shuffled samples) keep each client's slice a
    valid time series, which is required for reservoir state harvesting. The split
    is therefore *deterministic*: ``rng`` is accepted only for signature
    compatibility with shuffling partitioners and is ignored.
    """
    if len(u) != len(y):
        raise ValueError(f"u and y disagree on length: {len(u)} vs {len(y)}")
    if n_clients < 1:
        raise ValueError(f"n_clients must be >= 1, got {n_clients}")
    if n_clients > len(u):
        raise ValueError(
            f"cannot split {len(u)} steps across {n_clients} clients without "
            "producing empty partitions"
        )
    bounds = np.linspace(0, len(u), n_clients + 1, dtype=int)
    return [(u[a:b], y[a:b]) for a, b in zip(bounds[:-1], bounds[1:])]

def from_array(
    series, *, predict: str = "next", normalize: bool = True
) -> tuple[np.ndarray, np.ndarray]:
    """Turn a 1-D real series into a one-step-ahead forecasting task.

    Parameters
    ----------
    series
        1-D sequence of observations (list or array).
    predict
        ``"next"`` to forecast the next value (level), or ``"change"`` to
        forecast the next first difference.
    normalize
        If true, z-score the series (recommended for tanh reservoirs).

    Returns
    -------
    (u, y)
        Input/target arrays of shape (T-1, 1).

    Notes
    -----
    Non-finite observations are dropped, which *closes the gap* rather than
    interpolating it: the remaining samples become adjacent in the returned task,
    so the series is treated as evenly spaced. A warning reports how many were
    removed.
    """
    s = np.asarray(series, dtype=float).ravel()
    finite = np.isfinite(s)
    n_dropped = int(s.size - finite.sum())
    if n_dropped:
        warnings.warn(
            f"dropped {n_dropped} non-finite observation(s); the surrounding "
            "samples are now treated as consecutive in time",
            stacklevel=2,
        )
    s = s[finite]
    # "next" pairs consecutive levels (needs 2 points); "change" first differences
    # the series and then pairs those (needs 3).
    needed = {"next": 2, "change": 3}
    if predict not in needed:
        raise ValueError(
            f"predict must be 'next' or 'change', got {predict!r}"
        )
    if s.size < needed[predict]:
        raise ValueError(
            f"predict={predict!r} needs at least {needed[predict]} finite "
            f"observations, got {s.size}"
        )
    if normalize and s.std() > 0:
        s = (s - s.mean()) / s.std()
    if predict == "next":
        return s[:-1].reshape(-1, 1), s[1:].reshape(-1, 1)
    d = np.diff(s)
    return d[:-1].reshape(-1, 1), d[1:].reshape(-1, 1)

def load_csv(
    path, column=None, *, predict: str = "next", normalize: bool = True
) -> tuple[np.ndarray, np.ndarray]:
    """Load a column from a CSV file and build a forecasting task.

    ``column`` may be a header name, an integer index, or ``None`` (last column).
    Missing markers ('', '.', 'NA', 'NaN', 'null') are dropped.
    """
    text = Path(path).read_text(encoding="utf-8")
    series = _column_from_csv(text, column)
    return from_array(series, predict=predict, normalize=normalize)