use std::collections::{BTreeSet, HashMap, 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()
}
pub fn dot(self, other: Self) -> f32 {
self.x * other.x + self.y * other.y + self.z * other.z
}
pub fn cross(self, other: Self) -> Self {
Self::new(
self.y * other.z - self.z * other.y,
self.z * other.x - self.x * other.z,
self.x * other.y - self.y * other.x,
)
}
pub fn normalized(self) -> Self {
let length = self.length();
if length <= f32::EPSILON {
Self::default()
} else {
self * (1.0 / length)
}
}
pub fn lerp(self, other: Self, amount: f32) -> Self {
self * (1.0 - amount) + other * amount
}
}
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)
}
}
impl std::ops::Neg for Vec3 {
type Output = Self;
fn neg(self) -> Self::Output {
Self::new(-self.x, -self.y, -self.z)
}
}
#[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, Copy, Debug, Eq, PartialEq)]
pub enum SecondaryStructure {
Coil,
Helix,
Sheet,
}
impl SecondaryStructure {
pub const fn name(self) -> &'static str {
match self {
Self::Coil => "coil",
Self::Helix => "helix",
Self::Sheet => "sheet",
}
}
}
#[derive(Clone, Debug)]
pub struct Residue {
pub name: String,
pub chain: String,
pub number: i32,
pub insertion_code: char,
pub n: Option<usize>,
pub ca: Option<usize>,
pub c: Option<usize>,
pub o: Option<usize>,
pub secondary: SecondaryStructure,
}
#[derive(Clone, Debug)]
pub struct Protein {
pub source: PathBuf,
pub title: String,
pub atoms: Vec<Atom>,
pub residues: Vec<Residue>,
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 mut structured_residues = build_residues(&atoms);
assign_secondary_structure(&atoms, &mut structured_residues);
let alpha_carbons: Vec<usize> = structured_residues
.iter()
.filter_map(|residue| residue.ca)
.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,
residues: structured_residues,
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)
};
let helices = self
.residues
.iter()
.filter(|residue| residue.secondary == SecondaryStructure::Helix)
.count();
let sheets = self
.residues
.iter()
.filter(|residue| residue.secondary == SecondaryStructure::Sheet)
.count();
format!(
"{}\nsource: {}\natoms: {}\nresidues: {}\nchains: {} ({})\nC-alpha links: {}\nsecondary structure: {} helix / {} sheet / {} coil residues\n{}\npredicted structure: {}",
self.title,
self.source.display(),
self.atoms.len(),
self.residue_count,
self.chains.len(),
self.chains.join(", "),
self.trace.len(),
helices,
sheets,
self.residues.len().saturating_sub(helices + sheets),
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 build_residues(atoms: &[Atom]) -> Vec<Residue> {
let mut residues = Vec::new();
let mut lookup = HashMap::new();
for (atom_index, atom) in atoms.iter().enumerate() {
if atom.hetero {
continue;
}
let key = (atom.chain.clone(), atom.residue_number, atom.insertion_code);
let residue_index = *lookup.entry(key).or_insert_with(|| {
let index = residues.len();
residues.push(Residue {
name: atom.residue_name.clone(),
chain: atom.chain.clone(),
number: atom.residue_number,
insertion_code: atom.insertion_code,
n: None,
ca: None,
c: None,
o: None,
secondary: SecondaryStructure::Coil,
});
index
});
let residue = &mut residues[residue_index];
match atom.name.as_str() {
"N" => residue.n = Some(atom_index),
"CA" => residue.ca = Some(atom_index),
"C" => residue.c = Some(atom_index),
"O" | "OXT" => residue.o = Some(atom_index),
_ => {}
}
}
residues
}
fn assign_secondary_structure(atoms: &[Atom], residues: &mut [Residue]) {
if residues.len() < 3 {
return;
}
for index in 1..residues.len() - 1 {
let previous = &residues[index - 1];
let current = &residues[index];
let next = &residues[index + 1];
if previous.chain != current.chain || current.chain != next.chain {
continue;
}
let (Some(previous_c), Some(n), Some(ca), Some(c), Some(next_n)) =
(previous.c, current.n, current.ca, current.c, next.n)
else {
continue;
};
let phi = dihedral(
atoms[previous_c].position,
atoms[n].position,
atoms[ca].position,
atoms[c].position,
);
let psi = dihedral(
atoms[n].position,
atoms[ca].position,
atoms[c].position,
atoms[next_n].position,
);
residues[index].secondary =
if (-120.0..=-30.0).contains(&phi) && (-90.0..=45.0).contains(&psi) {
SecondaryStructure::Helix
} else if (-180.0..=-60.0).contains(&phi)
&& ((45.0..=180.0).contains(&psi) || (-180.0..=-150.0).contains(&psi))
{
SecondaryStructure::Sheet
} else {
SecondaryStructure::Coil
};
}
for index in 1..residues.len() - 1 {
if residues[index].secondary == SecondaryStructure::Coil
&& residues[index - 1].chain == residues[index].chain
&& residues[index + 1].chain == residues[index].chain
&& residues[index - 1].secondary == residues[index + 1].secondary
&& residues[index - 1].secondary != SecondaryStructure::Coil
{
residues[index].secondary = residues[index - 1].secondary;
}
}
remove_short_secondary_runs(residues, SecondaryStructure::Helix, 4);
remove_short_secondary_runs(residues, SecondaryStructure::Sheet, 3);
}
fn remove_short_secondary_runs(
residues: &mut [Residue],
structure: SecondaryStructure,
minimum: usize,
) {
let mut start = 0;
while start < residues.len() {
if residues[start].secondary != structure {
start += 1;
continue;
}
let chain = residues[start].chain.clone();
let mut end = start + 1;
while end < residues.len()
&& residues[end].chain == chain
&& residues[end].secondary == structure
{
end += 1;
}
if end - start < minimum {
for residue in &mut residues[start..end] {
residue.secondary = SecondaryStructure::Coil;
}
}
start = end;
}
}
fn dihedral(p0: Vec3, p1: Vec3, p2: Vec3, p3: Vec3) -> f32 {
let axis = (p2 - p1).normalized();
let first = p0 - p1;
let last = p3 - p2;
let first_plane = first - axis * first.dot(axis);
let last_plane = last - axis * last.dot(axis);
let x = first_plane.dot(last_plane);
let y = axis.cross(first_plane).dot(last_plane);
y.atan2(x).to_degrees()
}
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);
}
}