1import numpy as np
 2
 3
 4def fourier_modes(w, modes):
 5    """Circular-kernel coefficients sum_j w[j] exp(-2 pi i k j/N)."""
 6    w = np.asarray(w)
 7    modes = np.asarray(modes)
 8    return np.fft.fft(w)[modes % len(w)]
 9
10
11def newton_root(what, a=1.0, tau_r=1.0, tau_d=0.5, iterations=12, init=None):
12    """Solve tau_r*lambda = -1 + a*what*exp(-lambda*tau_d)."""
13    lam = complex((-1.0 + a * what) / tau_r if init is None else init)
14    for _ in range(iterations):
15        e = np.exp(-lam * tau_d)
16        f = tau_r * lam + 1.0 - a * what * e
17        df = tau_r + a * what * tau_d * e
18        if abs(df) < 1e-12:
19            break
20        lam -= f / df
21    return lam
22
23
24def mode_regularizer(w, modes=(-3, -2, -1, 1, 2, 3), a=1.0,
25                     tau_r=1.0, tau_d=0.5, epsilon=0.02, eta=0.0,
26                     target_velocity=0.0, T=1.0):
27    """Detached-root spectral penalty for selected Fourier modes."""
28    modes = tuple(modes)
29    hats = fourier_modes(w, modes)
30    roots = np.array([newton_root(h, a, tau_r, tau_d) for h in hats])
31    growth = roots.real
32    stable = np.logaddexp(0.0, growth + epsilon) ** 2
33    velocity = np.array([T * z.imag / (2.0 * np.pi * k)
34                         for z, k in zip(roots, modes)])
35    total = np.sum(stable + eta * (velocity - target_velocity) ** 2)
36    return float(total), roots, hats