sim.bend

bend-prot@0.1.0.0 · 271 lines · 43 definitions

1import Base
2import bend-prot@0.1.0.0/geom.bend as G
3import bend-prot@0.1.0.0/protein.bend as P
4import bend-prot@0.1.0.0/topology.bend as T
5import bend-prot@0.1.0.0/force.bend as F
6import bend-prot@0.1.0.0/bonded.bend as B
7import bend-prot@0.1.0.0/rng.bend as Rng
8
9# --- Simulation engine v0: velocity Verlet NVT (Berendsen) + steepest
10# descent with convergence driver, over explicit DynAtom state. ---
11# Units (AMBER-consistent): lengths A, masses amu, energies kcal/mol, time
12# fs, charges e, ke(Coulomb) = 332.0637. Accelerations need ACC_F = 4.184e-4
13# A/fs^2 per kcal/mol/A/amu; velocities are A/fs; temperature from KE via
14# kB = 0.001987 kcal/mol/K. Dynamics forces here are uniform-LJ with
15# 1-2/1-3 exclusions + Coulomb + bonds (zero charges and no angle forces
16# in v0; per-element reporting lives in pdbdemo).
17# State is one Data list (never affine pairs), so drivers reuse it freely;
18# loops are single self-recursions with computed flags passed as params.
19
20type DynAtom is Data:
21  DA{serial: U32, elem: U32, pos: G.Vec3, vel: G.Vec3, frc: G.Vec3, mass: F32}
22
23type SimParams is Data:
24  SP{eps: F32, sig2: F32, c2: F32, ke: F32, kbond: F32}
25
26def acc_conv() -> F32:
27  0.0004184
28
29def boltzmann() -> F32:
30  0.001987
31
32def ke_coul() -> F32:
33  332.0637
34
35def dyn_of_atom(h: P.Atom) -> DynAtom:
36  match h:
37    case P.Atom{serial, +elem, pos}:
38      DA{serial, elem, pos, G.V3{0.0, 0.0, 0.0}, G.V3{0.0, 0.0, 0.0}, F.mass(elem)}
39
40def dyn_of_atoms(xs: List<&2, P.Atom>) -> List<&2, DynAtom>:
41  match xs:
42    case Nil{}:
43      Nil{}
44    case h <> t:
45      dyn_of_atom(h) <> dyn_of_atoms(t)
46
47def dyn_atoms(ds: List<&2, DynAtom>) -> List<&2, P.Atom>:
48  match ds:
49    case Nil{}:
50      Nil{}
51    case h <> t:
52      match h:
53        case DA{serial, elem, pos, vel, frc, mass}:
54          P.Atom{serial, elem, pos} <> dyn_atoms(t)
55
56# Zero charges aligned with atoms (no charge model yet; Coulomb kernels
57# take external charges and are exercised by laws).
58def zero_charges(xs: List<&2, P.Atom>) -> List<&2, F32>:
59  match xs:
60    case Nil{}:
61      Nil{}
62    case h <> t:
63      0.0 <> zero_charges(t)
64
65# Uniform kicks in [-0.5, 0.5) per component, seeded per serial.
66def kick_vec(+s: U32) -> G.Vec3:
67  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)}
68
69def init_kick_atom(h: DynAtom) -> DynAtom:
70  match h:
71    case DA{+serial, elem, pos, vel, frc, mass}:
72      DA{serial, elem, pos, kick_vec(serial), frc, mass}
73
74def init_kicks(ds: List<&2, DynAtom>) -> List<&2, DynAtom>:
75  match ds:
76    case Nil{}:
77      Nil{}
78    case h <> t:
79      init_kick_atom(h) <> init_kicks(t)
80
81def sum_vels_ds(ds: List<&2, DynAtom>) -> G.Vec3:
82  match ds:
83    case Nil{}:
84      G.V3{0.0, 0.0, 0.0}
85    case h <> t:
86      match h:
87        case DA{serial, elem, pos, vel, frc, mass}:
88          G.Vec3.add(vel, sum_vels_ds(t))
89
90def sub_com_ds(ds: List<&2, DynAtom>, +com: G.Vec3) -> List<&2, DynAtom>:
91  match ds:
92    case Nil{}:
93      Nil{}
94    case h <> t:
95      match h:
96        case DA{serial, elem, pos, vel, frc, mass}:
97          DA{serial, elem, pos, G.Vec3.sub(vel, com), frc, mass} <> sub_com_ds(t, com)
98
99def ke_head(h: DynAtom) -> F32:
100  match h:
101    case DA{serial, elem, pos, vel, frc, mass}:
102      ((0.5 * mass : F32) * G.Vec3.norm2(vel) : F32)
103
104def ke_sim_of(ds: List<&2, DynAtom>) -> F32:
105  match ds:
106    case Nil{}:
107      0.0
108    case h <> t:
109      (ke_head(h) + ke_sim_of(t) : F32)
110
111def temp_of_ke(ke: F32, n: Nat) -> F32:
112  ((2.0 * (ke / acc_conv() : F32) : F32) / ((3.0 * F32.from_nat(n) : F32) * boltzmann() : F32) : F32)
113
114def temp_of_ds(ds: List<&2, DynAtom>, n: Nat) -> F32:
115  temp_of_ke(ke_sim_of(ds), n)
116
117def vv_scale_head(h: DynAtom, lam: F32) -> DynAtom:
118  match h:
119    case DA{serial, elem, pos, vel, frc, mass}:
120      DA{serial, elem, pos, G.Vec3.scale(lam, vel), frc, mass}
121
122def vv_scale(ds: List<&2, DynAtom>, +lam: F32) -> List<&2, DynAtom>:
123  match ds:
124    case Nil{}:
125      Nil{}
126    case h <> t:
127      vv_scale_head(h, lam) <> vv_scale(t, lam)
128
129def scale_to_temp(ds: List<&2, DynAtom>, +now: F32, want: F32) -> List<&2, DynAtom>:
130  vv_scale(ds, Bool.pick(F32, F32.is_eq(now, 0.0), 1.0, F32.sqrt((want / now : F32))))
131
132def com_of_n(+kicked: List<&2, DynAtom>, n: Nat) -> G.Vec3:
133  G.Vec3.scale((1.0 / F32.from_nat(n) : F32), sum_vels_ds(kicked))
134
135def init_scaled(+nocom: List<&2, DynAtom>, +n: Nat, temp0: F32) -> List<&2, DynAtom>:
136  scale_to_temp(nocom, temp_of_ds(nocom, n), temp0)
137
138def init_nocom(+kicked: List<&2, DynAtom>, +n: Nat, temp0: F32) -> List<&2, DynAtom>:
139  init_scaled(sub_com_ds(kicked, com_of_n(kicked, n)), n, temp0)
140
141def init_vels(+ds: List<&2, DynAtom>, +temp0: F32) -> List<&2, DynAtom>:
142  init_nocom(init_kicks(ds), List.length(&2, DynAtom, ds), temp0)
143
144# Set per-atom forces from a force list (truncates to the shorter).
145def dyn_set_force(ds: List<&2, DynAtom>, fs: List<&2, G.Vec3>) -> List<&2, DynAtom>:
146  match ds:
147    case Nil{}:
148      Nil{}
149    case h <> t:
150      match h:
151        case DA{serial, elem, pos, vel, frc, mass}:
152          match fs:
153            case Nil{}:
154              DA{serial, elem, pos, vel, G.V3{0.0, 0.0, 0.0}, mass} <> t
155            case f <> ft:
156              DA{serial, elem, pos, vel, f, mass} <> dyn_set_force(t, ft)
157
158def add_forces(a: List<&2, G.Vec3>, b: List<&2, G.Vec3>) -> List<&2, G.Vec3>:
159  match a b:
160    case Nil{} Nil{}:
161      Nil{}
162    case Nil{} bh <> bt:
163      Nil{}
164    case ah <> at Nil{}:
165      Nil{}
166    case ah <> at bh <> bt:
167      G.Vec3.add(ah, bh) <> add_forces(at, bt)
168
169# Fresh forces (uniform-LJ with exclusions + Coulomb + bonds) into state.
170def set_forces(+ds: List<&2, DynAtom>, +qs: List<&2, F32>, +bonds: List<&2, U32>, +p: SimParams) -> List<&2, DynAtom>:
171  match p:
172    case SP{+eps, +sig2, +c2, +ke, +kbond}:
173      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)))
174
175# Velocity Verlet halves over explicit state (forces must be fresh).
176def vv_half_head(h: DynAtom, +dt: F32) -> DynAtom:
177  match h:
178    case DA{serial, elem, pos, vel, +frc, +mass}:
179      acc = G.Vec3.scale(((1.0 / mass : F32) * acc_conv() : F32), frc)
180      +vh = G.Vec3.add(vel, G.Vec3.scale((0.5 * dt : F32), acc))
181      DA{serial, elem, G.Vec3.add(pos, G.Vec3.scale(dt, vh)), vh, frc, mass}
182
183def vv_half(ds: List<&2, DynAtom>, +dt: F32) -> List<&2, DynAtom>:
184  match ds:
185    case Nil{}:
186      Nil{}
187    case h <> t:
188      vv_half_head(h, dt) <> vv_half(t, dt)
189
190def vv_full_head(h: DynAtom, dt: F32) -> DynAtom:
191  match h:
192    case DA{serial, elem, pos, vel, +frc, +mass}:
193      acc = G.Vec3.scale(((1.0 / mass : F32) * acc_conv() : F32), frc)
194      DA{serial, elem, pos, G.Vec3.add(vel, G.Vec3.scale((0.5 * dt : F32), acc)), frc, mass}
195
196def vv_full(ds: List<&2, DynAtom>, +dt: F32) -> List<&2, DynAtom>:
197  match ds:
198    case Nil{}:
199      Nil{}
200    case h <> t:
201      vv_full_head(h, dt) <> vv_full(t, dt)
202
203# Berendsen thermostat folded into the step (T == 0 keeps velocities).
204def md_thermo(+ds: List<&2, DynAtom>, +dt: F32, +tau: F32, +temp0: F32, +n: Nat) -> List<&2, DynAtom>:
205  +now = temp_of_ds(ds, n)
206  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))))
207
208def 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>:
209  md_thermo(vv_full(set_forces(vv_half(ds, dt), qs, bonds, p), dt), dt, tau, temp0, n)
210
211def 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>:
212  match steps:
213    case 0n:
214      ds
215    case 1n+f:
216      md_run(f, md_step(ds, qs, bonds, p, dt, tau, temp0, n), qs, bonds, p, dt, tau, temp0, n)
217
218# Steepest descent over state (fixed fuel; CLI reports convergedness).
219# Normalized steps (fixed distance along each force) relax clashes that
220# raw F*dt would overshoot; raw forces still feed MD and max-force.
221def sd_norm_head(h: DynAtom, +step: F32) -> DynAtom:
222  match h:
223    case DA{serial, elem, pos, vel, +frc, mass}:
224      +n2 = G.Vec3.norm2(frc)
225      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}
226
227def sd_norm_sweep(ds: List<&2, DynAtom>, +step: F32) -> List<&2, DynAtom>:
228  match ds:
229    case Nil{}:
230      Nil{}
231    case h <> t:
232      sd_norm_head(h, step) <> sd_norm_sweep(t, step)
233
234def min_run(fuel: Nat, ds: List<&2, DynAtom>, +qs: List<&2, F32>, +bonds: List<&2, U32>, +p: SimParams, +dt: F32) -> List<&2, DynAtom>:
235  match fuel:
236    case 0n:
237      ds
238    case 1n+f:
239      fds = set_forces(ds, qs, bonds, p)
240      min_run(f, sd_norm_sweep(fds, dt), qs, bonds, p, dt)
241
242def maxforce2_of(ds: List<&2, DynAtom>) -> F32:
243  match ds:
244    case Nil{}:
245      0.0
246    case h <> t:
247      match h:
248        case DA{serial, elem, pos, vel, frc, mass}:
249          F32.max(G.Vec3.norm2(frc), maxforce2_of(t))
250
251def converged_ds(ds: List<&2, DynAtom>, +tol: F32) -> Bool:
252  F32.is_lt(maxforce2_of(ds), (tol * tol : F32))
253
254# Flat bonds back to pairs for the pair-list energies.
255def pairs_of_flat(bs: List<&2, U32>) -> List<&1, U32 & U32>:
256  match bs:
257    case Nil{}:
258      Nil{}
259    case x <> Nil{}:
260      Nil{}
261    case x <> y <> t:
262      (x, y) <> pairs_of_flat(t)
263
264# Reportable potential: excluded uniform LJ + Coulomb + bonds.
265def energy_sum(+xs: List<&2, P.Atom>, +qs: List<&2, F32>, +bonds: List<&2, U32>, +p: SimParams) -> F32:
266  match p:
267    case SP{+eps, +sig2, +c2, +ke, +kbond}:
268      ((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)
269
270def total_energy(+ds: List<&2, DynAtom>, +qs: List<&2, F32>, +bonds: List<&2, U32>, +p: SimParams) -> F32:
271  energy_sum(dyn_atoms(ds), qs, bonds, p)