219 lines
9.1 KiB
Lua
219 lines
9.1 KiB
Lua
-- ============================================================================
|
|
-- SPDX-License-Identifier: GPL-3.0-or-later
|
|
-- Copyright (C) 2026 Alexander Allan (MDMAchine) -- A&E Concepts
|
|
-- ============================================================================
|
|
|
|
-- MD Vortex V1 -- Vorticity Damping Sampler
|
|
-- MDMAchine | A&E Concepts (c) 2026
|
|
--
|
|
-- Multi-scale curl proxy detects rotational velocity energy. Enstrophy EMA
|
|
-- triggers targeted vortex shedding (direct damping of vortical component).
|
|
-- Euler advance with cleaned velocity. owns_loop = true. Single NFE.
|
|
-- ============================================================================
|
|
|
|
local C = require("md_solver_commons")
|
|
|
|
-- ── CURL PROXY (per-batch) ──────────────────────────────────────────────────
|
|
|
|
local function compute_curl_proxy(v_curr, v_prev, off, cnt, multi_scale)
|
|
local curl = {}
|
|
for i = 0, cnt - 1 do curl[i] = 0.0 end
|
|
|
|
for i = 1, cnt - 2 do
|
|
local grad_curr = (v_curr[off + i + 1] - v_curr[off + i - 1]) * 0.5
|
|
local grad_prev = (v_prev[off + i + 1] - v_prev[off + i - 1]) * 0.5
|
|
curl[i] = curl[i] + (grad_curr - grad_prev)
|
|
end
|
|
|
|
if multi_scale then
|
|
for i = 2, cnt - 3 do
|
|
local grad_curr = (v_curr[off + i + 2] - v_curr[off + i - 2]) * 0.25
|
|
local grad_prev = (v_prev[off + i + 2] - v_prev[off + i - 2]) * 0.25
|
|
curl[i] = curl[i] + 0.5 * (grad_curr - grad_prev)
|
|
end
|
|
for i = 4, cnt - 5 do
|
|
local grad_curr = (v_curr[off + i + 4] - v_curr[off + i - 4]) * 0.125
|
|
local grad_prev = (v_prev[off + i + 4] - v_prev[off + i - 4]) * 0.125
|
|
curl[i] = curl[i] + 0.25 * (grad_curr - grad_prev)
|
|
end
|
|
end
|
|
|
|
local enstrophy, curl_max = 0.0, 0.0
|
|
for i = 0, cnt - 1 do
|
|
enstrophy = enstrophy + curl[i] * curl[i]
|
|
local ac = math.abs(curl[i])
|
|
if ac > curl_max then curl_max = ac end
|
|
end
|
|
return curl, enstrophy / math.max(cnt, 1), curl_max
|
|
end
|
|
|
|
local function apply_shedding(v_out, off, cnt, curl, enstrophy_ema,
|
|
threshold, strength, progressive)
|
|
if strength < 1e-6 or enstrophy_ema <= threshold then return false end
|
|
|
|
local eff_strength = strength
|
|
if progressive then
|
|
local overshoot = C.clamp((enstrophy_ema - threshold) / (threshold + C.EPSILON), 0.0, 2.0)
|
|
eff_strength = strength * (overshoot / 2.0)
|
|
end
|
|
if eff_strength < 1e-6 then return false end
|
|
|
|
local correction = {}
|
|
correction[0] = 0.0
|
|
for i = 1, cnt - 1 do correction[i] = correction[i - 1] + curl[i] end
|
|
|
|
local corr_energy = 0.0
|
|
for i = 0, cnt - 1 do corr_energy = corr_energy + correction[i] * correction[i] end
|
|
corr_energy = math.sqrt(corr_energy / math.max(cnt, 1) + C.EPSILON)
|
|
|
|
local scale = eff_strength / (corr_energy + C.EPSILON)
|
|
local v_rms = 0.0
|
|
for i = 0, cnt - 1 do v_rms = v_rms + v_out[off + i] * v_out[off + i] end
|
|
v_rms = math.sqrt(v_rms / math.max(cnt, 1) + C.EPSILON)
|
|
local max_corr = 0.05 * v_rms
|
|
|
|
for i = 0, cnt - 1 do
|
|
v_out[off + i] = v_out[off + i] - C.clamp(correction[i] * scale, -max_corr, max_corr)
|
|
end
|
|
return true
|
|
end
|
|
|
|
-- ── SOLVER DEFINITION ───────────────────────────────────────────────────────
|
|
|
|
solver = {
|
|
name = "md_vortex_v1",
|
|
display = "MD Vortex V1",
|
|
description = "Vorticity damping sampler. Multi-scale curl proxy, enstrophy tracking, targeted shedding. Batch-aware, shared anchor stack.",
|
|
nfe = 1,
|
|
order = 1,
|
|
needs_model = false,
|
|
stateful = true,
|
|
stochastic = true,
|
|
owns_loop = true,
|
|
params = {
|
|
{ key = "shedding_strength", type = "slider", label = "Shedding Strength",
|
|
default = 0.30, min = 0.0, max = 1.0, step = 0.05,
|
|
hint = "Vortical energy damping. 0 = monitor only." },
|
|
{ key = "enstrophy_threshold", type = "slider", label = "Enstrophy Threshold",
|
|
default = 0.02, min = 0.001, max = 0.2, step = 0.001,
|
|
hint = "EMA level triggering shedding. Calibrate with verbose=true." },
|
|
{ key = "enstrophy_ema_alpha", type = "slider", label = "Enstrophy EMA Alpha",
|
|
default = 0.1, min = 0.02, max = 0.5, step = 0.02,
|
|
hint = "Tracker responsiveness." },
|
|
{ key = "multi_scale", type = "toggle", label = "Multi-Scale Curl",
|
|
default = true, hint = "Adjacent + skip-2 + skip-4 gradient changes." },
|
|
{ key = "sigma_gate", type = "slider", label = "Sigma Gate",
|
|
default = 0.9, min = 0.5, max = 1.0, step = 0.05,
|
|
hint = "Sigma fraction above which shedding is disabled." },
|
|
{ key = "progressive_shedding", type = "toggle", label = "Progressive Shedding",
|
|
default = true, hint = "Strength scales with enstrophy overshoot." },
|
|
},
|
|
}
|
|
|
|
C.append_common_params(solver.params)
|
|
|
|
-- ── SAMPLE ──────────────────────────────────────────────────────────────────
|
|
|
|
function sample(xt, vt_buf, schedule, n, model_fn)
|
|
local p = params or {}
|
|
local B, NPB = C.get_batch_routing(n)
|
|
|
|
local shed_str = C.num_param(p, "shedding_strength", 0.30)
|
|
local enst_thresh = C.num_param(p, "enstrophy_threshold", 0.02)
|
|
local enst_alpha = C.num_param(p, "enstrophy_ema_alpha", 0.1)
|
|
local f_multi = C.bool_param(p, "multi_scale", true)
|
|
local sigma_gate = C.num_param(p, "sigma_gate", 0.9)
|
|
local f_prog = C.bool_param(p, "progressive_shedding", true)
|
|
local opts = C.read_common_opts(p)
|
|
local state = C.new_state()
|
|
|
|
-- Engine schedule has NO trailing 0 (fix ported from 46c081e): iterate all ns
|
|
-- entries so the last iteration gets sigma_next = 0.0 and the terminal branch
|
|
-- performs the final x0 projection. With ns - 1 that branch is dead code and
|
|
-- the output keeps ~final-sigma noise.
|
|
local ns, n_steps = #schedule, #schedule
|
|
if n_steps < 1 then return end
|
|
|
|
local sigma_max = schedule[1]
|
|
local v_prev = nil
|
|
local enst_ema = {}
|
|
for b = 0, B - 1 do enst_ema[b] = 0.0 end
|
|
|
|
local x = C.fa_to_tbl(xt, n)
|
|
|
|
if opts.verbose then
|
|
print(string.format("[VORTEX V1] Schedule: %d steps | B=%d NPB=%d | shed=%.2f thresh=%.4f",
|
|
n_steps, B, NPB, shed_str, enst_thresh))
|
|
end
|
|
|
|
for i = 1, n_steps do
|
|
local sigma_curr = schedule[i]
|
|
local sigma_next = (i < ns) and schedule[i + 1] or 0.0
|
|
local step_idx = i - 1
|
|
local sigma_ratio = C.clamp(sigma_curr / math.max(sigma_max, C.EPSILON), 0.0, 1.0)
|
|
|
|
if sigma_next == 0.0 then
|
|
C.tbl_to_fa(x, xt, n)
|
|
model_fn(xt, sigma_curr)
|
|
local v_final = C.fa_to_tbl(vt_buf, n)
|
|
for j = 0, n - 1 do x[j] = x[j] - v_final[j] * sigma_curr end
|
|
break
|
|
end
|
|
|
|
C.tbl_to_fa(x, xt, n)
|
|
model_fn(xt, sigma_curr)
|
|
local v_curr = C.fa_to_tbl(vt_buf, n)
|
|
local dt = sigma_next - sigma_curr
|
|
|
|
-- Relational decomposition
|
|
if opts.rw > 0 then
|
|
C.apply_relational(v_curr, n, B, NPB, sigma_ratio, sigma_max,
|
|
opts.rw, opts.rw_sigma_pow, opts.drift_on, opts.drift_thr, x)
|
|
end
|
|
|
|
local shedding_active, max_enst, max_curl = false, 0.0, 0.0
|
|
|
|
if v_prev ~= nil and sigma_ratio < sigma_gate then
|
|
local v_shed = C.vec_clone(v_curr, n)
|
|
for b = 0, B - 1 do
|
|
local off = b * NPB
|
|
local curl, enstrophy, curl_max = compute_curl_proxy(v_curr, v_prev, off, NPB, f_multi)
|
|
enst_ema[b] = (1.0 - enst_alpha) * enst_ema[b] + enst_alpha * enstrophy
|
|
if enst_ema[b] > max_enst then max_enst = enst_ema[b] end
|
|
if curl_max > max_curl then max_curl = curl_max end
|
|
if apply_shedding(v_shed, off, NPB, curl, enst_ema[b], enst_thresh, shed_str, f_prog) then
|
|
shedding_active = true
|
|
end
|
|
end
|
|
if not C.has_nan_inf(v_shed, n) then v_curr = v_shed end
|
|
end
|
|
|
|
v_prev = C.vec_clone(v_curr, n)
|
|
|
|
local x_new = {}
|
|
for j = 0, n - 1 do x_new[j] = x[j] + dt * v_curr[j] end
|
|
if C.has_nan_inf(x_new, n) then
|
|
local v_raw = C.fa_to_tbl(vt_buf, n)
|
|
for j = 0, n - 1 do x_new[j] = x[j] + dt * v_raw[j] end
|
|
end
|
|
|
|
opts.sigma_next = sigma_next
|
|
opts.step_idx = step_idx
|
|
C.post_advance(x_new, n, B, NPB, sigma_ratio, opts, state)
|
|
|
|
if opts.verbose then
|
|
print(string.format("[VORTEX V1] step %02d | enst_ema=%.5f %s | curl_max=%.4f | rms=%.3f",
|
|
step_idx, enst_ema[0] or 0, shedding_active and "SHEDDING" or "quiet",
|
|
max_curl, C.rms(x_new, n)))
|
|
end
|
|
|
|
x = x_new
|
|
C.tbl_to_fa(x, xt, n)
|
|
C.tbl_to_fa(v_curr, vt_buf, n)
|
|
if on_step(step_idx, sigma_curr, sigma_next) then return end
|
|
x = C.fa_to_tbl(xt, n)
|
|
end
|
|
|
|
C.tbl_to_fa(x, xt, n)
|
|
end
|