Packing fixes

mbuild/packing.py

@@ -1,7 +1,9 @@
 from __future__ import division
 
 import sys
+import os
 import tempfile
+import warnings
 from distutils.spawn import find_executable
 from subprocess import Popen, PIPE
 
@@ -37,7 +39,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, overlap=0.2, seed=12345, write_tempfile=False):
     """Fill a box with a compound using packmol.
 
     Two arguments of `n_compounds, box, and density` must be specified.
@@ -52,20 +55,34 @@ 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
+        Compound or list of compounds to be put in box.
     n_compounds : int or list of int
+        Number of compounds to be put in box.
     box : mb.Box
-    overlap : float, units nm
+        Box to be filled by compounds.
+    aspect_ratio : list of float
+        If a non-cubic box is desired, the ratio of box lengths in the x, y,
+        and z directions.
+    overlap : float, units nm, default=0.2
+        Minimum separation between atoms of different molecules.
+    edge : float, units nm, default=0.2
+        Buffer at the edge of the box to not place molecules. This is necessary
+        in some systems because PACKMOL does not account for periodic boundary
+        conditions in its optimizaiton.
     density : float, units kg/m^3
+        Target density for the system.
 
     Returns
     -------
     filled : mb.Compound
 
-    TODO : Support aspect ratios in generated boxes
-    TODO : Support ratios of n_compounds
     """
     if not PACKMOL:
         msg = "Packmol not found."
@@ -74,6 +91,9 @@ def fill_box(compound, n_compounds=None, box=None, density=None, overlap=0.2, se
                          "packmol.exe is on the path.")
         raise IOError(msg)
 
