protein.bend

bend-prot@0.1.0.0 · 202 lines · 36 definitions

1import Base
2import bend-prot@0.1.0.0/geom.bend as G
3
4# --- Protein model v0: Atom / Residue / Chain / Complex. All Data. ---
5
6type Atom is Data:
7  Atom{serial: U32, elem: U32, pos: G.Vec3}
8
9type Residue is Data:
10  Residue{seq: Nat, atoms: List<&2, Atom>}
11
12type Chain is Data:
13  Chain{cid: U32, residues: List<&2, Residue>}
14
15type Complex is Data:
16  Complex{chains: List<&2, Chain>}
17
18# Atom position (drops serial/elem, keeps pos)
19def Atom.pos(a: Atom) -> G.Vec3:
20  match a:
21    case Atom{serial, elem, pos}:
22      pos
23
24# Atom element number (drops serial/pos, keeps elem)
25def Atom.elem(a: Atom) -> U32:
26  match a:
27    case Atom{serial, elem, pos}:
28      elem
29
30# Atom serial number (drops elem/pos, keeps serial)
31def Atom.serial(a: Atom) -> U32:
32  match a:
33    case Atom{serial, elem, pos}:
34      serial
35
36def Atom.dist2(+a: Atom, +b: Atom) -> F32:
37  match a b:
38    case Atom{serial1, elem1, +pa} Atom{serial2, elem2, +pb}:
39      G.Vec3.dist2(pa, pb)
40
41def Atom.dist(+a: Atom, +b: Atom) -> F32:
42  F32.sqrt(Atom.dist2(a, b))
43
44def count_atoms(xs: List<&2, Atom>) -> Nat:
45  match xs:
46    case Nil{}:
47      0n
48    case h <> t:
49      1n+count_atoms(t)
50
51def count_residues(xs: List<&2, Residue>) -> Nat:
52  match xs:
53    case Nil{}:
54      0n
55    case h <> t:
56      1n+count_residues(t)
57
58def sum_count_extend_go(pos: G.Vec3, rest: G.Vec3 & Nat) -> G.Vec3 & Nat:
59  match rest:
60    case (rest_sum, rest_n):
61      (G.Vec3.add(pos, rest_sum), 1n+rest_n)
62
63def sum_count_extend(h: Atom, rest: G.Vec3 & Nat) -> G.Vec3 & Nat:
64  match h:
65    case Atom{serial, elem, pos}:
66      sum_count_extend_go(pos, rest)
67
68# Sum positions + count in one pass (list consumed once)
69def sum_count_atoms(xs: List<&2, Atom>) -> G.Vec3 & Nat:
70  match xs:
71    case Nil{}:
72      (G.V3{0.0, 0.0, 0.0}, 0n)
73    case h <> t:
74      sum_count_extend(h, sum_count_atoms(t))
75
76def centroid_of_sum_go(sum: G.Vec3, n: Nat) -> G.Vec3:
77  match n:
78    case 0n:
79      sum
80    case 1n+p:
81      G.Vec3.scale((1.0 / (F32.from_nat(n)) : F32), sum)
82
83def centroid_of_sum(pair: G.Vec3 & Nat) -> G.Vec3:
84  match pair:
85    case (sum, n):
86      centroid_of_sum_go(sum, n)
87
88def Residue.centroid(r: Residue) -> G.Vec3:
89  match r:
90    case Residue{seq, atoms}:
91      centroid_of_sum(sum_count_atoms(atoms))
92
93def Residue.num_atoms(r: Residue) -> Nat:
94  match r:
95    case Residue{seq, atoms}:
96      count_atoms(atoms)
97
98def Chain.num_residues(c: Chain) -> Nat:
99  match c:
100    case Chain{cid, residues}:
101      count_residues(residues)
102
103def append_atoms(xs: List<&2, Atom>, ys: List<&2, Atom>) -> List<&2, Atom>:
104  match xs:
105    case Nil{}:
106      ys
107    case h <> t:
108      h <> append_atoms(t, ys)
109
110def flatten_res_extend(h: Residue, rest: List<&2, Atom>) -> List<&2, Atom>:
111  match h:
112    case Residue{seq, atoms}:
113      append_atoms(atoms, rest)
114
115def flatten_residues(xs: List<&2, Residue>) -> List<&2, Atom>:
116  match xs:
117    case Nil{}:
118      Nil{}
119    case h <> t:
120      flatten_res_extend(h, flatten_residues(t))
121
122def atoms_of(h: Residue) -> Nat:
123  match h:
124    case Residue{seq, atoms}:
125      count_atoms(atoms)
126
127def count_atoms_in(xs: List<&2, Residue>) -> Nat:
128  match xs:
129    case Nil{}:
130      0n
131    case h <> t:
132      Nat.add(count_atoms_in(t), atoms_of(h))
133
134def Chain.flatten(c: Chain) -> List<&2, Atom>:
135  match c:
136    case Chain{cid, residues}:
137      flatten_residues(residues)
138
139def Chain.num_atoms(c: Chain) -> Nat:
140  match c:
141    case Chain{cid, residues}:
142      count_atoms_in(residues)
143
144def Chain.centroid(c: Chain) -> G.Vec3:
145  match c:
146    case Chain{cid, residues}:
147      centroid_of_sum(sum_count_atoms(flatten_residues(residues)))
148
149def count_chains(xs: List<&2, Chain>) -> Nat:
150  match xs:
151    case Nil{}:
152      0n
153    case h <> t:
154      1n+count_chains(t)
155
156def Complex.num_chains(k: Complex) -> Nat:
157  match k:
158    case Complex{chains}:
159      count_chains(chains)
160
161def append_residues(
162  xs: List<&2, Residue>, ys: List<&2, Residue>
163) -> List<&2, Residue>:
164  match xs:
165    case Nil{}:
166      ys
167    case h <> t:
168      h <> append_residues(t, ys)
169
170def flatten_chains_extend(
171  h: Chain, rest: List<&2, Residue>
172) -> List<&2, Residue>:
173  match h:
174    case Chain{cid, residues}:
175      append_residues(residues, rest)
176
177def flatten_chains_residues(xs: List<&2, Chain>) -> List<&2, Residue>:
178  match xs:
179    case Nil{}:
180      Nil{}
181    case h <> t:
182      flatten_chains_extend(h, flatten_chains_residues(t))
183
184def Complex.num_residues(k: Complex) -> Nat:
185  match k:
186    case Complex{chains}:
187      count_residues(flatten_chains_residues(chains))
188
189def Complex.flatten(k: Complex) -> List<&2, Atom>:
190  match k:
191    case Complex{chains}:
192      flatten_residues(flatten_chains_residues(chains))
193
194def Complex.num_atoms(k: Complex) -> Nat:
195  match k:
196    case Complex{chains}:
197      count_atoms_in(flatten_chains_residues(chains))
198
199def Complex.centroid(k: Complex) -> G.Vec3:
200  match k:
201    case Complex{chains}:
202      centroid_of_sum(sum_count_atoms(flatten_residues(flatten_chains_residues(chains))))