bentopy 0.2.0a10__cp313-cp313-manylinux_2_34_x86_64.whl

solvate/main.rs

@@ -0,0 +1,177 @@
+use std::cmp::Reverse;
+use std::io::{self, Write};
+
+use anyhow::Context;
+use args::SortBehavior;
+use clap::Parser;
+use eightyseven::reader::ReadGro;
+use rand::SeedableRng;
+
+use crate::args::Args;
+pub(crate) use crate::args::{BoundaryMode, PeriodicMode};
+use crate::solvate::solvate;
+use crate::structure::{BoxVecsExtension, Structure, write_structure};
+use crate::substitute::substitute;
+
+mod args;
+mod convert;
+mod cookies;
+mod placement;
+mod solvate;
+mod structure;
+mod substitute;
+mod water;
+
+fn main() -> anyhow::Result<()> {
+    let config = Args::parse();
+    let cutoff = config.cutoff.unwrap_or(config.water_type.default_cutoff());
+    let solvent_cutoff = config
+        .solvent_cutoff
+        .unwrap_or(config.water_type.default_solvent_cutoff());
+    eprintln!("Solvent-structure cutoff is set to {cutoff} nm.");
+    eprintln!("Solvent-solvent cutoff is set to   {solvent_cutoff} nm.");
+
+    eprint!("Loading structure {:?}... ", config.input);
+    let start = std::time::Instant::now();
+    let input_path = &config.input;
+    let mut structure = Structure::open_gro(input_path)
+        .with_context(|| format!("Failed to open structure {input_path:?}"))?;
+    eprintln!("Took {:.3} s.", start.elapsed().as_secs_f32());
+
+    eprintln!("Solvating...");
+    let start = std::time::Instant::now();
+    let solvent = config.water_type.into();
+    let mut placemap = solvate(
+        &mut structure,
+        solvent,
+        cutoff,
+        solvent_cutoff,
+        &config.ignore,
+        config.center,
+        config.boundary_mode,
+        config.periodic_mode,
+    );
+    eprintln!("Took {:.3} s.", start.elapsed().as_secs_f32());
+
+    let volume = structure
+        .boxvecs
+        .as_vec3()
+        .as_dvec3()
+        .to_array()
+        .iter()
+        .product();
+    let neutralizing_ions = config.charge.and_then(|c| c.bake());
+    let substitutes = config
+        .substitutes
+        .into_iter()
+        .map(|sc| sc.bake(volume, placemap.unoccupied_count() as u64))
+        .chain(neutralizing_ions)
+        .collect::<Vec<_>>();
+
+    let substitutions = if !substitutes.is_empty() {
+        eprintln!("Making substitutions...");
+        let start = std::time::Instant::now();
+        let mut rng = match config.seed {
+            Some(seed) => rand::rngs::StdRng::seed_from_u64(seed),
+            None => rand::rngs::StdRng::from_os_rng(),
+        };
+
+        let mut subs = substitute(&mut rng, &mut placemap, &substitutes);
+
+        // If desired, glue the substitutes with identical names together.
+        // NOTE: Sorry for the double negative here ;)
+        if !config.no_combine_substitutes {
+            let mut piles: Vec<substitute::Substitution<'_>> = Vec::new();
+            for substitution in subs {
+                if let Some(pile) = piles
+                    .iter_mut()
+                    .find(|pile| pile.name() == substitution.name())
+                {
+                    // Tack this substitution's replacements onto the existing pile.
+                    pile.glue(&substitution)
+                } else {
+                    piles.push(substitution);
+                }
+            }
+
+            subs = piles;
+        }
+
+        // Order them.
+        match config.sort_substitutes {
+            SortBehavior::Size => subs.sort_by_cached_key(|s| Reverse(s.natoms())),
+            SortBehavior::RevSize => subs.sort_by_cached_key(|s| s.natoms()),
+            SortBehavior::Alphabetical => subs.sort_by_key(|s| s.name().to_string()),
+            SortBehavior::RevAlphabetical => subs.sort_by_key(|s| Reverse(s.name().to_string())),
+            SortBehavior::No => {} // Nothing to do.
+        }
+
+        eprintln!("Took {:.3} s.", start.elapsed().as_secs_f32());
+        subs
+    } else {
+        Vec::new()
+    };
+
+    eprintln!("Writing to {:?}...", config.output);
+    let start = std::time::Instant::now();
+    let output_path = &config.output;
+    let file = std::fs::File::create(output_path)
+        .with_context(|| format!("Failed to create output file {output_path:?}"))?;
+    let mut writer = io::BufWriter::new(file);
+    let buffer_size = config.buffer_size;
+    if config.no_write_parallel {
+        write_structure::<false>(
+            &mut writer,
+            &structure,
+            &placemap,
+            &substitutions,
+            buffer_size,
+        )
+        .with_context(|| {
+            format!("Encountered a problem while writing the output structure to {output_path:?}")
+        })?;
+    } else {
+        write_structure::<true>(
+            &mut writer,
+            &structure,
+            &placemap,
+            &substitutions,
+            buffer_size,
+        )
+        .with_context(|| {
+            format!("Encountered a problem while (in parallel) writing the output structure to {output_path:?}")
+        })?;
+    }
+    eprintln!("Writing took {:.3} s", start.elapsed().as_secs_f32());
+
+    let solvent_name = solvent.resname();
+    let nsolvent = placemap.unoccupied_count() as usize;
+    let mut topology = vec![(solvent_name, nsolvent)];
+    topology.extend(substitutions.iter().map(|s| (s.name(), s.natoms())));
+    let mut stdout = io::stdout();
+    match &config.append_topol {
+        Some(path) => eprintln!("Appending solvent topology lines to {path:?} and stdout:"),
+        None => eprintln!("Printing Solvent topology lines to stdout:"),
+    }
+    let mut topol = config
+        .append_topol
+        .map(|path| {
+            std::fs::OpenOptions::new()
+                .create(true)
+                .append(true)
+                .open(&path)
+                .with_context(|| format!("Failed to open topology file {path:?}"))
+        })
+        .transpose()?;
+    for (name, natoms) in topology {
+        let s = format!("{name}\t{natoms}\n");
+        stdout.write_all(s.as_bytes())?;
+        if let Some(topol) = &mut topol {
+            topol
+                .write_all(s.as_bytes())
+                .with_context(|| format!("Failed to append component '{name}' to topology file"))?;
+        }
+    }
+
+    Ok(())
+}

