molframe_pdb/writer/

pdb.rs

//! Writing the legacy fixed-column format, or refusing to.
//!
//! The format's fields have fixed widths, and a structure can outgrow every one
//! of them: a chain label of two characters, a serial past ninety-nine thousand,
//! a residue number past ten thousand, a coordinate too large for eight columns.
//!
//! Every one of those is checked before a single byte is written, and a
//! structure that does not fit is refused with a diagnostic naming the field and
//! offering a way out. Truncating a chain label from `AA` to `A` produces a file
//! that opens cleanly in any viewer and describes a different molecule, which is
//! the failure this library exists to prevent.

use super::capacity::{RequiredAtomFields, check_capacity, required_text, required_value};
use crate::hybrid36;
use crate::{PDB_HEADERS_EXTENSION, PdbHeaders};
use molframe_core::diagnostic::Diagnostic;
use molframe_core::io::{Select, SelectAll};
use molframe_core::structure::{AtomRef, ChainRef, ResidueRef, Structure};
use std::collections::HashMap;
use std::fmt;

#[path = "identity.rs"]
mod identity;
pub use identity::PdbIdentifierNamespace;
pub(crate) use identity::{atom_name, chain_name, component_name, residue_sequence};

/// How to write, and what to do about a structure that does not fit.
#[derive(Clone, Debug)]
pub struct PdbOptions {
    /// Rename chains on the way out, so a structure with labels the format
    /// cannot hold can still be written deliberately.
    chain_map: HashMap<Box<str>, Box<str>>,
    /// Write serials and residue numbers past the decimal range in hybrid-36.
    ///
    /// Off by default: a consumer that does not implement the scheme would read
    /// a wrong number without noticing, so it is offered rather than assumed.
    hybrid36: bool,
    /// Identifier namespace used consistently for every deposited identity.
    namespace: PdbIdentifierNamespace,
}

impl Default for PdbOptions {
    fn default() -> Self {
        Self {
            chain_map: HashMap::new(),
            hybrid36: false,
            namespace: PdbIdentifierNamespace::Label,
        }
    }
}

impl PdbOptions {
    /// Options that refuse anything the format cannot hold.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Renames one chain on the way out.
    #[must_use]
    pub fn chain_map(mut self, from: impl Into<Box<str>>, to: impl Into<Box<str>>) -> Self {
        self.chain_map.insert(from.into(), to.into());
        self
    }

    /// Writes counts past the decimal range in hybrid-36 rather than refusing.
    #[must_use]
    pub const fn hybrid36(mut self, enabled: bool) -> Self {
        self.hybrid36 = enabled;
        self
    }

    /// Selects one identifier namespace for chains, residues and atoms.
    #[must_use]
    pub const fn namespace(mut self, namespace: PdbIdentifierNamespace) -> Self {
        self.namespace = namespace;
        self
    }

    pub(crate) const fn identifier_namespace(&self) -> PdbIdentifierNamespace {
        self.namespace
    }

    pub(crate) const fn uses_hybrid36(&self) -> bool {
        self.hybrid36
    }

    /// The label a chain is written under.
    pub(crate) fn label_for<'a>(&'a self, label: &'a str) -> &'a str {
        match self.chain_map.get(label) {
            Some(replacement) => replacement,
            None => label,
        }
    }
}

/// Writes a structure, or explains why it cannot be written.
///
/// # Errors
///
/// Returns every capacity the structure exceeds, so a caller sees the whole
/// problem rather than fixing one field at a time.
pub fn write(structure: &Structure, options: &PdbOptions) -> Result<String, Vec<Diagnostic>> {
    write_selected(structure, options, &SelectAll)
}

/// Writes the parts of a structure a filter accepts.
///
/// # Errors
///
/// Returns every capacity the accepted parts exceed.
pub fn write_selected(
    structure: &Structure,
    options: &PdbOptions,
    select: &impl Select,
) -> Result<String, Vec<Diagnostic>> {
    let mut out = String::with_capacity(structure.atom_count() as usize * 81);
    render_selected(&mut out, structure, options, select)?;
    Ok(out)
}

pub(super) fn render_selected(
    out: &mut impl fmt::Write,
    structure: &Structure,
    options: &PdbOptions,
    select: &impl Select,
) -> Result<(), Vec<Diagnostic>> {
    let mut refusals = check_capacity(structure, options, select, RequiredAtomFields::Pdb);
    refusals.extend(super::secondary::preflight(structure, options, select));
    if !refusals.is_empty() {
        return Err(refusals);
    }

    let data = structure.data();

    if let Some(headers) = data.extensions.get::<PdbHeaders>(PDB_HEADERS_EXTENSION) {
        for record in headers.records() {
            if matches!(record.name(), "HELIX" | "SHEET")
                && structure
                    .secondary_structure()
                    .iter()
                    .any(|state| state.is_helix() || state.is_strand())
            {
                continue;
            }
            let _ = writeln!(out, "{}", record.line());
        }
    } else {
        write_generated_metadata(out, structure);
    }
    super::secondary::write(out, structure, options, select);

    let mut serial = 1i64;
    for chain in data.chains() {
        if !select.accept_chain(chain.index()) {
            continue;
        }
        let label = chain_label(&chain, options.namespace);
        let label = options.label_for(label);

        for residue in chain.residues() {
            if !select.accept_residue(residue.index()) {
                continue;
            }
            for atom in residue.atoms() {
                if !select.accept_atom(atom.index()) {
                    continue;
                }
                if let Err(finding) =
                    write_atom(out, structure, &atom, &residue, label, serial, options)
                {
                    return Err(vec![finding]);
                }
                serial += 1;
            }
        }
        let _ = writeln!(out, "TER   {:>5}", serial_field(serial, options));
        serial += 1;
    }
    let _ = out.write_str("END\n");
    Ok(())
}

