Add custom forcefields

environment-cpu.yml

@@ -1,6 +1,7 @@
 name: hoomd-organics
 channels:
   - conda-forge
+  - pytorch
 dependencies:
   - foyer
   - freud
@@ -16,6 +17,8 @@ dependencies:
   - pytest-cov
   - python=3.9
   - unyt
+  - pytorch
+  - cpuonly
   - pip:
       - git+https://github.com/cmelab/cmeutils.git@master
       - git+https://github.com/cmelab/grits.git@main

environment-dev.yml

@@ -1,6 +1,7 @@
 name: hoomd-organics-dev
 channels:
   - conda-forge
+  - pytorch
 dependencies:
   - foyer
   - freud
@@ -17,6 +18,8 @@ dependencies:
   - python=3.9
   - unyt
   - pre-commit
+  - pytorch
+  - cpuonly
   - pip:
       - git+https://github.com/cmelab/cmeutils.git@master
       - git+https://github.com/cmelab/grits.git@main

environment-gpu.yml

@@ -1,6 +1,8 @@
 name: hoomd-organics
 channels:
   - conda-forge
+  - pytorch
+  - nvidia
 dependencies:
   - foyer
   - freud
@@ -16,6 +18,8 @@ dependencies:
   - pytest-cov
   - python=3.9
   - unyt
+  - pytorch
+  - pytorch-cuda=11.7
   - pip:
       - git+https://github.com/cmelab/cmeutils.git@master
       - git+https://github.com/cmelab/grits.git@main

hoomd_organics/library/__init__.py

