OpenReservoirComputing 0.1.0__py3-none-any.whl

orc/tuning/__init__.py

@@ -0,0 +1,6 @@
+"""Hyperparameter tuning utilities for Reservoir Computers.
+
+This module is currently a placeholder for future tuning implementations.
+"""
+
+__all__ = []

orc/utils/__init__.py

@@ -0,0 +1,6 @@
+"""Common utility functions for numerics and visualization."""
+
+from orc.utils import numerics, regressions, visualization
+from orc.utils.numerics import max_eig_arnoldi
+
+__all__ = ["numerics", "max_eig_arnoldi", "visualization", "regressions"]

orc/utils/numerics.py

@@ -0,0 +1,115 @@
+"""Additional utility functions for ORC."""
+
+import functools
+
+import jax
+import jax.numpy as jnp
+from jaxtyping import Array
+
+
+@functools.partial(jax.jit, static_argnames=["max_iters"])
+def _arnoldi_iteration(A, max_iters=200, seed=0):
+    """Perform Arnoldi iteration to find an orthonormal basis for the Krylov subspace.
+
+    Parameters
+    ----------
+    A : Array
+        The input matrix (m x m) for which the Krylov subspace is computed.
+    max_iters : int, optional
+        The maximum number of Arnoldi iterations to perform. Default is 200. If the
+        number of iterations exceeds the size of the matrix, it will be capped at m.
+    seed : int, optional
+        Random seed for initializing the starting vector. Default is 0.
+
+    Returns
+    -------
+    Q : Array
+        An orthonormal basis for the Krylov subspace (m x n).
+    H : Array
+        The upper Hessenberg matrix (n x n) representing the coefficients of the Krylov
+        subspace.
+    """
+    # A is m x m; n is the size of the krylov basis
+    m = A.shape[0]
+    n = max_iters
+
+    # choose a random vector to start iterating on
+    key = jax.random.PRNGKey(seed)
+    b = jax.random.normal(key, (m,))
+    q0 = b / jnp.linalg.norm(b)
+
+    # init krylov basis Q and hessenberg matrix H
+    Q = jnp.zeros((m, n + 1), dtype=A.dtype)
+    H = jnp.zeros((n + 1, n), dtype=A.dtype)
+    Q = Q.at[:, 0].set(q0)
+
+    # modified gs step to form orth krylov basis
+    def gs_step(carry, j):
+        v, Q, mask = carry
+        # Only apply when j is less than or equal to k (mask is True)
+        h_jk = jnp.where(mask[j], jnp.dot(Q[:, j], v), 0.0)
+        v = jnp.where(mask[j], v - h_jk * Q[:, j], v)
+        return (v, Q, mask), h_jk
+
+    def arnoldi_step(carry, k):
+        A, Q, H = carry
+
+        # new candidate vector
+        v = A @ Q[:, k]
+
+        # Create a mask for valid indices (0 to k)
+        idx_mask = jnp.arange(n + 1) <= k
+
+        # run modified gs with fixed-size loop and masking
+        final_carry_gs, h_jk_vals = jax.lax.scan(
+            gs_step,
+            (v, Q, idx_mask),
+            jnp.arange(n + 1),  # fixed size scan
+        )
+        v = final_carry_gs[0]  # orthogonalized candidate vector
+
+        # Calculate subdiagonal
+        h_kplus1k = jnp.linalg.norm(v)
+
+        # Update H column k using a mask
+        col_indices = jnp.arange(n + 1)
+        mask = col_indices <= k  # For the first k+1 elements
+
+        # Apply h_jk_vals for the first k+1 elements using the mask
+        H_col = jnp.where(mask, h_jk_vals, H[:, k])
+
+        # Set the k+1 element to h_kplus1k
+        H_col = H_col.at[k + 1].set(h_kplus1k)
+
+        # Update the k-th column of H
+        H = H.at[:, k].set(H_col)
+
+        # Update Q[:, k+1]
+        Q = Q.at[:, k + 1].set(v / (h_kplus1k))
+
+        return (A, Q, H), None
+
+    final_carry_arnoldi, _ = jax.lax.scan(arnoldi_step, (A, Q, H), jnp.arange(n))
+    _, Q, H = final_carry_arnoldi
+
+    return Q[:, :n], H[:n, :n]
+
+
+@functools.partial(jax.jit, static_argnames=("max_iters",))
+def max_eig_arnoldi(A: Array, max_iters: int = 200, seed: int = 0):
+    """Compute the maximum eigenvalue of a matrix using the Arnoldi iteration method.
+
+    Parameters
+    ----------
+    A : Array
+        The input matrix (m x m) for which the maximum eigenvalue is computed.
+    max_iters : int, optional
+        The maximum number of Arnoldi iterations to perform. Default is 200. If the
+        number of iterations exceeds the size of the matrix, it will be capped at m.
+    seed : int, optional
+        Random seed for initializing the starting vector. Default is 0.
+    """
+    _, H = _arnoldi_iteration(A, max_iters, seed)
+    eigvals = jnp.linalg.eigvals(H)
+    lambda_max = eigvals[jnp.argmax(jnp.abs(eigvals))]
+    return lambda_max

