Add restricted insertions attribute for moves

mosdef_cassandra/core/moves.py

@@ -1,6 +1,7 @@
 from copy import deepcopy
 
 import parmed
+import warnings
 
 
 class Moves(object):
@@ -58,6 +59,10 @@ def __init__(self, ensemble, species_topologies):
         else:
             self._n_boxes = 2
 
+        # Set '_restricted_typed' and '_restricted_value'
+        self._restricted_type = None
+        self._restricted_value = None
+
         # Define default probabilities
         # Most are ensemble-dependent
         self.prob_angle = 0.0
@@ -173,6 +178,96 @@ def __init__(self, ensemble, species_topologies):
                 self.prob_translate += self.prob_rotate
                 self.prob_rotate = 0.0
 
+    def add_restricted_insertions(
+        self, species_topologies, restricted_type, restricted_value
+    ):
+        """Add restricted insertions for specific species and boxes
+
+        Parameters
+        ----------
+        species_topologies : list
+            list of ``parmed.Structures`` containing one list per box of species
+        restricted_type : list
+            list of restricted insertion types containing one list per box of species
+        restricted_value : list
+            list of restricted insertion values containing one list per box of species
+        """
+        if self._restricted_type and self._restricted_value:
+            warnings.warn(
+                "Restricted insertion has been previously"
+                " added and will be replaced."
+            )
+        if self.ensemble not in ["gcmc", "gemc", "gemc_npt"]:
+            raise ValueError(
+                "Restricted insertions are only valid for"
+                " 'gcmc', 'gemc', and 'gemc_npt' ensembles."
+            )
+        if len(restricted_type) != len(restricted_value):
+            raise ValueError(
+                "Length of 'restricted_type' and "
+                " 'restricted_value' must match."
+            )
+        for box in restricted_type:
+            if isinstance(box, (str, int, float)):
+                raise TypeError(
+                    "Restricted type must be passed as a list"
+                    " of lists corresponding to each box."
+                )
+            if len(box) != len(species_topologies):
+                raise ValueError(
+                    "Length of 'species' and "
+                    " length of box list in 'restricted_type'"
+                    " must match.  `species` has a length of {}"
+                    " and the box list in 'restricted_type' has a "
+                    " length of {}".format(len(species_topologies), len(box))
+                )
+        for box in restricted_value:
+            if isinstance(box, (str, int, float)):
+                raise TypeError(
+                    "Restricted value must be passed as a list"
+                    " of lists corresponding to each box."
+                )
+            if len(box) != len(species_topologies):
+                raise ValueError(
+                    "Length of 'species' and "
+                    " length of species list in 'restricted_value'"
+                    " must match.  `species` has a length of {}"
+                    " and the box list in 'restricted_value' has a "
+                    " length of {}".format(len(species_topologies), len(box))
+                )
+        if self.ensemble == "gcmc" and len(restricted_type) != 1:
+            raise ValueError(
+                "GCMC ensemble contains 1 box but"
+                " `restricted_type` of length {}"
+                " was passed.".format(len(restricted_type))
+            )
+        if self.ensemble in ["gemc", "gemc_npt"] and len(restricted_type) != 2:
+            raise ValueError(
+                "GEMC ensembles contain 2 boxes but"
+                " `restricted_type` of length {}"
+                " was passed.".format(len(restricted_type))
+            )
+
+        for types, values in zip(restricted_type, restricted_value):
+            for typ, val in zip(types, values):
+                if not typ and not val:
+                    pass
+                elif typ and not val:
+                    raise ValueError(
+                        "`restricted_type` {} was passed"
+                        " but `restricted_value` is None.".format(typ, val)
+                    )
+                elif val and not typ:
+                    raise ValueError(
+                        "`restricted_value` {} was passed"
+                        " but `restricted_type` is None.".format(val, typ)
+                    )
+                else:
+                    _check_restriction_type(typ, val)
+
+        self._restricted_type = restricted_type
+        self._restricted_value = restricted_value
+
     @property
     def ensemble(self):
         return self._ensemble
@@ -677,4 +772,65 @@ def print(self):
                 box=box + 1, max_vol=max_vol
             )
 
