bentopy 0.2.0a10__cp313-cp313-manylinux_2_34_x86_64.whl

solvate/solvate.rs

@@ -0,0 +1,244 @@
+use eightyseven::structure::ResName;
+use glam::{UVec3, Vec3};
+
+use crate::convert::Convert;
+use crate::cookies::Cookies;
+use crate::placement::PlaceMap;
+use crate::placement::iter_3d;
+use crate::structure::{BoxVecsExtension, Structure};
+use crate::water::Water;
+use crate::{BoundaryMode, PeriodicMode};
+
+/// Solvate a [`Structure`] with a template solvent box.
+#[allow(clippy::too_many_arguments)]
+pub fn solvate<'sol>(
+    structure: &mut Structure,
+    solvent: Water,
+    cutoff: f32,
+    solvent_cutoff: f32,
+    ignore: &[ResName],
+    center: bool,
+    boundary_mode: BoundaryMode,
+    periodic_mode: PeriodicMode,
+) -> PlaceMap<'sol> {
+    // Determine how many copies of the solvent cell are required to fill the input box for each
+    // direction.
+    let cookie_size = solvent.dimensions();
+    let d = structure.boxvecs.as_vec3() / cookie_size;
+    let dimensions = d.ceil().as_uvec3();
+
+    // Depending on the desired behavior, modify the structure.
+    match boundary_mode {
+        BoundaryMode::Cut => {}
+        BoundaryMode::Grow => {
+            eprintln!("\tGrowing output structure box to fit solvent precisely.");
+            // If the solvent box does not fit precisely (viz., the input box size is not an
+            // integer multiple of the solvent box size), the size will be overshot. This means that in
+            // many cases, the size of the output box may be greater than that of the input structure.
+            let remainder = d.fract() * cookie_size;
+            if center {
+                // Add half of the overshot to the positions in the structure to place it in the center of
+                // the final box.
+                let offset = remainder * 0.5;
+                for bead in &mut structure.atoms {
+                    bead.position += offset.convert();
+                }
+
+                eprintln!("\tCentered structure in expanded box.");
+            }
+
+            match &mut structure.boxvecs {
+                // We don't have to cut off the solvent structures at the boundaries, but we have to
+                // change the boundaries to fit.
+                eightyseven::structure::BoxVecs::Short(three) => {
+                    *three = (Vec3::from_array(*three) + remainder).to_array()
+                }
+                eightyseven::structure::BoxVecs::Full(_) => todo!(),
+            }
+        }
+    }
+
+    // Cut the input structure into cell-sized cookies that are ordered in the same manner as the
+    // solvent placement map is.
+    if !ignore.is_empty() {
+        let ignored = ignore.join(", ");
+        eprintln!("\tParticles with the following resnames will be ignored: {ignored}");
+    }
+    eprint!("\tCutting cookies... ");
+    let start = std::time::Instant::now();
+    let cookies = Cookies::new(
+        structure,
+        ignore,
+        cookie_size,
+        dimensions,
+        cutoff,
+        periodic_mode,
+    );
+    eprintln!("Took {:.3} s.", start.elapsed().as_secs_f32());
+
+    // For each location of a solvent box, go through each of the solvent bead positions and see
+    // whether it collides with a structure atom. If it does, we jot that down so we don't write
+    // that solvent bead out later.
+    eprint!("\tChecking solvent collisions... ");
+    let start = std::time::Instant::now();
+    let mut placemap = PlaceMap::new(solvent, dimensions);
+    let cutoff2 = cutoff.powi(2); // Square to avoid taking the square root of the distances.
+    for cell_pos in iter_3d(dimensions) {
+        let cookie = cookies.get(cell_pos).unwrap(); // We are sure there is a cookie here.
+        if cookie.is_empty() {
+            // If the cookie does not contain any structure beads, we can just skip
+            // measuring any distances.
+            continue;
+        }
+        let translation = cookies.offset(cell_pos);
+        // TODO: This is a bit dirty but necessary for the time being because we are iterating over
+        // something that we also need to modify.
+        let p = placemap.all_positions().to_vec();
+        for (idx, &solvent_bead) in p.iter().enumerate() {
+            let solvent_pos = solvent_bead + translation;
+            let mut collision = false;
+            for &cookie_bead in cookie {
+                // TODO: Consider applying this translation as a map before the iter over
+                // the solvent beads?
+                // Translate the unit cell solvent bead position to the reference frame of
+                // the bead from the cookie.
+                if solvent_pos.distance_squared(cookie_bead) < cutoff2 {
+                    collision = true;
+                    break;
+                }
+            }
+            if collision {
+                // This unwrap is safe, since we are sure there is a cell here.
+                let placement = &mut placemap.get_mut(cell_pos).unwrap();
+                placement.set_occupied(idx)
+            }
+        }
+    }
+    eprintln!("Took {:.3} s.", start.elapsed().as_secs_f32());
+
+    // Deal with the boundary conditions.
+    match boundary_mode {
+        BoundaryMode::Cut => {
+            let solvent_cutoff2 = solvent_cutoff.powi(2);
+
+            eprint!("\tCutting solvent to fit box... ");
+            let start = std::time::Instant::now();
+            for axis in [Axis::X, Axis::Y, Axis::Z] {
+                let max = dimensions - 1;
+                let box_dimensions = structure.boxvecs.as_vec3();
+                let (main, asz, bsz) = match axis {
+                    Axis::X => (Vec3::X, dimensions.y, dimensions.z),
+                    Axis::Y => (Vec3::Y, dimensions.x, dimensions.z),
+                    Axis::Z => (Vec3::Z, dimensions.x, dimensions.y),
+                };
+
+                // Set up a 3×3 layer of atoms as a crown on the box dimensions at the `this` box.
+                let limit = axis.select(box_dimensions);
+                let crown: Box<[Vec3]> = axis
+                    .crown(box_dimensions)
+                    .map(|translation| {
+                        placemap
+                            .all_positions()
+                            .iter()
+                            .map(move |&a| a + translation)
+                            // We are only interested in positions that lie at the periodic
+                            // interface. Anything beyond the solvent cutoff distance from that
+                            // interface is out of reach.
+                            .filter(|&position| axis.select(position) <= limit + solvent_cutoff)
+                    })
+                    .into_iter()
+                    .flatten()
+                    .collect();
+
+                let translation = cookies.offset(main.as_uvec3() * max);
+                let this = placemap
+                    .all_positions()
+                    .iter()
+                    .map(|&position| position + translation);
+
+                // We'll make a list of solvent atoms we need to reject by index.
+                let rejected: Box<[_]> = this
+                    .enumerate()
+                    .filter_map(|(idx, bead)| {
+                        let outside = axis.select(bead) > limit;
+                        let too_close = crown
+                            .iter()
+                            .any(|&crown_bead| bead.distance_squared(crown_bead) < solvent_cutoff2);
+                        if outside || too_close {
+                            Some(idx)
+                        } else {
+                            None
+                        }
+                    })
+                    .collect();
+
+                // Apply the rejections to all solvent boxes on this face.
+                for b in 0..bsz {
+                    for a in 0..asz {
+                        let pos = match axis {
+                            Axis::X => UVec3::new(max.x, a, b),
+                            Axis::Y => UVec3::new(a, max.y, b),
+                            Axis::Z => UVec3::new(a, b, max.z),
+                        };
+                        let mut pm = placemap.get_mut(pos).unwrap();
+                        for &idx in &rejected {
+                            pm.set_occupied(idx)
+                        }
+                    }
+                }
+            }
+            eprintln!("Took {:.6} s.", start.elapsed().as_secs_f32());
+        }
+        // We already set the box size at the top. Nothing to do here.
+        BoundaryMode::Grow => {}
+    }
+
+    placemap
+}
+
+#[derive(Clone, Copy)]
+enum Axis {
+    X,
+    Y,
+    Z,
+}
+
+impl Axis {
+    const fn select(&self, v: Vec3) -> f32 {
+        v.to_array()[*self as usize]
+    }
+
+    const fn roll(&self, mut v: Vec3) -> Vec3 {
+        let mut n = 0;
+        while n < *self as usize {
+            v = Vec3 {
+                x: v.z,
+                y: v.x,
+                z: v.y,
+            };
+            n += 1;
+        }
+        v
+    }
+
+    const fn crown(&self, box_dimensions: Vec3) -> [Vec3; 9] {
+        let d = self.select(box_dimensions);
+        let mut crown = [
+            Vec3::new(d, -1.0, -1.0),
+            Vec3::new(d, 0.0, -1.0),
+            Vec3::new(d, 1.0, -1.0),
+            Vec3::new(d, -1.0, 0.0),
+            Vec3::new(d, 0.0, 0.0),
+            Vec3::new(d, 1.0, 0.0),
+            Vec3::new(d, -1.0, 1.0),
+            Vec3::new(d, 0.0, 1.0),
+            Vec3::new(d, 1.0, 1.0),
+        ];
+        let mut i = 0;
+        while i < crown.len() {
+            crown[i] = self.roll(crown[i]);
+            i += 1;
+        }
+        crown
+    }
+}

