1	import Base
2	import ./geom.bend as G
3	import ./protein.bend as P
4	import ./topology.bend as T
5	import ./force.bend as F
6	import ./bonded.bend as B
7	import ./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	
20	type DynAtom is Data:
21	  DA{serial: U32, elem: U32, pos: G.Vec3, vel: G.Vec3, frc: G.Vec3, mass: F32}
22	
23	type SimParams is Data:
24	  SP{eps: F32, sig2: F32, c2: F32, ke: F32, kbond: F32}
25	
26	def acc_conv() -> F32:
27	  0.0004184
28	
29	def boltzmann() -> F32:
30	  0.001987
31	
32	def ke_coul() -> F32:
33	  332.0637
34	
35	def 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	
40	def 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	
47	def 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).
58	def 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.
66	def 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	
69	def 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	
74	def 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	
81	def 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	
90	def 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	
99	def 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	
104	def 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	
111	def 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	
114	def temp_of_ds(ds: List<&2, DynAtom>, n: Nat) -> F32:
115	  temp_of_ke(ke_sim_of(ds), n)
116	
117	def 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	
122	def 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	
129	def 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	
132	def 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	
135	def 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	
138	def 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	
141	def 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).
145	def 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	
158	def 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.
170	def 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).
176	def 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	
183	def 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	
190	def 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	
196	def 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).
204	def 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	
208	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>:
209	  md_thermo(vv_full(set_forces(vv_half(ds, dt), qs, bonds, p), dt), dt, tau, temp0, n)
210	
211	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>:
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.
221	def 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	
227	def 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	
234	def 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	
242	def 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	
251	def 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.
255	def 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.
265	def 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	
270	def 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)