bentopy 0.2.0a10__cp313-cp313-manylinux_2_34_x86_64.whl

render/render.rs

@@ -0,0 +1,393 @@
+use std::io::{self, Read};
+use std::path::PathBuf;
+
+use anyhow::{Context, Result};
+use bentopy::core::placement::{Placement, PlacementList};
+use bentopy::core::utilities::CLEAR_LINE;
+use glam::Vec3;
+use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
+
+use crate::args::{Args, Mode, ResnumMode};
+use crate::limits::Limits;
+use crate::structure::{self, Atom, Molecule, rotate_molecule};
+
+/// Read a placement list to return [`Placements`].
+fn read_placement_list(
+    placement_list: &str,
+    root: Option<PathBuf>,
+) -> serde_json::Result<PlacementList> {
+    let mut placements: PlacementList = serde_json::from_str(placement_list)?;
+    if let Some(root) = root {
+        for p in &mut placements.placements {
+            if p.path.is_relative() {
+                let mut path = root.clone();
+                path.push(&p.path);
+                p.path = path;
+            }
+        }
+    }
+    Ok(placements)
+}
+
+impl Mode {
+    /// Based on the [`Mode`], prepare the list of molecules.
+    ///
+    /// If necessary, a molecule is loaded from its structure path, centered, possible its atoms
+    /// are filtered down to a subset, or its residues or the whole molecule itself is represented
+    /// by single particles.
+    ///
+    /// The goal is to do as little work as necessary for `Mode`s that aim to make lightweight
+    /// representations for inspection.
+    fn prepare_molecules(
+        &self,
+        placements: &[Placement],
+        ignore_tags: bool,
+        verbose: bool,
+    ) -> Result<Vec<Molecule>> {
+        let load_molecule = |path| -> Result<_> {
+            let prefix = if verbose { "" } else { CLEAR_LINE };
+            let suffix = if verbose { "\n" } else { "\r" };
+            eprint!("{prefix}\tLoading {path:?}...{suffix}");
+            let mut molecule = structure::load_molecule(path)?;
+            molecule.translate_to_center();
+            Ok(molecule)
+        };
+        let apply_tag = |molecule: &mut Molecule, tag: Option<&str>| {
+            if !ignore_tags && let Some(tag) = tag {
+                molecule.apply_tag(tag)
+            }
+        };
+
+        match self {
+            Mode::Full => placements
+                .iter()
+                .map(|p| {
+                    let mut molecule = load_molecule(&p.path)?;
+                    apply_tag(&mut molecule, p.tag());
+                    Ok(molecule)
+                })
+                .collect(),
+            Mode::Backbone => placements
+                .iter()
+                .map(|p| {
+                    let mut molecule = load_molecule(&p.path)?;
+                    apply_tag(&mut molecule, p.tag());
+                    molecule
+                        .atoms
+                        .retain(|a| ["N", "CA", "C", "O"].contains(&a.name.as_str()));
+                    Ok(molecule)
+                })
+                .collect(),
+            Mode::Alpha => placements
+                .iter()
+                .map(|p| {
+                    let mut molecule = load_molecule(&p.path)?;
+                    apply_tag(&mut molecule, p.tag());
+                    molecule.atoms.retain(|a| a.name.as_str() == "CA");
+                    Ok(molecule)
+                })
+                .collect(),
+            Mode::Residue => placements
+                .iter()
+                .map(|p| {
+                    let mut molecule = load_molecule(&p.path)?;
+                    let mut residues = Vec::new();
+                    let mut residue_atoms = Vec::new();
+                    for atom in molecule.atoms {
+                        let Some(last) = residue_atoms.last() else {
+                            residue_atoms.push(atom);
+                            continue;
+                        };
+                        if last.resnum == atom.resnum {
+                            residue_atoms.push(atom);
+                        } else {
+                            // The residue we just collected is complete.
+                            let avg_pos =
+                                residue_atoms.iter().fold(Vec3::ZERO, |acc, a| acc + a.pos)
+                                    / residue_atoms.len() as f32;
+                            residues.push(Atom {
+                                name: "DUMMY".into(),
+                                num: residues.len() as u32,
+                                pos: avg_pos,
+                                ..last.clone()
+                            });
+
+                            // Prepare for the next resdue.
+                            residue_atoms.clear();
+                            residue_atoms.push(atom);
+                        }
+                    }
+                    molecule.atoms = residues;
+                    apply_tag(&mut molecule, p.tag());
+                    molecule.translate_to_center();
+                    Ok(molecule)
+                })
+                .collect(),
+            Mode::Instance => Ok(placements
+                .iter()
+                .map(|p| {
+                    let mut atom = Atom::dummy();
+                    if let Some(tag) = p.tag() {
+                        atom.resname = tag.into();
+                    }
+                    Molecule { atoms: vec![atom] }
+                })
+                .collect()),
+        }
+    }
+}
+
+/// Write out [`Placements`] as a `gro` file.
+///
+/// See the Gromacs manual entry on the [gro file][gro_manual].
+///
+/// Note that each [`Placement`] (i.e., kind of structure) is given a different `resnum`, such that
+/// the different kinds of structures can be targeted individually in programmes that use this
+/// `gro` file output. For example, each kind of molecule may be drawn in a different color in a
+/// molecular visualization program.
+///
+/// [gro_manual]: https://manual.gromacs.org/archive/5.0.3/online/gro.html
+fn write_gro(
+    writer: &mut impl io::Write,
+    placements: &PlacementList,
+    limits: Limits,
+    mode: Mode,
+    resnum_mode: ResnumMode,
+    ignore_tags: bool,
+    verbose: bool,
+) -> Result<Box<[(String, usize)]>> {
+    let min_limits = Vec3::new(
+        limits.minx.unwrap_or_default(),
+        limits.miny.unwrap_or_default(),
+        limits.minz.unwrap_or_default(),
+    );
+    let size = limits.box_size(Vec3::from(placements.size));
+    let placements = &placements.placements;
+    // Load the molecules and center them with respect to themselves.
+    let molecules = mode.prepare_molecules(placements, ignore_tags, verbose)?;
+    writeln!(
+        writer,
+        "{} (v{})",
+        env!("CARGO_BIN_NAME"),
+        bentopy::core::version::VERSION
+    )?;
+    let placed = placements.iter().zip(&molecules).map(|(p, m)| {
+        let n_placed = p
+            .batches
+            .iter()
+            .map(|b| {
+                b.positions
+                    .iter()
+                    .filter(|&&pos| limits.is_inside(pos))
+                    .count()
+            })
+            .sum::<usize>();
+        let n_atoms = n_placed * m.atoms.len();
+        let name = p.name.clone();
+        (name, n_placed, n_atoms)
+    });
+    let n_atoms: usize = placed.clone().map(|(_, _, n)| n).sum();
+    let placed = placed.map(|(name, n_placed, _)| (name, n_placed)).collect();
+    let mut instance_resnum = 0;
+    writeln!(writer, "{n_atoms}")?;
+    for (segidx, (placement, molecule)) in placements.iter().zip(molecules).enumerate() {
+        {
+            let prefix = if verbose { "" } else { CLEAR_LINE };
+            let suffix = if verbose { "\n" } else { "\r" };
+            let name = &placement.name;
+            let tag = placement
+                .tag()
+                .map(|tag| format!(":{tag}"))
+                .unwrap_or_default();
+            let path = &placement.path;
+            eprint!(
+                "{prefix}\tNow writing '{name}{tag}' from {path:?} data to output file.{suffix}",
+            );
+        }
+        // Make sure that the printed resnum cannot exceed the 5 allotted columns.
+        let segidx = segidx % 100000;
+
+        let mut inserts = Vec::new();
+        for atom in &molecule.atoms {
+            // HACK: For some reason, converting to string here is necessary for getting correct
+            // alignment of ArrayString<U5>.
+            let resnum = atom.resnum;
+            let resname = atom.resname.to_string();
+            let atomname = atom.name.to_string();
+            let atomnum = atom.num;
+            let insert = match resnum_mode {
+                ResnumMode::Instance => format!("{resname:<5}{atomname:>5}{atomnum:>5}"),
+                ResnumMode::Segment => format!("{segidx:>5}{resname:<5}{atomname:>5}{atomnum:>5}"),
+                ResnumMode::Keep => format!("{resnum:>5}{resname:<5}{atomname:>5}{atomnum:>5}"),
+            };
+            inserts.push(insert)
+        }
+
+        let mut output = String::new();
+        let t0 = std::time::Instant::now();
+        for batch in &placement.batches {
+            // This tends to take around half a ms.
+            let molecule = rotate_molecule(&molecule, &batch.rotation());
+
+            // Move over the positions so that they do not float somewhere in space, far from the
+            // origin, in case the limits are set. This puts them within the specified limits.
+            let min = molecule.min() + min_limits;
+            let included_positions: Vec<_> = batch
+                .positions
+                .iter()
+                .filter(|&&position| limits.is_inside(position))
+                .collect();
+            let n_included_instances = included_positions.len();
+            let out: String = included_positions
+                .par_iter()
+                .enumerate()
+                .map(|(batch_instance_resnum, position)| {
+                    let mut output = String::new();
+                    for (atom, insert) in molecule.atoms.iter().zip(&inserts) {
+                        let mut pos = atom.pos;
+                        pos.x += position[0] - min.x;
+                        pos.y += position[1] - min.y;
+                        pos.z += position[2] - min.z;
+
+                        match resnum_mode {
+                            ResnumMode::Instance => {
+                                let resnum = instance_resnum + batch_instance_resnum;
+                                // Make sure that the printed resnum cannot exceed the 5 allotted columns.
+                                let resnum = resnum % 100000;
+                                output.push_str(&format!("{resnum:>5}"))
+                            }
+                            ResnumMode::Segment => {}
+                            ResnumMode::Keep => {}
+                        };
+                        output.push_str(insert);
+                        output.push_str(&format!("{:8.3}{:8.3}{:8.3}\n", pos.x, pos.y, pos.z));
+                    }
+                    output
+                })
+                .collect();
+            instance_resnum += n_included_instances;
+            output.push_str(&out);
+        }
+        if verbose {
+            eprintln!("\t\tFormatting took {:.3} s.", t0.elapsed().as_secs_f32());
+        }
+
+        let t0 = std::time::Instant::now();
+        writer.write_all(output.as_bytes())?;
+        writer.flush()?;
+        if verbose {
+            eprintln!("\t\tWriting took    {:.3} s.", t0.elapsed().as_secs_f32());
+        }
+    }
+
+    // Write the box vectors.
+    let [v1x, v2y, v3z] = size.to_array();
+    writeln!(writer, "{v1x:.3} {v2y:.3} {v3z:.3}")?;
+
+    Ok(placed)
+}
+
+/// Write out the topology as a `top` file.
+///
+/// See the Gromacs manual entry on the [top file][top_manual].
+///
+/// [top_manual]: https://manual.gromacs.org/archive/5.0/online/top.html
+fn write_top(
+    writer: &mut impl io::Write,
+    includes: &[String],
+    title: &str,
+    placed: &[(String, usize)],
+) -> io::Result<()> {
+    // Write some meta stuff.
+    writeln!(
+        writer,
+        "; This topology file was generated by {} (v{}).",
+        env!("CARGO_BIN_NAME"),
+        bentopy::core::version::VERSION
+    )?;
+    writeln!(writer, "; {}", env!("CARGO_PKG_AUTHORS"))?;
+    writeln!(writer, "; {}", env!("CARGO_PKG_REPOSITORY"))?;
+    writeln!(writer)?;
+
+    // Dump the includes.
+    for include in includes {
+        writeln!(writer, "#include \"{include}\"")?;
+    }
+    writeln!(writer)?;
+
+    // Write the system title.
+    writeln!(writer, "[ system ]")?;
+    writeln!(writer, "{title}")?;
+    writeln!(writer)?;
+
+    // Write out the name of each structure and how many of them were written to the output file.
+    writeln!(writer, "[ molecules ]")?;
+    for (name, count) in placed {
+        if *count > 0 {
+            writeln!(writer, "{name}\t{count}")?;
+        }
+    }
+
+    Ok(())
+}
+
+/// Render a placement list to a gro file.
+///
+/// If `input_path` is "-", the placement list is read from stdin.
+pub fn render(args: Args) -> anyhow::Result<()> {
+    let Args {
+        input: input_path,
+        output: output_path,
+        topol: topol_path,
+        root,
+        limits,
+        mode,
+        resnum_mode,
+        ignore_tags,
+        verbose,
+    } = args;
+
+    eprint!("Reading from {input_path:?}... ");
+    let t0 = std::time::Instant::now();
+    let mut placement_list = String::new();
+    if input_path.to_str() == Some("-") {
+        std::io::stdin().read_to_string(&mut placement_list)?;
+    } else {
+        std::fs::File::open(&input_path)?.read_to_string(&mut placement_list)?;
+    }
+    let placements = read_placement_list(&placement_list, root)
+        .with_context(|| format!("Failed to read placement list from {input_path:?}"))?;
+    eprintln!("Done in {:.3} ms.", t0.elapsed().as_millis());
+
+    eprintln!("Writing structure to {output_path:?}... ");
+    let t0 = std::time::Instant::now();
+    let outfile = std::fs::File::create(&output_path)?;
+    let placed = write_gro(
+        &mut std::io::BufWriter::new(outfile),
+        &placements,
+        limits.unwrap_or_default(),
+        mode,
+        resnum_mode,
+        ignore_tags,
+        verbose,
+    )?;
+    eprintln!("{CLEAR_LINE}Done in {:.3} s.", t0.elapsed().as_secs_f32());
+    eprintln!(
+        "Wrote {} placements of {} different kinds of structures.",
+        placed.iter().map(|(_, c)| c).sum::<usize>(),
+        placed.iter().filter(|(_, c)| *c > 0).count(),
+    );
+
+    if let Some(topol_path) = topol_path {
+        let mut top = std::fs::File::create(&topol_path)?;
+        write_top(
+            &mut top,
+            &placements.topol_includes,
+            &placements.title,
+            &placed,
+        )?;
+        eprintln!("Wrote topology to {topol_path:?}.");
+    }
+
+    Ok(())
+}

