Datasets¶

class OUNoiseLayer(_BaseNoiseLayer):
    """Ornstein-Uhlenbeck (OU) noise layer for temporally correlated noise.

    Implements exact discretization of the OU process where ``sigma`` is the
    stationary standard deviation:

        n_{t+1} = alpha * n_t + beta * eps_t
        alpha = exp(-dt/tau)
        beta  = sigma * sqrt(1 - exp(-2*dt/tau))
        eps_t ~ N(0, 1)

    The layer supports both single-step (stateful) and multi-step (vectorized)
    modes. In stateful mode, the noise state persists across forward calls.

    Scale and bias are handled internally by
    :class:`btorch.models.linear.LearnableScale`.

    Learnable parameters:
        - ``scale``: Multiplicative scaling (default: 1.0)
        - ``bias``: Additive offset (default: 0.0)
        - ``sigma``: Stationary standard deviation
        - ``tau``: Time constant

    Tensor Conventions:
        - Multi-step input: Output shape is ``(T, *batch_dims, *n_neuron)``
          where neuron dims are trailing.
        - Single-step input: Output shape is ``(*batch_dims, *n_neuron)``.
        - Internal state ``self.noise`` stores the state BEFORE the current
          step/sequence (i.e., the initial condition).

    Multi-step Backend:
        - Scalar tau/sigma: Uses single conv1d (fast)
        - Per-neuron tau/sigma: Uses grouped conv1d, O(T^2 * D) complexity

    Determinism Note:
        Multi-step uses vectorized RNG and convolution, so exact equality with
        repeated single-step updates is not guaranteed even with a generator.
        Use single-step loops if you need step-by-step equivalence.

    Args:
        n_neuron: Number of neurons or shape of trailing neuron dims.
        sigma: Stationary standard deviation (scalar or per-neuron).
        tau: Time constant in same units as dt (scalar or per-neuron).
        step_mode: ``'s'`` for single-step, ``'m'`` for multi-step.
        trainable_param: Set of parameter names to make trainable, or True/False
            for all/none. Options: {"scale", "bias", "sigma", "tau"}.
        trainable_shape: Shape policy for trainable values:
            - ``"scalar"`` (default): Store as scalar, broadcast to neurons
            - ``"full"``: Store as full per-neuron tensor
        stateful: If True, maintain noise state between calls (required for
            single-step mode).
        tau_min: Minimum value for tau (clamped for numerical stability).
        scale: Initial multiplicative scaling.
        bias: Initial additive offset.

    Attributes:
        noise: Current noise state tensor (only if ``stateful=True``).
        scale: Output scaling (via LearnableScale).
        bias: Output offset (via LearnableScale).
        sigma: Stationary std dev (Parameter if trainable, else buffer).
        tau: Time constant (Parameter if trainable, else buffer).
    """

    noise: Tensor

    def __init__(
        self,
        n_neuron: int | Sequence[int],
        sigma: float | Tensor = 0.5,
        tau: float | Tensor = 10.0,
        step_mode: Literal["s", "m"] = "m",
        trainable_param: bool | set[str] = False,
        *,
        trainable_shape: str = "scalar",
        stateful: bool = False,
        tau_min: float = 1e-6,
        scale: float | Tensor = 1.0,
        bias: float | Tensor = 0.0,
    ):
        super().__init__(
            n_neuron,
            scale=scale,
            bias=bias,
            trainable_param=trainable_param,
            trainable_shape=trainable_shape,
            step_mode=step_mode,
            stateful=stateful,
        )
        if not stateful and step_mode == "s":
            raise ValueError("stateful must be True for single-step mode.")
        self.tau_min = float(tau_min)

        # Memory state: stored as "n_0" (state BEFORE current step/sequence).
        if stateful:
            self.register_memory("noise", 0.0, self.n_neuron)

        # OU-specific parameters
        self.def_param(
            "sigma",
            sigma,
            trainable_param=self.trainable_param,
            trainable_shape=trainable_shape,
        )
        self.def_param(
            "tau",
            tau,
            trainable_param=self.trainable_param,
            trainable_shape=trainable_shape,
        )

    def single_step_forward(
        self, dt: float | None = None, *, generator: torch.Generator | None = None
    ) -> Tensor:
        """Single-step update of OU noise.

        Args:
            dt: Timestep (defaults to ``environ.get("dt")``).
            generator: Optional RNG generator.

        Returns:
            Updated noise tensor with same shape as ``self.noise``.
        """
        assert self.stateful, "single_step_forward requires stateful=True"

        sigma, tau = self.sigma, self.tau
        dt: float = dt if dt is not None else environ.get("dt")

        alpha = torch.exp(-dt / tau.clamp(min=self.tau_min))
        beta = sigma * torch.sqrt(
            1.0 - torch.exp(-2.0 * dt / tau.clamp(min=self.tau_min))
        )

        self.noise = alpha * self.noise + beta * randn_like(
            self.noise, generator=generator
        )
        return self._apply_scale_bias(self.noise)

    def multi_step_forward(
        self,
        T: int,
        dt: float | None = None,
        *,
        generator: torch.Generator | None = None,
    ) -> Tensor:
        """Generate a multi-step OU noise sequence.

        Args:
            T: Number of timesteps.
            dt: Timestep (defaults to ``environ.get("dt")``).
            generator: Optional RNG generator.

        Returns:
            Noise sequence of shape ``(T, *noise_shape)`` where ``noise_shape``
            matches ``self.noise.shape``.

        Raises:
            RuntimeError: If stateful but noise buffer not initialized.
        """
        if T == 0:
            return torch.empty(
                (0,) + tuple(self.noise.shape),
                device=self.noise.device,
                dtype=self.noise.dtype,
            )

        if not hasattr(self, "noise"):
            raise RuntimeError(
                "OUNoiseLayer: noise buffer is not initialized. Call reset(...) first."
            )

        dt = dt if dt is not None else environ.get("dt")
        out = ou_noise(
            sigma=self.sigma,
            tau=self.tau.clamp(min=self.tau_min),
            T=T,
            dt=dt,
            noise0=self.noise,
            generator=generator,
        )

        if self.stateful:
            # Update self.noise to the final state after the sequence.
            self.noise = out[-1]
        return self._apply_scale_bias(out)

    def extra_repr(self) -> str:
        return (
            f"n_neuron={self.n_neuron}, sigma={self._format_repr_value(self.sigma)}, "
            f"tau={self._format_repr_value(self.tau)}, "
            f"scale={self._format_repr_value(self.scale)}, "
            f"bias={self._format_repr_value(self.bias)}, "
            f"step_mode={self.step_mode}, tau_min={self.tau_min}"
        )

