cds2.sde

Stochastic differential equations: Euler-Maruyama and Milstein ensembles.

Stochastic differential equation ensemble solvers.

SciPy integrates ODEs (solve_ivp) and quadratures (quad) but has no stochastic integrator: a dW term is not a parameter flip on a deterministic solver, because Ito calculus changes what convergence even means (strong vs weak order, and the Ito-Stratonovich choice). Users otherwise reach for unmaintained packages or pull in a full autodiff stack.

This module solves diagonal-noise systems

dy = drift(y, t) dt + diffusion(y, t) dW

with Euler-Maruyama (strong order 1/2) and Milstein (strong order 1), both as ensembles: every path is advanced together as one array operation, so the cost of 10,000 trajectories is one vectorised step per time point rather than 10,000 sequential integrations.

Drift and diffusion are called with a stacked state of shape (n_paths, dim) and must return the same shape. Anything written with NumPy operators satisfies this automatically::

def drift(y, t):      return 0.05 * y      # works for (dim,) and (n, dim)
def diffusion(y, t):  return 0.20 * y

That single convention is what keeps this module honest — it neither probes callables nor silently falls back to per-path loops, so a shape mistake surfaces as an error instead of a slow path.

SdeEnsemble dataclass

Ensemble of SDE trajectories.

Attributes: t: Time grid, shape (n_steps + 1,). paths: Trajectories, shape (n_paths, n_steps + 1, dim). method: Solver name, "euler" or "milstein". dt: Achieved step size. May be slightly below the requested dt because the grid is made to land exactly on t_span[1]. n_paths: Number of trajectories. seed: Seed used, or None when the generator was unseeded.

Source code in src/cds2/sde.py

@dataclass(frozen=True)
class SdeEnsemble:
    """Ensemble of SDE trajectories.

    Attributes:
        t: Time grid, shape ``(n_steps + 1,)``.
        paths: Trajectories, shape ``(n_paths, n_steps + 1, dim)``.
        method: Solver name, ``"euler"`` or ``"milstein"``.
        dt: Achieved step size. May be slightly below the requested ``dt``
            because the grid is made to land exactly on ``t_span[1]``.
        n_paths: Number of trajectories.
        seed: Seed used, or ``None`` when the generator was unseeded.
    """

    t: FloatArray
    paths: FloatArray
    method: str
    dt: float
    n_paths: int
    seed: int | None

    @property
    def terminal(self) -> FloatArray:
        """Final state of every path, shape ``(n_paths, dim)``."""
        return np.asarray(self.paths[:, -1, :], dtype=np.float64)

terminal property

Final state of every path, shape (n_paths, dim).

sde_euler_maruyama(drift, diffusion, y0, t_span, dt, n_paths=1024, seed=None)

Euler-Maruyama ensemble for dy = drift dt + diffusion dW.

Strong order 1/2, weak order 1. Each dimension is driven by its own independent Wiener process (diagonal noise).

Args: drift: drift(y, t) with y of shape (n_paths, dim), returning the same shape. diffusion: diffusion(y, t), same convention. y0: Initial state, shape (dim,), shared by every path. t_span: (t0, t1) with t0 < t1. dt: Target step size. The achieved step is (t1 - t0) / ceil((t1 - t0) / dt), never larger. n_paths: Number of trajectories to advance together. seed: Seed for the Wiener increments; None for a fresh stream.

Returns: The :class:SdeEnsemble.

Raises: TypeError: If drift or diffusion is not callable. ValueError: If y0, t_span, dt or n_paths is invalid, or a callable returns the wrong shape. FloatingPointError: If the state becomes non-finite.

Example: Geometric Brownian motion, where E[y_T] = y_0 exp(mu T) is known in closed form and so serves as a check on the solver::

    >>> import numpy as np
    >>> ens = sde_euler_maruyama(
    ...     lambda y, t: 0.05 * y,
    ...     lambda y, t: 0.20 * y,
    ...     y0=[100.0], t_span=(0.0, 1.0), dt=1e-3,
    ...     n_paths=20000, seed=7,
    ... )
    >>> bool(abs(ens.terminal.mean() - 100.0 * np.exp(0.05)) < 1.0)
    True

Source code in src/cds2/sde.py

