From 5550f6aa4dc14d35d91fd6aa4c79e550e01f63e1 Mon Sep 17 00:00:00 2001 From: Jan-Williams Date: Tue, 26 May 2026 10:33:35 -0700 Subject: [PATCH 1/6] Multi-sequence training of RCForecasters + formatting --- src/orc/__init__.py | 2 +- src/orc/classifier/train.py | 5 +- src/orc/control/models.py | 2 +- src/orc/drivers.py | 4 +- src/orc/forecaster/base.py | 2 +- src/orc/forecaster/train.py | 154 ++++++++--- tests/classifier/test_classifier_models.py | 10 +- tests/control/test_control_models.py | 4 +- tests/forecaster/test_forecast_models.py | 304 ++++++++++++++++++++- tests/test_drivers.py | 2 +- tests/test_readouts.py | 14 +- 11 files changed, 434 insertions(+), 69 deletions(-) diff --git a/src/orc/__init__.py b/src/orc/__init__.py index eba9d6e..8e48416 100644 --- a/src/orc/__init__.py +++ b/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, diff --git a/src/orc/classifier/train.py b/src/orc/classifier/train.py index 6a011b9..05015b9 100644 --- a/src/orc/classifier/train.py +++ b/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) diff --git a/src/orc/control/models.py b/src/orc/control/models.py index ecaa195..ead080a 100644 --- a/src/orc/control/models.py +++ b/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, ) diff --git a/src/orc/drivers.py b/src/orc/drivers.py index 276ce02..e1bcee7 100644 --- a/src/orc/drivers.py +++ b/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): diff --git a/src/orc/forecaster/base.py b/src/orc/forecaster/base.py index f59e8a2..d8a7365 100644 --- a/src/orc/forecaster/base.py +++ b/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. diff --git a/src/orc/forecaster/train.py b/src/orc/forecaster/train.py index ec72886..948ece2 100644 --- a/src/orc/forecaster/train.py +++ b/src/orc/forecaster/train.py @@ -25,6 +25,7 @@ def train_RCForecaster( initial_res_state: Array | None = None, beta: float = 8e-8, batch_size: int | None = None, + multi_sequence: bool = False, **force_kwargs, ) -> tuple[VariableRCForecaster, Array]: """Unified training function for reservoir computer forecasters. @@ -41,14 +42,21 @@ def train_RCForecaster( model : RCForecasterBase Reservoir computer forecaster 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. 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 @@ -56,58 +64,86 @@ def train_RCForecaster( regression. If None (default), processes all reservoirs at once. Use smaller values to reduce memory usage for large numbers of parallel reservoirs. Only used when readout.chunks > 0. + multi_sequence : bool + If True, treat train_seq (and target_seq) as a batch 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() (e.g. ts=t_train - for continuous models). + for continuous models). When ``multi_sequence=True``, array-valued + kwargs must have a leading trajectory axis (e.g. ts of shape + ``(n_traj, seq_len)``). Returns ------- model : RCForecasterBase 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 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 - target_seq = train_seq[1:, :] - train_seq = train_seq[:-1, :] + target_seq = train_seq[:, 1:, :] if multi_sequence else train_seq[1:, :] else: - tot_seq = jnp.vstack((train_seq, target_seq[-1:])) + if multi_sequence: + tot_seq = jnp.concatenate([train_seq, target_seq[:, -1:, :]], axis=1) + else: + tot_seq = jnp.vstack((train_seq, target_seq[-1:])) - tot_res_seq = model.force(tot_seq, initial_res_state, **force_kwargs) - res_seq = tot_res_seq[:-1] + if multi_sequence: + if force_kwargs: + ts = force_kwargs["ts"] - res_seq_train = model.readout.prepare_train(res_seq) - train_target = model.readout.prepare_target(target_seq) + def _force(seq, state, t): + return model.force(seq, state, ts=t) + + tot_res_seq = eqx.filter_vmap(_force)(tot_seq, initial_res_state, ts) + else: + tot_res_seq = eqx.filter_vmap(model.force)(tot_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, 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) @@ -123,6 +159,7 @@ def train_ESNForecaster( initial_res_state: Array | None = None, beta: float = 8e-8, batch_size: int | None = None, + multi_sequence: bool = False, ) -> tuple[ESNForecaster, Array]: """Training function for ESNForecaster. @@ -144,6 +181,9 @@ def train_ESNForecaster( Number of parallel reservoirs to process in each batch for ridge regression. If None (default), processes all reservoirs at once. Use smaller values to reduce memory usage for large numbers of parallel reservoirs. + multi_sequence : bool + If True, treat train_seq as a batch of trajectories with shape + (n_traj, seq_len, data_dim). See train_RCForecaster for details. Returns ------- @@ -156,7 +196,14 @@ def train_ESNForecaster( raise TypeError("Model must be an ESNForecaster.") return train_RCForecaster( - model, train_seq, target_seq, spinup, initial_res_state, beta, batch_size + model, + train_seq, + target_seq, + spinup, + initial_res_state, + beta, + batch_size, + multi_sequence, ) @@ -169,6 +216,7 @@ def train_CESNForecaster( initial_res_state: Array | None = None, beta: float = 8e-8, batch_size: int | None = None, + multi_sequence: bool = False, ) -> tuple[CESNForecaster, Array]: """Training function for CESNForecaster. @@ -177,11 +225,14 @@ def train_CESNForecaster( model : CESNForecaster CESNForecaster 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)) or + (shape=(n_traj, seq_len, data_dim)) when multi_sequence=True. t_train : Array - time vector corresponding to the training sequence, (shape=(seq_len,)). + Time vector corresponding to the training sequence, (shape=(seq_len,)) + or (shape=(n_traj, seq_len)) when multi_sequence=True. target_seq : Array - Target sequence for training reservoir, (shape=(seq_len, data_dim)). + Target sequence for training reservoir, (shape=(seq_len, data_dim)) or + (shape=(n_traj, seq_len, data_dim)) when multi_sequence=True. initial_res_state : Array Initial reservoir state, (shape=(chunks, res_dim,)). spinup : int @@ -192,6 +243,9 @@ def train_CESNForecaster( Number of parallel reservoirs to process in each batch for ridge regression. If None (default), processes all reservoirs at once. Use smaller values to reduce memory usage for large numbers of parallel reservoirs. + multi_sequence : bool + If True, treat train_seq as a batch of trajectories. t_train must then + have shape (n_traj, seq_len). See train_RCForecaster for details. Returns ------- @@ -203,11 +257,24 @@ def train_CESNForecaster( if not isinstance(model, CESNForecaster): raise TypeError("Model must be a CESNForecaster.") - if train_seq.shape[0] != t_train.shape[0]: - raise ValueError("train_seq and t_train must have the same length.") + if multi_sequence: + if train_seq.shape[:2] != t_train.shape[:2]: + raise ValueError( + "train_seq and t_train must have the same (n_traj, seq_len)." + ) + else: + if train_seq.shape[0] != t_train.shape[0]: + raise ValueError("train_seq and t_train must have the same length.") return train_RCForecaster( - model, train_seq, target_seq, spinup, initial_res_state, beta, batch_size, + model, + train_seq, + target_seq, + spinup, + initial_res_state, + beta, + batch_size, + multi_sequence, ts=t_train, ) @@ -220,6 +287,7 @@ def train_EnsembleESNForecaster( initial_res_state: Array | None = None, beta: float = 8e-8, batch_size: int | None = None, + multi_sequence: bool = False, ) -> tuple[EnsembleESNForecaster, Array]: """Training function for EnsembleESNForecaster. @@ -241,6 +309,9 @@ def train_EnsembleESNForecaster( Number of parallel reservoirs to process in each batch for ridge regression. If None (default), processes all reservoirs at once. Use smaller values to reduce memory usage for large numbers of parallel reservoirs. + multi_sequence : bool + If True, treat train_seq as a batch of trajectories with shape + (n_traj, seq_len, data_dim). See train_RCForecaster for details. Returns ------- @@ -253,5 +324,12 @@ def train_EnsembleESNForecaster( raise TypeError("Model must be an EnsembleESNForecaster.") return train_RCForecaster( - model, train_seq, target_seq, spinup, initial_res_state, beta, batch_size + model, + train_seq, + target_seq, + spinup, + initial_res_state, + beta, + batch_size, + multi_sequence, ) diff --git a/tests/classifier/test_classifier_models.py b/tests/classifier/test_classifier_models.py index 1673568..30bd1c2 100644 --- a/tests/classifier/test_classifier_models.py +++ b/tests/classifier/test_classifier_models.py @@ -291,9 +291,7 @@ def test_train_rcclassifier_quadratic(dummy_classification_data): @pytest.fixture def esn_classifier(): - return orc.classifier.ESNClassifier( - data_dim=4, n_classes=3, res_dim=200, seed=0 - ) + return orc.classifier.ESNClassifier(data_dim=4, n_classes=3, res_dim=200, seed=0) def test_classifier_transform_stability(esn_classifier): @@ -304,11 +302,11 @@ def test_classifier_transform_stability(esn_classifier): batch_size = 3 res_state = model.driver.default_state() - batch_seq = jax.random.normal( - key, shape=(batch_size, seq_len, model.data_dim) - ) + batch_seq = jax.random.normal(key, shape=(batch_size, seq_len, model.data_dim)) + def fwd(in_seq): return model.classify(in_seq, res_state) + vmap_fwd = eqx.filter_vmap(fwd) assert jnp.allclose( vmap_fwd(batch_seq), diff --git a/tests/control/test_control_models.py b/tests/control/test_control_models.py index 43a002e..6cd08a5 100644 --- a/tests/control/test_control_models.py +++ b/tests/control/test_control_models.py @@ -375,9 +375,7 @@ def test_train_rccontroller_quadratic(dummy_control_problem_params): @pytest.fixture def esn_controller(): - return orc.control.ESNController( - data_dim=3, control_dim=2, res_dim=200, seed=0 - ) + return orc.control.ESNController(data_dim=3, control_dim=2, res_dim=200, seed=0) def test_controller_transform_stability(esn_controller): diff --git a/tests/forecaster/test_forecast_models.py b/tests/forecaster/test_forecast_models.py index c2f5e5e..8d4033f 100644 --- a/tests/forecaster/test_forecast_models.py +++ b/tests/forecaster/test_forecast_models.py @@ -549,9 +549,7 @@ def test_train_rcforecaster_cesn(): U_train = U[:2000, :] ts_train = ts[:2000] - cesn = orc.forecaster.CESNForecaster( - data_dim=3, res_dim=res_dim, seed=0 - ) + cesn = orc.forecaster.CESNForecaster(data_dim=3, res_dim=res_dim, seed=0) cesn_old, R_old = orc.forecaster.train_CESNForecaster(cesn, U_train, ts_train) cesn_new, R_new = orc.forecaster.train_RCForecaster(cesn, U_train, ts=ts_train) @@ -627,6 +625,26 @@ class SimpleForecaster(orc.forecaster.RCForecasterBase): ##################### JAX TRANSFORM & AD TESTS ##################### +def test_train_rcforecaster_jax_jit_defaults(): + """jax.jit(train_RCForecaster)(model, U_train) works with only positional args. + + Default arguments (beta, spinup, multi_sequence, etc.) are Python literals, + so they are never traced by JAX and remain safe to use without static_argnums. + """ + tN, dt = 20, 0.01 + u0 = np.array([0.05, 1.0, 1.05]) + U, _ = orc.data.lorenz63(tN=tN, dt=dt, u0=u0) + U_train = U[:1000] + + esn = orc.forecaster.ESNForecaster(data_dim=3, res_dim=200, seed=0) + + esn_jit, R_jit = jax.jit(orc.forecaster.train_RCForecaster)(esn, U_train) + esn_ref, R_ref = orc.forecaster.train_RCForecaster(esn, U_train) + + assert jnp.allclose(esn_jit.readout.wout, esn_ref.readout.wout) + assert jnp.allclose(R_jit, R_ref) + + @pytest.fixture def esn_forecaster(): return orc.forecaster.ESNForecaster(data_dim=3, res_dim=200, seed=0) @@ -744,15 +762,287 @@ def loss_fn(m, s): def check_finite(g): if eqx.is_array(g) and jnp.issubdtype(g.dtype, jnp.inexact): assert jnp.all(jnp.isfinite(g)) + jax.tree_util.tree_map(check_finite, grads) +##################### MULTI-SEQUENCE TRAINING TESTS ##################### + + +@pytest.fixture +def lorenz_segments(): + """Three Lorenz63 trajectories of equal length for multi-sequence tests.""" + u0s = [ + np.array([0.05, 1.0, 1.05]), + np.array([1.0, 0.5, 2.0]), + np.array([-1.0, 2.0, 0.5]), + ] + segs = [orc.data.lorenz63(tN=5, dt=0.01, u0=u0)[0][:300] for u0 in u0s] + return jnp.stack(segs) # (3, 300, 3) + + +def test_multi_seq_n1_matches_single_seq(lorenz_segments): + """n_traj=1 multi_sequence training matches single-sequence training.""" + U = lorenz_segments[0] # (300, 3) + esn = orc.forecaster.ESNForecaster(data_dim=3, res_dim=200, seed=0) + + esn_single, _ = orc.forecaster.train_RCForecaster(esn, U, spinup=20) + esn_multi, _ = orc.forecaster.train_RCForecaster( + esn, U[None], spinup=20, multi_sequence=True + ) + + assert jnp.allclose(esn_single.readout.wout, esn_multi.readout.wout) + + +def test_multi_seq_n1_explicit_target(lorenz_segments): + """n_traj=1 multi_sequence with explicit target_seq matches single-sequence.""" + U = lorenz_segments[0] + train, target = U[:250], U[1:251] + esn = orc.forecaster.ESNForecaster(data_dim=3, res_dim=200, seed=0) + + esn_single, _ = orc.forecaster.train_RCForecaster( + esn, train, target_seq=target, spinup=20 + ) + esn_multi, _ = orc.forecaster.train_RCForecaster( + esn, train[None], target_seq=target[None], spinup=20, multi_sequence=True + ) + + assert jnp.allclose(esn_single.readout.wout, esn_multi.readout.wout) + + +def test_multi_seq_equiv_manual_concat(lorenz_segments): + """Multi-sequence ESN training matches per-trajectory training.""" + from orc.utils.regressions import _solve_all_ridge_reg + + segs = lorenz_segments # (3, 300, 3) + spinup = 20 + beta = 1e-6 + esn = orc.forecaster.ESNForecaster(data_dim=3, res_dim=200, seed=0) + + esn_multi, _ = orc.forecaster.train_RCForecaster( + esn, segs, spinup=spinup, beta=beta, multi_sequence=True + ) + + default_state = esn.driver.default_state() + res_parts, target_parts = [], [] + for i in range(segs.shape[0]): + U = segs[i] + tot_res = esn.force(U, default_state) + res = tot_res[:-1] + target = U[1:] + res_parts.append(esn.readout.prepare_train(res)[spinup:]) + target_parts.append(esn.readout.prepare_target(target)[spinup:]) + + cmat = _solve_all_ridge_reg( + jnp.concatenate(res_parts, axis=0), + jnp.concatenate(target_parts, axis=0), + beta, + ) + + assert jnp.allclose(esn_multi.readout.wout, cmat) + + +def test_multi_seq_tot_res_seq_shape(lorenz_segments): + """tot_res_seq has leading (n_traj, seq_len) axes when multi_sequence=True.""" + segs = lorenz_segments # (3, 300, 3) + n_traj, seq_len, _ = segs.shape + esn = orc.forecaster.ESNForecaster(data_dim=3, res_dim=200, seed=0) + + _, tot_res_seq = orc.forecaster.train_RCForecaster(esn, segs, multi_sequence=True) + + assert tot_res_seq.shape[0] == n_traj + assert tot_res_seq.shape[1] == seq_len + + +def test_multi_seq_spinup_validation(lorenz_segments): + """spinup >= seq_len raises ValueError when multi_sequence=True.""" + segs = lorenz_segments # (3, 300, 3) + esn = orc.forecaster.ESNForecaster(data_dim=3, res_dim=200, seed=0) + + with pytest.raises(ValueError, match="spinup"): + orc.forecaster.train_RCForecaster(esn, segs, spinup=300, multi_sequence=True) + + +def test_multi_seq_par_esn_n1_matches_single(lorenz_segments): + """n_traj=1 multi_sequence with parallel ESN matches single-sequence.""" + U = lorenz_segments[0] + chunks = 3 + esn = orc.forecaster.ESNForecaster(data_dim=3, res_dim=200, seed=0, chunks=chunks) + + esn_single, _ = orc.forecaster.train_RCForecaster(esn, U, spinup=20) + esn_multi, _ = orc.forecaster.train_RCForecaster( + esn, U[None], spinup=20, multi_sequence=True + ) + + assert jnp.allclose(esn_single.readout.wout, esn_multi.readout.wout) + + +def test_multi_seq_par_esn_equiv_manual_concat(lorenz_segments): + """Multi-sequence parallel ESN training matches per-trajectory training.""" + from orc.utils.regressions import _solve_all_ridge_reg + + segs = lorenz_segments + chunks = 3 + spinup = 20 + beta = 1e-6 + esn = orc.forecaster.ESNForecaster(data_dim=3, res_dim=200, seed=0, chunks=chunks) + + esn_multi, _ = orc.forecaster.train_RCForecaster( + esn, segs, spinup=spinup, beta=beta, multi_sequence=True + ) + + default_state = esn.driver.default_state() + res_parts, target_parts = [], [] + for i in range(segs.shape[0]): + U = segs[i] + tot_res = esn.force(U, default_state) + res = tot_res[:-1] + target = U[1:] + res_parts.append(esn.readout.prepare_train(res)[spinup:]) + target_parts.append(esn.readout.prepare_target(target)[spinup:]) + + cmat = _solve_all_ridge_reg( + jnp.concatenate(res_parts, axis=0), + jnp.concatenate(target_parts, axis=0), + beta, + ) + + assert jnp.allclose(esn_multi.readout.wout, cmat) + + +def test_multi_seq_cesn_n1_matches_single(lorenz_segments): + """n_traj=1 multi_sequence CESN training matches single-sequence.""" + U = lorenz_segments[0] # (300, 3) + dt = 0.01 + ts = jnp.arange(300, dtype=jnp.float64) * dt + + cesn = orc.forecaster.CESNForecaster( + data_dim=3, + res_dim=100, + seed=0, + solver=diffrax.Euler(), + stepsize_controller=diffrax.ConstantStepSize(), + ) + + cesn_single, _ = orc.forecaster.train_RCForecaster(cesn, U, spinup=10, ts=ts) + cesn_multi, _ = orc.forecaster.train_RCForecaster( + cesn, U[None], spinup=10, ts=ts[None], multi_sequence=True + ) + + assert jnp.allclose(cesn_single.readout.wout, cesn_multi.readout.wout) + + +def test_train_cesn_multi_seq_ts_shape_error(lorenz_segments): + """train_CESNForecaster raises ValueError with t_train/train_seq mismatch.""" + segs = lorenz_segments # (3, 300, 3) + wrong_ts = jnp.arange(300, dtype=jnp.float64) * 0.01 # (300,) instead of (3, 300) + + cesn = orc.forecaster.CESNForecaster( + data_dim=3, + res_dim=100, + seed=0, + solver=diffrax.Euler(), + stepsize_controller=diffrax.ConstantStepSize(), + ) + + with pytest.raises(ValueError): + orc.forecaster.train_CESNForecaster(cesn, segs, wrong_ts, multi_sequence=True) + + +def test_multi_seq_jax_jit_compatible(lorenz_segments): + """train_RCForecaster with multi_sequence=True is compatible with jax.jit.""" + segs = lorenz_segments + esn = orc.forecaster.ESNForecaster(data_dim=3, res_dim=200, seed=0) + + def train_fn(model, train_seq): + return orc.forecaster.train_RCForecaster( + model, train_seq, spinup=20, beta=1e-6, multi_sequence=True + ) + + esn_jit, _ = jax.jit(train_fn)(esn, segs) + esn_ref, _ = train_fn(esn, segs) + + assert jnp.allclose(esn_jit.readout.wout, esn_ref.readout.wout) + + +def test_multi_seq_filter_jit_compatible(lorenz_segments): + """train_RCForecaster with multi_sequence=True is compatible with eqx.filter_jit.""" + segs = lorenz_segments + esn = orc.forecaster.ESNForecaster(data_dim=3, res_dim=200, seed=0) + + @eqx.filter_jit + def train_fn(model, train_seq): + return orc.forecaster.train_RCForecaster( + model, train_seq, spinup=20, beta=1e-6, multi_sequence=True + ) + + esn_jit, _ = train_fn(esn, segs) + esn_ref, _ = orc.forecaster.train_RCForecaster( + esn, segs, spinup=20, beta=1e-6, multi_sequence=True + ) + + assert jnp.allclose(esn_jit.readout.wout, esn_ref.readout.wout) + + +def test_multi_seq_jax_jit_static_argnums(lorenz_segments): + """multi_sequence can be declared as static_argnums for jax.jit directly.""" + segs = lorenz_segments + esn = orc.forecaster.ESNForecaster(data_dim=3, res_dim=200, seed=0) + + # positional order: model, train_seq, target_seq, spinup, initial_res_state, + # beta, batch_size, multi_sequence + # beta (pos 5) must be static: _solve_all_ridge_reg uses eqx.if_array(1) which + # attempts to vmap over any array argument, including a traced scalar beta. + jit_train = jax.jit( + orc.forecaster.train_RCForecaster, + static_argnums=(3, 5, 6, 7), # spinup, beta, batch_size, multi_sequence + ) + + esn_jit, _ = jit_train(esn, segs, None, 20, None, 1e-6, None, True) + esn_ref, _ = orc.forecaster.train_RCForecaster( + esn, segs, spinup=20, beta=1e-6, multi_sequence=True + ) + + assert jnp.allclose(esn_jit.readout.wout, esn_ref.readout.wout) + + +def test_multi_seq_cesn_jit_compatible(lorenz_segments): + """Multi-sequence CESN training is compatible with jax.jit and eqx.filter_jit.""" + segs = lorenz_segments # (3, 300, 3) + dt = 0.01 + ts_batch = jnp.stack([jnp.arange(300, dtype=jnp.float64) * dt] * 3) # (3, 300) + + cesn = orc.forecaster.CESNForecaster( + data_dim=3, + res_dim=100, + seed=0, + solver=diffrax.Euler(), + stepsize_controller=diffrax.ConstantStepSize(), + ) + + def train_fn(model, train_seq, ts): + return orc.forecaster.train_RCForecaster( + model, train_seq, spinup=10, beta=1e-6, multi_sequence=True, ts=ts + ) + + cesn_ref, _ = train_fn(cesn, segs, ts_batch) + + # XLA may reorder floating-point ops under JIT, causing small numerical + # differences vs eager mode even with a deterministic solver. + tol = 1e-4 + + cesn_jit, _ = jax.jit(train_fn)(cesn, segs, ts_batch) + assert jnp.allclose(cesn_jit.readout.wout, cesn_ref.readout.wout, atol=tol) + + cesn_fjit, _ = eqx.filter_jit(train_fn)(cesn, segs, ts_batch) + assert jnp.allclose(cesn_fjit.readout.wout, cesn_ref.readout.wout, atol=tol) + + ##################### GRU FORECAST SMOKE TEST ##################### def test_gru_forecaster_smoke(): - """Built-in GRUDriver forecasts Rossler. - """ + """Built-in GRUDriver forecasts Rossler.""" class GRUForecaster(orc.forecaster.RCForecasterBase): pass @@ -776,9 +1066,7 @@ class GRUForecaster(orc.forecaster.RCForecasterBase): model, _ = orc.forecaster.train_RCForecaster( model, train_seq=U_train, spinup=spinup, beta=1e-7 ) - U_pred = model.forecast_from_IC( - fcast_len=fcast_len, spinup_data=U_train[-spinup:] - ) + U_pred = model.forecast_from_IC(fcast_len=fcast_len, spinup_data=U_train[-spinup:]) assert U_pred.shape == (fcast_len, data_dim) assert jnp.all(jnp.isfinite(U_pred)) diff --git a/tests/test_drivers.py b/tests/test_drivers.py index 0fcb07e..1efc940 100644 --- a/tests/test_drivers.py +++ b/tests/test_drivers.py @@ -1100,7 +1100,7 @@ def test_driver_transform_stability(driver_fixture, request): # jax.vmap'ed advance should match the built-in batch_advance assert jnp.allclose( vmap_advance(batch_proj, batch_state), - driver.batch_advance(batch_proj, batch_state) + driver.batch_advance(batch_proj, batch_state), ) jit_advance = eqx.filter_jit(driver.advance) diff --git a/tests/test_readouts.py b/tests/test_readouts.py index 5f91543..a95ab98 100644 --- a/tests/test_readouts.py +++ b/tests/test_readouts.py @@ -127,8 +127,11 @@ def test_parallel_nonlinear_prepare_train(chunks): res_dim = 12 out_dim = 12 readout = orc.readouts.ParallelNonlinearReadout( - out_dim=out_dim, res_dim=res_dim, nonlin_list=[lambda x: x**2], - chunks=chunks, dtype=jnp.float64, + out_dim=out_dim, + res_dim=res_dim, + nonlin_list=[lambda x: x**2], + chunks=chunks, + dtype=jnp.float64, ) res_seq = jax.random.normal(jax.random.key(0), shape=(50, chunks, res_dim)) result = readout.prepare_train(res_seq) @@ -142,8 +145,11 @@ def test_parallel_nonlinear_prepare_target(chunks): """Test that ParallelNonlinearReadout.prepare_target reshapes correctly.""" out_dim = 12 readout = orc.readouts.ParallelNonlinearReadout( - out_dim=out_dim, res_dim=12, nonlin_list=[lambda x: x**2], - chunks=chunks, dtype=jnp.float64, + out_dim=out_dim, + res_dim=12, + nonlin_list=[lambda x: x**2], + chunks=chunks, + dtype=jnp.float64, ) target_seq = jnp.ones((50, out_dim)) result = readout.prepare_target(target_seq) From e439dfe81b5b92e7a30bec9aa804964b3a196639 Mon Sep 17 00:00:00 2001 From: Jan-Williams Date: Tue, 26 May 2026 12:29:25 -0700 Subject: [PATCH 2/6] Added multi-sequence training for RCController training --- pyproject.toml | 3 + src/orc/control/train.py | 125 +++++++++++++++------- tests/control/test_control_models.py | 150 +++++++++++++++++++++++++++ 3 files changed, 240 insertions(+), 38 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 6fb025b..7eaf04f 100644 --- a/pyproject.toml +++ b/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"] diff --git a/src/orc/control/train.py b/src/orc/control/train.py index f1eb093..8c7ad4e 100644 --- a/src/orc/control/train.py +++ b/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, ) diff --git a/tests/control/test_control_models.py b/tests/control/test_control_models.py index 6cd08a5..87a3582 100644 --- a/tests/control/test_control_models.py +++ b/tests/control/test_control_models.py @@ -370,6 +370,156 @@ def test_train_rccontroller_quadratic(dummy_control_problem_params): assert jnp.all(jnp.isfinite(U_controlled)) +##################### MULTI-SEQUENCE TRAINING TESTS ##################### + + +@pytest.fixture +def lorenz_controller_segments(): + """Three Lorenz63 trajectories with control sequences for multi-sequence tests.""" + seq_len = 200 + u0s = [ + np.array([0.05, 1.0, 1.05]), + np.array([1.0, 0.5, 2.0]), + np.array([-1.0, 2.0, 0.5]), + ] + state_segs = jnp.stack( + [orc.data.lorenz63(tN=5, dt=0.01, u0=u0)[0][:seq_len] for u0 in u0s] + ) # (3, 200, 3) + time = jnp.arange(seq_len, dtype=jnp.float64).reshape(-1, 1) + control_segs = jnp.stack( + [ + 0.1 * jnp.sin(2 * jnp.pi * (i + 1) * time / seq_len * jnp.array([1.0, 1.5])) + for i in range(3) + ] + ) # (3, 200, 2) + return state_segs, control_segs + + +def test_multi_seq_controller_n1_matches_single_seq(lorenz_controller_segments): + """n_traj=1 multi_sequence controller training matches single-sequence.""" + state_segs, control_segs = lorenz_controller_segments + U, C = state_segs[0], control_segs[0] # (200, 3), (200, 2) + + ctrl = orc.control.ESNController(data_dim=3, control_dim=2, res_dim=200, seed=0) + + ctrl_single, _ = orc.control.train_RCController(ctrl, U, C, spinup=20) + ctrl_multi, _ = orc.control.train_RCController( + ctrl, U[None], C[None], spinup=20, multi_sequence=True + ) + + assert jnp.allclose(ctrl_single.readout.wout, ctrl_multi.readout.wout) + + +def test_multi_seq_controller_n1_explicit_target(lorenz_controller_segments): + """n_traj=1 multi_sequence with explicit target_seq matches single-sequence.""" + state_segs, control_segs = lorenz_controller_segments + U, C = state_segs[0], control_segs[0] + train, target, ctrl_in = U[:150], U[1:151], C[:150] + + ctrl = orc.control.ESNController(data_dim=3, control_dim=2, res_dim=200, seed=0) + + ctrl_single, _ = orc.control.train_RCController( + ctrl, train, ctrl_in, target_seq=target, spinup=20 + ) + ctrl_multi, _ = orc.control.train_RCController( + ctrl, + train[None], + ctrl_in[None], + target_seq=target[None], + spinup=20, + multi_sequence=True, + ) + + assert jnp.allclose(ctrl_single.readout.wout, ctrl_multi.readout.wout) + + +def test_multi_seq_controller_equiv_manual_concat(lorenz_controller_segments): + """Multi-sequence controller training matches per-trajectory training.""" + from orc.utils.regressions import _solve_all_ridge_reg + + state_segs, control_segs = lorenz_controller_segments + spinup = 20 + beta = 1e-6 + ctrl = orc.control.ESNController(data_dim=3, control_dim=2, res_dim=200, seed=0) + + ctrl_multi, _ = orc.control.train_RCController( + ctrl, state_segs, control_segs, spinup=spinup, beta=beta, multi_sequence=True + ) + + default_state = ctrl.driver.default_state() + res_parts, target_parts = [], [] + for i in range(state_segs.shape[0]): + U, C = state_segs[i], control_segs[i] + tot_res = ctrl.force(U, C, default_state) + res = tot_res[:-1] + target = U[1:] + res_parts.append(ctrl.readout.prepare_train(res)[spinup:]) + target_parts.append(ctrl.readout.prepare_target(target)[spinup:]) + + cmat = _solve_all_ridge_reg( + jnp.concatenate(res_parts, axis=0), + jnp.concatenate(target_parts, axis=0), + beta, + ) + + assert jnp.allclose(ctrl_multi.readout.wout, cmat) + + +def test_multi_seq_controller_tot_res_seq_shape(lorenz_controller_segments): + """tot_res_seq has leading (n_traj, seq_len) axes when multi_sequence=True.""" + state_segs, control_segs = lorenz_controller_segments + n_traj, seq_len, _ = state_segs.shape + ctrl = orc.control.ESNController(data_dim=3, control_dim=2, res_dim=200, seed=0) + + _, tot_res_seq = orc.control.train_RCController( + ctrl, state_segs, control_segs, multi_sequence=True + ) + + assert tot_res_seq.shape[0] == n_traj + assert tot_res_seq.shape[1] == seq_len + + +def test_multi_seq_controller_spinup_validation(lorenz_controller_segments): + """spinup >= seq_len raises ValueError when multi_sequence=True.""" + state_segs, control_segs = lorenz_controller_segments + ctrl = orc.control.ESNController(data_dim=3, control_dim=2, res_dim=200, seed=0) + + with pytest.raises(ValueError, match="spinup"): + orc.control.train_RCController( + ctrl, state_segs, control_segs, spinup=200, multi_sequence=True + ) + + +def test_multi_seq_controller_seq_length_mismatch(lorenz_controller_segments): + """Mismatched train_seq/control_seq seq_len raises ValueError.""" + state_segs, control_segs = lorenz_controller_segments + ctrl = orc.control.ESNController(data_dim=3, control_dim=2, res_dim=200, seed=0) + + with pytest.raises(ValueError): + orc.control.train_RCController( + ctrl, state_segs, control_segs[:, :-1, :], multi_sequence=True + ) + + +def test_multi_seq_controller_jax_jit_compatible(lorenz_controller_segments): + """train_RCController with multi_sequence=True is compatible with jits.""" + state_segs, control_segs = lorenz_controller_segments + ctrl = orc.control.ESNController(data_dim=3, control_dim=2, res_dim=200, seed=0) + + def train_fn(model, train_seq, ctrl_seq): + return orc.control.train_RCController( + model, train_seq, ctrl_seq, spinup=20, beta=1e-6, multi_sequence=True + ) + + ctrl_ref, _ = train_fn(ctrl, state_segs, control_segs) + + ctrl_jit, _ = jax.jit(train_fn)(ctrl, state_segs, control_segs) + assert jnp.allclose(ctrl_jit.readout.wout, ctrl_ref.readout.wout) + + ctrl_fjit, _ = eqx.filter_jit(train_fn)(ctrl, state_segs, control_segs) + assert jnp.allclose(ctrl_fjit.readout.wout, ctrl_ref.readout.wout) + + ##################### JAX TRANSFORM & AD TESTS ##################### From 17db146f043a4687a4cbfa798b84e23f6c573a92 Mon Sep 17 00:00:00 2001 From: Jan-Williams Date: Tue, 21 Jul 2026 08:57:33 -0700 Subject: [PATCH 3/6] Updated equivalence test to check on outputs, not C matrices --- tests/control/test_control_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/control/test_control_models.py b/tests/control/test_control_models.py index 87a3582..7e8de89 100644 --- a/tests/control/test_control_models.py +++ b/tests/control/test_control_models.py @@ -462,7 +462,7 @@ def test_multi_seq_controller_equiv_manual_concat(lorenz_controller_segments): beta, ) - assert jnp.allclose(ctrl_multi.readout.wout, cmat) + assert jnp.allclose(ctrl_multi.readout.wout @ res, cmat @ res) def test_multi_seq_controller_tot_res_seq_shape(lorenz_controller_segments): From 660705fa017da9cef6c1cb7d8ef0ae6c558ee944 Mon Sep 17 00:00:00 2001 From: Jan-Williams Date: Wed, 22 Jul 2026 18:36:39 -0700 Subject: [PATCH 4/6] Bumped tolerance for multi control sequence testing --- tests/control/test_control_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/control/test_control_models.py b/tests/control/test_control_models.py index 7e8de89..0f8df23 100644 --- a/tests/control/test_control_models.py +++ b/tests/control/test_control_models.py @@ -462,7 +462,7 @@ def test_multi_seq_controller_equiv_manual_concat(lorenz_controller_segments): beta, ) - assert jnp.allclose(ctrl_multi.readout.wout @ res, cmat @ res) + assert jnp.allclose(cmat, ctrl_multi.readout.wout, atol=1e-5) def test_multi_seq_controller_tot_res_seq_shape(lorenz_controller_segments): From cdc228aa11a2f1ef7b9b43246e7a5ff72f432a83 Mon Sep 17 00:00:00 2001 From: Jan-Williams Date: Wed, 22 Jul 2026 18:45:34 -0700 Subject: [PATCH 5/6] Bumped tolerance for multi control sequence testing --- tests/control/test_control_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/control/test_control_models.py b/tests/control/test_control_models.py index 0f8df23..f93c10f 100644 --- a/tests/control/test_control_models.py +++ b/tests/control/test_control_models.py @@ -462,7 +462,7 @@ def test_multi_seq_controller_equiv_manual_concat(lorenz_controller_segments): beta, ) - assert jnp.allclose(cmat, ctrl_multi.readout.wout, atol=1e-5) + assert jnp.allclose(cmat, ctrl_multi.readout.wout, atol=1e-4) def test_multi_seq_controller_tot_res_seq_shape(lorenz_controller_segments): From 0f432d0128cd5c01a0b99926cff014ceb6e588e5 Mon Sep 17 00:00:00 2001 From: Jan-Williams Date: Wed, 22 Jul 2026 18:54:06 -0700 Subject: [PATCH 6/6] Bumped tolerance for multi control sequence testing --- tests/control/test_control_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/control/test_control_models.py b/tests/control/test_control_models.py index f93c10f..e8aefe5 100644 --- a/tests/control/test_control_models.py +++ b/tests/control/test_control_models.py @@ -462,7 +462,7 @@ def test_multi_seq_controller_equiv_manual_concat(lorenz_controller_segments): beta, ) - assert jnp.allclose(cmat, ctrl_multi.readout.wout, atol=1e-4) + assert jnp.allclose(cmat, ctrl_multi.readout.wout, atol=1e-2) def test_multi_seq_controller_tot_res_seq_shape(lorenz_controller_segments):