OpenReservoirComputing 0.1.3a0__tar.gz → 0.2.0__tar.gz

{openreservoircomputing-0.1.3a0/src/OpenReservoirComputing.egg-info → openreservoircomputing-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: OpenReservoirComputing
-Version: 0.1.3a0
+Version: 0.2.0
 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>
@@ -111,7 +111,7 @@ t_test = t[split_idx:]
 
 # Initialize and train the ESN
 esn = orc.forecaster.ESNForecaster(data_dim=3, res_dim=400)
-esn, R = orc.forecaster.train_ESNForecaster(esn, U_train)
+esn, R = orc.forecaster.train_RCForecaster(esn, U_train)
 
 # Forecast! 
 U_pred = esn.forecast(fcast_len=U_test.shape[0], res_state=R[-1]) # feed in the last reservoir state seen in training

{openreservoircomputing-0.1.3a0 → openreservoircomputing-0.2.0}/README.md

@@ -52,7 +52,7 @@ t_test = t[split_idx:]
 
 # Initialize and train the ESN
 esn = orc.forecaster.ESNForecaster(data_dim=3, res_dim=400)
-esn, R = orc.forecaster.train_ESNForecaster(esn, U_train)
+esn, R = orc.forecaster.train_RCForecaster(esn, U_train)
 
 # Forecast! 
 U_pred = esn.forecast(fcast_len=U_test.shape[0], res_state=R[-1]) # feed in the last reservoir state seen in training

{openreservoircomputing-0.1.3a0 → openreservoircomputing-0.2.0}/docs/api/classifier.md

@@ -20,6 +20,11 @@ The `orc.classifier` subpackage provides reservoir computing models for time ser
 
 ## Training Functions
 
+::: orc.classifier.train.train_RCClassifier
+    options:
+      show_root_heading: true
+      show_source: true
+
 ::: orc.classifier.train.train_ESNClassifier
     options:
       show_root_heading: true

{openreservoircomputing-0.1.3a0 → openreservoircomputing-0.2.0}/docs/api/control.md

@@ -20,6 +20,11 @@ The `orc.control` subpackage provides reservoir computing models for dynamical s
 
 ## Training Functions
 
+::: orc.control.train.train_RCController
+    options:
+      show_root_heading: true
+      show_source: true
+
 ::: orc.control.train.train_ESNController
     options:
       show_root_heading: true

{openreservoircomputing-0.1.3a0 → openreservoircomputing-0.2.0}/docs/api/forecaster.md

@@ -38,6 +38,11 @@ The `orc.forecaster` subpackage provides reservoir computing models for time ser
 
 ## Training Functions
 
+::: orc.forecaster.train.train_RCForecaster
+    options:
+      show_root_heading: true
+      show_source: true
+
 ::: orc.forecaster.train.train_ESNForecaster
     options:
       show_root_heading: true

{openreservoircomputing-0.1.3a0 → openreservoircomputing-0.2.0}/docs/examples/classification.ipynb

@@ -207,7 +207,7 @@
     }
    ],
    "source": [
-    "trained_classifier = orc.classifier.train_ESNClassifier(\n",
+    "trained_classifier = orc.classifier.train_RCClassifier(\n",
     "    classifier,\n",
     "    train_seqs=train_seqs,\n",
     "    labels=train_labels,\n",
@@ -245,7 +245,7 @@
     "probs = jax.vmap(trained_classifier.classify)(train_seqs)\n",
     "\n",
     "\n",
-    "train_preds = jnp.array(jnp.argmax(probs, axis=1))\n",
+    "train_preds = jnp.argmax(probs, axis=1)\n",
     "train_acc = jnp.mean(train_preds == train_labels)\n",
     "print(f\"Training accuracy: {train_acc:.1%}\")"
    ]
@@ -277,7 +277,7 @@
    "source": [
     "probs = jax.vmap(trained_classifier.classify)(test_seqs)\n",
     "\n",
-    "test_preds = jnp.array(jnp.argmax(probs, axis=1))\n",
+    "test_preds = jnp.argmax(probs, axis=1)\n",
     "test_acc = jnp.mean(test_preds == test_labels)\n",
     "print(f\"Test accuracy: {test_acc:.1%}\")"
    ]
@@ -356,7 +356,7 @@
     "    state_repr=\"mean\",\n",
     ")\n",
     "\n",
