287 lines
13 KiB
Lua
287 lines
13 KiB
Lua
-- ============================================================================
|
|
-- SPDX-License-Identifier: GPL-3.0-or-later
|
|
-- Copyright (C) 2026 Alexander Allan (MDMAchine) -- A&E Concepts
|
|
-- ============================================================================
|
|
|
|
-- MD Hamiltonian V2 -- Energy-Conserving Momentum-Augmented Sampler
|
|
-- MDMAchine | A&E Concepts (c) 2026
|
|
--
|
|
-- Euler-primary architecture with momentum correction layer, sigma-adaptive
|
|
-- decay, confidence gating, spectral momentum, Hamiltonian energy tracking.
|
|
-- Two look-backs (primary + post-step). owns_loop = true. Single NFE.
|
|
-- ============================================================================
|
|
|
|
local C = require("md_solver_commons")
|
|
|
|
-- ── HAMILTONIAN ENERGY ──────────────────────────────────────────────────────
|
|
|
|
local function kinetic_energy(p, mass, n)
|
|
local s = 0.0
|
|
for i = 0, n - 1 do s = s + p[i] * p[i] end
|
|
return 0.5 * s / mass
|
|
end
|
|
|
|
local function potential_energy(x, v_curr, sigma_ratio, n)
|
|
return -C.vec_dot(v_curr, x, n) * sigma_ratio
|
|
end
|
|
|
|
-- ── SOLVER DEFINITION ───────────────────────────────────────────────────────
|
|
|
|
solver = {
|
|
name = "md_hamiltonian_v2",
|
|
display = "MD Hamiltonian V2",
|
|
description = "Energy-conserving momentum-augmented sampler. Euler + momentum correction, spectral weighting, Hamiltonian tracking. Shared anchor stack.",
|
|
nfe = 1,
|
|
order = 1,
|
|
needs_model = false,
|
|
stateful = true,
|
|
stochastic = true,
|
|
owns_loop = true,
|
|
params = {
|
|
-- Momentum
|
|
{ key = "momentum_weight", type = "slider", label = "Momentum Weight",
|
|
default = 0.20, min = 0.0, max = 0.8, step = 0.05,
|
|
hint = "Momentum blend. 0 = pure Euler. Scaled by confidence gating and sigma fadeout." },
|
|
{ key = "momentum_decay", type = "slider", label = "Momentum Decay",
|
|
default = 0.85, min = 0.0, max = 0.99, step = 0.01,
|
|
hint = "Step-to-step carry-over. Sigma-adaptive." },
|
|
{ key = "momentum_ema_alpha", type = "slider", label = "Momentum EMA Alpha",
|
|
default = 0.3, min = 0.05, max = 0.8, step = 0.05,
|
|
hint = "Velocity absorption rate. Sigma-adaptive." },
|
|
{ key = "mass", type = "slider", label = "Particle Mass",
|
|
default = 1.0, min = 0.1, max = 5.0, step = 0.1,
|
|
hint = "Inertial mass." },
|
|
-- Energy
|
|
{ key = "energy_tolerance", type = "slider", label = "Energy Tolerance",
|
|
default = 0.05, min = 0.005, max = 0.5, step = 0.005,
|
|
hint = "Hamiltonian drift before Metropolis correction." },
|
|
{ key = "correction_strength", type = "slider", label = "Correction Strength",
|
|
default = 0.7, min = 0.0, max = 1.0, step = 0.05,
|
|
hint = "Metropolis momentum rescale. 0 = monitor only." },
|
|
{ key = "energy_tracking", type = "select", label = "Energy Tracking",
|
|
default = "adaptive",
|
|
options = {
|
|
{ value = "fixed", label = "Fixed" },
|
|
{ value = "adaptive", label = "Adaptive" },
|
|
{ value = "monitor", label = "Monitor Only" },
|
|
},
|
|
hint = "How H reference evolves." },
|
|
-- Confidence
|
|
{ key = "confidence_floor", type = "slider", label = "Confidence Floor",
|
|
default = 0.2, min = 0.0, max = 0.8, step = 0.05, hint = "Min alignment for momentum." },
|
|
{ key = "confidence_ceiling", type = "slider", label = "Confidence Ceiling",
|
|
default = 0.7, min = 0.3, max = 1.0, step = 0.05, hint = "Full momentum alignment." },
|
|
-- Spectral momentum
|
|
{ key = "spectral_momentum", type = "toggle", label = "Spectral Momentum",
|
|
default = true, hint = "Per-batch 4-band momentum weighting." },
|
|
{ key = "spectral_hi_boost", type = "slider", label = "Spectral HF Boost",
|
|
default = 1.4, min = 1.0, max = 4.0, step = 0.1, hint = "HF momentum multiplier." },
|
|
{ key = "spectral_mid_cut", type = "slider", label = "Spectral Mid Cut",
|
|
default = 0.6, min = 0.1, max = 1.0, step = 0.05, hint = "Mid momentum multiplier." },
|
|
-- Post-step look-back (secondary)
|
|
{ key = "post_look_back", type = "slider", label = "Post-Step Look-Back",
|
|
default = 0.0, min = 0.0, max = 0.6, step = 0.05, hint = "Additional SNR-adaptive EMA. 0 = off (default)." },
|
|
{ key = "post_look_back_snr", type = "slider", label = "Post-Step LB SNR Power",
|
|
default = 1.0, min = 0.5, max = 3.0, step = 0.1, hint = "Falloff." },
|
|
},
|
|
}
|
|
|
|
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 mom_weight = C.num_param(p, "momentum_weight", 0.20)
|
|
local mom_decay = C.num_param(p, "momentum_decay", 0.85)
|
|
local mom_alpha = C.num_param(p, "momentum_ema_alpha", 0.3)
|
|
local mass = C.num_param(p, "mass", 1.0)
|
|
local energy_tol = C.num_param(p, "energy_tolerance", 0.05)
|
|
local corr_str = C.num_param(p, "correction_strength", 0.7)
|
|
local energy_mode = p.energy_tracking or "adaptive"
|
|
local conf_floor = C.num_param(p, "confidence_floor", 0.2)
|
|
local conf_ceil = C.num_param(p, "confidence_ceiling", 0.7)
|
|
local f_spec_mom = C.bool_param(p, "spectral_momentum", true)
|
|
local spec_hi = C.num_param(p, "spectral_hi_boost", 1.4)
|
|
local spec_mid = C.num_param(p, "spectral_mid_cut", 0.6)
|
|
local post_lb_lam = C.num_param(p, "post_look_back", 0.0)
|
|
local post_lb_snr = C.num_param(p, "post_look_back_snr", 1.0)
|
|
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 momentum = nil
|
|
local H_ref = nil
|
|
local post_lb_enabled = (post_lb_lam > 0)
|
|
|
|
local x = C.fa_to_tbl(xt, n)
|
|
|
|
if opts.verbose then
|
|
print(string.format("[HAMILTONIAN V2] Schedule: %d steps | B=%d NPB=%d | weight=%.2f decay=%.2f mass=%.1f",
|
|
n_steps, B, NPB, mom_weight, mom_decay, mass))
|
|
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
|
|
|
|
-- 1. Euler advance
|
|
local x_euler = {}
|
|
for j = 0, n - 1 do x_euler[j] = x[j] + dt * v_curr[j] end
|
|
|
|
-- 2. Momentum update (sigma-adaptive)
|
|
local decay_power = 1.0 + 2.0 * (1.0 - sigma_ratio)
|
|
local effective_decay = mom_decay ^ decay_power
|
|
local effective_alpha = mom_alpha + (1.0 - mom_alpha) * 0.5 * (1.0 - sigma_ratio)
|
|
|
|
if momentum == nil then
|
|
momentum = {}
|
|
for j = 0, n - 1 do momentum[j] = v_curr[j] * mass end
|
|
else
|
|
for j = 0, n - 1 do momentum[j] = momentum[j] * effective_decay end
|
|
for j = 0, n - 1 do
|
|
momentum[j] = (1.0 - effective_alpha) * momentum[j] + effective_alpha * v_curr[j] * mass
|
|
end
|
|
end
|
|
|
|
-- 3. Momentum-predicted position
|
|
local x_mom = {}
|
|
local inv_mass = 1.0 / mass
|
|
for j = 0, n - 1 do x_mom[j] = x[j] + dt * momentum[j] * inv_mass end
|
|
|
|
-- 4. Confidence gating + sigma fadeout (linear)
|
|
local mom_norm = C.vec_norm(momentum, n)
|
|
local v_norm = C.vec_norm(v_curr, n)
|
|
local alignment = 0.0
|
|
if mom_norm > C.EPSILON and v_norm > C.EPSILON then
|
|
alignment = C.vec_dot(momentum, v_curr, n) / (mom_norm * v_norm)
|
|
end
|
|
local confidence = C.smoothstep(alignment, conf_floor, conf_ceil)
|
|
local sigma_fade = sigma_ratio
|
|
local eff_weight = mom_weight * confidence * sigma_fade
|
|
|
|
-- Re-alignment when fighting
|
|
if mom_norm > C.EPSILON and v_norm > C.EPSILON and alignment < 0.3 then
|
|
local blend = 0.3 * (1.0 - alignment)
|
|
for j = 0, n - 1 do
|
|
momentum[j] = (1.0 - blend) * momentum[j] + blend * v_curr[j] * mass * math.abs(dt)
|
|
end
|
|
end
|
|
|
|
-- 5. Blend (per-batch spectral awareness)
|
|
local x_new = {}
|
|
if f_spec_mom and eff_weight > 1e-6 then
|
|
local band_mults = { 1.0, spec_mid, spec_mid, spec_hi }
|
|
local bsize = math.floor(NPB / 4)
|
|
for j = 0, n - 1 do
|
|
local local_idx = j % NPB
|
|
local band = math.min(math.floor(local_idx / bsize), 3)
|
|
local local_w = C.clamp(eff_weight * band_mults[band + 1], 0.0, 0.95)
|
|
x_new[j] = (1.0 - local_w) * x_euler[j] + local_w * x_mom[j]
|
|
end
|
|
else
|
|
for j = 0, n - 1 do
|
|
x_new[j] = (1.0 - eff_weight) * x_euler[j] + eff_weight * x_mom[j]
|
|
end
|
|
end
|
|
|
|
if C.has_nan_inf(x_new, n) then
|
|
for j = 0, n - 1 do x_new[j] = x_euler[j] end
|
|
end
|
|
|
|
-- 6. Hamiltonian energy tracking
|
|
local T = kinetic_energy(momentum, mass, n)
|
|
local V = potential_energy(x_new, v_curr, sigma_ratio, n)
|
|
local H = T + V
|
|
local corrected = false
|
|
|
|
if H_ref == nil then
|
|
H_ref = H
|
|
else
|
|
local rel_drift = math.abs(H - H_ref) / (math.abs(H_ref) + C.EPSILON)
|
|
if energy_mode ~= "monitor" and rel_drift > energy_tol and corr_str > 0 then
|
|
local T_target = H_ref - V
|
|
if T_target < 0.01 then T_target = 0.01 end
|
|
local scale = math.sqrt(T_target / (T + C.EPSILON))
|
|
scale = 1.0 + corr_str * (scale - 1.0)
|
|
scale = C.clamp(scale, 0.5, 2.0)
|
|
for j = 0, n - 1 do momentum[j] = momentum[j] * scale end
|
|
corrected = true
|
|
end
|
|
if energy_mode == "adaptive" then H_ref = 0.95 * H_ref + 0.05 * H end
|
|
end
|
|
|
|
-- 7-8. Identity + tonal anchor (via commons)
|
|
if opts.f_id_anchor then
|
|
C.apply_identity_anchor(x_new, n, sigma_ratio, opts.anchor_sigma, opts.anchor_blend, state)
|
|
end
|
|
if opts.f_tonal then
|
|
C.apply_tonal_anchor(x_new, n, B, NPB, sigma_ratio, opts.anchor_sigma, opts.tonal_str, state)
|
|
end
|
|
|
|
-- 9. Primary look-back (via commons)
|
|
if opts.f_lookback then
|
|
C.apply_look_back(x_new, n, sigma_ratio, opts.lb_lambda, opts.lb_snr_power, state, "lb_prev")
|
|
end
|
|
|
|
-- 10. RMS servo (via commons)
|
|
if opts.f_rms then
|
|
C.apply_rms_servo(x_new, n, B, NPB, sigma_ratio, opts.rms_tgt_min, opts.rms_tgt_max, opts.rms_gain)
|
|
end
|
|
|
|
-- 11. Post-step look-back (secondary, via commons)
|
|
if post_lb_enabled then
|
|
C.apply_look_back(x_new, n, sigma_ratio, post_lb_lam, post_lb_snr, state, "lb2_prev")
|
|
end
|
|
|
|
-- 12. SDE noise + safety clamp (via commons)
|
|
C.apply_sde_noise(x_new, n, sigma_next, opts.eta, opts.seed, step_idx)
|
|
C.apply_safety_clamp(x_new, n, opts.sclamp)
|
|
|
|
if opts.verbose then
|
|
print(string.format(
|
|
"[HAMILTONIAN V2] step %02d | H=%.2f %s | align=%.3f conf=%.2f ew=%.3f | rms=%.3f",
|
|
step_idx, H, corrected and "CORR" or "ok",
|
|
alignment, confidence, eff_weight, 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
|