OpenReservoirComputing 0.3.2__tar.gz → 0.4.0__tar.gz

{openreservoircomputing-0.3.2 → openreservoircomputing-0.4.0}/.gitignore

@@ -97,10 +97,8 @@ ipython_config.py
 #Pipfile.lock
 
 # UV
-#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
-#   This is especially recommended for binary packages to ensure reproducibility, and is more
-#   commonly ignored for libraries.
-#uv.lock
+#   uv.lock IS COMMITTED to this repository on purpose: it gives contributors reproducible
+#   environments and CI runs `uv sync --locked` against it. Do not add it here.
 
 # poetry
 #   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.

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

@@ -1,63 +1,29 @@
 Metadata-Version: 2.4
 Name: OpenReservoirComputing
-Version: 0.3.2
+Version: 0.4.0
 Summary: GPU accelerated implementations of common RC architectures
+Project-URL: Repository, https://github.com/Jan-Williams/OpenReservoirComputing
 Author-email: "Jan P. Williams" <jmpw1@uw.edu>, Dima Tretiak <dtretiak@uw.edu>
 Maintainer-email: "Jan P. Williams" <jmpw1@uw.edu>
-Project-URL: Repository, https://github.com/Jan-Williams/OpenReservoirComputing
+License-File: LICENSE
 Keywords: reservoir computing
-Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Requires-Python: <3.14,>=3.10
-Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: jax
+Requires-Dist: diffrax
 Requires-Dist: equinox
+Requires-Dist: jax
 Requires-Dist: matplotlib
-Requires-Dist: setuptools_scm>=8.1
-Requires-Dist: diffrax
-Provides-Extra: all
-Requires-Dist: jax[cuda12]; extra == "all"
-Requires-Dist: ruff; extra == "all"
-Requires-Dist: ty; extra == "all"
-Requires-Dist: pytest; extra == "all"
-Requires-Dist: coverage; extra == "all"
-Requires-Dist: pytest-cov; extra == "all"
-Requires-Dist: notebook; extra == "all"
-Requires-Dist: ipykernel; extra == "all"
-Requires-Dist: pytest-env; extra == "all"
-Requires-Dist: mkdocs; extra == "all"
-Requires-Dist: mkdocs-material; extra == "all"
-Requires-Dist: mkdocs-autorefs; extra == "all"
-Requires-Dist: mkdocstrings[python]; extra == "all"
-Requires-Dist: mkdocs-jupyter; extra == "all"
-Provides-Extra: dev
-Requires-Dist: ruff; extra == "dev"
-Requires-Dist: ty; extra == "dev"
-Requires-Dist: pytest; extra == "dev"
-Requires-Dist: coverage; extra == "dev"
-Requires-Dist: pytest-cov; extra == "dev"
-Requires-Dist: pytest-env; extra == "dev"
-Requires-Dist: mkdocs; extra == "dev"
-Requires-Dist: mkdocs-material; extra == "dev"
-Requires-Dist: mkdocs-autorefs; extra == "dev"
-Requires-Dist: mkdocstrings[python]; extra == "dev"
-Provides-Extra: notebooks
-Requires-Dist: notebook; extra == "notebooks"
-Requires-Dist: ipykernel; extra == "notebooks"
 Provides-Extra: gpu
-Requires-Dist: jax[cuda12]; extra == "gpu"
-Provides-Extra: docs
-Requires-Dist: mkdocs; extra == "docs"
-Requires-Dist: mkdocs-material; extra == "docs"
-Requires-Dist: mkdocs-autorefs; extra == "docs"
-Requires-Dist: mkdocstrings[python]; extra == "docs"
-Requires-Dist: mkdocs-jupyter; extra == "docs"
-Dynamic: license-file
+Requires-Dist: jax[cuda12]; (sys_platform == 'linux') and extra == 'gpu'
+Provides-Extra: notebooks
+Requires-Dist: ipykernel; extra == 'notebooks'
+Requires-Dist: notebook; extra == 'notebooks'
+Description-Content-Type: text/markdown
 
 <div align="center">
   <img src="imgs/ORC_logo_cropped.png" alt="ORC Logo" width="200px" />
