Reference

Models & ensembles

Five factory functions that return ready-to-use ESNModel instances (two of them — linear_esn and headless_esn — deliberately ship without a readout), plus coupled_ensemble_esn, which builds a CoupledEnsembleESNModel: the coupled-ensemble wrapper that runs N sub-models against a shared feedback signal, with the aggregators that merge their forecasts.

models

Premade ESN Architectures

This module provides pre-configured ESN model architectures that can be used directly or customized for specific tasks.

FUNCTION	DESCRIPTION
classic_esn	Traditional ESN with input concatenation.
ott_esn	Ott's ESN with state augmentation (squared even units).
power_augmented	Power Augmented ESN with state augmentation (exponentiated reservoir states).
deep_esn	Deep (hierarchical) ESN — a stack of reservoir layers whose readout sees the concatenation of every layer's state.
hybrid_esn	Hybrid physics-prior ESN (Pathak et al. 2018) — a reservoir combined with a knowledge-based expert wired in before the reservoir and before the readout.
headless_esn	Reservoir only (no readout) for analysis.
linear_esn	Linear reservoir for baseline comparison.
lsm	Spiking Liquid State Machine — leaky integrate-and-fire reservoir with a filtered-spike-trace readout (input-driven; not a chaos forecaster).
coupled_ensemble_esn	Build a :class:~resdag.ensemble.CoupledEnsembleESNModel of N coupled ESN sub-models sharing an averaged feedback signal.

CLASS	DESCRIPTION
KnowledgeModel	:class:~torch.nn.Module adapter wrapping a knowledge-based expert f(u) -> u_next for use inside :func:hybrid_esn.
Each architecture accepts individual parameters for full customization
while providing sensible defaults for quick experimentation.

Examples:

Quick start with Ott's ESN:

>>> from resdag.models import ott_esn
>>> model = ott_esn(reservoir_size=500, feedback_size=3, output_size=3)
>>> predictions = model.forecast(warmup_data, horizon=100)

Classic ESN with custom topology:

>>> from resdag.models import classic_esn
>>> from resdag.init.topology import get_topology
>>>
>>> model = classic_esn(
...     reservoir_size=500,
...     feedback_size=3,
...     output_size=3,
...     topology=get_topology("watts_strogatz", k=4, p=0.3),
...     spectral_radius=0.95,
... )

See Also

resdag.core.ESNModel : Base ESN model class. resdag.layers.ESNLayer : ESN layer used by these models. resdag.training.ESNTrainer : Trainer for fitting readout layers.

classic_esn

classic_esn(reservoir_size: int, feedback_size: int, output_size: int, input_size: int | None = None, input_initializer: InitializerSpec | None = None, topology: TopologySpec | None = None, spectral_radius: float = 0.9, leak_rate: float = 1.0, noise: float = 0.0, feedback_initializer: InitializerSpec | None = None, activation: str = 'tanh', bias: bool = True, trainable: bool = False, readout: ReadoutSpec = 'ridge', readout_alpha: float = 1e-06, readout_bias: bool = True, readout_name: str = 'output', **reservoir_kwargs: Any) -> ESNModel

Build a classic Echo State Network (ESN) model.

This architecture concatenates the input with the reservoir output before passing to the readout layer, following the traditional ESN design.

Architecture::

Input -> Reservoir -> Concatenate(Input, Reservoir) -> Readout

The readout sees both the raw input and the reservoir's nonlinear transformation, which can improve performance on many tasks.

PARAMETER	DESCRIPTION
reservoir_size	Number of units in the reservoir. TYPE: int
feedback_size	Number of feedback features (input dimension). TYPE: int
output_size	Number of output features. TYPE: int
input_size	Dimension of an optional driving (exogenous) input. When given, the model takes two inputs (feedback, driver) and the driver feeds the reservoir alongside the autoregressive feedback. The driver is kept out of the input concatenation, so the readout in_features stays feedback_size + reservoir_size. None (default) builds a feedback-only model, unchanged from previous behavior. TYPE: int or None DEFAULT: None
input_initializer	Initializer for the driving-input weights. Same accepted forms as feedback_initializer. Only used when input_size is given. TYPE: InitializerSpec DEFAULT: None
topology	Topology for recurrent weights. Accepts: str: Registry name (e.g., "erdos_renyi") tuple: (name, params) like ("watts_strogatz", {"k": 6, "p": 0.1}) :class:~resdag.init.topology.TopologyInitializer: Pre-configured object TYPE: TopologySpec DEFAULT: None
spectral_radius	Desired spectral radius for recurrent weights. TYPE: float DEFAULT: 0.9
leak_rate	Leaky integration rate (1.0 = no leak). TYPE: float DEFAULT: 1.0
noise	Standard deviation of additive Gaussian state noise injected into the reservoir after the activation. Active only in training mode (a no-op under :meth:~torch.nn.Module.eval); 0.0 disables it. Forwarded to :class:~resdag.layers.ESNLayer. Must be non-negative. TYPE: float DEFAULT: 0.0
feedback_initializer	Initializer for feedback weights. Accepts: str: Registry name (e.g., "pseudo_diagonal") tuple: (name, params) like ("chebyshev", {"p": 0.5}) :class:~resdag.init.input_feedback.InputFeedbackInitializer: Pre-configured object TYPE: InitializerSpec DEFAULT: None
activation	Activation function ("tanh", "relu", "sigmoid", "identity"). TYPE: str DEFAULT: "tanh"
bias	Whether to use bias in the reservoir. TYPE: bool DEFAULT: True
trainable	Whether reservoir weights are trainable. TYPE: bool DEFAULT: False
readout	Which readout solver to attach. "ridge" (default) is the exact direct Cholesky ridge solve (:class:~resdag.layers.RidgeReadoutLayer); "cg" selects the legacy iterative :class:~resdag.layers.CGReadoutLayer. A :class:~resdag.layers.readouts.ReadoutLayer subclass may also be passed to use a custom readout. TYPE: str or type DEFAULT: "ridge"
readout_alpha	Ridge regression regularization for readout. TYPE: float DEFAULT: 1e-6
readout_bias	Whether to use bias in the readout. TYPE: bool DEFAULT: True
readout_name	Name for the readout layer (used in training targets). TYPE: str DEFAULT: "output"
**reservoir_kwargs	Additional keyword arguments passed to :class:~resdag.layers.ESNLayer. TYPE: Any DEFAULT: {}

RETURNS	DESCRIPTION
class:`~resdag.core.ESNModel`	Configured ESN model ready for training and inference.

Examples:

Simple usage with defaults:

>>> from resdag.models import classic_esn
>>> import torch
>>> model = classic_esn(100, 1, 1)

With custom topology and initializer:

>>> model = classic_esn(
...     reservoir_size=400,
...     feedback_size=1,
...     output_size=1,
...     topology=("watts_strogatz", {"k": 6, "p": 0.1}),
...     feedback_initializer="pseudo_diagonal",
...     spectral_radius=0.9,
...     leak_rate=0.5,
... )

Forward pass:

>>> x = torch.randn(4, 50, 1)  # (batch, time, features)
>>> y = model(x)

See Also

:func:resdag.models.ott_esn : OTT ESN variant :func:resdag.models.linear_esn : Linear ESN variant :class:resdag.training.ESNTrainer : Trainer for fitting readouts

Source code in src/resdag/models/classic_esn.py

def classic_esn(
    reservoir_size: int,
    feedback_size: int,
    output_size: int,
    input_size: int | None = None,
    input_initializer: InitializerSpec | None = None,
    # Reservoir params
    topology: TopologySpec | None = None,
    spectral_radius: float = 0.9,
    leak_rate: float = 1.0,
    noise: float = 0.0,
    feedback_initializer: InitializerSpec | None = None,
    activation: str = "tanh",
    bias: bool = True,
    trainable: bool = False,
    # Readout params
    readout: ReadoutSpec = "ridge",
    readout_alpha: float = 1e-6,
    readout_bias: bool = True,
    readout_name: str = "output",
    # Extra reservoir kwargs
    **reservoir_kwargs: Any,
) -> ESNModel:
    """Build a classic Echo State Network (ESN) model.

    This architecture concatenates the input with the reservoir output before
    passing to the readout layer, following the traditional ESN design.

    Architecture::

        Input -> Reservoir -> Concatenate(Input, Reservoir) -> Readout

    The readout sees both the raw input and the reservoir's nonlinear
    transformation, which can improve performance on many tasks.

    Parameters
    ----------
    reservoir_size : int
        Number of units in the reservoir.
    feedback_size : int
        Number of feedback features (input dimension).
    output_size : int
        Number of output features.
    input_size : int or None, optional
        Dimension of an optional driving (exogenous) input.  When given, the
        model takes two inputs ``(feedback, driver)`` and the driver feeds the
        reservoir alongside the autoregressive feedback.  The driver is kept out
        of the input concatenation, so the readout ``in_features`` stays
        ``feedback_size + reservoir_size``.  ``None`` (default) builds a
        feedback-only model, unchanged from previous behavior.
    input_initializer : InitializerSpec, optional
        Initializer for the driving-input weights.  Same accepted forms as
        ``feedback_initializer``.  Only used when ``input_size`` is given.
    topology : TopologySpec, optional
        Topology for recurrent weights. Accepts:

        - str: Registry name (e.g., ``"erdos_renyi"``)
        - tuple: (name, params) like ``("watts_strogatz", {"k": 6, "p": 0.1})``
        - :class:`~resdag.init.topology.TopologyInitializer`: Pre-configured object
    spectral_radius : float, default=0.9
        Desired spectral radius for recurrent weights.
    leak_rate : float, default=1.0
        Leaky integration rate (1.0 = no leak).
    noise : float, default=0.0
        Standard deviation of additive Gaussian state noise injected into the
        reservoir after the activation.  Active only in training mode (a no-op
        under :meth:`~torch.nn.Module.eval`); ``0.0`` disables it.  Forwarded to
        :class:`~resdag.layers.ESNLayer`.  Must be non-negative.
    feedback_initializer : InitializerSpec, optional
        Initializer for feedback weights. Accepts:

        - str: Registry name (e.g., ``"pseudo_diagonal"``)
        - tuple: (name, params) like ``("chebyshev", {"p": 0.5})``
        - :class:`~resdag.init.input_feedback.InputFeedbackInitializer`: Pre-configured object
    activation : str, default="tanh"
        Activation function (``"tanh"``, ``"relu"``, ``"sigmoid"``, ``"identity"``).
    bias : bool, default=True
        Whether to use bias in the reservoir.
    trainable : bool, default=False
        Whether reservoir weights are trainable.
    readout : str or type, default="ridge"
        Which readout solver to attach.  ``"ridge"`` (default) is the exact
        direct Cholesky ridge solve (:class:`~resdag.layers.RidgeReadoutLayer`);
        ``"cg"`` selects the legacy iterative
        :class:`~resdag.layers.CGReadoutLayer`.  A
        :class:`~resdag.layers.readouts.ReadoutLayer` subclass may also be passed
        to use a custom readout.
    readout_alpha : float, default=1e-6
        Ridge regression regularization for readout.
    readout_bias : bool, default=True
        Whether to use bias in the readout.
    readout_name : str, default="output"
        Name for the readout layer (used in training targets).
    **reservoir_kwargs : Any
        Additional keyword arguments passed to :class:`~resdag.layers.ESNLayer`.

    Returns
    -------
    :class:`~resdag.core.ESNModel`
        Configured ESN model ready for training and inference.

    Examples
    --------
    Simple usage with defaults:

    >>> from resdag.models import classic_esn
    >>> import torch
    >>> model = classic_esn(100, 1, 1)

    With custom topology and initializer:

    >>> model = classic_esn(
    ...     reservoir_size=400,
    ...     feedback_size=1,
    ...     output_size=1,
    ...     topology=("watts_strogatz", {"k": 6, "p": 0.1}),
    ...     feedback_initializer="pseudo_diagonal",
    ...     spectral_radius=0.9,
    ...     leak_rate=0.5,
    ... )

    Forward pass:

    >>> x = torch.randn(4, 50, 1)  # (batch, time, features)
    >>> y = model(x)

    See Also
    --------
    :func:`resdag.models.ott_esn` : OTT ESN variant
    :func:`resdag.models.linear_esn` : Linear ESN variant
    :class:`resdag.training.ESNTrainer` : Trainer for fitting readouts
    """
    # Classic ESN: no augmentation, concatenate the raw input back in.
    return _esn_builder(
        reservoir_size=reservoir_size,
        feedback_size=feedback_size,
        output_size=output_size,
        augment=None,
        concat_input=True,
        input_size=input_size,
        input_initializer=input_initializer,
        topology=topology,
        spectral_radius=spectral_radius,
        leak_rate=leak_rate,
        noise=noise,
        feedback_initializer=feedback_initializer,
        activation=activation,
        bias=bias,
        trainable=trainable,
        readout=readout,
        readout_alpha=readout_alpha,
        readout_bias=readout_bias,
        readout_name=readout_name,
        **reservoir_kwargs,
    )