def sde_euler_maruyama(
    drift: EnsembleFn,
    diffusion: EnsembleFn,
    y0: object,
    t_span: tuple[float, float] | Sequence[float],
    dt: float,
    n_paths: int = 1024,
    seed: int | None = None,
) -> SdeEnsemble:
    """Euler-Maruyama ensemble for ``dy = drift dt + diffusion dW``.

    Strong order 1/2, weak order 1. Each dimension is driven by its own
    independent Wiener process (diagonal noise).

    Args:
        drift: ``drift(y, t)`` with ``y`` of shape ``(n_paths, dim)``,
            returning the same shape.
        diffusion: ``diffusion(y, t)``, same convention.
        y0: Initial state, shape ``(dim,)``, shared by every path.
        t_span: ``(t0, t1)`` with ``t0 < t1``.
        dt: Target step size. The achieved step is
            ``(t1 - t0) / ceil((t1 - t0) / dt)``, never larger.
        n_paths: Number of trajectories to advance together.
        seed: Seed for the Wiener increments; ``None`` for a fresh stream.

    Returns:
        The :class:`SdeEnsemble`.

    Raises:
        TypeError: If ``drift`` or ``diffusion`` is not callable.
        ValueError: If ``y0``, ``t_span``, ``dt`` or ``n_paths`` is invalid,
            or a callable returns the wrong shape.
        FloatingPointError: If the state becomes non-finite.

    Example:
        Geometric Brownian motion, where ``E[y_T] = y_0 exp(mu T)`` is known
        in closed form and so serves as a check on the solver::

            >>> import numpy as np
            >>> ens = sde_euler_maruyama(
            ...     lambda y, t: 0.05 * y,
            ...     lambda y, t: 0.20 * y,
            ...     y0=[100.0], t_span=(0.0, 1.0), dt=1e-3,
            ...     n_paths=20000, seed=7,
            ... )
            >>> bool(abs(ens.terminal.mean() - 100.0 * np.exp(0.05)) < 1.0)
            True
    """
    return _solve("euler", drift, diffusion, y0, t_span, dt, n_paths, seed, None)

sde_milstein(drift, diffusion, y0, t_span, dt, n_paths=1024, seed=None, jacobian=None)

Milstein ensemble: Euler-Maruyama plus the 0.5 b b' (dW^2 - dt) term.

Strong order 1 for diagonal noise, against 1/2 for Euler-Maruyama. The gain is in pathwise accuracy; both schemes have weak order 1, so an expectation like a European option price converges at the same rate while a path-dependent quantity (barrier crossing, maximum drawdown) does not.

Args: drift: drift(y, t), ensemble convention as in :func:sde_euler_maruyama. diffusion: diffusion(y, t), same convention. y0: Initial state, shape (dim,). t_span: (t0, t1) with t0 < t1. dt: Target step size. n_paths: Number of trajectories. seed: Seed for the Wiener increments. jacobian: Diagonal of d diffusion / d y, same call convention. When None it is approximated by central differences, which costs 2 * dim extra evaluations per step and limits accuracy to about 1e-7 relative — pass it explicitly when known.

Returns: The :class:SdeEnsemble.

Raises: TypeError: If drift, diffusion or jacobian is not callable. ValueError: On invalid arguments or a wrong return shape. FloatingPointError: If the state becomes non-finite.

Example: For b(y) = sigma * y the derivative is just sigma, so the analytic Jacobian is available and the finite-difference step can be skipped::

    >>> ens = sde_milstein(
    ...     lambda y, t: 0.05 * y,
    ...     lambda y, t: 0.20 * y,
    ...     y0=[100.0], t_span=(0.0, 1.0), dt=1e-3,
    ...     n_paths=4096, seed=7,
    ...     jacobian=lambda y, t: 0.20 * np.ones_like(y),
    ... )
    >>> ens.method
    'milstein'

Source code in src/cds2/sde.py

