Add restricted insertions attribute for moves

mosdef_cassandra/core/moves.py

@@ -58,6 +58,14 @@ def __init__(self, ensemble, species_topologies):
         else:
             self._n_boxes = 2
 
+        # Set 'restricted_typed' and 'restricted_value'
+        if self.ensemble in ['gcmc', 'gemc', 'gemc_npt']:
+            self.restricted_type = [[None] * self._n_species] * self._n_boxes
+            self.restricted_value = [[None] * self._n_species] * self._n_boxes
+        else:
+            self.restricted_type = None
+            self.restricted_value = None
+
         # Define default probabilities
         # Most are ensemble-dependent
         self.prob_angle = 0.0
@@ -173,6 +181,94 @@ def __init__(self, ensemble, species_topologies):
                 self.prob_translate += self.prob_rotate
                 self.prob_rotate = 0.0
 
+    @property
+    def restricted_type(self):
+        return self._restricted_type
+
+    @restricted_type.setter
+    def restricted_type(self, restricted_type):
+        if restricted_type:
+            if self.ensemble == 'gcmc':
+                if len(restricted_type) != 1:
+                    raise ValueError("{} ensemble only has 1 box"
+                    ' but restricted_insertion of length {}'
+                    ' was passed.'.format(
+                        self.ensemble,
+                        len(restricted_type)))
+                if not isinstance(restricted_type[0], (list, tuple)):
+                    raise TypeError("'restricted_type' specified for"
+                    " a box must be an iterable.")
+                if len(restricted_type[0]) != self._n_species:
+                    raise ValueError("{} species specified"
+                    ' but restricted_type for {}'
+                    ' species was passed.'.format(
+                        self._n_species,
+                        len(restricted_type[0])))
+            elif self.ensemble in ['gemc', 'gemc_npt']:
+                if len(restricted_type) != 2:
+                    raise ValueError("{} ensemble requires 2 boxes"
+                    ' but restricted_type of length {}'
+                    ' was passed.'.format(
+                        self.ensemble,
+                        len(restricted_type)))
+                if len(restricted_type[0]) != self._n_species:
+                    raise ValueError("{} species specified"
+                    ' but restricted_typed for only {}'
+                    ' species was passed.'.format(
+                        self._n_species,
+                        len(restricted_type[0])))
+            else:
+                raise ValueError("Restricted insertions are only valid"
+                    ' for GCMC, GEMC, and GEMC-NPT ensembles.'
+                    ' The {} ensemble was passed.'.format(
+                        self.ensemble))
+
+            self._restricted_type = restricted_type
+        else:
+            self._restricted_type = restricted_type
+
+    @property
+    def restricted_value(self):
+        return self._restricted_value
+
+    @restricted_value.setter
+    def restricted_value(self, restricted_value):
+        if restricted_value:
+            if len(self.restricted_type) != len(restricted_value):
+                raise ValueError("'restricted_type' is a list of length"
+                        " {} but 'restricted_value' is a list of length"
+                        " {}.".format(
+                            len(self.restricted_type),
+                            len(restricted_value))
+                        )
+            if not isinstance(restricted_value[0], (list, tuple)):
+                raise TypeError("'restricted_value' specified for"
+                " a box must be an iterable.")
+            if len(restricted_value[0]) != self._n_species:
+                raise ValueError("{} species specified"
+                ' but restricted_value for only {}'
+                ' species was passed.'.format(
+                    self._n_species,
+                    len(restricted_value[0])))
+            else:
+                for box in range(self._n_boxes):
+                    for typ, value in zip(self.restricted_type[box], restricted_value[box]):
+                        if typ and value:
+                            _check_restriction_type(typ, value)
+                        elif typ == None and value == None:
+                            pass
+                        else:
+                            raise ValueError("'None' must be passed to both"
+                            " 'restricted_value' and 'restricted_type'.")
+
+                self._restricted_value = restricted_value
+        else:
+            if self.restricted_type:
+                raise ValueError("'None' must be passed to both"
+                " 'restricted_value' and 'restricted_type'.")
+            else:
+                self._restricted_value = restricted_value
+
     @property
     def ensemble(self):
         return self._ensemble