+        if self._restricted_type != None:
+            contents += "\nRestricted Insertions (Ang):\n"
+            for box in range(self._n_boxes):
+                for species, (typ, value) in enumerate(
+                    zip(
+                        self._restricted_type[box], self._restricted_value[box]
+                    )
+                ):
+                    if typ == "sphere":
+                        contents += "Box {box}, Species {species}: sphere, R = {r_value}\n".format(
+                            box=box + 1, species=species + 1, r_value=value
+                        )
+                    elif typ == "cylinder":
+                        contents += "Box {box}, Species {species}: cylinder, R = {r_value}\n".format(
+                            box=box + 1, species=species + 1, r_value=value
+                        )
+                    elif typ == "slitpore":
+                        contents += "Box {box}, Species {species}: slitpore, z_max = {z_max}\n".format(
+                            box=box + 1, species=species + 1, z_max=value
+                        )
+                    elif typ == "interface":
+                        contents += "Box {box}, Species {species}: interface, z_min = {z_min}, z_max = {z_max}\n".format(
+                            box=box + 1,
+                            species=species + 1,
+                            z_min=value[0],
+                            z_max=value[1],
+                        )
+                    else:
+                        contents += "Box {box}, Species {species}: None\n".format(
+                            box=box + 1, species=species + 1
+                        )
+
         print(contents)
+
+
+def _check_restriction_type(restriction_type, restriction_value):
+    valid_restrict_types = ["sphere", "cylinder", "slitpore", "interface"]
+    # Check restriction insertion type
+    if restriction_type not in valid_restrict_types:
+        raise ValueError(
+            'Invalid restriction type "{}".  Supported '
+            "restriction types include {}".format(
+                restriction_type, valid_restrict_types
+            )
+        )
+    # Check if correct number of arguments passed
+    if restriction_type == "interface":
+        if len(restriction_value) != 2:
+            raise ValueError(
+                "Invalid number of arguments passed."
+                "{} arguments for restriction type {}"
+                "were passed.  2 are required".format(
+                    len(restriction_value), restriction_type
+                )
+            )
+    else:
+        if not isinstance(restriction_value, (float, int)):
+            raise TypeError(
+                "Restriction type is {}. A"
+                ' single argument of type "int"'
+                'or "float" should be passed'.format(restriction_type)
+            )

mosdef_cassandra/examples/__init__.py

@@ -5,3 +5,4 @@
 
 from .nvt_mixture import run_nvt_mixture
 from .gcmc_adsorption import run_gcmc_adsorption
+from .gcmc_restricted import run_gcmc_restricted

mosdef_cassandra/examples/gcmc_restricted.py

@@ -0,0 +1,44 @@
+import mbuild
+import foyer
+import mosdef_cassandra as mc
+
+
+def run_gcmc_restricted():
+
+    # Use mbuild to create molecules
+    methane = mbuild.load("C", smiles=True)
+
+    # Create an empty mbuild.Box
+    box = mbuild.Box(lengths=[5.0, 5.0, 5.0])
+
+    # Load forcefields
+    oplsaa = foyer.forcefields.load_OPLSAA()
+
+    # Use foyer to apply forcefields
+    methane_ff = oplsaa.apply(methane)
+
+    # Create box and species list
+    box_list = [box]
+    species_list = [methane_ff]
+
+    mols_to_add = [[10]]
+
+    system = mc.System(box_list, species_list, mols_to_add=mols_to_add)
+    moves = mc.Moves("gcmc", species_list)
+
+    # Specify restricted insertions
+    moves.add_restricted_insertions(species_list, [["sphere"]], [[20]])
+
+    mc.run(
+        system=system,
+        moves=moves,
+        run_type="equilibration",
+        run_length=100,
+        temperature=300.0,
+        chemical_potentials=[-35.0],
+        prop_freq=10,
+    )
+
+
+if __name__ == "__main__":
+    run_gcmc_restricted()

mosdef_cassandra/tests/test_examples.py

@@ -146,3 +146,26 @@ def test_run_gcmc_adsorption(self):
                     if "Cassandra simulation complete" in line:
                         completed = True
                 assert completed
+
+    def test_run_gcmc_restricted(self):
+        with temporary_directory() as tmp_dir:
+            with temporary_cd(tmp_dir):
+                ex.run_gcmc_restricted()
+                log_files = sorted(
+                    glob.glob("./mosdef_cassandra*.log"), key=os.path.getmtime
+                )
+                log_file = log_files[-1]
+                log_data = []
+                save_data = False
+                with open(log_file) as log:
+                    for line in log:
+                        if "CASSANDRA STANDARD" in line:
+                            save_data = True
+                        if save_data == True:
+                            log_data.append(line)
+
+                completed = False
+                for line in log_data:
+                    if "Cassandra simulation complete" in line:
+                        completed = True
+                assert completed

mosdef_cassandra/tests/test_moves.py

@@ -1,6 +1,7 @@
 import pytest
 
 import mosdef_cassandra as mc
+import warnings
 from mosdef_cassandra.tests.base_test import BaseTest
 
 
@@ -115,6 +116,22 @@ def test_ensemble_gcmc(self, methane_oplsaa):
         assert moves.sp_insertable[0] == True
         assert moves.sp_prob_regrow[0] == 1.0
 