class PinkNoiseLayer(_BaseNoiseLayer):
    """Pink (1/f) noise layer using causal FIR filtering.

    Generates colored noise with PSD ~ 1/frequency by filtering white noise
    through a fractional integration FIR kernel. Supports both single-step
    (stateful, with history preservation) and multi-step (vectorized) modes.

    In stateful mode, the FIR history is preserved across calls for seamless
    continuation of noise sequences.

    Scale and bias are handled internally by
    :class:`btorch.models.linear.LearnableScale`.

    Learnable parameters:
        - ``scale``: Multiplicative scaling (default: 1.0)
        - ``bias``: Additive offset (default: 0.0)

    Args:
        n_neuron: Number of neurons or shape of trailing neuron dims.
        fir_order: Length of the FIR filter kernel (default 64).
        step_mode: ``'s'`` for single-step, ``'m'`` for multi-step.
        trainable_param: Set of parameter names to make trainable, or True/False
            for all/none. Options: {"scale", "bias"}.
        trainable_shape: Shape policy for trainable values:
            - ``"scalar"`` (default): Store as scalar, broadcast to neurons
            - ``"full"``: Store as full per-neuron tensor
        stateful: If True, maintain FIR history between calls (required for
            single-step mode).
        scale: Initial multiplicative scaling.
        bias: Initial additive offset.

    Attributes:
        noise: Current noise value (only if ``stateful=True``).
        white_history: Previous white noise samples for FIR continuity
            (shape ``(*n_neuron, fir_order-1)``).
        fir_order: Length of the FIR kernel.
        scale: Output scaling (via LearnableScale).
        bias: Output offset (via LearnableScale).
    """

    noise: Tensor
    white_history: Tensor

    def __init__(
        self,
        n_neuron: int | Sequence[int],
        fir_order: int = 64,
        step_mode: Literal["s", "m"] = "m",
        trainable_param: bool | set[str] = False,
        *,
        trainable_shape: str = "scalar",
        stateful: bool = False,
        scale: float | Tensor = 1.0,
        bias: float | Tensor = 0.0,
    ):
        super().__init__(
            n_neuron,
            scale=scale,
            bias=bias,
            trainable_param=trainable_param,
            trainable_shape=trainable_shape,
            step_mode=step_mode,
            stateful=stateful,
        )
        if not stateful and step_mode == "s":
            raise ValueError("stateful must be True for single-step mode.")
        if fir_order < 1:
            raise ValueError(f"fir_order must be >= 1, got {fir_order}.")

        self.fir_order = int(fir_order)

        if stateful:
            self.register_memory("noise", 0.0, self.n_neuron)
            self.register_memory(
                "white_history", 0.0, self.n_neuron + (self.fir_order - 1,)
            )

    def single_step_forward(
        self, *, generator: torch.Generator | None = None
    ) -> Tensor:
        """Single-step pink-noise update using FIR history.

        Args:
            generator: Optional RNG generator.

        Returns:
            Single noise sample with shape ``(*n_neuron)``.
        """
        assert self.stateful, "single_step_forward requires stateful=True"
        out, new_hist = pink_noise(
            T=1,
            fir_order=self.fir_order,
            white_history=self.white_history,
            generator=generator,
            return_white_history=True,
        )
        self.noise = out[0]
        self.white_history = new_hist
        return self._apply_scale_bias(self.noise)

    def multi_step_forward(
        self,
        T: int,
        *,
        generator: torch.Generator | None = None,
    ) -> Tensor:
        """Vectorized multi-step pink-noise generation.

        Args:
            T: Number of timesteps.
            generator: Optional RNG generator.

        Returns:
            Noise sequence of shape ``(T, *n_neuron)``.
        """
        if T == 0:
            return torch.empty(
                (0,) + tuple(self.noise.shape),
                device=self.noise.device,
                dtype=self.noise.dtype,
            )
        if not hasattr(self, "noise") or not hasattr(self, "white_history"):
            raise RuntimeError(
                "PinkNoiseLayer memories are not initialized. "
                "Call init_state(...) first."
            )

        out, new_hist = pink_noise(
            T=T,
            fir_order=self.fir_order,
            white_history=self.white_history,
            generator=generator,
            return_white_history=True,
        )
        if self.stateful:
            self.noise = out[-1]
            self.white_history = new_hist
        return self._apply_scale_bias(out)

    def extra_repr(self) -> str:
        return (
            f"n_neuron={self.n_neuron}, fir_order={self.fir_order}, "
            f"scale={self._format_repr_value(self.scale)}, "
            f"bias={self._format_repr_value(self.bias)}, step_mode={self.step_mode}"
        )

