use crate::atom::{element_symbol, AtomStr, ATOM_NAME_EXPECT, ATOM_RESNAME_EXPECT};
use crate::prelude::*;
use std::{
fs::File,
io::{BufRead, BufReader, BufWriter, Write},
num::{ParseFloatError, ParseIntError},
path::Path,
};
use thiserror::Error;
pub struct PdbFileHandler {
reader: Option<BufReader<File>>,
bonds: Vec<[usize; 2]>, stored_topology: Option<Topology>,
stored_state: Option<State>,
at_least_one_state_read: bool,
writer: Option<BufWriter<File>>,
}
#[derive(Debug, Error)]
pub enum PdbHandlerError {
#[error("can't open pdb file for reading")]
OpenRead(#[source] std::io::Error),
#[error("can't open pdb file for writing")]
OpenWrite(#[source] std::io::Error),
#[error("pdb file is empty or contains no ATOM/HETATM records")]
Empty,
#[error("unexpected end of file reached")]
Eof,
#[error(transparent)]
ParseInt(#[from] ParseIntError),
#[error(transparent)]
ParseFloat(#[from] ParseFloatError),
#[error("invalid periodic box")]
Pbc(#[from] PeriodicBoxError),
#[error("io error")]
Io(#[from] std::io::Error),
}
fn parse_f32(line: &str, start: usize, end: usize) -> Result<f32, PdbHandlerError> {
line.get(start..end)
.unwrap_or("")
.trim()
.parse::<f32>()
.map_err(PdbHandlerError::ParseFloat)
}
fn parse_f32_opt(line: &str, start: usize, end: usize) -> Option<f32> {
line.get(start..end)?.trim().parse::<f32>().ok()
}
fn parse_i32_opt(line: &str, start: usize, end: usize) -> Option<i32> {
line.get(start..end)?.trim().parse::<i32>().ok()
}
fn parse_serial_opt(line: &str, start: usize, end: usize) -> Option<usize> {
let s = line.get(start..end)?.trim();
if s.is_empty() {
return None;
}
s.parse::<usize>().ok().map(|n| n.saturating_sub(1))
}
fn format_atom_name(name: &str) -> String {
if name.len() >= 4 {
format!("{:<4}", &name[..4])
} else {
format!(" {:<3}", name)
}
}
impl FileFormatHandler for PdbFileHandler {
fn open(fname: impl AsRef<Path>) -> Result<Self, FileFormatError>
where
Self: Sized,
{
Ok(Self {
reader: Some(BufReader::new(
File::open(fname).map_err(PdbHandlerError::OpenRead)?,
)),
writer: None,
bonds: Vec::new(),
stored_topology: None,
stored_state: None,
at_least_one_state_read: false,
})
}
fn create(fname: impl AsRef<Path>) -> Result<Self, FileFormatError>
where
Self: Sized,
{
Ok(Self {
reader: None,
writer: Some(BufWriter::new(
File::create(fname).map_err(PdbHandlerError::OpenWrite)?,
)),
bonds: Vec::new(),
stored_topology: None,
stored_state: None,
at_least_one_state_read: false,
})
}
fn read(&mut self) -> Result<(Topology, State), FileFormatError> {
let buf = self.reader.as_mut().unwrap();
if buf.fill_buf()?.is_empty() {
return if self.at_least_one_state_read {
Err(PdbHandlerError::Eof)?
} else {
Err(PdbHandlerError::Empty)?
};
}
self.bonds.clear();
let mut atoms: Vec<Atom> = Vec::new();
let mut coords: Vec<Pos> = Vec::new();
let mut current_box: Option<PeriodicBox> = None;
let mut has_atoms = false;
let mut line = String::new();
loop {
line.clear();
let n = buf.read_line(&mut line)?;
if n == 0 {
break;
}
if line.starts_with("ATOM ") || line.starts_with("HETATM") {
has_atoms = true;
let x = parse_f32(&line, 30, 38).unwrap_or(0.0);
let y = parse_f32(&line, 38, 46).unwrap_or(0.0);
let z = parse_f32(&line, 46, 54).unwrap_or(0.0);
coords.push(Pos::new(x * 0.1, y * 0.1, z * 0.1));
let name = line.get(12..16).unwrap_or("").trim();
let resname = line.get(17..20).unwrap_or("").trim();
let chain = line
.get(21..22)
.and_then(|s| s.chars().next())
.unwrap_or(' ');
let resid = parse_i32_opt(&line, 22, 26).unwrap_or(0);
let occupancy = parse_f32_opt(&line, 54, 60).unwrap_or(1.0);
let bfactor = parse_f32_opt(&line, 60, 66).unwrap_or(0.0);
let mut at = Atom {
name: AtomStr::try_from_str(name).expect(ATOM_NAME_EXPECT),
resname: AtomStr::try_from_str(resname).expect(ATOM_RESNAME_EXPECT),
resid,
chain,
occupancy,
bfactor,
type_name: AtomStr::try_from_str("").unwrap(),
..Default::default()
};
at.guess_element_and_mass_from_name();
atoms.push(at);
} else if line.starts_with("CRYST1") {
let a = parse_f32(&line, 6, 15).unwrap_or(0.0);
let b = parse_f32(&line, 15, 24).unwrap_or(0.0);
let c = parse_f32(&line, 24, 33).unwrap_or(0.0);
let alpha = parse_f32(&line, 33, 40).unwrap_or(90.0);
let beta = parse_f32(&line, 40, 47).unwrap_or(90.0);
let gamma = parse_f32(&line, 47, 54).unwrap_or(90.0);
current_box = PeriodicBox::from_vectors_angles(
a * 0.1,
b * 0.1,
c * 0.1,
alpha,
beta,
gamma,
)
.ok();
} else if line.starts_with("MODEL") {
if has_atoms {
break;
}
} else if line.starts_with("ENDMDL") {
break;
} else if line.starts_with("CONECT") {
if let Some(a) = parse_serial_opt(&line, 6, 11) {
for (s, e) in [(11usize, 16usize), (16, 21), (21, 26), (26, 31)] {
if let Some(b) = parse_serial_opt(&line, s, e) {
if a != b {
let mut pair = [a, b];
pair.sort();
self.bonds.push(pair);
}
}
}
}
} else if line.starts_with("END") && !line.starts_with("ENDMDL") {
break;
}
}
if !has_atoms {
return if self.at_least_one_state_read {
Err(PdbHandlerError::Eof)?
} else {
Err(PdbHandlerError::Empty)?
};
}
self.bonds.sort();
self.bonds.dedup();
let mut top = Topology::default();
top.atoms = atoms;
top.bonds = self.bonds.clone();
top.assign_resindex();
self.at_least_one_state_read = true;
Ok((top, State { coords, time: 0.0, pbox: current_box }))
}
fn read_topology(&mut self) -> Result<Topology, FileFormatError> {
if self.stored_topology.is_some() {
Ok(self.stored_topology.take().unwrap())
} else {
let (top, st) = self.read()?;
self.stored_state.get_or_insert(st);
Ok(top)
}
}
fn read_state(&mut self) -> Result<State, FileFormatError> {
if self.stored_state.is_some() {
Ok(self.stored_state.take().unwrap())
} else {
let (top, st) = self.read()?;
self.stored_topology.get_or_insert(top);
Ok(st)
}
}
fn write(&mut self, data: &dyn SaveTopologyState) -> Result<(), FileFormatError> {
let w = self.writer.as_mut().ok_or(FileFormatError::NotWritable)?;
write_cryst1(w, data.get_box())?;
for (i, (at, pos)) in data.iter_atoms_dyn().zip(data.iter_pos_dyn()).enumerate() {
write_atom_record(w, (i % 99999) + 1, at, pos)?;
}
writeln!(w, "END")?;
Ok(())
}
}
fn write_cryst1(
w: &mut BufWriter<File>,
pbox: Option<&PeriodicBox>,
) -> Result<(), FileFormatError> {
if let Some(b) = pbox {
let (lengths, angles) = b.to_vectors_angles();
writeln!(
w,
"CRYST1{:9.3}{:9.3}{:9.3}{:7.2}{:7.2}{:7.2} P 1 1",
lengths[0] * 10.0,
lengths[1] * 10.0,
lengths[2] * 10.0,
angles[0],
angles[1],
angles[2]
)?;
}
Ok(())
}
fn write_atom_record(
w: &mut BufWriter<File>,
serial: usize,
at: &Atom,
pos: &Pos,
) -> Result<(), FileFormatError> {
let name = format_atom_name(at.name.as_str());
let elem = element_symbol(at.atomic_number);
writeln!(
w,
"ATOM {:>5} {:<4} {:<3} {:1}{:>4} {:>8.3}{:>8.3}{:>8.3}{:>6.2}{:>6.2} {:>2}",
serial,
name,
at.resname,
at.chain,
at.resid % 9999,
pos.x * 10.0,
pos.y * 10.0,
pos.z * 10.0,
at.occupancy,
at.bfactor,
elem
)?;
Ok(())
}