OpenReservoirComputing 0.3.2__tar.gz → 0.3.3__tar.gz

{openreservoircomputing-0.3.2 → openreservoircomputing-0.3.3}/JOSS/paper.md

@@ -44,7 +44,7 @@ OpenReservoirComputing (ORC) is a Python library for reservoir computing (RC) bu
 
 # Statement of Need
 
-Time-series prediction, classification, and control are fundamental tasks across science and engineering, arising in applications from climate modeling and fluid dynamics to robotics and neuroscience. Deep learning approaches to these tasks typically require large datasets, long training times, and expensive tuning of optimization hyperparameters. RC offers a compelling alternative. Since only the readout layer is trained via a single ridge regression, RC models can be trained in a fraction of the time required by comparable recurrent neural networks, often with less data and fewer hyperparameters to tune [@lukosevicius2009reservoir,@pathak2018model,@wyder2025common]. This makes RC particularly attractive for rapid prototyping, real-time applications, and data-limited settings. However, realizing these benefits in practice requires software that is both efficient and adaptable.
+Time-series prediction, classification, and control are fundamental tasks across science and engineering, arising in applications from climate modeling and fluid dynamics to robotics and neuroscience. Deep learning approaches to these tasks typically require large datasets, long training times, and expensive tuning of optimization hyperparameters. RC offers a compelling alternative. Since only the readout layer is trained via a single ridge regression, RC models can be trained in a fraction of the time required by comparable recurrent neural networks, often with less data and fewer hyperparameters to tune [@lukosevicius2009reservoir;@pathak2018model;@wyder2025common]. This makes RC particularly attractive for rapid prototyping, real-time applications, and data-limited settings. However, realizing these benefits in practice requires software that is both efficient and adaptable.
 
 ORC's built-in models provide an easy entry point for users new to the field. In particular, a new user can supply their own time-series data to instantiate, train, and forecast in three simple lines of code. 
 
@@ -58,7 +58,7 @@ U_pred = esn.forecast(fcast_len=1000, res_state=R[-1])
 
 Built-in visualization tools make it easy to evaluate model performance. Varying the hyperparameters of built-in models lets users explore how RC performance depends on configuration choices. ORC's JAX foundation makes scaling to higher-dimensional parallel reservoir architectures equally simple.
 
-Much RC research is aimed at designing performant reservoir architectures. ORC makes this easy through its use of abstract base classes for Embedding, Driver, and Readout layers. Users need only define forward pass logic to integrate a new reservoir topology or readout strategy, while reusing the rest of the framework. This modular approach also enables ablation studies on how different components affect RC performance. Because of ORC's functional approach in JAX, built-in and user-created models provide end-to-end differentiability by default. This enables gradient-based optimization of input sequences for control problems. This also makes ORC well suited to integrate with deep learning models such as those presented in [@ozalp2023reconstruction; @ozalp2025real].
+Much RC research is aimed at designing performant reservoir architectures. ORC makes this easy through its use of abstract base classes for Embedding, Driver, and Readout layers. Users need only define forward pass logic to integrate a new reservoir topology or readout strategy, while reusing the rest of the framework. This modular approach also enables ablation studies on how different components affect RC performance. Because of ORC's functional approach in JAX, built-in and user-created models provide end-to-end differentiability by default. This enables gradient-based optimization of input sequences for control problems. This also makes ORC well suited to integrate with deep learning models such as those presented by Özalp et al. [-@ozalp2023reconstruction; -@ozalp2025real].
 
 # State of the Field
 | | **ORC** | **ReservoirPy** | **RC.jl** |
