1	import Base
2	import ./geom.bend as G
3	import ./protein.bend as P
4	
5	# --- Topology v0: contacts, clashes, spread. Cutoffs are SQUARED (no sqrt). ---
6	# Parallel note: each pair test is independent; a future pass can evaluate
7	# `count_below` / `contacts_between` with balanced fork-join + GPU (`!`).
8	
9	def hit_nat_go(hit: Bool) -> Nat:
10	  match hit:
11	    case True{}:
12	      1n
13	    case False{}:
14	      0n
15	
16	def head_hit_nat(h: P.Atom, +qp: G.Vec3, +c2: F32) -> Nat:
17	  match h:
18	    case P.Atom{serial, elem, +pos}:
19	      hit_nat_go(F32.is_lt(G.Vec3.dist2(pos, qp), c2))
20	
21	def count_below(xs: List<&2, P.Atom>, +qp: G.Vec3, +c2: F32) -> Nat:
22	  match xs:
23	    case Nil{}:
24	      0n
25	    case h <> t:
26	      Nat.add(count_below(t, qp, c2), head_hit_nat(h, qp, c2))
27	
28	# Plain (non-squared) cutoff: squares it once, then counts.
29	def count_within(xs: List<&2, P.Atom>, +qp: G.Vec3, +cutoff: F32) -> Nat:
30	  count_below(xs, qp, (cutoff * cutoff : F32))
31	
32	def head_is_close(h: P.Atom, +qp: G.Vec3, +min2: F32) -> Bool:
33	  match h:
34	    case P.Atom{serial, elem, +pos}:
35	      F32.is_lt(G.Vec3.dist2(pos, qp), min2)
36	
37	def has_close(xs: List<&2, P.Atom>, +qp: G.Vec3, +min2: F32) -> Bool:
38	  match xs:
39	    case Nil{}:
40	      False{}
41	    case h <> t:
42	      Bool.or(has_close(t, qp, min2), head_is_close(h, qp, min2))
43	
44	def head_dist2(h: P.Atom, +qp: G.Vec3) -> F32:
45	  match h:
46	    case P.Atom{serial, elem, +pos}:
47	      G.Vec3.dist2(pos, qp)
48	
49	def sum_dist2(xs: List<&2, P.Atom>, +qp: G.Vec3) -> F32:
50	  match xs:
51	    case Nil{}:
52	      0.0
53	    case h <> t:
54	      (sum_dist2(t, qp) + head_dist2(h, qp) : F32)
55	
56	def sumd2_extend(a: F32, rest: F32 & Nat) -> F32 & Nat:
57	  match rest:
58	    case (s, n):
59	      ((a + s : F32), 1n+n)
60	
61	# One pass: summed squared distances about qp, plus atom count.
62	def sumd2_count(xs: List<&2, P.Atom>, +qp: G.Vec3) -> F32 & Nat:
63	  match xs:
64	    case Nil{}:
65	      (0.0, 0n)
66	    case h <> t:
67	      sumd2_extend(head_dist2(h, qp), sumd2_count(t, qp))
68	
69	def rg_of_go(s: F32, n: Nat) -> F32:
70	  match n:
71	    case 0n:
72	      0.0
73	    case 1n+p:
74	      F32.sqrt((s / (F32.from_nat(n)) : F32))
75	
76	def rg_of(pair: F32 & Nat) -> F32:
77	  match pair:
78	    case (s, n):
79	      rg_of_go(s, n)
80	
81	# Radius of gyration of an atom list about a given center.
82	def rg_about(xs: List<&2, P.Atom>, +qp: G.Vec3) -> F32:
83	  rg_of(sumd2_count(xs, qp))
84	
85	def contacts_of_head(h: P.Atom, +ys: List<&2, P.Atom>, +c2: F32) -> Nat:
86	  match h:
87	    case P.Atom{serial, elem, +pos}:
88	      count_below(ys, pos, c2)
89	
90	# Interface size: pairs (x in xs, y in ys) with dist2 < c2. Consumes xs once;
91	# ys is reusable so every head reuses it (reference counted).
92	def contacts_between(xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>, +c2: F32) -> Nat:
93	  match xs:
94	    case Nil{}:
95	      0n
96	    case h <> t:
97	      Nat.add(contacts_between(t, ys, c2), contacts_of_head(h, ys, c2))
98	
99	def min_cons_go(d: F32, rest: Maybe<&2, F32>) -> Maybe<&2, F32>:
100	  match rest:
101	    case None{}:
102	      Some{d}
103	    case Some{best}:
104	      Some{F32.min(d, best)}
105	
106	def min_cons(h: P.Atom, rest: Maybe<&2, F32>, +qp: G.Vec3) -> Maybe<&2, F32>:
107	  match h:
108	    case P.Atom{serial, elem, +pos}:
109	      min_cons_go(G.Vec3.dist2(pos, qp), rest)
110	
111	# Closest squared distance from qp to any atom in xs. None{} when empty.
112	def min_dist2_to(xs: List<&2, P.Atom>, +qp: G.Vec3) -> Maybe<&2, F32>:
113	  match xs:
114	    case Nil{}:
115	      None{}
116	    case h <> t:
117	      min_cons(h, min_dist2_to(t, qp), qp)
118	
119	def min_maybe_some(x: F32, b: Maybe<&2, F32>) -> Maybe<&2, F32>:
120	  match b:
121	    case None{}:
122	      Some{x}
123	    case Some{y}:
124	      Some{F32.min(x, y)}
125	
126	def min_maybe(a: Maybe<&2, F32>, b: Maybe<&2, F32>) -> Maybe<&2, F32>:
127	  match a:
128	    case None{}:
129	      b
130	    case Some{x}:
131	      min_maybe_some(x, b)
132	
133	def head_link(h: P.Atom, +ys: List<&2, P.Atom>) -> Maybe<&2, F32>:
134	  match h:
135	    case P.Atom{serial, elem, +pos}:
136	      min_dist2_to(ys, pos)
137	
138	# Closest squared approach between two atom sets. None{} when either is empty.
139	def min_link2(xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>) -> Maybe<&2, F32>:
140	  match xs:
141	    case Nil{}:
142	      None{}
143	    case h <> t:
144	      min_maybe(min_link2(t, ys), head_link(h, ys))
145	
146	def min_link(link: Maybe<&2, F32>) -> Maybe<&2, F32>:
147	  match link:
148	    case None{}:
149	      None{}
150	    case Some{d2}:
151	      Some{F32.sqrt(d2)}
152	
153	# Per-atom neighbor counts: for each x in xs, atoms of ys within dist2 < c2.
154	def coord_counts(
155	  xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>, +c2: F32
156	) -> List<&2, Nat>:
157	  match xs:
158	    case Nil{}:
159	      Nil{}
160	    case h <> t:
161	      contacts_of_head(h, ys, c2) <> coord_counts(t, ys, c2)