fn write_generated_metadata(out: &mut impl fmt::Write, structure: &Structure) {
    let data = structure.data();
    if let Some(id) = &data.entry.id {
        let _ = writeln!(out, "HEADER    {:<40}{:<9}   {:>4}", "", "", id);
    }
    if let Some(title) = &data.entry.title {
        let _ = writeln!(out, "TITLE     {title}");
    }
    if let Some(method) = &data.entry.method {
        let _ = writeln!(out, "EXPDTA    {method}");
    }
    if let Some(cell) = data.cell {
        let _ = writeln!(
            out,
            "CRYST1{:9.3}{:9.3}{:9.3}{:7.2}{:7.2}{:7.2} {:<11}",
            cell.lengths[0],
            cell.lengths[1],
            cell.lengths[2],
            cell.angles[0],
            cell.angles[1],
            cell.angles[2],
            // CRYST1 names P 1 when the entry records no space group.
            match data.entry.space_group.as_deref() {
                Some(symbol) => symbol,
                None => "P 1",
            },
        );
    }
}

fn write_atom(
    out: &mut impl fmt::Write,
    structure: &Structure,
    atom: &AtomRef<'_>,
    residue: &ResidueRef<'_>,
    chain: &str,
    serial: i64,
    options: &PdbOptions,
) -> Result<(), Diagnostic> {
    let record = if residue.is_het() { "HETATM" } else { "ATOM  " };
    let name = required_text(
        atom_name(*atom, options.namespace),
        "atom identifier",
        *atom,
        options,
    )?;
    let element = match atom.element() {
        Some(element) => element.symbol(),
        None => "",
    };
    // An atom whose position was never recorded still has a row here, but the
    // format has no way to say "unrecorded" — so it is left out of the file
    // rather than written as an origin the experiment never observed.
    let position = required_value(atom.position(), "coordinates", *atom, options)?;
    let seq = i64::from(required_value(
        residue_sequence(*residue, options.namespace),
        "residue sequence identifier",
        *atom,
        options,
    )?);
    let component = required_text(
        component_name(*atom, *residue, options.namespace),
        "component identifier",
        *atom,
        options,
    )?;
    let occupancy = required_value(atom.occupancy(), "occupancy", *atom, options)?;
    let b_factor = required_value(atom.b_factor(), "B factor", *atom, options)?;

    let _ = writeln!(
        out,
        "{record}{serial:>5} {name:<4}{alt:1}{comp:>3} {chain:>1}{seq}{ins:1}   \
         {x:8.3}{y:8.3}{z:8.3}{occ:6.2}{b:6.2}          {element:>2}",
        serial = serial_field(serial, options),
        name = atom_name_field(name),
        alt = alt_field(structure, atom),
        comp = component,
        seq = residue_field(seq, options),
        ins = insertion_code(*residue),
        x = f64::from(position[0]),
        y = f64::from(position[1]),
        z = f64::from(position[2]),
        occ = f64::from(occupancy),
        b = f64::from(b_factor),
    );
    // A tensor recorded for this atom follows under the same serial and the
    // same identity columns, as the record pairs by serial.
    if let Some(u) = structure.data().anisotropy.for_atom(atom.index()) {
        super::aniso::Record {
            atom,
            structure,
            residue,
            name,
            component,
            chain,
            seq: &residue_field(seq, options),
            serial: &serial_field(serial, options),
            u: &u,
        }
        .write(out);
    }
    Ok(())
}

/// The four-column name field.
///
/// A name shorter than four characters starts in the second column, which is the
/// convention that lets a reader tell a one-letter element from a two-letter one.
pub(crate) fn atom_name_field(name: &str) -> String {
    if name.len() >= 4 {
        name.to_owned()
    } else {
        format!(" {name:<3}")
    }
}

/// The one-column alternate-location field.
///
/// A label longer than one character has already been refused by the capacity
/// check, so anything reaching here fits.
pub(crate) fn alt_field<'a>(structure: &'a Structure, atom: &AtomRef<'a>) -> &'a str {
    let Some(alt) = atom.alt_id() else { return " " };
    let Some(symbol) = alt.symbol() else {
        return " ";
    };
    match structure.resolve(symbol) {
        Some(label) => label,
        None => " ",
    }
}

pub(crate) fn serial_field(serial: i64, options: &PdbOptions) -> String {
    match options
        .hybrid36
        .then(|| hybrid36::encode(serial, 5))
        .flatten()
    {
        Some(field) => field,
        None => format!("{serial:>5}"),
    }
}

pub(crate) fn residue_field(seq: i64, options: &PdbOptions) -> String {
    match options.hybrid36.then(|| hybrid36::encode(seq, 4)).flatten() {
        Some(field) => field,
        None => format!("{seq:>4}"),
    }
}

pub(crate) fn chain_label<'a>(
    chain: &'a ChainRef<'a>,
    namespace: PdbIdentifierNamespace,
) -> &'a str {
    let Some(name) = chain_name(*chain, namespace) else {
        return "";
    };
    name
}

pub(crate) fn insertion_code(residue: ResidueRef<'_>) -> &str {
    let Some(code) = residue.ins_code() else {
        return " ";
    };
    code
}

#[cfg(test)]
#[path = "pdb_tests.rs"]
mod tests;