Spaces:
Running on Zero
Running on Zero
| import math | |
| def phi(j: int, neg_h: float) -> float: | |
| """ | |
| Compute φⱼ(z) where z = -h (negative step size in log-space) | |
| φ₁(z) = (e^z - 1) / z | |
| φ₂(z) = (e^z - 1 - z) / z² | |
| φⱼ(z) = (e^z - Σₖ₌₀^(j-1) zᵏ/k!) / zʲ | |
| These functions naturally appear when solving: | |
| dx/dt = A*x + g(x,t) (linear drift + nonlinear part) | |
| """ | |
| if abs(neg_h) < 1e-10: | |
| # Taylor series for small h to avoid division by zero | |
| # φⱼ(0) = 1/j! | |
| return 1.0 / math.factorial(j) | |
| # Compute the "remainder" sum: Σₖ₌₀^(j-1) z^k/k! | |
| remainder = sum(neg_h**k / math.factorial(k) for k in range(j)) | |
| # φⱼ(z) = (e^z - remainder) / z^j | |
| return (math.exp(neg_h) - remainder) / (neg_h**j) | |
| def get_res2s_coefficients(h: float, phi_cache: dict, c2: float = 0.5) -> tuple[float, float, float]: | |
| """ | |
| Compute res_2s Runge-Kutta coefficients for a given step size. | |
| Args: | |
| h: Step size in log-space = log(sigma / sigma_next) | |
| phi_cache: Dictionary to cache phi function results. Cache key: (j, neg_h) | |
| c2: Substep position (default 0.5 = midpoint) | |
| Returns: | |
| a21: Coefficient for computing intermediate x | |
| b1, b2: Coefficients for final combination | |
| """ | |
| def get_phi(j: int, neg_h: float) -> float: | |
| """Get phi value with caching.""" | |
| cache_key = (j, neg_h) | |
| if cache_key in phi_cache: | |
| return phi_cache[cache_key] | |
| result = phi(j, neg_h) | |
| phi_cache[cache_key] = result | |
| return result | |
| # Substep coefficient: how much of ε₁ to use for intermediate point | |
| # a21 = c2 * φ₁(-h*c2) | |
| neg_h_c2 = -h * c2 | |
| phi_1_c2 = get_phi(1, neg_h_c2) | |
| a21 = c2 * phi_1_c2 | |
| # Final combination weights | |
| # b2 = φ₂(-h) / c2 | |
| neg_h_full = -h | |
| phi_2_full = get_phi(2, neg_h_full) | |
| b2 = phi_2_full / c2 | |
| # b1 = φ₁(-h) - b2 | |
| phi_1_full = get_phi(1, neg_h_full) | |
| b1 = phi_1_full - b2 | |
| return a21, b1, b2 | |