quickstart.py

"""
memory_esn quickstart -- one example per class in the hierarchy.

Run from the project root:
    python examples/quickstart.py
"""

import numpy as np

from memory_esn import (
    BaseESN,
    MultiESN,
    DoubleReservoirESN,
    fESN,
    wESN,
    TimeSeriesDataset,
)


def make_data(n=800, seed=0):
    rng = np.random.RandomState(seed)
    # Non-stationary (random-walk) inputs; smooth nonlinear target
    X = np.cumsum(rng.randn(n, 3), axis=0) * 0.1
    y = np.sin(X[:, :1]) + 0.1 * rng.randn(n, 1)
    split = int(0.8 * n)
    return X[:split], X[split:], y[:split], y[split:]


def mse(a, b):
    return float(np.mean((a - b) ** 2))


def main():
    Xtr, Xte, ytr, yte = make_data()

    print("1) BaseESN (single reservoir)")
    esn = BaseESN(n_reservoir=200, spectral_radius=0.95, random_state=0)
    esn.fit(Xtr, ytr, washout=100)
    print("   test MSE:", mse(esn.predict(Xte), yte))

    print("2) MultiESN (3 reservoirs, separate inputs)")
    multi = MultiESN(n_reservoirs=3, n_reservoir=[120, 100, 80], random_state=0)
    multi.fit([Xtr, Xtr, Xtr], ytr, washout=100)
    print("   test MSE:", mse(multi.predict([Xte, Xte, Xte]), yte))

    print("3) DoubleReservoirESN (two explicit inputs)")
    dbl = DoubleReservoirESN(n_reservoir=(150, 100), random_state=0)
    dbl.fit([Xtr, Xtr], ytr, washout=100)
    print("   test MSE:", mse(dbl.predict([Xte, Xte]), yte))

    print("4) fESN  (memory reservoir <- ((1-B)^d - 1) u)")
    # Defaults follow the paper: uniform init, 90% sparse reservoirs, pure-past
    # fractional memory filter. Optional state noise via noise=(sigma_x, sigma_m).
    fesn = fESN(n_reservoir=(200, 150), d=0.5, K=100,
                         noise=(0.0, 0.01), random_state=0)
    fesn.fit(Xtr, ytr, washout=100)
    y_pred = fesn.predict(Xte, continuation=True)  # continuation keeps length
    print("   univariate  (d=0.5)      test MSE:", mse(y_pred, yte))
    # multivariate: several differencing orders stacked as channels
    fesn_mv = fESN(n_reservoir=(200, 150), d=[0.3, 0.5, 0.8], K=100, random_state=0)
    fesn_mv.fit(Xtr, ytr, washout=100)
    print("   multivariate (d=[.3,.5,.8]) test MSE:",
          mse(fesn_mv.predict(Xte, continuation=True), yte),
          "| reservoir-2 channels:", fesn_mv.reservoirs_[1].n_inputs_)

    print("5) wESN  (memory reservoir <- ((1-B)^d - 1) MODWT_smooth(u))")
    wesn = wESN(n_reservoir=(200, 150), d=0.25, K=100,
                      wavelet="db4", wavelet_level=2, wavelet_norm="modwt", random_state=0)
    wesn.fit(Xtr, ytr, washout=100)
    print("   univariate  (smooth)     test MSE:", mse(wesn.predict(Xte, continuation=True), yte))
    # multivariate: detail levels 1 & 2 plus the smooth
    wesn_mv = wESN(n_reservoir=(200, 150), d=0.25, K=100, wavelet="db4",
                         wavelet_level=3, wavelet_components=[1, 2, "smooth"], random_state=0)
    wesn_mv.fit(Xtr, ytr, washout=100)
    print("   multivariate (D1,D2,A3)  test MSE:",
          mse(wesn_mv.predict(Xte, continuation=True), yte),
          "| reservoir-2 channels:", wesn_mv.reservoirs_[1].n_inputs_)

    print("6) TimeSeriesDataset (sliding-window splits)")
    series = np.sin(np.linspace(0, 40, 1000)) + 0.1 * np.random.RandomState(0).randn(1000)
    ds = TimeSeriesDataset(series, lookback=20, lookahead=5, test_size=50, scaling="standard")
    X_train, y_train = ds.get_full_train_data()
    print("   ", ds)
    print("    train windows:", X_train.shape, "->", y_train.shape)

    print("7) Paper-style multi-horizon forecasting (feed the raw series u(t))")
    # The reservoir processes u(t) sequentially; the readout maps phi(t) to the
    # next H values. Targets y[t] = [u(t+1), ..., u(t+H)] -> one ridge per horizon.
    H = 3
    u = (np.cumsum(np.random.RandomState(1).randn(600)) * 0.1).reshape(-1, 1)
    Y = np.column_stack([np.roll(u[:, 0], -h) for h in range(1, H + 1)])
    fesn_h = fESN(n_reservoir=(200, 150), d=0.4, K=80, random_state=0)
    fesn_h.fit(u, Y, washout=100)                      # y has H columns
    preds = fesn_h.predict(u, continuation=True)       # (T, H): one column per horizon
    print(f"    horizons H={H}: prediction shape {preds.shape},",
          f"readout W_out {fesn_h.readout_.coef_.shape} = (H, p+q+2)")


if __name__ == "__main__":
    main()