@@ -676,5 +772,58 @@ def print(self):
             contents += "Box {box}: {max_vol}\n".format(
                 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 {}.  Only'
+                'a single argument of type "int"'
+                'or "float" should be passed'.format(
+                    restriction_type))

mosdef_cassandra/tests/test_moves.py

@@ -115,6 +115,16 @@ 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.restricted_type = [[typ]]
+        moves.restricted_value = [[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 +159,16 @@ 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.restricted_type = [[None], [typ]]
+        moves.restricted_value = [[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 +204,14 @@ 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.restricted_type = [[None], ['slitpore']]
+        moves.restricted_value = [[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])

mosdef_cassandra/tests/test_writers.py

@@ -1557,3 +1557,77 @@ def test_cbmc_info(self, onecomp_system, twobox_system):
                 cbmc_kappa_dih=5,
                 cbmc_rcut=[],
             )
+
+    @pytest.mark.parametrize('typ,value', [('slitpore', 1), ('cylinder', 1), ('sphere',
+        1), ('interface', [1,2])])
+    def test_write_restricted_gcmc(self, gcmc_system, typ, value):
+        (system, moves) = gcmc_system
+        moves.restricted_type = [[None, typ]]
+        moves.restricted_value = [[None, value]]
+        inp_data = generate_input(
+            system=system,
+            moves=moves,
+            run_type="equilibration",
+            run_length=500,
+            temperature=300.0,
+            chemical_potentials=["none", 10.0],
+        )
+
+        if typ == 'interface':
+            assert "\nrestricted_insertion {} {} {}\n".format(typ, value[0], value[1]) in inp_data
+        else:
+            assert "\nrestricted_insertion {} {}\n".format(typ, value) in inp_data
+
+    @pytest.mark.parametrize('typ,value', [('slitpore', 3), ('cylinder', 3), ('sphere',
+        3), ('interface', [3,5])])
+    def test_fail_restricted_gcmc(self, gcmc_system, typ, value):
+        (system, moves) = gcmc_system
+        moves.restricted_type = [[None, typ]]
+        moves.restricted_value = [[None, value]]
+        with pytest.raises(ValueError, match=r"Restricted insertion"):
+            inp_data = generate_input(
+                system=system,
+                moves=moves,
+                run_type="equilibration",
+                run_length=500,
+                temperature=300.0,
+                chemical_potentials=["none", 10.0],
+            )
+
+    @pytest.mark.parametrize('typ,value', [('slitpore', 1), ('cylinder', 1), ('sphere',
+        1), ('interface', [1,2])])
+    def test_write_restricted_gemc_npt(self, twocomptwobox_system, typ, value):
+        (system, moves) = twocomptwobox_system
+        moves.restricted_type = [[None, None], [None, typ]]
+        moves.restricted_value = [[None, None], [None, value]]
+        inp_data = generate_input(
+            system=system,
+            moves=moves,
+            run_type="equilibration",
+            run_length=500,
+            temperature=300.0,
+            pressure=1,
+            chemical_potentials=["none", 10.0],
+        )
+
+        if typ == 'interface':
+            assert "\nrestricted_insertion {} {} {}\n".format(typ, value[0], value[1]) in inp_data
+        else:
+            assert "\nrestricted_insertion {} {}\n".format(typ, value) in inp_data
+
+    @pytest.mark.parametrize('typ,value', [('slitpore', 6), ('cylinder', 6), ('sphere',
+        6), ('interface', [1,7])])
+    def test_fail_restricted_gemc_npt(self, twocomptwobox_system, typ, value):
+        (system, moves) = twocomptwobox_system
+        moves.restricted_type = [[None, None], [None, typ]]
+        moves.restricted_value = [[None, None], [None, value]]
+        with pytest.raises(ValueError, match=r"Restricted insertion"):
+            inp_data = generate_input(
+                system=system,
+                moves=moves,
+                run_type="equilibration",
+                run_length=500,
+                temperature=300.0,
+                pressure=1,
+                chemical_potentials=["none", 10.0],
+            )