def sde_milstein(
    drift: EnsembleFn,
    diffusion: EnsembleFn,
    y0: object,
    t_span: tuple[float, float] | Sequence[float],
    dt: float,
    n_paths: int = 1024,
    seed: int | None = None,
    jacobian: EnsembleFn | None = None,
) -> SdeEnsemble:
    """Milstein ensemble: Euler-Maruyama plus the ``0.5 b b' (dW^2 - dt)`` term.

    Strong order 1 for diagonal noise, against 1/2 for Euler-Maruyama. The
    gain is in *pathwise* accuracy; both schemes have weak order 1, so an
    expectation like a European option price converges at the same rate while
    a path-dependent quantity (barrier crossing, maximum drawdown) does not.

    Args:
        drift: ``drift(y, t)``, ensemble convention as in
            :func:`sde_euler_maruyama`.
        diffusion: ``diffusion(y, t)``, same convention.
        y0: Initial state, shape ``(dim,)``.
        t_span: ``(t0, t1)`` with ``t0 < t1``.
        dt: Target step size.
        n_paths: Number of trajectories.
        seed: Seed for the Wiener increments.
        jacobian: Diagonal of ``d diffusion / d y``, same call convention.
            When ``None`` it is approximated by central differences, which
            costs ``2 * dim`` extra evaluations per step and limits accuracy
            to about ``1e-7`` relative — pass it explicitly when known.

    Returns:
        The :class:`SdeEnsemble`.

    Raises:
        TypeError: If ``drift``, ``diffusion`` or ``jacobian`` is not callable.
        ValueError: On invalid arguments or a wrong return shape.
        FloatingPointError: If the state becomes non-finite.

    Example:
        For ``b(y) = sigma * y`` the derivative is just ``sigma``, so the
        analytic Jacobian is available and the finite-difference step can be
        skipped::

            >>> ens = sde_milstein(
            ...     lambda y, t: 0.05 * y,
            ...     lambda y, t: 0.20 * y,
            ...     y0=[100.0], t_span=(0.0, 1.0), dt=1e-3,
            ...     n_paths=4096, seed=7,
            ...     jacobian=lambda y, t: 0.20 * np.ones_like(y),
            ... )
            >>> ens.method
            'milstein'
    """
    return _solve("milstein", drift, diffusion, y0, t_span, dt, n_paths, seed, jacobian)

ensemble_stats(ensemble, quantiles=None)

Summarise an ensemble across paths at every time point.

Args: ensemble: Result of :func:sde_euler_maruyama or :func:sde_milstein. quantiles: Levels in the open interval (0, 1). Defaults to (0.05, 0.5, 0.95).

Returns: Mapping with "t", "mean", "std" and "quantiles". Each statistic has shape (n_steps + 1, dim), or (n_steps + 1,) when dim == 1 so that a scalar SDE plots without squeezing.

Raises: TypeError: If ensemble is not an :class:SdeEnsemble. ValueError: If any quantile lies outside (0, 1).

Example: >>> ens = sde_euler_maruyama( ... lambda y, t: -y, lambda y, t: 0.1 * np.ones_like(y), ... y0=[1.0], t_span=(0.0, 1.0), dt=1e-2, n_paths=512, seed=1, ... ) >>> stats = ensemble_stats(ens, quantiles=[0.5]) >>> stats["mean"].shape == ens.t.shape True

Source code in src/cds2/sde.py

def ensemble_stats(
    ensemble: SdeEnsemble,
    quantiles: Sequence[float] | None = None,
) -> dict[str, Any]:
    """Summarise an ensemble across paths at every time point.

    Args:
        ensemble: Result of :func:`sde_euler_maruyama` or :func:`sde_milstein`.
        quantiles: Levels in the open interval ``(0, 1)``. Defaults to
            ``(0.05, 0.5, 0.95)``.

    Returns:
        Mapping with ``"t"``, ``"mean"``, ``"std"`` and ``"quantiles"``. Each
        statistic has shape ``(n_steps + 1, dim)``, or ``(n_steps + 1,)`` when
        ``dim == 1`` so that a scalar SDE plots without squeezing.

    Raises:
        TypeError: If ``ensemble`` is not an :class:`SdeEnsemble`.
        ValueError: If any quantile lies outside ``(0, 1)``.

    Example:
        >>> ens = sde_euler_maruyama(
        ...     lambda y, t: -y, lambda y, t: 0.1 * np.ones_like(y),
        ...     y0=[1.0], t_span=(0.0, 1.0), dt=1e-2, n_paths=512, seed=1,
        ... )
        >>> stats = ensemble_stats(ens, quantiles=[0.5])
        >>> stats["mean"].shape == ens.t.shape
        True
    """
    if not isinstance(ensemble, SdeEnsemble):
        msg = f"ensemble must be an SdeEnsemble (got {type(ensemble).__name__})"
        raise TypeError(msg)

    levels = (0.05, 0.5, 0.95) if quantiles is None else tuple(float(q) for q in quantiles)
    for q in levels:
        if not 0.0 < q < 1.0:
            msg = f"quantiles must lie in the open interval (0, 1) (got {q})"
            raise ValueError(msg)

    paths = ensemble.paths
    scalar = paths.shape[2] == 1

    def _shape(arr: FloatArray) -> FloatArray:
        return np.asarray(arr[:, 0] if scalar else arr, dtype=np.float64)

    return {
        "t": ensemble.t,
        "mean": _shape(paths.mean(axis=0)),
        "std": _shape(paths.std(axis=0, ddof=0)),
        "quantiles": {q: _shape(np.quantile(paths, q, axis=0)) for q in levels},
    }