render/structure.rs

@@ -0,0 +1,264 @@
+use std::{io::Read, path::Path};
+
+use anyhow::{Context, Result};
+use eightyseven::structure::{AtomName, AtomNum, ResName, ResNum};
+use glam::{Mat3, Vec3};
+
+/// A molecule as derived from PDB 'ATOM' records.
+#[derive(Clone)]
+pub struct Molecule {
+    pub atoms: Vec<Atom>,
+}
+
+impl Molecule {
+    /// Create a new molecule from a PDB file.
+    fn from_pdb(pdb: &str) -> Result<Self> {
+        let mut atoms = Vec::new();
+        for (ln, line) in pdb.lines().enumerate() {
+            if line.starts_with("ATOM") || line.starts_with("HETATM") {
+                let ln = ln + 1;
+                let atom = Atom::from_pdb_atom_line(line)
+                    .context(format!("could not parse atom on line {ln}"))?;
+                atoms.push(atom);
+            }
+        }
+
+        Ok(Self { atoms })
+    }
+
+    /// Create a new molecule from a gro file.
+    fn from_gro(gro: &str) -> Result<Self> {
+        let mut atoms = Vec::new();
+        let mut lines = gro.lines().enumerate();
+        let (_ln, _title) = lines.next().context("expected a title line")?;
+        let n_atoms: usize = lines
+            .next()
+            .context("expected the number of atoms")?
+            .1
+            .trim()
+            .parse()
+            .context("could not parse the number of atoms")?;
+        // Read the atoms.
+        for (ln, line) in lines.take(n_atoms) {
+            let ln = ln + 1;
+            let atom = Atom::from_gro_atom_line(line)
+                .context(format!("could not parse atom on line {ln}"))?;
+            atoms.push(atom);
+        }
+        // We don't check for the presence and correctness of the box vectors, even though they
+        // should be there. We just don't care here.
+
+        Ok(Self { atoms })
+    }
+
+    /// Apply a `tag` to the atoms in this [`Molecule`].
+    ///
+    /// This will replace the `resname` for each atom with the provided `tag`.
+    pub fn apply_tag(&mut self, tag: impl Into<ResName>) {
+        let tag = tag.into();
+        for Atom { resname, .. } in &mut self.atoms {
+            *resname = tag;
+        }
+    }
+
+    /// Returns the center of this [`Molecule`].
+    ///
+    /// Here, center refers to mean position. This can be understood as a center of mass where all
+    /// atoms are considered to have the same weight.
+    pub fn center(&self) -> Vec3 {
+        if self.atoms.is_empty() {
+            return Vec3::ZERO;
+        }
+        let mut mean = Vec3::ZERO;
+        for atom in &self.atoms {
+            mean += atom.pos
+        }
+        mean / self.atoms.len() as f32
+    }
+
+    /// Translate all atoms such that their center lies at the origin.
+    ///
+    /// Here, center refers to mean position. This can be understood as a center of mass where all
+    /// atoms are considered to have the same weight.
+    pub fn translate_to_center(&mut self) {
+        let center = self.center();
+        for atom in &mut self.atoms {
+            atom.pos -= center;
+        }
+    }
+
+    /// Returns the minimum _x_, _y_, and _z_ values of this [`Molecule`] as a [`Vec3`].
+    ///
+    /// Note that this is a virtual point, not an actual point in the atoms point cloud.
+    /// In other words, it is not the lowest point in a sorted list of existing positions.
+    pub fn min(&self) -> Vec3 {
+        let mut min = Vec3::ZERO;
+        for atom in &self.atoms {
+            let pos = atom.pos;
+            if pos.x < min.x {
+                min.x = pos.x
+            }
+            if pos.y < min.y {
+                min.y = pos.y
+            }
+            if pos.z < min.z {
+                min.z = pos.z
+            }
+        }
+        min
+    }
+}
+
+/// Assumes the [Molecule]'s atoms have been centered.
+pub fn rotate_molecule(molecule: &Molecule, rotmat: &Mat3) -> Molecule {
+    let mut molecule = molecule.clone();
+    for atom in &mut molecule.atoms {
+        atom.pos = rotmat.mul_vec3(atom.pos)
+    }
+
+    molecule
+}
+
+/// The representation of an atom.
+#[derive(Clone)]
+pub struct Atom {
+    /// Atom name.
+    pub name: AtomName,
+    /// Residue name.
+    pub resname: ResName,
+    /// Residue sequence number.
+    pub resnum: ResNum,
+    /// Number of the atom within its structure.
+    ///
+    /// Also understood as its 'serial' number.
+    pub num: AtomNum,
+    /// Position in nanometers.
+    pub pos: Vec3,
+}
+
+impl Atom {
+    // COLUMNS        DATA  TYPE    FIELD        DEFINITION
+    // -------------------------------------------------------------------------------------
+    //  1 -  6        Record name   "ATOM  "
+    //  7 - 11        Integer       serial       Atom  serial number.
+    // 13 - 16        Atom          name         Atom name.
+    // 17             Character     altLoc       Alternate location indicator.
+    // 18 - 20        Residue name  resName      Residue name.
+    // 22             Character     chainID      Chain identifier.
+    // 23 - 26        Integer       resSeq       Residue sequence number.
+    // 27             AChar         iCode        Code for insertion of residues.
+    // 31 - 38        Real(8.3)     x            Orthogonal coordinates for X in Angstroms.
+    // 39 - 46        Real(8.3)     y            Orthogonal coordinates for Y in Angstroms.
+    // 47 - 54        Real(8.3)     z            Orthogonal coordinates for Z in Angstroms.
+    // 55 - 60        Real(6.2)     occupancy    Occupancy.
+    // 61 - 66        Real(6.2)     tempFactor   Temperature  factor.
+    // 77 - 78        LString(2)    element      Element symbol, right-justified.
+    // 79 - 80        LString(2)    charge       Charge  on the atom.
+    //
+    // Examples:
+    // ATOM      1  N   ALA A   1      11.104   6.134  -6.504
+    // ATOM      1  N   MET     1      48.048  69.220  58.803
+    // ATOM      2  C13 DPPCM 367      31.671 -46.874  39.426  1.00  0.00      MEMB
+    /// Read a single "ATOM" or "HETATM" record from a PDB and return an Atom.
+    fn from_pdb_atom_line(line: &str) -> Result<Atom> {
+        let serial = line[6..11]
+            .trim()
+            .parse()
+            .context("could not parse atom serial")?;
+        let name = line[12..16].trim().into();
+        // NOTE: Even though the PDB specification only regards columns 18..21 as constituting the
+        // resname, in practice the character directly after that is also included. This column is
+        // not defined by the spec. Especially for telling apart lipids like DPPC and DPPG, it's
+        // quite important to include that by-convention resname character. Thanks Abby!
+        let resname = line[17..21].trim().into();
+        let resnum = line[22..26]
+            .trim()
+            .parse()
+            .context("could not parse atom resnum")?;
+        let x = line[30..38]
+            .trim()
+            .parse()
+            .context("could not parse x coordinate")?;
+        let y = line[38..46]
+            .trim()
+            .parse()
+            .context("could not parse y coordinate")?;
+        let z = line[46..54]
+            .trim()
+            .parse()
+            .context("could not parse z coordinate")?;
+        Ok(Atom {
+            name,
+            resname,
+            resnum,
+            num: serial,
+            pos: Vec3::new(x, y, z) / 10.0, // Convert from Å to nm.
+        })
+    }
+
+    // `    2WATER  HW3    6   1.326   0.120   0.568  1.9427 -0.8216 -0.0244`
+    /// Read a single atom line from a gro file and return an Atom.
+    fn from_gro_atom_line(line: &str) -> Result<Atom> {
+        let resnum = line[0..5]
+            .trim()
+            .parse()
+            .context("could not parse resnum")?;
+        let resname = line[5..10].trim().into();
+        let name = line[10..15].trim().into(); // Atom name.
+        let num = line[15..20]
+            .trim()
+            .parse()
+            .context("could not parse atomnum")?; // Atom number.
+        let x = line[20..28]
+            .trim()
+            .parse()
+            .context("could not parse x coordinate")?;
+        let y = line[28..36]
+            .trim()
+            .parse()
+            .context("could not parse y coordinate")?;
+        let z = line[36..44]
+            .trim()
+            .parse()
+            .context("could not parse z coordinate")?;
+        Ok(Atom {
+            name,
+            resname,
+            resnum,
+            num,
+            pos: Vec3::new(x, y, z), // Values are already in nm.
+        })
+    }
+
+    /// Create a dummy atom.
+    pub fn dummy() -> Self {
+        Self {
+            name: "DUMMY".into(),
+            resname: "DUMMY".into(),
+            resnum: 0,
+            num: 0,
+            pos: Vec3::ZERO,
+        }
+    }
+}
+
+/// Load a [`Molecule`] from a pdb file.
+pub fn load_molecule<P: AsRef<Path> + std::fmt::Debug>(path: P) -> Result<Molecule> {
+    let mut data = String::new();
+    std::fs::File::open(&path)?.read_to_string(&mut data)?;
+
+    let molecule =
+        match path.as_ref().extension().and_then(|s| s.to_str()) {
+            Some("gro") => Molecule::from_gro(&data)
+                .context(format!("could not parse .gro file at {path:?}"))?,
+            Some("pdb") => Molecule::from_pdb(&data)
+                .context(format!("could not parse .pdb file at {path:?}"))?,
+            None | Some(_) => {
+                eprintln!("WARNING: Assuming {path:?} is a pdb file.");
+                Molecule::from_pdb(&data)
+                    .context(format!("could not parse file at {path:?} as .pdb"))?
+            }
+        };
+
+    Ok(molecule)
+}