@@ -81,12 +81,12 @@ First, ORC was *designed* on top of JAX [@jax2018github] and Equinox [@kidger202
 Second, ORC has a different built-in feature set. ORC supports continuous-time reservoir dynamics via Diffrax [@kidger2021on], allowing users to define reservoir equations as ordinary differential equations solved with adaptive-step integrators. ORC also supports novel architectures such as Taylor-expanded and GRU-based drivers alongside standard echo state networks. ReservoirPy does not provide these functionalities. 
 Third, ORC models compose seamlessly with other deep learning models in Equinox, whereas ReservoirPy is suited to standalone RC tasks. 
 
-Other open source libraries for RC include Pytorch-ESN [@nardo2018pytorchesn] and ReservoirComputing.jl [@martinuzzi2022reservoircomputing]. Pytorch-ESN allows for integration with other PyTorch models, but is not as widely adopted and offers a much more limited feature set than ORC or ReservoirPy. For instance, it does not support sparse matrix operations or parallel reservoirs. ReservoirComputing.jl is a widely used Julia library for RC that influenced many of ORC's design choices, but lacks Python ML ecosystem integration. Table 1 summarizes the functionality of ORC, ReservoirPy, and ReservoirComputing.jl. 
+Other open source libraries for RC include PyTorch-ESN [@nardo2018pytorchesn] and ReservoirComputing.jl [@martinuzzi2022reservoircomputing]. PyTorch-ESN allows for integration with other PyTorch models, but is not as widely adopted and offers a much more limited feature set than ORC or ReservoirPy. For instance, it does not support sparse matrix operations or parallel reservoirs. ReservoirComputing.jl is a widely used Julia library for RC that influenced many of ORC's design choices, but lacks Python ML ecosystem integration. Table 1 summarizes the functionality of ORC, ReservoirPy, and ReservoirComputing.jl. 
 
 
 # Software Design
 
-![ORC three-layer pipeline architecture. Each reservoir computer (RC) is decomposed into (i) an ``embedding'' function that lifts a low-dimensional signal $u_t$ to the reservoir dimension, (ii) a ``driver'' function that propagates the reservoir state, and (iii) a ``readout'' that maps back to a target $y$. For control and forecasting RCs, the target $y$ is typically $u_{t+1}$ (either in the presence of a forcing term or not) and for classification the target $y$ is a label. \label{fig:architecture}](../imgs/architecture_paper.pdf)
+![ORC three-layer pipeline architecture. Each reservoir computer (RC) is decomposed into (i) an embedding function that lifts a low-dimensional signal $u_t$ to the reservoir dimension, (ii) a driver function that propagates the reservoir state, and (iii) a readout that maps back to a target $y$. For control and forecasting RCs, the target $y$ is typically $u_{t+1}$ (either in the presence of a forcing term or not) and for classification the target $y$ is a label. \label{fig:architecture}](../imgs/architecture_paper.pdf)
 
 ORC models are decomposed into three components, illustrated in \autoref{fig:architecture}: (i) an embedding $f_E$ that lifts a low-dimensional input signal $u_t$ to a high-dimensional space, (ii) a driver $f_R$ that propagates the high-dimensional state $r_t$, and (iii) a readout $f_O$ that maps the latent state back to an approximation of some low-dimensional signal $y_t$. Depending on the task at hand, $y_t$ may be a future time-step of $u_t$, a label associated with input data, or some other target signal. ORC differs from many existing approaches that unify (i) and (ii). Separating the embedding from the reservoir state propagation allows for cleaner application of RC to non-standard tasks, such as acting as a surrogate model for model predictive control.
 
@@ -104,12 +104,14 @@ ORC also includes a data generation module with ODE and PDE integrators for stan
 
 ![ORC with GPU acceleration enables significantly faster performance than ReservoirPy, even when using ReservoirPy's JAX backend. Panel (a) shows the time per forecast step of RC models with varying reservoir dimension trained to forecast the Lorenz system, while panel (b) shows training time for RC models with fixed reservoir dimension of 2000 but a varying number of training samples. Performance of the two libraries with and without GPU acceleration are shown. GPU results were obtained on an NVIDIA A40 GPU and CPU results were obtained with an Apple M2 chip. \label{fig:comp}](../imgs/benchmark_results_combined.png) 
 
-ORC enables easy reimplementation of architectures that integrate RC with larger neural networks [@ozalp2023reconstruction; @ozalp2025real], and makes it easier to iterate on these ideas than any existing library. We also benchmark ORC against ReservoirPy across reservoir sizes and the number of training samples \autoref{fig:comp}. We find that with GPU acceleration, ORC scales far more favorably than ReservoirPy. Moreover, ORC is the only current package that supports training RC for control tasks, enabled by its efficiency and the end-to-end differentiability that JAX and Equinox provide. RC-based control has recently been shown to be advantageous over other forms of RNN-based MPC [@williams2024reservoir]. ORC's performance also makes it well suited for the parallel architectures needed to extend RC to higher-dimensional settings.
+ORC enables easy reimplementation of architectures that integrate RC with larger neural networks [@ozalp2023reconstruction; @ozalp2025real], and makes it easier to iterate on these ideas than any existing library. We also benchmark ORC against ReservoirPy across reservoir sizes and numbers of training samples, as shown in \autoref{fig:comp}. We find that with GPU acceleration, ORC scales far more favorably than ReservoirPy. Moreover, ORC is the only current package that supports training RC for control tasks, enabled by its efficiency and the end-to-end differentiability that JAX and Equinox provide. RC-based control has recently been shown to be advantageous over other forms of RNN-based MPC [@williams2024reservoir]. ORC's performance also makes it well suited for the parallel architectures needed to extend RC to higher-dimensional settings.
 
 
 # AI Usage Disclosure
 
-Claude (Anthropic) was used for code assistance during code debugging, proofreading this paper, and the generation of the architecture visualization SVG. ChatGPT (OpenAI) was used to generate the ORC logo. All generated code was reviewed, tested, and validated by the authors.
+Claude (Anthropic, Opus 4-4.8) was used for code assistance during debugging, for proofreading and editing the text of this paper, and for generating the architecture visualization SVG (\autoref{fig:architecture}). ChatGPT (OpenAI, GPT-4o) was used to generate the ORC logo. All AI-assisted outputs, including code, paper text, and figures, were reviewed, edited, and validated by the authors before inclusion.
 
 # Acknowledgements
 The authors acknowledge support from the NSF AI Institute in Dynamic Systems (grant number 2112085). The authors also thank Anastasia Bizyaeva, Noa Kaplan, and Ling-Wei Kong for insightful conversations.
+
+# References

{openreservoircomputing-0.3.2 → openreservoircomputing-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: OpenReservoirComputing
-Version: 0.3.2
+Version: 0.3.3
 Summary: GPU accelerated implementations of common RC architectures
 Author-email: "Jan P. Williams" <jmpw1@uw.edu>, Dima Tretiak <dtretiak@uw.edu>
 Maintainer-email: "Jan P. Williams" <jmpw1@uw.edu>

{openreservoircomputing-0.3.2 → openreservoircomputing-0.3.3}/pyproject.toml

@@ -52,6 +52,9 @@ testpaths = "tests"
 addopts = "--ignore=scripts --ignore=docs"
 env = ["MPLBACKEND=Agg"]
 
+[tool.ruff]
+exclude = ["examples"]
+
 [tool.ruff.lint]
 select = ["E", "W", "F", "D", "UP", "B", "A", "C4", "N", "SIM", "I"]
 

{openreservoircomputing-0.3.2 → openreservoircomputing-0.3.3}/src/OpenReservoirComputing.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: OpenReservoirComputing
-Version: 0.3.2
+Version: 0.3.3
 Summary: GPU accelerated implementations of common RC architectures
 Author-email: "Jan P. Williams" <jmpw1@uw.edu>, Dima Tretiak <dtretiak@uw.edu>
 Maintainer-email: "Jan P. Williams" <jmpw1@uw.edu>

{openreservoircomputing-0.3.2 → openreservoircomputing-0.3.3}/src/OpenReservoirComputing.egg-info/SOURCES.txt

@@ -46,11 +46,14 @@ examples/lorenz.ipynb
 examples/rc_background.ipynb
 imgs/ORC_logo_cropped.png
 imgs/architecture_paper.pdf
+imgs/benchmark_results_combined.png
 imgs/readme_example_forecast.png
 src/OpenReservoirComputing.egg-info/PKG-INFO
 src/OpenReservoirComputing.egg-info/SOURCES.txt
 src/OpenReservoirComputing.egg-info/dependency_links.txt
 src/OpenReservoirComputing.egg-info/requires.txt
+src/OpenReservoirComputing.egg-info/scm_file_list.json
+src/OpenReservoirComputing.egg-info/scm_version.json
 src/OpenReservoirComputing.egg-info/top_level.txt
 src/orc/__init__.py
 src/orc/drivers.py

openreservoircomputing-0.3.3/src/OpenReservoirComputing.egg-info/scm_file_list.json

@@ -0,0 +1,86 @@
+{
+  "files": [
+    "mkdocs.yml",
+    "pyproject.toml",
+    "LICENSE",
+    "README.md",
+    ".gitignore",
+    ".github/CODEOWNERS",
+    ".github/workflows/draft-pdf.yml",
+    ".github/workflows/tests.yml",
+    ".github/workflows/publish.yml",
+    ".github/workflows/docs.yml",
+    "JOSS/paper.bib",
+    "JOSS/CODE_OF_CONDUCT.md",
+    "JOSS/joss-paper.pdf",
+    "JOSS/CITATION.cff",
+    "JOSS/paper.md",
+    "examples/jit_compatibility.ipynb",
+    "examples/benchmark_combined.py",
+    "examples/continuous_rc.ipynb",
+    "examples/classification.ipynb",
+    "examples/control.ipynb",
+    "examples/data_library.ipynb",
+    "examples/lorenz.ipynb",
+    "examples/rc_background.ipynb",
+    "examples/ks.ipynb",
+    "tests/test_drivers.py",
+    "tests/test_embeddings.py",
+    "tests/test_readouts.py",
+    "tests/utils/test_numerics.py",
+    "tests/utils/test_visualization.py",
+    "tests/utils/test_regressions.py",
+    "tests/control/test_control_models.py",
+    "tests/classifier/test_classifier_models.py",
+    "tests/data/test_integrators.py",
+    "tests/forecaster/test_forecast_models.py",
+    "imgs/readme_example_forecast.png",
+    "imgs/benchmark_results_combined.png",
+    "imgs/architecture_paper.pdf",
+    "imgs/ORC_logo_cropped.png",
+    "src/orc/readouts.py",
+    "src/orc/embeddings.py",
+    "src/orc/__init__.py",
+    "src/orc/drivers.py",
+    "src/orc/utils/regressions.py",
+    "src/orc/utils/__init__.py",
+    "src/orc/utils/numerics.py",
+    "src/orc/utils/visualization.py",
+    "src/orc/control/train.py",
+    "src/orc/control/__init__.py",
+    "src/orc/control/models.py",
+    "src/orc/control/base.py",
+    "src/orc/classifier/train.py",
+    "src/orc/classifier/__init__.py",
+    "src/orc/classifier/models.py",
+    "src/orc/classifier/base.py",
+    "src/orc/data/integrators.py",
+    "src/orc/data/__init__.py",
+    "src/orc/forecaster/train.py",
+    "src/orc/forecaster/__init__.py",
+    "src/orc/forecaster/models.py",
+    "src/orc/forecaster/base.py",
+    "docs/contributing.md",
+    "docs/index.md",
+    "docs/examples",
+    "docs/getting-started/installation.md",
+    "docs/getting-started/quickstart.md",
+    "docs/user-guide/overview.md",
+    "docs/user-guide/readouts.md",
+    "docs/user-guide/data.md",
+    "docs/user-guide/embeddings.md",
+    "docs/user-guide/models.md",
+    "docs/user-guide/drivers.md",
+    "docs/stylesheets/extra.css",
+    "docs/imgs/ORC_logo_cropped.png",
+    "docs/api/forecaster.md",
+    "docs/api/classifier.md",
+    "docs/api/readouts.md",
+    "docs/api/control.md",
+    "docs/api/data.md",
+    "docs/api/embeddings.md",
+    "docs/api/utils.md",
+    "docs/api/drivers.md",
+    "docs/javascripts/mathjax.js"
+  ]
+}

openreservoircomputing-0.3.3/src/OpenReservoirComputing.egg-info/scm_version.json

@@ -0,0 +1,8 @@
+{
+  "tag": "0.3.3",
+  "distance": 0,
+  "node": "gcfeb341f30a1ad48016b4daee4b959efb2d13a5a",
+  "dirty": false,
+  "branch": "HEAD",
+  "node_date": "2026-07-23"
+}

{openreservoircomputing-0.3.2 → openreservoircomputing-0.3.3}/src/orc/__init__.py

@@ -17,7 +17,7 @@ from orc import (
 
 if not getattr(jax.config, "jax_enable_x64", False):
     warnings.warn(
-        "For good performance, orc often requires float64 precision. Enable it " \
+        "For good performance, orc often requires float64 precision. Enable it "
         "before importing orc with: "
         "jax.config.update('jax_enable_x64', True)",
         UserWarning,

{openreservoircomputing-0.3.2 → openreservoircomputing-0.3.3}/src/orc/classifier/train.py

@@ -17,6 +17,7 @@ from orc.utils.regressions import (
 
 VariableRCClassifier = TypeVar(name="VariableRCClassifier", bound=RCClassifierBase)
 
+
 def train_RCClassifier(
     model: VariableRCClassifier,
     train_seqs: Array,
@@ -66,9 +67,7 @@ def train_RCClassifier(
         raise ValueError("Number of training sequences must match number of labels.")
 
     n_samples = train_seqs.shape[0]
-    initial_res_states = jnp.tile(
-        model.driver.default_state(), (n_samples, 1)
-    )
+    initial_res_states = jnp.tile(model.driver.default_state(), (n_samples, 1))
 
     # Force all sequences through the reservoir in parallel via vmap
     all_res_seqs = jax.vmap(model.force)(train_seqs, initial_res_states)

{openreservoircomputing-0.3.2 → openreservoircomputing-0.3.3}/src/orc/control/models.py

@@ -152,5 +152,5 @@ class ESNController(RCControllerBase):
             seed=seed,
             alpha_1=alpha_1,
             alpha_2=alpha_2,
-            alpha_3=alpha_3
+            alpha_3=alpha_3,
         )

openreservoircomputing-0.3.3/src/orc/control/train.py

@@ -0,0 +1,213 @@
+"""Training functions for reservoir computer controllers."""
+
+from typing import TypeVar
+
+import equinox as eqx
+import jax.numpy as jnp
+from jaxtyping import Array
+
+from orc.control.base import RCControllerBase
+from orc.control.models import ESNController
+from orc.utils.regressions import (
+    _solve_all_ridge_reg,
+    _solve_all_ridge_reg_batched,
+    ridge_regression,
+)
+
+VariableRCController = TypeVar(name="VariableRCController", bound=RCControllerBase)
+
+
+def train_RCController(
+    model: VariableRCController,
+    train_seq: Array,
+    control_seq: Array,
+    target_seq: Array | None = None,
+    spinup: int = 0,
+    initial_res_state: Array | None = None,
+    beta: float = 8e-8,
+    batch_size: int | None = None,
+    multi_sequence: bool = False,
+    **force_kwargs,
+) -> tuple[VariableRCController, Array]:
+    """Unified training function for reservoir computer controllers.
+
+    Works with any model inheriting from RCControllerBase, including
+    ESNController and custom models with user-defined readout layers.
+
+    Parameters
+    ----------
+    model : RCControllerBase
+        Reservoir computer controller model to train.
+    train_seq : Array
+        Training input sequence for reservoir. Shape ``(seq_len, data_dim)``
+        for a single trajectory, or ``(n_traj, seq_len, data_dim)`` when
+        ``multi_sequence=True``. All trajectories must have the same length.
+    control_seq : Array
+        Control input sequence for reservoir. Shape ``(seq_len, control_dim)``
+        for a single trajectory, or ``(n_traj, seq_len, control_dim)`` when
+        ``multi_sequence=True``.
+    target_seq : Array, optional
+        Target sequence for training reservoir. Shape ``(seq_len, data_dim)``
+        for a single trajectory, or ``(n_traj, seq_len, data_dim)`` when
+        ``multi_sequence=True``. If None, defaults to train_seq[1:] (or
+        train_seq[:, 1:] for multi-sequence).
+    spinup : int
+        Initial transient of reservoir states to discard per trajectory.
+    initial_res_state : Array, optional
+        Initial reservoir state. If None, uses model.driver.default_state()
+        (tiled across trajectories when ``multi_sequence=True``). When
+        ``multi_sequence=True`` and not None, must have a leading trajectory
+        axis, i.e. shape ``(n_traj, ...)``.
+    beta : float
+        Tikhonov regularization parameter.
+    batch_size : int, optional
+        Number of parallel reservoirs to process in each batch for ridge
+        regression. If None (default), processes all reservoirs at once.
+        Only used when readout.chunks > 0.
+    multi_sequence : bool
+        If True, treat train_seq, control_seq, and target_seq as batches of
+        trajectories with a leading trajectory axis. Reservoir states from all
+        trajectories are concatenated (after spinup) before solving the
+        regression. This flag is a static boolean suitable for use as
+        ``static_argnums`` when wrapping with ``jax.jit``.
+    **force_kwargs
+        Additional keyword arguments passed to model.force().
+
+    Returns
+    -------
+    model : RCControllerBase
+        Trained model.
+    tot_res_seq : Array
+        Full reservoir state sequence from teacher forcing. Shape
+        ``(seq_len, ...)`` for single-sequence or ``(n_traj, seq_len, ...)``
+        for multi-sequence.
+    """
+    if multi_sequence:
+        if train_seq.shape[:2] != control_seq.shape[:2]:
+            raise ValueError(
+                "train_seq and control_seq must have the same (n_traj, seq_len)."
+            )
+    else:
+        if train_seq.shape[0] != control_seq.shape[0]:
+            raise ValueError("train_seq and control_seq must have the same length.")
+
+    seq_len = train_seq.shape[1] if multi_sequence else train_seq.shape[0]
+    if spinup >= seq_len:
+        raise ValueError(
+            "spinup must be less than the length of the training sequence."
+        )
+
+    if initial_res_state is None:
+        default = model.driver.default_state()
+        if multi_sequence:
+            n_traj = train_seq.shape[0]
+            initial_res_state = jnp.broadcast_to(
+                default[None], (n_traj,) + default.shape
+            )
+        else:
+            initial_res_state = default
+
+    if target_seq is None:
+        tot_seq = train_seq
+        tot_control_seq = control_seq
+        target_seq = train_seq[:, 1:, :] if multi_sequence else train_seq[1:, :]
+    else:
+        if multi_sequence:
+            tot_seq = jnp.concatenate([train_seq, target_seq[:, -1:, :]], axis=1)
+            tot_control_seq = jnp.concatenate(
+                [control_seq, control_seq[:, -1:, :]], axis=1
+            )
+        else:
+            tot_seq = jnp.vstack((train_seq, target_seq[-1:]))
+            tot_control_seq = jnp.vstack((control_seq, control_seq[-1:]))
+
+    if multi_sequence:
+        tot_res_seq = eqx.filter_vmap(model.force)(
+            tot_seq, tot_control_seq, initial_res_state
+        )
+        res_seq = tot_res_seq[:, :-1]
+        res_seq_train = eqx.filter_vmap(model.readout.prepare_train)(res_seq)
+        train_target = eqx.filter_vmap(model.readout.prepare_target)(target_seq)
+        res_seq_train = res_seq_train[:, spinup:].reshape(
+            (-1,) + res_seq_train.shape[2:]
+        )
+        train_target = train_target[:, spinup:].reshape((-1,) + train_target.shape[2:])
+    else:
+        tot_res_seq = model.force(
+            tot_seq, tot_control_seq, initial_res_state, **force_kwargs
+        )
+        res_seq = tot_res_seq[:-1]
+        res_seq_train = model.readout.prepare_train(res_seq)
+        train_target = model.readout.prepare_target(target_seq)
+        res_seq_train = res_seq_train[spinup:]
+        train_target = train_target[spinup:]
+
+    if model.readout.chunks > 0:
+        if batch_size is None:
+            cmat = _solve_all_ridge_reg(res_seq_train, train_target, beta)
+        else:
+            cmat = _solve_all_ridge_reg_batched(
+                res_seq_train, train_target, beta, batch_size
+            )
+    else:
+        cmat = ridge_regression(res_seq_train, train_target, beta)
+
+    new_readout = model.readout.set_wout(cmat)
+    model = eqx.tree_at(lambda m: m.readout, model, new_readout)
+
+    return model, tot_res_seq
+
+
+def train_ESNController(
+    model: ESNController,
+    train_seq: Array,
+    control_seq: Array,
+    target_seq: Array | None = None,
+    spinup: int = 0,
+    initial_res_state: Array | None = None,
+    beta: float = 8e-8,
+    multi_sequence: bool = False,
+) -> tuple[ESNController, Array]:
+    """Training function for ESNController.
+
+    Parameters
+    ----------
+    model : ESNController
+        ESNController model to train.
+    train_seq : Array
+        Training input sequence for reservoir, (shape=(seq_len, data_dim)).
+    control_seq : Array
+        Control input sequence for reservoir, (shape=(seq_len, control_dim)).
+    target_seq : Array
+        Target sequence for training reservoir, (shape=(seq_len, data_dim)).
+        If None, defaults to train_seq[1:].
+    initial_res_state : Array
+        Initial reservoir state, (shape=(res_dim,)).
+    spinup : int
+        Initial transient of reservoir states to discard.
+    beta : float
+        Tikhonov regularization parameter.
+    multi_sequence : bool
+        If True, treat train_seq and control_seq as batches of trajectories
+        with shape (n_traj, seq_len, ...). See train_RCController for details.
+
+    Returns
+    -------
+    model : ESNController
+        Trained ESN controller model.
+    res_seq : Array
+        Training sequence of reservoir states.
+    """
+    if not isinstance(model, ESNController):
+        raise TypeError("Model must be an ESNController.")
+
+    return train_RCController(
+        model,
+        train_seq,
+        control_seq,
+        target_seq,
+        spinup,
+        initial_res_state,
+        beta,
+        multi_sequence=multi_sequence,
+    )

{openreservoircomputing-0.3.2 → openreservoircomputing-0.3.3}/src/orc/drivers.py

@@ -869,9 +869,7 @@ class GRUDriver(DriverBase):
         super().__init__(res_dim=res_dim)
         key = jax.random.key(seed)
         gru = eqx.nn.GRUCell(res_dim, res_dim, key=key)
-        gru = eqx.tree_at(
-            lambda c: c.weight_ih, gru, gru.weight_ih * input_rescaling
-        )
+        gru = eqx.tree_at(lambda c: c.weight_ih, gru, gru.weight_ih * input_rescaling)
         self.gru = gru
 
     def advance(self, proj_vars, res_state):

{openreservoircomputing-0.3.2 → openreservoircomputing-0.3.3}/src/orc/forecaster/base.py

@@ -391,7 +391,7 @@ class CRCForecasterBase(RCForecasterBase, ABC):
         res_seq = sol.ys
         return res_seq
 
-    def __call__(self, in_seq: Array, res_state: Array, ts: Array) -> Array: # ty: ignore[invalid-method-override]
+    def __call__(self, in_seq: Array, res_state: Array, ts: Array) -> Array:  # ty: ignore[invalid-method-override]
         """
         Teacher forces the reservoir, wrapper for `force` method.