Merge pull request #501 from justinGilmer/new_plugin

mattwthompson · web-flow · commit d80f38ba7ab4 · 2019-04-11T10:33:21.000-05:00
Add recipe/plugin architecture
diff --git a/mbuild/__init__.py b/mbuild/__init__.py
@@ -5,7 +5,6 @@
 from mbuild.pattern import *
 from mbuild.packing import *
 from mbuild.port import Port
-from mbuild.recipes import *
 from mbuild.lattice import Lattice
-
+from mbuild.recipes import recipes
 from mbuild.version import version
diff --git a/mbuild/examples/alkane/alkane.py b/mbuild/examples/alkane/alkane.py
@@ -3,6 +3,7 @@
 
 from mbuild.lib.moieties import CH2
 from mbuild.lib.moieties import CH3
+from mbuild.lib.recipes import Polymer
 
 
 class Alkane(mb.Compound):
@@ -24,7 +25,7 @@ def __init__(self, n=3, cap_front=True, cap_end=True):
             n += 1
         if not cap_end:
             n += 1
-        chain = mb.Polymer(CH2(), n=n-2, port_labels=('up', 'down'))
+        chain = mb.recipes.Polymer(CH2(), n=n-2, port_labels=('up', 'down'))
         self.add(chain, 'chain')
 
         if cap_front:
@@ -43,4 +44,4 @@ def __init__(self, n=3, cap_front=True, cap_end=True):
             # Hoist port label to Alkane level.
             self.add(chain['down'], 'down', containment=False)
 
-# -- ==alkane== --
+# -- ==alkane== --
diff --git a/mbuild/examples/alkane_monolayer/alkane_monolayer.py b/mbuild/examples/alkane_monolayer/alkane_monolayer.py
@@ -6,10 +6,11 @@
 
 from mbuild.lib.surfaces import Betacristobalite
 from mbuild.lib.atoms import H
+from mbuild.lib.recipes import Monolayer
 from mbuild.examples.alkane_monolayer.alkylsilane import AlkylSilane
 
 
-class AlkaneMonolayer(mb.Monolayer):
+class AlkaneMonolayer(mb.recipes.Monolayer):
     """An akylsilane monolayer on beta-cristobalite. """
 
     def __init__(self, pattern, tile_x=1, tile_y=1, chain_length=10):
@@ -34,4 +35,4 @@ def __init__(self, pattern, tile_x=1, tile_y=1, chain_length=10):
                                               pattern=pattern, tile_x=tile_x,
                                               tile_y=tile_y)
 
-# -- ==alkane_monolayer== --
+# -- ==alkane_monolayer== --
diff --git a/mbuild/examples/pmpc/brush.py b/mbuild/examples/pmpc/brush.py
@@ -4,6 +4,7 @@
 
 from mbuild.lib.moieties import Silane
 from mbuild.lib.moieties import CH3
+from mbuild.lib.recipes import Polymer
 from mbuild.examples.pmpc.mpc import MPC
 from mbuild.examples.pmpc.initiator import Initiator
 
@@ -16,7 +17,7 @@ def __init__(self, chain_length=4, alpha=pi/4):
         # Add parts
         self.add(Silane(), label='silane')
         self.add(Initiator(), label='initiator')