ott_esn

ott_esn(reservoir_size: int, feedback_size: int, output_size: int, input_size: int | None = None, input_initializer: InitializerSpec | None = None, topology: TopologySpec | None = None, spectral_radius: float = 0.9, leak_rate: float = 1.0, feedback_initializer: InitializerSpec | None = None, activation: str = 'tanh', bias: bool = True, trainable: bool = False, readout: ReadoutSpec = 'ridge', readout_alpha: float = 1e-06, readout_bias: bool = True, readout_name: str = 'output', **reservoir_kwargs: Any) -> ESNModel

Build Ott's ESN model with state augmentation.

This model follows the architecture proposed by Edward Ott, which augments reservoir states by squaring even-indexed units and concatenating with input. This augmentation helps capture higher-order dynamics in chaotic systems.

Architecture::

Input -> Reservoir -> SelectiveExponentiation -> Concatenate(Input, Augmented) -> Readout

PARAMETER	DESCRIPTION
reservoir_size	Number of units in the reservoir. TYPE: int
feedback_size	Dimension of feedback signal (input features). TYPE: int
output_size	Dimension of output signal. TYPE: int
input_size	Dimension of an optional driving (exogenous) input. When given, the model takes two inputs (feedback, driver) and the driver feeds the reservoir alongside the autoregressive feedback. The driver is kept out of the input concatenation, so the readout in_features stays feedback_size + reservoir_size. None (default) builds a feedback-only model, unchanged from previous behavior. TYPE: int or None DEFAULT: None
input_initializer	Initializer for the driving-input weights. Same format as feedback_initializer. Only used when input_size is given. TYPE: str, tuple, or InputFeedbackInitializer DEFAULT: None
topology	Topology for recurrent weights. Accepts: str: Registry name (e.g., "erdos_renyi") tuple: (name, params) like ("watts_strogatz", {"k": 6, "p": 0.1}) TopologyInitializer: Pre-configured object TYPE: str, tuple, or TopologyInitializer DEFAULT: None
spectral_radius	Target spectral radius for recurrent weights. TYPE: float DEFAULT: 0.9
leak_rate	Leaky integration rate. 1.0 = no leaking (standard ESN). TYPE: float DEFAULT: 1.0
feedback_initializer	Initializer for feedback weights. Same format as topology. TYPE: str, tuple, or InputFeedbackInitializer DEFAULT: None
activation	Activation function for reservoir neurons. TYPE: (tanh, relu, sigmoid, identity) DEFAULT: 'tanh'
bias	Whether to include bias in reservoir. TYPE: bool DEFAULT: True
trainable	If True, reservoir weights are trainable via backpropagation. TYPE: bool DEFAULT: False
readout	Which readout solver to attach. "ridge" (default) is the exact direct Cholesky ridge solve (:class:~resdag.layers.RidgeReadoutLayer); "cg" selects the legacy iterative :class:~resdag.layers.CGReadoutLayer. A :class:~resdag.layers.readouts.ReadoutLayer subclass may also be passed to use a custom readout. TYPE: str or type DEFAULT: "ridge"
readout_alpha	L2 regularization strength for ridge regression in readout. TYPE: float DEFAULT: 1e-6
readout_bias	Whether to include bias in readout layer. TYPE: bool DEFAULT: True
readout_name	Name for the readout layer. Used as target key in training. TYPE: str DEFAULT: 'output'
**reservoir_kwargs	Additional keyword arguments passed to :class:ESNLayer. TYPE: Any DEFAULT: {}

RETURNS	DESCRIPTION
ESNModel	Configured Ott ESN model ready for training and inference.

Examples:

Basic usage:

>>> from resdag.models import ott_esn
>>> model = ott_esn(
...     reservoir_size=500,
...     feedback_size=3,
...     output_size=3,
... )
>>> model.summary()

With custom topology:

>>> from resdag.init.topology import get_topology
>>> model = ott_esn(
...     reservoir_size=500,
...     feedback_size=3,
...     output_size=3,
...     topology=get_topology("watts_strogatz", k=4, p=0.3),
...     spectral_radius=0.95,
... )

Training and forecasting:

>>> from resdag.training import ESNTrainer
>>> trainer = ESNTrainer(model)
>>> trainer.fit(
...     warmup_inputs=(warmup,),
...     train_inputs=(train,),
...     targets={"output": target},
... )
>>> predictions = model.forecast(forecast_warmup, horizon=100)

See Also

classic_esn : Traditional ESN without state augmentation. resdag.training.ESNTrainer : Trainer for fitting readout. resdag.init.topology.get_topology : Get topology by name.

Source code in src/resdag/models/ott_esn.py

def ott_esn(
    reservoir_size: int,
    feedback_size: int,
    output_size: int,
    input_size: int | None = None,
    input_initializer: InitializerSpec | None = None,
    # Reservoir params
    topology: TopologySpec | None = None,
    spectral_radius: float = 0.9,
    leak_rate: float = 1.0,
    feedback_initializer: InitializerSpec | None = None,
    activation: str = "tanh",
    bias: bool = True,
    trainable: bool = False,
    # Readout params
    readout: ReadoutSpec = "ridge",
    readout_alpha: float = 1e-6,
    readout_bias: bool = True,
    readout_name: str = "output",
    # Extra reservoir kwargs
    **reservoir_kwargs: Any,
) -> ESNModel:
    """
    Build Ott's ESN model with state augmentation.

    This model follows the architecture proposed by Edward Ott, which augments
    reservoir states by squaring even-indexed units and concatenating with input.
    This augmentation helps capture higher-order dynamics in chaotic systems.

    Architecture::

        Input -> Reservoir -> SelectiveExponentiation -> Concatenate(Input, Augmented) -> Readout

    Parameters
    ----------
    reservoir_size : int
        Number of units in the reservoir.
    feedback_size : int
        Dimension of feedback signal (input features).
    output_size : int
        Dimension of output signal.
    input_size : int or None, optional
        Dimension of an optional driving (exogenous) input.  When given, the
        model takes two inputs ``(feedback, driver)`` and the driver feeds the
        reservoir alongside the autoregressive feedback.  The driver is kept out
        of the input concatenation, so the readout ``in_features`` stays
        ``feedback_size + reservoir_size``.  ``None`` (default) builds a
        feedback-only model, unchanged from previous behavior.
    input_initializer : str, tuple, or InputFeedbackInitializer, optional
        Initializer for the driving-input weights.  Same format as
        ``feedback_initializer``.  Only used when ``input_size`` is given.
    topology : str, tuple, or TopologyInitializer, optional
        Topology for recurrent weights. Accepts:

        - str: Registry name (e.g., ``"erdos_renyi"``)
        - tuple: ``(name, params)`` like ``("watts_strogatz", {"k": 6, "p": 0.1})``
        - TopologyInitializer: Pre-configured object

    spectral_radius : float, default=0.9
        Target spectral radius for recurrent weights.
    leak_rate : float, default=1.0
        Leaky integration rate. 1.0 = no leaking (standard ESN).
    feedback_initializer : str, tuple, or InputFeedbackInitializer, optional
        Initializer for feedback weights. Same format as ``topology``.
    activation : {'tanh', 'relu', 'sigmoid', 'identity'}, default='tanh'
        Activation function for reservoir neurons.
    bias : bool, default=True
        Whether to include bias in reservoir.
    trainable : bool, default=False
        If True, reservoir weights are trainable via backpropagation.
    readout : str or type, default="ridge"
        Which readout solver to attach.  ``"ridge"`` (default) is the exact
        direct Cholesky ridge solve (:class:`~resdag.layers.RidgeReadoutLayer`);
        ``"cg"`` selects the legacy iterative
        :class:`~resdag.layers.CGReadoutLayer`.  A
        :class:`~resdag.layers.readouts.ReadoutLayer` subclass may also be passed
        to use a custom readout.
    readout_alpha : float, default=1e-6
        L2 regularization strength for ridge regression in readout.
    readout_bias : bool, default=True
        Whether to include bias in readout layer.
    readout_name : str, default='output'
        Name for the readout layer. Used as target key in training.
    **reservoir_kwargs
        Additional keyword arguments passed to :class:`ESNLayer`.

    Returns
    -------
    ESNModel
        Configured Ott ESN model ready for training and inference.

    Examples
    --------
    Basic usage:

    >>> from resdag.models import ott_esn
    >>> model = ott_esn(
    ...     reservoir_size=500,
    ...     feedback_size=3,
    ...     output_size=3,
    ... )
    >>> model.summary()

    With custom topology:

    >>> from resdag.init.topology import get_topology
    >>> model = ott_esn(
    ...     reservoir_size=500,
    ...     feedback_size=3,
    ...     output_size=3,
    ...     topology=get_topology("watts_strogatz", k=4, p=0.3),
    ...     spectral_radius=0.95,
    ... )

    Training and forecasting:

    >>> from resdag.training import ESNTrainer
    >>> trainer = ESNTrainer(model)
    >>> trainer.fit(
    ...     warmup_inputs=(warmup,),
    ...     train_inputs=(train,),
    ...     targets={"output": target},
    ... )
    >>> predictions = model.forecast(forecast_warmup, horizon=100)

    See Also
    --------
    classic_esn : Traditional ESN without state augmentation.
    resdag.training.ESNTrainer : Trainer for fitting readout.
    resdag.init.topology.get_topology : Get topology by name.
    """
    # Ott augmentation: square even-indexed reservoir units, concatenate input.
    return _esn_builder(
        reservoir_size=reservoir_size,
        feedback_size=feedback_size,
        output_size=output_size,
        augment=lambda: SelectiveExponentiation(index=0, exponent=2.0),
        concat_input=True,
        input_size=input_size,
        input_initializer=input_initializer,
        topology=topology,
        spectral_radius=spectral_radius,
        leak_rate=leak_rate,
        feedback_initializer=feedback_initializer,
        activation=activation,
        bias=bias,
        trainable=trainable,
        readout=readout,
        readout_alpha=readout_alpha,
        readout_bias=readout_bias,
        readout_name=readout_name,
        **reservoir_kwargs,
    )

power_augmented

power_augmented(reservoir_size: int, feedback_size: int, output_size: int, exponent: float = 3.0, input_size: int | None = None, input_initializer: InitializerSpec | None = None, topology: TopologySpec | None = None, spectral_radius: float = 0.9, leak_rate: float = 1.0, feedback_initializer: InitializerSpec | None = None, activation: str = 'tanh', bias: bool = True, trainable: bool = False, readout: ReadoutSpec = 'ridge', readout_alpha: float = 1e-06, readout_bias: bool = True, readout_name: str = 'output', **reservoir_kwargs: Any) -> ESNModel

Build Power Augmented ESN model.

This model augments reservoir states by exponentiating to a power and concatenating with input. This augmentation helps capture higher-order dynamics in chaotic systems.

Architecture::

Input -> Reservoir -> Power -> Concatenate(Input, Augmented) -> Readout

PARAMETER	DESCRIPTION
reservoir_size	Number of units in the reservoir. TYPE: int
feedback_size	Dimension of feedback signal (input features). TYPE: int
output_size	Dimension of output signal. TYPE: int
exponent	Exponent applied to every reservoir state. Tanh reservoir states live in [-1, 1] and routinely include negative and zero values, so the exponent must be chosen with that signed range in mind (see Notes). The default 3.0 is an odd integer and therefore sign-preserving; an even exponent such as 2.0 maps every state to a non-negative value and so discards the sign of the reservoir activations. TYPE: float DEFAULT: 3.0
input_size	Dimension of an optional driving (exogenous) input. When given, the model takes two inputs (feedback, driver) and the driver feeds the reservoir alongside the autoregressive feedback. The driver is kept out of the input concatenation, so the readout in_features stays feedback_size + reservoir_size. None (default) builds a feedback-only model, unchanged from previous behavior. TYPE: int or None DEFAULT: None
input_initializer	Initializer for the driving-input weights. Same format as feedback_initializer. Only used when input_size is given. TYPE: str, tuple, or InputFeedbackInitializer DEFAULT: None
topology	Topology for recurrent weights. Accepts: str: Registry name (e.g., "erdos_renyi") tuple: (name, params) like ("watts_strogatz", {"k": 6, "p": 0.1}) TopologyInitializer: Pre-configured object TYPE: str, tuple, or TopologyInitializer DEFAULT: None
spectral_radius	Target spectral radius for recurrent weights. TYPE: float DEFAULT: 0.9
leak_rate	Leaky integration rate. 1.0 = no leaking (standard ESN). TYPE: float DEFAULT: 1.0
feedback_initializer	Initializer for feedback weights. Same format as topology. TYPE: str, tuple, or InputFeedbackInitializer DEFAULT: None
activation	Activation function for reservoir neurons. TYPE: (tanh, relu, sigmoid, identity) DEFAULT: 'tanh'
bias	Whether to include bias in reservoir. TYPE: bool DEFAULT: True
trainable	If True, reservoir weights are trainable via backpropagation. TYPE: bool DEFAULT: False
readout	Which readout solver to attach. "ridge" (default) is the exact direct Cholesky ridge solve (:class:~resdag.layers.RidgeReadoutLayer); "cg" selects the legacy iterative :class:~resdag.layers.CGReadoutLayer. A :class:~resdag.layers.readouts.ReadoutLayer subclass may also be passed to use a custom readout. TYPE: str or type DEFAULT: "ridge"
readout_alpha	L2 regularization strength for ridge regression in readout. TYPE: float DEFAULT: 1e-6
readout_bias	Whether to include bias in readout layer. TYPE: bool DEFAULT: True
readout_name	Name for the readout layer. Used as target key in training. TYPE: str DEFAULT: 'output'
**reservoir_kwargs	Additional keyword arguments passed to :class:ESNLayer. TYPE: Any DEFAULT: {}

