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	
7	# --- Bonded terms + exclusions v0. ---
8	# Harmonic bonds/angles with UNIFORM parameters over explicit serial lists;
9	# 1-2 pairs excluded by bond membership, 1-3 by shared neighbors, all over
10	# flat (Data, reusable) U32 bond lists with (a, b) pairs laid as [a, b, ..].
11	# A flat list is reusable where a pair list (always affine) is not, which is
12	# what lets exclusion tests re-scan bonds inside rows. Missing atoms
13	# contribute 0.0 (documented skip: every lookup is a Maybe pinned by
14	# pos_of_serial laws). Angles need explicit triples; deriving them needs
15	# residue templates (CONECT alone only bonds what it lists). PBC, 1-4
16	# scaling and bonded FORCES are follow-ups.
17	
18	def pos_of_serial_go(xs: List<&2, P.Atom>, +s: U32, best: Maybe<&2, G.Vec3>) -> Maybe<&2, G.Vec3>:
19	  match xs:
20	    case Nil{}:
21	      best
22	    case h <> t:
23	      match h:
24	        case P.Atom{serial, elem, +pos}:
25	          pos_of_serial_go(t, s, Bool.pick(Maybe<&2, G.Vec3>, U32.is_eq(serial, s), Some{pos}, best))
26	
27	def pos_of_serial(xs: List<&2, P.Atom>, s: U32) -> Maybe<&2, G.Vec3>:
28	  pos_of_serial_go(xs, s, None{})
29	
30	# Harmonic bond energy k*(d - (ri+rj))^2 with covalent equilibrium; missing
31	# atoms answer 0.0. Lookups return (position & element) together so the
32	# equilibrium needs no extra pass.
33	def cov_rad(e: U32) -> F32:
34	  match e:
35	    case 1:
36	      0.31
37	    case 6:
38	      0.76
39	    case 7:
40	      0.71
41	    case 8:
42	      0.66
43	    case 15:
44	      1.07
45	    case 16:
46	      1.05
47	    case _:
48	      1.0
49	
50	def atom_of_serial_go(xs: List<&2, P.Atom>, +s: U32, best: Maybe<&1, G.Vec3 & U32>) -> Maybe<&1, G.Vec3 & U32>:
51	  match xs:
52	    case Nil{}:
53	      best
54	    case h <> t:
55	      match h:
56	        case P.Atom{serial, elem, +pos}:
57	          atom_of_serial_go(t, s, Bool.pick(Maybe<&1, G.Vec3 & U32>, U32.is_eq(serial, s), Some{(pos, elem)}, best))
58	
59	def atom_of_serial(xs: List<&2, P.Atom>, s: U32) -> Maybe<&1, G.Vec3 & U32>:
60	  atom_of_serial_go(xs, s, None{})
61	
62	def bond_e(+pa: G.Vec3, +pb: G.Vec3, +k: F32, +r0: F32) -> F32:
63	  (k * F.lj_sq((G.Vec3.dist(pa, pb) - r0 : F32)) : F32)
64	
65	def bond_cov_look1(pa: G.Vec3, ea: U32, mb: Maybe<&1, G.Vec3 & U32>, +k: F32) -> F32:
66	  match mb:
67	    case None{}:
68	      0.0
69	    case Some{(pb, eb)}:
70	      bond_e(pa, pb, k, (cov_rad(ea) + cov_rad(eb) : F32))
71	
72	def bond_cov_look(ma: Maybe<&1, G.Vec3 & U32>, mb: Maybe<&1, G.Vec3 & U32>, +k: F32) -> F32:
73	  match ma:
74	    case None{}:
75	      0.0
76	    case Some{(pa, ea)}:
77	      bond_cov_look1(pa, ea, mb, k)
78	
79	def bond_cov_pair(+xs: List<&2, P.Atom>, s1: U32, s2: U32, +k: F32) -> F32:
80	  bond_cov_look(atom_of_serial(xs, s1), atom_of_serial(xs, s2), k)
81	
82	def bond_total(+xs: List<&2, P.Atom>, bonds: List<&1, U32 & U32>, +k: F32) -> F32:
83	  match bonds:
84	    case Nil{}:
85	      0.0
86	    case (s1, s2) <> t:
87	      (bond_cov_pair(xs, s1, s2, k) + bond_total(xs, t, k) : F32)
88	
89	# Harmonic angle energy k*(theta - eq)^2 about the middle serial.
90	def angle_e(+pa: G.Vec3, +pb: G.Vec3, +pc: G.Vec3, +k: F32, +eq: F32) -> F32:
91	  (k * F.lj_sq((G.Vec3.angle(G.Vec3.sub(pa, pb), G.Vec3.sub(pc, pb)) - eq : F32)) : F32)
92	
93	def angle_e0(pa: G.Vec3, pb: G.Vec3, mc: Maybe<&2, G.Vec3>, +k: F32, +eq: F32) -> F32:
94	  match mc:
95	    case None{}:
96	      0.0
97	    case Some{pc}:
98	      angle_e(pa, pb, pc, k, eq)
99	
100	def angle_e1(pa: G.Vec3, mb: Maybe<&2, G.Vec3>, mc: Maybe<&2, G.Vec3>, +k: F32, +eq: F32) -> F32:
101	  match mb:
102	    case None{}:
103	      0.0
104	    case Some{pb}:
105	      angle_e0(pa, pb, mc, k, eq)
106	
107	def angle_e2(ma: Maybe<&2, G.Vec3>, mb: Maybe<&2, G.Vec3>, mc: Maybe<&2, G.Vec3>, +k: F32, +eq: F32) -> F32:
108	  match ma:
109	    case None{}:
110	      0.0
111	    case Some{pa}:
112	      angle_e1(pa, mb, mc, k, eq)
113	
114	def angle_triple(+xs: List<&2, P.Atom>, sa: U32, sb: U32, sc: U32, +k: F32, +eq: F32) -> F32:
115	  angle_e2(pos_of_serial(xs, sa), pos_of_serial(xs, sb), pos_of_serial(xs, sc), k, eq)
116	
117	def angle_total(+xs: List<&2, P.Atom>, triples: List<&1, U32 & U32 & U32>, +k: F32, +eq: F32) -> F32:
118	  match triples:
119	    case Nil{}:
120	      0.0
121	    case (sa, sb, sc) <> t:
122	      (angle_triple(xs, sa, sb, sc, k, eq) + angle_total(xs, t, k, eq) : F32)
123	
124	# Pair bonds -> flat [a, b, ..] so tests can re-scan (Data reuses).
125	def flatten_bonds(bs: List<&1, U32 & U32>) -> List<&2, U32>:
126	  match bs:
127	    case Nil{}:
128	      Nil{}
129	    case (x, y) <> t:
130	      x <> y <> flatten_bonds(t)
131	
132	def bond_match(+a: U32, +b: U32, +x: U32, +y: U32) -> Bool:
133	  Bool.or(Bool.and(U32.is_eq(a, x), U32.is_eq(b, y)), Bool.and(U32.is_eq(a, y), U32.is_eq(b, x)))
134	
135	def bonded12_flat(bs: List<&2, U32>, +a: U32, +b: U32) -> Bool:
136	  match bs:
137	    case Nil{}:
138	      False{}
139	    case x <> Nil{}:
140	      False{}
141	    case x <> y <> t:
142	      Bool.or(bond_match(a, b, x, y), bonded12_flat(t, a, b))
143	
144	def bond_neighbors_go(bs: List<&2, U32>, +s: U32, +acc: List<&2, U32>) -> List<&2, U32>:
145	  match bs:
146	    case Nil{}:
147	      List.reverse(&2, U32, acc)
148	    case x <> Nil{}:
149	      List.reverse(&2, U32, acc)
150	    case x <> y <> t:
151	      +x2 = x
152	      +y2 = y
153	      bond_neighbors_go(t, s, Bool.pick(List<&2, U32>, U32.is_eq(x2, s), y2 <> acc, Bool.pick(List<&2, U32>, U32.is_eq(y2, s), x2 <> acc, acc)))
154	
155	def bond_neighbors(s: U32, bs: List<&2, U32>) -> List<&2, U32>:
156	  bond_neighbors_go(bs, s, Nil{})
157	
158	def list_member(+x: U32, ys: List<&2, U32>) -> Bool:
159	  match ys:
160	    case Nil{}:
161	      False{}
162	    case h <> t:
163	      Bool.or(U32.is_eq(x, h), list_member(x, t))
164	
165	def share_one(xs: List<&2, U32>, +ys: List<&2, U32>) -> Bool:
166	  match xs:
167	    case Nil{}:
168	      False{}
169	    case h <> t:
170	      Bool.or(list_member(h, ys), share_one(t, ys))
171	
172	def excluded12(a: U32, b: U32, +bonds: List<&2, U32>) -> Bool:
173	  bonded12_flat(bonds, a, b)
174	
175	def excluded13(a: U32, b: U32, +bonds: List<&2, U32>) -> Bool:
176	  share_one(bond_neighbors(a, bonds), bond_neighbors(b, bonds))
177	
178	def excluded(+a: U32, +b: U32, +bonds: List<&2, U32>) -> Bool:
179	  Bool.or(excluded12(a, b, bonds), excluded13(a, b, bonds))
180	
181	# Nonbonded rows/totals that skip 1-2 and 1-3 pairs by serial.
182	def lj_row_excl(+si: U32, +pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, ys: List<&2, P.Atom>, +bonds: List<&2, U32>) -> F32:
183	  match ys:
184	    case Nil{}:
185	      0.0
186	    case h <> t:
187	      match h:
188	        case P.Atom{serial, elem, +pj}:
189	          (Bool.pick(F32, excluded(si, serial, bonds), 0.0, F.lj_pair_at(pi, pj, eps, sig2, c2)) + lj_row_excl(si, pi, eps, sig2, c2, t, bonds) : F32)
190	
191	def lj_total_excl(xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32, +bonds: List<&2, U32>) -> F32:
192	  match xs:
193	    case Nil{}:
194	      0.0
195	    case h <> t:
196	      match h:
197	        case P.Atom{+serial, elem, +pos}:
198	          (lj_total_excl(t, ys, eps, sig2, c2, bonds) + lj_row_excl(serial, pos, eps, sig2, c2, ys, bonds) : F32)
199	
200	def coul_row_excl(+si: U32, +pi: G.Vec3, +qi: F32, +ke: F32, +c2: F32, ys: List<&2, P.Atom>, qs: List<&2, F32>, +bonds: List<&2, U32>) -> F32:
201	  match ys qs:
202	    case Nil{} Nil{}:
203	      0.0
204	    case Nil{} qh <> qt:
205	      0.0
206	    case h <> t Nil{}:
207	      0.0
208	    case h <> t qh <> qt:
209	      +h2 = h
210	      (Bool.pick(F32, excluded(si, P.Atom.serial(h2), bonds), 0.0, F.coul_at(pi, P.Atom.pos(h2), (qi * qh : F32), ke, c2)) + coul_row_excl(si, pi, qi, ke, c2, t, qt, bonds) : F32)
211	
212	def coul_total_excl(xs: List<&2, P.Atom>, +qs_o: List<&2, F32>, +ys: List<&2, P.Atom>, +qs_e: List<&2, F32>, +ke: F32, +c2: F32, +bonds: List<&2, U32>) -> F32:
213	  match xs qs_o:
214	    case Nil{} Nil{}:
215	      0.0
216	    case Nil{} qh <> qt:
217	      0.0
218	    case h <> t Nil{}:
219	      0.0
220	    case h <> t qh <> qt:
221	      +h2 = h
222	      (coul_total_excl(t, qt, ys, qs_e, ke, c2, bonds) + coul_row_excl(P.Atom.serial(h2), P.Atom.pos(h2), qh, ke, c2, ys, qs_e, bonds) : F32)
223	
224	# Per-element rows/totals that also skip 1-2 and 1-3 pairs: the physically
225	# meaningful combination (uniform rows blow up on bonded neighbors).
226	def lj_row_elem_excl(+si: U32, +se: U32, +pi: G.Vec3, +ys: List<&2, P.Atom>, +c2: F32, +bonds: List<&2, U32>) -> F32:
227	  match ys:
228	    case Nil{}:
229	      0.0
230	    case h <> t:
231	      match h:
232	        case P.Atom{serial, +elem, +pj}:
233	          (Bool.pick(F32, excluded(si, serial, bonds), 0.0, F.lj_pair_at(pi, pj, F.lj_combine_eps(F.lj_eps(se), F.lj_eps(elem)), F.lj_combine_sig2(F.lj_sig(se), F.lj_sig(elem)), c2)) + lj_row_elem_excl(si, se, pi, t, c2, bonds) : F32)
234	
235	def lj_total_elem_excl(xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>, +c2: F32, +bonds: List<&2, U32>) -> F32:
236	  match xs:
237	    case Nil{}:
238	      0.0
239	    case h <> t:
240	      match h:
241	        case P.Atom{serial, elem, pos}:
242	          (lj_total_elem_excl(t, ys, c2, bonds) + lj_row_elem_excl(serial, elem, pos, ys, c2, bonds) : F32)
243	
244	# Excluded LJ forces mirror lj_forces, skipping 1-2/1-3 pairs by serial.
245	def lj_frow_excl(+si: U32, +pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, ys: List<&2, P.Atom>, +bonds: List<&2, U32>) -> G.Vec3:
246	  match ys:
247	    case Nil{}:
248	      G.V3{0.0, 0.0, 0.0}
249	    case h <> t:
250	      match h:
251	        case P.Atom{serial, elem, +pj}:
252	          G.Vec3.add(lj_frow_excl(si, pi, eps, sig2, c2, t, bonds), Bool.pick(G.Vec3, excluded(si, serial, bonds), G.V3{0.0, 0.0, 0.0}, F.lj_fat(pi, pj, eps, sig2, c2)))
253	
254	def lj_fself_atom_excl(h: P.Atom, si: U32, +ys: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32, +bonds: List<&2, U32>) -> G.Vec3:
255	  match h:
256	    case P.Atom{serial, elem, +pi}:
257	      lj_frow_excl(serial, pi, eps, sig2, c2, ys, bonds)
258	
259	def lj_forces_excl(xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32, +bonds: List<&2, U32>) -> List<&2, G.Vec3>:
260	  match xs:
261	    case Nil{}:
262	      Nil{}
263	    case +h <> t:
264	      lj_fself_atom_excl(h, P.Atom.serial(h), ys, eps, sig2, c2, bonds) <> lj_forces_excl(t, ys, eps, sig2, c2, bonds)
265	
266	# Bond forces: F = -2k(d-r0) * unit(pa-pb) on a, opposite on b; zero for
267	# coincident or missing atoms. Per-atom loop over flat bonds.
268	def bond_fpair(+d: G.Vec3, +d2: F32, +k: F32, +r0: F32) -> G.Vec3 & G.Vec3:
269	  +r = F32.sqrt(d2)
270	  +mag = Bool.pick(F32, F32.is_eq(d2, 0.0), 0.0, (((0.0 - (2.0 * k : F32) : F32) * (r - r0 : F32) : F32) / r : F32))
271	  +fa = G.Vec3.scale(mag, d)
272	  (fa, G.Vec3.neg(fa))
273	
274	def bond_fpick(eqa: Bool, eqb: Bool, fab: G.Vec3 & G.Vec3) -> G.Vec3:
275	  match fab:
276	    case (fa, fb):
277	      Bool.pick(G.Vec3, eqa, fa, Bool.pick(G.Vec3, eqb, fb, G.V3{0.0, 0.0, 0.0}))
278	
279	def bond_fsas(+s: U32, sa: U32, sb: U32, +pa: G.Vec3, +pb: G.Vec3, +k: F32, +r0: F32) -> G.Vec3:
280	  bond_fpick(U32.is_eq(s, sa), U32.is_eq(s, sb), bond_fpair(G.Vec3.sub(pa, pb), G.Vec3.dist2(pa, pb), k, r0))
281	
282	def bond_force_look1(s: U32, pa: G.Vec3, ea: U32, mb: Maybe<&1, G.Vec3 & U32>, a: U32, b: U32, +k: F32) -> G.Vec3:
283	  match mb:
284	    case None{}:
285	      G.V3{0.0, 0.0, 0.0}
286	    case Some{(pb, eb)}:
287	      bond_fsas(s, a, b, pa, pb, k, (cov_rad(ea) + cov_rad(eb) : F32))
288	
289	def bond_force_look(s: U32, ma: Maybe<&1, G.Vec3 & U32>, mb: Maybe<&1, G.Vec3 & U32>, a: U32, b: U32, +k: F32) -> G.Vec3:
290	  match ma:
291	    case None{}:
292	      G.V3{0.0, 0.0, 0.0}
293	    case Some{(pa, ea)}:
294	      bond_force_look1(s, pa, ea, mb, a, b, k)
295	
296	def bond_force_bond(+xs: List<&2, P.Atom>, s: U32, +a: U32, +b: U32, +k: F32) -> G.Vec3:
297	  bond_force_look(s, atom_of_serial(xs, a), atom_of_serial(xs, b), a, b, k)
298	
299	def bond_force_list(+s: U32, bs: List<&2, U32>, +xs: List<&2, P.Atom>, +k: F32) -> G.Vec3:
300	  match bs:
301	    case Nil{}:
302	      G.V3{0.0, 0.0, 0.0}
303	    case x <> Nil{}:
304	      G.V3{0.0, 0.0, 0.0}
305	    case x <> y <> t:
306	      G.Vec3.add(bond_force_bond(xs, s, x, y, k), bond_force_list(s, t, xs, k))
307	
308	def bond_forces(+xs: List<&2, P.Atom>, +bonds: List<&2, U32>, +k: F32) -> List<&2, G.Vec3>:
309	  match xs:
310	    case Nil{}:
311	      Nil{}
312	    case +h <> t:
313	      bond_force_list(P.Atom.serial(h), bonds, xs, k) <> bond_forces(t, bonds, k)
314	
315	def bond_cut2(ei: U32, ej: U32) -> F32:
316	  +s = (cov_rad(ei) + cov_rad(ej) : F32)
317	  +t = (s * 1.25 : F32)
318	  (t * t : F32)
319	
320	def infer_one(si: U32, pi: G.Vec3, ei: U32, sj: U32, ej: U32, pj: G.Vec3) -> List<&1, U32 & U32>:
321	  Bool.pick(List<&1, U32 & U32>, F32.is_lt(G.Vec3.dist2(pi, pj), bond_cut2(ei, ej)), [(si, sj)], Nil{})
322	
323	def append_infer(a: List<&1, U32 & U32>, b: List<&1, U32 & U32>) -> List<&1, U32 & U32>:
324	  match a:
325	    case Nil{}:
326	      b
327	    case h <> t:
328	      h <> append_infer(t, b)
329	
330	def infer_row(+si: U32, +pi: G.Vec3, +ei: U32, ys: List<&2, P.Atom>) -> List<&1, U32 & U32>:
331	  match ys:
332	    case Nil{}:
333	      Nil{}
334	    case +h <> t:
335	      append_infer(infer_one(si, pi, ei, P.Atom.serial(h), P.Atom.elem(h), P.Atom.pos(h)), infer_row(si, pi, ei, t))
336	
337	def infer_total(+xs: List<&2, P.Atom>) -> List<&1, U32 & U32>:
338	  match xs:
339	    case Nil{}:
340	      Nil{}
341	    case +h <> t:
342	      append_infer(infer_row(P.Atom.serial(h), P.Atom.pos(h), P.Atom.elem(h), t), infer_total(t))
343	
344	# Angle auto-derivation from flat bonds: for each bond, fan over the other
345	# neighbors of each end. Needs no templates; duplicates across overlapping
346	# fans are possible in rings (documented, harmless for boolean use).
347	def bond_neighbors_except_go(bs: List<&2, U32>, +s: U32, +excl: U32, +acc: List<&2, U32>) -> List<&2, U32>:
348	  match bs:
349	    case Nil{}:
350	      List.reverse(&2, U32, acc)
351	    case x <> Nil{}:
352	      List.reverse(&2, U32, acc)
353	    case x <> y <> t:
354	      +x2 = x
355	      +y2 = y
356	      bond_neighbors_except_go(t, s, excl, Bool.pick(List<&2, U32>, U32.is_eq(x2, s), Bool.pick(List<&2, U32>, U32.is_eq(y2, excl), acc, y2 <> acc), Bool.pick(List<&2, U32>, U32.is_eq(y2, s), Bool.pick(List<&2, U32>, U32.is_eq(x2, excl), acc, x2 <> acc), acc)))
357	
358	def bond_neighbors_except(s: U32, excl: U32, bs: List<&2, U32>) -> List<&2, U32>:
359	  bond_neighbors_except_go(bs, s, excl, Nil{})
360	
361	def fan_triples(+x: U32, +y: U32, ns: List<&2, U32>) -> List<&1, U32 & U32 & U32>:
362	  match ns:
363	    case Nil{}:
364	      Nil{}
365	    case c <> t:
366	      (x, y, c) <> fan_triples(x, y, t)
367	
368	def append_triples(a: List<&1, U32 & U32 & U32>, b: List<&1, U32 & U32 & U32>) -> List<&1, U32 & U32 & U32>:
369	  match a:
370	    case Nil{}:
371	      b
372	    case h <> t:
373	      h <> append_triples(t, b)
374	
375	def angle_fan_both(+a: U32, +b: U32, +flat: List<&2, U32>) -> List<&1, U32 & U32 & U32>:
376	  append_triples(fan_triples(a, b, bond_neighbors_except(b, a, flat)), fan_triples(b, a, bond_neighbors_except(a, b, flat)))
377	
378	def angle_triples_flat_go(bs: List<&2, U32>, +flat: List<&2, U32>) -> List<&1, U32 & U32 & U32>:
379	  match bs:
380	    case Nil{}:
381	      Nil{}
382	    case x <> Nil{}:
383	      Nil{}
384	    case x <> y <> t:
385	      append_triples(angle_fan_both(x, y, flat), angle_triples_flat_go(t, flat))
386	
387	def angle_triples_auto(+flat: List<&2, U32>) -> List<&1, U32 & U32 & U32>:
388	  angle_triples_flat_go(flat, flat)