bentopy 0.2.0a10__cp313-cp313-manylinux_2_34_x86_64.whl

solvate/args.rs

@@ -0,0 +1,324 @@
+use std::path::PathBuf;
+use std::str::FromStr;
+
+use clap::{Parser, ValueEnum};
+use eightyseven::structure::ResName;
+
+/// Solvate.
+#[derive(Debug, Parser)]
+#[command(about, version = bentopy::core::version::VERSION)]
+pub struct Args {
+    /// Structure input path.
+    #[arg(short, long)]
+    pub input: PathBuf,
+    /// Output path.
+    #[arg(short, long)]
+    pub output: PathBuf,
+
+    /// Lowest allowable distance between solvent and structure beads (nm).
+    ///
+    /// This is the minimum allowable center-to-center distance when checking for a collision
+    /// between a solvent bead and a structure bead.
+    ///
+    /// For the solvent-solvent cutoff distance, see `--solvent-cutoff`.
+    ///
+    /// For Martini solvation, the default cutoff is 0.43 nm. For atomistic solvation, the default
+    /// cutoff is 0.28.
+    #[arg(long)]
+    pub cutoff: Option<f32>,
+
+    /// Lowest allowable distance between solvent beads (nm).
+    ///
+    /// This is the minimum allowable center-to-center distance between solvent beads when cutting
+    /// the sides to fit in the specified output structure box.
+    ///
+    /// For Martini solvation, the default cutoff is 0.21 nm. For atomistic solvation, the default
+    /// cutoff is 0.21.
+    #[arg(long)]
+    pub solvent_cutoff: Option<f32>,
+
+    /// List of resnames to ignore when checking against structure-solvent collisions.
+    #[arg(long, value_delimiter = ',')]
+    pub ignore: Vec<ResName>,
+
+    /// The type of water written to the output file.
+    #[arg(long, value_enum, default_value_t)]
+    pub water_type: WaterType,
+
+    /// Center the structure in the new box.
+    ///
+    /// Note that this option only has an effect if the boundary mode is set to `grow`.
+    #[arg(short, long)]
+    pub center: bool,
+
+    /// Set how the boundaries of the box will be treated.
+    #[arg(short, long, value_enum, default_value_t)]
+    pub boundary_mode: BoundaryMode,
+
+    /// Set how periodic images of the input structure are considered.
+    ///
+    /// Note that only orthorhombic (cuboid) periodicity is considered, currently.
+    // TODO: Consider whether supporting other forms of periodicity is worthwhile.
+    #[arg(short, long, value_enum, default_value_t)]
+    pub periodic_mode: PeriodicMode,
+
+    /// Define substitutes for solvent residues, such as ions.
+    #[arg(short, long = "substitute")]
+    pub substitutes: Vec<SubstituteConfig>,
+
+    /// Set the charge to neutralize with additional ions.
+    ///
+    /// In essence, this is a shorthand for explicitly providing ions to compensate the charge as
+    /// substitutes. This can be helpful for automation purposes.
+    ///
+    /// By default, NA (positive) and CL (negative) ions are used, but a different
+    /// positive-negative substitution pair can be specified by prepending a colon followed by the
+    /// positive and negative substitute name separated by one comma.
+    ///
+    ///     <charge>:<positive ion name>,<negative ion name>
+    ///
+    /// So, by default, some `<charge>` is interpreted as
+    ///
+    ///     <charge>:NA,CL
+    #[arg(long, allow_hyphen_values = true)]
+    pub charge: Option<ChargeConfig>,
+
+    /// Combine substitutes with identical names into one block.
+    #[arg(long, default_value_t)]
+    pub no_combine_substitutes: bool,
+
+    /// Set whether and in what way substitutes are sorted.
+    #[arg(long, value_enum, default_value_t)]
+    pub sort_substitutes: SortBehavior,
+
+    /// Random number generator seed for solvent substitution, such as ion placement.
+    #[arg(long)]
+    pub seed: Option<u64>,
+
+    /// If the solvent template contains velocity information, write these velocities to the output
+    /// file.
+    #[arg(long, default_value_t)]
+    pub write_velocities: bool,
+
+    /// Append solvation topology lines to a path.
+    #[arg(short = 't', long)]
+    pub append_topol: Option<PathBuf>,
+
+    #[arg(long)]
+    pub no_write_parallel: bool,
+    /// The suggested number of atoms to format at once.
+    ///
+    /// Setting this to a larger value will allow more atoms to be formatted at once, at the cost
+    /// of higher memory consumption. Setting this to a smaller value will lower the memory
+    /// footprint at a possible cost of efficiency.
+    ///
+    /// Note that depending on the number of residues in the template box, the provided value may
+    /// be overshot by at most that number.
+    #[arg(long, default_value_t = 10000000)]
+    pub buffer_size: usize,
+}
+
+#[derive(Debug, Default, Clone, Copy, ValueEnum, PartialEq, Eq)]
+#[clap(rename_all = "lowercase")]
+pub enum WaterType {
+    #[default]
+    Martini,
+    Tip3P,
+}
+
+impl WaterType {
+    pub const fn default_cutoff(&self) -> f32 {
+        match self {
+            WaterType::Martini => 0.43,
+            WaterType::Tip3P => 0.28,
+        }
+    }
+
+    pub const fn default_solvent_cutoff(&self) -> f32 {
+        match self {
+            WaterType::Martini => 0.21,
+            WaterType::Tip3P => 0.21,
+        }
+    }
+}
+
+#[derive(Debug, Default, Clone, ValueEnum, PartialEq, Eq)]
+pub enum BoundaryMode {
+    /// Cut at the structure box size and remove solvent residues that overlap with the periodic
+    /// neighbors.
+    #[default]
+    Cut,
+    /// If necessary, grow the box of the output structure to fit a whole number of template boxes.
+    Grow,
+}
+
+#[derive(Debug, Default, Clone, ValueEnum, PartialEq, Eq)]
+pub enum PeriodicMode {
+    /// Include the periodic images of the structure when checking whether a solvent spot is
+    /// occupied.
+    #[default]
+    Periodic,
+    /// Ignore the periodic images of the structure for the solvent placement check.
+    Ignore,
+    /// Treat any structure atoms outside of the output box as an error and exit.
+    Deny,
+}
+
+#[derive(Debug, Clone)]
+pub struct SubstituteConfig {
+    name: String,
+    quantity: Quantity,
+}
+
+impl SubstituteConfig {
+    /// Bake into a [`Substitute`] according to a final volume in nm³ and some number of valid
+    /// solvent beads.
+    pub fn bake(self, volume: f64, solvent_beads: u64) -> Substitute {
+        Substitute {
+            name: self.name,
+            number: self.quantity.number(volume, solvent_beads),
+        }
+    }
+}
+
+impl FromStr for SubstituteConfig {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let mut words = s.split(':');
+        let name = words
+            .next()
+            .ok_or("expected a name, then a colon, followed by a quantity".to_string())?
+            .to_string();
+        let quantity = words
+            .next()
+            .ok_or("expected a quantifier after the colon".to_string())?
+            .parse()?;
+        Ok(Self { name, quantity })
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct ChargeConfig {
+    charge: i64,
+    positive: String,
+    negative: String,
+}
+
+impl FromStr for ChargeConfig {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let mut words = s.split(':');
+        let charge = words
+            .next()
+            .ok_or("expected a charge")?
+            .parse::<i64>()
+            .map_err(|err| err.to_string())?;
+        let (positive, negative) = if let Some(names) = words.next() {
+            names
+                .split_once(',')
+                .ok_or("expected two substitute names separated by a comma")?
+        } else {
+            ("NA", "CL")
+        };
+        Ok(Self {
+            charge,
+            positive: positive.to_string(),
+            negative: negative.to_string(),
+        })
+    }
+}
+
+impl ChargeConfig {
+    /// Bake the [`ChargeConfig`] into a [`Substitute`] if applicable.
+    ///
+    /// If the charge is 0, no `Substitute` is returned since there is no charge to neutralize.
+    pub fn bake(self) -> Option<Substitute> {
+        // Choose the name of the ion to neutralize with. If the charge to compensate is negative,
+        // we compensate with the positive ion substitute, and vice versa.
+        let name = match self.charge {
+            ..0 => self.positive,
+            0 => return None,
+            1.. => self.negative,
+        };
+
+        Some(Substitute {
+            name,
+            number: self.charge.unsigned_abs(),
+        })
+    }
+}
+
+#[derive(Debug, Clone)]
+enum Quantity {
+    /// A fixed number.
+    Number(u64),
+    /// A fraction of the valid solvent beads.
+    Ratio(f64),
+    /// Molarity in mol/L.
+    Molarity(f64),
+}
+
+impl FromStr for Quantity {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        if let Some(molarity) = s.strip_suffix('M') {
+            // Molarity.
+            return Ok(Self::Molarity(
+                molarity.parse::<f64>().map_err(|err| err.to_string())?,
+            ));
+        }
+
+        if let Ok(number) = s.parse::<u64>() {
+            return Ok(Self::Number(number));
+        }
+
+        if let Ok(ratio) = s.parse::<f64>() {
+            if !(0.0..=1.0).contains(&ratio) {
+                return Err(format!(
+                    "ratio value must satisfy `1.0 >= r >= 0.0`, found {ratio}"
+                ));
+            }
+            return Ok(Self::Ratio(ratio));
+        }
+
+        Err(format!("could not parse quantity {s:?}"))
+    }
+}
+
+impl Quantity {
+    /// Given some volume in nm³ and some number of initial valid solvent beads, determine the
+    /// number of beads according to this [`Quantity`].
+    fn number(self, volume: f64, solvent_beads: u64) -> u64 {
+        const N_AVOGADRO: f64 = 6.0221415e23; // per mol.
+        let volume_liter = volume * 1e-24;
+        let number = match self {
+            Quantity::Number(number) => number,
+            Quantity::Ratio(ratio) => (solvent_beads as f64 * ratio).round() as u64,
+            Quantity::Molarity(molarity) => (volume_liter * molarity * N_AVOGADRO).round() as u64,
+        };
+        assert!(
+            number <= solvent_beads,
+            "the number of substitutions that are specified ({number}) exceeds the number of \
+                solvent positions that can be substituted ({solvent_beads})"
+        );
+        number
+    }
+}
+
+pub struct Substitute {
+    pub name: String,
+    pub number: u64,
+}
+
+#[derive(ValueEnum, Default, Clone, Debug)]
+pub enum SortBehavior {
+    #[default]
+    Size,
+    RevSize,
+    Alphabetical,
+    RevAlphabetical,
+    No,
+}