RETURNS	DESCRIPTION
ESNModel	Configured Power Augmented ESN model ready for training and inference.

Notes

The augmentation applies :class:~resdag.layers.Power to the raw reservoir states, which for a tanh activation lie in [-1, 1] — negatives and zeros included. Choose exponent accordingly:

• Even integers (2.0, 4.0, …) are always safe: every state maps to a non-negative value with no nan/inf.
• Odd integers (3.0, 5.0, …) are safe and sign-preserving: negative states stay negative.
• Non-integer exponents (e.g. 0.5, 1.5) applied to a negative state produce nan under the default :func:torch.pow, silently corrupting the readout inputs. Negative exponents (e.g. -1.0) produce inf on the zeros that tanh states pass through.

If you need a non-integer exponent on signed states, wire the model by hand with Power(exponent, sign_preserving=True), which applies sign(x) * abs(x) ** exponent and stays finite for negative bases. This factory always uses the default (non-sign-preserving) Power to keep the well-established integer-exponent behaviour unchanged.

Examples:

Basic usage:

>>> from resdag.models import power_augmented
>>> model = power_augmented(
...     reservoir_size=500,
...     feedback_size=3,
...     output_size=3,
...     exponent=3.0,
... )
>>> model.summary()

With custom topology:

>>> from resdag.init.topology import get_topology
>>> model = power_augmented(
...     reservoir_size=500,
...     feedback_size=3,
...     output_size=3,
...     topology=get_topology("watts_strogatz", k=4, p=0.3),
...     spectral_radius=0.95,
...     exponent=3.0,
... )

Training and forecasting:

>>> from resdag.training import ESNTrainer
>>> trainer = ESNTrainer(model)
>>> trainer.fit(
...     warmup_inputs=(warmup,),
...     train_inputs=(train,),
...     targets={"output": target},
... )
>>> predictions = model.forecast(forecast_warmup, horizon=100)

See Also

classic_esn : Traditional ESN without state augmentation. resdag.training.ESNTrainer : Trainer for fitting readout. resdag.init.topology.get_topology : Get topology by name.

Source code in src/resdag/models/power_augmented.py

def power_augmented(
    reservoir_size: int,
    feedback_size: int,
    output_size: int,
    exponent: float = 3.0,
    input_size: int | None = None,
    input_initializer: InitializerSpec | None = None,
    # Reservoir params
    topology: TopologySpec | None = None,
    spectral_radius: float = 0.9,
    leak_rate: float = 1.0,
    feedback_initializer: InitializerSpec | None = None,
    activation: str = "tanh",
    bias: bool = True,
    trainable: bool = False,
    # Readout params
    readout: ReadoutSpec = "ridge",
    readout_alpha: float = 1e-6,
    readout_bias: bool = True,
    readout_name: str = "output",
    # Extra reservoir kwargs
    **reservoir_kwargs: Any,
) -> ESNModel:
    """
    Build Power Augmented ESN model.

    This model augments reservoir states by exponentiating to a power and concatenating with input.
    This augmentation helps capture higher-order dynamics in chaotic systems.

    Architecture::

        Input -> Reservoir -> Power -> Concatenate(Input, Augmented) -> Readout

    Parameters
    ----------
    reservoir_size : int
        Number of units in the reservoir.
    feedback_size : int
        Dimension of feedback signal (input features).
    output_size : int
        Dimension of output signal.
    exponent : float, default=3.0
        Exponent applied to **every** reservoir state. Tanh reservoir states
        live in ``[-1, 1]`` and routinely include negative and zero values, so
        the exponent must be chosen with that signed range in mind (see Notes).
        The default ``3.0`` is an odd integer and therefore *sign-preserving*;
        an even exponent such as ``2.0`` maps every state to a non-negative
        value and so discards the sign of the reservoir activations.
    input_size : int or None, optional
        Dimension of an optional driving (exogenous) input.  When given, the
        model takes two inputs ``(feedback, driver)`` and the driver feeds the
        reservoir alongside the autoregressive feedback.  The driver is kept out
        of the input concatenation, so the readout ``in_features`` stays
        ``feedback_size + reservoir_size``.  ``None`` (default) builds a
        feedback-only model, unchanged from previous behavior.
    input_initializer : str, tuple, or InputFeedbackInitializer, optional
        Initializer for the driving-input weights.  Same format as
        ``feedback_initializer``.  Only used when ``input_size`` is given.
    topology : str, tuple, or TopologyInitializer, optional
        Topology for recurrent weights. Accepts:

        - str: Registry name (e.g., ``"erdos_renyi"``)
        - tuple: ``(name, params)`` like ``("watts_strogatz", {"k": 6, "p": 0.1})``
        - TopologyInitializer: Pre-configured object

    spectral_radius : float, default=0.9
        Target spectral radius for recurrent weights.
    leak_rate : float, default=1.0
        Leaky integration rate. 1.0 = no leaking (standard ESN).
    feedback_initializer : str, tuple, or InputFeedbackInitializer, optional
        Initializer for feedback weights. Same format as ``topology``.
    activation : {'tanh', 'relu', 'sigmoid', 'identity'}, default='tanh'
        Activation function for reservoir neurons.
    bias : bool, default=True
        Whether to include bias in reservoir.
    trainable : bool, default=False
        If True, reservoir weights are trainable via backpropagation.
    readout : str or type, default="ridge"
        Which readout solver to attach.  ``"ridge"`` (default) is the exact
        direct Cholesky ridge solve (:class:`~resdag.layers.RidgeReadoutLayer`);
        ``"cg"`` selects the legacy iterative
        :class:`~resdag.layers.CGReadoutLayer`.  A
        :class:`~resdag.layers.readouts.ReadoutLayer` subclass may also be passed
        to use a custom readout.
    readout_alpha : float, default=1e-6
        L2 regularization strength for ridge regression in readout.
    readout_bias : bool, default=True
        Whether to include bias in readout layer.
    readout_name : str, default='output'
        Name for the readout layer. Used as target key in training.
    **reservoir_kwargs
        Additional keyword arguments passed to :class:`ESNLayer`.

    Returns
    -------
    ESNModel
        Configured Power Augmented ESN model ready for training and inference.

    Notes
    -----
    The augmentation applies :class:`~resdag.layers.Power` to the raw reservoir
    states, which for a ``tanh`` activation lie in ``[-1, 1]`` — negatives and
    zeros included. Choose ``exponent`` accordingly:

    - **Even integers** (``2.0``, ``4.0``, …) are always safe: every state maps
      to a non-negative value with no ``nan``/``inf``.
    - **Odd integers** (``3.0``, ``5.0``, …) are safe and *sign-preserving*:
      negative states stay negative.
    - **Non-integer exponents** (e.g. ``0.5``, ``1.5``) applied to a *negative*
      state produce ``nan`` under the default :func:`torch.pow`, silently
      corrupting the readout inputs. **Negative exponents** (e.g. ``-1.0``)
      produce ``inf`` on the zeros that ``tanh`` states pass through.

    If you need a non-integer exponent on signed states, wire the model by hand
    with ``Power(exponent, sign_preserving=True)``, which applies
    ``sign(x) * abs(x) ** exponent`` and stays finite for negative bases. This
    factory always uses the default (non-sign-preserving) ``Power`` to keep the
    well-established integer-exponent behaviour unchanged.

    Examples
    --------
    Basic usage:

    >>> from resdag.models import power_augmented
    >>> model = power_augmented(
    ...     reservoir_size=500,
    ...     feedback_size=3,
    ...     output_size=3,
    ...     exponent=3.0,
    ... )
    >>> model.summary()

    With custom topology:

    >>> from resdag.init.topology import get_topology
    >>> model = power_augmented(
    ...     reservoir_size=500,
    ...     feedback_size=3,
    ...     output_size=3,
    ...     topology=get_topology("watts_strogatz", k=4, p=0.3),
    ...     spectral_radius=0.95,
    ...     exponent=3.0,
    ... )

    Training and forecasting:

    >>> from resdag.training import ESNTrainer
    >>> trainer = ESNTrainer(model)
    >>> trainer.fit(
    ...     warmup_inputs=(warmup,),
    ...     train_inputs=(train,),
    ...     targets={"output": target},
    ... )
    >>> predictions = model.forecast(forecast_warmup, horizon=100)

    See Also
    --------
    classic_esn : Traditional ESN without state augmentation.
    resdag.training.ESNTrainer : Trainer for fitting readout.
    resdag.init.topology.get_topology : Get topology by name.
    """
    # Power augmentation: raise reservoir states to ``exponent``, concat input.
    return _esn_builder(
        reservoir_size=reservoir_size,
        feedback_size=feedback_size,
        output_size=output_size,
        augment=lambda: Power(exponent=exponent),
        concat_input=True,
        input_size=input_size,
        input_initializer=input_initializer,
        topology=topology,
        spectral_radius=spectral_radius,
        leak_rate=leak_rate,
        feedback_initializer=feedback_initializer,
        activation=activation,
        bias=bias,
        trainable=trainable,
        readout=readout,
        readout_alpha=readout_alpha,
        readout_bias=readout_bias,
        readout_name=readout_name,
        **reservoir_kwargs,
    )

linear_esn

linear_esn(reservoir_size: int, feedback_size: int, topology: TopologySpec | None = None, spectral_radius: float = 0.9, leak_rate: float = 1.0, feedback_initializer: InitializerSpec | None = None, bias: bool = True, trainable: bool = False, **reservoir_kwargs: Any) -> ESNModel

Build an ESN model with no readout layer and a linear activation function.

This model uses a linear activation function in the reservoir, which can be useful for studying linear dynamics or as a baseline for comparison with nonlinear reservoirs.

Architecture: Input -> Reservoir(activation='identity') (output)

PARAMETER	DESCRIPTION
reservoir_size	Number of units in the reservoir. TYPE: int
feedback_size	Number of feedback features. TYPE: int
topology	Topology for recurrent weights. Accepts: - str: Registry name (e.g., "erdos_renyi") - tuple: (name, params) like ("watts_strogatz", {"k": 6, "p": 0.1}) - GraphTopology: Pre-configured object TYPE: TopologySpec DEFAULT: None
spectral_radius	Desired spectral radius for recurrent weights. TYPE: float DEFAULT: 0.9
leak_rate	Leaky integration rate (1.0 = no leak). TYPE: float DEFAULT: 1.0
feedback_initializer	Initializer for feedback weights. TYPE: InitializerSpec DEFAULT: None
bias	Whether to use bias in the reservoir. TYPE: bool DEFAULT: True
trainable	Whether reservoir weights are trainable. TYPE: bool DEFAULT: False
**reservoir_kwargs	Additional keyword arguments passed to ESNLayer. activation is not accepted here: linear_esn forces activation="identity" by design. TYPE: Any DEFAULT: {}

RETURNS	DESCRIPTION
ESNModel	ESN model with linear reservoir activation.

RAISES	DESCRIPTION
ValueError	If activation is passed via reservoir_kwargs. linear_esn forces a linear (identity) activation by design, so overriding it is not allowed.

Examples:

>>> from resdag.models import linear_esn
>>> model = linear_esn(100, 1)
>>> linear_states = model(input_data)

Source code in src/resdag/models/linear_esn.py