@@ -1,10 +1,10 @@
-from .forcefields import (
+from .forcefields.custom_forcefields import BeadSpring, TorchCustomForce
+from .forcefields.xml_forcefields import (
     GAFF,
     OPLS_AA,
     OPLS_AA_BENZENE,
     OPLS_AA_DIMETHYLETHER,
     OPLS_AA_PPS,
-    BeadSpring,
     FF_from_file,
 )
 from .polymers import (

hoomd_organics/library/forcefields/custom_forcefields.py

@@ -0,0 +1,166 @@
+import itertools
+
+import hoomd
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+class BeadSpring:
+    def __init__(
+        self,
+        r_cut,
+        beads,
+        bonds=None,
+        angles=None,
+        dihedrals=None,
+        exclusions=["bond", "1-3"],
+    ):
+        self.beads = beads
+        self.bonds = bonds
+        self.angles = angles
+        self.dihedrals = dihedrals
+        self.r_cut = r_cut
+        self.exclusions = exclusions
+        self.hoomd_forcefield = self._create_forcefield()
+
+    def _create_forcefield(self):
+        forces = []
+        # Create pair force:
+        nlist = hoomd.md.nlist.Cell(buffer=0.40, exclusions=self.exclusions)
+        lj = hoomd.md.pair.LJ(nlist=nlist)
+        bead_types = [key for key in self.beads.keys()]
+        all_pairs = list(itertools.combinations_with_replacement(bead_types, 2))
+        for pair in all_pairs:
+            epsilon0 = self.beads[pair[0]]["epsilon"]
+            epsilon1 = self.beads[pair[1]]["epsilon"]
+            pair_epsilon = (epsilon0 + epsilon1) / 2
+
+            sigma0 = self.beads[pair[0]]["sigma"]
+            sigma1 = self.beads[pair[1]]["sigma"]
+            pair_sigma = (sigma0 + sigma1) / 2
+
+            lj.params[pair] = dict(epsilon=pair_epsilon, sigma=pair_sigma)
+            lj.r_cut[pair] = self.r_cut
+        forces.append(lj)
+        # Create bond-stretching force:
+        if self.bonds:
+            harmonic_bond = hoomd.md.bond.Harmonic()
+            for bond_type in self.bonds:
+                harmonic_bond.params[bond_type] = self.bonds[bond_type]
+            forces.append(harmonic_bond)
+        # Create bond-bending force:
+        if self.angles:
+            harmonic_angle = hoomd.md.angle.Harmonic()
+            for angle_type in self.angles:
+                harmonic_angle.params[angle_type] = self.angles[angle_type]
+            forces.append(harmonic_angle)
+        # Create torsion force:
+        if self.dihedrals:
+            periodic_dihedral = hoomd.md.dihedral.Periodic()
+            for dih_type in self.dihedrals:
+                periodic_dihedral.params[dih_type] = self.dihedrals[dih_type]
+            forces.append(periodic_dihedral)
+        return forces
+
+
+class NN(nn.Module):
+    def __init__(self, hidden_dim, n_layers, act_fn="Tanh", dropout=0.5):
+        super(NN, self).__init__()
+        self.in_dim = 4  # (relative position vector, center-to-center distance)
+        self.out_dim = 1  # predicted energy
+        self.hidden_dim = hidden_dim
+
+        self.n_layers = n_layers
+        self.act_fn = act_fn
+        self.dropout = dropout
+
+        self.net = nn.Sequential(*self._get_net())
+
+    def _get_act_fn(self):
+        act = getattr(nn, self.act_fn)
+        return act()
+
+    def _get_net(self):
+        layers = [nn.Linear(self.in_dim, self.hidden_dim), self._get_act_fn()]
+        for i in range(self.n_layers - 1):
+            layers.append(nn.Linear(self.hidden_dim, self.hidden_dim))
+            layers.append(self._get_act_fn())
+            layers.append(nn.Dropout(p=self.dropout))
+        layers.append(nn.Linear(self.hidden_dim, self.out_dim))
+        return layers
+
+    def forward(self, pos_1, pos_2):
+        x = pos_1 - pos_2
+        r = torch.norm(x, dim=1, keepdim=True)
+        x = torch.cat((x, r), dim=1)
+        energy = self.net(x)
+        force = (-1.0) * torch.autograd.grad(
+            energy, pos_1, retain_graph=True, create_graph=True
+        )[0][0]
+        return force
+
+
+class TorchCustomForce(hoomd.md.force.Custom):
+    """
+    Custom force that uses a PyTorch model to predict the forces from particle
+    positions.
+    """
+
+    def __init__(self, model):
+        """
+
+        Parameters
+        ----------
+        model: torch.nn.Module
+            pretrained model that predicts forces
+        """
+        super().__init__()
+        # load ML model
+        self.device = torch.device(
+            "cuda:0" if torch.cuda.is_available() else "cpu"
+        )
+        self.model = model
+        self.model.to(self.device)
+
+    def set_forces(self, timestep):
+        """
+        Set the forces on all particles in the system.
+        """
+        # get positions for all particles
+        with self._state.cpu_local_snapshot as snap:
+            particle_rtags = np.copy(snap.particles.rtag)
+            positions = np.array(
+                snap.particles.position[particle_rtags], copy=True
+            )
+
+        num_particles = len(positions)
+        particle_forces = []
+        for i, pos_1 in enumerate(positions):
+            pos_1_tensor = (
+                torch.from_numpy(pos_1)
+                .type(torch.FloatTensor)
+                .unsqueeze(0)
+                .to(self.device)
+            )
+            other_particles_idx = list(range(num_particles))
+            other_particles_idx.remove(i)
+            particle_force = 0
+            for pos_2 in positions[other_particles_idx]:
+                pos_2_tensor = (
+                    torch.from_numpy(pos_2)
+                    .type(torch.FloatTensor)
+                    .unsqueeze(0)
+                    .to(self.device)
+                )
+                pos_1_tensor.requires_grad = True
+                particle_force += (
+                    self.model(pos_1_tensor, pos_2_tensor)
+                    .cpu()
+                    .detach()
+                    .numpy()
+                )
+            particle_forces.append(particle_force)
+
+        with self.cpu_local_force_arrays as arrays:
+            arrays.force[particle_rtags] = np.asarray(particle_forces)

hoomd_organics/library/forcefields.py → hoomd_organics/library/forcefields/xml_forcefields.py

@@ -1,8 +1,5 @@
-import itertools
-
 import forcefield_utilities as ffutils
 import foyer
-import hoomd
 
 from hoomd_organics.assets import FF_DIR
 
@@ -66,61 +63,3 @@ class FF_from_file(foyer.Forcefield):
     def __init__(self, xml_file):
         super(FF_from_file, self).__init__(forcefield_files=xml_file)
         self.gmso_ff = ffutils.FoyerFFs().load(xml_file).to_gmso_ff()
-
-
-class BeadSpring:
-    def __init__(
-        self,
-        r_cut,
-        beads,
-        bonds=None,
-        angles=None,
-        dihedrals=None,
-        exclusions=["bond", "1-3"],
-    ):
-        self.beads = beads
-        self.bonds = bonds
-        self.angles = angles
-        self.dihedrals = dihedrals
-        self.r_cut = r_cut
-        self.exclusions = exclusions
-        self.hoomd_forcefield = self._create_forcefield()
-
-    def _create_forcefield(self):
-        forces = []
-        # Create pair force:
-        nlist = hoomd.md.nlist.Cell(buffer=0.40, exclusions=self.exclusions)
-        lj = hoomd.md.pair.LJ(nlist=nlist)
-        bead_types = [key for key in self.beads.keys()]
-        all_pairs = list(itertools.combinations_with_replacement(bead_types, 2))
-        for pair in all_pairs:
-            epsilon0 = self.beads[pair[0]]["epsilon"]
-            epsilon1 = self.beads[pair[1]]["epsilon"]
-            pair_epsilon = (epsilon0 + epsilon1) / 2
-
-            sigma0 = self.beads[pair[0]]["sigma"]
-            sigma1 = self.beads[pair[1]]["sigma"]
-            pair_sigma = (sigma0 + sigma1) / 2
-
-            lj.params[pair] = dict(epsilon=pair_epsilon, sigma=pair_sigma)
-            lj.r_cut[pair] = self.r_cut
-        forces.append(lj)
-        # Create bond-stretching force:
-        if self.bonds:
-            harmonic_bond = hoomd.md.bond.Harmonic()
-            for bond_type in self.bonds:
-                harmonic_bond.params[bond_type] = self.bonds[bond_type]
-            forces.append(harmonic_bond)
-        # Create bond-bending force:
-        if self.angles:
-            harmonic_angle = hoomd.md.angle.Harmonic()
-            for angle_type in self.angles:
-                harmonic_angle.params[angle_type] = self.angles[angle_type]
-            forces.append(harmonic_angle)
-        # Create torsion force:
-        if self.dihedrals:
-            periodic_dihedral = hoomd.md.dihedral.Periodic()
-            for dih_type in self.dihedrals:
-                periodic_dihedral.params[dih_type] = self.dihedrals[dih_type]
-            forces.append(periodic_dihedral)
-        return forces

hoomd_organics/tests/library/test_forcefields.py

@@ -11,10 +11,10 @@
     BeadSpring,
     FF_from_file,
 )
-from hoomd_organics.tests.base_test import ASSETS_DIR
+from hoomd_organics.tests.base_test import ASSETS_DIR, BaseTest
 
 
-class TestForceFields:
+class TestXMLForceFields:
     def test_GAFF(self):
         ff = GAFF()
         assert ff.gmso_ff is not None
@@ -40,6 +40,8 @@ def test_FF_from_file(self):
         ff = FF_from_file(xml_file)
         assert ff.gmso_ff is not None
 
+
+class TestCustomForceFields(BaseTest):
     def test_BeadSpring(self):
         ff = BeadSpring(
             r_cut=2.5,
@@ -89,3 +91,24 @@ def test_BeadSpring(self):
             assert ff.hoomd_forcefield[3].params[param]["k"] == 100
             assert ff.hoomd_forcefield[3].params[param]["d"] == -1
             assert ff.hoomd_forcefield[3].params[param]["n"] == 1
+
+    def test_TorchCustomForce(self, benzene_smiles):
+        pass
+        # molecule = Molecule(num_mols=2, smiles=benzene_smiles)
+        # molecule.coarse_grain(beads={"A": benzene_smiles})
+        # system = Pack(
+        #     molecules=[molecule],
+        #     density=0.8,
+        #     r_cut=2.5,
+        # )
+        # model = NN(hidden_dim=32, n_layers=2, act_fn="Tanh", dropout=0.5)
+        # device = torch.device("cuda:0" if torch.cuda.is_available()
+        # else "cpu")
+        # model.load_state_dict(torch.load("best_model.pth",
+        # map_location=device))
+        # custom_force = TorchCustomForce(model)
+        # sim = Simulation(
+        #     initial_state=system.hoomd_snapshot,
+        #     forcefield=[custom_force],
+        # )
+        # sim.run_NVT(n_steps=10, kT=1.0, tau_kt=1.0)