solvate/structure.rs

@@ -0,0 +1,160 @@
+use std::io;
+
+use eightyseven::structure::BoxVecs;
+use eightyseven::writer::WriteGro;
+use glam::Vec3;
+use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
+
+use crate::placement::PlaceMap;
+use crate::substitute::Substitution;
+
+pub type Structure = eightyseven::structure::Structure;
+
+pub trait BoxVecsExtension {
+    fn as_vec3(&self) -> glam::Vec3;
+}
+
+impl BoxVecsExtension for BoxVecs {
+    fn as_vec3(&self) -> Vec3 {
+        match self {
+            // FIX: When a gro file with cuboid boundaries represented as a nine-valued vector
+            // (trailing six values are zero) is loaded, we can canonicalize it as a three-valued
+            // cuboid. This should probably be fixed in eightyseven, but I'll want to reconsider
+            // this and related issues properly at another moment.
+            &BoxVecs::Short(three) | &BoxVecs::Full([three @ .., 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) => {
+                Vec3::from_array(three)
+            }
+            BoxVecs::Full(nine) => todo!("nine-value boxvecs ({nine:?}) are currently unsupported"),
+        }
+    }
+}
+
+pub fn write_structure<const PAR: bool>(
+    writer: &mut impl io::Write,
+    structure: &Structure,
+    solvent_placemap: &PlaceMap,
+    substitutions: &[Substitution],
+    buffer_size: usize,
+) -> io::Result<()> {
+    let solvent = solvent_placemap.solvent;
+    let natoms_structure = structure.natoms();
+    let natoms_per_residue = solvent.residue_points();
+    let natoms_solvent = solvent_placemap.unoccupied_count() as usize * natoms_per_residue;
+    let natoms_substitutes = substitutions.iter().map(|s| s.natoms()).sum::<usize>();
+    let natoms = natoms_structure + natoms_solvent + natoms_substitutes;
+
+    // Write the gro header.
+    writeln!(
+        writer,
+        "solvated by {} (v{})",
+        env!("CARGO_BIN_NAME"),
+        bentopy::core::version::VERSION
+    )?;
+    writeln!(writer, "{natoms}")?;
+
+    // First, we write the original structure.
+    let start = std::time::Instant::now();
+    eprint!("Writing original structure ({natoms_structure:>9}/{natoms} total atoms)... ");
+    for line in structure.format_atom_lines_iter() {
+        writer.write_all(line.as_bytes())?;
+    }
+    eprintln!("Done, took {:.3} s.", start.elapsed().as_secs_f32());
+
+    // Now for the solvent. We will go over the cells described by the placemap, and format and
+    // write the atom lines per cell. This means that we can avoid creating the whole solvent
+    // structure at once, which would be tremendously memory intensive.
+    let start = std::time::Instant::now();
+    eprintln!("Writing solvent structure  ({natoms_solvent:>9}/{natoms} total atoms)...");
+    let mut placements = solvent_placemap.iter_atoms_chunks();
+    let mut particle_buffer = Vec::new();
+    let mut n = 0;
+    'write_atoms: loop {
+        particle_buffer.clear();
+
+        // TODO: I think it's scary that buffer_size could feasibly be zero.
+        'fill_buffer: while particle_buffer.len() < buffer_size {
+            match placements.next() {
+                Some(sp) => particle_buffer.extend(sp),
+                None if particle_buffer.is_empty() => break 'write_atoms,
+                None => break 'fill_buffer,
+            }
+        }
+
+        if PAR {
+            let lines: String = particle_buffer
+                .par_iter()
+                .map(Structure::format_atom_line)
+                .collect();
+            writer.write_all(lines.as_bytes())?;
+        } else {
+            // TODO: Even this could be done with a buffer of lines that fills up for a while,
+            // and then we write it. So non-parallel but still buffered. But in a way that is
+            // equivalent to making sure we use a BufWriter. Maybe first benchmark if requiring a
+            // BufWriter in the function signature (we already use it at the call site) makes a
+            // difference (it shouldn't because the impl Write makes it generic over our
+            // BufWriter). Only then try to see if rolling our own helps out.
+            for atom in &particle_buffer {
+                let line = Structure::format_atom_line(atom);
+                writer.write_all(line.as_bytes())?;
+            }
+        }
+
+        // Report the progress.
+        n += particle_buffer.len();
+        let progress = n as f32 / natoms_solvent as f32;
+        let percentage = progress * 100.0;
+        let delta = std::time::Instant::now() - start;
+        let time_left = delta.as_secs_f32() * (progress.recip() - 1.0);
+        eprint!(
+            "\r({percentage:>4.1}%) Wrote {n:>9}/{natoms_solvent} solvent atoms. (ETA {time_left:>4.0} s) "
+        );
+    }
+    eprintln!();
+    eprintln!(
+        "Done writing solvent atoms, took {:.3} s.",
+        start.elapsed().as_secs_f32()
+    );
+
+    if !substitutions.is_empty() {
+        // Write out the substitutes at the end of the structure.
+        let start = std::time::Instant::now();
+        eprintln!("Writing substitutes        ({natoms_substitutes:>9}/{natoms} total atoms)...");
+        for substitution in substitutions {
+            let name = substitution.name();
+            let number = substitution.natoms();
+            eprint!("\t{name}, {number} atoms... ");
+            let start = std::time::Instant::now();
+            let mut placements = substitution.iter_atoms();
+            loop {
+                particle_buffer.clear();
+                particle_buffer.extend(placements.by_ref().take(buffer_size));
+                if particle_buffer.is_empty() {
+                    break;
+                }
+
+                let beads = &particle_buffer;
+                if PAR {
+                    let lines: String = beads.par_iter().map(Structure::format_atom_line).collect();
+                    writer.write_all(lines.as_bytes())?;
+                } else {
+                    for atom in beads {
+                        let line = Structure::format_atom_line(atom);
+                        writer.write_all(line.as_bytes())?;
+                    }
+                }
+            }
+            eprintln!("Took {:.3} s.", start.elapsed().as_secs_f32());
+        }
+        eprintln!(
+            "Done writing substitutes, took {:.3} s.",
+            start.elapsed().as_secs_f32()
+        );
+    }
+
+    // Finally, we write the box vectors.
+    writeln!(writer, "{}", structure.boxvecs())?;
+
+    writer.flush()?;
+
+    Ok(())
+}

