1	import Base
2	import ./geom.bend as G
3	import ./protein.bend as P
4	import ./force.bend as F
5	
6	# --- Periodic boundaries v0: orthorhombic box, minimum image. ---
7	# The box rides as a Vec3 of edge lengths (Lx, Ly, Lz). Minimum image per
8	# component: d - L*round(d/L). Vacuum (non-periodic) runs ignore this
9	# module; solvated follow-ups use these kernels. No PME: Coulomb under PBC
10	# would still be cutoff (like everything here), so only LJ ships PBC
11	# variants for now.
12	
13	def mic1d(+d: F32, +len: F32) -> F32:
14	  (d - (len * F32.round((d / len : F32)) : F32) : F32)
15	
16	def mic_vec(+d: G.Vec3, +box: G.Vec3) -> G.Vec3:
17	  match d box:
18	    case G.V3{+dx, +dy, +dz} G.V3{+lx, +ly, +lz}:
19	      G.V3{mic1d(dx, lx), mic1d(dy, ly), mic1d(dz, lz)}
20	
21	def mic_dist2(+a: G.Vec3, +b: G.Vec3, +box: G.Vec3) -> F32:
22	  G.Vec3.norm2(mic_vec(G.Vec3.sub(a, b), box))
23	
24	def lj_pair_pbc(+pi: G.Vec3, +pj: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, +box: G.Vec3) -> F32:
25	  +d2 = mic_dist2(pi, pj, box)
26	  F.lj_pair_go(d2, eps, sig2, F32.is_lt(d2, c2))
27	
28	def lj_head_pbc(+pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, +box: G.Vec3, h: P.Atom) -> F32:
29	  match h:
30	    case P.Atom{serial, elem, +pj}:
31	      lj_pair_pbc(pi, pj, eps, sig2, c2, box)
32	
33	def lj_row_pbc(+pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, +box: G.Vec3, ys: List<&2, P.Atom>) -> F32:
34	  match ys:
35	    case Nil{}:
36	      0.0
37	    case h <> t:
38	      (lj_row_pbc(pi, eps, sig2, c2, box, t) + lj_head_pbc(pi, eps, sig2, c2, box, h) : F32)
39	
40	def lj_self_atom_pbc(h: P.Atom, +ys: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32, +box: G.Vec3) -> F32:
41	  match h:
42	    case P.Atom{serial, elem, +pi}:
43	      lj_row_pbc(pi, eps, sig2, c2, box, ys)
44	
45	def lj_total_pbc(xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32, +box: G.Vec3) -> F32:
46	  match xs:
47	    case Nil{}:
48	      0.0
49	    case h <> t:
50	      (lj_total_pbc(t, ys, eps, sig2, c2, box) + lj_self_atom_pbc(h, ys, eps, sig2, c2, box) : F32)
51	
52	def lj_fat_pbc(+pi: G.Vec3, +pj: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, +box: G.Vec3) -> G.Vec3:
53	  +dv = mic_vec(G.Vec3.sub(pi, pj), box)
54	  +d2 = G.Vec3.norm2(dv)
55	  F.lj_fgo(d2, dv, eps, sig2, F32.is_lt(d2, c2))
56	
57	def lj_fhead_pbc(+pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, +box: G.Vec3, h: P.Atom) -> G.Vec3:
58	  match h:
59	    case P.Atom{serial, elem, +pj}:
60	      lj_fat_pbc(pi, pj, eps, sig2, c2, box)
61	
62	def lj_frow_pbc(+pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, +box: G.Vec3, ys: List<&2, P.Atom>) -> G.Vec3:
63	  match ys:
64	    case Nil{}:
65	      G.V3{0.0, 0.0, 0.0}
66	    case h <> t:
67	      G.Vec3.add(lj_frow_pbc(pi, eps, sig2, c2, box, t), lj_fhead_pbc(pi, eps, sig2, c2, box, h))
68	
69	def lj_fself_atom_pbc(h: P.Atom, +ys: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32, +box: G.Vec3) -> G.Vec3:
70	  match h:
71	    case P.Atom{serial, elem, +pi}:
72	      lj_frow_pbc(pi, eps, sig2, c2, box, ys)
73	
74	def lj_forces_pbc(xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32, +box: G.Vec3) -> List<&2, G.Vec3>:
75	  match xs:
76	    case Nil{}:
77	      Nil{}
78	    case h <> t:
79	      lj_fself_atom_pbc(h, ys, eps, sig2, c2, box) <> lj_forces_pbc(t, ys, eps, sig2, c2, box)