def linear_esn(
    reservoir_size: int,
    feedback_size: int,
    # Reservoir params
    topology: TopologySpec | None = None,
    spectral_radius: float = 0.9,
    leak_rate: float = 1.0,
    feedback_initializer: InitializerSpec | None = None,
    bias: bool = True,
    trainable: bool = False,
    # Extra reservoir kwargs
    **reservoir_kwargs: Any,
) -> ESNModel:
    """Build an ESN model with no readout layer and a linear activation function.

    This model uses a linear activation function in the reservoir, which can be
    useful for studying linear dynamics or as a baseline for comparison with
    nonlinear reservoirs.

    Architecture:
        Input -> Reservoir(activation='identity') (output)

    Parameters
    ----------
    reservoir_size : int
        Number of units in the reservoir.
    feedback_size : int
        Number of feedback features.
    topology : TopologySpec, optional
        Topology for recurrent weights. Accepts:
        - str: Registry name (e.g., "erdos_renyi")
        - tuple: (name, params) like ("watts_strogatz", {"k": 6, "p": 0.1})
        - GraphTopology: Pre-configured object
    spectral_radius : float, default=0.9
        Desired spectral radius for recurrent weights.
    leak_rate : float, default=1.0
        Leaky integration rate (1.0 = no leak).
    feedback_initializer : InitializerSpec, optional
        Initializer for feedback weights.
    bias : bool, default=True
        Whether to use bias in the reservoir.
    trainable : bool, default=False
        Whether reservoir weights are trainable.
    **reservoir_kwargs
        Additional keyword arguments passed to ESNLayer.  ``activation`` is not
        accepted here: ``linear_esn`` forces ``activation="identity"`` by design.

    Returns
    -------
    ESNModel
        ESN model with linear reservoir activation.

    Raises
    ------
    ValueError
        If ``activation`` is passed via ``reservoir_kwargs``.  ``linear_esn``
        forces a linear (identity) activation by design, so overriding it is
        not allowed.

    Examples
    --------
    >>> from resdag.models import linear_esn
    >>> model = linear_esn(100, 1)
    >>> linear_states = model(input_data)
    """
    # Linear: reservoir output only, no readout, identity activation forced.
    # Forbid an explicit ``activation`` override rather than letting it collide
    # with the hardcoded ``activation="identity"`` below (which would otherwise
    # raise a cryptic "multiple values for keyword argument 'activation'").
    if "activation" in reservoir_kwargs:
        raise ValueError(
            "linear_esn forces activation='identity'; do not pass activation="
            f"{reservoir_kwargs['activation']!r}. Use a different premade model "
            "(e.g. headless_esn) for a nonlinear reservoir without a readout."
        )

    return _esn_builder(
        reservoir_size=reservoir_size,
        feedback_size=feedback_size,
        output_size=None,
        augment=None,
        topology=topology,
        spectral_radius=spectral_radius,
        leak_rate=leak_rate,
        feedback_initializer=feedback_initializer,
        activation="identity",  # Force linear activation
        bias=bias,
        trainable=trainable,
        **reservoir_kwargs,
    )

headless_esn

headless_esn(reservoir_size: int, feedback_size: int, topology: TopologySpec | None = None, spectral_radius: float = 0.9, leak_rate: float = 1.0, feedback_initializer: InitializerSpec | None = None, activation: str = 'tanh', bias: bool = True, trainable: bool = False, **reservoir_kwargs: Any) -> ESNModel

Build an ESN model with no readout layer.

This model can be used to study the dynamics of the reservoir by applying different transformations to the reservoir states without a readout layer. Useful for analyzing reservoir dynamics, state space properties, and feature extraction.

Architecture: Input -> Reservoir (output)

The reservoir is not connected to a readout layer, allowing direct access to reservoir states for analysis or custom processing.

PARAMETER	DESCRIPTION
reservoir_size	Number of units in the reservoir. TYPE: int
feedback_size	Number of feedback features. TYPE: int
topology	Topology for recurrent weights. Accepts: - str: Registry name (e.g., "erdos_renyi") - tuple: (name, params) like ("watts_strogatz", {"k": 6, "p": 0.1}) - GraphTopology: Pre-configured object TYPE: TopologySpec DEFAULT: None
spectral_radius	Desired spectral radius for recurrent weights. TYPE: float DEFAULT: 0.9
leak_rate	Leaky integration rate (1.0 = no leak). TYPE: float DEFAULT: 1.0
feedback_initializer	Initializer for feedback weights. TYPE: InitializerSpec DEFAULT: None
activation	Activation function ("tanh", "relu", "sigmoid", "identity"). TYPE: str DEFAULT: "tanh"
bias	Whether to use bias in the reservoir. TYPE: bool DEFAULT: True
trainable	Whether reservoir weights are trainable. TYPE: bool DEFAULT: False
**reservoir_kwargs	Additional keyword arguments passed to ESNLayer. TYPE: Any DEFAULT: {}

RETURNS	DESCRIPTION
ESNModel	ESN model with reservoir output only.

Examples:

>>> from resdag.models import headless_esn
>>> model = headless_esn(100, 1)
>>> reservoir_states = model(input_data)  # Direct reservoir output

Source code in src/resdag/models/headless_esn.py

def headless_esn(
    reservoir_size: int,
    feedback_size: int,
    # Reservoir params
    topology: TopologySpec | None = None,
    spectral_radius: float = 0.9,
    leak_rate: float = 1.0,
    feedback_initializer: InitializerSpec | None = None,
    activation: str = "tanh",
    bias: bool = True,
    trainable: bool = False,
    # Extra reservoir kwargs
    **reservoir_kwargs: Any,
) -> ESNModel:
    """Build an ESN model with no readout layer.

    This model can be used to study the dynamics of the reservoir by applying
    different transformations to the reservoir states without a readout layer.
    Useful for analyzing reservoir dynamics, state space properties, and
    feature extraction.

    Architecture:
        Input -> Reservoir (output)

    The reservoir is not connected to a readout layer, allowing direct
    access to reservoir states for analysis or custom processing.

    Parameters
    ----------
    reservoir_size : int
        Number of units in the reservoir.
    feedback_size : int
        Number of feedback features.
    topology : TopologySpec, optional
        Topology for recurrent weights. Accepts:
        - str: Registry name (e.g., "erdos_renyi")
        - tuple: (name, params) like ("watts_strogatz", {"k": 6, "p": 0.1})
        - GraphTopology: Pre-configured object
    spectral_radius : float, default=0.9
        Desired spectral radius for recurrent weights.
    leak_rate : float, default=1.0
        Leaky integration rate (1.0 = no leak).
    feedback_initializer : InitializerSpec, optional
        Initializer for feedback weights.
    activation : str, default="tanh"
        Activation function ("tanh", "relu", "sigmoid", "identity").
    bias : bool, default=True
        Whether to use bias in the reservoir.
    trainable : bool, default=False
        Whether reservoir weights are trainable.
    **reservoir_kwargs
        Additional keyword arguments passed to ESNLayer.

    Returns
    -------
    ESNModel
        ESN model with reservoir output only.

    Examples
    --------
    >>> from resdag.models import headless_esn
    >>> model = headless_esn(100, 1)
    >>> reservoir_states = model(input_data)  # Direct reservoir output
    """
    # Headless: reservoir output only, no augmentation and no readout.
    return _esn_builder(
        reservoir_size=reservoir_size,
        feedback_size=feedback_size,
        output_size=None,
        augment=None,
        topology=topology,
        spectral_radius=spectral_radius,
        leak_rate=leak_rate,
        feedback_initializer=feedback_initializer,
        activation=activation,
        bias=bias,
        trainable=trainable,
        **reservoir_kwargs,
    )

coupled_ensemble_esn

coupled_ensemble_esn(n_models: int, model_factory: Callable[..., ESNModel] = ott_esn, aggregate: str | Module = 'mean', seed: int | None = None, **model_kwargs: Any) -> CoupledEnsembleESNModel

Build a coupled ensemble of N independently-initialized ESN models.

Each sub-model is created by calling model_factory(**model_kwargs) with an independent random reservoir initialization, providing the diversity needed for ensemble averaging to be beneficial.

During autoregressive forecasting all N sub-models receive the same aggregated output (e.g. the mean across models) as their next feedback input — they are coupled through the shared signal.

PARAMETER	DESCRIPTION
n_models	Number of sub-models in the ensemble. TYPE: int
model_factory	Factory that creates one :class:~resdag.core.ESNModel. All keyword arguments not consumed by coupled_ensemble_esn itself (i.e. model_kwargs) are forwarded verbatim to every call of this factory. The factory must return a readout-bearing model whose first output dimension equals the feedback input dimension, so the aggregated output can be fed back autoregressively. The readout-bearing premade factories — classic_esn, ott_esn and power_augmented — satisfy this. The headless factories (headless_esn, linear_esn) emit raw reservoir states rather than a feedback-dimensioned output and are rejected by :class:~resdag.ensemble.CoupledEnsembleESNModel (see its __init__ for the exact check). TYPE: callable DEFAULT: :func:`ott_esn`
aggregate	Aggregation strategy applied to the N model outputs at each autoregressive step: "mean" — arithmetic mean. "median" — median. nn.Module — any module accepting a stacked tensor of shape (N, batch, timesteps, features) and returning (batch, timesteps, features). E.g. :class:~resdag.ensemble.aggregators.OutliersFilteredMean. TYPE: str or Module DEFAULT: ``"mean"``
seed	Master seed for deterministic sub-model construction. Sub-model i is seeded with seed + i immediately before the factory call. If model_factory itself accepts a seed keyword argument, seed + i is also forwarded as that argument so initialisers that take an explicit seed (e.g. graph topologies) become reproducible. The process-global default RNG is not clobbered: its state is captured before the construction loop and restored on exit (even if a factory raises), so building a seeded ensemble does not perturb a subsequent global torch.randn draw. Ensemble construction is therefore composable inside an otherwise-reproducible pipeline. TYPE: int DEFAULT: None
**model_kwargs	All remaining keyword arguments are forwarded verbatim to each model_factory call. This includes architecture parameters such as reservoir_size, feedback_size, output_size, spectral_radius, readout_alpha, readout_name, etc. The exact set of valid keys depends on the chosen model_factory. TYPE: Any DEFAULT: {}

RETURNS	DESCRIPTION
CoupledEnsembleESNModel	Ensemble ready for training and forecasting.

Examples:

Quick start with default Ott's ESN sub-models:

>>> import resdag as rd
>>>
>>> ensemble = rd.coupled_ensemble_esn(
...     n_models=5,
...     reservoir_size=300,
...     feedback_size=3,
...     output_size=3,
... )
>>> ensemble.fit((warmup,), (train,), {"output": targets})
>>> ensemble.reset_reservoirs()
>>> preds = ensemble.forecast(forecast_warmup, horizon=200)
>>> preds.shape
torch.Size([1, 200, 3])

Using classic_esn as the base architecture:

>>> from resdag.models import classic_esn
>>> ensemble = rd.coupled_ensemble_esn(
...     n_models=8,
...     model_factory=classic_esn,
...     reservoir_size=500,
...     feedback_size=3,
...     output_size=3,
...     spectral_radius=0.95,
...     readout_alpha=1e-7,
... )

Outlier-robust aggregation:

>>> from resdag.ensemble.aggregators import OutliersFilteredMean
>>> ensemble = rd.coupled_ensemble_esn(
...     n_models=10,
...     reservoir_size=300,
...     feedback_size=3,
...     output_size=3,
...     aggregate=OutliersFilteredMean(method="z_score", threshold=2.0),
... )

Recovering individual sub-model trajectories for post-hoc analysis:

>>> ensemble.reset_reservoirs()
>>> preds, individuals = ensemble.forecast(
...     forecast_warmup, horizon=200, return_individuals=True
... )
>>> preds.shape          # averaged forecast
torch.Size([1, 200, 3])
>>> len(individuals)     # one tensor per sub-model
10
>>> individuals[0].shape
torch.Size([1, 200, 3])

See Also

CoupledEnsembleESNModel : The ensemble class with full API documentation. resdag.ensemble.aggregators.OutliersFilteredMean : Outlier-robust aggregation layer. resdag.models.ott_esn : Default sub-model factory.

Source code in src/resdag/models/coupled_ensemble_esn.py