+    @pytest.mark.parametrize(
+        "typ,value",
+        [
+            ("slitpore", 1),
+            ("cylinder", 1),
+            ("sphere", 1),
+            ("interface", [1, 2]),
+        ],
+    )
+    def test_restricted_gcmc(self, methane_oplsaa, typ, value):
+        moves = mc.Moves("gcmc", [methane_oplsaa])
+        moves.add_restricted_insertions([methane_oplsaa], [[typ]], [[value]])
+
+        assert moves._restricted_type == [[typ]]
+        assert moves._restricted_value == [[value]]
+
     def test_ensemble_gemc(self, methane_oplsaa):
         moves = mc.Moves("gemc", [methane_oplsaa])
         assert moves.ensemble == "gemc"
@@ -149,6 +166,24 @@ def test_ensemble_gemc(self, methane_oplsaa):
         assert moves.sp_insertable[0] == True
         assert moves.sp_prob_regrow[0] == 1.0
 
+    @pytest.mark.parametrize(
+        "typ,value",
+        [
+            ("slitpore", 1),
+            ("cylinder", 1),
+            ("sphere", 1),
+            ("interface", [1, 2]),
+        ],
+    )
+    def test_restricted_gemc(self, methane_oplsaa, typ, value):
+        moves = mc.Moves("gemc", [methane_oplsaa])
+        moves.add_restricted_insertions(
+            [methane_oplsaa], [[None], [typ]], [[None], [value]]
+        )
+
+        assert moves._restricted_type == [[None], [typ]]
+        assert moves._restricted_value == [[None], [value]]
+
     def test_ensemble_gemcnpt(self, methane_oplsaa):
         moves = mc.Moves("gemc_npt", [methane_oplsaa])
         assert moves.ensemble == "gemc_npt"
@@ -184,6 +219,15 @@ def test_ensemble_gemcnpt(self, methane_oplsaa):
         assert moves.sp_insertable[0] == True
         assert moves.sp_prob_regrow[0] == 1.0
 
+    def test_restricted_gemc_npt(self, methane_oplsaa):
+        moves = mc.Moves("gemc_npt", [methane_oplsaa])
+        moves.add_restricted_insertions(
+            [methane_oplsaa], [[None], ["slitpore"]], [[None], [3]]
+        )
+
+        assert moves._restricted_type == [[None], ["slitpore"]]
+        assert moves._restricted_value == [[None], [3]]
+
     def test_single_site_nvt(self, methane_trappe):
 
         moves = mc.Moves("nvt", [methane_trappe])
@@ -468,3 +512,45 @@ def test_print_moves(self, methane_oplsaa):
 
         moves = mc.Moves("gemc", [methane_oplsaa])
         moves.print()
+
+    def test_invalid_ensemble_and_restriction(self, methane_oplsaa):
+        moves = mc.Moves("nvt", [methane_oplsaa])
+        with pytest.raises(ValueError, match=r"only valid"):
+            moves.add_restricted_insertions(
+                [methane_oplsaa], [["slitpore"]], [[1]]
+            )
+
+    @pytest.mark.parametrize(
+        "typ,value",
+        [("slitpore", [[1], [2]]), ("cylinder", [[None]]), (None, [[1]]),],
+    )
+    def test_value_error_restricted_type_and_value(
+        self, methane_oplsaa, typ, value
+    ):
+        moves = mc.Moves("gcmc", [methane_oplsaa])
+        with pytest.raises(ValueError):
+            moves.add_restricted_insertions([methane_oplsaa], [[typ]], value)
+
+    def test_type_error_restricted_type_and_value(self, methane_oplsaa):
+        moves = mc.Moves("gcmc", [methane_oplsaa])
+        with pytest.raises(TypeError):
+            moves.add_restricted_insertions(
+                [methane_oplsaa], ["cylinder"], [3]
+            )
+
+    def test_invalid_restricted_type_and_species(self, methane_oplsaa):
+        moves = mc.Moves("gcmc", [methane_oplsaa])
+        with pytest.raises(ValueError, match=r"Length of 'species'"):
+            moves.add_restricted_insertions(
+                [methane_oplsaa], [["slitpore", None]], [[1, None]]
+            )
+
+    def test_add_multiple_restricted_insertions(self, methane_oplsaa):
+        moves = mc.Moves("gcmc", [methane_oplsaa])
+        moves.add_restricted_insertions(
+            [methane_oplsaa], [["slitpore"]], [[3]]
+        )
+        with pytest.warns(None) as record:
+            moves.add_restricted_insertions(
+                [methane_oplsaa], [["cylinder"]], [[3]]
+            )