Packing fixes

mbuild/packing.py

@@ -37,7 +37,8 @@
 """
 
 
-def fill_box(compound, n_compounds=None, box=None, density=None, overlap=0.2, seed=12345):
+def fill_box(compound, n_compounds=None, box=None, aspect_ratio=None,
+        density=None, compound_ratio=None, overlap=0.2, edge=0.2, seed=12345):
     """Fill a box with a compound using packmol.
 
     Two arguments of `n_compounds, box, and density` must be specified.
@@ -52,20 +53,24 @@ def fill_box(compound, n_compounds=None, box=None, density=None, overlap=0.2, se
     If `box` and `density` are not None, the corresponding number of
     compounds will be calculated internally.
 
+    For the cases in which `box` is not specified but generated internally,
+    the default behavior is to calculate a cubic box. Optionally,
+    `aspect_ratio` can be passed to generate a non-cubic box.
+
     Parameters
     ----------
     compound : mb.Compound or list of mb.Compound
     n_compounds : int or list of int
     box : mb.Box
+    aspect_ratio : list of float
     overlap : float, units nm
+    edge : float, units nm
     density : float, units kg/m^3
 
     Returns
     -------
     filled : mb.Compound
 
-    TODO : Support aspect ratios in generated boxes
-    TODO : Support ratios of n_compounds
     """
     if not PACKMOL:
         msg = "Packmol not found."
@@ -93,22 +98,42 @@ def fill_box(compound, n_compounds=None, box=None, density=None, overlap=0.2, se
                 for c,n in zip(compound, n_compounds)])
             # Conversion from (amu/(kg/m^3))**(1/3) to nm
             L = (total_mass/density)**(1/3)*1.1841763
-            box = _validate_box(Box(3*[L]))
-        if n_compounds is None and box is not None:
-            if len(compound) > 1:
-                msg = ("Determing `n_compounds` from `density` and `box` "
-                    "currently only supported for systems with one "
-                    "compound type.")
-                raise ValueError(msg)
+            if aspect_ratio is None:
+                box = _validate_box(Box(3*[L]))
             else:
+                L *= np.prod(aspect_ratio) ** (-1/3)
+                box = _validate_box(Box([val*L for val in aspect_ratio]))
+        if n_compounds is None and box is not None:
+            if len(compound) == 1:
                 compound_mass = np.sum([a.mass for a in compound[0].to_parmed().atoms])
                 # Conversion from kg/m^3 / amu * nm^3 to dimensionless units
                 n_compounds = [int(density/compound_mass*np.prod(box.lengths)*.60224)]
+            else:
+                if compound_ratio is None:
+                    msg = ("Determing `n_compounds` from `density` and `box` "
+                    "for systems with more than one compound type requires"
+                    "`compound_ratio`")
+                    raise ValueError(msg)
+                if len(compound) != len(compound_ratio):
+                    msg = ("Length of `compound_ratio` must equal length of "
+                        "`compound`")
+                    raise ValueError(msg)
+                prototype_mass = 0
+                for c, r in zip(compound, compound_ratio):
+                    prototype_mass += r * np.sum([a.mass for a in c.to_parmed().atoms])
+                # Conversion from kg/m^3 / amu * nm^3 to dimensionless units
+                n_prototypes = int(density/prototype_mass*np.prod(box.lengths)*.60224)
+                print(n_prototypes)
+                print(compound_ratio)
+                n_compounds = n_prototypes * compound_ratio
 
     # In angstroms for packmol.
     box_mins = box.mins * 10
     box_maxs = box.maxs * 10
     overlap *= 10
+
+    # Apply edge buffer 
+    box_maxs = box_maxs - edge * 10
 
     # Build the input file for each compound and call packmol.
     filled_pdb = tempfile.mkstemp(suffix='.pdb')[1]
@@ -137,7 +162,7 @@ def fill_box(compound, n_compounds=None, box=None, density=None, overlap=0.2, se
     return filled
 
 
-def fill_region(compound, n_compounds, region, overlap=0.2, seed=12345):
+def fill_region(compound, n_compounds, region, overlap=0.2, edge=0.2, seed=12345):
     """Fill a region of a box with a compound using packmol.
 
     Parameters
@@ -203,7 +228,7 @@ def fill_region(compound, n_compounds, region, overlap=0.2, seed=12345):
     return filled
 
 
-def solvate(solute, solvent, n_solvent, box, overlap=0.2, seed=12345):
+def solvate(solute, solvent, n_solvent, box, overlap=0.2, edge=0.2, seed=12345):
     """Solvate a compound in a box of solvent using packmol.
 
     Parameters
@@ -233,6 +258,9 @@ def solvate(solute, solvent, n_solvent, box, overlap=0.2, seed=12345):
     box_maxs = box.maxs * 10
     overlap *= 10
     center_solute = (box_maxs + box_mins) / 2
+
+    # Apply edge buffer 
+    box_maxs = box_maxs - edge * 10
 
     # Build the input file for each compound and call packmol.
     solvated_pdb = tempfile.mkstemp(suffix='.pdb')[1]

mbuild/tests/test_packing.py

@@ -16,11 +16,21 @@ def test_fill_box_density_box(self, h2o):
         filled = mb.fill_box(h2o, n_compounds=1000, density=1000)
         assert [3.1042931 < period < 3.1042932 for period in filled.periodicity]
 
+    def test_fill_box_aspect_ratio(self, h2o):
+        filled = mb.fill_box(h2o, n_compounds=1000,
+                density=1000, aspect_ratio=[1, 2, 1])
+        assert filled.periodicity[0]/filled.periodicity[1] == 0.5
+        assert filled.periodicity[1]/filled.periodicity[2] == 2
+
     def test_fill_box_density_n_compounds(self, h2o):
         filled = mb.fill_box(h2o, density=1000,
                              box=mb.Box([3.1042931, 3.1042931, 3.1042931]))
         assert filled.n_particles == 3000
 
+    def test_fill_box_compound_ratio(self, h2o, ethane):
+        filled = mb.fill_box(compound=[h2o, ethane], density=800,
+                compound_ratio=[2, 1], box=[2, 2, 2, 4, 4, 4])
+
     def test_fill_region(self, h2o):
         filled = mb.fill_region(h2o, n_compounds=50,
                                 region=[3, 2, 2, 4, 4, 3])