1"""Nonlinear noise-tightening drift utilities."""
 2import numpy as np
 3
 4
 5def linear_drift(x, lam):
 6    return -lam * np.asarray(x)
 7
 8
 9def cubic_drift(x, lam, beta):
10    x = np.asarray(x)
11    return -lam * x - beta * x**3
12
13
14def radial_cubic_drift(x, lam, beta):
15    """Vector extension: -lambda*x - beta*x*||x||^2."""
16    x = np.asarray(x)
17    norm2 = np.sum(x * x, axis=-1, keepdims=True)
18    return -lam * x - beta * x * norm2
19
20
21def scalar_pairwise_rate(x, y, lam, beta):
22    """r=(v(x)-v(y))/(y-x) for x != y, extended continuously at x=y."""
23    x, y = np.asarray(x), np.asarray(y)
24    return lam + beta * (x*x + x*y + y*y)
25
26
27def euler_maruyama(drift, x0, dt, steps, D, rng):
28    x = np.array(x0, dtype=float, copy=True)
29    noise_scale = np.sqrt(2.0 * D * dt)
30    for _ in range(steps):
31        x += drift(x) * dt + noise_scale * rng.normal(size=x.shape)
32    return x