solvate/placement.rs

@@ -0,0 +1,380 @@
+use eightyseven::structure::Atom;
+use glam::{UVec3, Vec3};
+
+use crate::{convert::Convert, water::Water};
+
+/// Return an index into a linear array that represents items on a 3-dimensional grid in a z-major
+/// ordering.
+pub fn index_3d(pos: UVec3, dimensions: UVec3) -> usize {
+    // Casts to usize here serve to avoid wrapping when indexing _very_ large sizes. You never know!
+    pos.x as usize
+        + pos.y as usize * dimensions.x as usize
+        + pos.z as usize * dimensions.x as usize * dimensions.y as usize
+}
+
+/// Iterate over all positions between `(0, 0, 0)` and `dimensions`.
+pub fn iter_3d(dimensions: UVec3) -> impl Iterator<Item = UVec3> {
+    let [dx, dy, dz] = dimensions.to_array();
+    (0..dz).flat_map(move |z| (0..dy).flat_map(move |y| (0..dx).map(move |x| UVec3::new(x, y, z))))
+}
+
+#[derive(Clone)]
+pub struct PlaceMap<'s> {
+    pub solvent: Water,
+    positions: Box<[Vec3]>,
+    dimensions: UVec3,
+    // Invariant: length = (solvent.natoms() / 8).ceil() * dimensions.product()
+    // Invariant: Any wasted bits are always set to zero.
+    // Value is 0 if the place of this solvent bead is not otherwise occupied.
+    // Conversely, if the value of a bit is 1, a solvent bead cannot be placed there.
+    //
+    // To illustrate what these wasted bits are, the figure below shows meaningful bits as `*`, and
+    // the wasted bits as `0`. The bytes are stored in first to last order, and the bits in the
+    // bytes are just treated according to their significance.
+    //
+    // ******** ******** 0000****
+    //                   ^^^^-- Wasted bits.
+    placements: Box<[u8]>,
+    _remove: std::marker::PhantomData<&'s ()>,
+}
+
+impl<'s> PlaceMap<'s> {
+    /// Creates a new [`PlaceMap`].
+    pub fn new(solvent: Water, size: UVec3) -> Self {
+        let positions = solvent.positions().collect::<Box<[_]>>();
+        let n_cells = size.x as usize * size.y as usize * size.z as usize;
+        let n_bytes = positions.len().div_ceil(8);
+        let placements = vec![0x00; n_bytes * n_cells].into_boxed_slice();
+        Self {
+            dimensions: size,
+            solvent,
+            positions,
+            placements,
+            _remove: std::marker::PhantomData,
+        }
+    }
+
+    /// Returns the number of cells that make up this [`PlaceMap`]
+    pub fn n_cells(&self) -> usize {
+        self.dimensions.x as usize * self.dimensions.y as usize * self.dimensions.z as usize
+    }
+
+    /// Returns the number of positions to be considered for solvent-structure collisions.
+    const fn npositions(&self) -> usize {
+        self.positions.len()
+    }
+
+    /// Returns a reference to all solvent positions available in this [`PlaceMap`].
+    pub fn all_positions(&self) -> &[Vec3] {
+        &self.positions
+    }
+
+    // TODO: Revisit, because there's some wasteful go-around, here. We're first creating Atoms
+    // from positions which then get turned back into positions.
+    /// Returns the positions of the solvent particles in this [`PlaceMap`].
+    pub fn iter_positions(&self) -> impl Iterator<Item = eightyseven::structure::Vec3> + '_ {
+        // TODO: Consider if these are hard-equivalent or soft-equivalent in terms of performance.
+        //      self.iter_atoms_chunks().flatten()
+        // My worry is that the flatten version will still allocate more than the one below.
+        self.iter_placements()
+            .flat_map(|(translation, placement)| {
+                self.solvent.substitute_positions(placement, translation)
+            })
+            .map(|pos| pos.convert())
+    }
+
+    pub fn iter_atoms_chunks(&self) -> impl Iterator<Item = Box<[Atom]>> + '_ {
+        self.iter_placements()
+            .map(|(translation, placement)| self.solvent.spray(placement, translation).collect())
+    }
+
+    /// Returns iterator over the [`Placement`]s with their associated translations stored in this
+    /// [`PlaceMap`].
+    pub fn iter_placements(&'_ self) -> impl Iterator<Item = (Vec3, Placement<'_>)> {
+        iter_3d(self.dimensions).map(|cell_pos| {
+            let translation = cell_pos.as_vec3() * self.solvent.dimensions();
+            let placement = self.get(cell_pos).unwrap();
+            (translation, placement)
+        })
+    }
+
+    pub fn iter_placements_mut(&'_ mut self) -> impl Iterator<Item = PlacementMut<'_>> {
+        let n_beads = self.npositions();
+        let n_bytes = n_beads.div_ceil(8);
+        self.placements
+            .chunks_exact_mut(n_bytes)
+            .map(move |bytes| PlacementMut::new(bytes, n_beads))
+    }
+
+    pub fn get_mut(&'_ mut self, pos: UVec3) -> Option<PlacementMut<'_>> {
+        if pos.x >= self.dimensions.x || pos.y >= self.dimensions.y || pos.z >= self.dimensions.z {
+            return None;
+        }
+
+        // TODO: Is it wasteful to 'compute' this here? Or does it reduce down well?
+        let n_beads = self.npositions();
+        let n_bytes = n_beads.div_ceil(8);
+        let idx = index_3d(pos, self.dimensions);
+        // This unwrap should be safe, since we checked at the start of the function.
+        let bytes = self.placements.chunks_exact_mut(n_bytes).nth(idx).unwrap();
+
+        Some(PlacementMut::new(bytes, n_beads))
+    }
+
+    pub fn get(&'_ self, pos: UVec3) -> Option<Placement<'_>> {
+        if pos.x >= self.dimensions.x || pos.y >= self.dimensions.y || pos.z >= self.dimensions.z {
+            return None;
+        }
+
+        let n_beads = self.npositions();
+        let n_bytes = n_beads.div_ceil(8);
+        let idx = index_3d(pos, self.dimensions);
+        // This unwrap should be safe, since we checked at the start of the function.
+        let bytes = self.placements.chunks_exact(n_bytes).nth(idx).unwrap();
+
+        Some(Placement::new(bytes, n_beads))
+    }
+
+    /// Return the count of the positions that are marked as occupied.
+    pub fn occupied_count(&self) -> u64 {
+        // Since a bit is marked as 1 iff the position it represents is occupied, we can just
+        // perform a popcount over the whole bit array.
+        // Note that we don't have to worry about the wasted bits, since these are always set to
+        // zero.
+        self.placements.iter().map(|b| b.count_ones() as u64).sum()
+    }
+
+    /// Return the count of unoccupied positions.
+    pub fn unoccupied_count(&self) -> u64 {
+        let total = self.npositions() * self.n_cells();
+        total as u64 - self.occupied_count()
+    }
+
+    /// Repair all wasted bits in the type.
+    fn repair(&mut self) {
+        let n_beads = self.npositions();
+        let n_bytes = n_beads.div_ceil(8);
+        self.placements
+            .chunks_exact_mut(n_bytes)
+            .map(|chunk| PlacementMut::new(chunk, n_beads))
+            .for_each(|mut p| p.repair())
+    }
+}
+
+impl std::ops::BitOrAssign for PlaceMap<'_> {
+    // Invariant: The or operation will not break the invariant as long as the two operands are
+    // also valid.
+    fn bitor_assign(&mut self, rhs: Self) {
+        assert_eq!(self.dimensions, rhs.dimensions);
+        assert_eq!(self.n_cells(), rhs.n_cells());
+        assert_eq!(self.solvent, rhs.solvent);
+
+        self.placements
+            .iter_mut()
+            .zip(rhs.placements)
+            .for_each(|(s, r)| {
+                *s |= r;
+            });
+    }
+}
+
+impl std::ops::BitAndAssign for PlaceMap<'_> {
+    // Invariant: The and operation will not break the invariant as long as the two operands are
+    // also valid.
+    fn bitand_assign(&mut self, rhs: Self) {
+        assert_eq!(self.dimensions, rhs.dimensions);
+        assert_eq!(self.n_cells(), rhs.n_cells());
+        assert_eq!(self.solvent, rhs.solvent);
+
+        self.placements
+            .iter_mut()
+            .zip(rhs.placements)
+            .for_each(|(s, r)| {
+                *s &= r;
+            });
+    }
+}
+
+impl std::ops::BitXorAssign for PlaceMap<'_> {
+    // Invariant: The xor operation will not break the invariant as long as the two operands are
+    // also valid.
+    fn bitxor_assign(&mut self, rhs: Self) {
+        assert_eq!(self.dimensions, rhs.dimensions);
+        assert_eq!(self.n_cells(), rhs.n_cells());
+        assert_eq!(self.solvent, rhs.solvent);
+
+        self.placements
+            .iter_mut()
+            .zip(rhs.placements)
+            .for_each(|(s, r)| {
+                *s ^= r;
+            });
+    }
+}
+
+impl std::ops::Not for PlaceMap<'_> {
+    type Output = Self;
+
+    fn not(mut self) -> Self::Output {
+        // In order to preserve the invariant that the wasted bits are always zero, we need to be
+        // careful, here.
+        // First, we will go over all the placement bytes and apply a not.
+        self.placements.iter_mut().for_each(|s| *s = !(*s));
+
+        // After that, we go and repair the wasted bits, setting them to zero.
+        self.repair();
+
+        self
+    }
+}
+
+pub struct Placement<'ps> {
+    pub(crate) bytes: &'ps [u8],
+    pub(crate) len: usize,
+    idx: usize,
+}
+
+impl<'ps> Placement<'ps> {
+    fn new(bytes: &'ps [u8], len: usize) -> Self {
+        Self { bytes, len, idx: 0 }
+    }
+}
+
+impl Iterator for Placement<'_> {
+    type Item = bool;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.idx >= self.len {
+            return None;
+        }
+
+        let byte = self.bytes[self.idx / 8];
+        let bit = (byte >> (self.idx % 8)) & 1;
+        self.idx += 1;
+        Some(bit > 0)
+    }
+}
+
+pub struct PlacementMut<'ps> {
+    bytes: &'ps mut [u8],
+    len: usize,
+}
+
+impl<'ps> PlacementMut<'ps> {
+    pub(crate) fn new(bytes: &'ps mut [u8], len: usize) -> Self {
+        Self { bytes, len }
+    }
+
+    pub fn len(&self) -> usize {
+        self.len
+    }
+
+    /// Set the bit value for a solvent bead to represent that its spot is occupied.
+    ///
+    /// This communicates that a solvent particle cannot be rendered at this spot.
+    pub fn set_occupied(&mut self, idx: usize) {
+        if idx >= self.len {
+            panic!(
+                "index out of range (length was {} but index was {idx})",
+                self.len
+            )
+        }
+
+        let byte = &mut self.bytes[idx / 8];
+        *byte |= 1 << (idx % 8);
+    }
+
+    /// Get the value of the bit at `idx`.
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if the provided `idx` exceeds the range represented by this value.
+    pub fn get(&self, idx: usize) -> bool {
+        if idx >= self.len {
+            panic!(
+                "index out of range (length was {} but index was {idx})",
+                self.len
+            )
+        }
+
+        let byte = self.bytes[idx / 8];
+        byte >> (idx % 8) & 1 > 0
+    }
+
+    /// Repair the wasted bits at the end of this [`PlacementMut`] by setting them to zero.
+    ///
+    /// This operation is idempotent.
+    fn repair(&mut self) {
+        let n_bytes = self.bytes.len();
+        let n_bits = n_bytes * u8::BITS as usize;
+        if n_bits == self.len {
+            return;
+        }
+
+        let Some(last) = self.bytes.last_mut() else {
+            return;
+        };
+        let n_wasted = n_bits - self.len;
+
+        let mask = (1u8 << (u8::BITS as usize - n_wasted)) - 1;
+        *last &= mask;
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn repair_placement_mut() {
+        let mut bytes = [0xff, 0xff, 0xff];
+        let mut p = PlacementMut::new(&mut bytes, 20);
+
+        assert_eq!(*p.bytes.last().unwrap(), 0b11111111);
+        p.repair();
+        assert_eq!(*p.bytes.last().unwrap(), 0b00001111);
+        p.repair();
+        assert_eq!(*p.bytes.last().unwrap(), 0b00001111);
+    }
+
+    /// Verify that the repair that occurs within the not operation works.
+    ///
+    /// If the repair is incorrect, we would expect to see `unoccupied_count > natoms_total` after
+    /// the not.
+    #[test]
+    fn check_not_repair() {
+        let solvent = Water::Martini;
+        // TODO: Currently, we're testing a trivially easy case.
+        // // Make sure we have a structure size that does not neatly fit in some number of bytes.
+        // if solvent.positions().count() % 8 == 0 {
+        //     solvent.atoms.pop();
+        // }
+        let natoms = solvent.positions().count();
+        let size = UVec3::new(3, 5, 7);
+        let mut placemap = PlaceMap::new(solvent, size);
+        let natoms_total = (natoms * placemap.n_cells()) as u64;
+        let n_occupied = 0;
+
+        // Set some atom to be occupied.
+        placemap
+            .get_mut(UVec3::new(1, 3, 5))
+            .unwrap()
+            .set_occupied(161);
+        // Update natoms_total and n_occupied to reflect this.
+        let natoms_total = natoms_total - 1;
+        let n_occupied = n_occupied + 1;
+
+        assert_eq!(placemap.unoccupied_count(), natoms_total);
+        assert_eq!(placemap.occupied_count(), n_occupied);
+
+        placemap = !placemap;
+
+        assert_eq!(placemap.unoccupied_count(), n_occupied);
+        assert_eq!(placemap.occupied_count(), natoms_total);
+
+        placemap = !placemap;
+
+        assert_eq!(placemap.unoccupied_count(), natoms_total);
+        assert_eq!(placemap.occupied_count(), n_occupied);
+    }
+}