-    "trained_mean = orc.classifier.train_ESNClassifier(\n",
+    "trained_mean = orc.classifier.train_RCClassifier(\n",
     "    classifier_mean,\n",
     "    train_seqs=train_seqs,\n",
     "    labels=train_labels,\n",

{openreservoircomputing-0.1.3a0 → openreservoircomputing-0.2.0}/docs/examples/continuous_rc.ipynb

@@ -98,18 +98,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 4,
    "id": "13d6be2b",
    "metadata": {},
    "outputs": [],
    "source": [
     "esn = orc.forecaster.CESNForecaster(data_dim=U_train.shape[1], res_dim=1000, time_const=40.0)\n",
-    "esn, R = orc.forecaster.train_CESNForecaster(model=esn, train_seq=U_train, t_train=t_train)"
+    "esn, R = orc.forecaster.train_RCForecaster(model=esn, train_seq=U_train, ts=t_train)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 5,
    "id": "6d3630b8",
    "metadata": {},
    "outputs": [
@@ -146,7 +146,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 6,
    "id": "8e549270",
    "metadata": {},
    "outputs": [],
@@ -159,7 +159,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 7,
    "id": "9aaa840e",
    "metadata": {},
    "outputs": [
@@ -196,7 +196,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 8,
    "id": "ad1c9f18",
    "metadata": {},
    "outputs": [
@@ -226,7 +226,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 9,
    "id": "618df66f",
    "metadata": {},
    "outputs": [

{openreservoircomputing-0.1.3a0 → openreservoircomputing-0.2.0}/docs/examples/control.ipynb

@@ -190,13 +190,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 4,
    "id": "65a44d34",
    "metadata": {},
    "outputs": [],
    "source": [
     "control_model = orc.control.ESNController(data_dim=2, control_dim=1, res_dim=1000)\n",
-    "control_model, R = orc.control.train_ESNController(control_model, outputs_damped[:-1], u[1:])"
+    "control_model, R = orc.control.train_RCController(control_model, outputs_damped[:-1], u[1:])"
    ]
   },
   {
@@ -217,7 +217,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 5,
    "id": "28b6b864",
    "metadata": {},
    "outputs": [],
@@ -257,7 +257,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 6,
    "id": "39b5601a",
    "metadata": {},
    "outputs": [],
@@ -267,7 +267,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 7,
    "id": "be2c5103",
    "metadata": {},
    "outputs": [

{openreservoircomputing-0.1.3a0 → openreservoircomputing-0.2.0}/docs/examples/ks.ipynb

@@ -164,7 +164,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "esn, R = orc.forecaster.train_ESNForecaster(\n",
+    "esn, R = orc.forecaster.train_RCForecaster(\n",
     "    esn,\n",
     "    U_train,\n",
     "    beta=beta\n",

{openreservoircomputing-0.1.3a0 → openreservoircomputing-0.2.0/docs}/examples/lorenz.ipynb

@@ -143,7 +143,7 @@
     "# init + train ESN\n",
     "NR = 1000\n",
     "esn = orc.forecaster.ESNForecaster(data_dim=3, res_dim=NR, seed=0)\n",
-    "esn, R = orc.forecaster.train_ESNForecaster(esn, U_train)"
+    "esn, R = orc.forecaster.train_RCForecaster(esn, U_train)"
    ]
   },
   {

{openreservoircomputing-0.1.3a0 → openreservoircomputing-0.2.0/docs}/examples/rc_background.ipynb

@@ -163,14 +163,14 @@
     "class Readout(orc.readouts.ReadoutBase):\n",
     "    res_dim: int\n",
     "    out_dim: int\n",
-    "    W_O: jnp.ndarray\n",
+    "    wout: jnp.ndarray\n",
     "    \n",
     "    def __init__(self, out_dim, res_dim):\n",
     "        super().__init__(out_dim, res_dim)\n",
-    "        self.W_O = jnp.zeros((out_dim, res_dim))\n",
+    "        self.wout = jnp.zeros((out_dim, res_dim))\n",
     "    \n",
     "    def readout(self, res_state):\n",
-    "        return self.W_O @ res_state"
+    "        return self.wout @ res_state"
    ]
   },
   {
@@ -299,7 +299,7 @@
     "# set transient to discard\n",
     "spinup = 200\n",
     "\n",
-    "# learn W_O\n",
+    "# learn wout\n",
     "readout_mat = ridge_regression(forced_seq[spinup:], target_seq[spinup:], beta=1e-7)"
    ]
   },
@@ -318,10 +318,20 @@
    "source": [
     "# define where in the forecaster model we need to update\n",
     "def where(model: Forecaster):\n",
-    "    return model.readout.W_O\n",
+    "    return model.readout.wout\n",
     "model = eqx.tree_at(where, model, readout_mat)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Alternatively\n",
+    "model, forced_seq = orc.forecaster.train_RCForecaster(model, train_seq=U_train, spinup=spinup, beta=1e-7)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -331,7 +341,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [
     {

{openreservoircomputing-0.1.3a0 → openreservoircomputing-0.2.0}/docs/user-guide/drivers.md

@@ -12,12 +12,16 @@ The `DriverBase` class defines the core interface that all driver implementation
 
 ### Key Attributes
 - `res_dim`: Reservoir dimension
+- `chunks`: Number of parallel reservoirs (default `0`, meaning no chunks dimension)
 - `dtype`: JAX array dtype (jnp.float32 or jnp.float64)
 
 ### Core Methods
 - `advance(proj_vars, res_state)`: Advance single reservoir state by one time step
 - `batch_advance(proj_vars, res_state)`: Advance batch of reservoir states by one time step
 - `__call__(proj_vars, res_state)`: Flexible interface supporting both single and batch updates
+- `default_state()`: Create a zero initial reservoir state with the appropriate shape —
+  `(chunks, res_dim)` when `chunks > 0`, or `(res_dim,)` for single-reservoir drivers.
+  Used by the unified training functions when no `initial_res_state` is provided.
 
 ## ParallelESNDriver
 
@@ -142,6 +146,11 @@ class CustomDriver(DriverBase):
 ### Requirements
 - Inherit from `DriverBase`
 - Implement the abstract `advance()` method
+- Override `default_state()` if your driver uses a non-standard state shape (e.g., a single
+  reservoir driver should return `jnp.zeros((res_dim,), dtype=self.dtype)`). The base class
+  implementation returns `(chunks, res_dim)` when `chunks > 0` and `(res_dim,)` otherwise.
+  The unified training functions (`train_RCForecaster`, etc.) call this method to construct
+  the initial reservoir state when none is provided by the caller.
 - Choose your parallel reservoir support level:
 
 #### Single Reservoir Only

{openreservoircomputing-0.1.3a0 → openreservoircomputing-0.2.0}/docs/user-guide/models.md

@@ -176,9 +176,116 @@ model = CESNForecaster(
 )
 ```
 
+## ESN Controller
+
+The `ESNController` implements an Echo State Network for control tasks, where the model tracks a reference trajectory by optimizing a control input sequence.
+
+### Basic Usage
+
+```python
+from orc.control import ESNController, train_ESNController
+import jax.numpy as jnp
+
+# Create ESN controller
+model = ESNController(
+    data_dim=3,           # System input/output dimension
+    control_dim=1,        # Control input dimension
+    res_dim=500,          # Reservoir dimension
+    seed=42
+)
+
+# Train with input and control sequences
+trained_model, res_states = train_ESNController(
+    model,
+    train_seq,            # shape: (seq_len, data_dim)
+    control_seq,          # shape: (seq_len, control_dim)
+    target_seq=target_seq,
+    spinup=100,
+    beta=8e-8
+)
+
+# Compute optimal control to track a reference trajectory
+control_opt = trained_model.compute_control(
+    control_seq=initial_guess,   # shape: (fcast_len, control_dim)
+    res_state=res_states[-1],
+    ref_traj=reference           # shape: (fcast_len, data_dim)
+)
+```
+
+### Control Penalty Weights
+
+The controller optimizes a penalty function with three terms, each weighted by a tunable parameter:
+
+- `alpha_1` (default 100): **Trajectory deviation** — penalizes squared error between the controlled forecast and the reference trajectory. Higher values enforce tighter tracking.
+
+- `alpha_2` (default 1): **Control magnitude** — penalizes the squared norm of control inputs. Higher values favor smaller control effort.
+
+- `alpha_3` (default 5): **Control smoothness** — penalizes the squared norm of consecutive control input differences. Higher values produce smoother control signals.
+
+```python
+# Custom penalty weights
+model = ESNController(
+    data_dim=3,
+    control_dim=1,
+    res_dim=500,
+    alpha_1=200,    # Stricter trajectory tracking
+    alpha_2=0.5,    # Allow larger control inputs
+    alpha_3=10,     # Enforce smoother control
+    seed=42
+)
+```
+
+### Key Parameters
+
+In addition to the standard ESN parameters (`leak_rate`, `bias`, `embedding_scaling`, `Wr_density`, `Wr_spectral_radius`), the controller accepts:
+
+- `control_dim`: Dimension of the control input
+- `alpha_1`, `alpha_2`, `alpha_3`: Control penalty weights (see above)
+
+For detailed API documentation, see the [Control API Reference](../api/control.md).
+
 ## Training Functions
 
-Both ESN models use ridge regression for training the readout layer.
+All models use ridge regression for training the readout layer. ORC provides unified training functions that work with any model inheriting from the base classes, as well as convenience wrappers for each specific model type.
+
+### Unified Training Functions (Recommended)
+
+The primary training API uses generic functions that work with any `RCForecasterBase`, `RCClassifierBase`, or `RCControllerBase` subclass, including custom models:
+
+```python
+from orc.forecaster import train_RCForecaster
+from orc.classifier import train_RCClassifier
+from orc.control import train_RCController
+
+# Forecaster: works with ESNForecaster, CESNForecaster, EnsembleESNForecaster,
+# and any custom RCForecasterBase subclass
+trained_model, res_states = train_RCForecaster(model, train_seq, target_seq=target_seq)
+
+# Classifier: works with ESNClassifier and custom RCClassifierBase subclasses
+trained_model = train_RCClassifier(model, train_seqs, labels)
+
+# Controller: works with ESNController and custom RCControllerBase subclasses
+trained_model, res_states = train_RCController(model, train_seq, control_seq)
+```
+
+For continuous-time models, pass the time vector via `**force_kwargs`:
+
+```python
+# CESNForecaster via the unified function
+trained_model, res_states = train_RCForecaster(
+    model, train_seq, target_seq=target_seq, ts=t_train
+)
+```
+
+### Model-Specific Convenience Functions
+
+The legacy model-specific functions remain available and delegate to the unified functions above:
+
+- `train_ESNForecaster` — equivalent to `train_RCForecaster` with an `ESNForecaster` type check
+- `train_CESNForecaster` — equivalent to `train_RCForecaster(..., ts=t_train)` with a `CESNForecaster` type check
+- `train_EnsembleESNForecaster` — equivalent to `train_RCForecaster` with an `EnsembleESNForecaster` type check
+- `train_ESNClassifier` — equivalent to `train_RCClassifier` with an `ESNClassifier` type check
+- `train_ESNController` — equivalent to `train_RCController` with an `ESNController` type check
 
 ### Training Parameters
 
@@ -187,19 +294,22 @@ Both ESN models use ridge regression for training the readout layer.
 
 - `beta`: Tikhonov regularization parameter (typically 1e-8 to 1e-6)
 
-- `initial_res_state`: Custom initial reservoir state (default: zeros)
+- `initial_res_state`: Custom initial reservoir state (default: uses `model.driver.default_state()`)
 
 **Memory Management:**
-- `batch_size`: Process parallel reservoirs in batches to reduce memory usage
+- `batch_size`: Process parallel reservoirs in batches to reduce memory usage (available for all three unified functions)
+
+**Extra Force Arguments:**
+- `**force_kwargs`: Additional keyword arguments forwarded to `model.force()` (e.g., `ts=t_train` for continuous models)
 
 ### Advanced Training
 
 ```python
 # Batched training for memory efficiency with many parallel reservoirs
-trained_model, res_states = train_ESNForecaster(
+trained_model, res_states = train_RCForecaster(
     model,
     train_seq,
-    target_seq,
+    target_seq=target_seq,
     spinup=100,
     beta=1e-7,
     batch_size=10    # Process 10 parallel reservoirs at a time
@@ -249,9 +359,9 @@ The models module demonstrates the power of ORC's modular design:
 To create your own model, follow the ESN pattern:
 
 ```python
-from orc.rc import RCForecasterBase
+from orc.forecaster import RCForecasterBase, train_RCForecaster
 from orc.drivers import CustomDriver
-from orc.embeddings import CustomEmbedding  
+from orc.embeddings import CustomEmbedding
 from orc.readouts import CustomReadout
 
 class CustomForecaster(RCForecasterBase):
@@ -260,11 +370,22 @@ class CustomForecaster(RCForecasterBase):
         driver = CustomDriver(res_dim=res_dim, **kwargs)
         embedding = CustomEmbedding(in_dim=data_dim, res_dim=res_dim, **kwargs)
         readout = CustomReadout(out_dim=data_dim, res_dim=res_dim, **kwargs)
-        
+
         # Initialize base class
         super().__init__(driver, readout, embedding, **kwargs)
+
+# Train using the unified function — no model-specific training code needed
+trained_model, res_states = train_RCForecaster(
+    CustomForecaster(data_dim=3, res_dim=500),
+    train_seq,
+    target_seq=target_seq,
+    spinup=100,
+    beta=8e-8,
+)
 ```
 
+The unified `train_RCForecaster` function delegates shape handling to the readout's `prepare_train`, `prepare_target`, and `set_wout` methods and the driver's `default_state` method, so custom components that implement those methods integrate automatically with training. See the [Readouts](readouts.md) and [Drivers](drivers.md) user guides for details.
+
 ## Performance Considerations
 
 - Use `jnp.float64` for numerical precision (default)

{openreservoircomputing-0.1.3a0 → openreservoircomputing-0.2.0}/docs/user-guide/readouts.md

@@ -13,13 +13,26 @@ The `ReadoutBase` class defines the core interface that all readout implementati
 ### Key Attributes
 - `out_dim`: Output dimension
 - `res_dim`: Reservoir dimension
+- `chunks`: Number of parallel reservoirs (default `0`, meaning no chunks dimension)
 - `dtype`: JAX array dtype (jnp.float32 or jnp.float64)
 
 ### Core Methods
+
+**Inference:**
 - `readout(res_state)`: Map a single reservoir state to output
-- `batch_readout(res_state)`: Map multiple reservoir states to outputs efficiently  
+- `batch_readout(res_state)`: Map multiple reservoir states to outputs efficiently
 - `__call__(res_state)`: Flexible interface supporting both single and batch inputs
 
+**Training interface** (used by `train_RCForecaster`, `train_RCClassifier`, `train_RCController`):
+- `prepare_train(res_seq)`: Prepare reservoir states for ridge regression. Default: identity
+  (returns `res_seq` unchanged). Override to apply nonlinear transforms or reshape for
+  chunk structure.
+- `prepare_target(target_seq)`: Prepare target sequence for ridge regression. Default:
+  identity. Override to reshape for chunk structure (e.g., adding or repeating the chunks axis).
+- `set_wout(cmat)`: Return a **new** readout instance with the updated weight matrix `cmat`
+  stored in `wout`. Override if your readout uses a different attribute name for the output
+  weights.
+
 ## ParallelLinearReadout Example
 
 The `ParallelLinearReadout` class demonstrates a standard linear transformation from reservoir states to outputs with support for parallel reservoirs.
@@ -72,7 +85,7 @@ The readout weights (`wout`) are initialized to zero and trained externally usin
 
 ## Custom Readouts
 
-The primary purpose of the readout system is to enable you to create your own readout layers tailored to your specific needs. While RC architectures often are trained via regression, this also provies more flexibility for more complicated readout architectures and training protocols.
+The primary purpose of the readout system is to enable you to create your own readout layers tailored to your specific needs. While RC architectures often are trained via regression, this also provides more flexibility for more complicated readout architectures and training protocols.
 
 ```python
 from orc.readouts import ReadoutBase
@@ -80,26 +93,39 @@ import equinox as eqx
 import jax.numpy as jnp
 
 class CustomReadout(ReadoutBase):
-    # Define your parameters
+    # Define your parameters — `chunks` is inherited from ReadoutBase
     wout: Array
-    chunks: int
-    
+
     def __init__(self, out_dim, res_dim, chunks=1, **kwargs):
-        super().__init__(out_dim, res_dim)
-        self.chunks = chunks
-        # Initialize your parameters (e.g., weights, nonlinearities)
-        self.wout = jnp.zeros((chunks, int(out_dim/chunks), res_dim))
-    
+        super().__init__(out_dim, res_dim, chunks=chunks)
+        # Initialize weights
+        self.wout = jnp.zeros((chunks, int(out_dim / chunks), res_dim))
+
     def readout(self, res_state):
         # Implement your custom readout logic
-        # Apply transformations, nonlinearities, etc.
         transformed_state = your_transformation(res_state)
         return jnp.ravel(eqx.filter_vmap(jnp.matmul)(self.wout, transformed_state))
+
+    # Training interface — override to integrate with train_RCForecaster etc.
+    def prepare_train(self, res_seq):
+        # Apply any custom transform and/or reshape for chunk structure
+        return your_transformation(res_seq)  # shape: (seq_len, chunks, res_dim)
+
+    def prepare_target(self, target_seq):
+        # Reshape target to match chunks structure expected by ridge regression
+        return target_seq.reshape(target_seq.shape[0], self.chunks, -1)
+
+    # set_wout is inherited and works as long as the weight attribute is named `wout`
 ```
 
 ### Requirements for Custom Readouts
 - Inherit from `ReadoutBase`
 - Implement the abstract `readout()` method
+- Do **not** redefine `chunks` as a class attribute — it is inherited from `ReadoutBase`
+- Override `prepare_train`, `prepare_target`, and/or `set_wout` to integrate with the
+  unified training functions (`train_RCForecaster`, `train_RCClassifier`,
+  `train_RCController`). The base class defaults (identity transforms) are sufficient
+  for simple single-reservoir readouts with a `wout` attribute.
 - Choose your parallel reservoir support level:
 
 #### Single Reservoir Only
@@ -137,10 +163,19 @@ def __call__(self, res_state):
 ## Design Considerations
 
 ### Training Integration
-- Readout weights are learned separately via ridge regression (see training utilities)
-- The readout layer structure determines the feature space for regression
-- More complex readouts can capture richer output relationships but may require more training data or different training procedures.
 
+The unified training functions (`train_RCForecaster`, `train_RCClassifier`,
+`train_RCController`) interact with readout layers entirely through the training interface:
+
+1. `prepare_train(res_seq)` — called to transform/reshape reservoir states before ridge
+   regression
+2. `prepare_target(target_seq)` — called to reshape targets to match the expected ridge
+   regression input
+3. `set_wout(cmat)` — called after solving regression to install the learned weights
+
+Custom readouts that override these three methods integrate automatically with training
+without any changes to the training code. Readouts whose `chunks > 0` receive the batched
+ridge regression path; readouts with `chunks == 0` use the standard scalar ridge regression.
 
 ### Integration with RC Models
 - Readout layers receive reservoir states from drivers

{openreservoircomputing-0.1.3a0 → openreservoircomputing-0.2.0}/examples/classification.ipynb

@@ -207,7 +207,7 @@
     }
    ],
    "source": [
-    "trained_classifier = orc.classifier.train_ESNClassifier(\n",
+    "trained_classifier = orc.classifier.train_RCClassifier(\n",
     "    classifier,\n",
     "    train_seqs=train_seqs,\n",
     "    labels=train_labels,\n",
@@ -245,7 +245,7 @@
     "probs = jax.vmap(trained_classifier.classify)(train_seqs)\n",
     "\n",
     "\n",
-    "train_preds = jnp.array(jnp.argmax(probs, axis=1))\n",
+    "train_preds = jnp.argmax(probs, axis=1)\n",
     "train_acc = jnp.mean(train_preds == train_labels)\n",
     "print(f\"Training accuracy: {train_acc:.1%}\")"
    ]
@@ -277,7 +277,7 @@
    "source": [
     "probs = jax.vmap(trained_classifier.classify)(test_seqs)\n",
     "\n",
-    "test_preds = jnp.array(jnp.argmax(probs, axis=1))\n",
+    "test_preds = jnp.argmax(probs, axis=1)\n",
     "test_acc = jnp.mean(test_preds == test_labels)\n",
     "print(f\"Test accuracy: {test_acc:.1%}\")"
    ]
@@ -356,7 +356,7 @@
     "    state_repr=\"mean\",\n",
     ")\n",
     "\n",
-    "trained_mean = orc.classifier.train_ESNClassifier(\n",
+    "trained_mean = orc.classifier.train_RCClassifier(\n",
     "    classifier_mean,\n",
     "    train_seqs=train_seqs,\n",
     "    labels=train_labels,\n",