solvate/convert.rs

@@ -0,0 +1,25 @@
+/// Convert between `Vec3` types.
+///
+/// In this program we use both [`glam::Vec3`] and the [`eightyseven::structure::Vec3`] type, which
+/// is an export of some version of `glam::Vec3` used in `eightyseven`.
+///
+/// This is very annoying, because we can't just use them interchangeably despite the types being
+/// otherwise identical.
+pub trait Convert<V> {
+    /// Convert between `Vec3` types.
+    fn convert(&self) -> V;
+}
+
+impl Convert<glam::Vec3> for eightyseven::structure::Vec3 {
+    #[inline(always)]
+    fn convert(&self) -> glam::Vec3 {
+        glam::Vec3::from_array(self.to_array())
+    }
+}
+
+impl Convert<eightyseven::structure::Vec3> for glam::Vec3 {
+    #[inline(always)]
+    fn convert(&self) -> eightyseven::structure::Vec3 {
+        eightyseven::structure::Vec3::from_array(self.to_array())
+    }
+}

solvate/cookies.rs

@@ -0,0 +1,185 @@
+use eightyseven::structure::ResName;
+use glam::{BVec3, IVec3, UVec3, Vec3, Vec3A};
+
+use crate::PeriodicMode;
+use crate::convert::Convert;
+use crate::placement::{index_3d, iter_3d};
+use crate::structure::{BoxVecsExtension, Structure};
+
+/// A simple spatial data structure that is used to accelerate solvent-structure collision checks.
+///
+/// Each _cookie_ represents one solvent box, and is constructed in such a manner that all
+/// positions that are in or sufficiently close to the solvent beads represented by a cookie are
+/// stored in the cookie.
+pub struct Cookies {
+    dimensions: UVec3,
+    cookie_size: Vec3,
+    cookies: Box<[Box<[Vec3]>]>,
+}
+
+impl Cookies {
+    /// Creates a new [`Cookies`].
+    ///
+    /// Note that the box size for the `structure` are treated as the dimensions when considering
+    /// the treatment of atoms according to the [`PeriodicMode`].
+    pub fn new(
+        structure: &Structure,
+        ignore: &[ResName],
+        cookie_size: Vec3,
+        dimensions: UVec3,
+        cutoff: f32,
+        mode: PeriodicMode,
+    ) -> Self {
+        let n_cells = dimensions.x as usize * dimensions.y as usize * dimensions.z as usize;
+        let mut cookies = vec![Vec::new(); n_cells];
+
+        // Place the bead's position in the appropriate cookie.
+        // Like raisins or pieces of chocolate. But in our case the cookies aren't round but
+        // cuboid. Look, I don't even remember why I picked this name.
+        let box_dimensions = structure.boxvecs.as_vec3();
+        for (i, bead) in structure.atoms.iter().enumerate() {
+            let mut pos = bead.position.convert();
+            if ignore.contains(&bead.resname) {
+                // We don't even consider this particle in setting up our Cookies because we want
+                // to ignore it in our solvent-water check.
+                continue;
+            }
+            match mode {
+                PeriodicMode::Periodic => pos = pos.rem_euclid(box_dimensions),
+                PeriodicMode::Ignore => {
+                    if pos.cmplt(Vec3::ZERO).any() || pos.cmpgt(box_dimensions).any() {
+                        continue;
+                    }
+                }
+                PeriodicMode::Deny => {
+                    if pos.cmplt(Vec3::ZERO).any() || pos.cmpgt(box_dimensions).any() {
+                        eprintln!(
+                            "ERROR: Atom {i} with position {pos} lies outside the box {box_dimensions}."
+                        );
+                        // FIXME: It's rather gross to me to exit from inside this function. A nice
+                        // improvement could be to bubble it up as an Err to the caller.
+                        std::process::exit(2);
+                    }
+                }
+            }
+            // Note that rem_euclid or our PeriodicMode::Deny check may return negative zeros,
+            // which may subtly break our assumptions for determining cookie membership.
+            // Perhaps overly strict but better be safe.
+            let pos = pos.abs();
+            let cell_pos = (pos / cookie_size).floor().as_uvec3();
+            let idx = index_3d(cell_pos, dimensions);
+            cookies[idx].push(pos);
+        }
+
+        // Now, we have to convolve the contents of the cookies to their neighbors. Otherwise, a
+        // neighboring bead may be too close to some solvent bead but go unnoticed in the collision
+        // check.
+        let idimensions = dimensions.as_ivec3();
+        let da = Vec3A::from(box_dimensions);
+        // NOTE: If we are ever _really_ pressed for memory, we can win quite a bit here. Cloning
+        // the cookies here saves some time gluing the neighbors onto the original positions. But,
+        // there are ways of removing this clone.
+        let mut convolved_cookies = cookies.clone();
+        for cell_pos in iter_3d(dimensions) {
+            // This closure returns true if the provided position is in cutoff-range of the present
+            // cookie.
+            let is_in_range = {
+                let start = Vec3A::from(cookie_size * cell_pos.as_vec3() - cutoff);
+                let end = Vec3A::from(cookie_size * (cell_pos + 1).as_vec3() + cutoff);
+                move |v: &Vec3A| v.cmpge(start).all() && v.cmple(end).all()
+            };
+            let idx = index_3d(cell_pos, dimensions);
+
+            for offset in NEIGHBORS {
+                let neighbor_pos = cell_pos.as_ivec3() + offset;
+                // Note that these jumps are mutually exclusive.
+                let backward_jumps = neighbor_pos.cmplt(IVec3::ZERO);
+                let forward_jumps = neighbor_pos.cmpge(idimensions);
+                // Sanity check of their mutual exclusivity.
+                assert_eq!(backward_jumps & forward_jumps, BVec3::FALSE);
+
+                if !(backward_jumps | forward_jumps).any() {
+                    // Default case, no complicated stuff with periodicity to consider.
+                    // TODO: Take only the positions that are closer than `cutoff` to the
+                    // central cell.
+                    let neighbor_pos = neighbor_pos.as_uvec3();
+                    let neighbor_idx = index_3d(neighbor_pos, dimensions);
+                    let neighbor_content = cookies[neighbor_idx]
+                        .iter()
+                        .map(|&v| Vec3A::from(v))
+                        .filter(is_in_range)
+                        .map(Vec3::from);
+                    convolved_cookies[idx].extend(neighbor_content);
+                } else {
+                    fn apply(b: BVec3, v: Vec3A) -> Vec3A {
+                        UVec3::new(b.x as u32, b.y as u32, b.z as u32).as_vec3a() * v
+                    }
+
+                    let wrapped_neighbor_pos = neighbor_pos.rem_euclid(idimensions).as_uvec3();
+                    let wrapped_neighbor_idx = index_3d(wrapped_neighbor_pos, dimensions);
+                    let wrapped_neighbor_content = cookies[wrapped_neighbor_idx]
+                        .iter()
+                        .map(|&v| {
+                            let translation = apply(backward_jumps, -da) + apply(forward_jumps, da);
+                            Vec3A::from(v) + translation
+                        })
+                        .filter(is_in_range)
+                        .map(Vec3::from);
+                    convolved_cookies[idx].extend(wrapped_neighbor_content);
+                }
+            }
+        }
+
+        Self {
+            dimensions,
+            cookie_size,
+            cookies: convolved_cookies
+                .into_iter()
+                .map(|cookie| cookie.into_boxed_slice())
+                .collect(),
+        }
+    }
+
+    pub fn get(&self, cell_pos: UVec3) -> Option<&[Vec3]> {
+        let idx = index_3d(cell_pos, self.dimensions);
+        self.cookies.get(idx).map(|cookie| &cookie[..])
+    }
+
+    /// Calculate the offset of a cookie at some position from the system origin.
+    pub fn offset(&self, cell_pos: UVec3) -> Vec3 {
+        cell_pos.as_vec3() * self.cookie_size()
+    }
+
+    pub fn cookie_size(&self) -> Vec3 {
+        self.cookie_size
+    }
+}
+
+const NEIGHBORS: [IVec3; 26] = [
+    IVec3::new(-1, -1, -1),
+    IVec3::new(0, -1, -1),
+    IVec3::new(1, -1, -1),
+    IVec3::new(-1, 0, -1),
+    IVec3::new(0, 0, -1),
+    IVec3::new(1, 0, -1),
+    IVec3::new(-1, 1, -1),
+    IVec3::new(0, 1, -1),
+    IVec3::new(1, 1, -1),
+    IVec3::new(-1, -1, 0),
+    IVec3::new(0, -1, 0),
+    IVec3::new(1, -1, 0),
+    IVec3::new(-1, 0, 0),
+    IVec3::new(1, 0, 0),
+    IVec3::new(-1, 1, 0),
+    IVec3::new(0, 1, 0),
+    IVec3::new(1, 1, 0),
+    IVec3::new(-1, -1, 1),
+    IVec3::new(0, -1, 1),
+    IVec3::new(1, -1, 1),
+    IVec3::new(-1, 0, 1),
+    IVec3::new(0, 0, 1),
+    IVec3::new(1, 0, 1),
+    IVec3::new(-1, 1, 1),
+    IVec3::new(0, 1, 1),
+    IVec3::new(1, 1, 1),
+];