sim.bend source
sim.bend on the hub · documented module
import Baseimport ./geom.bend as Gimport ./protein.bend as Pimport ./topology.bend as Timport ./force.bend as Fimport ./bonded.bend as Bimport ./rng.bend as Rng# --- Simulation engine v0: velocity Verlet NVT (Berendsen) + steepest# descent with convergence driver, over explicit DynAtom state. ---# Units (AMBER-consistent): lengths A, masses amu, energies kcal/mol, time# fs, charges e, ke(Coulomb) = 332.0637. Accelerations need ACC_F = 4.184e-4# A/fs^2 per kcal/mol/A/amu; velocities are A/fs; temperature from KE via# kB = 0.001987 kcal/mol/K. Dynamics forces here are uniform-LJ with# 1-2/1-3 exclusions + Coulomb + bonds (zero charges and no angle forces# in v0; per-element reporting lives in pdbdemo).# State is one Data list (never affine pairs), so drivers reuse it freely;# loops are single self-recursions with computed flags passed as params.type DynAtom is Data: DA{serial: U32, elem: U32, pos: G.Vec3, vel: G.Vec3, frc: G.Vec3, mass: F32}type SimParams is Data: SP{eps: F32, sig2: F32, c2: F32, ke: F32, kbond: F32}def acc_conv() -> F32: 0.0004184def boltzmann() -> F32: 0.001987def ke_coul() -> F32: 332.0637def dyn_of_atom(h: P.Atom) -> DynAtom: match h: case P.Atom{serial, +elem, pos}: DA{serial, elem, pos, G.V3{0.0, 0.0, 0.0}, G.V3{0.0, 0.0, 0.0}, F.mass(elem)}def dyn_of_atoms(xs: List<&2, P.Atom>) -> List<&2, DynAtom>: match xs: case Nil{}: Nil{} case h <> t: dyn_of_atom(h) <> dyn_of_atoms(t)def dyn_atoms(ds: List<&2, DynAtom>) -> List<&2, P.Atom>: match ds: case Nil{}: Nil{} case h <> t: match h: case DA{serial, elem, pos, vel, frc, mass}: P.Atom{serial, elem, pos} <> dyn_atoms(t)# Zero charges aligned with atoms (no charge model yet; Coulomb kernels# take external charges and are exercised by laws).def zero_charges(xs: List<&2, P.Atom>) -> List<&2, F32>: match xs: case Nil{}: Nil{} case h <> t: 0.0 <> zero_charges(t)# Uniform kicks in [-0.5, 0.5) per component, seeded per serial.def kick_vec(+s: U32) -> G.Vec3: G.V3{(Rng.rng01(s) - 0.5 : F32), (Rng.rng01(U32.add(s, 99991)) - 0.5 : F32), (Rng.rng01(U32.add(s, 199999)) - 0.5 : F32)}def init_kick_atom(h: DynAtom) -> DynAtom: match h: case DA{+serial, elem, pos, vel, frc, mass}: DA{serial, elem, pos, kick_vec(serial), frc, mass}def init_kicks(ds: List<&2, DynAtom>) -> List<&2, DynAtom>: match ds: case Nil{}: Nil{} case h <> t: init_kick_atom(h) <> init_kicks(t)def sum_vels_ds(ds: List<&2, DynAtom>) -> G.Vec3: match ds: case Nil{}: G.V3{0.0, 0.0, 0.0} case h <> t: match h: case DA{serial, elem, pos, vel, frc, mass}: G.Vec3.add(vel, sum_vels_ds(t))def sub_com_ds(ds: List<&2, DynAtom>, +com: G.Vec3) -> List<&2, DynAtom>: match ds: case Nil{}: Nil{} case h <> t: match h: case DA{serial, elem, pos, vel, frc, mass}: DA{serial, elem, pos, G.Vec3.sub(vel, com), frc, mass} <> sub_com_ds(t, com)def ke_head(h: DynAtom) -> F32: match h: case DA{serial, elem, pos, vel, frc, mass}: ((0.5 * mass : F32) * G.Vec3.norm2(vel) : F32)def ke_sim_of(ds: List<&2, DynAtom>) -> F32: match ds: case Nil{}: 0.0 case h <> t: (ke_head(h) + ke_sim_of(t) : F32)def temp_of_ke(ke: F32, n: Nat) -> F32: ((2.0 * (ke / acc_conv() : F32) : F32) / ((3.0 * F32.from_nat(n) : F32) * boltzmann() : F32) : F32)def temp_of_ds(ds: List<&2, DynAtom>, n: Nat) -> F32: temp_of_ke(ke_sim_of(ds), n)def vv_scale_head(h: DynAtom, lam: F32) -> DynAtom: match h: case DA{serial, elem, pos, vel, frc, mass}: DA{serial, elem, pos, G.Vec3.scale(lam, vel), frc, mass}def vv_scale(ds: List<&2, DynAtom>, +lam: F32) -> List<&2, DynAtom>: match ds: case Nil{}: Nil{} case h <> t: vv_scale_head(h, lam) <> vv_scale(t, lam)def scale_to_temp(ds: List<&2, DynAtom>, +now: F32, want: F32) -> List<&2, DynAtom>: vv_scale(ds, Bool.pick(F32, F32.is_eq(now, 0.0), 1.0, F32.sqrt((want / now : F32))))def com_of_n(+kicked: List<&2, DynAtom>, n: Nat) -> G.Vec3: G.Vec3.scale((1.0 / F32.from_nat(n) : F32), sum_vels_ds(kicked))def init_scaled(+nocom: List<&2, DynAtom>, +n: Nat, temp0: F32) -> List<&2, DynAtom>: scale_to_temp(nocom, temp_of_ds(nocom, n), temp0)def init_nocom(+kicked: List<&2, DynAtom>, +n: Nat, temp0: F32) -> List<&2, DynAtom>: init_scaled(sub_com_ds(kicked, com_of_n(kicked, n)), n, temp0)def init_vels(+ds: List<&2, DynAtom>, +temp0: F32) -> List<&2, DynAtom>: init_nocom(init_kicks(ds), List.length(&2, DynAtom, ds), temp0)# Set per-atom forces from a force list (truncates to the shorter).def dyn_set_force(ds: List<&2, DynAtom>, fs: List<&2, G.Vec3>) -> List<&2, DynAtom>: match ds: case Nil{}: Nil{} case h <> t: match h: case DA{serial, elem, pos, vel, frc, mass}: match fs: case Nil{}: DA{serial, elem, pos, vel, G.V3{0.0, 0.0, 0.0}, mass} <> t case f <> ft: DA{serial, elem, pos, vel, f, mass} <> dyn_set_force(t, ft)def add_forces(a: List<&2, G.Vec3>, b: List<&2, G.Vec3>) -> List<&2, G.Vec3>: match a b: case Nil{} Nil{}: Nil{} case Nil{} bh <> bt: Nil{} case ah <> at Nil{}: Nil{} case ah <> at bh <> bt: G.Vec3.add(ah, bh) <> add_forces(at, bt)# Fresh forces (uniform-LJ with exclusions + Coulomb + bonds) into state.def set_forces(+ds: List<&2, DynAtom>, +qs: List<&2, F32>, +bonds: List<&2, U32>, +p: SimParams) -> List<&2, DynAtom>: match p: case SP{+eps, +sig2, +c2, +ke, +kbond}: dyn_set_force(ds, add_forces(add_forces(B.lj_forces_excl(dyn_atoms(ds), dyn_atoms(ds), eps, sig2, c2, bonds), F.coul_forces(dyn_atoms(ds), qs, dyn_atoms(ds), qs, ke, c2)), B.bond_forces(dyn_atoms(ds), bonds, kbond)))# Velocity Verlet halves over explicit state (forces must be fresh).def vv_half_head(h: DynAtom, +dt: F32) -> DynAtom: match h: case DA{serial, elem, pos, vel, +frc, +mass}: acc = G.Vec3.scale(((1.0 / mass : F32) * acc_conv() : F32), frc) +vh = G.Vec3.add(vel, G.Vec3.scale((0.5 * dt : F32), acc)) DA{serial, elem, G.Vec3.add(pos, G.Vec3.scale(dt, vh)), vh, frc, mass}def vv_half(ds: List<&2, DynAtom>, +dt: F32) -> List<&2, DynAtom>: match ds: case Nil{}: Nil{} case h <> t: vv_half_head(h, dt) <> vv_half(t, dt)def vv_full_head(h: DynAtom, dt: F32) -> DynAtom: match h: case DA{serial, elem, pos, vel, +frc, +mass}: acc = G.Vec3.scale(((1.0 / mass : F32) * acc_conv() : F32), frc) DA{serial, elem, pos, G.Vec3.add(vel, G.Vec3.scale((0.5 * dt : F32), acc)), frc, mass}def vv_full(ds: List<&2, DynAtom>, +dt: F32) -> List<&2, DynAtom>: match ds: case Nil{}: Nil{} case h <> t: vv_full_head(h, dt) <> vv_full(t, dt)# Berendsen thermostat folded into the step (T == 0 keeps velocities).def md_thermo(+ds: List<&2, DynAtom>, +dt: F32, +tau: F32, +temp0: F32, +n: Nat) -> List<&2, DynAtom>: +now = temp_of_ds(ds, n) vv_scale(ds, Bool.pick(F32, F32.is_eq(now, 0.0), 1.0, F32.sqrt((1.0 + (dt / tau : F32) * ((temp0 / now : F32) - 1.0 : F32) : F32))))def md_step(+ds: List<&2, DynAtom>, +qs: List<&2, F32>, +bonds: List<&2, U32>, +p: SimParams, +dt: F32, +tau: F32, +temp0: F32, +n: Nat) -> List<&2, DynAtom>: md_thermo(vv_full(set_forces(vv_half(ds, dt), qs, bonds, p), dt), dt, tau, temp0, n)def md_run(steps: Nat, ds: List<&2, DynAtom>, +qs: List<&2, F32>, +bonds: List<&2, U32>, +p: SimParams, +dt: F32, +tau: F32, +temp0: F32, +n: Nat) -> List<&2, DynAtom>: match steps: case 0n: ds case 1n+f: md_run(f, md_step(ds, qs, bonds, p, dt, tau, temp0, n), qs, bonds, p, dt, tau, temp0, n)# Steepest descent over state (fixed fuel; CLI reports convergedness).# Normalized steps (fixed distance along each force) relax clashes that# raw F*dt would overshoot; raw forces still feed MD and max-force.def sd_norm_head(h: DynAtom, +step: F32) -> DynAtom: match h: case DA{serial, elem, pos, vel, +frc, mass}: +n2 = G.Vec3.norm2(frc) DA{serial, elem, G.Vec3.add(pos, G.Vec3.scale(step, Bool.pick(G.Vec3, F32.is_eq(n2, 0.0), G.V3{0.0, 0.0, 0.0}, G.Vec3.scale((1.0 / F32.sqrt(n2) : F32), frc)))), vel, frc, mass}def sd_norm_sweep(ds: List<&2, DynAtom>, +step: F32) -> List<&2, DynAtom>: match ds: case Nil{}: Nil{} case h <> t: sd_norm_head(h, step) <> sd_norm_sweep(t, step)def min_run(fuel: Nat, ds: List<&2, DynAtom>, +qs: List<&2, F32>, +bonds: List<&2, U32>, +p: SimParams, +dt: F32) -> List<&2, DynAtom>: match fuel: case 0n: ds case 1n+f: fds = set_forces(ds, qs, bonds, p) min_run(f, sd_norm_sweep(fds, dt), qs, bonds, p, dt)def maxforce2_of(ds: List<&2, DynAtom>) -> F32: match ds: case Nil{}: 0.0 case h <> t: match h: case DA{serial, elem, pos, vel, frc, mass}: F32.max(G.Vec3.norm2(frc), maxforce2_of(t))def converged_ds(ds: List<&2, DynAtom>, +tol: F32) -> Bool: F32.is_lt(maxforce2_of(ds), (tol * tol : F32))# Flat bonds back to pairs for the pair-list energies.def pairs_of_flat(bs: List<&2, U32>) -> List<&1, U32 & U32>: match bs: case Nil{}: Nil{} case x <> Nil{}: Nil{} case x <> y <> t: (x, y) <> pairs_of_flat(t)# Reportable potential: excluded uniform LJ + Coulomb + bonds.def energy_sum(+xs: List<&2, P.Atom>, +qs: List<&2, F32>, +bonds: List<&2, U32>, +p: SimParams) -> F32: match p: case SP{+eps, +sig2, +c2, +ke, +kbond}: ((B.lj_total_excl(xs, xs, eps, sig2, c2, bonds) + B.coul_total_excl(xs, qs, xs, qs, ke, c2, bonds) : F32) + B.bond_total(xs, pairs_of_flat(bonds), kbond) : F32)def total_energy(+ds: List<&2, DynAtom>, +qs: List<&2, F32>, +bonds: List<&2, U32>, +p: SimParams) -> F32: energy_sum(dyn_atoms(ds), qs, bonds, p)