topology.bend

bend-prot@0.1.0.0 · 161 lines · 24 definitions

1import Base
2import bend-prot@0.1.0.0/geom.bend as G
3import bend-prot@0.1.0.0/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
9def hit_nat_go(hit: Bool) -> Nat:
10  match hit:
11    case True{}:
12      1n
13    case False{}:
14      0n
15
16def 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
21def 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.
29def count_within(xs: List<&2, P.Atom>, +qp: G.Vec3, +cutoff: F32) -> Nat:
30  count_below(xs, qp, (cutoff * cutoff : F32))
31
32def 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
37def 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
44def 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
49def 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
56def 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.
62def 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
69def 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
76def 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.
82def rg_about(xs: List<&2, P.Atom>, +qp: G.Vec3) -> F32:
83  rg_of(sumd2_count(xs, qp))
84
85def 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).
92def 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
99def 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
106def 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.
112def 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
119def 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
126def 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
133def 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.
139def 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
146def 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.
154def 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)