def coupled_ensemble_esn(
    n_models: int,
    model_factory: Callable[..., ESNModel] = ott_esn,
    aggregate: str | nn.Module = "mean",
    seed: int | None = None,
    **model_kwargs: Any,
) -> CoupledEnsembleESNModel:
    """Build a coupled ensemble of N independently-initialized ESN models.

    Each sub-model is created by calling ``model_factory(**model_kwargs)``
    with an independent random reservoir initialization, providing the
    diversity needed for ensemble averaging to be beneficial.

    During autoregressive forecasting all N sub-models receive the **same
    aggregated output** (e.g. the mean across models) as their next feedback
    input — they are coupled through the shared signal.

    Parameters
    ----------
    n_models : int
        Number of sub-models in the ensemble.
    model_factory : callable, default :func:`ott_esn`
        Factory that creates one :class:`~resdag.core.ESNModel`.
        All keyword arguments not consumed by ``coupled_ensemble_esn`` itself
        (i.e. ``model_kwargs``) are forwarded verbatim to every call of this
        factory.  The factory must return a **readout-bearing** model whose
        first output dimension equals the feedback input dimension, so the
        aggregated output can be fed back autoregressively.  The readout-bearing
        premade factories — ``classic_esn``, ``ott_esn`` and ``power_augmented``
        — satisfy this.  The headless factories (``headless_esn``,
        ``linear_esn``) emit raw reservoir states rather than a
        feedback-dimensioned output and are rejected by
        :class:`~resdag.ensemble.CoupledEnsembleESNModel` (see its ``__init__``
        for the exact check).
    aggregate : str or nn.Module, default ``"mean"``
        Aggregation strategy applied to the N model outputs at each
        autoregressive step:

        - ``"mean"`` — arithmetic mean.
        - ``"median"`` — median.
        - ``nn.Module`` — any module accepting a stacked tensor of shape
          ``(N, batch, timesteps, features)`` and returning
          ``(batch, timesteps, features)``.  E.g.
          :class:`~resdag.ensemble.aggregators.OutliersFilteredMean`.
    seed : int, optional
        Master seed for deterministic sub-model construction.  Sub-model
        ``i`` is seeded with ``seed + i`` immediately before the factory call.
        If ``model_factory`` itself accepts a ``seed`` keyword argument,
        ``seed + i`` is also forwarded as that argument so initialisers that
        take an explicit seed (e.g. graph topologies) become reproducible.

        The process-global default RNG is **not** clobbered: its state is
        captured before the construction loop and restored on exit (even if a
        factory raises), so building a seeded ensemble does not perturb a
        subsequent global ``torch.randn`` draw.  Ensemble construction is
        therefore composable inside an otherwise-reproducible pipeline.

    **model_kwargs
        All remaining keyword arguments are forwarded verbatim to each
        ``model_factory`` call.  This includes architecture parameters such as
        ``reservoir_size``, ``feedback_size``, ``output_size``,
        ``spectral_radius``, ``readout_alpha``, ``readout_name``, etc.
        The exact set of valid keys depends on the chosen ``model_factory``.

    Returns
    -------
    CoupledEnsembleESNModel
        Ensemble ready for training and forecasting.

    Examples
    --------
    Quick start with default Ott's ESN sub-models:

    >>> import resdag as rd
    >>>
    >>> ensemble = rd.coupled_ensemble_esn(
    ...     n_models=5,
    ...     reservoir_size=300,
    ...     feedback_size=3,
    ...     output_size=3,
    ... )
    >>> ensemble.fit((warmup,), (train,), {"output": targets})
    >>> ensemble.reset_reservoirs()
    >>> preds = ensemble.forecast(forecast_warmup, horizon=200)
    >>> preds.shape
    torch.Size([1, 200, 3])

    Using ``classic_esn`` as the base architecture:

    >>> from resdag.models import classic_esn
    >>> ensemble = rd.coupled_ensemble_esn(
    ...     n_models=8,
    ...     model_factory=classic_esn,
    ...     reservoir_size=500,
    ...     feedback_size=3,
    ...     output_size=3,
    ...     spectral_radius=0.95,
    ...     readout_alpha=1e-7,
    ... )

    Outlier-robust aggregation:

    >>> from resdag.ensemble.aggregators import OutliersFilteredMean
    >>> ensemble = rd.coupled_ensemble_esn(
    ...     n_models=10,
    ...     reservoir_size=300,
    ...     feedback_size=3,
    ...     output_size=3,
    ...     aggregate=OutliersFilteredMean(method="z_score", threshold=2.0),
    ... )

    Recovering individual sub-model trajectories for post-hoc analysis:

    >>> ensemble.reset_reservoirs()
    >>> preds, individuals = ensemble.forecast(
    ...     forecast_warmup, horizon=200, return_individuals=True
    ... )
    >>> preds.shape          # averaged forecast
    torch.Size([1, 200, 3])
    >>> len(individuals)     # one tensor per sub-model
    10
    >>> individuals[0].shape
    torch.Size([1, 200, 3])

    See Also
    --------
    CoupledEnsembleESNModel : The ensemble class with full API documentation.
    resdag.ensemble.aggregators.OutliersFilteredMean : Outlier-robust aggregation layer.
    resdag.models.ott_esn : Default sub-model factory.
    """
    # Detect whether the factory accepts a ``seed`` kwarg so we can forward
    # the per-model seed to initialisers that take an explicit one (graph
    # topologies, etc.). When it doesn't, the per-model RNG seed below still
    # controls every non-seeded random draw.
    factory_accepts_seed = (
        seed is not None and "seed" in inspect.signature(model_factory).parameters
    )

    models: list[ESNModel] = []
    if seed is None:
        # No seeding requested: nothing touches the global RNG beyond what the
        # factory itself draws, so no save/restore dance is needed.
        for _ in range(n_models):
            models.append(model_factory(**model_kwargs))
        return CoupledEnsembleESNModel(models, aggregator=aggregate)

    # Seeded path. Most reservoir/feedback initialisers draw from PyTorch's
    # process-global default generator, so we still seed it per sub-model with
    # ``seed + i``. To keep ensemble construction composable inside a larger
    # reproducible pipeline we snapshot the global RNG state up front and
    # restore it in ``finally`` — building the ensemble leaves the global
    # generator exactly where it was, even if a factory call raises.
    global_rng_state = torch.get_rng_state()
    try:
        for i in range(n_models):
            sub_seed = seed + i
            torch.manual_seed(sub_seed)
            kwargs = dict(model_kwargs)
            if factory_accepts_seed:
                kwargs.setdefault("seed", sub_seed)
            models.append(model_factory(**kwargs))
    finally:
        torch.set_rng_state(global_rng_state)
    return CoupledEnsembleESNModel(models, aggregator=aggregate)

---

ensemble

Ensemble Reservoir Computing Models

This module provides ensemble wrappers for groups of independently-trained :class:~resdag.core.ESNModel instances.

CLASS	DESCRIPTION
CoupledEnsembleESNModel	N independent ESN sub-models whose autoregressive forecasts are coupled through a shared averaged feedback signal at every timestep.

See Also

resdag.models.coupled_ensemble_esn : Factory function for quick construction.

CoupledEnsembleESNModel

CoupledEnsembleESNModel(models: list[ESNModel], aggregator: str | Module = 'mean')

Bases: Module

Ensemble of N independently-trained ESN models with coupled feedback.

Each sub-model is a complete :class:~resdag.core.ESNModel built via the pytorch_symbolic API. The models are trained independently but coupled during forecasting: at every autoregressive step every model receives the same aggregated output (e.g. the mean across models) as its next feedback input.

PARAMETER	DESCRIPTION
models	N pre-built ESN sub-models. Diversity comes from their independent random reservoir initialization. TYPE: list of ESNModel
aggregator	How to combine the N per-model outputs into a single feedback tensor. "mean" — arithmetic mean across models. "median" — interpolated (statistical) median across models; for an even number of models the two central values are averaged. Any nn.Module that accepts a stacked tensor of shape (N, batch, timesteps, features) and returns (batch, timesteps, features). E.g. :class:~resdag.ensemble.aggregators.OutliersFilteredMean. TYPE: str or Module DEFAULT: ``"mean"``

Examples:

Typical use via the factory:

>>> from resdag.models import coupled_ensemble_esn
>>> ensemble = coupled_ensemble_esn(n_models=5, reservoir_size=300,
...                                feedback_size=3, output_size=3)
>>> ensemble.fit((warmup,), (train,), {"output": targets})
>>> ensemble.reset_reservoirs()
>>> preds = ensemble.forecast(forecast_warmup, horizon=200)
>>> preds.shape
torch.Size([1, 200, 3])

With a custom aggregator:

>>> from resdag.ensemble.aggregators import OutliersFilteredMean
>>> ensemble = coupled_ensemble_esn(
...     n_models=10, reservoir_size=300, feedback_size=3, output_size=3,
...     aggregate=OutliersFilteredMean(method="z_score", threshold=2.0),
... )

See Also

resdag.models.coupled_ensemble_esn : Convenience factory function. resdag.ensemble.aggregators.OutliersFilteredMean : Outlier-robust aggregation layer.

PARAMETER	DESCRIPTION
models	Sub-models. Must be non-empty. TYPE: list of ESNModel
aggregator	Aggregation strategy. See class docstring for options. TYPE: str or Module DEFAULT: ``"mean"``

RAISES

DESCRIPTION

ValueError

If models is empty, if aggregator is an unknown string, if any sub-model is multi-output (multi-head), or if any sub-model's output dimension does not match its feedback (first input) dimension. Multi-output sub-models are rejected because the ensemble aggregates the members into a single shared feedback signal and has no meaningful coupled aggregation for auxiliary heads. The dimension check rejects headless / linear sub-models, whose output is the raw reservoir state rather than a feedback-dimensioned readout, since the coupled autoregressive loop feeds the aggregated output back as the next feedback input.

Source code in src/resdag/ensemble/coupled.py

def __init__(
    self,
    models: list[ESNModel],
    aggregator: str | nn.Module = "mean",
) -> None:
    """Initialize the coupled ensemble.

    Parameters
    ----------
    models : list of ESNModel
        Sub-models.  Must be non-empty.
    aggregator : str or nn.Module, default ``"mean"``
        Aggregation strategy.  See class docstring for options.

    Raises
    ------
    ValueError
        If ``models`` is empty, if ``aggregator`` is an unknown string, if
        any sub-model is multi-output (multi-head), or if any sub-model's
        output dimension does not match its feedback (first input)
        dimension.  Multi-output sub-models are rejected because the
        ensemble aggregates the members into a single shared feedback signal
        and has no meaningful coupled aggregation for auxiliary heads.  The
        dimension check rejects headless / linear sub-models, whose output is
        the raw reservoir state rather than a feedback-dimensioned readout,
        since the coupled autoregressive loop feeds the aggregated output
        back as the next feedback input.
    """
    super().__init__()
    if len(models) == 0:
        raise ValueError("CoupledEnsembleESNModel requires at least one sub-model.")
    if isinstance(aggregator, str) and aggregator not in ("mean", "median"):
        raise ValueError(
            f"aggregator must be 'mean', 'median', or an nn.Module; got '{aggregator}'."
        )

    self._validate_feedback_output_dims(models)

    self.models = nn.ModuleList(models)

    if isinstance(aggregator, nn.Module):
        self.aggregator_module: nn.Module | None = aggregator
        self._aggregator_str: str | None = None
    else:
        self.aggregator_module = None
        self._aggregator_str = aggregator

n_models property

n_models: int

Number of sub-models in the ensemble.

forward

forward(*inputs: Tensor) -> Tensor

Run all sub-models on the same inputs and return the aggregated output.

PARAMETER	DESCRIPTION
*inputs	Same input tensors passed to every sub-model. TYPE: Tensor DEFAULT: ()

RETURNS	DESCRIPTION
Tensor	Aggregated output of shape (batch, timesteps, output_size).

Source code in src/resdag/ensemble/coupled.py

def forward(self, *inputs: torch.Tensor) -> torch.Tensor:
    """Run all sub-models on the same inputs and return the aggregated output.

    Parameters
    ----------
    *inputs : torch.Tensor
        Same input tensors passed to every sub-model.

    Returns
    -------
    torch.Tensor
        Aggregated output of shape ``(batch, timesteps, output_size)``.
    """
    outputs = [model(*inputs) for model in self.models]
    return self._aggregate(outputs)

warmup

warmup(*inputs: Tensor) -> None

Teacher-forced warmup: synchronize every sub-model's reservoir state.

PARAMETER	DESCRIPTION
*inputs	Warmup sequences. Passed identically to each sub-model. TYPE: Tensor DEFAULT: ()

Source code in src/resdag/ensemble/coupled.py

@torch.no_grad()
def warmup(self, *inputs: torch.Tensor) -> None:
    """Teacher-forced warmup: synchronize every sub-model's reservoir state.

    Parameters
    ----------
    *inputs : torch.Tensor
        Warmup sequences.  Passed identically to each sub-model.
    """
    for model in self._iter_models():
        model.warmup(*inputs)

reset_reservoirs

reset_reservoirs() -> None

Reset reservoir states in all sub-models to zero / None.

Source code in src/resdag/ensemble/coupled.py

def reset_reservoirs(self) -> None:
    """Reset reservoir states in all sub-models to zero / None."""
    for model in self._iter_models():
        model.reset_reservoirs()

get_reservoir_states

get_reservoir_states() -> list[dict[str, Tensor]]

Return reservoir states for all sub-models.

RETURNS	DESCRIPTION
list of dict	One dict per sub-model, mapping layer name to state tensor.

Source code in src/resdag/ensemble/coupled.py

