Specify individual axes to constrain within PACKMOL in packing.py

mbuild/packing.py

@@ -44,16 +44,49 @@
 end structure
 """
 
-PACKMOL_CONSTRAIN = """
-constrain_rotation x 0. 0.
-constrain_rotation y 0. 0.
-constrain_rotation z 0. 0.
-"""
+def packmol_constrain(fix_orientation):
+    """
+    Provides information to PACKMOL about which axes to constrain rotation.
+    fix_orientation is generated within the `fill_box` function
+    Parameters
+    ----------
+    fix_orientation : iterable of booleans, directions (x, y, z) about
+                      which to constrain rotation. If True, no rotation
+                      in that direction
+    Returns
+    -------
+    None if fix_orientation = (False, False, False)
+    string
+        rotation constraints to be added to PACKMOL input text
+    """
+    constraints = ['constrain_rotation x 0. 0.',
+                   'constrain_rotation y 0. 0.',
+                   'constrain_rotation z 0. 0.']
+    CONSTRAIN = """
+    {}
+    {}
+    {}
+    """
+    # Handles instances that are not iterable; defaults to True/Fales for all axes
+    if fix_orientation == True:
+        fix_orientation=(True, True, True)
+    if fix_orientation == False:
+        fix_orientation=(False, False, False)
+
+    if not any(fix_orientation):
+        return None
+    final_constraints = list()
+    for i, val in enumerate(fix_orientation):
+        if val:
+            final_constraints.append(constraints[i])
+        else:
+            final_constraints.append("")
+    return CONSTRAIN.format(*final_constraints)
 
 
 def fill_box(compound, n_compounds=None, box=None, density=None, overlap=0.2,
              seed=12345, edge=0.2, compound_ratio=None,
-             aspect_ratio=None, fix_orientation=False, temp_file=None,
+             aspect_ratio=None, fix_orientation=(False, False, False), temp_file=None,
              update_port_locations=False):
     """Fill a box with a `mbuild.compound` or `Compound`s using PACKMOL.
 
@@ -104,9 +137,10 @@ def fill_box(compound, n_compounds=None, box=None, density=None, overlap=0.2,
     aspect_ratio : list of float
         If a non-cubic box is desired, the ratio of box lengths in the x, y,
         and z directions.
-    fix_orientation : bool or list of bools
+    fix_orientation : iterable of booleans or list of iterables
         Specify that compounds should not be rotated when filling the box,
-        default=False.
+        Rotation contraints are specified for each individual axis (x, y, z)
+        default=(False, False, False)
     temp_file : str, default=None
         File name to write PACKMOL's raw output to.
     update_port_locations : bool, default=False
@@ -120,7 +154,6 @@ def fill_box(compound, n_compounds=None, box=None, density=None, overlap=0.2,
     """
     # check that the user has the PACKMOL binary on their PATH
     _check_packmol(PACKMOL)
-
     arg_count = 3 - [n_compounds, box, density].count(None)
     if arg_count != 2:
         msg = ("Exactly 2 of `n_compounds`, `box`, and `density` "
@@ -133,8 +166,16 @@ def fill_box(compound, n_compounds=None, box=None, density=None, overlap=0.2,
         compound = [compound]
     if n_compounds is not None and not isinstance(n_compounds, (list, set)):
         n_compounds = [n_compounds]
+
     if not isinstance(fix_orientation, (list, set)):
         fix_orientation = [fix_orientation]*len(compound)
+    try:
+        iter(fix_orientation[0])
+    except:
+        warnings.warn("fix_orientation can be given as an iterable of True/False values "
+             "for each compound to fix rotations about each x,y,z axis individually. "
+             "Using a single instance of True/False defaults to (True,True,True) "
+             "and (Fale,False,False) respectively")
 
     if compound is not None and n_compounds is not None:
         if len(compound) != len(n_compounds):
@@ -205,11 +246,12 @@ def fill_box(compound, n_compounds=None, box=None, density=None, overlap=0.2,
             compound_xyz_list.append(compound_xyz)
 
             comp.save(compound_xyz.name, overwrite=True)
+            PACKMOL_CONSTRAIN = packmol_constrain(rotate) 
             input_text += PACKMOL_BOX.format(compound_xyz.name, m_compounds,
                                              box_mins[0], box_mins[1],
                                              box_mins[2], box_maxs[0],
                                              box_maxs[1], box_maxs[2],
-                                             PACKMOL_CONSTRAIN if rotate else "")
+                                             PACKMOL_CONSTRAIN if PACKMOL_CONSTRAIN else "")
 
         _run_packmol(input_text, filled_xyz, temp_file)
         # Create the topology and update the coordinates.

mbuild/tests/test_packing.py

@@ -216,6 +216,22 @@ def test_no_rotate(self, h2o):
         w0 -= w0.sum(0) / len(w0)
         w1 -= w1.sum(0) / len(w1)
         assert np.isclose(w0, w1).all() is not True
+
+    def specify_axis(self):
+        arguments = [(True, False, False),
+                    (False, True, False),
+                    (False, False, True),
+                    (True, True, False),
+                    (False, True, True),
+                    (True, False, True)]
+        constraints = [["constrain_rotation x"],
+                    ["constrain_rotation y"],
+                    ["constrain_rotation z"],
+                    ["constrain_rotation x", "constrain_rotation y"],
+                    ["constrain_rotation y", "constrain_rotation z"],
+                    ["constrain_rotation x", "constrain_rotation z"]]        
+        for i, arg in enumerate(arguments):
+            assert all(c in mb.packing.packmol_constrain(arg) for c in constraints[i])
 
     def test_remove_port(self):
         from mbuild.lib.recipes import Alkane