orc/utils/regressions.py

@@ -0,0 +1,73 @@
+"""Implements common regressions used to train RC models."""
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from jaxtyping import Array
+
+
+def ridge_regression(res_seq: Array, target_seq: Array, beta: float = 1e-7):
+    """Solve a single matrix ridge regression problem.
+
+    Parameters
+    ----------
+    res_seq : Array
+        Sequence of training reservoir states, (shape=(seq_len, res_dim)).
+    target_seq : Array
+        Sequence of training targe states, (shape=(seq_len, out_dim)).
+    beta : float
+        Tikhonov regularization parameter.
+
+    Returns
+    -------
+    Array
+        Solution to ridge regression s.t. cmat @ res_seq = target_seq.
+    """
+    lhs = res_seq.T @ res_seq + beta * jnp.eye(res_seq.shape[1], dtype=res_seq.dtype)
+    rhs = res_seq.T @ target_seq
+    cmat = jax.scipy.linalg.solve(lhs, rhs, assume_a="sym").T
+    return cmat
+
+
+_solve_all_ridge_reg = eqx.filter_vmap(ridge_regression, in_axes=eqx.if_array(1))
+
+
+def _solve_all_ridge_reg_batched(
+    res_seq_train: Array, target_seq: Array, beta: float, batch_size: int
+) -> Array:
+    """Solve ridge regression for all parallel reservoirs using batched vmap.
+
+    This function processes the parallel reservoirs in batches to reduce memory
+    usage for large numbers of parallel reservoirs.
+
+    Parameters
+    ----------
+    res_seq_train : Array
+        Training reservoir states, shape=(seq_len, chunks, res_dim).
+    target_seq : Array
+        Target sequence, shape=(seq_len, chunks, out_dim).
+    beta : float
+        Tikhonov regularization parameter.
+    batch_size : int
+        Number of parallel reservoirs to process in each batch.
+
+    Returns
+    -------
+    Array
+        Ridge regression solution for all chunks, shape=(chunks, out_dim, res_dim).
+    """
+    chunks = res_seq_train.shape[1]
+
+    if batch_size >= chunks:
+        return _solve_all_ridge_reg(res_seq_train, target_seq, beta)
+
+    results = []
+    for i in range(0, chunks, batch_size):
+        end_idx = min(i + batch_size, chunks)
+        batch_res = res_seq_train[:, i:end_idx, :]
+        batch_target = target_seq[:, i:end_idx, :]
+
+        batch_result = _solve_all_ridge_reg(batch_res, batch_target, beta)
+        results.append(batch_result)
+
+    return jnp.concatenate(results, axis=0)

orc/utils/visualization.py

