civ/hot-step-cpp-ROCm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial release

Browse Source
This commit is contained in:
civ
2026-08-16 18:24:52 +07:00
commit 876886a39a
13244 changed files with 2353959 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/plugins/solvers/aflops.lua
+57
View File
@@ -0,0 +1,57 @@
-- aflops.lua: A-FloPS (1 NFE, stateful multistep)
-- Adaptive Flow Path Sampler with velocity decomposition.
solver = {
name = "aflops",
display = "A-FloPS (1 NFE)",
description = "Exponential integrator with residual tracking",
nfe = 1,
order = 2,
needs_model = false,
stateful = true,
stochastic = false,
}
local prev_w = nil
local prev_t = 0
local prev_t_dst = 0
local function clamp_alpha(t)
local a = 1 - t
return math.max(1e-6, math.min(a, 1 - 1e-6))
end
function step(xt, vt, t_curr, t_prev, n)
if (step_index or 0) == 0 then prev_w = nil; prev_t = 0; prev_t_dst = 0 end
local alpha_curr = clamp_alpha(t_curr)
local alpha_prev = clamp_alpha(t_prev)
local alpha_ratio = alpha_prev / alpha_curr
local log_ratio = math.log(alpha_ratio)
-- Compute residual velocity: w = v + x/(1-t)
local w = {}
local inv_alpha = 1 / alpha_curr
for i = 0, n-1 do w[i] = vt[i] + xt[i] * inv_alpha end
if prev_w then
-- 2nd order: AB-like correction
local dt_curr = t_curr - t_prev
local dt_prev = prev_t - prev_t_dst
local r = (dt_prev > 1e-8) and (dt_curr / dt_prev) or 1
local c1 = 1 + 0.5 * r
local c0 = 0.5 * r
for i = 0, n-1 do
local w_eff = c1 * w[i] - c0 * prev_w[i]
xt[i] = alpha_ratio * xt[i] - alpha_prev * w_eff * log_ratio
end
else
-- 1st order: exponential Euler
for i = 0, n-1 do
xt[i] = alpha_ratio * xt[i] - alpha_prev * w[i] * log_ratio
end
end
prev_w = w
prev_t = t_curr
prev_t_dst = t_prev
end
engine/plugins/solvers/aflops2.lua
+65
View File
@@ -0,0 +1,65 @@
-- aflops2.lua: A-FloPS Midpoint (2 NFE, stateless)
-- Midpoint-corrected exponential integrator.
solver = {
name = "aflops2",
display = "A-FloPS Midpoint (2 NFE)",
description = "Midpoint-corrected exponential integrator",
nfe = 2,
order = 2,
needs_model = true,
stateful = false,
stochastic = false,
}
local function clamp_alpha(t)
local a = 1 - t
return math.max(1e-6, math.min(a, 1 - 1e-6))
end
function step(xt, vt, t_curr, t_prev, n, model_fn, vt_buf)
local dt = t_curr - t_prev
if t_curr < 1e-8 then
for i = 0, n-1 do xt[i] = xt[i] - vt[i] * dt end
return
end
local alpha_curr = clamp_alpha(t_curr)
local alpha_prev = clamp_alpha(t_prev)
-- Save v_curr
local v_curr = {}
for i = 0, n-1 do v_curr[i] = vt[i] end
-- Euler half-step to midpoint
local half_dt = dt * 0.5
local t_mid = t_curr - half_dt
local alpha_mid = clamp_alpha(t_mid)
local x_mid = {}
for i = 0, n-1 do x_mid[i] = xt[i] - v_curr[i] * half_dt end
-- Evaluate at midpoint
model_fn(xt, t_mid) -- xt used as scratch, but we need x_mid...
-- Actually we need to pass x_mid to model_fn. Fix:
-- Store xt, use x_mid for model_fn
local xt_save = {}
for i = 0, n-1 do xt_save[i] = xt[i]; xt[i] = x_mid[i] end
model_fn(xt, t_mid)
-- Compute w_mid from midpoint
local inv_alpha_mid = 1 / alpha_mid
local w_mid = {}
for i = 0, n-1 do
w_mid[i] = vt_buf[i] + x_mid[i] * inv_alpha_mid
end
-- Full step using midpoint residual
local alpha_ratio = alpha_prev / alpha_curr
local log_ratio = math.log(alpha_ratio)
for i = 0, n-1 do
xt[i] = alpha_ratio * xt_save[i] - alpha_prev * w_mid[i] * log_ratio
end
end
engine/plugins/solvers/dop853.lua
+77
View File
@@ -0,0 +1,77 @@
-- dop853.lua: Dormand-Prince 8th order (13 NFE, fixed step)
solver = {
name = "dop853",
display = "DOP853 (13 NFE)",
description = "8th order Dormand-Prince (maximum accuracy)",
nfe = 13,
order = 8,
needs_model = true,
stateful = false,
stochastic = false,
}
local C = {
1/18, 1/12, 1/8, 5/16, 3/8,
59/400, 93/200, 5490023248/9719169821,
13/20, 1201146811/1299019798,
1, 1,
}
local A = {
{1/18},
{1/48, 1/16},
{1/32, 0, 3/32},
{5/16, 0, -75/64, 75/64},
{3/80, 0, 0, 3/16, 3/20},
{29443841/614563906, 0, 0, 77736538/692538347, -28693883/1125000000, 23124283/1800000000},
{16016141/946692911, 0, 0, 61564180/158732637, 22789713/633445777, 545815736/2771057229, -180193667/1043307555},
{39632708/573591083, 0, 0, -433636366/683701615, -421739975/2616292301, 100302831/723423059, 790204164/839813087, 800635310/3783071287},
{246121993/1340847787, 0, 0, -37695042795/15268766246, -309121744/1061227803, -12992083/490766935, 6005943493/2108947869, 393006217/1396673457, 123872331/1001029789},
{-1028468189/846180014, 0, 0, 8478235783/508512852, 1311729495/1432422823, -10304129995/1701304382, -48777925059/3047939560, 15336726248/1032824649, -45442868181/3398467696, 3065993473/597172653},
{185892177/718116043, 0, 0, -3185094517/667107341, -477755414/1098053517, -703635378/230739211, 5731566787/1027545527, 5232866602/850066563, -4093664535/808688257, 3962137247/1805957418, 65686358/487910083},
{403863854/491063109, 0, 0, -5068492393/434740067, -411421997/543043805, 652783627/914296604, 11173962825/925320556, -13158990841/6184727034, 3936647629/1978049680, -160528059/685178525, 248638103/1413531060, 0},
}
local B = {
14005451/335480064, 0, 0, 0, 0,
-59238493/1068277825, 181606767/758867731, 561292985/797845732,
-1041891430/1371343529, 760417239/1151165299, 118820643/751138087,
-528747749/2220607170, 1/4,
}
function step(xt, vt, t_curr, t_prev, n, model_fn, vt_buf)
local dt = t_curr - t_prev
local xt_orig = {}
local k1 = {}
for i = 0, n-1 do xt_orig[i] = xt[i]; k1[i] = vt[i] end
local ks = {k1}
local x_tmp = {}
-- 12 extra stages
for s = 1, 12 do
local a_row = A[s]
for i = 0, n-1 do
local combo = 0
for j = 1, #a_row do
if a_row[j] ~= 0 and ks[j] then
combo = combo + a_row[j] * ks[j][i]
end
end
x_tmp[i] = xt_orig[i] - dt * combo
end
-- Write x_tmp to xt for model_fn
for i = 0, n-1 do xt[i] = x_tmp[i] end
model_fn(xt, t_curr - C[s] * dt)
ks[s+1] = {}
for i = 0, n-1 do ks[s+1][i] = vt_buf[i] end
end
-- Combine
for i = 0, n-1 do
local sol = 0
for j = 1, 13 do
if B[j] ~= 0 and ks[j] then sol = sol + B[j] * ks[j][i] end
end
xt[i] = xt_orig[i] - dt * sol
end
end
engine/plugins/solvers/dopri5.lua
+120
View File
@@ -0,0 +1,120 @@
-- dopri5.lua: Dormand-Prince 5(4) adaptive solver (7+ NFE)
-- Adaptive sub-stepping with error estimation for optimal accuracy.
solver = {
name = "dopri5",
display = "DOPRI5 (7+ NFE)",
description = "Adaptive Dormand-Prince 5th order with error control",
nfe = 0, -- variable
order = 5,
needs_model = true,
stateful = false,
stochastic = false,
}
-- Butcher tableau constants
local C = {1/5, 3/10, 4/5, 8/9, 1, 1}
local A = {
{1/5},
{3/40, 9/40},
{44/45, -56/15, 32/9},
{19372/6561, -25360/2187, 64448/6561, -212/729},
{9017/3168, -355/33, 46732/5247, 49/176, -5103/18656},
{35/384, 0, 500/1113, 125/192, -2187/6784, 11/84},
}
local B = {35/384, 0, 500/1113, 125/192, -2187/6784, 11/84, 0}
local E = {
35/384 - 1951/21600,
0,
500/1113 - 22642/50085,
125/192 - 451/720,
-2187/6784 + 12231/42400,
11/84 - 649/6300,
-1/60,
}
-- Generic ERK step: compute all stages, return result in xt_out
-- xt_fa is a FloatArray used as scratch space for model_fn calls
local function erk_step(x, k1, t, h, n_elem, model_fn, vt_buf, xt_fa, num_extra, a_rows, c_vals, b_vals)
local ks = {k1}
for s = 1, num_extra do
local a_row = a_rows[s]
for i = 0, n_elem-1 do
local combo = 0
for j = 1, #a_row do
if a_row[j] ~= 0 then combo = combo + a_row[j] * ks[j][i] end
end
xt_fa[i] = x[i] - h * combo
end
-- model_fn expects a FloatArray, writes result into vt_buf
model_fn(xt_fa, t - c_vals[s] * h)
ks[s+1] = {}
for i = 0, n_elem-1 do ks[s+1][i] = vt_buf[i] end
end
local result = {}
for i = 0, n_elem-1 do
local sol = 0
for j = 1, #b_vals do
if b_vals[j] ~= 0 and ks[j] then sol = sol + b_vals[j] * ks[j][i] end
end
result[i] = x[i] - h * sol
end
return result, ks
end
function step(xt, vt, t_curr, t_prev, n, model_fn, vt_buf)
local atol = 1e-3
local rtol = 1e-2
local max_sub = 8
local safety = 0.9
local t = t_curr
local t_end = t_prev
local h = t - t_end
-- Working copy (Lua tables for intermediate math)
local x_cur = {}
local v_cur = {}
for i = 0, n-1 do x_cur[i] = xt[i]; v_cur[i] = vt[i] end
local sub = 0
while sub < max_sub and (t - t_end) > 1e-10 do
h = math.min(h, t - t_end)
local k1 = v_cur
-- Full DOPRI5 step (6 extra stages for 7 total including FSAL)
-- Pass xt as scratch FloatArray for model_fn calls
local x_next, ks = erk_step(x_cur, k1, t, h, n, model_fn, vt_buf, xt, 6, A, C, B)
-- Error estimate
local err_sq_sum = 0
for i = 0, n-1 do
local err_i = 0
for j = 1, 7 do
if E[j] ~= 0 and ks[j] then err_i = err_i + E[j] * ks[j][i] end
end
err_i = err_i * h
local scale = atol + rtol * math.max(math.abs(x_cur[i]), math.abs(x_next[i]))
local ratio = err_i / scale
err_sq_sum = err_sq_sum + ratio * ratio
end
local err_norm = math.sqrt(err_sq_sum / n)
if err_norm <= 1 then
t = t - h
x_cur = x_next
v_cur = ks[7] -- FSAL
if err_norm > 1e-10 then
h = h * math.min(5, safety * err_norm ^ (-0.2))
else
h = h * 5
end
else
h = h * math.max(0.2, safety * err_norm ^ (-0.2))
end
sub = sub + 1
end
for i = 0, n-1 do xt[i] = x_cur[i] end
end
engine/plugins/solvers/dpm2m.lua
+39
View File
@@ -0,0 +1,39 @@
-- dpm2m.lua: DPM++ 2M solver (Adams-Bashforth 2, 1 NFE, stateful)
-- Uses previous step's velocity for 2nd order correction.
-- First step falls back to Euler.
solver = {
name = "dpm2m",
display = "DPM++ 2M",
description = "2nd order Adams-Bashforth multistep (1 NFE)",
nfe = 1,
order = 2,
needs_model = false,
stateful = true,
stochastic = false,
}
-- Persistent state across steps (Lua tables survive between calls)
local prev_vt = nil
function step(xt, vt, t_curr, t_prev, n)
if (step_index or 0) == 0 then prev_vt = nil end
local dt = t_curr - t_prev
if prev_vt then
-- AB2: v_eff = 1.5 * vt - 0.5 * prev_vt
for i = 0, n - 1 do
local v_eff = 1.5 * vt[i] - 0.5 * prev_vt[i]
xt[i] = xt[i] - v_eff * dt
end
else
-- First step: Euler
for i = 0, n - 1 do
xt[i] = xt[i] - vt[i] * dt
end
end
-- Save velocity for next step
prev_vt = {}
for i = 0, n - 1 do prev_vt[i] = vt[i] end
end
engine/plugins/solvers/dpm2m_ada.lua
+41
View File
@@ -0,0 +1,41 @@
-- dpm2m_ada.lua: DPM++ 2M Adaptive (step-ratio-corrected AB2)
-- Adjusts AB2 coefficients based on step size ratio for non-uniform schedules.
solver = {
name = "dpm2m_ada",
display = "DPM++ 2M Adaptive",
description = "Step-ratio-corrected AB2 for non-uniform schedules",
nfe = 1,
order = 2,
needs_model = false,
stateful = true,
stochastic = false,
}
local prev_vt = nil
local prev_dt = 0
function step(xt, vt, t_curr, t_prev, n)
if (step_index or 0) == 0 then prev_vt = nil; prev_dt = 0 end
local dt = t_curr - t_prev
if prev_vt and prev_dt > 0 then
-- Step-ratio corrected AB2
local r = dt / prev_dt
local c1 = 1.0 + r / 2.0
local c0 = r / 2.0
for i = 0, n - 1 do
local v_eff = c1 * vt[i] - c0 * prev_vt[i]
xt[i] = xt[i] - v_eff * dt
end
else
-- First step: Euler
for i = 0, n - 1 do
xt[i] = xt[i] - vt[i] * dt
end
end
prev_vt = {}
for i = 0, n - 1 do prev_vt[i] = vt[i] end
prev_dt = dt
end
engine/plugins/solvers/dpm3m.lua
+45
View File
@@ -0,0 +1,45 @@
-- dpm3m.lua: DPM++ 3M solver (Adams-Bashforth 3, 1 NFE, stateful)
-- Uses two previous velocities for 3rd order correction.
solver = {
name = "dpm3m",
display = "DPM++ 3M",
description = "3rd order Adams-Bashforth multistep (1 NFE)",
nfe = 1,
order = 3,
needs_model = false,
stateful = true,
stochastic = false,
}
local prev_vt = nil
local prev_prev_vt = nil
function step(xt, vt, t_curr, t_prev, n)
if (step_index or 0) == 0 then prev_vt = nil; prev_prev_vt = nil end
local dt = t_curr - t_prev
if prev_vt and prev_prev_vt then
-- AB3: v_eff = (23/12)*vt - (16/12)*prev - (5/12)*prev_prev
for i = 0, n - 1 do
local v_eff = (23/12) * vt[i] - (16/12) * prev_vt[i] + (5/12) * prev_prev_vt[i]
xt[i] = xt[i] - v_eff * dt
end
elseif prev_vt then
-- AB2 fallback
for i = 0, n - 1 do
local v_eff = 1.5 * vt[i] - 0.5 * prev_vt[i]
xt[i] = xt[i] - v_eff * dt
end
else
-- Euler fallback
for i = 0, n - 1 do
xt[i] = xt[i] - vt[i] * dt
end
end
-- Shift history
prev_prev_vt = prev_vt
prev_vt = {}
for i = 0, n - 1 do prev_vt[i] = vt[i] end
end
engine/plugins/solvers/euler.lua
+21
View File
@@ -0,0 +1,21 @@
-- euler.lua: First-order Euler ODE solver
-- Single evaluation, simplest possible solver.
-- xt_next = xt - vt * (t_curr - t_prev)
solver = {
name = "euler",
display = "Euler (ODE)",
description = "First-order Euler step (default)",
nfe = 1,
order = 1,
needs_model = false,
stateful = false,
stochastic = false,
}
function step(xt, vt, t_curr, t_prev, n)
local dt = t_curr - t_prev
for i = 0, n - 1 do
xt[i] = xt[i] - vt[i] * dt
end
end
engine/plugins/solvers/gl2s.lua
+60
View File
@@ -0,0 +1,60 @@
-- gl2s.lua: Gauss-Legendre 2-stage implicit Runge-Kutta (4th order, 6 NFE)
-- A-stable, symplectic. Fixed-point iteration solves the implicit system.
solver = {
name = "gl2s",
display = "Gauss-Legendre 2s (6 NFE)",
description = "Implicit 4th-order A-stable symplectic integrator",
nfe = 6,
order = 4,
needs_model = true,
stateful = false,
stochastic = false,
}
-- Butcher tableau
local SQRT3_6 = 0.28867513459481287
local C1 = 0.5 - SQRT3_6 -- ≈ 0.2113
local C2 = 0.5 + SQRT3_6 -- ≈ 0.7887
local A11 = 0.25
local A12 = 0.25 - SQRT3_6 -- ≈ -0.0387
local A21 = 0.25 + SQRT3_6 -- ≈ 0.5387
local A22 = 0.25
local ITERATIONS = 3
function step(xt, vt, t_curr, t_prev, n, model_fn, vt_buf)
local dt = t_curr - t_prev
-- Initialize k1 = k2 = vt
local k1 = {}
local k2 = {}
for i = 0, n-1 do k1[i] = vt[i]; k2[i] = vt[i] end
local xt_orig = {}
for i = 0, n-1 do xt_orig[i] = xt[i] end
local t1 = t_curr - C1 * dt
local t2 = t_curr - C2 * dt
-- Fixed-point iteration
for iter = 1, ITERATIONS do
-- Stage 1: x1 = xt - dt*(A11*k1 + A12*k2)
for i = 0, n-1 do
xt[i] = xt_orig[i] - dt * (A11 * k1[i] + A12 * k2[i])
end
model_fn(xt, t1)
for i = 0, n-1 do k1[i] = vt_buf[i] end
-- Stage 2: x2 = xt - dt*(A21*k1 + A22*k2)
for i = 0, n-1 do
xt[i] = xt_orig[i] - dt * (A21 * k1[i] + A22 * k2[i])
end
model_fn(xt, t2)
for i = 0, n-1 do k2[i] = vt_buf[i] end
end
-- Final: xt = xt_orig - dt * 0.5 * (k1 + k2)
for i = 0, n-1 do
xt[i] = xt_orig[i] - dt * 0.5 * (k1[i] + k2[i])
end
end
engine/plugins/solvers/heun.lua
+34
View File
@@ -0,0 +1,34 @@
-- heun.lua: Heun's method (improved Euler / explicit trapezoidal)
-- 2 NFE: evaluate at t_curr, predict, evaluate at t_prev, average.
solver = {
name = "heun",
display = "Heun (2 NFE)",
description = "Second-order predictor-corrector",
nfe = 2,
order = 2,
needs_model = true,
stateful = false,
stochastic = false,
}
function step(xt, vt, t_curr, t_prev, n, model_fn, vt_buf)
local dt = t_curr - t_prev
-- Predict: xt_pred = xt - vt * dt
for i = 0, n - 1 do
xt[i] = xt[i] - vt[i] * dt
end
-- Correct: evaluate at (xt_pred, t_prev)
model_fn(xt, t_prev)
-- Average: xt = xt + 0.5 * (vt_buf - vt) * dt
-- Note: xt is already xt_pred = xt_orig - vt * dt
-- We want: xt_orig - 0.5*(vt + vt_buf)*dt
-- = (xt + vt*dt) - 0.5*(vt + vt_buf)*dt
-- = xt + 0.5*(vt - vt_buf)*dt
for i = 0, n - 1 do
xt[i] = xt[i] + 0.5 * (vt[i] - vt_buf[i]) * dt
end
end
engine/plugins/solvers/jkass_fast.lua
+111
View File
@@ -0,0 +1,111 @@
-- jkass_fast.lua: JKASS Fast solver (1 NFE, stateful)
-- Euler with momentum blending, frequency damping, and temporal smoothing.
-- Port from jeankassio/JK-AceStep-Nodes.
solver = {
name = "jkass_fast",
display = "JKASS Fast",
description = "Euler with beat stability, frequency damping, and temporal smoothing",
accent = "amber",
nfe = 1,
order = 1,
needs_model = false,
stateful = true,
stochastic = false,
params = {
{ key = "beat_stability", type = "slider", label = "Beat Stability",
default = 0.25, min = 0, max = 1, step = 0.01,
hint = "Momentum blend with previous step (0=off, 1=full momentum)" },
{ key = "frequency_damping", type = "slider", label = "Frequency Damping",
default = 0.4, min = 0, max = 5, step = 0.1,
hint = "Attenuate high-frequency bins (0=off)" },
{ key = "temporal_smoothing", type = "slider", label = "Temporal Smoothing",
default = 0.13, min = 0, max = 1, step = 0.01,
hint = "1D blur across time axis (0=off)" },
},
}
local prev_delta = nil
-- Frequency damping: exponential decay across channel dimension
local function apply_frequency_damping(data, offset, T, Oc, damping)
if damping <= 0 then return end
local freq_mult = {}
for c = 0, Oc - 1 do
local freq = c / (Oc - 1)
freq_mult[c] = math.exp(-damping * freq * freq)
end
for t = 0, T - 1 do
for c = 0, Oc - 1 do
local idx = offset + t * Oc + c
data[idx] = data[idx] * freq_mult[c]
end
end
end
-- Temporal smoothing: [0.25, 0.5, 0.25] blur across time axis
local function apply_temporal_smoothing(data, offset, T, Oc, strength)
if strength <= 0 or T < 3 then return end
local smoothed = {}
for c = 0, Oc - 1 do
for t = 0, T - 1 do
local t_prev = (t > 0) and (t - 1) or 1
local t_next = (t < T - 1) and (t + 1) or (T - 2)
local v_prev = data[offset + t_prev * Oc + c]
local v_curr = data[offset + t * Oc + c]
local v_next = data[offset + t_next * Oc + c]
smoothed[t * Oc + c] = 0.25 * v_prev + 0.5 * v_curr + 0.25 * v_next
end
end
for i = 0, T * Oc - 1 do
data[offset + i] = (1 - strength) * data[offset + i] + strength * smoothed[i]
end
end
function step(xt, vt, t_curr, t_prev, n)
if (step_index or 0) == 0 then prev_delta = nil end
local dt = t_curr - t_prev
-- Read params (injected by C++ before each call)
local bs = params and params.beat_stability or 0.25
local fd = params and params.frequency_damping or 0.4
local ts = params and params.temporal_smoothing or 0.13
-- Copy velocity as working delta
local delta = {}
for i = 0, n - 1 do delta[i] = vt[i] end
-- Beat stability: momentum blend
if prev_delta and bs > 0 then
for i = 0, n - 1 do
delta[i] = (1 - bs) * delta[i] + bs * prev_delta[i]
end
end
-- Save for next step
prev_delta = {}
for i = 0, n - 1 do prev_delta[i] = delta[i] end
-- Frequency damping (per batch item, Oc=64 for ACE-Step)
if fd > 0 and n_per and n_per > 0 then
local Oc = 64
local T = n_per / Oc
for b = 0, batch_n - 1 do
apply_frequency_damping(delta, b * n_per, T, Oc, fd)
end
end
-- Temporal smoothing (per batch item)
if ts > 0 and n_per and n_per > 0 then
local Oc = 64
local T = n_per / Oc
for b = 0, batch_n - 1 do
apply_temporal_smoothing(delta, b * n_per, T, Oc, ts)
end
end
-- Euler step with modified delta
for i = 0, n - 1 do
xt[i] = xt[i] - delta[i] * dt
end
end
engine/plugins/solvers/jkass_quality.lua
+40
View File
@@ -0,0 +1,40 @@
-- jkass_quality.lua: JKASS Quality solver (2 NFE)
-- Heun with derivative averaging.
solver = {
name = "jkass_quality",
display = "JKASS Quality (2 NFE)",
description = "Heun predictor-corrector with derivative averaging",
accent = "amber",
nfe = 2,
order = 2,
needs_model = true,
stateful = false,
stochastic = false,
}
function step(xt, vt, t_curr, t_prev, n, model_fn, vt_buf)
local dt = t_curr - t_prev
if t_prev <= 0 then
-- Fallback to Euler for final step
for i = 0, n - 1 do xt[i] = xt[i] - vt[i] * dt end
return
end
-- Save k1
local k1 = {}
for i = 0, n - 1 do k1[i] = vt[i] end
-- Euler predictor
for i = 0, n - 1 do xt[i] = xt[i] - vt[i] * dt end
-- Second evaluation
model_fn(xt, t_prev)
-- Correct: undo Euler, apply averaged derivative
for i = 0, n - 1 do
xt[i] = xt[i] + k1[i] * dt -- undo Euler
xt[i] = xt[i] - 0.5 * (k1[i] + vt_buf[i]) * dt -- Heun average
end
end
engine/plugins/solvers/md_confluence_v3.lua
+1054
View File
File diff suppressed because it is too large Load Diff
engine/plugins/solvers/rfsolver.lua
+40
View File
@@ -0,0 +1,40 @@
-- rfsolver.lua: RF-Solver (2 NFE) — Rectified Flow specific
-- Exploits the RF ODE structure for higher accuracy than generic midpoint.
solver = {
name = "rfsolver",
display = "RF-Solver (2 NFE)",
description = "Rectified-flow-aware midpoint solver",
nfe = 2,
order = 2,
needs_model = true,
stateful = false,
stochastic = false,
}
function step(xt, vt, t_curr, t_prev, n, model_fn, vt_buf)
local dt = t_curr - t_prev
if t_curr < 1e-8 then
for i = 0, n-1 do xt[i] = xt[i] - vt[i] * dt end
return
end
-- Save v_t
local v_t = {}
for i = 0, n-1 do v_t[i] = vt[i] end
-- Half-step to midpoint
local half_dt = dt * 0.5
local t_mid = t_curr - half_dt
for i = 0, n-1 do xt[i] = xt[i] - v_t[i] * half_dt end
-- Evaluate at midpoint
model_fn(xt, t_mid)
-- RF-specific: reconstruct via x_0 prediction from midpoint
for i = 0, n-1 do
local x_0_mid = xt[i] - t_mid * vt_buf[i]
xt[i] = x_0_mid + t_prev * vt_buf[i]
end
end
engine/plugins/solvers/rk4.lua
+54
View File
@@ -0,0 +1,54 @@
-- rk4.lua: Classic 4th-order Runge-Kutta solver
-- 4 NFE per step, excellent accuracy for smooth flows.
solver = {
name = "rk4",
display = "RK4 (4 NFE)",
description = "Classic 4th-order Runge-Kutta",
nfe = 4,
order = 4,
needs_model = true,
stateful = false,
stochastic = false,
}
function step(xt, vt, t_curr, t_prev, n, model_fn, vt_buf)
local dt = t_curr - t_prev
local t_mid = t_curr - 0.5 * dt
-- k1 = vt (already evaluated)
-- Save k1 and original xt
local k1 = {}
local xt_orig = {}
for i = 0, n - 1 do
k1[i] = vt[i]
xt_orig[i] = xt[i]
end
-- k2: evaluate at midpoint using k1
for i = 0, n - 1 do
xt[i] = xt_orig[i] - 0.5 * k1[i] * dt
end
model_fn(xt, t_mid)
local k2 = {}
for i = 0, n - 1 do k2[i] = vt_buf[i] end
-- k3: evaluate at midpoint using k2
for i = 0, n - 1 do
xt[i] = xt_orig[i] - 0.5 * k2[i] * dt
end
model_fn(xt, t_mid)
local k3 = {}
for i = 0, n - 1 do k3[i] = vt_buf[i] end
-- k4: evaluate at endpoint using k3
for i = 0, n - 1 do
xt[i] = xt_orig[i] - k3[i] * dt
end
model_fn(xt, t_prev)
-- Combine: xt = xt_orig - (k1 + 2*k2 + 2*k3 + k4) * dt / 6
for i = 0, n - 1 do
xt[i] = xt_orig[i] - (k1[i] + 2*k2[i] + 2*k3[i] + vt_buf[i]) * dt / 6.0
end
end
engine/plugins/solvers/rk5.lua
+64
View File
@@ -0,0 +1,64 @@
-- rk5.lua: 5th-order Runge-Kutta (Cash-Karp, 6 NFE per step)
solver = {
name = "rk5",
display = "RK5 (6 NFE)",
description = "5th-order Cash-Karp Runge-Kutta",
nfe = 6,
order = 5,
needs_model = true,
stateful = false,
stochastic = false,
}
-- Cash-Karp Butcher tableau
local a2 = 1/5
local a3 = 3/10
local a4 = 3/5
local a5 = 1
local a6 = 7/8
local b21 = 1/5
local b31 = 3/40; local b32 = 9/40
local b41 = 3/10; local b42 = -9/10; local b43 = 6/5
local b51 = -11/54; local b52 = 5/2; local b53 = -70/27; local b54 = 35/27
local b61 = 1631/55296; local b62 = 175/512; local b63 = 575/13824; local b64 = 44275/110592; local b65 = 253/4096
-- 5th order weights
local c1 = 37/378; local c3 = 250/621; local c4 = 125/594; local c6 = 512/1771
function step(xt, vt, t_curr, t_prev, n, model_fn, vt_buf)
local dt = t_curr - t_prev
local xt_orig = {}
local k1 = {}
for i = 0, n-1 do xt_orig[i] = xt[i]; k1[i] = vt[i] end
-- k2
for i = 0, n-1 do xt[i] = xt_orig[i] - dt * b21 * k1[i] end
model_fn(xt, t_curr - a2 * dt)
local k2 = {}; for i = 0, n-1 do k2[i] = vt_buf[i] end
-- k3
for i = 0, n-1 do xt[i] = xt_orig[i] - dt * (b31*k1[i] + b32*k2[i]) end
model_fn(xt, t_curr - a3 * dt)
local k3 = {}; for i = 0, n-1 do k3[i] = vt_buf[i] end
-- k4
for i = 0, n-1 do xt[i] = xt_orig[i] - dt * (b41*k1[i] + b42*k2[i] + b43*k3[i]) end
model_fn(xt, t_curr - a4 * dt)
local k4 = {}; for i = 0, n-1 do k4[i] = vt_buf[i] end
-- k5
for i = 0, n-1 do xt[i] = xt_orig[i] - dt * (b51*k1[i] + b52*k2[i] + b53*k3[i] + b54*k4[i]) end
model_fn(xt, t_curr - a5 * dt)
local k5 = {}; for i = 0, n-1 do k5[i] = vt_buf[i] end
-- k6
for i = 0, n-1 do xt[i] = xt_orig[i] - dt * (b61*k1[i] + b62*k2[i] + b63*k3[i] + b64*k4[i] + b65*k5[i]) end
model_fn(xt, t_curr - a6 * dt)
-- Final 5th-order result
for i = 0, n-1 do
xt[i] = xt_orig[i] - dt * (c1*k1[i] + c3*k3[i] + c4*k4[i] + c6*vt_buf[i])
end
end
engine/plugins/solvers/sde.lua
+28
View File
@@ -0,0 +1,28 @@
-- sde.lua: Stochastic Differential Equation Euler solver (1 NFE)
-- Predicts x0 then re-noises with Philox random numbers.
-- Note: requires philox_randn C helper exposed to Lua.
-- Falls back to ODE Euler if no seed info available.
solver = {
name = "sde",
display = "SDE (Stochastic)",
description = "SDE Euler with Philox re-noising for diversity",
nfe = 1,
order = 1,
needs_model = false,
stateful = false,
stochastic = true,
}
function step(xt, vt, t_curr, t_prev, n)
-- SDE needs philox_randn which is a C helper.
-- For now, fall back to ODE Euler. The C++ wrapper will handle
-- the SDE-specific logic (philox noise injection) when it detects
-- this plugin is stochastic and seeds are available.
local dt = t_curr - t_prev
for i = 0, n - 1 do
local x0 = xt[i] - vt[i] * t_curr
-- Without philox, do deterministic ODE step
xt[i] = xt[i] - vt[i] * dt
end
end
engine/plugins/solvers/stork2.lua
+139
View File
@@ -0,0 +1,139 @@
-- stork2.lua: STORK 2 — Stabilized Taylor Orthogonal RK (2nd order, 1 NFE)
-- Uses RKG2 Chebyshev sub-stepping with velocity derivatives from history.
solver = {
name = "stork2",
display = "STORK 2",
description = "2nd-order stabilized Taylor-Chebyshev (1 NFE)",
nfe = 1,
order = 2,
needs_model = false,
stateful = true,
stochastic = false,
params = {
{ key = "substeps", type = "slider", label = "Substeps",
default = 10, min = 2, max = 50, step = 1,
hint = "Number of Chebyshev sub-steps (more = more stable)" },
},
}
local velocity_history = {} -- {vt={}, dt=float}
local function rms(data, n)
local sum = 0
for i = 0, n-1 do sum = sum + data[i] * data[i] end
return math.sqrt(sum / n)
end
local function has_nan_inf(data, n)
for i = 0, n-1 do
local v = data[i]
if v ~= v or v == math.huge or v == -math.huge then return true end
end
return false
end
local function compute_derivatives(vt, n)
local hist = velocity_history
if #hist == 0 then return 0, nil, nil end
local v_prev = hist[#hist].vt
local h1 = hist[#hist].dt
if h1 == 0 then return 0, nil, nil end
local dv = {}
for i = 0, n-1 do dv[i] = (v_prev[i] - vt[i]) / h1 end
local vt_rms = rms(vt, n)
local dv_rms = rms(dv, n)
if vt_rms > 0 and dv_rms * math.abs(h1) > 5 * vt_rms then return 0, nil, nil end
if #hist < 2 then return 1, dv, nil end
local v_prev2 = hist[#hist - 1].vt
local h2 = hist[#hist - 1].dt
if h2 == 0 then return 1, dv, nil end
local denom = h1 * h2 * (h1 + h2)
if math.abs(denom) < 1e-30 then return 1, dv, nil end
local coeff = 2 / denom
local d2v = {}
for i = 0, n-1 do
d2v[i] = coeff * (v_prev2[i] * h1 - v_prev[i] * (h1 + h2) + vt[i] * h2)
end
local d2v_rms = rms(d2v, n)
if vt_rms > 0 and d2v_rms * h1 * h1 > 5 * vt_rms then return 1, dv, nil end
return 2, dv, d2v
end
local function taylor_approx(vt, deriv_order, diff, dv, d2v, n)
local out = {}
if deriv_order >= 2 and d2v then
local half_d2 = 0.5 * diff * diff
for i = 0, n-1 do out[i] = vt[i] + diff * dv[i] + half_d2 * d2v[i] end
elseif deriv_order >= 1 and dv then
for i = 0, n-1 do out[i] = vt[i] + diff * dv[i] end
else
for i = 0, n-1 do out[i] = vt[i] end
end
return out
end
local function rkg2_b(j)
if j <= 0 then return 1 end
if j == 1 then return 1/3 end
return 4 * (j - 1) * (j + 4) / (3 * j * (j + 1) * (j + 2) * (j + 3))
end
local function rkg2_substep(xt, vt, s, t_curr, t_prev, deriv_order, dv, d2v, n)
local dt = t_curr - t_prev
local Y_j_2 = {}; local Y_j_1 = {}; local Y_j = {}
for i = 0, n-1 do Y_j_2[i] = xt[i]; Y_j_1[i] = xt[i]; Y_j[i] = xt[i] end
local s2ps = s * s + s - 2
for j = 1, s do
if j == 1 then
local mu_t = 6 / ((s + 4) * (s - 1))
for i = 0, n-1 do Y_j[i] = Y_j_1[i] - dt * mu_t * vt[i] end
else
local frac = (j == 2) and (4 / (3 * s2ps)) or ((j-1)*(j-1)+(j-1)-2) / s2ps
local bj = rkg2_b(j); local bj1 = rkg2_b(j-1); local bj2 = rkg2_b(j-2)
local mu = (2*j+1) * bj / (j * bj1)
local nu = -(j+1) * bj / (j * bj2)
local mu_t = mu * 6 / ((s+4)*(s-1))
local gamma_t = -mu_t * (1 - j*(j+1)*bj1/2)
local diff = -frac * dt
local vel = taylor_approx(vt, deriv_order, diff, dv, d2v, n)
for i = 0, n-1 do
Y_j[i] = mu*Y_j_1[i] + nu*Y_j_2[i] + (1-mu-nu)*xt[i]
- dt*mu_t*vel[i] - dt*gamma_t*vt[i]
end
end
for i = 0, n-1 do Y_j_2[i] = Y_j_1[i]; Y_j_1[i] = Y_j[i] end
end
return Y_j, not has_nan_inf(Y_j, n)
end
local function update_history(vt, n, dt)
local rec = {vt = {}, dt = dt}
for i = 0, n-1 do rec.vt[i] = vt[i] end
table.insert(velocity_history, rec)
while #velocity_history > 3 do table.remove(velocity_history, 1) end
end
function step(xt, vt, t_curr, t_prev, n)
local dt = t_curr - t_prev
if step_index == 0 then
velocity_history = {}
for i = 0, n-1 do xt[i] = xt[i] - vt[i] * dt end
update_history(vt, n, dt)
return
end
local deriv_order, dv, d2v = compute_derivatives(vt, n)
local s = math.max((params and params.substeps) or 10, 2)
local success = false; local result
while s >= 2 do
result, success = rkg2_substep(xt, vt, s, t_curr, t_prev, deriv_order, dv, d2v, n)
if success then break end
s = math.floor(s / 2)
end
if success then
for i = 0, n-1 do xt[i] = result[i] end
else
for i = 0, n-1 do xt[i] = xt[i] - vt[i] * dt end
end
update_history(vt, n, dt)
end
engine/plugins/solvers/stork4.lua
+186
View File
@@ -0,0 +1,186 @@
-- stork4.lua: STORK 4 — Stabilized Taylor Orthogonal RK (4th order, 1 NFE)
-- Uses ROCK4 Chebyshev sub-stepping with precomputed coefficients.
-- Requires companion stork4_constants.lua.
local C = require("stork4_constants")
solver = {
name = "stork4",
display = "STORK 4",
description = "4th-order stabilized Taylor-ROCK4 (1 NFE)",
nfe = 1,
order = 4,
needs_model = false,
stateful = true,
stochastic = false,
params = {
{ key = "substeps", type = "slider", label = "Substeps",
default = 10, min = 2, max = 50, step = 1,
hint = "Number of ROCK4 sub-steps (more = more stable)" },
},
}
local velocity_history = {}
local function rms(data, n)
local sum = 0
for i = 0, n-1 do sum = sum + data[i] * data[i] end
return math.sqrt(sum / n)
end
local function has_nan_inf(data, n)
for i = 0, n-1 do
local v = data[i]
if v ~= v or v == math.huge or v == -math.huge then return true end
end
return false
end
local function compute_derivatives(vt, n)
local hist = velocity_history
if #hist == 0 then return 0, nil, nil end
local v_prev = hist[#hist].vt
local h1 = hist[#hist].dt
if h1 == 0 then return 0, nil, nil end
local dv = {}
for i = 0, n-1 do dv[i] = (v_prev[i] - vt[i]) / h1 end
local vt_rms = rms(vt, n)
local dv_rms = rms(dv, n)
if vt_rms > 0 and dv_rms * math.abs(h1) > 5 * vt_rms then return 0, nil, nil end
if #hist < 2 then return 1, dv, nil end
local v_prev2 = hist[#hist - 1].vt
local h2 = hist[#hist - 1].dt
if h2 == 0 then return 1, dv, nil end
local denom = h1 * h2 * (h1 + h2)
if math.abs(denom) < 1e-30 then return 1, dv, nil end
local coeff = 2 / denom
local d2v = {}
for i = 0, n-1 do
d2v[i] = coeff * (v_prev2[i] * h1 - v_prev[i] * (h1 + h2) + vt[i] * h2)
end
local d2v_rms = rms(d2v, n)
if vt_rms > 0 and d2v_rms * h1 * h1 > 5 * vt_rms then return 1, dv, nil end
return 2, dv, d2v
end
local function taylor_approx(vt, order, diff, dv, d2v, n)
local out = {}
if order >= 2 and d2v then
local half_d2 = 0.5 * diff * diff
for i = 0, n-1 do out[i] = vt[i] + diff * dv[i] + half_d2 * d2v[i] end
elseif order >= 1 and dv then
for i = 0, n-1 do out[i] = vt[i] + diff * dv[i] end
else
for i = 0, n-1 do out[i] = vt[i] end
end
return out
end
local function rock4_mdegr(s)
local mp1 = 1
for i = 1, C.MS_LEN do
if C.MS[i] >= s then
return C.MS[i], i, mp1 - 1
end
mp1 = mp1 + C.MS[i] * 2 - 1
end
return C.MS[C.MS_LEN], C.MS_LEN, mp1 - 1
end
local function rock4_substep(xt, vt, s, t_curr, t_prev, deriv_order, dv, d2v, n)
local max_rock4 = C.MS[C.MS_LEN]
if s > max_rock4 then s = max_rock4 end
local mdeg, mz, mr = rock4_mdegr(s)
local dt = t_curr - t_prev
local Y_j_2 = {}; local Y_j_1 = {}; local Y_j = {}
for i = 0, n-1 do Y_j_2[i] = xt[i]; Y_j_1[i] = xt[i] end
local ci1 = t_curr
-- ROCK4 Chebyshev recurrence (1-indexed RECF)
for j = 1, mdeg do
if j == 1 then
local temp1 = -dt * C.RECF[mr + 1]
ci1 = t_curr + temp1
for i = 0, n-1 do Y_j_1[i] = xt[i] + temp1 * vt[i] end
else
local diff = ci1 - t_curr
local vel = taylor_approx(vt, deriv_order, diff, dv, d2v, n)
local idx1 = mr + 2 * (j - 2) + 2
local idx2 = mr + 2 * (j - 2) + 3
local temp1 = -dt * C.RECF[idx1]
local temp3 = -C.RECF[idx2]
local temp2 = 1 - temp3
for i = 0, n-1 do
Y_j[i] = temp1 * vel[i] + temp2 * Y_j_1[i] + temp3 * Y_j_2[i]
end
for i = 0, n-1 do Y_j_2[i] = Y_j_1[i]; Y_j_1[i] = Y_j[i] end
ci1 = temp1 + temp2 * ci1 + temp3 * ci1 -- simplified
end
end
-- ROCK4 finishing procedure (4 stages)
local Y_base = Y_j_1
local fpa = C.FPA[mz]; local fpb = C.FPB[mz]
-- F1
local diff1 = ci1 - t_curr
local F1 = taylor_approx(vt, deriv_order, diff1, dv, d2v, n)
local fpa0 = -dt * fpa[1]
local Yf = {}
for i = 0, n-1 do Yf[i] = Y_base[i] + fpa0 * F1[i] end
-- F2
local diff2 = ci1 + fpa0 - t_curr
local F2 = taylor_approx(vt, deriv_order, diff2, dv, d2v, n)
local fpa1 = -dt * fpa[2]; local fpa2 = -dt * fpa[3]
for i = 0, n-1 do Yf[i] = Y_base[i] + fpa1 * F1[i] + fpa2 * F2[i] end
-- F3
local diff3 = ci1 + fpa1 + fpa2 - t_curr
local F3 = taylor_approx(vt, deriv_order, diff3, dv, d2v, n)
local fpa3 = -dt * fpa[4]; local fpa4 = -dt * fpa[5]; local fpa5 = -dt * fpa[6]
-- F4
local diff4 = ci1 + fpa3 + fpa4 + fpa5 - t_curr
local F4 = taylor_approx(vt, deriv_order, diff4, dv, d2v, n)
local fpb0 = -dt * fpb[1]; local fpb1 = -dt * fpb[2]; local fpb2 = -dt * fpb[3]; local fpb3 = -dt * fpb[4]
local result = {}
for i = 0, n-1 do
result[i] = Y_base[i] + fpb0*F1[i] + fpb1*F2[i] + fpb2*F3[i] + fpb3*F4[i]
end
return result, not has_nan_inf(result, n)
end
local function update_history(vt, n, dt)
local rec = {vt = {}, dt = dt}
for i = 0, n-1 do rec.vt[i] = vt[i] end
table.insert(velocity_history, rec)
while #velocity_history > 3 do table.remove(velocity_history, 1) end
end
function step(xt, vt, t_curr, t_prev, n)
local dt = t_curr - t_prev
if step_index == 0 then
velocity_history = {}
for i = 0, n-1 do xt[i] = xt[i] - vt[i] * dt end
update_history(vt, n, dt)
return
end
local deriv_order, dv, d2v = compute_derivatives(vt, n)
local s = math.max((params and params.substeps) or 10, 2)
local success = false; local result
while s >= 2 do
result, success = rock4_substep(xt, vt, s, t_curr, t_prev, deriv_order, dv, d2v, n)
if success then break end
s = math.floor(s / 2)
end
if success then
for i = 0, n-1 do xt[i] = result[i] end
else
for i = 0, n-1 do xt[i] = xt[i] - vt[i] * dt end
end
update_history(vt, n, dt)
end
engine/plugins/solvers/stork4_constants.lua
+1216
View File
File diff suppressed because it is too large Load Diff
engine/plugins/solvers/unipc.lua
+158
View File
@@ -0,0 +1,158 @@
-- unipc.lua: UniPC (Unified Predictor-Corrector, 2 NFE)
-- B(h)1 variant with data prediction in log-SNR space.
solver = {
name = "unipc",
display = "UniPC (2 NFE)",
description = "Unified predictor-corrector in log-SNR space",
nfe = 2,
order = 2,
needs_model = true,
stateful = true,
stochastic = false,
}
local history = {} -- {model_output={}, t=float}
local max_order = 2
local function lambda(t)
t = math.max(t, 1e-7); t = math.min(t, 1 - 1e-7)
return math.log((1 - t) / t)
end
local function expm1(x) return math.exp(x) - 1 end
local function solve_1x1(R, b) return {b[1] / (math.abs(R[1]) > 1e-12 and R[1] or 1)} end
local function solve_2x2(R, b)
local det = R[1]*R[4] - R[2]*R[3]
if math.abs(det) < 1e-12 then return {0, 0} end
local inv = 1 / det
return {(R[4]*b[1] - R[2]*b[2]) * inv, (R[1]*b[2] - R[3]*b[1]) * inv}
end
local function solve(K, R, b)
if K == 1 then return solve_1x1(R, b)
elseif K == 2 then return solve_2x2(R, b)
else return {0} end
end
local function bh1_update(xt, vt, t_curr, t_next, n, model_fn, vt_buf, use_corrector)
-- Data prediction: D_n = x - t * v
local D_n = {}
for i = 0, n-1 do D_n[i] = xt[i] - t_curr * vt[i] end
local lam_curr = lambda(t_curr)
local lam_next = lambda(t_next)
local h = lam_next - lam_curr
local alpha_next = 1 - t_next
local sigma_next = t_next
local sigma_curr = math.max(t_curr, 1e-7)
local hh = -h
local B_h = hh
local h_phi_1 = expm1(hh)
local avail = #history
local order = math.min(max_order, avail + 1)
local K = order
local n_D1 = order - 1
-- h_phi_k sequence
local h_phi_k_vals = {h_phi_1}
local fact = 1; local hpk = h_phi_1
for k = 1, order do
hpk = hpk / hh - 1 / fact
h_phi_k_vals[k+1] = hpk
fact = fact * (k + 1)
end
-- rks, R matrix, b vector
local rks = {}
for i = 1, n_D1 do
local hist_idx = avail - i + 1
local lam_hist = lambda(history[hist_idx].t)
rks[i] = (lam_hist - lam_curr) / h
end
rks[n_D1 + 1] = 1
local R_mat = {}
for row = 1, K do
for col = 1, K do
R_mat[(row-1)*K + col] = rks[col] ^ (row - 1)
end
end
local b_vec = {}
fact = 1; hpk = h_phi_1
for i = 1, K do
hpk = hpk / hh - 1 / fact
b_vec[i] = hpk * fact / B_h
fact = fact * (i + 1)
end
-- D1 differences
local d1 = {}
for i = 1, n_D1 do
local hist_idx = avail - i + 1
local D_hist = history[hist_idx].model_output
local rk_inv = (math.abs(rks[i]) > 1e-12) and (1 / rks[i]) or 0
d1[i] = {}
for j = 0, n-1 do d1[i][j] = (D_hist[j] - D_n[j]) * rk_inv end
end
-- Base term
local sigma_ratio = (math.abs(sigma_curr) > 1e-7) and (sigma_next / sigma_curr) or 0
local x_t_ = {}
for i = 0, n-1 do
x_t_[i] = sigma_ratio * xt[i] - alpha_next * h_phi_1 * D_n[i]
end
-- Predictor
if n_D1 > 0 then
local rhos_p
if order == 2 then rhos_p = {0.5}
else
local Kp = K - 1; local R_p = {}
for row = 1, Kp do for col = 1, Kp do R_p[(row-1)*Kp+col] = R_mat[(row-1)*K+col] end end
rhos_p = solve(Kp, R_p, b_vec)
end
for i = 0, n-1 do
local pred = 0
for k = 1, n_D1 do pred = pred + rhos_p[k] * d1[k][i] end
xt[i] = x_t_[i] - alpha_next * B_h * pred
end
else
for i = 0, n-1 do xt[i] = x_t_[i] end
end
-- Corrector
if use_corrector and model_fn then
model_fn(xt, t_next)
local D_corr_diff = {}
for i = 0, n-1 do
local D_corr = xt[i] - t_next * vt_buf[i]
D_corr_diff[i] = D_corr - D_n[i]
end
local rhos_c
if order == 1 then rhos_c = {0.5}
else rhos_c = solve(K, R_mat, b_vec) end
for i = 0, n-1 do
local corr = 0
for k = 1, n_D1 do corr = corr + rhos_c[k] * d1[k][i] end
corr = corr + rhos_c[K] * D_corr_diff[i]
xt[i] = x_t_[i] - alpha_next * B_h * corr
end
end
-- Update history
table.insert(history, {model_output = D_n, t = t_curr})
while #history > max_order do table.remove(history, 1) end
end
function step(xt, vt, t_curr, t_prev, n, model_fn, vt_buf)
-- Reset state on first step of a new generation
if (step_index or 0) == 0 then history = {} end
bh1_update(xt, vt, t_curr, t_prev, n, model_fn, vt_buf, true)
end
engine/plugins/solvers/unipc_p.lua
+89
View File
@@ -0,0 +1,89 @@
-- unipc_p.lua: UniPC Predictor only (1 NFE, stateful)
-- Same as UniPC but without the corrector step.
solver = {
name = "unipc_p",
display = "UniPC Predictor (1 NFE)",
description = "UniPC predictor-only (no corrector, 1 NFE)",
nfe = 1,
order = 2,
needs_model = false,
stateful = true,
stochastic = false,
}
-- Shares the same logic as unipc.lua but with use_corrector = false
-- For brevity, we duplicate the core with corrector disabled.
local history = {}
local max_order = 2
local function lambda(t)
t = math.max(t, 1e-7); t = math.min(t, 1 - 1e-7)
return math.log((1 - t) / t)
end
local function expm1(x) return math.exp(x) - 1 end
local function solve_1x1(R, b) return {b[1] / (math.abs(R[1]) > 1e-12 and R[1] or 1)} end
local function solve_2x2(R, b)
local det = R[1]*R[4] - R[2]*R[3]
if math.abs(det) < 1e-12 then return {0, 0} end
local inv = 1 / det
return {(R[4]*b[1] - R[2]*b[2]) * inv, (R[1]*b[2] - R[3]*b[1]) * inv}
end
local function solve(K, R, b)
if K == 1 then return solve_1x1(R, b) else return solve_2x2(R, b) end
end
function step(xt, vt, t_curr, t_prev, n)
-- Reset state on first step of a new generation
if (step_index or 0) == 0 then history = {} end
local D_n = {}
for i = 0, n-1 do D_n[i] = xt[i] - t_curr * vt[i] end
local lam_curr = lambda(t_curr)
local lam_next = lambda(t_prev)
local h = lam_next - lam_curr
local alpha_next = 1 - t_prev
local sigma_next = t_prev
local sigma_curr = math.max(t_curr, 1e-7)
local hh = -h
local h_phi_1 = expm1(hh)
local avail = #history
local order = math.min(max_order, avail + 1)
local n_D1 = order - 1
local rks = {}
for i = 1, n_D1 do
local hist_idx = avail - i + 1
rks[i] = (lambda(history[hist_idx].t) - lam_curr) / h
end
local d1 = {}
for i = 1, n_D1 do
local D_hist = history[avail - i + 1].model_output
local rk_inv = (math.abs(rks[i]) > 1e-12) and (1 / rks[i]) or 0
d1[i] = {}
for j = 0, n-1 do d1[i][j] = (D_hist[j] - D_n[j]) * rk_inv end
end
local sigma_ratio = (math.abs(sigma_curr) > 1e-7) and (sigma_next / sigma_curr) or 0
for i = 0, n-1 do
xt[i] = sigma_ratio * xt[i] - alpha_next * h_phi_1 * D_n[i]
end
if n_D1 > 0 then
local rhos_p = (order == 2) and {0.5} or solve(n_D1, {1}, {0.5})
for i = 0, n-1 do
local pred = 0
for k = 1, n_D1 do pred = pred + rhos_p[k] * d1[k][i] end
xt[i] = xt[i] - alpha_next * hh * pred
end
end
table.insert(history, {model_output = D_n, t = t_curr})
while #history > max_order do table.remove(history, 1) end
end