def get_reservoir_states(self) -> list[dict[str, torch.Tensor]]:
    """Return reservoir states for all sub-models.

    Returns
    -------
    list of dict
        One dict per sub-model, mapping layer name to state tensor.
    """
    return [model.get_reservoir_states() for model in self._iter_models()]

set_reservoir_states

set_reservoir_states(states: list[dict[str, Tensor]], strict: bool = True) -> None

Restore reservoir states in all sub-models.

PARAMETER	DESCRIPTION
states	One dict per sub-model as returned by :meth:get_reservoir_states. TYPE: list of dict
strict	Forwarded to :meth:ESNModel.set_reservoir_states. When True (default), each sub-model raises on missing/unexpected reservoir keys. TYPE: bool DEFAULT: True

RAISES	DESCRIPTION
ValueError	If the number of dicts in states does not match self.n_models.

Source code in src/resdag/ensemble/coupled.py

def set_reservoir_states(
    self,
    states: list[dict[str, torch.Tensor]],
    strict: bool = True,
) -> None:
    """Restore reservoir states in all sub-models.

    Parameters
    ----------
    states : list of dict
        One dict per sub-model as returned by :meth:`get_reservoir_states`.
    strict : bool, default=True
        Forwarded to :meth:`ESNModel.set_reservoir_states`.  When
        ``True`` (default), each sub-model raises on missing/unexpected
        reservoir keys.

    Raises
    ------
    ValueError
        If the number of dicts in ``states`` does not match
        ``self.n_models``.
    """
    if len(states) != self.n_models:
        raise ValueError(
            f"Expected {self.n_models} state dict(s) (one per sub-model), "
            f"got {len(states)}."
        )
    for model, state_dict in zip(self._iter_models(), states):
        model.set_reservoir_states(state_dict, strict=strict)

fit

fit(warmup_inputs: tuple[Tensor, ...], train_inputs: tuple[Tensor, ...], targets: dict[str, Tensor], n_workers: int = 1, coerce: bool = False) -> None

Train all sub-models independently using :class:~resdag.training.ESNTrainer.

Each sub-model is trained separately on the same warmup/train data. Ensemble diversity comes from the different random reservoir initialisations of each sub-model.

Before dispatching to :class:~resdag.training.ESNTrainer, every warmup_inputs / train_inputs tensor and every targets value is checked against the ensemble's reference device/dtype (the first sub-model's first floating-point parameter; see :meth:_reference_device_dtype). A mismatch raises a clear, named :class:ValueError — e.g. CPU targets fed to a GPU ensemble — instead of a raw cross-device RuntimeError deep inside the readout solve. Pass coerce=True to .to()-coerce the data to the reference instead of raising.

PARAMETER	DESCRIPTION
warmup_inputs	Warmup sequences (feedback, driver1, ...). Shape of each: (batch, warmup_steps, features). TYPE: tuple of torch.Tensor
train_inputs	Training sequences (feedback, driver1, ...). Shape of each: (batch, train_steps, features). TYPE: tuple of torch.Tensor
targets	Mapping from readout name to target tensor. Shape of each target: (batch, train_steps, out_features). TYPE: dict of str to torch.Tensor
n_workers	Number of worker threads used to fit sub-models concurrently. 1 (the default) runs sequentially in the calling thread. Larger values dispatch fits through a :class:concurrent.futures.ThreadPoolExecutor — PyTorch releases the GIL during BLAS, so threading gives a real speed-up for CPU-bound CG ridge solves without the pickling cost of multiprocessing. On GPU, all workers share one device and may interfere via the same CUDA stream; benchmark before raising n_workers above 1 in that case. TYPE: int DEFAULT: ``1``
coerce	If False (default), a device/dtype mismatch between any input or target and the ensemble's sub-models raises a clear, named :class:ValueError. If True, such tensors are coerced with .to(device=..., dtype=...) to the reference instead. TYPE: bool DEFAULT: ``False``

RAISES	DESCRIPTION
ValueError	If n_workers < 1, or if coerce=False and any input/target tensor's device or dtype does not match the ensemble's sub-models.

Source code in src/resdag/ensemble/coupled.py

def fit(
    self,
    warmup_inputs: tuple[torch.Tensor, ...],
    train_inputs: tuple[torch.Tensor, ...],
    targets: dict[str, torch.Tensor],
    n_workers: int = 1,
    coerce: bool = False,
) -> None:
    """Train all sub-models independently using :class:`~resdag.training.ESNTrainer`.

    Each sub-model is trained separately on the same warmup/train data.
    Ensemble diversity comes from the different random reservoir
    initialisations of each sub-model.

    Before dispatching to :class:`~resdag.training.ESNTrainer`, every
    ``warmup_inputs`` / ``train_inputs`` tensor and every ``targets`` value
    is checked against the ensemble's reference device/dtype (the first
    sub-model's first floating-point parameter; see
    :meth:`_reference_device_dtype`).  A mismatch raises a clear, named
    :class:`ValueError` — e.g. CPU targets fed to a GPU ensemble — instead
    of a raw cross-device ``RuntimeError`` deep inside the readout solve.
    Pass ``coerce=True`` to ``.to()``-coerce the data to the reference
    instead of raising.

    Parameters
    ----------
    warmup_inputs : tuple of torch.Tensor
        Warmup sequences ``(feedback, driver1, ...)``.
        Shape of each: ``(batch, warmup_steps, features)``.
    train_inputs : tuple of torch.Tensor
        Training sequences ``(feedback, driver1, ...)``.
        Shape of each: ``(batch, train_steps, features)``.
    targets : dict of str to torch.Tensor
        Mapping from readout name to target tensor.
        Shape of each target: ``(batch, train_steps, out_features)``.
    n_workers : int, default ``1``
        Number of worker threads used to fit sub-models concurrently.
        ``1`` (the default) runs sequentially in the calling thread.
        Larger values dispatch fits through a
        :class:`concurrent.futures.ThreadPoolExecutor` — PyTorch releases
        the GIL during BLAS, so threading gives a real speed-up for
        CPU-bound CG ridge solves without the pickling cost of
        multiprocessing.  On GPU, all workers share one device and may
        interfere via the same CUDA stream; benchmark before raising
        ``n_workers`` above 1 in that case.
    coerce : bool, default ``False``
        If ``False`` (default), a device/dtype mismatch between any input or
        target and the ensemble's sub-models raises a clear, named
        :class:`ValueError`.  If ``True``, such tensors are coerced with
        ``.to(device=..., dtype=...)`` to the reference instead.

    Raises
    ------
    ValueError
        If ``n_workers < 1``, or if ``coerce=False`` and any input/target
        tensor's device or dtype does not match the ensemble's sub-models.
    """
    if n_workers < 1:
        raise ValueError(f"n_workers must be >= 1, got {n_workers}.")

    ref_device, ref_dtype = self._reference_device_dtype()
    warmup_inputs = tuple(
        self._coerce_tensor_to_reference(
            t, ref_device, ref_dtype, coerce=coerce, label=f"warmup_inputs[{i}]"
        )
        for i, t in enumerate(warmup_inputs)
    )
    train_inputs = tuple(
        self._coerce_tensor_to_reference(
            t, ref_device, ref_dtype, coerce=coerce, label=f"train_inputs[{i}]"
        )
        for i, t in enumerate(train_inputs)
    )
    targets = {
        name: self._coerce_tensor_to_reference(
            t, ref_device, ref_dtype, coerce=coerce, label=f"targets[{name!r}]"
        )
        for name, t in targets.items()
    }

    if n_workers == 1 or self.n_models == 1:
        for model in self._iter_models():
            ESNTrainer(model).fit(warmup_inputs, train_inputs, targets)
        return

    # Thread-pool path.  Each thread gets its own ESNTrainer over a
    # distinct sub-model, so there is no shared mutable state to guard.
    from concurrent.futures import ThreadPoolExecutor

    def _fit_one(model: ESNModel) -> None:
        ESNTrainer(model).fit(warmup_inputs, train_inputs, targets)

    with ThreadPoolExecutor(max_workers=min(n_workers, self.n_models)) as pool:
        list(pool.map(_fit_one, self.models))

forecast

forecast(warmup_inputs: tuple[Tensor, ...] | Tensor, forecast_inputs: tuple[Tensor, ...] | None = None, *, horizon: int, return_warmup: bool = False, return_individuals: bool = False, reset: bool = True) -> Tensor | tuple[Tensor, list[Tensor]]

Coupled autoregressive forecast.

Phase 1 — Warmup: all sub-models are teacher-forced with the same warmup data, advancing their reservoir states independently.