class PoissonNoiseLayer(_BaseNoiseLayer):
    """Poisson noise layer serving as both generator and encoder.

    Generates Poisson-distributed event counts with rate scaled by ``dt``:
    ``lambda = rate * dt`` per timestep. Because Poisson processes are
    memoryless, this layer does not maintain internal state.

    The layer can operate in two modes:
    - **Generator mode**: uses the ``rate`` provided at construction.
    - **Encoder mode**: accepts an external ``rate`` tensor via ``forward()``
      or ``multi_step_forward()``, similar to SpikingJelly's encoder pattern.

    Scale and bias are handled internally by
    :class:`btorch.models.linear.LearnableScale`.

    Learnable parameters:
        - ``scale``: Multiplicative scaling (default: 1.0)
        - ``bias``: Additive offset (default: 0.0)
        - ``rate``: Events per unit time (default: 1.0)

    Args:
        n_neuron: Number of neurons or shape of trailing neuron dims.
        rate: Default events per unit time (scalar or per-neuron,
            broadcastable). Used when ``rate`` is not provided to ``forward``.
        step_mode: ``'s'`` for single-step, ``'m'`` for multi-step.
        trainable_param: Set of parameter names to make trainable, or True/False
            for all/none. Options: {"scale", "bias", "rate"}.
        trainable_shape: Shape policy for trainable values:
            - ``"scalar"`` (default): Store as scalar, broadcast to neurons
            - ``"full"``: Store as full per-neuron tensor
        stateful: Kept for API compatibility but ignored (Poisson is
            memoryless).
        scale: Initial multiplicative scaling.
        bias: Initial additive offset.

    Attributes:
        scale: Output scaling (via LearnableScale).
        bias: Output offset (via LearnableScale).
        rate: Event rate (Parameter if trainable, else buffer).
    """

    def __init__(
        self,
        n_neuron: int | Sequence[int],
        rate: float | Tensor = 1.0,
        step_mode: Literal["s", "m"] = "m",
        trainable_param: bool | set[str] = False,
        *,
        trainable_shape: str = "scalar",
        stateful: bool = False,
        scale: float | Tensor = 1.0,
        bias: float | Tensor = 0.0,
    ):
        super().__init__(
            n_neuron,
            scale=scale,
            bias=bias,
            trainable_param=trainable_param,
            trainable_shape=trainable_shape,
            step_mode=step_mode,
            stateful=stateful,
        )
        # stateful is a no-op for Poisson (memoryless), kept for compatibility.

        self.def_param(
            "rate",
            rate,
            trainable_param=self.trainable_param,
            trainable_shape=trainable_shape,
        )

    def forward(
        self,
        rate: Tensor | None = None,
        dt: float | None = None,
        *,
        generator: torch.Generator | None = None,
    ) -> Tensor:
        """Single-step Poisson sampling.

        Args:
            rate: Optional rate tensor with shape ``(*n_neuron)`` (or any
                batch-prefixed variant). If None, uses ``self.rate``.
            dt: Timestep (defaults to ``environ.get("dt")``).
            generator: Optional RNG generator.

        Returns:
            Event counts tensor broadcastable to the rate shape.
        """
        dt = dt if dt is not None else environ.get("dt")
        rate_tensor = self.rate if rate is None else rate
        rate_tensor = torch.as_tensor(rate_tensor)

        lam = rate_tensor * float(dt)
        if torch.any(lam < 0):
            raise ValueError("Poisson rate * dt must be non-negative.")

        # Broadcast to n_neuron so shape is always well-defined.
        base = torch.zeros(self.n_neuron, dtype=lam.dtype, device=lam.device)
        lam = base + lam

        out = torch.poisson(lam, generator=generator)
        return self._apply_scale_bias(out)

    def multi_step_forward(
        self,
        T: int,
        rate: Tensor | None = None,
        dt: float | None = None,
        *,
        generator: torch.Generator | None = None,
    ) -> Tensor:
        """Vectorized multi-step Poisson sampling.

        Args:
            T: Number of timesteps.
            rate: Optional rate tensor. If None, uses ``self.rate``.
                For shape ``(*n_neuron)``, the same rate is used at every
                timestep. For shape ``(T, *n_neuron)``, per-timestep rates
                are used.
            dt: Timestep (defaults to ``environ.get("dt")``).
            generator: Optional RNG generator.

        Returns:
            Event counts of shape ``(T, *n_neuron)``.
        """
        if T < 0:
            raise ValueError(f"T must be non-negative, got {T}.")
        if T == 0:
            return torch.empty((0,) + self.n_neuron)

        dt = dt if dt is not None else environ.get("dt")
        rate_tensor = self.rate if rate is None else rate
        rate_tensor = torch.as_tensor(rate_tensor)

        lam = rate_tensor * float(dt)
        if torch.any(lam < 0):
            raise ValueError("Poisson rate * dt must be non-negative.")

        if lam.ndim >= len(self.n_neuron) + 1 and lam.shape[0] == T:
            # Per-timestep rates provided directly.
            out = torch.poisson(lam, generator=generator)
        else:
            # Scalar, per-neuron, or batch-prefixed rates: replicate across T.
            base = torch.zeros(self.n_neuron, dtype=lam.dtype, device=lam.device)
            lam_full = base + lam
            lam_seq = lam_full.unsqueeze(0).expand((T,) + lam_full.shape)
            out = torch.poisson(lam_seq, generator=generator)

        return self._apply_scale_bias(out)

    def extra_repr(self) -> str:
        return (
            f"n_neuron={self.n_neuron}, rate={self._format_repr_value(self.rate)}, "
            f"scale={self._format_repr_value(self.scale)}, "
            f"bias={self._format_repr_value(self.bias)}, step_mode={self.step_mode}"
        )