+    if write_tempfile:
+        original_dir = os.getcwd() # Avoid saving to a temp dir
+
     arg_count = 3 - [n_compounds, box, density].count(None)
     if arg_count != 2:
         msg = ("Exactly 2 of `n_compounds`, `box`, and `density` "
@@ -87,28 +107,53 @@ def fill_box(compound, n_compounds=None, box=None, density=None, overlap=0.2, se
     if n_compounds is not None and not isinstance(n_compounds, (list, set)):
         n_compounds = [n_compounds]
 
+    if compound is not None and n_compounds is not None:
+        if len(compound) != len(n_compounds):
+            msg = ("`compound` and `n_compounds` must be of equal length.")
+            raise ValueError(msg)
+
     if density is not None:
         if box is None and n_compounds is not None:
             total_mass = np.sum([n*np.sum([a.mass for a in c.to_parmed().atoms])
                 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)
+                n_compounds = list()
+                for c in compound_ratio:
+                    n_compounds.append(int(n_prototypes * c))
 
     # 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]
@@ -124,9 +169,20 @@ def fill_box(compound, n_compounds=None, box=None, density=None, overlap=0.2, se
 
     proc = Popen(PACKMOL, stdin=PIPE, stdout=PIPE, stderr=PIPE, universal_newlines=True)
     out, err = proc.communicate(input=input_text)
-    if err:
+    #import pdb; pdb.set_trace()
+    if 'WITHOUT PERFECT PACKING' in out:
+        msg = ("Packmol finished with imperfect packing. Using "
+               "the .pdb_FORCED file instead. This may not be a "
+               "sufficient packing result.")
+        warnings.warn(msg)
+        os.system('cp {0}_FORCED {0}'.format(filled_pdb))
+    if 'ERROR' in out:
         _packmol_error(out, err)
 
+    if write_tempfile:
+        os.system('cp {0} {1}'.format(filled_pdb,
+                                      os.path.join(original_dir, 'packmol_temp.pdb')))
+
     # Create the topology and update the coordinates.
     filled = Compound()
     for comp, m_compounds in zip(compound, n_compounds):
@@ -137,15 +193,20 @@ 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, write_tempfile=False):
     """Fill a region of a box with a compound using packmol.
 
     Parameters
     ----------
     compound : mb.Compound or list of mb.Compound
+        Compound or list of compounds to be put in box.
     n_compounds : int or list of int
+        Number of compounds to be put in box.
     region : mb.Box or list of mb.Box
+        Region to be filled by compounds.
     overlap : float
+        Minimum separation between atoms of different molecules.
 
     Returns
     -------
@@ -161,10 +222,19 @@ def fill_region(compound, n_compounds, region, overlap=0.2, seed=12345):
                          "packmol.exe is on the path.")
         raise IOError(msg)
 
+    if write_tempfile:
+        original_dir = os.getcwd() # Avoid saving to a temp dir
+
     if not isinstance(compound, (list, set)):
         compound = [compound]
     if not isinstance(n_compounds, (list, set)):
         n_compounds = [n_compounds]
+
+    if compound is not None and n_compounds is not None:
+        if len(compound) != len(n_compounds):
+            msg = ("`compound` and `n_compounds` must be of equal length.")
+            raise ValueError(msg)
+
     # See if region is a single region or list
     if isinstance(region, Box): # Cannot iterate over boxes
         region = [region]
@@ -194,6 +264,10 @@ def fill_region(compound, n_compounds, region, overlap=0.2, seed=12345):
     if err:
         _packmol_error(out, err)
 
+    if write_tempfile:
+        os.system('cp {0} {1}'.format(solvated_pdb,
+                                      os.path.join(original_dir, 'packmol_temp.pdb')))
+
     # Create the topology and update the coordinates.
     filled = Compound()
     for comp, m_compounds in zip(compound, n_compounds):
@@ -203,16 +277,22 @@ 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, write_tempfile=False):
     """Solvate a compound in a box of solvent using packmol.
 
     Parameters
     ----------
     solute : mb.Compound
+        Compound to be placed in a box and solvated.
     solvent : mb.Compound
+        Compound to solvate the box.
     n_solvent : int
+        Number of solvents to be put in box.
     box : mb.Box
+        Box to be filled by compounds.
     overlap : float
+        Minimum separation between atoms of different molecules.
 
     Returns
     -------
@@ -222,17 +302,27 @@ def solvate(solute, solvent, n_solvent, box, overlap=0.2, seed=12345):
     if not PACKMOL:
         raise IOError("Packmol not found")
 
+    if write_tempfile:
+        original_dir = os.getcwd() # Avoid saving to a temp dir
+
     box = _validate_box(box)
     if not isinstance(solvent, (list, set)):
         solvent = [solvent]
     if not isinstance(n_solvent, (list, set)):
         n_solvent = [n_solvent]
 
+    if len(solvent) != len(n_solvent):
+        msg = ("`n_solvent` and `n_solvent` must be of equal length.")
+        raise ValueError(msg)
+
     # In angstroms for packmol.
     box_mins = box.mins * 10
     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]
@@ -254,6 +344,10 @@ def solvate(solute, solvent, n_solvent, box, overlap=0.2, seed=12345):
     if err:
         _packmol_error(out, err)
 
+    if write_tempfile:
+        os.system('cp {0} {1}'.format(solvated_pdb,
+                                      os.path.join(original_dir, 'packmol_temp.pdb')))
+
     # Create the topology and update the coordinates.
     solvated = Compound()
     solvated.add(solute)

mbuild/tests/test_packing.py

@@ -1,3 +1,5 @@
+import os
+
 import pytest
 import numpy as np
 
@@ -16,11 +18,24 @@ 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])
+        n_ethane = len([c for c in filled.children if c.name == 'Ethane'])
+        n_water = len([c for c in filled.children if c.name == 'H2O'])
+        assert n_water / n_ethane == 2
+
     def test_fill_region(self, h2o):
         filled = mb.fill_region(h2o, n_compounds=50,
                                 region=[3, 2, 2, 4, 4, 3])
@@ -101,3 +116,24 @@ def test_wrong_box(self, h2o):
             filled = mb.fill_box(h2o, n_compounds=50, box=[2, 2])
         with pytest.raises(MBuildError):
             filled = mb.fill_box(h2o, n_compounds=50, box=[2, 2, 2, 2])
+
+    def test_bad_args(self, h2o):
+        with pytest.raises(ValueError):
+            mb.fill_box(h2o, n_compounds=10)
+        with pytest.raises(ValueError):
+            mb.fill_box(h2o, density=1000)
+        with pytest.raises(ValueError):
+            mb.fill_box(h2o, box=[2, 2, 2])
+        with pytest.raises(ValueError):
+            mb.fill_box(h2o, n_compounds=10, density=1000, box=[2, 2, 2])
+        with pytest.raises(ValueError):
+            mb.fill_box(compound=[h2o, h2o], n_compounds=[10], density=1000)
+        with pytest.raises(ValueError):
+            mb.solvate(solute=h2o, solvent=[h2o], n_solvent=[10, 10], box=[2, 2, 2])
+        with pytest.raises(ValueError):
+            mb.fill_region(h2o, n_compounds=[10, 10], region=[2, 2, 2, 4, 4, 4])
+
+    def test_write_tempfile(self, h2o):
+        cwd = os.getcwd() # Must keep track of the temp dir that pytest creates
+        filled = mb.fill_box(h2o, n_compounds=50, box=[2, 2, 2, 4, 4, 4], write_tempfile=True)
+        assert os.path.isfile(os.path.join(cwd, 'packmol_temp.pdb'))