diff --git a/src/LanPaint/lanpaint.py b/src/LanPaint/lanpaint.py index 7df06bf..10a6344 100644 --- a/src/LanPaint/lanpaint.py +++ b/src/LanPaint/lanpaint.py @@ -115,23 +115,70 @@ class LanPaint(): x_t = x_t.to(dtype) v = v.to(dtype) return x_t, v - if args is None: - #v = torch.zeros_like(x_t) - v = None - C = Coef_C(x_t) - #print(torch.squeeze(dtx), torch.squeeze(dty)) - x_t, v = advance_time(x_t, v, dt, Gamma, A, C, D) - else: - v, C = args - x_t, v = advance_time(x_t, v, dt/2, Gamma, A, C, D) + def advance_time_overdamped(x_t, dt, A, C, D): + """ + Overdamped (Gamma -> infinity) limit: + dx = -A x dt + C dt + D dW_t + with C treated as constant over this substep. + """ + dtype = x_t.dtype + with torch.autocast(device_type=x_t.device.type, dtype=torch.float32): + A_dt = A * dt + exp_neg = torch.exp(-A_dt) - C_new = Coef_C(x_t) - v = v + Gamma**0.5 * ( C_new - C) *dt + eps = 1e-8 + abs_A = torch.abs(A) + # k = (1 - exp(-A dt)) / A -> dt when A -> 0 + k = torch.where(abs_A < eps, dt, (-torch.expm1(-A_dt)) / A) + # k2 = (1 - exp(-2 A dt)) / (2 A) -> dt when A -> 0 + k2 = torch.where(abs_A < eps, dt, (-torch.expm1(-2 * A_dt)) / (2 * A)) - x_t, v = advance_time(x_t, v, dt/2, Gamma, A, C, D) + mean = exp_neg * x_t + k * C + var = (D ** 2) * k2 + noise = torch.randn_like(x_t) * torch.sqrt(torch.clamp(var, min=0.0)) + x_t = mean + noise + return x_t.to(dtype) - C = C_new + def run_damped(x_t, args): + if args is None: + v = None + C = Coef_C(x_t) + x_t, v = advance_time(x_t, v, dt, Gamma, A, C, D) + else: + v, C = args + x_t, v = advance_time(x_t, v, dt/2, Gamma, A, C, D) + C_new = Coef_C(x_t) + v = v + Gamma**0.5 * ( C_new - C) *dt + x_t, v = advance_time(x_t, v, dt/2, Gamma, A, C, D) + C = C_new + return x_t, (v, C) + + def run_overdamped(x_t, args): + if args is None: + C = Coef_C(x_t) + x_t = advance_time_overdamped(x_t, dt, A, C, D) + else: + _, C = args + x_t = advance_time_overdamped(x_t, dt / 2, A, C, D) + C_new = Coef_C(x_t) + x_t = x_t + (C_new - C) * dt + x_t = advance_time_overdamped(x_t, dt / 2, A, C, D) + C = C_new + return x_t, (None, C) + + try: + x_t_next, (v_next, C_next) = run_damped(x_t, args) + + if torch.isnan(x_t_next).any() or torch.isnan(v_next).any(): + raise ValueError("NaN detected") + + x_t = x_t_next + v = v_next + C = C_next + + except Exception: + x_t, (v, C) = run_overdamped(x_t, args) return x_t, (v, C)