1	import Base
2	import ./geom.bend as G
3	
4	# --- Protein model v0: Atom / Residue / Chain / Complex. All Data. ---
5	
6	type Atom is Data:
7	  Atom{serial: U32, elem: U32, pos: G.Vec3}
8	
9	type Residue is Data:
10	  Residue{seq: Nat, atoms: List<&2, Atom>}
11	
12	type Chain is Data:
13	  Chain{cid: U32, residues: List<&2, Residue>}
14	
15	type Complex is Data:
16	  Complex{chains: List<&2, Chain>}
17	
18	# Atom position (drops serial/elem, keeps pos)
19	def 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)
25	def 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)
31	def Atom.serial(a: Atom) -> U32:
32	  match a:
33	    case Atom{serial, elem, pos}:
34	      serial
35	
36	def 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	
41	def Atom.dist(+a: Atom, +b: Atom) -> F32:
42	  F32.sqrt(Atom.dist2(a, b))
43	
44	def 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	
51	def 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	
58	def 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	
63	def 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)
69	def 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	
76	def 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	
83	def centroid_of_sum(pair: G.Vec3 & Nat) -> G.Vec3:
84	  match pair:
85	    case (sum, n):
86	      centroid_of_sum_go(sum, n)
87	
88	def Residue.centroid(r: Residue) -> G.Vec3:
89	  match r:
90	    case Residue{seq, atoms}:
91	      centroid_of_sum(sum_count_atoms(atoms))
92	
93	def Residue.num_atoms(r: Residue) -> Nat:
94	  match r:
95	    case Residue{seq, atoms}:
96	      count_atoms(atoms)
97	
98	def Chain.num_residues(c: Chain) -> Nat:
99	  match c:
100	    case Chain{cid, residues}:
101	      count_residues(residues)
102	
103	def 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	
110	def 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	
115	def 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	
122	def atoms_of(h: Residue) -> Nat:
123	  match h:
124	    case Residue{seq, atoms}:
125	      count_atoms(atoms)
126	
127	def 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	
134	def Chain.flatten(c: Chain) -> List<&2, Atom>:
135	  match c:
136	    case Chain{cid, residues}:
137	      flatten_residues(residues)
138	
139	def Chain.num_atoms(c: Chain) -> Nat:
140	  match c:
141	    case Chain{cid, residues}:
142	      count_atoms_in(residues)
143	
144	def 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	
149	def 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	
156	def Complex.num_chains(k: Complex) -> Nat:
157	  match k:
158	    case Complex{chains}:
159	      count_chains(chains)
160	
161	def 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	
170	def 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	
177	def 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	
184	def Complex.num_residues(k: Complex) -> Nat:
185	  match k:
186	    case Complex{chains}:
187	      count_residues(flatten_chains_residues(chains))
188	
189	def Complex.flatten(k: Complex) -> List<&2, Atom>:
190	  match k:
191	    case Complex{chains}:
192	      flatten_residues(flatten_chains_residues(chains))
193	
194	def Complex.num_atoms(k: Complex) -> Nat:
195	  match k:
196	    case Complex{chains}:
197	      count_atoms_in(flatten_chains_residues(chains))
198	
199	def 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))))