StochasticForceInference 2.0.0__py3-none-any.whl

SFI/langevin/chunked.py

@@ -0,0 +1,225 @@
+"""Chunked simulation with periodic neighbor-list rebuilds.
+
+Provides :func:`simulate_chunked`, a thin wrapper around
+:meth:`OverdampedProcess.simulate` that breaks a long simulation into
+shorter chunks and rebuilds the CSR neighbor list (via
+:func:`~SFI.utils.neighbors.build_neighbor_csr`) between chunks.
+
+This avoids the O(N²) cost of ``AutoPairs`` for large particle systems
+while keeping the extras interface exactly as-is.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from SFI.trajectory.collection import TrajectoryCollection
+from SFI.utils.neighbors import build_neighbor_csr, pad_neighbor_csr
+
+
+def simulate_chunked(
+    proc,
+    dt: float,
+    Nsteps: int,
+    key,
+    *,
+    cutoff: float,
+    box: np.ndarray,
+    skin: float = 0.0,
+    rebuild_every: int = 50,
+    save_every: Optional[int] = None,
+    spatial_dims: slice = slice(None, 2),
+    indptr_key: str = "indptr",
+    indices_key: str = "indices",
+    nnz_safety: float = 1.25,
+    oversampling: int = 1,
+    prerun: int = 0,
+    compute_observables: bool = False,
+    jit_compile: bool = True,
+    verbose: bool = False,
+) -> TrajectoryCollection:
+    """Run a chunked overdamped simulation with periodic neighbor rebuilds.
+
+    Parameters
+    ----------
+    proc : OverdampedProcess
+        An initialized process whose force uses
+        ``dispatch_pairs_from_extras(indptr_key, indices_key)``.
+    dt : float
+        Time step.
+    Nsteps : int
+        Total number of steps.
+    key : jax PRNG key
+        Random key for the simulation.
+    cutoff : float
+        Cutoff radius for neighbor list construction.
+    box : array-like, shape ``(d,)``
+        Periodic box lengths.
+    skin : float
+        Verlet skin width.  The neighbor list is built with radius
+        ``cutoff + skin`` so that particles drifting into range between
+        rebuilds are already included.  After each chunk the maximum
+        particle displacement is checked; a warning is printed if it
+        exceeds ``skin / 2`` (the Verlet safety threshold).
+    rebuild_every : int
+        Number of simulation steps between neighbor-list rebuilds.
+    save_every : int, optional
+        Number of simulation steps per output dataset.  If *None*,
+        defaults to ``rebuild_every``.  Must be a multiple of
+        ``rebuild_every``.
+    spatial_dims : slice
+        Slice into the state vector that selects spatial coordinates
+        (default: first two components ``[:2]``).
+    indptr_key, indices_key : str
+        Extras keys for the CSR neighbor list.
+    nnz_safety : float
+        Fraction by which ``max_nnz`` is enlarged beyond the initial
+        neighbor count to absorb fluctuations (default 1.25 = 25%).
+    oversampling, prerun, compute_observables, jit_compile
+        Forwarded to ``proc.simulate()``.
+    verbose : bool
+        Print progress info.
+
+    Returns
+    -------
+    TrajectoryCollection
+        Concatenated trajectory from all chunks.
+    """
+    # Time-dependent extras are not supported here: simulate() is re-invoked
+    # per rebuild chunk, which would mis-slice frame-aligned schedules.
+    from SFI.trajectory.dataset import TimeSeriesExtra
+
+    for src in (proc.extras_global, proc.extras_local):
+        for k, v in (src or {}).items():
+            if isinstance(v, TimeSeriesExtra) or (callable(v) and not hasattr(v, "func")):
+                raise NotImplementedError(
+                    f"simulate_chunked does not support time-dependent extras (got {k!r}); "
+                    "use proc.simulate() directly."
+                )
+
+    box = np.asarray(box, dtype=np.float64)
+    cutoff_list = cutoff + skin  # Verlet list radius
+
+    if save_every is None:
+        save_every = rebuild_every
+    if save_every % rebuild_every != 0:
+        raise ValueError(f"save_every ({save_every}) must be a multiple of rebuild_every ({rebuild_every})")
+    rebuilds_per_save = save_every // rebuild_every
+
+    # --- initial neighbor list ---
+    positions = np.asarray(proc._x)  # (P, d) or (d,)
+    # Assumes layout (P, d) for ndim==2 or (d,) for ndim==1; (d, P) would silently mislabel.
+    pos_spatial = positions[:, spatial_dims] if positions.ndim == 2 else positions[spatial_dims]
+    indptr, indices = build_neighbor_csr(pos_spatial, cutoff_list, box)
+    pos_at_rebuild = pos_spatial.copy()  # track displacements
+
+    # Fixed max_nnz with safety margin
+    max_nnz = max(int(len(indices) * nnz_safety), len(indices) + 1)
+    indptr_pad, indices_pad = pad_neighbor_csr(indptr, indices, max_nnz)
+
+    # Merge CSR into existing extras (don't clobber other keys)
+    base_extras = dict(proc.extras_global or {})
+    base_extras[indptr_key] = jnp.array(indptr_pad)
+    base_extras[indices_key] = jnp.array(indices_pad)
+    proc.set_extras(extras_global=base_extras)
+
+    # --- chunk the simulation ---
+    n_rebuilds = max(1, (Nsteps + rebuild_every - 1) // rebuild_every)
+    remaining = Nsteps
+    collections = []  # final output datasets
+    sub_collections = []  # accumulate rebuild-chunks within a save-chunk
+    step_done = 0
+    max_disp_in_save = 0.0  # track across rebuilds within a save-chunk
+
+    for rebuild_i in range(n_rebuilds):
+        chunk_steps = min(rebuild_every, remaining)
+        if chunk_steps <= 0:
+            break
+
+        key, sub_key = jax.random.split(key)
+
+        coll = proc.simulate(
+            dt,
+            Nsteps=chunk_steps,
+            key=sub_key,
+            oversampling=oversampling,
+            prerun=prerun if rebuild_i == 0 else 0,
+            compute_observables=compute_observables,
+            jit_compile=jit_compile,
+        )
+        sub_collections.append(coll)
+        remaining -= chunk_steps
+        step_done += chunk_steps
+
+        # Emit a save-chunk when we've accumulated enough rebuilds
+        if (rebuild_i + 1) % rebuilds_per_save == 0 or remaining <= 0:
+            # Merge sub-collection X arrays into one contiguous dataset.
+            # Each sub-collection has one dataset with X shape (T_sub, N, d)
+            # where T_sub = rebuild_every (no duplicate frames).
+            X_parts = [np.asarray(sc.datasets[0].X) for sc in sub_collections]
+            X_merged = np.concatenate(X_parts, axis=0)
+            if verbose:
+                print(f"    merged {len(sub_collections)} sub-chunks → X shape {X_merged.shape}")
+            merged = TrajectoryCollection.from_arrays(
+                X=X_merged,
+                dt=dt,
+            )
+            collections.append(merged)
+            sub_collections = []
+
+            if verbose:
+                print(f"  chunk {len(collections)}: {step_done}/{Nsteps} steps")
+
+        # Rebuild neighbors from updated positions
+        if remaining > 0:
+            positions = np.asarray(proc._x)  # (P, d) or (d,)
+            pos_spatial = positions[:, spatial_dims] if positions.ndim == 2 else positions[spatial_dims]
+
+            # --- Verlet displacement check ---
+            disp = pos_spatial - pos_at_rebuild
+            if box is not None:
+                disp = disp - box * np.round(disp / box)
+            max_disp = float(np.sqrt((disp * disp).sum(axis=-1)).max())
+            max_disp_in_save = max(max_disp_in_save, max_disp)
+            if skin > 0 and max_disp > skin / 2:
+                import warnings
+
+                warnings.warn(
+                    f"Rebuild {rebuild_i}: max displacement {max_disp:.3f} "
+                    f"> skin/2 = {skin / 2:.3f}.  Neighbor list may have "
+                    f"missed interactions.  Increase skin or decrease "
+                    f"rebuild_every.",
+                    stacklevel=2,
+                )
+            if verbose and (rebuild_i + 1) % rebuilds_per_save == 0:
+                print(f"    max displacement = {max_disp_in_save:.3f}  (skin/2 = {skin / 2:.3f})")
+                max_disp_in_save = 0.0
+
+            indptr, indices = build_neighbor_csr(pos_spatial, cutoff_list, box)
+            pos_at_rebuild = pos_spatial.copy()
+
+            # Grow max_nnz if needed
+            if len(indices) > max_nnz:
+                max_nnz = int(len(indices) * nnz_safety)
+                if verbose:
+                    print(f"    max_nnz grew to {max_nnz}")
+
+            indptr_pad, indices_pad = pad_neighbor_csr(indptr, indices, max_nnz)
+
+            # Update extras, preserving all non-CSR keys
+            updated_extras = dict(proc.extras_global or {})
+            updated_extras[indptr_key] = jnp.array(indptr_pad)
+            updated_extras[indices_key] = jnp.array(indices_pad)
+            proc.set_extras(extras_global=updated_extras)
+
+            # Free old JIT caches to prevent GPU memory buildup
+            jax.clear_caches()
+
+    # --- concatenate chunks ---
+    if len(collections) == 1:
+        return collections[0]
+    return collections[0].concat(collections[1:], weights="pool")

SFI/langevin/noise.py

@@ -0,0 +1,446 @@
+"""
+Noise models for Langevin simulators.
+
+This module provides a hierarchy of noise models that go beyond the
+simple diagonal diffusion ``D = σ I``.  In particular, it supports
+**conserved noise** needed by SPDE models such as Active Model B+, where
+the noise takes the form ``∇·(σ η)`` and preserves the spatial integral
+of the field.
+
+Class hierarchy
+---------------
+- :class:`NoiseModel` — abstract base;
+- :class:`WhiteNoise` — i.i.d. per-site Gaussian (recovers current behaviour);
+- :class:`ConservedNoise` — ``sqrt(-σ² ∇²) ξ`` via FFT on a periodic grid;
+- :class:`CompositeNoise` — different noise models on different field components.
+
+Usage
+-----
+Pass a ``NoiseModel`` instance as ``D=`` to :class:`~SFI.langevin.OverdampedProcess`
+(or any ``LangevinBase`` subclass).  The simulator detects the noise model and
+delegates noise generation to it instead of the traditional ``sqrt(2D)·ξ`` path.
+
+.. code-block:: python
+
+   from SFI.langevin.noise import ConservedNoise
+
+   noise = ConservedNoise(sigma=0.3, grid_shape=(64, 64), dx=1.0)
+   proc  = OverdampedProcess(BASIS, D=noise)
+   proc.set_params(theta_F=theta)
+   proc.set_extras(extras_global=box_extras)
+   proc.initialize(X0)
+   coll = proc.simulate(dt=0.02, Nsteps=3000, key=key, oversampling=4)
+"""
+
+from __future__ import annotations
+
+import abc
+from typing import List, Sequence, Tuple, Union
+
+import jax.numpy as jnp
+from jax import random
+
+Array = jnp.ndarray
+
+
+# ============================================================================
+# Abstract base
+# ============================================================================
+
+
+class NoiseModel(abc.ABC):
+    """Abstract base for noise models used by Langevin simulators.
+
+    Subclasses must implement :meth:`sample` and :meth:`effective_D_per_site`.
+    The simulator calls ``sample(key, x, extras)`` once per Euler–Maruyama
+    substep to obtain the noise increment (already scaled by ``sqrt(2)`` so
+    that the step becomes ``x += dt*F + sqrt(dt)*sample(key, x, extras)``).
+
+    Parameters
+    ----------
+    n_fields : int
+        Number of field components per grid site (= ``dim`` in the force
+        contract).  E.g. 1 for a scalar field, 2 for a 2-component system.
+    """
+
+    def __init__(self, *, n_fields: int) -> None:
+        self._n_fields = n_fields
+
+    @property
+    def n_fields(self) -> int:
+        """Number of field components per site."""
+        return self._n_fields
+
+    # Alias for the force-contract convention
+    @property
+    def dim(self) -> int:
+        return self._n_fields
+
+    @abc.abstractmethod
+    def sample(self, key: Array, x: Array, extras: dict) -> Array:
+        r"""Draw one noise increment.
+
+        Parameters
+        ----------
+        key : PRNG key
+        x : Array, shape ``(P, d)`` or ``(d,)``
+            Current state (used only for shape; not accessed by white/conserved
+            noise, but may be needed by multiplicative noise subclasses).
+        extras : dict
+            Process extras (contains grid geometry, etc.).
+
+        Returns
+        -------
+        Array, same shape as *x*
+            Noise increment already multiplied by ``sqrt(2)`` so the
+            integrator applies ``x += dt*F + sqrt(dt) * sample(...)``.
+        """
+        ...
+
+    @abc.abstractmethod
+    def effective_D_per_site(self, extras: dict) -> Array:
+        r"""Return an approximate per-site diffusion matrix ``(d, d)``.
+
+        This is used by the inference pipeline as a *pragmatic approximation*
+        when the noise is not white-in-space.  For ``WhiteNoise(σ)`` this
+        returns exactly ``σ·I``.  For ``ConservedNoise`` it returns the
+        spatially-averaged effective variance per site.
+
+        Returns
+        -------
+        Array, shape ``(d, d)``
+        """
+        ...
+
+    @property
+    def noise_kind(self) -> str:
+        """Short string tag for the noise type."""
+        return self.__class__.__name__
+
+
+# ============================================================================
+# White noise  (recovers current constant-scalar behaviour)
+# ============================================================================
+
+
+class WhiteNoise(NoiseModel):
+    r"""Spatially uncorrelated Gaussian noise: ``B = sqrt(2σ) I``.
+
+    Each grid site receives i.i.d. ``N(0, 2σ dt)`` noise per component.
+    This recovers the current ``D = σ`` (scalar constant) behaviour.
+
+    Parameters
+    ----------
+    sigma : float
+        Scalar diffusion coefficient (the *D* in ``dx = F dt + sqrt(2D) dW``).
+    n_fields : int
+        Number of field components per site.
+    """
+
+    def __init__(self, sigma: float, *, n_fields: int = 1) -> None:
+        super().__init__(n_fields=n_fields)
+        if sigma < 0:
+            raise ValueError(f"sigma must be non-negative, got {sigma}")
+        self._sigma = float(sigma)
+        # Precompute sqrt(2σ) for efficiency
+        self._sqrt2sigma = float(jnp.sqrt(2.0 * sigma))
+
+    @property
+    def sigma(self) -> float:
+        return self._sigma
+
+    def sample(self, key: Array, x: Array, extras: dict) -> Array:
+        xi = random.normal(key, shape=x.shape)
+        return self._sqrt2sigma * xi
+
+    def effective_D_per_site(self, extras: dict) -> Array:
+        return self._sigma * jnp.eye(self._n_fields)
+
+    def __repr__(self) -> str:
+        return f"WhiteNoise(sigma={self._sigma}, n_fields={self._n_fields})"
+
+
+# ============================================================================
+# Conserved noise  (sqrt(-σ² ∇²) via FFT on periodic grids)
+# ============================================================================
+
+
+def _build_freq_amplitudes(
+    grid_shape: Sequence[int],
+    dx: Union[float, Sequence[float]],
+    ndim: int,
+) -> Array:
+    r"""Build the Fourier-space multiplier ``|k|`` for conserved noise.
+
+    For a periodic grid with spacing *dx*, the wavenumbers along axis α are
+
+    .. math::
+
+        k_\alpha = \frac{2\pi\,n_\alpha}{N_\alpha\,\Delta x_\alpha}
+
+    and the multiplier is ``|k| = sqrt(sum_α k_α²)``, which corresponds to
+    the operator ``sqrt(-∇²)`` in Fourier space.
+
+    We use ``rfft`` along the last spatial axis, so the returned array has
+    shape ``(N_0, N_1, ..., N_{d-2}, N_{d-1}//2+1)`` for an *ndim*-D grid.
+
+    The ``k = 0`` mode is set to zero (conserved noise has zero mean).
+
+    Returns
+    -------
+    Array, real, shape matching rfft output
+        ``|k|`` on the half-complex grid.
+    """
+    grid_shape = tuple(int(n) for n in grid_shape)
+    if isinstance(dx, (int, float)):
+        dx_arr = [float(dx)] * ndim
+    else:
+        dx_arr = [float(d) for d in dx]
+
+    # Build k² = sum_α k_α²
+    k_sq = jnp.zeros(grid_shape[:-1] + (grid_shape[-1] // 2 + 1,))
+    for axis in range(ndim):
+        N = grid_shape[axis]
+        if axis < ndim - 1:
+            # Full-size axis: use fftfreq
+            freq = jnp.fft.fftfreq(N, d=dx_arr[axis])  # shape (N,)
+        else:
+            # Last axis: rfft convention → only non-negative frequencies
+            freq = jnp.fft.rfftfreq(N, d=dx_arr[axis])  # shape (N//2+1,)
+
+        k_alpha = 2 * jnp.pi * freq  # angular wavenumber
+
+        # Broadcast to the full rfft grid shape
+        shape = [1] * ndim
+        if axis < ndim - 1:
+            shape[axis] = N
+        else:
+            shape[axis] = grid_shape[-1] // 2 + 1
+        k_alpha = k_alpha.reshape(shape)
+
+        k_sq = k_sq + k_alpha**2
+
+    # |k| = sqrt(k²), with k=0 mode zeroed out
+    k_abs = jnp.sqrt(k_sq)
+    return k_abs
+
+
+class ConservedNoise(NoiseModel):
+    r"""Conserved (divergence-form) noise on a periodic square grid.
+
+    Implements noise of the form
+
+    .. math::
+
+        \eta(x, t) = \nabla \cdot \bigl(\sigma\, \vec{\Lambda}(x,t)\bigr)
+
+    where :math:`\vec{\Lambda}` is spatiotemporal white vector noise.
+    In Fourier space this is equivalent to multiplying each mode by
+    :math:`|k|`:
+
+    .. math::
+
+        \hat{\eta}_k = \sigma\,|k|\,\hat{\xi}_k
+
+    This noise **conserves the spatial average** of the field
+    (:math:`\sum_i \phi_i` is a constant of the noise process), as
+    required by Model B / Active Model B+ dynamics.
+
+    Parameters
+    ----------
+    sigma : float
+        Noise amplitude (the :math:`\sigma` in the equations above).
+        This is the *continuum* amplitude; the grid discretisation is
+        handled internally.
+    grid_shape : sequence of int
+        Grid dimensions ``(Nx, Ny, ...)`` — must match the simulation grid.
+    dx : float or sequence of float
+        Grid spacing (uniform or per-axis).
+    n_fields : int
+        Number of field components per site.
+
+    Notes
+    -----
+    The ``sample`` method uses real FFT (``rfftn`` / ``irfftn``) for
+    efficiency.  It draws white noise in real space, transforms to
+    Fourier space, multiplies by :math:`\sigma\,|k|\,\sqrt{2/\Delta V}`
+    (where :math:`\Delta V = \prod \Delta x_\alpha` is the cell volume),
+    and transforms back.
+
+    The factor :math:`1/\sqrt{\Delta V}` provides the correct continuum
+    limit: the noise covariance
+    :math:`\langle\eta_i\,\eta_j\rangle = -\sigma^2 \nabla^2 \delta_{ij} / \Delta V`
+    is independent of grid resolution when *sigma* is held fixed.
+    """
+
+    def __init__(
+        self,
+        sigma: float,
+        *,
+        grid_shape: Sequence[int],
+        dx: Union[float, Sequence[float]] = 1.0,
+        n_fields: int = 1,
+    ) -> None:
+        super().__init__(n_fields=n_fields)
+        if sigma < 0:
+            raise ValueError(f"sigma must be non-negative, got {sigma}")
+        self._sigma = float(sigma)
+        self._grid_shape = tuple(int(n) for n in grid_shape)
+        self._ndim = len(self._grid_shape)
+
+        if isinstance(dx, (int, float)):
+            self._dx = tuple([float(dx)] * self._ndim)
+        else:
+            self._dx = tuple(float(d) for d in dx)
+
+        self._P = 1
+        for n in self._grid_shape:
+            self._P *= n
+
+        # Cell volume for continuum-limit normalisation
+        dV = 1.0
+        for d in self._dx:
+            dV *= d
+        self._dV = dV
+
+        # Precompute |k| array for rfft
+        self._k_abs = _build_freq_amplitudes(self._grid_shape, self._dx, self._ndim)
+
+        # Combined multiplier: sigma * |k| * sqrt(2 / dV)
+        # The sqrt(2) enters because the integrator step is
+        #   x += sqrt(dt) * sample(...)
+        # and we need <sample_i sample_j> = 2 * D_eff * delta_{ij}
+        # where D_eff is the noise operator.
+        self._multiplier = self._sigma * self._k_abs * jnp.sqrt(2.0 / self._dV)
+
+        # Effective per-site D for inference approximation:
+        # Var(noise_i) = sigma² * <|k|²> / dV
+        # where <|k|²> = (1/N) sum_k |k|² (excluding k=0)
+        k_sq_mean = float(jnp.mean(self._k_abs**2))
+        self._D_eff = float(self._sigma**2 * k_sq_mean / self._dV)
+
+    @property
+    def sigma(self) -> float:
+        return self._sigma
+
+    @property
+    def grid_shape(self) -> Tuple[int, ...]:
+        return self._grid_shape
+
+    def sample(self, key: Array, x: Array, extras: dict) -> Array:
+        r"""Draw one conserved-noise increment.
+
+        Steps:
+        1. Draw white noise ξ ~ N(0,1) on the grid, shape (Nx, Ny, ..., d)
+        2. rFFT along spatial axes
+        3. Multiply by :math:`\sigma\,|k|\,\sqrt{2/\Delta V}` (the ``_multiplier``)
+        4. iFFT back to real space
+        5. Flatten back to (P, d)
+
+        The k=0 mode of ``_multiplier`` is zero, so sum_i η_i = 0 exactly.
+        """
+        d = self._n_fields
+        grid_d = self._grid_shape + (d,)
+
+        # Draw white noise on the grid
+        xi = random.normal(key, shape=grid_d)
+
+        # FFT axes = spatial only (not the field axis)
+        fft_axes = tuple(range(self._ndim))
+
+        # Forward real FFT along spatial axes
+        xi_hat = jnp.fft.rfftn(xi, axes=fft_axes)
+
+        # Multiply by the precomputed |k|*sigma*sqrt(2/dV)
+        # multiplier shape: (Nx, ..., Ny//2+1) — broadcast over field axis d
+        eta_hat = xi_hat * self._multiplier[..., None]
+
+        # Inverse real FFT
+        eta = jnp.fft.irfftn(eta_hat, s=self._grid_shape, axes=fft_axes)
+
+        # Flatten spatial axes: (Nx, Ny, ..., d) → (P, d)
+        return eta.reshape(self._P, d)
+
+    def effective_D_per_site(self, extras: dict) -> Array:
+        return self._D_eff * jnp.eye(self._n_fields)
+
+    def __repr__(self) -> str:
+        return (
+            f"ConservedNoise(sigma={self._sigma}, "
+            f"grid_shape={self._grid_shape}, dx={self._dx}, "
+            f"n_fields={self._n_fields})"
+        )
+
+
+# ============================================================================
+# Composite noise (different models on different field components)
+# ============================================================================
+
+
+class CompositeNoise(NoiseModel):
+    r"""Apply different noise models to different field components.
+
+    Useful when some fields have conserved dynamics (e.g. concentration)
+    and others have non-conserved dynamics (e.g. velocity).
+
+    Parameters
+    ----------
+    components : list of ``(NoiseModel, field_indices)`` pairs
+        Each element specifies a noise model and the field indices it
+        applies to.  ``field_indices`` is a list of ints.  Together the
+        indices must cover ``range(n_fields)`` exactly once.
+
+    Example
+    -------
+    >>> conserved = ConservedNoise(sigma=0.3, grid_shape=(64, 64), n_fields=1)
+    >>> white = WhiteNoise(sigma=0.1, n_fields=2)
+    >>> composite = CompositeNoise(
+    ...     components=[(conserved, [0]), (white, [1, 2])],
+    ...     n_fields=3,
+    ... )
+    """
+
+    def __init__(
+        self,
+        *,
+        components: List[Tuple[NoiseModel, List[int]]],
+        n_fields: int,
+    ) -> None:
+        super().__init__(n_fields=n_fields)
+
+        # Validate coverage
+        all_indices: set[int] = set()
+        for model, indices in components:
+            idx_set = set(indices)
+            if idx_set & all_indices:
+                raise ValueError(f"Overlapping field indices: {idx_set & all_indices}")
+            all_indices |= idx_set
+        if all_indices != set(range(n_fields)):
+            raise ValueError(f"Field indices must cover range({n_fields}), got {sorted(all_indices)}")
+
+        self._components = components
+
+    def sample(self, key: Array, x: Array, extras: dict) -> Array:
+        result = jnp.zeros_like(x)
+        for i, (model, indices) in enumerate(self._components):
+            key, sub = random.split(key)
+            # Extract the sub-state for this component's fields
+            idx = jnp.array(indices)
+            x_sub = x[..., idx]
+            noise_sub = model.sample(sub, x_sub, extras)
+            # Place back into the full field
+            result = result.at[..., idx].set(noise_sub)
+        return result
+
+    def effective_D_per_site(self, extras: dict) -> Array:
+        D = jnp.zeros((self._n_fields, self._n_fields))
+        for model, indices in self._components:
+            D_sub = model.effective_D_per_site(extras)
+            for i_local, i_global in enumerate(indices):
+                for j_local, j_global in enumerate(indices):
+                    D = D.at[i_global, j_global].set(D_sub[i_local, j_local])
+        return D
+
+    def __repr__(self) -> str:
+        parts = [f"({m!r}, {idx})" for m, idx in self._components]
+        return f"CompositeNoise(components=[{', '.join(parts)}], n_fields={self._n_fields})"