1"""Small-jump Gaussian compensation for an Euler neural-SDE update.
 2
 3The implementation uses a symmetric Levy measure (random +/- signs), so the
 4retained large jumps are already centered in distribution. For a one-sided
 5uncompensated process with alpha < 1, ``small_mean`` gives the optional exact
 6mean correction described in the idea.
 7"""
 8import math
 9import numpy as np
10
11
12def small_variance(epsilon, alpha, c=1.0):
13    if not (0.0 < alpha < 2.0 and epsilon > 0.0 and c > 0.0):
14        raise ValueError("require 0 < alpha < 2, epsilon > 0, c > 0")
15    return c * epsilon ** (2.0 - alpha) / (2.0 - alpha)
16
17
18def small_mean(epsilon, alpha, c=1.0):
19    """Mean of discarded positive jumps; finite only for alpha < 1."""
20    if alpha >= 1.0:
21        raise ValueError("positive-jump mean diverges for alpha >= 1")
22    return c * epsilon ** (1.0 - alpha) / (1.0 - alpha)
23
24
25def compensated_euler_step(x, drift, h, epsilon, alpha, c=1.0,
26                           rng=None, one_sided=False, add_small_mean=False):
27    """Advance a batch of states by one Levy-perturbed Euler step.
28
29    Args:
30      x: array [batch, ...]; drift has the same shape and is f_theta(x,t).
31      h, epsilon, alpha, c: Euler time step and Levy parameters.
32      one_sided: use positive jumps instead of symmetric signs.
33      add_small_mean: for one-sided alpha<1, add h times the exact discarded
34        small-jump mean (useful for an uncompensated positive process).
35    Returns: (new_x, number_of_large_jumps_sampled).
36    """
37    if rng is None:
38        rng = np.random.default_rng()
39    x = np.asarray(x)
40    drift = np.asarray(drift)
41    if x.shape != drift.shape or h <= 0:
42        raise ValueError("x and drift must have equal shape; h must be positive")
43    batch = x.shape[0]
44    rate = h * c * epsilon ** (-alpha) / alpha
45    counts = rng.poisson(rate, size=batch)
46    total = int(counts.sum())
47    large = np.zeros_like(x, dtype=float)
48    if total:
49        # Conditional law of U given U >= epsilon:
50        # P(U > u)=(u/epsilon)^(-alpha).
51        u = epsilon * rng.random(total) ** (-1.0 / alpha)
52        signs = np.ones(total) if one_sided else rng.choice([-1.0, 1.0], total)
53        owner = np.repeat(np.arange(batch), counts)
54        sums = np.bincount(owner, weights=signs * u, minlength=batch)
55        large[:, ...] = sums.reshape((batch,) + (1,) * (x.ndim - 1))
56    small = rng.normal(0.0, math.sqrt(h * small_variance(epsilon, alpha, c)),
57                       size=x.shape)
58    if one_sided and add_small_mean:
59        small += h * small_mean(epsilon, alpha, c)
60    return x + h * drift + large + small, total