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."""
    d = A.shape[0]
    return np.linalg.solve(A + ridge * np.eye(d), B)

def random_reservoir(n: int, density: float = 0.1, rng=None) -> np.ndarray:
    """Erdos-Renyi reservoir: each directed edge present with prob. ``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.
    """
    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.
    """
    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).
    """
    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.
    """
    adj = (W != 0).astype(int)
    g = nx.from_numpy_array(adj, create_using=nx.DiGraph)
    ug = g.to_undirected()
    n = adj.shape[0]
    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)``.
    """
    rng = np.random.default_rng(rng) if not isinstance(rng, np.random.Generator) else 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}")

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.
    """
    rng = np.random.default_rng(rng) if not isinstance(rng, np.random.Generator) else rng
    types = list(types)
    p = None if weights is None else np.asarray(weights, float) / np.sum(weights)
    return rng.choice(types, size=n, p=p)

def narma10(n: int, rng=None, warmup: int = 50) -> tuple[np.ndarray, np.ndarray]:
    """Generate a NARMA-10 input/target sequence of length ``n``."""
    rng = np.random.default_rng(rng) if not isinstance(rng, np.random.Generator) else rng
    total = n + warmup
    u = rng.uniform(0.0, 0.5, size=total)
    y = np.zeros(total)
    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
        )
    u, y = u[warmup:], y[warmup:]
    return u.reshape(-1, 1), y.reshape(-1, 1)

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)."""
    cut = int(len(u) * train_frac)
    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.
    """
    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).
    """
    s = np.asarray(series, dtype=float).ravel()
    s = s[np.isfinite(s)]
    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)
    if predict == "change":
        d = np.diff(s)
        return d[:-1].reshape(-1, 1), d[1:].reshape(-1, 1)
    raise ValueError("predict must be 'next' or 'change'")

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)

def load_ted_spread(
    *, predict: str = "next", normalize: bool = True, raw: bool = False
):
    """Load the bundled TED spread series (interbank/counterparty-risk measure).

    The TED spread is the difference between the 3-month interbank rate and the
    3-month Treasury bill; it is a classic gauge of perceived counterparty/credit
    risk in the banking system. Daily, 1986--2022.

    Source: Federal Reserve Bank of St. Louis, FRED, series ``TEDRATE``.

    With ``raw=True`` returns the unprocessed 1-D level array instead of (u, y).
    """
    series = _column_from_csv(
        (_DATA_DIR / "ted_spread.csv").read_text(encoding="utf-8"), "ted_spread"
    )
    if raw:
        return series
    return from_array(series, predict=predict, normalize=normalize)

def load_fred(
    series_id: str,
    *,
    predict: str = "next",
    normalize: bool = True,
    cache_dir=None,
    timeout: float = 30.0,
    raw: bool = False,
):
    """Download a series from FRED (no API key) and build a forecasting task.

    Useful counterparty/credit-risk series include ``TEDRATE`` (TED spread),
    ``BAMLH0A0HYM2`` (US high-yield credit spread), ``BAA10Y`` (Baa credit
    spread) and ``STLFSI4`` / ``NFCI`` (financial stress/conditions). The CSV is
    cached on disk so repeated calls are offline.

    Data are retrieved from FRED (Federal Reserve Bank of St. Louis) and are
    subject to FRED's terms of use.
    """
    cache = Path(cache_dir) if cache_dir else (Path.home() / ".cache" / "esnfed")
    cache.mkdir(parents=True, exist_ok=True)
    cached = cache / f"fred_{series_id}.csv"
    if cached.exists():
        text = cached.read_text(encoding="utf-8")
    else:
        url = f"https://fred.stlouisfed.org/graph/fredgraph.csv?id={series_id}"
        with _urlopen(url, timeout) as resp:
            text = resp.read().decode("utf-8")
        cached.write_text(text, encoding="utf-8")
    series = _column_from_csv(text, series_id)
    if raw:
        return series
    return from_array(series, predict=predict, normalize=normalize)

def load_fred_matrix(
    series_ids,
    *,
    target=None,
    predict: str = "next",
    normalize: bool = True,
    cache_dir=None,
    timeout: float = 30.0,
    raw: bool = False,
):
    """Download several FRED series and align them into a *multivariate* task.

    The series are intersected on their common dates into a matrix ``(T, d)``;
    the task is to forecast the next-step value of the ``target`` series (default:
    the first) from the full d-dimensional vector. This exercises high-dimensional
    multivariate scaling --- e.g. forecasting counterparty risk (``TEDRATE``) from
    a panel such as ``["TEDRATE", "VIXCLS", "DGS10", "DFF"]``.

    Returns ``(u, y)`` with ``u`` of shape ``(T-1, d)`` and ``y`` of shape
    ``(T-1, 1)`` (or the raw ``(T, d)`` matrix if ``raw=True``).
    """
    if isinstance(series_ids, str):
        series_ids = [series_ids]
    cache = _cache_dir(cache_dir)
    dicts = [_fred_series_dict(_fred_text(s, cache, timeout)) for s in series_ids]
    common = sorted(set.intersection(*[set(d) for d in dicts]))
    if not common:
        raise ValueError("the requested FRED series have no overlapping dates")
    M = np.array([[d[dt] for d in dicts] for dt in common], dtype=float)
    if normalize:
        std = M.std(0)
        M = (M - M.mean(0)) / np.where(std > 0, std, 1.0)
    if raw:
        return M
    ti = series_ids.index(target) if target is not None else 0
    if predict == "next":
        return M[:-1], M[1:, ti : ti + 1]
    if predict == "change":
        D = np.diff(M, axis=0)
        return D[:-1], D[1:, ti : ti + 1]
    raise ValueError("predict must be 'next' or 'change'")

def load_japanese_vowels(*, cache_dir=None, timeout: float = 60.0) -> SequenceDataset:
    """Japanese Vowels (UCI 128): 9 male speakers uttering ``/ae/``, encoded as
    12-dimensional LPC cepstra, with variable-length utterances (7--29 frames).
    The task is speaker identification (9 classes). The natural federation unit is
    the speaker --- which is *also* the label, giving an extreme label-skew split
    in which local-only training cannot work and federation is essential.
    """
    cache = _cache_dir(cache_dir)
    zp = cache / "japanese_vowels.zip"
    if not zp.exists():
        url = "https://archive.ics.uci.edu/static/public/128/japanese+vowels.zip"
        with _urlopen(url, timeout) as r:
            zp.write_bytes(r.read())
    with zipfile.ZipFile(zp) as z:
        X_train = _parse_jv_blocks(z.read("ae.train").decode())
        X_test = _parse_jv_blocks(z.read("ae.test").decode())
        y_train = _counts_to_labels(z.read("size_ae.train").decode())
        y_test = _counts_to_labels(z.read("size_ae.test").decode())
    return SequenceDataset(X_train, y_train, y_train.copy(),
                           X_test, y_test, y_test.copy(), 9, 12)

def load_har(*, cache_dir=None, timeout: float = 180.0) -> SequenceDataset:
    """Human Activity Recognition Using Smartphones (UCI 240): 30 subjects, 6
    activities, raw 3-axial accelerometer + gyroscope signals windowed into
    128-timestep, 9-channel sequences at 50 Hz. The natural federation unit is the
    subject (``groups``): every subject performs all activities, so the split is
    feature-non-i.i.d. (people move differently) without label skew --- ideal for
    comparing the prediction ensemble against parameter aggregation.

    The ~60 MB archive is downloaded once and the parsed arrays are cached as a
    compressed ``.npz`` for fast re-loading.
    """
    cache = _cache_dir(cache_dir)
    npz = cache / "har_parsed.npz"
    if npz.exists():
        d = np.load(npz)
        return SequenceDataset(d["Xtr"], d["ytr"], d["gtr"],
                               d["Xte"], d["yte"], d["gte"], 6, 9)
    zp = cache / "har.zip"
    if not zp.exists():
        url = ("https://archive.ics.uci.edu/static/public/240/"
               "human+activity+recognition+using+smartphones.zip")
        with _urlopen(url, timeout) as r:
            zp.write_bytes(r.read())
    with zipfile.ZipFile(zp) as outer:
        inner = outer.read("UCI HAR Dataset.zip")
    with zipfile.ZipFile(io.BytesIO(inner)) as z:
        def load_split(split):
            base = f"UCI HAR Dataset/{split}/"
            chans = [np.loadtxt(io.BytesIO(
                z.read(f"{base}Inertial Signals/{s}_{split}.txt")))
                for s in _HAR_SIGNALS]
            X = np.stack(chans, axis=2).astype(np.float32)  # (n, 128, 9)
            y = np.loadtxt(io.BytesIO(z.read(f"{base}y_{split}.txt"))).astype(int) - 1
            g = np.loadtxt(io.BytesIO(z.read(f"{base}subject_{split}.txt"))).astype(int)
            return X, y, g
        Xtr, ytr, gtr = load_split("train")
        Xte, yte, gte = load_split("test")
    np.savez_compressed(npz, Xtr=Xtr, ytr=ytr, gtr=gtr, Xte=Xte, yte=yte, gte=gte)
    return SequenceDataset(Xtr, ytr, gtr, Xte, yte, gte, 6, 9)

def group_clients(X, y, groups):
    """Split labelled sequences into per-group client datasets.

    Returns a list of ``(sequences, labels)`` pairs, one per unique value in
    ``groups`` (e.g. one client per speaker or per subject).
    """
    clients = []
    for g in np.unique(groups):
        idx = np.where(np.asarray(groups) == g)[0]
        clients.append(([X[i] for i in idx], np.asarray(y)[idx]))
    return clients

def mse(y_true: np.ndarray, y_pred: np.ndarray) -> float:
    """Mean squared error."""
    y_true, y_pred = np.asarray(y_true), np.asarray(y_pred)
    return float(np.mean((y_true - y_pred) ** 2))

def rmse(y_true: np.ndarray, y_pred: np.ndarray) -> float:
    """Root mean squared error."""
    return float(np.sqrt(mse(y_true, y_pred)))

def nrmse(y_true: np.ndarray, y_pred: np.ndarray) -> float:
    """Normalised RMSE, scaled by the standard deviation of the target.

    NRMSE = 1 corresponds to a trivial predictor outputting the target mean.
    This is the standard figure of merit for NARMA and Mackey-Glass tasks.
    """
    y_true = np.asarray(y_true)
    sigma = y_true.std()
    if sigma == 0:
        return float("nan")
    return rmse(y_true, y_pred) / sigma