use std::collections::{BTreeSet, HashSet};
use std::fs;
use std::path::{Path, PathBuf};
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub struct Vec3 {
pub x: f32,
pub y: f32,
pub z: f32,
}
impl Vec3 {
pub const fn new(x: f32, y: f32, z: f32) -> Self {
Self { x, y, z }
}
pub fn length(self) -> f32 {
(self.x * self.x + self.y * self.y + self.z * self.z).sqrt()
}
pub fn distance(self, other: Self) -> f32 {
(self - other).length()
}
}
impl std::ops::Add for Vec3 {
type Output = Self;
fn add(self, rhs: Self) -> Self::Output {
Self::new(self.x + rhs.x, self.y + rhs.y, self.z + rhs.z)
}
}
impl std::ops::Sub for Vec3 {
type Output = Self;
fn sub(self, rhs: Self) -> Self::Output {
Self::new(self.x - rhs.x, self.y - rhs.y, self.z - rhs.z)
}
}
impl std::ops::Mul<f32> for Vec3 {
type Output = Self;
fn mul(self, rhs: f32) -> Self::Output {
Self::new(self.x * rhs, self.y * rhs, self.z * rhs)
}
}
#[derive(Clone, Debug)]
pub struct Atom {
pub serial: i32,
pub name: String,
pub residue_name: String,
pub chain: String,
pub residue_number: i32,
pub insertion_code: char,
pub position: Vec3,
pub occupancy: f32,
pub b_factor: f32,
pub element: String,
pub hetero: bool,
}
#[derive(Clone, Debug)]
pub struct TraceSegment {
pub from: usize,
pub to: usize,
}
#[derive(Clone, Debug)]
pub struct Protein {
pub source: PathBuf,
pub title: String,
pub atoms: Vec<Atom>,
pub trace: Vec<TraceSegment>,
pub center: Vec3,
pub radius: f32,
pub residue_count: usize,
pub chains: Vec<String>,
pub mean_b_factor: f32,
pub is_prediction: bool,
}
impl Protein {
pub fn from_path(path: impl AsRef<Path>) -> Result<Self, String> {
let path = path.as_ref();
let text = fs::read_to_string(path)
.map_err(|error| format!("cannot read {}: {error}", path.display()))?;
Self::from_pdb_str(&text, path)
}
pub fn from_pdb_str(text: &str, source: impl AsRef<Path>) -> Result<Self, String> {
let source = source.as_ref().to_path_buf();
let mut atoms = Vec::new();
let mut titles = Vec::new();
let mut saw_model = false;
let mut in_first_model = true;
let mut is_prediction = false;
for (line_number, line) in text.lines().enumerate() {
let record = field(line, 0, 6).trim();
match record {
"MODEL" => {
if saw_model {
in_first_model = false;
} else {
saw_model = true;
}
continue;
}
"ENDMDL" if saw_model && in_first_model => break,
"TITLE" => {
let part = field(line, 10, line.len()).trim();
if !part.is_empty() {
titles.push(part.to_owned());
}
if line.to_ascii_uppercase().contains("ESMFOLD") {
is_prediction = true;
}
continue;
}
"HEADER" => {
if line.to_ascii_uppercase().contains("ESMFOLD") {
is_prediction = true;
}
continue;
}
"ATOM" | "HETATM" if in_first_model => {}
_ => continue,
}
if line.len() < 54 {
return Err(format!(
"{}:{}: truncated {record} record",
source.display(),
line_number + 1
));
}
let alternate = char_at(line, 16);
if alternate != ' ' && alternate != 'A' {
continue;
}
let position = Vec3::new(
parse_required(line, 30, 38, "x", line_number, &source)?,
parse_required(line, 38, 46, "y", line_number, &source)?,
parse_required(line, 46, 54, "z", line_number, &source)?,
);
let raw_atom_name = field(line, 12, 16);
let atom_name = raw_atom_name.trim().to_owned();
let element_field = field(line, 76, 78).trim();
let element = if element_field.is_empty() {
infer_element(raw_atom_name)
} else {
normalize_element(element_field)
};
let chain = field(line, 21, 22).trim();
atoms.push(Atom {
serial: parse_or(line, 6, 11, atoms.len() as i32 + 1),
name: atom_name,
residue_name: field(line, 17, 20).trim().to_owned(),
chain: if chain.is_empty() { "_" } else { chain }.to_owned(),
residue_number: parse_or(line, 22, 26, 0),
insertion_code: char_at(line, 26),
position,
occupancy: parse_or(line, 54, 60, 1.0),
b_factor: parse_or(line, 60, 66, 0.0),
element,
hetero: record == "HETATM",
});
}
if atoms.is_empty() {
return Err(format!(
"{} contains no readable ATOM or HETATM records",
source.display()
));
}
let mut minimum = atoms[0].position;
let mut maximum = atoms[0].position;
let mut residues = HashSet::new();
let mut chains = BTreeSet::new();
let mut b_factor_sum = 0.0;
for atom in &atoms {
minimum.x = minimum.x.min(atom.position.x);
minimum.y = minimum.y.min(atom.position.y);
minimum.z = minimum.z.min(atom.position.z);
maximum.x = maximum.x.max(atom.position.x);
maximum.y = maximum.y.max(atom.position.y);
maximum.z = maximum.z.max(atom.position.z);
residues.insert((atom.chain.clone(), atom.residue_number, atom.insertion_code));
chains.insert(atom.chain.clone());
b_factor_sum += atom.b_factor;
}
let center = (minimum + maximum) * 0.5;
let radius = atoms
.iter()
.map(|atom| (atom.position - center).length())
.fold(0.0_f32, f32::max)
.max(0.001);
let alpha_carbons: Vec<usize> = atoms
.iter()
.enumerate()
.filter(|(_, atom)| !atom.hetero && atom.name == "CA")
.map(|(index, _)| index)
.collect();
let trace = alpha_carbons
.windows(2)
.filter_map(|pair| {
let from = &atoms[pair[0]];
let to = &atoms[pair[1]];
(from.chain == to.chain && from.position.distance(to.position) <= 4.5).then_some(
TraceSegment {
from: pair[0],
to: pair[1],
},
)
})
.collect();
let title = if titles.is_empty() {
source
.file_name()
.and_then(|name| name.to_str())
.unwrap_or("protein")
.to_owned()
} else {
titles.join(" ")
};
let atom_count = atoms.len() as f32;
Ok(Self {
source,
title,
atoms,
trace,
center,
radius,
residue_count: residues.len(),
chains: chains.into_iter().collect(),
mean_b_factor: b_factor_sum / atom_count,
is_prediction,
})
}
pub fn summary(&self) -> String {
let metric = if self.is_prediction {
format!("mean pLDDT: {:.2}", self.mean_confidence())
} else {
format!("mean B-factor: {:.2}", self.mean_b_factor)
};
format!(
"{}\nsource: {}\natoms: {}\nresidues: {}\nchains: {} ({})\nC-alpha links: {}\n{}\npredicted structure: {}",
self.title,
self.source.display(),
self.atoms.len(),
self.residue_count,
self.chains.len(),
self.chains.join(", "),
self.trace.len(),
metric,
if self.is_prediction { "yes" } else { "no" }
)
}
pub fn confidence(&self, b_factor: f32) -> f32 {
if self.is_prediction && self.mean_b_factor <= 1.5 {
b_factor * 100.0
} else {
b_factor
}
}
pub fn mean_confidence(&self) -> f32 {
self.confidence(self.mean_b_factor)
}
}
fn field(line: &str, start: usize, end: usize) -> &str {
if start >= line.len() {
""
} else {
&line[start..end.min(line.len())]
}
}
fn char_at(line: &str, index: usize) -> char {
line.as_bytes()
.get(index)
.copied()
.map(char::from)
.unwrap_or(' ')
}
fn parse_required(
line: &str,
start: usize,
end: usize,
label: &str,
line_number: usize,
source: &Path,
) -> Result<f32, String> {
field(line, start, end).trim().parse().map_err(|_| {
format!(
"{}:{}: invalid {label} coordinate",
source.display(),
line_number + 1
)
})
}
fn parse_or<T: std::str::FromStr>(line: &str, start: usize, end: usize, fallback: T) -> T {
field(line, start, end).trim().parse().unwrap_or(fallback)
}
fn infer_element(raw_atom_name: &str) -> String {
let atom_name = raw_atom_name.trim();
let letters: String = atom_name
.chars()
.filter(|character| character.is_ascii_alphabetic())
.collect();
if letters.is_empty() {
return "?".to_owned();
}
let upper = letters.to_ascii_uppercase();
const TWO_LETTER: &[&str] = &[
"BR", "CL", "FE", "MG", "MN", "ZN", "CA", "NA", "CU", "CO", "NI", "SE",
];
if raw_atom_name.starts_with(' ')
|| atom_name.chars().next().is_some_and(|c| c.is_ascii_digit())
|| !TWO_LETTER.iter().any(|item| upper.starts_with(item))
{
upper[..1].to_owned()
} else {
upper[..2].to_owned()
}
}
fn normalize_element(element: &str) -> String {
element.trim().to_ascii_uppercase()
}
#[cfg(test)]
mod tests {
use super::*;
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";
#[test]
fn parses_structure_and_trace() {
let protein = Protein::from_pdb_str(MINI_PDB, "test.pdb").unwrap();
assert_eq!(protein.atoms.len(), 4);
assert_eq!(protein.residue_count, 3);
assert_eq!(protein.chains, ["A", "B"]);
assert_eq!(protein.trace.len(), 1);
assert!(protein.is_prediction);
assert!((protein.mean_b_factor - 80.0).abs() < 0.01);
assert_eq!(protein.atoms[3].element, "ZN");
}
#[test]
fn rejects_empty_pdb() {
let error = Protein::from_pdb_str("HEADER EMPTY\n", "empty.pdb").unwrap_err();
assert!(error.contains("no readable"));
}
#[test]
fn ignores_non_primary_altloc() {
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";
let protein = Protein::from_pdb_str(pdb, "alt.pdb").unwrap();
assert_eq!(protein.atoms.len(), 1);
assert_eq!(protein.atoms[0].position.x, 1.0);
}
#[test]
fn infers_alpha_carbon_as_carbon_without_element_column() {
assert_eq!(infer_element(" CA "), "C");
assert_eq!(infer_element("ZN "), "ZN");
}
#[test]
fn normalizes_early_atlas_confidence_scale() {
let pdb = "TITLE ESMFOLD V0 TEST\nATOM 1 CA ALA A 1 0.000 0.000 0.000 1.00 0.75 C \n";
let protein = Protein::from_pdb_str(pdb, "atlas.pdb").unwrap();
assert!((protein.mean_confidence() - 75.0).abs() < 0.01);
}
}