1use std::collections::{BTreeSet, HashSet};
2use std::fs;
3use std::path::{Path, PathBuf};
4
5#[derive(Clone, Copy, Debug, Default, PartialEq)]
6pub struct Vec3 {
7 pub x: f32,
8 pub y: f32,
9 pub z: f32,
10}
11
12impl Vec3 {
13 pub const fn new(x: f32, y: f32, z: f32) -> Self {
14 Self { x, y, z }
15 }
16
17 pub fn length(self) -> f32 {
18 (self.x * self.x + self.y * self.y + self.z * self.z).sqrt()
19 }
20
21 pub fn distance(self, other: Self) -> f32 {
22 (self - other).length()
23 }
24}
25
26impl std::ops::Add for Vec3 {
27 type Output = Self;
28
29 fn add(self, rhs: Self) -> Self::Output {
30 Self::new(self.x + rhs.x, self.y + rhs.y, self.z + rhs.z)
31 }
32}
33
34impl std::ops::Sub for Vec3 {
35 type Output = Self;
36
37 fn sub(self, rhs: Self) -> Self::Output {
38 Self::new(self.x - rhs.x, self.y - rhs.y, self.z - rhs.z)
39 }
40}
41
42impl std::ops::Mul<f32> for Vec3 {
43 type Output = Self;
44
45 fn mul(self, rhs: f32) -> Self::Output {
46 Self::new(self.x * rhs, self.y * rhs, self.z * rhs)
47 }
48}
49
50#[derive(Clone, Debug)]
51pub struct Atom {
52 pub serial: i32,
53 pub name: String,
54 pub residue_name: String,
55 pub chain: String,
56 pub residue_number: i32,
57 pub insertion_code: char,
58 pub position: Vec3,
59 pub occupancy: f32,
60 pub b_factor: f32,
61 pub element: String,
62 pub hetero: bool,
63}
64
65#[derive(Clone, Debug)]
66pub struct TraceSegment {
67 pub from: usize,
68 pub to: usize,
69}
70
71#[derive(Clone, Debug)]
72pub struct Protein {
73 pub source: PathBuf,
74 pub title: String,
75 pub atoms: Vec<Atom>,
76 pub trace: Vec<TraceSegment>,
77 pub center: Vec3,
78 pub radius: f32,
79 pub residue_count: usize,
80 pub chains: Vec<String>,
81 pub mean_b_factor: f32,
82 pub is_prediction: bool,
83}
84
85impl Protein {
86 pub fn from_path(path: impl AsRef<Path>) -> Result<Self, String> {
87 let path = path.as_ref();
88 let text = fs::read_to_string(path)
89 .map_err(|error| format!("cannot read {}: {error}", path.display()))?;
90 Self::from_pdb_str(&text, path)
91 }
92
93 pub fn from_pdb_str(text: &str, source: impl AsRef<Path>) -> Result<Self, String> {
94 let source = source.as_ref().to_path_buf();
95 let mut atoms = Vec::new();
96 let mut titles = Vec::new();
97 let mut saw_model = false;
98 let mut in_first_model = true;
99 let mut is_prediction = false;
100
101 for (line_number, line) in text.lines().enumerate() {
102 let record = field(line, 0, 6).trim();
103 match record {
104 "MODEL" => {
105 if saw_model {
106 in_first_model = false;
107 } else {
108 saw_model = true;
109 }
110 continue;
111 }
112 "ENDMDL" if saw_model && in_first_model => break,
113 "TITLE" => {
114 let part = field(line, 10, line.len()).trim();
115 if !part.is_empty() {
116 titles.push(part.to_owned());
117 }
118 if line.to_ascii_uppercase().contains("ESMFOLD") {
119 is_prediction = true;
120 }
121 continue;
122 }
123 "HEADER" => {
124 if line.to_ascii_uppercase().contains("ESMFOLD") {
125 is_prediction = true;
126 }
127 continue;
128 }
129 "ATOM" | "HETATM" if in_first_model => {}
130 _ => continue,
131 }
132
133 if line.len() < 54 {
134 return Err(format!(
135 "{}:{}: truncated {record} record",
136 source.display(),
137 line_number + 1
138 ));
139 }
140 let alternate = char_at(line, 16);
141 if alternate != ' ' && alternate != 'A' {
142 continue;
143 }
144 let position = Vec3::new(
145 parse_required(line, 30, 38, "x", line_number, &source)?,
146 parse_required(line, 38, 46, "y", line_number, &source)?,
147 parse_required(line, 46, 54, "z", line_number, &source)?,
148 );
149 let raw_atom_name = field(line, 12, 16);
150 let atom_name = raw_atom_name.trim().to_owned();
151 let element_field = field(line, 76, 78).trim();
152 let element = if element_field.is_empty() {
153 infer_element(raw_atom_name)
154 } else {
155 normalize_element(element_field)
156 };
157 let chain = field(line, 21, 22).trim();
158 atoms.push(Atom {
159 serial: parse_or(line, 6, 11, atoms.len() as i32 + 1),
160 name: atom_name,
161 residue_name: field(line, 17, 20).trim().to_owned(),
162 chain: if chain.is_empty() { "_" } else { chain }.to_owned(),
163 residue_number: parse_or(line, 22, 26, 0),
164 insertion_code: char_at(line, 26),
165 position,
166 occupancy: parse_or(line, 54, 60, 1.0),
167 b_factor: parse_or(line, 60, 66, 0.0),
168 element,
169 hetero: record == "HETATM",
170 });
171 }
172
173 if atoms.is_empty() {
174 return Err(format!(
175 "{} contains no readable ATOM or HETATM records",
176 source.display()
177 ));
178 }
179
180 let mut minimum = atoms[0].position;
181 let mut maximum = atoms[0].position;
182 let mut residues = HashSet::new();
183 let mut chains = BTreeSet::new();
184 let mut b_factor_sum = 0.0;
185 for atom in &atoms {
186 minimum.x = minimum.x.min(atom.position.x);
187 minimum.y = minimum.y.min(atom.position.y);
188 minimum.z = minimum.z.min(atom.position.z);
189 maximum.x = maximum.x.max(atom.position.x);
190 maximum.y = maximum.y.max(atom.position.y);
191 maximum.z = maximum.z.max(atom.position.z);
192 residues.insert((atom.chain.clone(), atom.residue_number, atom.insertion_code));
193 chains.insert(atom.chain.clone());
194 b_factor_sum += atom.b_factor;
195 }
196 let center = (minimum + maximum) * 0.5;
197 let radius = atoms
198 .iter()
199 .map(|atom| (atom.position - center).length())
200 .fold(0.0_f32, f32::max)
201 .max(0.001);
202
203 let alpha_carbons: Vec<usize> = atoms
204 .iter()
205 .enumerate()
206 .filter(|(_, atom)| !atom.hetero && atom.name == "CA")
207 .map(|(index, _)| index)
208 .collect();
209 let trace = alpha_carbons
210 .windows(2)
211 .filter_map(|pair| {
212 let from = &atoms[pair[0]];
213 let to = &atoms[pair[1]];
214 (from.chain == to.chain && from.position.distance(to.position) <= 4.5).then_some(
215 TraceSegment {
216 from: pair[0],
217 to: pair[1],
218 },
219 )
220 })
221 .collect();
222
223 let title = if titles.is_empty() {
224 source
225 .file_name()
226 .and_then(|name| name.to_str())
227 .unwrap_or("protein")
228 .to_owned()
229 } else {
230 titles.join(" ")
231 };
232 let atom_count = atoms.len() as f32;
233 Ok(Self {
234 source,
235 title,
236 atoms,
237 trace,
238 center,
239 radius,
240 residue_count: residues.len(),
241 chains: chains.into_iter().collect(),
242 mean_b_factor: b_factor_sum / atom_count,
243 is_prediction,
244 })
245 }
246
247 pub fn summary(&self) -> String {
248 let metric = if self.is_prediction {
249 format!("mean pLDDT: {:.2}", self.mean_confidence())
250 } else {
251 format!("mean B-factor: {:.2}", self.mean_b_factor)
252 };
253 format!(
254 "{}\nsource: {}\natoms: {}\nresidues: {}\nchains: {} ({})\nC-alpha links: {}\n{}\npredicted structure: {}",
255 self.title,
256 self.source.display(),
257 self.atoms.len(),
258 self.residue_count,
259 self.chains.len(),
260 self.chains.join(", "),
261 self.trace.len(),
262 metric,
263 if self.is_prediction { "yes" } else { "no" }
264 )
265 }
266
267 pub fn confidence(&self, b_factor: f32) -> f32 {
271 if self.is_prediction && self.mean_b_factor <= 1.5 {
272 b_factor * 100.0
273 } else {
274 b_factor
275 }
276 }
277
278 pub fn mean_confidence(&self) -> f32 {
279 self.confidence(self.mean_b_factor)
280 }
281}
282
283fn field(line: &str, start: usize, end: usize) -> &str {
284 if start >= line.len() {
285 ""
286 } else {
287 &line[start..end.min(line.len())]
288 }
289}
290
291fn char_at(line: &str, index: usize) -> char {
292 line.as_bytes()
293 .get(index)
294 .copied()
295 .map(char::from)
296 .unwrap_or(' ')
297}
298
299fn parse_required(
300 line: &str,
301 start: usize,
302 end: usize,
303 label: &str,
304 line_number: usize,
305 source: &Path,
306) -> Result<f32, String> {
307 field(line, start, end).trim().parse().map_err(|_| {
308 format!(
309 "{}:{}: invalid {label} coordinate",
310 source.display(),
311 line_number + 1
312 )
313 })
314}
315
316fn parse_or<T: std::str::FromStr>(line: &str, start: usize, end: usize, fallback: T) -> T {
317 field(line, start, end).trim().parse().unwrap_or(fallback)
318}
319
320fn infer_element(raw_atom_name: &str) -> String {
321 let atom_name = raw_atom_name.trim();
322 let letters: String = atom_name
323 .chars()
324 .filter(|character| character.is_ascii_alphabetic())
325 .collect();
326 if letters.is_empty() {
327 return "?".to_owned();
328 }
329 let upper = letters.to_ascii_uppercase();
330 const TWO_LETTER: &[&str] = &[
331 "BR", "CL", "FE", "MG", "MN", "ZN", "CA", "NA", "CU", "CO", "NI", "SE",
332 ];
333 if raw_atom_name.starts_with(' ')
334 || atom_name.chars().next().is_some_and(|c| c.is_ascii_digit())
335 || !TWO_LETTER.iter().any(|item| upper.starts_with(item))
336 {
337 upper[..1].to_owned()
338 } else {
339 upper[..2].to_owned()
340 }
341}
342
343fn normalize_element(element: &str) -> String {
344 element.trim().to_ascii_uppercase()
345}
346
347#[cfg(test)]
348mod tests {
349 use super::*;
350
351 const MINI_PDB: &str = "TITLE ESMFOLD TEST\nATOM 1 N ALA A 1 -1.000 0.000 0.000 1.00 95.00 N \nATOM 2 CA ALA A 1 0.000 0.000 0.000 1.00 95.00 C \nATOM 3 CA GLY A 2 3.800 0.000 0.000 1.00 80.00 C \nHETATM 4 ZN ZN B 3 5.000 1.000 0.000 1.00 50.00 ZN\nEND\n";
352
353 #[test]
354 fn parses_structure_and_trace() {
355 let protein = Protein::from_pdb_str(MINI_PDB, "test.pdb").unwrap();
356 assert_eq!(protein.atoms.len(), 4);
357 assert_eq!(protein.residue_count, 3);
358 assert_eq!(protein.chains, ["A", "B"]);
359 assert_eq!(protein.trace.len(), 1);
360 assert!(protein.is_prediction);
361 assert!((protein.mean_b_factor - 80.0).abs() < 0.01);
362 assert_eq!(protein.atoms[3].element, "ZN");
363 }
364
365 #[test]
366 fn rejects_empty_pdb() {
367 let error = Protein::from_pdb_str("HEADER EMPTY\n", "empty.pdb").unwrap_err();
368 assert!(error.contains("no readable"));
369 }
370
371 #[test]
372 fn ignores_non_primary_altloc() {
373 let pdb = "ATOM 1 CA BALA A 1 0.000 0.000 0.000 1.00 20.00 C \nATOM 2 CA AALA A 1 1.000 0.000 0.000 1.00 20.00 C \n";
374 let protein = Protein::from_pdb_str(pdb, "alt.pdb").unwrap();
375 assert_eq!(protein.atoms.len(), 1);
376 assert_eq!(protein.atoms[0].position.x, 1.0);
377 }
378
379 #[test]
380 fn infers_alpha_carbon_as_carbon_without_element_column() {
381 assert_eq!(infer_element(" CA "), "C");
382 assert_eq!(infer_element("ZN "), "ZN");
383 }
384
385 #[test]
386 fn normalizes_early_atlas_confidence_scale() {
387 let pdb = "TITLE ESMFOLD V0 TEST\nATOM 1 CA ALA A 1 0.000 0.000 0.000 1.00 0.75 C \n";
388 let protein = Protein::from_pdb_str(pdb, "atlas.pdb").unwrap();
389 assert!((protein.mean_confidence() - 75.0).abs() < 0.01);
390 }
391}