Merge pull request cmelab#113 from cmelab/combine-slabs

We can now use two GSD files in `welding.Interface`

flowermd/modules/welding/utils.py

@@ -32,18 +32,37 @@ def add_void_particles(
     """
     void_axis = np.asarray(void_axis)
     snapshot.particles.N += num_voids
-    snapshot.particles.position[-1] = (
-        void_axis * snapshot.configuration.box[0:3] / 2
+    # Set updated positions
+    void_pos = void_axis * snapshot.configuration.box[0:3] / 2
+    init_pos = snapshot.particles.position
+    new_pos = np.empty((init_pos.shape[0] + 1, 3))
+    new_pos[: init_pos.shape[0]] = init_pos
+    new_pos[-1] = void_pos
+    snapshot.particles.position = np.concatenate(
+        (init_pos, void_pos.reshape(1, 3)), axis=0
     )
-    snapshot.particles.types = snapshot.particles.types + ["VOID"]
-    snapshot.particles.typeid[-1] = len(snapshot.particles.types) - 1
-    snapshot.particles.mass[-1] = 1
+    # Set updated types and type IDs
+    snapshot.particles.types.append("VOID")
+    void_id = len(snapshot.particles.types) - 1
+    init_ids = snapshot.particles.typeid
+    snapshot.particles.typeid = np.concatenate(
+        (init_ids, np.array([void_id])), axis=None
+    )
+    # Set updated mass and charges
+    init_mass = snapshot.particles.mass
+    snapshot.particles.mass = np.concatenate(
+        (init_mass, np.array([1])), axis=None
+    )
+    init_charges = snapshot.particles.charge
+    snapshot.particles.charge = np.concatenate(
+        (init_charges, np.array([0])), axis=None
+    )
+    # Updated LJ params
     lj = [i for i in forcefield if isinstance(i, hoomd.md.pair.LJ)][0]
     for ptype in snapshot.particles.types:
         lj.params[(ptype, "VOID")] = {
             "sigma": void_diameter,
             "epsilon": epsilon,
         }
         lj.r_cut[(ptype, "VOID")] = r_cut
-
     return snapshot, forcefield

flowermd/modules/welding/welding.py

@@ -5,6 +5,7 @@
 import numpy as np
 
 from flowermd.base.simulation import Simulation
+from flowermd.internal import check_return_iterable
 from flowermd.utils import HOOMDThermostats
 
 
@@ -25,112 +26,146 @@ class Interface:
 
     """
 
-    def __init__(self, gsd_file, interface_axis, gap, wall_sigma=1.0):
-        self.gsd_file = gsd_file
+    def __init__(
+        self,
+        gsd_files,
+        interface_axis,
+        gap,
+        wall_sigma=1.0,
+        remove_void_particles=True,
+    ):
+        self.gsd_files = check_return_iterable(gsd_files)
         self.interface_axis = interface_axis
         self.gap = gap
         self.wall_sigma = wall_sigma
+        self._remove_void_particles = remove_void_particles
         self.hoomd_snapshot = self._build()
 
     def _build(self):
         """Duplicates the slab and builds the interface."""
-        gsd_file = gsd.hoomd.open(self.gsd_file)
-        snap = gsd_file[-1]
-        gsd_file.close()
-        axis_index = np.where(self.interface_axis != 0)[0]
+        if len(self.gsd_files) == 1:
+            gsd_file_L = gsd.hoomd.open(self.gsd_files[0])
+            gsd_file_R = gsd.hoomd.open(self.gsd_files[0])
+        else:
+            gsd_file_L = gsd.hoomd.open(self.gsd_files[0])
+            gsd_file_R = gsd.hoomd.open(self.gsd_files[1])
+        # Get snapshots
+        snap_L = gsd_file_L[-1]
+        snap_R = gsd_file_R[-1]
+        gsd_file_L.close()
+        gsd_file_R.close()
+        # Remove VOID partilces if needed
+        if self._remove_void_particles:
+            for snap in [snap_L, snap_R]:
+                if "VOID" in snap.particles.types:
+                    snap.particles.N -= 1
+                    void_id = snap.particles.types.index("VOID")
+                    snap.particles.types.remove("VOID")
+                    keep_indices = np.where(snap.particles.typeid != void_id)[0]
+                    snap.particles.position = snap.particles.position[
+                        keep_indices
+                    ]
+                    snap.particles.mass = snap.particles.mass[keep_indices]
+                    snap.particles.charge = snap.particles.charge[keep_indices]
+                    snap.particles.orientation = snap.particles.orientation[
+                        keep_indices
+                    ]
+                    snap.particles.typeid = snap.particles.typeid[keep_indices]
 
         interface = gsd.hoomd.Frame()
-        interface.particles.N = snap.particles.N * 2
-        interface.bonds.N = snap.bonds.N * 2
-        interface.bonds.M = snap.bonds.M
-        interface.angles.N = snap.angles.N * 2
-        interface.angles.M = snap.angles.M
-        interface.dihedrals.N = snap.dihedrals.N * 2
-        interface.dihedrals.M = snap.dihedrals.M
-        interface.pairs.N = snap.pairs.N * 2
-
+        interface.particles.N = snap_L.particles.N + snap_R.particles.N
+        interface.bonds.N = snap_L.bonds.N + snap_R.bonds.N
+        interface.bonds.M = 2
+        interface.angles.N = snap_L.angles.N + snap_R.angles.N
+        interface.angles.M = 3
+        interface.dihedrals.N = snap_L.dihedrals.N + snap_R.dihedrals.N
+        interface.dihedrals.M = 4
+        interface.pairs.N = snap_L.pairs.N + snap_R.pairs.N
         # Set up box. Box edge is doubled along the interface axis direction,
         # plus the gap
-        interface.configuration.box = np.copy(snap.configuration.box)
+        axis_index = np.where(self.interface_axis != 0)[0]
+        interface.configuration.box = np.copy(snap_L.configuration.box)
         interface.configuration.box[axis_index] *= 2
         interface.configuration.box[axis_index] += self.gap - self.wall_sigma
 
         # Set up snapshot.particles info:
         # Get set of new coordiantes, shifted along interface axis
         shift = (
-            snap.configuration.box[axis_index] + self.gap - self.wall_sigma
+            snap_L.configuration.box[axis_index] + self.gap - self.wall_sigma
         ) / 2
-        right_pos = np.copy(snap.particles.position)
+        right_pos = np.copy(snap_R.particles.position)
         right_pos[:, axis_index] += shift
-        left_pos = np.copy(snap.particles.position)
+        left_pos = np.copy(snap_L.particles.position)
         left_pos[:, axis_index] -= shift
 
         pos = np.concatenate((left_pos, right_pos), axis=None)
         mass = np.concatenate(
-            (snap.particles.mass, snap.particles.mass), axis=None
+            (snap_L.particles.mass, snap_R.particles.mass), axis=None
         )
         charges = np.concatenate(
-            (snap.particles.charge, snap.particles.charge), axis=None
+            (snap_L.particles.charge, snap_R.particles.charge), axis=None
         )
         type_ids = np.concatenate(
-            (snap.particles.typeid, snap.particles.typeid), axis=None
+            (snap_L.particles.typeid, snap_R.particles.typeid), axis=None
         )
         interface.particles.position = pos
         interface.particles.mass = mass
         interface.particles.charge = charges
-        interface.particles.types = snap.particles.types
+        interface.particles.types = snap_L.particles.types
         interface.particles.typeid = type_ids
 
         # Set up bonds:
-        bond_group_left = np.copy(snap.bonds.group)
-        bond_group_right = np.copy(snap.bonds.group) + snap.particles.N
+        bond_group_left = np.copy(snap_L.bonds.group)
+        bond_group_right = np.copy(snap_R.bonds.group) + snap_R.particles.N
         bond_group = np.concatenate(
             (bond_group_left, bond_group_right), axis=None
         )
         bond_type_ids = np.concatenate(
-            (snap.bonds.typeid, snap.bonds.typeid), axis=None
+            (snap_L.bonds.typeid, snap_R.bonds.typeid), axis=None
         )
         interface.bonds.group = bond_group
         interface.bonds.typeid = bond_type_ids
-        interface.bonds.types = snap.bonds.types
+        interface.bonds.types = snap_L.bonds.types
 
         # Set up angles:
-        angle_group_left = np.copy(snap.angles.group)
-        angle_group_right = np.copy(snap.angles.group) + snap.particles.N
+        angle_group_left = np.copy(snap_L.angles.group)
+        angle_group_right = np.copy(snap_R.angles.group) + snap_L.particles.N
         angle_group = np.concatenate(
             (angle_group_left, angle_group_right), axis=None
         )
         angle_type_ids = np.concatenate(
-            (snap.angles.typeid, snap.angles.typeid), axis=None
+            (snap_L.angles.typeid, snap_R.angles.typeid), axis=None
         )
         interface.angles.group = angle_group
         interface.angles.typeid = angle_type_ids
-        interface.angles.types = snap.angles.types
+        interface.angles.types = snap_L.angles.types
 
         # Set up dihedrals:
-        dihedral_group_left = np.copy(snap.dihedrals.group)
-        dihedral_group_right = np.copy(snap.dihedrals.group) + snap.particles.N
+        dihedral_group_left = np.copy(snap_L.dihedrals.group)
+        dihedral_group_right = (
+            np.copy(snap_R.dihedrals.group) + snap_L.particles.N
+        )
         dihedral_group = np.concatenate(
             (dihedral_group_left, dihedral_group_right), axis=None
         )
         dihedral_type_ids = np.concatenate(
-            (snap.dihedrals.typeid, snap.dihedrals.typeid), axis=None
+            (snap_L.dihedrals.typeid, snap_R.dihedrals.typeid), axis=None
         )
         interface.dihedrals.group = dihedral_group
         interface.dihedrals.typeid = dihedral_type_ids
-        interface.dihedrals.types = snap.dihedrals.types
+        interface.dihedrals.types = snap_L.dihedrals.types
 
         # Set up pairs:
-        if snap.pairs.N > 0:
-            pair_group_left = np.copy(snap.pairs.group)
-            pair_group_right = np.copy(snap.pairs.group) + snap.particles.N
+        if snap_L.pairs.N > 0:
+            pair_group_left = np.copy(snap_L.pairs.group)
+            pair_group_right = np.copy(snap_R.pairs.group) + snap_L.particles.N
             pair_group = np.concatenate((pair_group_left, pair_group_right))
             pair_type_ids = np.concatenate(
-                (snap.pairs.typeid, snap.pairs.typeid), axis=None
+                (snap_L.pairs.typeid, snap_R.pairs.typeid), axis=None
             )
             interface.pairs.group = pair_group
             interface.pairs.typeid = pair_type_ids
-            interface.pairs.types = snap.pairs.types
+            interface.pairs.types = snap_L.pairs.types
         return interface
 
 
@@ -242,3 +277,6 @@ def __init__(
             wall_r_cut,
             wall_r_extrap,
         )
+        snap = self.state.get_snapshot()
+        integrate_types = [i for i in snap.particles.types if i != "VOID"]
+        self.integrate_group = hoomd.filter.Type(integrate_types)

flowermd/tests/modules/test_weld.py

@@ -2,6 +2,7 @@
 
 import gsd.hoomd
 import hoomd
+import numpy as np
 
 from flowermd import Simulation
 from flowermd.modules.welding import Interface, SlabSimulation, WeldSimulation
@@ -26,7 +27,45 @@ def test_interface(self, polyethylene_system):
         )
         sim.save_restart_gsd()
         interface = Interface(
-            gsd_file="restart.gsd", interface_axis=(1, 0, 0), gap=0.1
+            gsd_files="restart.gsd", interface_axis=(1, 0, 0), gap=0.1
+        )
+        interface_snap = interface.hoomd_snapshot
+        with gsd.hoomd.open("restart.gsd", "r") as traj:
+            slab_snap = traj[0]
+
+        assert interface_snap.particles.N == slab_snap.particles.N * 2
+        assert interface_snap.bonds.N == slab_snap.bonds.N * 2
+        assert interface_snap.bonds.M == slab_snap.bonds.M
+        assert interface_snap.angles.N == slab_snap.angles.N * 2
+        assert interface_snap.angles.M == slab_snap.angles.M
+        assert interface_snap.dihedrals.N == slab_snap.dihedrals.N * 2
+        assert interface_snap.dihedrals.M == slab_snap.dihedrals.M
+        assert interface_snap.pairs.N == slab_snap.pairs.N * 2
+        assert interface_snap.pairs.M == slab_snap.pairs.M
+
+        if os.path.isfile("restart.gsd"):
+            os.remove("restart.gsd")
+
+    def test_interface_2_files(self, polyethylene_system):
+        sim = Simulation(
+            initial_state=polyethylene_system.hoomd_snapshot,
+            forcefield=polyethylene_system.hoomd_forcefield,
+            log_write_freq=2000,
+        )
+        sim.add_walls(wall_axis=(1, 0, 0), sigma=1, epsilon=1, r_cut=2)
+        sim.run_update_volume(
+            n_steps=1000,
+            period=10,
+            kT=2.0,
+            tau_kt=0.01,
+            final_box_lengths=sim.box_lengths / 2,
+        )
+        sim.save_restart_gsd("restart.gsd")
+        sim.save_restart_gsd("restart2.gsd")
+        interface = Interface(
+            gsd_files=["restart.gsd", "restart2.gsd"],
+            interface_axis=(1, 0, 0),
+            gap=0.1,
         )
         interface_snap = interface.hoomd_snapshot
         with gsd.hoomd.open("restart.gsd", "r") as traj:
@@ -90,7 +129,7 @@ def test_weld_sim(self, polyethylene_system):
         sim.save_restart_gsd()
         # Create interface system from slab restart.gsd file
         interface = Interface(
-            gsd_file="restart.gsd", interface_axis="x", gap=0.1
+            gsd_files="restart.gsd", interface_axis="x", gap=0.1
         )
         sim = WeldSimulation(
             initial_state=interface.hoomd_snapshot,
@@ -117,3 +156,53 @@ def test_void_particle(self, polyethylene_system):
         for p_type in init_types:
             assert lj.params[(p_type, "VOID")]["sigma"] == 0.4
             assert lj.params[(p_type, "VOID")]["epsilon"] == 1
+
+    def test_interface_with_void_particle(self, polyethylene_system):
+        init_snap = polyethylene_system.hoomd_snapshot
+        init_N = np.copy(init_snap.particles.N)
+        void_snap, ff = add_void_particles(
+            init_snap,
+            polyethylene_system.hoomd_forcefield,
+            void_diameter=0.10,
+            num_voids=1,
+            void_axis=(1, 0, 0),
+            epsilon=0.7,
+            r_cut=0.7,
+        )
+        assert void_snap.particles.N == init_N + 1
+        sim = SlabSimulation(
+            initial_state=void_snap,
+            forcefield=ff,
+            log_write_freq=2000,
+        )
+        sim.run_update_volume(
+            n_steps=1000,
+            period=10,
+            kT=2.0,
+            tau_kt=0.01,
+            final_box_lengths=sim.box_lengths / 2,
+        )
+        sim.save_restart_gsd("restart.gsd")
+        interface = Interface(
+            gsd_files=["restart.gsd"],
+            interface_axis=(1, 0, 0),
+            gap=0.1,
+        )
+
+        interface_snap = interface.hoomd_snapshot
+        with gsd.hoomd.open("restart.gsd", "r") as traj:
+            slab_snap = traj[0]
+
+        assert "VOID" not in interface_snap.particles.types
+        assert interface_snap.particles.N == (slab_snap.particles.N * 2) - 2
+        assert interface_snap.bonds.N == slab_snap.bonds.N * 2
+        assert interface_snap.bonds.M == slab_snap.bonds.M
+        assert interface_snap.angles.N == slab_snap.angles.N * 2
+        assert interface_snap.angles.M == slab_snap.angles.M
+        assert interface_snap.dihedrals.N == slab_snap.dihedrals.N * 2
+        assert interface_snap.dihedrals.M == slab_snap.dihedrals.M
+        assert interface_snap.pairs.N == slab_snap.pairs.N * 2
+        assert interface_snap.pairs.M == slab_snap.pairs.M
+
+        if os.path.isfile("restart.gsd"):
+            os.remove("restart.gsd")