@@ -79,6 +45,11 @@ The easiest way to get started with ORC is to install from PyPI:
 pip install OpenReservoirComputing
 ```
 
+or, if you use [uv](https://docs.astral.sh/uv/):
+```bash
+uv add OpenReservoirComputing
+```
+
 If you're interested in the latest, unreleased version or in contributing, you can install from source. Please see the Contribution guidelines below for more details. 
 
 ## Quick start example
@@ -131,35 +102,41 @@ plt.show()
 ORC models are built on top of [Equinox](https://docs.kidger.site/equinox/), and as a result we strongly recommend the use of Equinox transforms `eqx.filter_{jit, vmap, grad}` over `jax.{jit, vmap, grad}`. For more details, please check out the [JAX JIT Compatibility](https://Jan-Williams.github.io/OpenReservoirComputing/examples/jit_compatibility/) example notebook.
 
 ## Contribution guidelines
-First off, thanks for helping out! We appreciate your willingness to contribute! To get started, clone the repo and install the developer dependencies of ORC.
+First off, thanks for helping out! We appreciate your willingness to contribute! ORC uses [uv](https://docs.astral.sh/uv/) to manage development environments, which keeps setup to a single command.
+
+First, [install uv](https://docs.astral.sh/uv/getting-started/installation/) if you don't have it. Then clone the repo and sync:
 
 ```bash
 git clone https://github.com/Jan-Williams/OpenReservoirComputing.git
+cd OpenReservoirComputing
+uv sync
 ```
 
-From the root directory of the repository, create an editable install for your given hardware.
+`uv sync` creates a `.venv`, installs ORC in editable mode, and installs the development dependencies at the exact versions recorded in `uv.lock`. There is no need to activate the environment — just prefix commands with `uv run`.
 
-CPU:
+For NVIDIA GPU support (Linux only — see the [installation docs](https://Jan-Williams.github.io/OpenReservoirComputing/getting-started/installation/) for why):
 ```bash
-pip install -e ".[dev]"
+uv sync --extra gpu
 ```
 
-GPU:
+To also install the documentation toolchain:
 ```bash
-pip install -e ".[dev, gpu]"
+uv sync --all-groups
 ```
 
 The main branch is protected from direct changes. If you would like to make a change please create a new branch and work on your new feature. After you are satisfied with your changes, please run our testing suite to ensure all is working well. We also expect new tests to be written for all changes if additions are made. The tests can be simply run from the root directory of the repository with
 ```bash
-pytest
+uv run pytest
 ```
 Followed by a formatting check 
 ```bash
-ruff check
+uv run ruff check
 ```
 and a type annotation check
 ```bash
-ty check
+uv run ty check src/
 ```
 
-Finally, submit your changes as a pull request! When you submit the PR, please request reviews from both @dtretiak and @Jan-Williams, we will try to get back to you as soon as possible. When you submit the PR, the above tests will automatically be run on your proposed changes through Github Actions, so it is best to get everything tested first before submitting!
+If you add or change a dependency, use `uv add <package>` (or `uv add --group dev <package>` for a development tool) and **commit the updated `uv.lock`** — CI runs `uv sync --locked` and will fail if the lockfile is out of date.
+
+Finally, submit your changes as a pull request! When you submit the PR, please request reviews from both @dtretiak and @Jan-Williams, we will try to get back to you as soon as possible. When you submit the PR, the above tests will automatically be run on your proposed changes through Github Actions, so it is best to get everything tested first before submitting!

{openreservoircomputing-0.3.2 → openreservoircomputing-0.4.0}/README.md

@@ -18,6 +18,11 @@ The easiest way to get started with ORC is to install from PyPI:
 pip install OpenReservoirComputing
 ```
 
