Multi time series training

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"]
 

src/orc/__init__.py

@@ -17,7 +17,7 @@
 
 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,

src/orc/classifier/train.py

@@ -17,6 +17,7 @@
 
 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)

src/orc/control/models.py

@@ -152,5 +152,5 @@ def __init__(
             seed=seed,
             alpha_1=alpha_1,
             alpha_2=alpha_2,
-            alpha_3=alpha_3
+            alpha_3=alpha_3,
         )

src/orc/control/train.py

@@ -26,6 +26,7 @@ def train_RCController(
     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.
@@ -38,22 +39,37 @@ def train_RCController(
     model : RCControllerBase
         Reservoir computer controller model to train.
     train_seq : Array
-        Training input sequence for reservoir, (shape=(seq_len, data_dim)).
+        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)).
+        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)).
-        If None, defaults to train_seq[1:].
+        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.
+        Initial transient of reservoir states to discard per trajectory.
     initial_res_state : Array, optional
-        Initial reservoir state. If None, uses model.driver.default_state().
+        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().
 
@@ -62,57 +78,79 @@ def train_RCController(
     model : RCControllerBase
         Trained model.
     tot_res_seq : Array
-        Full sequence of reservoir states from teacher forcing.
+        Full reservoir state sequence from teacher forcing. Shape
+        ``(seq_len, ...)`` for single-sequence or ``(n_traj, seq_len, ...)``
+        for multi-sequence.
     """
-    if train_seq.shape[0] != control_seq.shape[0]:
-        raise ValueError("train_seq and control_seq must have the same length.")
+    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.")
 
-    if spinup >= train_seq.shape[0]:
+    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:
-        initial_res_state = model.driver.default_state()
+        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:, :]
-        train_seq = train_seq[:-1, :]
-        control_seq = control_seq[:-1, :]
+        target_seq = train_seq[:, 1:, :] if multi_sequence else train_seq[1:, :]
     else:
-        tot_seq = jnp.vstack((train_seq, target_seq[-1:]))
-        tot_control_seq = control_seq
-
-    tot_res_seq = model.force(
-        tot_seq, tot_control_seq, initial_res_state, **force_kwargs
-    )
-    res_seq = tot_res_seq[:-1]
+        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:]))
 
-    res_seq_train = model.readout.prepare_train(res_seq)
-    train_target = model.readout.prepare_target(target_seq)
+    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[spinup:],
-                train_target[spinup:],
-                beta,
-            )
+            cmat = _solve_all_ridge_reg(res_seq_train, train_target, beta)
         else:
             cmat = _solve_all_ridge_reg_batched(
-                res_seq_train[spinup:],
-                train_target[spinup:],
-                beta,
-                batch_size,
+                res_seq_train, train_target, beta, batch_size
             )
     else:
-        cmat = ridge_regression(
-            res_seq_train[spinup:],
-            train_target[spinup:],
-            beta,
-        )
+        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)
@@ -128,6 +166,7 @@ def train_ESNController(
     spinup: int = 0,
     initial_res_state: Array | None = None,
     beta: float = 8e-8,
+    multi_sequence: bool = False,
 ) -> tuple[ESNController, Array]:
     """Training function for ESNController.
 
@@ -148,6 +187,9 @@ def train_ESNController(
         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
     -------
@@ -160,5 +202,12 @@ def train_ESNController(
         raise TypeError("Model must be an ESNController.")
 
     return train_RCController(
-        model, train_seq, control_seq, target_seq, spinup, initial_res_state, beta
+        model,
+        train_seq,
+        control_seq,
+        target_seq,
+        spinup,
+        initial_res_state,
+        beta,
+        multi_sequence=multi_sequence,
     )

src/orc/drivers.py

@@ -869,9 +869,7 @@ def __init__(self, res_dim, input_rescaling=15.0, *, seed):
         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):

src/orc/forecaster/base.py

@@ -391,7 +391,7 @@ def res_ode(t, r, args):
         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.