Merge pull request #16 from mosdef-hub/main

Water box (mosdef-hub#1115)

environment-dev-win.yml

@@ -20,7 +20,7 @@ dependencies:
   - parmed>=3.4.3
   - pip
   - pre-commit
-  - protobuf
+  - protobuf<4
   - py3Dmol
   - pycifrw
   - pytest

environment-dev.yml

@@ -21,7 +21,7 @@ dependencies:
   - parmed>=3.4.3
   - pip
   - pre-commit
-  - protobuf
+  - protobuf<4
   - py3Dmol
   - pycifrw
   - pytest

mbuild/compound.py

@@ -1982,11 +1982,13 @@ def condense(self, inplace=True):
         """
         # temporary list of components
         comp_list = []
+        comp_list_id = []
         connected_subgraph = self.root.bond_graph.connected_components()
-
+        unbound_particles = []
         for molecule in connected_subgraph:
             if len(molecule) == 1:
                 ancestors = [molecule[0]]
+                unbound_particles.append(molecule[0])
             else:
                 ancestors = IndexedSet(molecule[0].ancestors())
                 for particle in molecule[1:]:
@@ -2000,9 +2002,22 @@ def condense(self, inplace=True):
                         IndexedSet(particle.ancestors())
                     )
 
-            """Parse molecule information"""
-            molecule_tag = ancestors[0]
-            comp_list.append(clone(molecule_tag))
+                """Parse molecule information"""
+                molecule_tag = ancestors[0]
+                comp_list.append(clone(molecule_tag))
+                # generate a list of particle ids within the compound
+                # this will also include any particles that are part of the
+                # Compound but are not bonded
+
+                pids = [id(p) for p in molecule_tag.particles()]
+                comp_list_id += pids
+        # loop over particles without bonds
+        # if any of the particles exist in a compound
+        # don't add them, as they already exist
+        for ubp in unbound_particles:
+            if id(ubp) not in comp_list_id:
+                comp_list.append(clone(ubp))
+
         if inplace:
             for child in [self.children]:
                 # Need to handle the case when child is a port

mbuild/lib/recipes/water_box.py

@@ -0,0 +1,220 @@
+"""mBuild recipe for building a water box."""
+import itertools
+import math as math
+from collections.abc import Iterable
+
+import numpy as np
+
+import mbuild.lib.molecules.water as water_models
+from mbuild import Compound, clone, force_overlap, load
+from mbuild.exceptions import MBuildError
+
+__all__ = ["Water3SiteBox"]
+
+
+def _flatten_list(c_list):
+    """Flatten a list.
+
+    Helper function to flatten a list that may be nested, e.g. [comp1, [comp2, comp3]].
+    """
+    if isinstance(c_list, Iterable) and not isinstance(c_list, str):
+        for c in c_list:
+            if isinstance(c, Iterable) and not isinstance(c, str):
+                yield from _flatten_list(c)
+            else:
+                yield c
+
+
+class Water3SiteBox(Compound):
+    """Generate a box of 3-site water molecules.
+
+    Efficiently create an mbuild Compound containing water at density ~1000 kg/m^3
+    where local molecule orientations should exist in relatively low energy states.
+    This loads in a configuration previously generated with packmol and relaxed with
+    GROMACS via NVT simulation at 305K using tip3p model, simulated in a 4 nm^3 box.
+    The code will duplicate/truncate the configuration as necessary to satisify the
+    given box dimensions.
+
+    Parameters
+    ----------
+    box : mb.Box
+        The desired box to fill with water
+    edge : float or list of floats, default=0.1 (nm)
+        Specifies the gutter around the system to avoid overlaps at boundaries
+    model : mb.Compound, optional, default=water_models.WaterTIP3P()
+        The specified 3-site water model to be used. This uses the force overlap
+        command to translate and orient the specified water model to the given coordinates.
+        See mbuild/lib/molecules/water.py for available water models or extend the base model.
+    mask : mb.Compound, optional, default=None
+        Remove water molecules from the final configuration that overlap with the Compound
+        specified by the mask. If the element field is set, the sum of the particle radii
+        will be used.
+    radii_dict : dict, optional, None
+        User defined radii values based on the name field of a Compound.  This will supercede
+        the built-in elements' radii.
+    radii_overlap : float, optional, default=0.15 (nm)
+        Default value if the radii_dict or element field are not defined for a particle.
+    radii_scaling : float, optional, default=1.0
+        Radii are multiplied by this factor during the masking routine. This allows the
+        space between the masking particles and the water to be adjusted. This will apply to
+        values in radii_dict, radii based on element, and radii_overlap.
+
+
+    """
+
+    def __init__(
+        self,
+        box,
+        edge=0.1,
+        model=water_models.WaterTIP3P(),
+        mask=None,
+        radii_dict=None,
+        radii_overlap=0.15,
+        radii_scaling=1.0,
+    ):
+        super(Water3SiteBox, self).__init__()
+
+        # if we do not define a dictionary, create an empty one
+        if radii_dict is None:
+            radii_dict = {}
+        else:
+            if not isinstance(radii_dict, dict):
+                raise ValueError(f"radii_dict should be dictionary.")
+
+        # check if we are given a list or single value
+        if isinstance(edge, (list, tuple)):
+            if len(edge) != 3:
+                raise ValueError(
+                    f"edge should either be a single float/int or a list of length 3, not a list of {len(edge)}."
+                )
+            edges = np.array(edge)
+        else:
+            assert isinstance(edge, (float, int))
+            edges = np.array([edge, edge, edge])
+
+        # If a model is specified, we will check to ensure that
+        # the first particle in the compound corresponds to Oxygen.
+
+        if model is not None:
+            if not isinstance(model, Compound):
+                raise MBuildError(f"Model must be a compound.")
+            particles = [p for p in model.particles()]
+            if particles[0].element.symbol != "O":
+                raise MBuildError(
+                    "The first particle in the model needs to correspond to oxygen."
+                )
+            if len(particles) != 3:
+                raise MBuildError("The only works with 3-site models of water.")
+
+        # check if mask is set
+        msg = "Mask must be a Compound or a list/tuple of Compounds."
+        if mask is not None:
+            if not isinstance(mask, (list, tuple, Compound)):
+                raise MBuildError(msg)
+            elif isinstance(mask, (list, tuple)):
+                # in case we are specified a list of Compounds,
+                # we will make sure it is a 1d list.
+                mask = [e for e in _flatten_list(mask)]
+                if not all([isinstance(entry, Compound) for entry in mask]):
+                    raise MBuildError(msg)
+
+        # read in our propotype, a 4.0x4.0x4.0 nm box
+        # our prototype was relaxed in GROMACs at 305 K, density 1000 kg/m^3 using tip3p
+        aa_waters = load(
+            "water_proto.gro",
+            relative_to_module=self.__module__,
+        )
+
+        # loop over each water in our configuration
+        # add in the necessary bonds missing from the .gro file
+        # rename particles/Compound according to the given water model
+        for water in aa_waters.children:
+            water.add_bond((water.children[0], water.children[1]))
+            water.add_bond((water.children[0], water.children[2]))
+
+            temp = clone(model)
+            force_overlap(temp, temp, water, add_bond=False)
+            water.name = model.name
+            water.children[0].name = model.children[0].name
+            water.children[1].name = model.children[1].name
+            water.children[2].name = model.children[2].name
+            water.xyz = temp.xyz
+
+        # scaling parameters for the new box
+        scale_Lx = math.ceil(box.Lx / aa_waters.box.Lx)
+        scale_Ly = math.ceil(box.Ly / aa_waters.box.Ly)
+        scale_Lz = math.ceil(box.Lz / aa_waters.box.Lz)
+
+        water_system_list = []
+
+        # we will create a list of particles for the mask
+        # if specified now to save time later
+        if mask is not None:
+            if isinstance(mask, Compound):
+                p_mask = [p for p in mask.particles()]
+            else:
+                p_mask = []
+                for entry in mask:
+                    p_mask = p_mask + [p for p in entry.particles()]
+
+        # add water molecules to a list
+        # note we add to a list first, as this is more efficient than calling
+        # the Compound.add function repeatedly as the Compound size grows.
+        for water in aa_waters.children:
+            for i, j, k in itertools.product(
+                range(scale_Lx), range(scale_Ly), range(scale_Lz)
+            ):
+                shift = np.array(
+                    [
+                        i * aa_waters.box.Lx,
+                        j * aa_waters.box.Ly,
+                        k * aa_waters.box.Lz,
+                    ]
+                )
+                if all(water.pos + shift < (box.lengths - edges)):
+                    if mask is not None:
+                        particles = [p for p in water.particles()]
+                        status = True
+
+                        # note this could be sped up using a cell list
+                        # will have to wait until that PR is merged
+                        for p1, p2 in itertools.product(particles, p_mask):
+                            dist = np.linalg.norm(p1.pos - p2.pos)
+
+                            if p1.name in radii_dict:
+                                c1 = radii_scaling * radii_dict[p1.name]
+                            elif p1.element is not None:
+                                c1 = (
+                                    radii_scaling
+                                    * p1.element.radius_alvarez
+                                    / 10.0
+                                )
+                            else:
+                                c1 = radii_scaling * radii_overlap
+
+                            if p2.name in radii_dict:
+                                c2 = radii_scaling * radii_dict[p2.name]
+                            elif p2.element is not None:
+                                c2 = (
+                                    radii_scaling
+                                    * p2.element.radius_alvarez
+                                    / 10.0
+                                )
+                            else:
+                                c2 = radii_scaling * radii_overlap
+
+                            cut_value = c1 + c2
+                            if dist <= cut_value:
+                                status = False
+                        if status:
+                            temp = clone(water)
+                            temp.translate(shift)
+                            water_system_list.append(temp)
+                    else:
+                        temp = clone(water)
+                        temp.translate(shift)
+                        water_system_list.append(temp)
+
+        # add to the Compound and set box size
+        self.add(water_system_list)
+        self.box = box