Phase 2 — Coupled autoregression: at every timestep t every sub-model receives the same current_feedback (the aggregation of all models' outputs at t-1). The averaged output couples all models through a shared feedback signal.

PARAMETER	DESCRIPTION
warmup_inputs	Warmup sequences (feedback, driver1, ...). A single tensor is accepted for the common feedback-only case. TYPE: tuple of torch.Tensor or torch.Tensor
forecast_inputs	Exogenous driver inputs for the autoregressive phase (feedback is generated by the ensemble itself). Autoregressive step t consumes forecast_inputs[i][:, t, :] — pass the driver series continuing exactly where the warmup drivers ended. Each tensor must have at least horizon timesteps (one per autoregressive step); extra trailing steps are ignored. This matches :meth:resdag.core.ESNModel.forecast exactly. TYPE: tuple of torch.Tensor DEFAULT: None
horizon	Number of autoregressive steps to generate. TYPE: (int, keyword - only)
return_warmup	If True, prepend the averaged warmup outputs to the returned aggregated forecast, giving shape (batch, warmup_steps + horizon, output_size). Individual trajectories (when return_individuals=True) always cover only the autoregressive phase. TYPE: bool DEFAULT: ``False``
return_individuals	If True, also return per-sub-model autoregressive trajectories. One buffer of shape (batch, horizon, output_size) is pre-allocated per sub-model on the same device at the start of the forecast loop — only set this flag when you actually need the individual sequences, to avoid allocating N extra GPU buffers. TYPE: bool DEFAULT: ``False``
reset	If True, every sub-model's reservoir state is reset before warmup. TYPE: bool DEFAULT: ``True``

RETURNS	DESCRIPTION
Tensor	Aggregated forecast of shape (batch, horizon, output_size), or (batch, warmup_steps + horizon, output_size) when return_warmup=True. Returned alone when return_individuals=False.
tuple of (torch.Tensor, list of torch.Tensor)	When return_individuals=True: a 2-tuple whose first element is the aggregated forecast (as above) and whose second element is a list of n_models tensors each of shape (batch, horizon, output_size), in the same order as self.models.

RAISES	DESCRIPTION
ValueError	If horizon is not a positive integer, if driver arguments are inconsistent with the number of warmup inputs, or if a forecast_inputs tensor has fewer than horizon timesteps.

Source code in src/resdag/ensemble/coupled.py

@torch.no_grad()
def forecast(
    self,
    warmup_inputs: tuple[torch.Tensor, ...] | torch.Tensor,
    forecast_inputs: tuple[torch.Tensor, ...] | None = None,
    *,
    horizon: int,
    return_warmup: bool = False,
    return_individuals: bool = False,
    reset: bool = True,
) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
    """Coupled autoregressive forecast.

    **Phase 1 — Warmup**: all sub-models are teacher-forced with the same
    warmup data, advancing their reservoir states independently.

    **Phase 2 — Coupled autoregression**: at every timestep ``t`` every
    sub-model receives the *same* ``current_feedback`` (the aggregation of
    all models' outputs at ``t-1``).  The averaged output couples all models
    through a shared feedback signal.

    Parameters
    ----------
    warmup_inputs : tuple of torch.Tensor or torch.Tensor
        Warmup sequences ``(feedback, driver1, ...)``.  A single tensor
        is accepted for the common feedback-only case.
    forecast_inputs : tuple of torch.Tensor, optional
        Exogenous driver inputs for the autoregressive phase (feedback is
        generated by the ensemble itself).  Autoregressive step ``t``
        consumes ``forecast_inputs[i][:, t, :]`` — pass the driver series
        continuing exactly where the warmup drivers ended.  Each tensor
        must have **at least** ``horizon`` timesteps (one per
        autoregressive step); extra trailing steps are ignored.  This
        matches :meth:`resdag.core.ESNModel.forecast` exactly.
    horizon : int, keyword-only
        Number of autoregressive steps to generate.
    return_warmup : bool, default ``False``
        If ``True``, prepend the averaged warmup outputs to the returned
        aggregated forecast, giving shape
        ``(batch, warmup_steps + horizon, output_size)``.
        Individual trajectories (when ``return_individuals=True``) always
        cover only the autoregressive phase.
    return_individuals : bool, default ``False``
        If ``True``, also return per-sub-model autoregressive trajectories.
        One buffer of shape ``(batch, horizon, output_size)`` is
        pre-allocated per sub-model on the same device at the start of the
        forecast loop — only set this flag when you actually need the
        individual sequences, to avoid allocating N extra GPU buffers.
    reset : bool, default ``True``
        If ``True``, every sub-model's reservoir state is reset before
        warmup.

    Returns
    -------
    torch.Tensor
        Aggregated forecast of shape ``(batch, horizon, output_size)``, or
        ``(batch, warmup_steps + horizon, output_size)`` when
        ``return_warmup=True``.  Returned alone when
        ``return_individuals=False``.
    tuple of (torch.Tensor, list of torch.Tensor)
        When ``return_individuals=True``: a 2-tuple whose first element is
        the aggregated forecast (as above) and whose second element is a
        list of ``n_models`` tensors each of shape
        ``(batch, horizon, output_size)``, in the same order as
        ``self.models``.

    Raises
    ------
    ValueError
        If ``horizon`` is not a positive integer, if driver arguments are
        inconsistent with the number of warmup inputs, or if a
        ``forecast_inputs`` tensor has fewer than ``horizon`` timesteps.
    """
    if isinstance(warmup_inputs, torch.Tensor):
        warmup_inputs = (warmup_inputs,)
    else:
        warmup_inputs = tuple(warmup_inputs)
    if len(warmup_inputs) == 0:
        raise ValueError("At least one warmup input (feedback) is required.")

    # Validate the horizon up front so callers get a clear error instead of
    # a downstream IndexError / negative-dimension RuntimeError.
    if horizon < 1:
        raise ValueError(f"horizon must be a positive integer, got {horizon}")

    num_drivers = len(warmup_inputs) - 1
    has_drivers = num_drivers > 0

    if has_drivers:
        if forecast_inputs is None:
            raise ValueError(
                f"Model has {num_drivers} driving input(s). "
                "forecast_inputs must be provided for the autoregressive phase."
            )
        forecast_inputs = tuple(forecast_inputs)
        if len(forecast_inputs) != num_drivers:
            raise ValueError(
                f"Expected {num_drivers} forecast driver(s), got {len(forecast_inputs)}."
            )
        for i, driver in enumerate(forecast_inputs):
            if driver.shape[1] < horizon:
                raise ValueError(
                    f"forecast_inputs[{i}] has {driver.shape[1]} steps, expected at "
                    f"least {horizon} (one per autoregressive step). forecast_inputs "
                    f"must hold the driver series for the forecast window, starting "
                    f"right after the warmup window."
                )

    batch_size = warmup_inputs[0].shape[0]
    # Build the output buffers from the sub-models' own device/dtype, not
    # from ``warmup_inputs`` — the sub-models compute on their parameter
    # device regardless of where the warmup data lives, so this is the
    # device/dtype every aggregated step actually has.  Deriving the buffer
    # from the inputs (the old behaviour) risked a cross-device in-place
    # write or a silent dtype cast when a custom aggregator changed dtype.
    device, dtype = self._reference_device_dtype()

    # Phase 1: warmup — all models independently, same teacher-forced data
    warmup_outputs_per_model: list[torch.Tensor] = []
    for model in self._iter_models():
        # ``return_outputs=True`` yields the warmup outputs. Sub-models are
        # guaranteed single-output by ``_validate_feedback_output_dims`` at
        # construction (multi-output models are rejected there), so this is
        # always a single tensor rather than a tuple of per-head tensors.
        out = cast(torch.Tensor, model.warmup(*warmup_inputs, return_outputs=True, reset=reset))
        warmup_outputs_per_model.append(out)  # (batch, W, output_size)

    output_size = warmup_outputs_per_model[0].shape[-1]

    # Aggregate the last warmup step as the initial autoregressive feedback.
    # This seeds ``current_feedback`` but is *not* written into any output
    # buffer — slot 0 is a genuine forecast step, matching ESNModel.forecast.
    # Conform it to the reference device/dtype so a dtype-changing custom
    # aggregator does not feed a float-mismatched tensor into the sub-models
    # on the very first autoregressive step.
    last_steps = [out[:, -1:, :] for out in warmup_outputs_per_model]
    current_feedback = self._coerce_step_for_buffer(
        self._aggregate(last_steps), device, dtype
    )  # (batch, 1, output_size)

    # Pre-allocate aggregated forecast buffer
    forecast_outputs = torch.empty(batch_size, horizon, output_size, device=device, dtype=dtype)

    # Per-model buffers: only allocated when the caller explicitly requests them
    individual_outputs: list[torch.Tensor] | None = None
    if return_individuals:
        individual_outputs = [
            torch.empty(batch_size, horizon, output_size, device=device, dtype=dtype)
            for _ in self.models
        ]

    # Resolve each sub-model's flattened single-step executor and its
    # reservoir states once, *before* the loop. Stepping every member through
    # its ``FlatStep`` (a straight-line, capture-the-layers step function)
    # instead of re-invoking ``model(*step_inputs)`` sheds the per-step
    # symbolic-graph re-walk, ``nn.Module.__call__`` dispatch, and reservoir
    # sequence-loop bookkeeping — the same overhead the core forecast engine
    # eliminates for single models. The step is numerically identical to the
    # graph path (proven bit-exact upstream), so the coupled forecast output
    # is unchanged. States are threaded explicitly per member rather than
    # mutated on the module buffers mid-loop; they are persisted back to each
    # reservoir after the loop so ``reset=False`` / ``get_reservoir_states``
    # continue to observe the advanced states.
    flats: list[FlatStep] = [m._get_flat_step() for m in self._iter_models()]
    states_per_model: list[list[torch.Tensor]] = [
        self._resolve_member_states(flat, batch_size, device, dtype) for flat in flats
    ]

    # Phase 2: coupled autoregressive loop. Every slot ``0..horizon-1`` is a
    # genuine coupled step; no teacher-forced warmup echo is emitted. The
    # members are coupled through the aggregator: every member is stepped on
    # the *same* ``current_feedback`` (the aggregation of all members' outputs
    # at ``t-1``), and the aggregated output becomes the shared feedback for
    # the next step — so the members cannot be rolled out independently and
    # the aggregation must be interleaved between single steps.
    for t in range(horizon):
        if has_drivers:
            # ``has_drivers`` implies ``forecast_inputs`` was validated as a
            # non-empty tuple above; restate it for the type-checker.
            assert forecast_inputs is not None
            # Same pairing as ESNModel.forecast: step ``t`` consumes the
            # driver at the ``t``-th step after warmup = forecast_inputs[:, t].
            driver_slice = tuple(d[:, t : t + 1, :] for d in forecast_inputs)
            step_inputs: tuple[torch.Tensor, ...] = (current_feedback,) + driver_slice
        else:
            step_inputs = (current_feedback,)

        # Step every member through its flat executor on the shared feedback,
        # threading each member's own reservoir states in and out. Sub-models
        # are single-output (enforced at construction), so ``out_tuple`` holds
        # exactly one ``(batch, 1, output_size)`` slice — identical to what
        # ``model(*step_inputs)`` returned on the graph path.
        step_outputs: list[torch.Tensor] = []
        for i, flat in enumerate(flats):
            out_tuple, states_per_model[i] = flat.step(step_inputs, states_per_model[i])
            step_outputs.append(out_tuple[0])

        if individual_outputs is not None:
            for buf, out in zip(individual_outputs, step_outputs):
                buf[:, t, :] = self._coerce_step_for_buffer(out.squeeze(1), device, dtype)

        # Conform the aggregated feedback to the reference device/dtype once,
        # deliberately. A custom aggregator that changes dtype would
        # otherwise (a) be silently down/upcast by the in-place buffer write
        # and (b) re-enter the float-mismatched sub-models on the *next*
        # step, raising an opaque ``F.linear`` dtype error. Coercing here
        # fixes both and keeps the autoregressive loop type-stable.
        current_feedback = self._coerce_step_for_buffer(
            self._aggregate(step_outputs), device, dtype
        )  # (batch, 1, output_size)
        forecast_outputs[:, t, :] = current_feedback.squeeze(1)

    # Persist each member's advanced reservoir states back onto its layers so
    # a follow-up ``forecast(reset=False)`` / ``get_reservoir_states`` sees
    # them, mirroring the in-place mutation of the legacy per-step loop and
    # honouring the ``detach_state_between_calls`` contract exactly as
    # ``ESNModel.forecast`` does.
    for flat, states in zip(flats, states_per_model):
        for reservoir, state in zip(flat.reservoir_layers, states):
            if reservoir.detach_state_between_calls and state.grad_fn is not None:
                reservoir.state = state.detach().clone()
            else:
                reservoir.state = state.clone()

    if return_warmup:
        warmup_stacked = torch.stack(warmup_outputs_per_model, dim=0)
        warmup_avg = self._aggregate_stacked(warmup_stacked)
        # Conform the aggregated warmup to the forecast buffer before the
        # concat so a custom aggregator that changes dtype/device produces a
        # clear error or deliberate cast rather than a raw ``torch.cat``
        # mismatch.  ``warmup_avg`` is (batch, W, F); reuse the per-step
        # coercion over the flattened time axis by checking the tensor as a
        # whole (device + dtype only, shape is preserved).
        if warmup_avg.device != device:
            raise ValueError(
                f"Aggregated warmup output is on device={warmup_avg.device}, but "
                f"the forecast output buffer is on device={device}. A custom "
                f"aggregator must keep its output on the sub-models' device."
            )
        if warmup_avg.dtype != dtype:
            warmup_avg = warmup_avg.to(dtype=dtype)
        aggregated = torch.cat([warmup_avg, forecast_outputs], dim=1)
    else:
        aggregated = forecast_outputs

    if return_individuals:
        return aggregated, individual_outputs  # type: ignore[return-value]
    return aggregated

save

save(path: str, include_states: bool = False, **metadata: Any) -> None

Save ensemble weights (and optionally reservoir states) to a file.

PARAMETER	DESCRIPTION
path	Destination file path. TYPE: str
include_states	If True, also save current reservoir states for all sub-models. TYPE: bool DEFAULT: ``False``
**metadata	Arbitrary key-value pairs stored alongside the weights. TYPE: Any DEFAULT: {}

Source code in src/resdag/ensemble/coupled.py

def save(
    self,
    path: str,
    include_states: bool = False,
    **metadata: Any,
) -> None:
    """Save ensemble weights (and optionally reservoir states) to a file.

    Parameters
    ----------
    path : str
        Destination file path.
    include_states : bool, default ``False``
        If ``True``, also save current reservoir states for all sub-models.
    **metadata
        Arbitrary key-value pairs stored alongside the weights.
    """
    save_dict: dict[str, Any] = {
        "state_dicts": [model.state_dict() for model in self.models],
        "metadata": metadata,
    }
    if include_states:
        save_dict["reservoir_states"] = self.get_reservoir_states()
    torch.save(save_dict, path)

load

load(path: str, strict: bool = True, load_states: bool = False) -> None

Load ensemble weights from a file created by :meth:save.

PARAMETER	DESCRIPTION
path	Source file path. TYPE: str
strict	Passed to load_state_dict for each sub-model. TYPE: bool DEFAULT: ``True``
load_states	If True, also restore reservoir states. TYPE: bool DEFAULT: ``False``

RAISES	DESCRIPTION
ValueError	If the checkpoint contains a different number of sub-models.

Source code in src/resdag/ensemble/coupled.py

def load(self, path: str, strict: bool = True, load_states: bool = False) -> None:
    """Load ensemble weights from a file created by :meth:`save`.

    Parameters
    ----------
    path : str
        Source file path.
    strict : bool, default ``True``
        Passed to ``load_state_dict`` for each sub-model.
    load_states : bool, default ``False``
        If ``True``, also restore reservoir states.

    Raises
    ------
    ValueError
        If the checkpoint contains a different number of sub-models.
    """
    checkpoint = torch.load(path, weights_only=False)
    if not isinstance(checkpoint, dict) or "state_dicts" not in checkpoint:
        raise ValueError(
            f"{path} is not a save() state-dict checkpoint. If it was written "
            f"by save_full() (or torch.save(ensemble)), use "
            f"CoupledEnsembleESNModel.load_full()."
        )
    state_dicts: list[dict] = checkpoint["state_dicts"]
    if len(state_dicts) != len(self.models):
        raise ValueError(
            f"Checkpoint has {len(state_dicts)} sub-model(s), "
            f"but this ensemble has {len(self.models)}."
        )
    for model, sd in zip(self.models, state_dicts):
        model.load_state_dict(sd, strict=strict)

    if load_states:
        if "reservoir_states" in checkpoint:
            self.set_reservoir_states(checkpoint["reservoir_states"])
        else:
            warnings.warn(
                "load_states=True but checkpoint has no 'reservoir_states' key.",
                UserWarning,
                stacklevel=2,
            )

save_full

save_full(path: str, **metadata: Any) -> None

Serialize the entire ensemble — every sub-model's architecture, weights, and reservoir states — to a single file.

Unlike :meth:save (state dicts only, requires rebuilding the ensemble before :meth:load), this pickles the whole ensemble object and is restored with :meth:load_full without rebuilding anything. Relies on the pickle support added in pytorch-symbolic 1.2.

PARAMETER	DESCRIPTION
path	Destination file path. TYPE: str
**metadata	Arbitrary key-value pairs stored alongside the ensemble. TYPE: Any DEFAULT: {}

Notes

Loaded back with weights_only=False (arbitrary unpickling), so only open files you trust. Custom callable topology/initializer/activation specs must be importable (module-level, not lambdas) to be picklable; otherwise use the lighter state-dict :meth:save.

See Also

load_full : Reconstruct an ensemble saved with this method. save : Lighter, state-dict-only persistence (architecture not stored).

Source code in src/resdag/ensemble/coupled.py

def save_full(self, path: str, **metadata: Any) -> None:
    """Serialize the entire ensemble — every sub-model's architecture,
    weights, and reservoir states — to a single file.

    Unlike :meth:`save` (state dicts only, requires rebuilding the ensemble
    before :meth:`load`), this pickles the whole ensemble object and is
    restored with :meth:`load_full` without rebuilding anything.  Relies on
    the pickle support added in ``pytorch-symbolic`` 1.2.

    Parameters
    ----------
    path : str
        Destination file path.
    **metadata
        Arbitrary key-value pairs stored alongside the ensemble.

    Notes
    -----
    Loaded back with ``weights_only=False`` (arbitrary unpickling), so only
    open files you trust.  Custom callable topology/initializer/activation
    specs must be importable (module-level, not lambdas) to be picklable;
    otherwise use the lighter state-dict :meth:`save`.

    See Also
    --------
    load_full : Reconstruct an ensemble saved with this method.
    save : Lighter, state-dict-only persistence (architecture not stored).
    """
    torch.save({"resdag_full_ensemble": self, "metadata": metadata}, path)

load_full classmethod

load_full(path: str, return_metadata: bool = False, map_location: Any = None) -> CoupledEnsembleESNModel | tuple[CoupledEnsembleESNModel, dict[str, Any]]

Reconstruct a complete ensemble saved with :meth:save_full.

No pre-built ensemble is required — every sub-model is restored intact.

PARAMETER	DESCRIPTION
path	File path written by :meth:save_full. TYPE: str
return_metadata	If True, return (ensemble, metadata) instead of just the ensemble. TYPE: bool DEFAULT: ``False``
map_location	Passed to torch.load to remap storage devices. TYPE: optional DEFAULT: None

RAISES	DESCRIPTION
ValueError	If the file does not contain a whole ensemble (e.g. a state-dict checkpoint from :meth:save).

Warnings

Loads with weights_only=False, which unpickles arbitrary Python objects. Only call this on files from a source you trust.

See Also

save_full : Serialize a complete ensemble.

Source code in src/resdag/ensemble/coupled.py

@classmethod
def load_full(
    cls,
    path: str,
    return_metadata: bool = False,
    map_location: Any = None,
) -> "CoupledEnsembleESNModel | tuple[CoupledEnsembleESNModel, dict[str, Any]]":
    """Reconstruct a complete ensemble saved with :meth:`save_full`.

    No pre-built ensemble is required — every sub-model is restored intact.

    Parameters
    ----------
    path : str
        File path written by :meth:`save_full`.
    return_metadata : bool, default ``False``
        If ``True``, return ``(ensemble, metadata)`` instead of just the
        ensemble.
    map_location : optional
        Passed to ``torch.load`` to remap storage devices.

    Raises
    ------
    ValueError
        If the file does not contain a whole ensemble (e.g. a state-dict
        checkpoint from :meth:`save`).

    Warnings
    --------
    Loads with ``weights_only=False``, which unpickles arbitrary Python
    objects.  Only call this on files from a source you trust.

    See Also
    --------
    save_full : Serialize a complete ensemble.
    """
    payload = torch.load(path, weights_only=False, map_location=map_location)
    ensemble: CoupledEnsembleESNModel
    metadata: dict[str, Any]
    if isinstance(payload, dict) and "resdag_full_ensemble" in payload:
        ensemble = payload["resdag_full_ensemble"]
        metadata = payload.get("metadata", {})
    elif isinstance(payload, CoupledEnsembleESNModel):
        # A bare ``torch.save(ensemble)`` file, with no metadata wrapper.
        ensemble = payload
        metadata = {}
    else:
        raise ValueError(
            f"{path} does not contain a full ensemble. Write one with "
            f"save_full() (or torch.save(ensemble)); for a state-dict "
            f"checkpoint from save(), rebuild the ensemble and use load()."
        )
    if return_metadata:
        return ensemble, metadata
    return ensemble

aggregators

Ensemble Aggregators

This module hosts nn.Module layers that combine N per-sub-model outputs into a single aggregated tensor. Aggregators expect a 4-D input of shape (samples, batch, timesteps, features) (or a list of length samples each of shape (batch, timesteps, features)) and return a 3-D tensor of shape (batch, timesteps, features).

They plug into :class:~resdag.ensemble.CoupledEnsembleESNModel via the aggregator argument.

CLASS	DESCRIPTION
OutliersFilteredMean	Mean across the samples dimension after dropping outliers (Z-score or IQR).

See Also

resdag.ensemble.CoupledEnsembleESNModel : Coupled ensemble that consumes these aggregators.

OutliersFilteredMean

OutliersFilteredMean(method: str = 'z_score', threshold: float | None = None)

Bases: Module

Mean over ensemble members after removing outlier norms.

For each (batch, timestep) location, computes the L2 norm of each member's feature vector, flags outliers via a robust modified Z-score or the IQR rule, then averages the inlier members. If every member is an outlier at a location, falls back to the plain mean.

The "z_score" method uses the modified Z-score 0.6745 * (x - median) / MAD (Iglewicz & Hoaglin, 1993), where MAD is the median absolute deviation. Unlike the classic mean/standard-deviation Z-score, the median and MAD are robust estimators: a single extreme member does not inflate the location/scale used to judge it, so a genuine outlier can be flagged even in small ensembles. (The classic non-robust Z-score saturates at sqrt(N - 1) for a lone outlier among N members, so a threshold=3.0 cutoff could never flag anything until N >= 11.)

PARAMETER	DESCRIPTION
method	Outlier detection rule applied to per-member norms. TYPE: (z_score, iqr) DEFAULT: "z_score"
threshold	Detection sensitivity. For "z_score" it is the modified-Z-score cutoff; for "iqr" it is the Tukey fence multiplier. When None, a method-appropriate default is used: 3.5 for "z_score" (the conventional modified-Z-score cutoff) and 1.5 for "iqr" (the conventional Tukey fence). Pass an explicit value to override. TYPE: float or None DEFAULT: None

Examples:

>>> layer = OutliersFilteredMean(method="z_score")
>>> x = torch.randn(10, 3, 5, 4)  # samples, batch, time, features
>>> layer(x).shape
torch.Size([3, 5, 4])

RAISES	DESCRIPTION
ValueError	If method is not "z_score" or "iqr".

References

Iglewicz, B. and Hoaglin, D. C. (1993). How to Detect and Handle Outliers. ASQC Quality Press.

PARAMETER	DESCRIPTION
method	How to label outlier members along the samples axis. TYPE: (z_score, iqr) DEFAULT: "z_score"
threshold	Detection sensitivity (modified-Z-score bound or IQR factor). When None, defaults to 3.5 for "z_score" and 1.5 for "iqr". TYPE: float or None DEFAULT: None

RAISES	DESCRIPTION
ValueError	If method is not supported.

Source code in src/resdag/ensemble/aggregators/outliers_filtered_mean.py

def __init__(self, method: str = "z_score", threshold: float | None = None) -> None:
    """Configure outlier detection.

    Parameters
    ----------
    method : {"z_score", "iqr"}, default="z_score"
        How to label outlier members along the samples axis.
    threshold : float or None, default=None
        Detection sensitivity (modified-Z-score bound or IQR factor). When
        ``None``, defaults to ``3.5`` for ``"z_score"`` and ``1.5`` for
        ``"iqr"``.

    Raises
    ------
    ValueError
        If *method* is not supported.
    """
    super().__init__()

    if method not in ["z_score", "iqr"]:
        raise ValueError(f"Unsupported method: {method}. Choose 'z_score' or 'iqr'.")

    self.method = method
    if threshold is None:
        threshold = _DEFAULT_Z_THRESHOLD if method == "z_score" else _DEFAULT_IQR_THRESHOLD
    self.threshold = threshold

forward

forward(input: Tensor | list[Tensor]) -> Tensor

Aggregate ensemble members with outlier filtering.

Outlier detection is norm-based and whole-member: each member's L2 feature-vector norm is scored at every (batch, timestep) location, and a flagged member is dropped in its entirety (the whole feature vector is excluded from that location's mean). A member that is an outlier in only one feature but normal in norm is therefore not caught, and one that is off in a single feature enough to inflate its norm is dropped wholesale rather than per feature.

