1	import Base
2	import ./geom.bend as G
3	import ./protein.bend as P
4	
5	# --- Nonbonded kernels v0: uniform Lennard-Jones + Coulomb, squared cutoffs.
6	# Self pairs (d2 == 0) and pairs beyond c2 contribute nothing.
7	# Convention: for self energy/forces, pass the same reusable list twice,
8	# e.g. lj_total(+xs, xs, eps, sig2, c2) style call sites. That call is a
9	# DOUBLE sum: every unordered pair {i, j} is counted twice (i,j and j,i),
10	# so the physical total is half the answer. The diagonal is dropped by
11	# POSITION (d2 == 0.0), not by index: callers must give distinct atoms
12	# distinct positions, since a coincident distinct pair is silently dropped
13	# instead of diverging.
14	
15	def lj_sq(+x: F32) -> F32:
16	  (x * x : F32)
17	
18	def lj_cube(+inv: F32) -> F32:
19	  ((inv * inv : F32) * inv : F32)
20	
21	def lj_eval(+u6: F32, +eps: F32) -> F32:
22	  ((4.0 * eps : F32) * ((lj_sq(u6) - u6 : F32)) : F32)
23	
24	def lj_from_inv(inv: F32, +eps: F32) -> F32:
25	  lj_eval(lj_sq(lj_cube(inv)), eps)
26	
27	# LJ energy from squared distance: 4e*((s2/d2)^12 - (s2/d2)^6), no sqrt.
28	def lj_of_d2(d2: F32, +eps: F32, +sig2: F32) -> F32:
29	  lj_from_inv((sig2 / d2 : F32), eps)
30	
31	def lj_self_go(d2: F32, +eps: F32, +sig2: F32, self: Bool) -> F32:
32	  match self:
33	    case True{}:
34	      0.0
35	    case False{}:
36	      lj_of_d2(d2, eps, sig2)
37	
38	def lj_pair_go(+d2: F32, +eps: F32, +sig2: F32, below: Bool) -> F32:
39	  match below:
40	    case True{}:
41	      lj_self_go(d2, eps, sig2, F32.is_eq(d2, 0.0))
42	    case False{}:
43	      0.0
44	
45	def lj_pair_at(+pi: G.Vec3, +pj: G.Vec3, +eps: F32, +sig2: F32, +c2: F32) -> F32:
46	  +d2 = G.Vec3.dist2(pi, pj)
47	  lj_pair_go(d2, eps, sig2, F32.is_lt(d2, c2))
48	
49	def lj_head(+pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, h: P.Atom) -> F32:
50	  match h:
51	    case P.Atom{serial, elem, +pj}:
52	      lj_pair_at(pi, pj, eps, sig2, c2)
53	
54	def lj_row(
55	  +pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, ys: List<&2, P.Atom>
56	) -> F32:
57	  match ys:
58	    case Nil{}:
59	      0.0
60	    case h <> t:
61	      (lj_row(pi, eps, sig2, c2, t) + lj_head(pi, eps, sig2, c2, h) : F32)
62	
63	def lj_self_atom(
64	  h: P.Atom, +ys: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32
65	) -> F32:
66	  match h:
67	    case P.Atom{serial, elem, +pi}:
68	      lj_row(pi, eps, sig2, c2, ys)
69	
70	def lj_total(
71	  xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32
72	) -> F32:
73	  match xs:
74	    case Nil{}:
75	      0.0
76	    case h <> t:
77	      (lj_total(t, ys, eps, sig2, c2) + lj_self_atom(h, ys, eps, sig2, c2) : F32)
78	
79	# --- Coulomb with per-atom charges in an aligned List<&2, F32>. ---
80	# The charge lists zip with the atom lists: a shorter charge list silently
81	# drops the remaining atoms (the Nil{} arms answer zero), so callers must
82	# keep charges aligned with atoms; see LAWS coul_trunc_charges.
83	
84	def coul_formula(d2: F32, +qq: F32, +ke: F32) -> F32:
85	  ((ke * qq : F32) / F32.sqrt(d2) : F32)
86	
87	def coul_self_go(d2: F32, +qq: F32, +ke: F32, self: Bool) -> F32:
88	  match self:
89	    case True{}:
90	      0.0
91	    case False{}:
92	      coul_formula(d2, qq, ke)
93	
94	def coul_pair_go(+d2: F32, +qq: F32, +ke: F32, below: Bool) -> F32:
95	  match below:
96	    case True{}:
97	      coul_self_go(d2, qq, ke, F32.is_eq(d2, 0.0))
98	    case False{}:
99	      0.0
100	
101	def coul_at(+pi: G.Vec3, +pj: G.Vec3, +qq: F32, +ke: F32, +c2: F32) -> F32:
102	  +d2 = G.Vec3.dist2(pi, pj)
103	  coul_pair_go(d2, qq, ke, F32.is_lt(d2, c2))
104	
105	def coul_head(
106	  +pi: G.Vec3, +qi: F32, +ke: F32, +c2: F32, h: P.Atom, qj: F32
107	) -> F32:
108	  match h:
109	    case P.Atom{serial, elem, +pj}:
110	      coul_at(pi, pj, (qi * qj : F32), ke, c2)
111	
112	def coul_row(
113	  +pi: G.Vec3, +qi: F32, +ke: F32, +c2: F32,
114	  ys: List<&2, P.Atom>, qs: List<&2, F32>
115	) -> F32:
116	  match ys qs:
117	    case Nil{} Nil{}:
118	      0.0
119	    case Nil{} qh <> qt:
120	      0.0
121	    case h <> t Nil{}:
122	      0.0
123	    case h <> t qh <> qt:
124	      (coul_row(pi, qi, ke, c2, t, qt) + coul_head(pi, qi, ke, c2, h, qh) : F32)
125	
126	def coul_self(
127	  h: P.Atom, +qh: F32, +ys: List<&2, P.Atom>, +qs_e: List<&2, F32>,
128	  +ke: F32, +c2: F32
129	) -> F32:
130	  match h:
131	    case P.Atom{serial, elem, +pi}:
132	      coul_row(pi, qh, ke, c2, ys, qs_e)
133	
134	def coul_total(
135	  xs: List<&2, P.Atom>, +qs_o: List<&2, F32>, +ys: List<&2, P.Atom>,
136	  +qs_e: List<&2, F32>, +ke: F32, +c2: F32
137	) -> F32:
138	  match xs qs_o:
139	    case Nil{} Nil{}:
140	      0.0
141	    case Nil{} qh <> qt:
142	      0.0
143	    case h <> t Nil{}:
144	      0.0
145	    case h <> t qh <> qt:
146	      (coul_total(t, qt, ys, qs_e, ke, c2) + coul_self(h, qh, ys, qs_e, ke, c2) : F32)
147	
148	# --- LJ forces: F_i = 24e*(2*u6^2 - u6)/d2 * (pi - pj), u6 = (sig2/d2)^3^2. ---
149	
150	def lj_k(+u6: F32, +eps: F32, +d2: F32) -> F32:
151	  (((24.0 * eps : F32) * (((2.0 * lj_sq(u6) : F32) - u6 : F32)) : F32) / d2 : F32)
152	
153	def lj_fvec(+d: G.Vec3, +k: F32) -> G.Vec3:
154	  G.Vec3.scale(k, d)
155	
156	def lj_fpair(+d2: F32, +d: G.Vec3, +eps: F32, +sig2: F32) -> G.Vec3:
157	  lj_fvec(d, lj_k(lj_sq(lj_cube((sig2 / d2 : F32))), eps, d2))
158	
159	def lj_fself(+d2: F32, +d: G.Vec3, +eps: F32, +sig2: F32, self: Bool) -> G.Vec3:
160	  match self:
161	    case True{}:
162	      G.V3{0.0, 0.0, 0.0}
163	    case False{}:
164	      lj_fpair(d2, d, eps, sig2)
165	
166	def lj_fgo(+d2: F32, +d: G.Vec3, +eps: F32, +sig2: F32, below: Bool) -> G.Vec3:
167	  match below:
168	    case True{}:
169	      lj_fself(d2, d, eps, sig2, F32.is_eq(d2, 0.0))
170	    case False{}:
171	      G.V3{0.0, 0.0, 0.0}
172	
173	def lj_fat(+pi: G.Vec3, +pj: G.Vec3, +eps: F32, +sig2: F32, +c2: F32) -> G.Vec3:
174	  +d2 = G.Vec3.dist2(pi, pj)
175	  lj_fgo(d2, G.Vec3.sub(pi, pj), eps, sig2, F32.is_lt(d2, c2))
176	
177	def lj_fhead(+pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, h: P.Atom) -> G.Vec3:
178	  match h:
179	    case P.Atom{serial, elem, +pj}:
180	      lj_fat(pi, pj, eps, sig2, c2)
181	
182	def lj_frow(
183	  +pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, ys: List<&2, P.Atom>
184	) -> G.Vec3:
185	  match ys:
186	    case Nil{}:
187	      G.V3{0.0, 0.0, 0.0}
188	    case h <> t:
189	      G.Vec3.add(lj_frow(pi, eps, sig2, c2, t), lj_fhead(pi, eps, sig2, c2, h))
190	
191	def lj_fself_atom(
192	  h: P.Atom, +ys: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32
193	) -> G.Vec3:
194	  match h:
195	    case P.Atom{serial, elem, +pi}:
196	      lj_frow(pi, eps, sig2, c2, ys)
197	
198	def lj_forces(
199	  xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32
200	) -> List<&2, G.Vec3>:
201	  match xs:
202	    case Nil{}:
203	      Nil{}
204	    case h <> t:
205	      lj_fself_atom(h, ys, eps, sig2, c2) <> lj_forces(t, ys, eps, sig2, c2)
206	
207	# --- Steepest descent: x += F * dt, fuel-bounded. Equal-length xs/fs. ---
208	
209	def sd_move(h: P.Atom, f: G.Vec3, +dt: F32) -> P.Atom:
210	  match h f:
211	    case P.Atom{serial, elem, +pos} G.V3{+fx, +fy, +fz}:
212	      P.Atom{serial, elem, G.Vec3.add(pos, G.Vec3.scale(dt, G.V3{fx, fy, fz}))}
213	
214	def sd_sweep(
215	  xs: List<&2, P.Atom>, fs: List<&2, G.Vec3>, +dt: F32
216	) -> List<&2, P.Atom>:
217	  match xs fs:
218	    case Nil{} Nil{}:
219	      Nil{}
220	    case Nil{} f <> ft:
221	      Nil{}
222	    case h <> t Nil{}:
223	      h <> t
224	    case h <> t f <> ft:
225	      sd_move(h, f, dt) <> sd_sweep(t, ft, dt)
226	
227	def minimize(
228	  fuel: Nat, +xs: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32, +dt: F32
229	) -> List<&2, P.Atom>:
230	  match fuel:
231	    case 0n:
232	      xs
233	    case 1n+f:
234	      minimize(f, sd_sweep(xs, lj_forces(xs, xs, eps, sig2, c2), dt), eps, sig2, c2, dt)
235	
236	# --- Per-element Lennard-Jones: illustrative eps (kcal/mol) and sigma (A)
237	# by element number; unknown elements take a generic (0.20, 3.00). These
238	# are placeholders, not a force field. Lorentz-Berthelot combining:
239	# eps = sqrt(ei*ej), sig2 = ((si+sj)/2)^2.
240	def lj_eps(e: U32) -> F32:
241	  match e:
242	    case 1:
243	      0.03
244	    case 6:
245	      0.066
246	    case 7:
247	      0.17
248	    case 8:
249	      0.21
250	    case 16:
251	      0.25
252	    case _:
253	      0.2
254	
255	def lj_sig(e: U32) -> F32:
256	  match e:
257	    case 1:
258	      2.5
259	    case 6:
260	      3.5
261	    case 7:
262	      3.25
263	    case 8:
264	      2.96
265	    case 16:
266	      3.55
267	    case _:
268	      3.0
269	
270	def lj_combine_eps(ei: F32, ej: F32) -> F32:
271	  F32.sqrt((ei * ej : F32))
272	
273	def lj_combine_sig2(si: F32, sj: F32) -> F32:
274	  lj_sq(((si + sj : F32) / 2.0 : F32))
275	
276	def lj_head_elem(+si: U32, +pi: G.Vec3, +c2: F32, h: P.Atom) -> F32:
277	  match h:
278	    case P.Atom{serial, +elem, +pj}:
279	      lj_pair_at(pi, pj, lj_combine_eps(lj_eps(si), lj_eps(elem)), lj_combine_sig2(lj_sig(si), lj_sig(elem)), c2)
280	
281	def lj_row_elem(+si: U32, +pi: G.Vec3, +ys: List<&2, P.Atom>, +c2: F32) -> F32:
282	  match ys:
283	    case Nil{}:
284	      0.0
285	    case h <> t:
286	      (lj_row_elem(si, pi, t, c2) + lj_head_elem(si, pi, c2, h) : F32)
287	
288	def lj_total_elem(xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>, +c2: F32) -> F32:
289	  match xs:
290	    case Nil{}:
291	      0.0
292	    case +h <> t:
293	      (lj_total_elem(t, ys, c2) + lj_row_elem(P.Atom.elem(h), P.Atom.pos(h), ys, c2) : F32)
294	
295	# --- Coulomb forces: F_i = ke*qi*qj/d2^1.5 * (pi - pj), repulsive for
296	# like charges along +d. Self pairs (d2 == 0) and pairs beyond c2 give zero.
297	def coul_k(+qq: F32, +ke: F32, +d2: F32) -> F32:
298	  ((ke * qq : F32) / ((d2 * F32.sqrt(d2) : F32)) : F32)
299	
300	def coul_fpair(+d2: F32, +d: G.Vec3, +qq: F32, +ke: F32) -> G.Vec3:
301	  G.Vec3.scale(coul_k(qq, ke, d2), d)
302	
303	def coul_fself(+d2: F32, +d: G.Vec3, +qq: F32, +ke: F32, self: Bool) -> G.Vec3:
304	  match self:
305	    case True{}:
306	      G.V3{0.0, 0.0, 0.0}
307	    case False{}:
308	      coul_fpair(d2, d, qq, ke)
309	
310	def coul_fgo(+d2: F32, +d: G.Vec3, +qq: F32, +ke: F32, below: Bool) -> G.Vec3:
311	  match below:
312	    case True{}:
313	      coul_fself(d2, d, qq, ke, F32.is_eq(d2, 0.0))
314	    case False{}:
315	      G.V3{0.0, 0.0, 0.0}
316	
317	def coul_fat(+pi: G.Vec3, +pj: G.Vec3, +qq: F32, +ke: F32, +c2: F32) -> G.Vec3:
318	  +d2 = G.Vec3.dist2(pi, pj)
319	  coul_fgo(d2, G.Vec3.sub(pi, pj), qq, ke, F32.is_lt(d2, c2))
320	
321	def coul_fhead(+pi: G.Vec3, +qi: F32, +ke: F32, +c2: F32, h: P.Atom, qj: F32) -> G.Vec3:
322	  match h:
323	    case P.Atom{serial, elem, +pj}:
324	      coul_fat(pi, pj, (qi * qj : F32), ke, c2)
325	
326	def coul_frow(+pi: G.Vec3, +qi: F32, +ke: F32, +c2: F32, ys: List<&2, P.Atom>, qs: List<&2, F32>) -> G.Vec3:
327	  match ys qs:
328	    case Nil{} Nil{}:
329	      G.V3{0.0, 0.0, 0.0}
330	    case Nil{} qh <> qt:
331	      G.V3{0.0, 0.0, 0.0}
332	    case h <> t Nil{}:
333	      G.V3{0.0, 0.0, 0.0}
334	    case h <> t qh <> qt:
335	      G.Vec3.add(coul_frow(pi, qi, ke, c2, t, qt), coul_fhead(pi, qi, ke, c2, h, qh))
336	
337	def coul_fself_atom(h: P.Atom, +qh: F32, +ys: List<&2, P.Atom>, +qs_e: List<&2, F32>, +ke: F32, +c2: F32) -> G.Vec3:
338	  match h:
339	    case P.Atom{serial, elem, +pi}:
340	      coul_frow(pi, qh, ke, c2, ys, qs_e)
341	
342	def coul_forces(xs: List<&2, P.Atom>, +qs_o: List<&2, F32>, +ys: List<&2, P.Atom>, +qs_e: List<&2, F32>, +ke: F32, +c2: F32) -> List<&2, G.Vec3>:
343	  match xs qs_o:
344	    case Nil{} Nil{}:
345	      Nil{}
346	    case Nil{} qh <> qt:
347	      Nil{}
348	    case h <> t Nil{}:
349	      Nil{}
350	    case h <> t qh <> qt:
351	      coul_fself_atom(h, qh, ys, qs_e, ke, c2) <> coul_forces(t, qt, ys, qs_e, ke, c2)
352	
353	# --- Atomic masses (amu) by element; unknown -> 12.0. ---
354	def mass(e: U32) -> F32:
355	  match e:
356	    case 1:
357	      1.008
358	    case 6:
359	      12.011
360	    case 7:
361	      14.007
362	    case 8:
363	      15.999
364	    case 15:
365	      30.974
366	    case 16:
367	      32.06
368	    case _:
369	      12.0