| 1 | import math |
| 2 | |
| 3 | |
| 4 | def phi(j: int, neg_h: float) -> float: |
| 5 | """ |
| 6 | Compute φⱼ(z) where z = -h (negative step size in log-space) |
| 7 | φ₁(z) = (e^z - 1) / z |
| 8 | φ₂(z) = (e^z - 1 - z) / z² |
| 9 | φⱼ(z) = (e^z - Σₖ₌₀^(j-1) zᵏ/k!) / zʲ |
| 10 | These functions naturally appear when solving: |
| 11 | dx/dt = A*x + g(x,t) (linear drift + nonlinear part) |
| 12 | """ |
| 13 | if abs(neg_h) < 1e-10: |
| 14 | # Taylor series for small h to avoid division by zero |
| 15 | # φⱼ(0) = 1/j! |
| 16 | return 1.0 / math.factorial(j) |
| 17 | |
| 18 | # Compute the "remainder" sum: Σₖ₌₀^(j-1) z^k/k! |
| 19 | remainder = sum(neg_h**k / math.factorial(k) for k in range(j)) |
| 20 | |
| 21 | # φⱼ(z) = (e^z - remainder) / z^j |
| 22 | return (math.exp(neg_h) - remainder) / (neg_h**j) |
| 23 | |
| 24 | |
| 25 | def get_res2s_coefficients(h: float, phi_cache: dict, c2: float = 0.5) -> tuple[float, float, float]: |
| 26 | """ |
| 27 | Compute res_2s Runge-Kutta coefficients for a given step size. |
| 28 | Args: |
| 29 | h: Step size in log-space = log(sigma / sigma_next) |
| 30 | phi_cache: Dictionary to cache phi function results. Cache key: (j, neg_h) |
| 31 | c2: Substep position (default 0.5 = midpoint) |
| 32 | Returns: |
| 33 | a21: Coefficient for computing intermediate x |
| 34 | b1, b2: Coefficients for final combination |
| 35 | """ |
| 36 | |
| 37 | def get_phi(j: int, neg_h: float) -> float: |
| 38 | """Get phi value with caching.""" |
| 39 | cache_key = (j, neg_h) |
| 40 | if cache_key in phi_cache: |
| 41 | return phi_cache[cache_key] |
| 42 | result = phi(j, neg_h) |
| 43 | phi_cache[cache_key] = result |
| 44 | return result |
| 45 | |
| 46 | # Substep coefficient: how much of ε₁ to use for intermediate point |
| 47 | # a21 = c2 * φ₁(-h*c2) |
| 48 | neg_h_c2 = -h * c2 |
| 49 | phi_1_c2 = get_phi(1, neg_h_c2) |
| 50 | a21 = c2 * phi_1_c2 |
| 51 | |
| 52 | # Final combination weights |
| 53 | # b2 = φ₂(-h) / c2 |
| 54 | neg_h_full = -h |
| 55 | phi_2_full = get_phi(2, neg_h_full) |
| 56 | b2 = phi_2_full / c2 |
| 57 | |
| 58 | # b1 = φ₁(-h) - b2 |
| 59 | phi_1_full = get_phi(1, neg_h_full) |
| 60 | b1 = phi_1_full - b2 |
| 61 | |
| 62 | return a21, b1, b2 |
| 63 |