PARAMETER	DESCRIPTION
input	One of: a 4-D tensor of shape (samples, batch, timesteps, features); a 3-D tensor of shape (batch, timesteps, features), treated as a single-member ensemble (a samples axis of length one is inserted); a list of length samples whose members are each 3-D (batch, timesteps, features) tensors. TYPE: torch.Tensor or list of torch.Tensor

RETURNS	DESCRIPTION
Tensor	Shape (batch, timesteps, features).

RAISES	DESCRIPTION
ValueError	If input is a tensor whose rank is neither 3-D nor 4-D (for example a 2-D (samples, features) tensor, which would otherwise collapse silently to the wrong shape), or a list whose stacked members are not 4-D.

Source code in src/resdag/ensemble/aggregators/outliers_filtered_mean.py

def forward(self, input: torch.Tensor | list[torch.Tensor]) -> torch.Tensor:
    """Aggregate ensemble members with outlier filtering.

    Outlier detection is **norm-based and whole-member**: each member's L2
    feature-vector norm is scored at every ``(batch, timestep)`` location,
    and a flagged member is dropped in its entirety (the whole feature
    vector is excluded from that location's mean). A member that is an
    outlier in only one feature but normal in norm is therefore *not*
    caught, and one that is off in a single feature enough to inflate its
    norm is dropped wholesale rather than per feature.

    Parameters
    ----------
    input : torch.Tensor or list of torch.Tensor
        One of:

        - a 4-D tensor of shape ``(samples, batch, timesteps, features)``;
        - a 3-D tensor of shape ``(batch, timesteps, features)``, treated
          as a single-member ensemble (a ``samples`` axis of length one is
          inserted);
        - a list of length *samples* whose members are each 3-D
          ``(batch, timesteps, features)`` tensors.

    Returns
    -------
    torch.Tensor
        Shape ``(batch, timesteps, features)``.

    Raises
    ------
    ValueError
        If *input* is a tensor whose rank is neither 3-D nor 4-D (for
        example a 2-D ``(samples, features)`` tensor, which would otherwise
        collapse silently to the wrong shape), or a list whose stacked
        members are not 4-D.
    """
    if isinstance(input, list):
        input = torch.stack(input, dim=0)
        if input.dim() != 4:
            raise ValueError(
                "OutliersFilteredMean expects a list of 3D "
                "(batch, timesteps, features) tensors, but stacking the "
                f"members produced a {input.dim()}D tensor with shape "
                f"{tuple(input.shape)}"
            )
    elif input.dim() == 3:
        input = input.unsqueeze(0)
    elif input.dim() != 4:
        raise ValueError(
            "OutliersFilteredMean expects a 4D "
            "(samples, batch, timesteps, features) tensor, a 3D "
            "(batch, timesteps, features) tensor, or a list of 3D tensors, "
            f"but got a {input.dim()}D tensor with shape {tuple(input.shape)}"
        )

    norms = torch.norm(input, p=2, dim=-1)

    if self.method == "z_score":
        mask = self._mask_modified_z_score(norms)
    else:
        mask = self._mask_iqr(norms)

    mask_expanded = mask.unsqueeze(-1)
    plain_mean = input.mean(dim=0)

    masked_input = input * mask_expanded
    sum_inliers = masked_input.sum(dim=0)
    count_inliers = mask_expanded.float().sum(dim=0).expand_as(sum_inliers)

    mean_result = torch.where(
        count_inliers > 0,
        sum_inliers / count_inliers.clamp(min=1),
        plain_mean,
    )

    return mean_result

extra_repr

extra_repr() -> str

String representation of layer configuration.

Source code in src/resdag/ensemble/aggregators/outliers_filtered_mean.py

def extra_repr(self) -> str:
    """String representation of layer configuration."""
    return f"method='{self.method}', threshold={self.threshold}"