@@ -0,0 +1,193 @@
+"""Visualization utilities for plotting time series and spatiotemporal data."""
+
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def plot_time_series(
+    U_lst,
+    t=None,
+    time_series_labels=None,
+    line_formats=None,
+    state_var_names=None,
+    t_lim=None,
+    figsize=(20, 8),
+    x_label=r"$t$",
+    title=None,
+    **plot_kwargs,
+):
+    """Plot time series data with separate panels for each state variable.
+
+    Parameters
+    ----------
+    U_lst : 2D array or list of 2D arrays
+        If a 2D array, shape should be (Nt, Nu) where Nu is the number of state
+        variables and Nt is the number of time points.If a list of 2D arrays,
+        each array should have shape (Nt, Nu) and represent different time series.
+    t : 1D array, optional
+        1D array of time points. If None, the time points will be assumed to be
+        evenly spaced from 0 to Nt-1.
+    time_series_labels : list of strings, optional
+        List of strings containing the labels for each time series to be shown in
+        a legend. If None, no labels will be shown.
+    line_formats : list of strings, optional
+        List of strings containing the line formats for each time series. If None,
+        default line format will be used.
+    state_var_names : list of strings, optional
+        List of strings containing the names of the state variables. If None,
+        no y-axis labels will be shown.
+    t_lim : tuple, optional
+        Limit for the x-axis. If None, the x-axis will be set to the full
+        range of time points.
+    figsize : tuple, optional
+        Size of the figure to be created. Default is (20, 8).
+    x_label : string, optional
+        Label for the x-axis. Default is r'$t$'.
+    title : string, optional
+        Title of the plot. If None, no title is shown.
+    plot_kwargs : dict, optional
+        Additional arguments to pass to the plot function.
+    """
+    # Input validation
+    if not isinstance(U_lst, list):
+        if not isinstance(U_lst, jnp.ndarray | np.ndarray) or U_lst.ndim != 2:
+            raise TypeError(
+                "U_lst must be a 2D JAX or NumPy array or a list of \
+                            2D JAX/NumPy arrays."
+            )
+        U_lst = [U_lst]
+    else:
+        if not all(
+            isinstance(U, jnp.ndarray | np.ndarray) and U.ndim == 2 for U in U_lst
+        ):
+            raise TypeError("All elements in U_lst must be 2D JAX or NumPy arrays.")
+        if not all(U.shape == U_lst[0].shape for U in U_lst):
+            raise ValueError("All arrays in U_lst must have the same shape.")
+
+    Nu = U_lst[0].shape[1]
+    Nt = U_lst[0].shape[0]
+
+    if t is not None:
+        if not isinstance(t, jnp.ndarray | np.ndarray) or t.ndim != 1:
+            raise TypeError("t must be a 1D JAX or NumPy array.")
+        if len(t) != Nt:
+            raise ValueError(
+                f"Length of t ({len(t)}) must match the number of time\
+                             points in U_lst ({Nt})."
+            )
+
+    if time_series_labels is not None:
+        if not isinstance(time_series_labels, list):
+            raise TypeError("time_series_labels must be a list of strings.")
+        if len(time_series_labels) != len(U_lst):
+            raise ValueError(
+                f"Length of time_series_labels ({len(time_series_labels)})\
+                             must match the number of time series ({len(U_lst)})."
+            )
+
+    if line_formats is not None:
+        if not isinstance(line_formats, list):
+            raise TypeError("line_formats must be a list of strings.")
+        if len(line_formats) != len(U_lst):
+            raise ValueError(
+                f"Length of line_formats ({len(line_formats)}) must \
+                             match the number of time series ({len(U_lst)})."
+            )
+
+    if state_var_names is not None:
+        if not isinstance(state_var_names, list):
+            raise TypeError("state_var_names must be a list of strings.")
+        if len(state_var_names) != Nu:
+            raise ValueError(
+                f"Length of state_var_names ({len(state_var_names)}) \
+                             must match the number of state variables ({Nu})."
+            )
+
+    if t_lim is not None and not isinstance(t_lim, int | float):
+        raise TypeError("t_lim must be a number (int or float).")
+
+    # defaults
+    plot_kwargs.setdefault("linewidth", 2)
+
+    # setup time vectors
+    if t is None:
+        t = jnp.arange(Nt)
+    if t_lim is None:
+        t_lim = t[-1]
+
+    # handle optional inputs
+    if time_series_labels is None:
+        time_series_labels = [None for _ in range(len(U_lst))]
+    if line_formats is None:
+        line_formats = ["-" for _ in range(len(U_lst))]
+
+    # plot
+    fig, axs = plt.subplots(Nu, figsize=figsize)
+    # Ensure axs is always iterable, even if Nu=1
+    if Nu == 1:
+        axs = [axs]
+    for i in range(Nu):
+        for j, Y in enumerate(U_lst):
+            axs[i].plot(
+                t, Y[:, i], line_formats[j], label=time_series_labels[j], **plot_kwargs
+            )
+            axs[i].set_xlim([0, t_lim])
+        if state_var_names is not None:
+            axs[i].set(ylabel=state_var_names[i])
+    if time_series_labels[0] is not None:
+        axs[0].legend(loc="upper right")
+    axs[-1].set(xlabel=x_label)
+    if title is not None:
+        axs[0].set_title(title, fontsize=14)
+    plt.show()
+
+
+def imshow_1D_spatiotemp(
+    U, tN, domain=(0, 1), figsize=(20, 6), title=None, x_label=r"$t$", **imshow_kwargs
+):
+    """
+    Plot 1D spatiotemporal data using imshow.
+
+    Parameters
+    ----------
+    U: 2D array
+        Data to be plotted, shape should be (Nt, Nx) where Nt is the number of time
+        points and Nx is the number of spatial points
+    tN: float
+        Final time of the simulation
+    domain: tuple of length 2
+        Bounds of the spatial domain, default is (0, 1)
+    figsize: tuple
+        Size of the figure to be created, default is (20, 6)
+    title: string, optional
+        Title of the plot, if None no title is shown
+    x_label: string, optional
+        Label for the x-axis, default is r'$t$'
+    **imshow_kwargs: additional arguments to pass to imshow
+    """
+    # Input validation
+    if not isinstance(U, jnp.ndarray | np.ndarray) or U.ndim != 2:
+        raise TypeError("U must be a 2D JAX or NumPy array.")
+    if not isinstance(domain, tuple) or len(domain) != 2:
+        raise TypeError("domain must be a tuple of length 2.")
+    if not all(isinstance(x, int | float) for x in domain):
+        raise TypeError("Both elements of domain must be numbers (int or float).")
+
+    # set defaults for imshow
+    imshow_kwargs.setdefault("aspect", "auto")
+    imshow_kwargs.setdefault("origin", "lower")
+    imshow_kwargs.setdefault("cmap", "RdGy")
+    imshow_kwargs.setdefault("extent", [0, tN, domain[0], domain[1]])
+
+    plt.figure(figsize=figsize, dpi=200)
+    plt.imshow(U.T, **imshow_kwargs)
+    plt.ylabel("x")
+    plt.xlabel(x_label)
+    if title is not None:
+        plt.title(title)
+    plt.colorbar(pad=0.01, label=r"$u$")
+    plt.show()
+
+
+# TODO: add plot_attrator function to visualize 2D/3D attractors in state space