+or, if you use [uv](https://docs.astral.sh/uv/):
+```bash
+uv add OpenReservoirComputing
+```
+
 If you're interested in the latest, unreleased version or in contributing, you can install from source. Please see the Contribution guidelines below for more details. 
 
 ## Quick start example
@@ -70,35 +75,41 @@ plt.show()
 ORC models are built on top of [Equinox](https://docs.kidger.site/equinox/), and as a result we strongly recommend the use of Equinox transforms `eqx.filter_{jit, vmap, grad}` over `jax.{jit, vmap, grad}`. For more details, please check out the [JAX JIT Compatibility](https://Jan-Williams.github.io/OpenReservoirComputing/examples/jit_compatibility/) example notebook.
 
 ## Contribution guidelines
-First off, thanks for helping out! We appreciate your willingness to contribute! To get started, clone the repo and install the developer dependencies of ORC.
+First off, thanks for helping out! We appreciate your willingness to contribute! ORC uses [uv](https://docs.astral.sh/uv/) to manage development environments, which keeps setup to a single command.
+
+First, [install uv](https://docs.astral.sh/uv/getting-started/installation/) if you don't have it. Then clone the repo and sync:
 
 ```bash
 git clone https://github.com/Jan-Williams/OpenReservoirComputing.git
+cd OpenReservoirComputing
+uv sync
 ```
 
-From the root directory of the repository, create an editable install for your given hardware.
+`uv sync` creates a `.venv`, installs ORC in editable mode, and installs the development dependencies at the exact versions recorded in `uv.lock`. There is no need to activate the environment — just prefix commands with `uv run`.
 
-CPU:
+For NVIDIA GPU support (Linux only — see the [installation docs](https://Jan-Williams.github.io/OpenReservoirComputing/getting-started/installation/) for why):
 ```bash
-pip install -e ".[dev]"
+uv sync --extra gpu
 ```
 
-GPU:
+To also install the documentation toolchain:
 ```bash
-pip install -e ".[dev, gpu]"
+uv sync --all-groups
 ```
 
 The main branch is protected from direct changes. If you would like to make a change please create a new branch and work on your new feature. After you are satisfied with your changes, please run our testing suite to ensure all is working well. We also expect new tests to be written for all changes if additions are made. The tests can be simply run from the root directory of the repository with
 ```bash
-pytest
+uv run pytest
 ```
 Followed by a formatting check 
 ```bash
-ruff check
+uv run ruff check
 ```
 and a type annotation check
 ```bash
-ty check
+uv run ty check src/
 ```
 
+If you add or change a dependency, use `uv add <package>` (or `uv add --group dev <package>` for a development tool) and **commit the updated `uv.lock`** — CI runs `uv sync --locked` and will fail if the lockfile is out of date.
+
 Finally, submit your changes as a pull request! When you submit the PR, please request reviews from both @dtretiak and @Jan-Williams, we will try to get back to you as soon as possible. When you submit the PR, the above tests will automatically be run on your proposed changes through Github Actions, so it is best to get everything tested first before submitting!

openreservoircomputing-0.4.0/pyproject.toml

@@ -0,0 +1,132 @@
+[build-system]
+requires = ["hatchling>=1.27", "hatch-vcs>=0.4"]
+build-backend = "hatchling.build"
+
+[project]
+name = "OpenReservoirComputing"
+authors = [
+    { name = "Jan P. Williams", email = "jmpw1@uw.edu" },
+    { name = "Dima Tretiak", email = "dtretiak@uw.edu" },
+]
+
+maintainers = [
+    { name = "Jan P. Williams", email = "jmpw1@uw.edu" },
+]
+
+description = "GPU accelerated implementations of common RC architectures"
+readme = "README.md"
+dynamic = ["version"]
+dependencies = [
+    "jax",
+    "equinox",
+    "matplotlib",
+    "diffrax",
+]
+
+requires-python = ">=3.10, <3.14"
+keywords = ["reservoir computing"]
+classifiers = [
+    "License :: OSI Approved :: Apache Software License",
+    "Intended Audience :: Science/Research",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+]
+
+[project.urls]
+Repository = "https://github.com/Jan-Williams/OpenReservoirComputing"
+
+[project.optional-dependencies]
+# JAX publishes CUDA wheels for Linux only (see JAX's installation docs). The marker
+# keeps this extra a no-op on macOS and native Windows instead of an install error.
+# Windows users with an NVIDIA GPU should install under WSL2, which reports
+# sys_platform == "linux" and therefore does get the CUDA build.
+gpu = ["jax[cuda12]; sys_platform == 'linux'"]
+notebooks = ["notebook", "ipykernel"]
+
+# Development and documentation dependencies live in PEP 735 dependency groups rather
+# than extras: they are never published in the package metadata, and `uv sync` installs
+# them automatically. They are therefore NOT reachable via `pip install -e ".[dev]"`.
+[dependency-groups]
+dev = [
+    "ruff",
+    "ty",
+    "pytest",
+    "pytest-cov",
+    "pytest-env",
+    "coverage",
+]
+docs = [
+    "mkdocs",
+    "mkdocs-material",
+    "mkdocs-autorefs",
+    "mkdocstrings[python]",
+    "mkdocs-jupyter",
+    # mkdocs.yml runs mkdocs-jupyter with execute: true / allow_errors: false, so the
+    # docs build executes every example notebook and needs a Jupyter kernel.
+    "notebook",
+    "ipykernel",
+]
+
+[tool.uv]
+default-groups = ["dev"]
+required-version = ">=0.9.0"
+
+[tool.hatch.version]
+source = "vcs"
+
+[tool.hatch.build.targets.wheel]
+# Required explicitly: hatchling's autodetection looks for the normalized project name
+# ("openreservoircomputing"), which does not match the import package ("orc").
+packages = ["src/orc"]
+
+[tool.hatch.build.targets.sdist]
+# only-include (not include): docs/examples is a symlink to ../examples, and hatchling's
+# tree walk follows symlinks while de-duplicating on (st_dev, st_ino) -- so whichever of
+# examples/ or docs/examples/ is visited first wins and the other is dropped, making the
+# sdist non-deterministic across machines. only-include bypasses the tree walk entirely.
+only-include = [
+    "src/orc",
+    "tests",
+    "examples",
+    "pyproject.toml",
+    "README.md",
+    "LICENSE",
+]
+
+[tool.pytest.ini_options]
+testpaths = "tests"
+addopts = "--ignore=scripts --ignore=docs"
+env = ["MPLBACKEND=Agg"]
+
+[tool.coverage.run]
+source_pkgs = ["orc"]
+branch = true
+
+[tool.coverage.paths]
+# Canonicalize editable (src/orc) and non-editable (site-packages/orc) layouts so
+# reports are comparable regardless of how the project was installed.
+source = ["src/orc", "*/site-packages/orc"]
+
+[tool.ruff]
+exclude = ["examples"]
+
+[tool.ruff.lint]
+select = ["E", "W", "F", "D", "UP", "B", "A", "C4", "N", "SIM", "I"]
+
+ignore = ["N806", "N802", "N803", "F722"]
+
+[tool.ruff.lint.per-file-ignores]
+"tests/*" = ["D", "E402"]
+"docs/*.py" = ["D", "A"]
+"*.ipynb" = ["E", "W", "F", "D", "UP", "B", "A", "C4", "N", "SIM", "I"]
+
+[tool.ruff.lint.pycodestyle]
+max-line-length = 88
+
+[tool.ruff.lint.pydocstyle]
+convention = "numpy"
+
+[tool.ty.src]
+exclude = ["tests", "examples"]

{openreservoircomputing-0.3.2 → openreservoircomputing-0.4.0}/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.4.0}/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.4.0}/src/orc/control/base.py

@@ -179,9 +179,7 @@ class RCControllerBase(eqx.Module, ABC):
         return self.force(in_seq, control_seq, res_state)
 
     @eqx.filter_jit
-    def apply_control(
-        self, control_seq: Array, res_state: Array
-    ) -> tuple[Array, Array]:
+    def apply_control(self, control_seq: Array, res_state: Array) -> Array:
         """Apply a predefined control sequence in closed-loop.
 
         The readout feeds back as the next input: u(t+1) = readout(x(t)).

{openreservoircomputing-0.3.2 → openreservoircomputing-0.4.0}/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.4.0/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.4.0}/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.4.0}/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.