solvate/substitute.rs

@@ -0,0 +1,113 @@
+use eightyseven::structure::{Atom, AtomName};
+
+use crate::{args::Substitute, placement::PlaceMap};
+
+pub struct Substitution<'sol> {
+    name: AtomName,
+    placemap: PlaceMap<'sol>,
+}
+
+impl<'sol> Substitution<'sol> {
+    pub fn new(name: impl Into<AtomName>, placemap: PlaceMap<'sol>) -> Self {
+        Self {
+            name: name.into(),
+            placemap,
+        }
+    }
+
+    pub fn natoms(&self) -> usize {
+        self.placemap.unoccupied_count() as usize
+    }
+
+    pub fn name(&self) -> &str {
+        &self.name
+    }
+
+    /// Iterate over all atoms in this [`Substitution`].
+    pub fn iter_atoms(&self) -> impl Iterator<Item = Atom> + '_ {
+        self.placemap.iter_positions().enumerate().map(|(i, pos)| {
+            let num = i as u32 + 1;
+            Atom {
+                resnum: num,
+                resname: self.name,
+                atomname: self.name,
+                atomnum: num,
+                position: pos,
+                velocity: Default::default(),
+            }
+        })
+    }
+
+    /// Fuse the contents of two [`Substitution`] maps.
+    pub fn glue(&mut self, substitution: &Substitution<'sol>) {
+        self.placemap &= substitution.placemap.clone();
+    }
+}
+
+pub fn substitute<'sol>(
+    rng: &mut rand::rngs::StdRng,
+    placemap: &mut PlaceMap<'sol>,
+    substitutes: &[Substitute],
+) -> Vec<Substitution<'sol>> {
+    use rand::seq::SliceRandom;
+
+    // FIXME: Make this into a nicer error. Already explains the situation well, but the assert
+    // failure is ugly in the output.
+    let n_free = placemap.unoccupied_count();
+    let n_subs = substitutes.iter().map(|s| s.number).sum::<u64>();
+    assert!(
+        n_subs <= n_free,
+        "the number of substitutions that are specified ({n_subs}) exceeds the number of \
+            solvent positions that can be substituted ({n_free})"
+    );
+
+    let mut substitutions = Vec::new();
+    let mut zebra = Vec::new();
+    for substitute in substitutes {
+        eprint!(
+            "\tSubstituting {} solvent locations with {}... ",
+            substitute.number, substitute.name
+        );
+        let start = std::time::Instant::now();
+        if substitute.number == 0 {
+            eprintln!("Skipping.");
+            continue;
+        }
+
+        let n_free = placemap.unoccupied_count();
+        if n_free == 0 {
+            eprintln!("No free solvent spots left.");
+            break;
+        }
+
+        zebra.clear();
+        zebra.resize(n_free as usize, false);
+        zebra[..substitute.number as usize].fill(true);
+        zebra.shuffle(rng);
+        let mut zebra = zebra.iter();
+
+        let mut old_placemap = placemap.clone(); // Becomes the map for our substitute.
+        for mut placement in placemap.iter_placements_mut() {
+            for i in 0..placement.len() {
+                let is_occupied = placement.get(i);
+                // If this position is not already occupied by a non-solvent bead, and if our
+                // shuffled values indicate that we need to place a substitute here, we set this
+                // bead to occupied in the solvent placemap.
+                if !is_occupied && *zebra.next().unwrap() {
+                    placement.set_occupied(i)
+                }
+            }
+        }
+        assert_eq!(zebra.len(), 0);
+
+        old_placemap ^= placemap.clone();
+        let substitute_placemap = !old_placemap;
+        substitutions.push(Substitution::new(
+            substitute.name.as_str(),
+            substitute_placemap,
+        ));
+        eprintln!("Took {:.3} s.", start.elapsed().as_secs_f32());
+    }
+
+    substitutions
+}