-        self.add(mb.Polymer(MPC(alpha=alpha), n=chain_length,
+        self.add(mb.recipes.Polymer(MPC(alpha=alpha), n=chain_length,
                             port_labels=('up', 'down')), label='pmpc')
         self.add(CH3(), label='methyl')
 
diff --git a/mbuild/examples/pmpc/pmpc_brush_layer.py b/mbuild/examples/pmpc/pmpc_brush_layer.py
@@ -7,13 +7,13 @@
 from numpy import pi
 
 import mbuild as mb
-
 from mbuild.lib.atoms import H
 from mbuild.lib.surfaces import Betacristobalite
+from mbuild.lib.recipes import Monolayer
 from mbuild.examples.pmpc.brush import Brush
 
 
-class PMPCLayer(mb.Monolayer):
+class PMPCLayer(mb.recipes.Monolayer):
     """Create a layer of grafted pMPC brushes on a beta-cristobalite surface."""
     def __init__(self, pattern, tile_x=1, tile_y=1, chain_length=4, alpha=pi / 4):
         surface = Betacristobalite()
@@ -23,4 +23,4 @@ def __init__(self, pattern, tile_x=1, tile_y=1, chain_length=4, alpha=pi / 4):
                                         pattern=pattern, tile_x=tile_x,
                                         tile_y=tile_y)
 
-# -- ==pmpc_brush_layer== --
+# -- ==pmpc_brush_layer== --
diff --git a/mbuild/lib/recipes/__init__.py b/mbuild/lib/recipes/__init__.py
@@ -0,0 +1,4 @@
+from mbuild.lib.recipes.polymer import Polymer
+from mbuild.lib.recipes.monolayer import Monolayer
+from mbuild.lib.recipes.tiled_compound import TiledCompound
+from mbuild.lib.recipes.silica_interface import SilicaInterface
diff --git a/mbuild/lib/recipes/monolayer.py b/mbuild/lib/recipes/monolayer.py
@@ -4,6 +4,7 @@
 import numpy as np
 
 import mbuild as mb
+from mbuild.lib.recipes.tiled_compound import TiledCompound
 
 __all__ = ['Monolayer']
 
@@ -37,15 +38,15 @@ def __init__(self, surface, chains, fractions=None, backfill=None, pattern=None,
         super(Monolayer, self).__init__()
 
         # Replicate the surface.
-        tiled_compound = mb.TiledCompound(surface, n_tiles=(tile_x, tile_y, 1))
+        tiled_compound = mb.recipes.TiledCompound(surface, n_tiles=(tile_x, tile_y, 1))
         self.add(tiled_compound, label='tiled_surface')
 
         if pattern is None:  # Fill the surface.
             pattern = mb.Random2DPattern(len(tiled_compound.referenced_ports()))
 
         if isinstance(chains, mb.Compound):
             chains = [chains]
-        
+
         if fractions:
             fractions = list(fractions)
             if len(chains) != len(fractions):
diff --git a/mbuild/lib/recipes/polymer.py b/mbuild/lib/recipes/polymer.py
diff --git a/mbuild/lib/recipes/silica_interface.py b/mbuild/lib/recipes/silica_interface.py
@@ -5,6 +5,7 @@
 
 import mbuild as mb
 import numpy as np
+from mbuild.lib.recipes import TiledCompound
 
 
 class SilicaInterface(mb.Compound):
@@ -60,13 +61,13 @@ def _cleave_interface(self, bulk_silica, tile_x, tile_y, thickness):
         """
         O_buffer = self._O_buffer
         tile_z = int(math.ceil((thickness + 2*O_buffer) / bulk_silica.periodicity[2]))
-        bulk = mb.TiledCompound(bulk_silica, n_tiles=(tile_x, tile_y, tile_z))
+        bulk = mb.recipes.TiledCompound(bulk_silica, n_tiles=(tile_x, tile_y, tile_z))
 
         interface = mb.Compound(periodicity=(bulk.periodicity[0],
                                              bulk.periodicity[1],
                                              0.0))
         for i, particle in enumerate(bulk.particles()):
-            if ((particle.name == 'Si' and O_buffer < particle.pos[2] < (thickness + O_buffer)) or 
+            if ((particle.name == 'Si' and O_buffer < particle.pos[2] < (thickness + O_buffer)) or
                     (particle.name == 'O' and particle.pos[2] < (thickness + 2*O_buffer))):
                 interface_particle = mb.Compound(name=particle.name, pos=particle.pos)
                 interface.add(interface_particle, particle.name + "_{}".format(i))
@@ -153,5 +154,5 @@ def _adjust_stoichiometry(self):
 
 if __name__ == "__main__":
     from mbuild.lib.bulk_materials import AmorphousSilica
-    silica_interface = mb.SilicaInterface(bulk_silica=AmorphousSilica(), thickness=1.2)
+    silica_interface = mb.recipes.SilicaInterface(bulk_silica=AmorphousSilica(), thickness=1.2)
     silica_interface.save('silica_interface.mol2', show_ports=True)
diff --git a/mbuild/lib/recipes/tiled_compound.py b/mbuild/lib/recipes/tiled_compound.py
diff --git a/mbuild/lib/surfaces/betacristobalite.py b/mbuild/lib/surfaces/betacristobalite.py
@@ -1,6 +1,7 @@
 import numpy as np
 
 import mbuild as mb
+from mbuild.lib.recipes import TiledCompound
 
 
 class Betacristobalite(mb.Compound):
@@ -31,7 +32,7 @@ def __init__(self):
         for particle in self.particles():
             if particle.name.startswith('O') and particle.pos[2] > 1.0:
                 count += 1
-                port = mb.Port(anchor=particle, orientation=[0, 0, 1], 
+                port = mb.Port(anchor=particle, orientation=[0, 0, 1],
                                separation=0.1)
                 self.add(port, 'port_{}'.format(count))
                 particle.name = 'O'  # Strip numbers required in .mol2 files.
@@ -40,5 +41,5 @@ def __init__(self):
 
 if __name__ == "__main__":
     single = Betacristobalite()
-    multiple = mb.TiledCompound(single, n_tiles=(2, 1, 1), name="tiled")
+    multiple = mb.recipes.TiledCompound(single, n_tiles=(2, 1, 1), name="tiled")
     multiple.save('betacristobalite.mol2')
diff --git a/mbuild/recipes/__init__.py b/mbuild/recipes/__init__.py
@@ -1,7 +1,8 @@
-__all__ = ['Polymer', 'TiledCompound', 'Monolayer', 'SilicaInterface']
+class Recipes(object):
+    pass
 
-
-from mbuild.recipes.polymer import Polymer
-from mbuild.recipes.tiled_compound import TiledCompound
-from mbuild.recipes.monolayer import Monolayer
-from mbuild.recipes.silica_interface import SilicaInterface
+recipes = Recipes()
+from pkg_resources import iter_entry_points
+available_methods = []
+for entry_point in iter_entry_points(group='mbuild.plugins', name=None):
+    setattr(recipes, entry_point.name, entry_point.load())
diff --git a/mbuild/tests/test_monolayer.py b/mbuild/tests/test_monolayer.py
@@ -3,6 +3,7 @@
 
 from mbuild.lib.surfaces import Betacristobalite
 from mbuild.lib.atoms import H
+from mbuild.lib.recipes import Polymer, Monolayer
 
 
 class TestMonolayer(BaseTest):
@@ -11,9 +12,12 @@ def test_monolayer(self, ch2):
         m = 8
         pattern = mb.Grid2DPattern(n, m)
 
-        chain = mb.Polymer(ch2, n=10)
-        monolayer = mb.Monolayer(surface=Betacristobalite(), chains=chain,
-                                 backfill=H(), pattern=pattern)
+        chain = mb.recipes.Polymer(ch2, n=10)
+        monolayer = mb.recipes.Monolayer(
+                                            surface=Betacristobalite(),
+                                            chains=chain,
+                                            backfill=H(),
+                                            pattern=pattern)
 
         assert monolayer.n_particles == 1900 + n * m * (10*3) + (100 - n*m)
         assert monolayer.n_bonds == 2400 + n * m * (10 * 2 + 9 + 1) + (100 - n * m)
@@ -23,10 +27,13 @@ def test_pattern_kwargs(self, ch2):
         m = 8
         pattern = mb.Grid2DPattern(n, m)
 
-        chain = mb.Polymer(ch2, n=10)
-        monolayer = mb.Monolayer(surface=Betacristobalite(), chains=H(),
-                                 guest_port_name='up', backfill=chain,
-                                 backfill_port_name='down', pattern=pattern)
+        chain = mb.recipes.Polymer(ch2, n=10)
+        monolayer = mb.recipes.Monolayer(surface=Betacristobalite(),
+                                            chains=H(),
+                                            guest_port_name='up',
+                                            backfill=chain,
+                                            backfill_port_name='down',
+                                            pattern=pattern)
 
         chains = 100 - (n*m)
 
@@ -39,9 +46,9 @@ def test_mixed_monolayer(self, ch2):
         pattern = mb.Grid2DPattern(n, m)
         fractions = [0.75,0.25]
 
-        chain_a = mb.Polymer(ch2, n=5)
-        chain_b = mb.Polymer(ch2, n=15)
-        monolayer = mb.Monolayer(surface=Betacristobalite(),
+        chain_a = mb.recipes.Polymer(ch2, n=5)
+        chain_b = mb.recipes.Polymer(ch2, n=15)
+        monolayer = mb.recipes.Monolayer(surface=Betacristobalite(),
                                  chains=[chain_a, chain_b],
                                  fractions=fractions,
                                  backfill=H(),
diff --git a/mbuild/tests/test_plugins.py b/mbuild/tests/test_plugins.py
@@ -0,0 +1,16 @@
+import pytest
+
+from mbuild.tests.base_test import BaseTest
+
+
+class TestPlugins(BaseTest):
+    def test_basic_import(self):
+        import mbuild.recipes
+
+    @pytest.mark.parametrize(
+        'recipe_name',
+        ['Monolayer', 'Polymer', 'SilicaInterface', 'TiledCompound'],
+    )
+    def test_recipes_contents(self, recipe_name):
+        import mbuild.recipes
+        recipe_name in dir(mbuild.recipes)
diff --git a/mbuild/tests/test_polymer.py b/mbuild/tests/test_polymer.py
@@ -2,20 +2,21 @@
 
 import mbuild as mb
 from mbuild.tests.base_test import BaseTest
+#from mbuild.lib.recipes import Polymer
 
 
 class TestPolymer(BaseTest):
 
     def test_polymer(self, ch2):
         n = 6
-        c6 = mb.Polymer(ch2, n=n)
+        c6 = mb.recipes.Polymer(ch2, n=n)
         assert c6.n_particles == n * 3
         assert c6.n_bonds == n * 2 + (n - 1)
 
     def test_block_copolymer(self, ch2, ester):
         n = 2
         sequence = 'ABBA'
-        abba = mb.Polymer([ch2, ester], sequence=sequence, n=n)
+        abba = mb.recipes.Polymer([ch2, ester], sequence=sequence, n=n)
 
         assert abba.n_particles == n * 3 * len(sequence)
         assert len(abba.children) == len(sequence) * n
diff --git a/mbuild/tests/test_silica_interface.py b/mbuild/tests/test_silica_interface.py
@@ -4,6 +4,7 @@
 
 import mbuild as mb
 from mbuild.lib.bulk_materials import AmorphousSilica
+from mbuild.lib.recipes import SilicaInterface
 from mbuild.tests.base_test import BaseTest
 
 
@@ -14,10 +15,10 @@ def test_silica_interface(self):
         tile_y = 1
         thickness = 0.6
 
-        interface = mb.SilicaInterface(bulk_silica=AmorphousSilica(),
-                                       tile_x=tile_x,
-                                       tile_y=tile_y,
-                                       thickness=thickness)
+        interface = mb.recipes.SilicaInterface(bulk_silica=AmorphousSilica(),
+                                                  tile_x=tile_x,
+                                                  tile_y=tile_y,
+                                                  thickness=thickness)
 
         thickness_tolerance = 0.05
         z = [atom.pos[2] for atom in interface.particles()
@@ -35,14 +36,14 @@ def test_seed(self):
         thickness = 0.6
         seed = 12345
 
-        interface1 = mb.SilicaInterface(bulk_silica=AmorphousSilica(),
+        interface1 = mb.recipes.SilicaInterface(bulk_silica=AmorphousSilica(),
                                         tile_x=tile_x,
                                         tile_y=tile_y,
                                         thickness=thickness,
                                         seed=seed)
         atom_names1 = np.array([atom.name for atom in interface1.particles()])
 
-        interface2 = mb.SilicaInterface(bulk_silica=AmorphousSilica(),
+        interface2 = mb.recipes.SilicaInterface(bulk_silica=AmorphousSilica(),
                                         tile_x=tile_x,
                                         tile_y=tile_y,
                                         thickness=thickness,
diff --git a/mbuild/tests/test_tiled_compound.py b/mbuild/tests/test_tiled_compound.py
@@ -2,6 +2,7 @@
 
 import mbuild as mb
 from mbuild.tests.base_test import BaseTest
+from mbuild.lib.recipes import TiledCompound
 
 
 class TestTiledCompound(BaseTest):
@@ -10,7 +11,7 @@ def test_2d_replication(self, betacristobalite):
         nx = 2
         ny = 3
         nz = 1
-        tiled = mb.TiledCompound(betacristobalite, [nx, ny, nz])
+        tiled = mb.recipes.TiledCompound(betacristobalite, [nx, ny, nz])
         assert tiled.n_particles == 1900 * nx * ny
         assert tiled.n_bonds == 2400 * nx * ny
         for at in tiled.particles():
@@ -29,7 +30,7 @@ def test_no_replication(self, betacristobalite):
         nx = 1
         ny = 1
         nz = 1
-        tiled = mb.TiledCompound(betacristobalite, [nx, ny, nz])
+        tiled = mb.recipes.TiledCompound(betacristobalite, [nx, ny, nz])
         assert tiled.n_particles == 1900 * nx * ny
         assert tiled.n_bonds == 2400 * nx * ny
 
@@ -38,13 +39,13 @@ def test_incorrect_periodicity(self, betacristobalite):
         ny = 3
         nz = 2
         with pytest.raises(ValueError):
-            mb.TiledCompound(betacristobalite, [nx, ny, nz])
+            mb.recipes.TiledCompound(betacristobalite, [nx, ny, nz])
 
     def test_negative_periodicity(self, betacristobalite):
         nx = -2
         ny = 3
         nz = 2
         with pytest.raises(ValueError):
-            mb.TiledCompound(betacristobalite, [nx, ny, nz])
+            mb.recipes.TiledCompound(betacristobalite, [nx, ny, nz])
 
 
diff --git a/setup.py b/setup.py
@@ -99,6 +99,14 @@ def write_version_py(version, isreleased, filename):
     package_data={'mbuild': ['utils/reference/*.{pdb,mol2}',
                              'lib/*.{pdb,mol2}',
                              ]},
+    entry_points={
+        'mbuild.plugins':[
+            "Monolayer = mbuild.lib.recipes.monolayer:Monolayer",
+            "Polymer = mbuild.lib.recipes.polymer:Polymer",
+            "SilicaInterface = mbuild.lib.recipes.silica_interface:SilicaInterface",
+            "TiledCompound = mbuild.lib.recipes.tiled_compound:TiledCompound",
+        ]
+    },
     package_dir={'mbuild': 'mbuild'},
     include_package_data=True,
     license="MIT",
