diff --git a/environment-cpu.yml b/environment-cpu.yml index cd9b7c77..95da24c6 100644 --- a/environment-cpu.yml +++ b/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 diff --git a/environment-dev.yml b/environment-dev.yml index ff931fd0..ec071f47 100644 --- a/environment-dev.yml +++ b/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 diff --git a/environment-gpu.yml b/environment-gpu.yml index c3e25a2b..0593a143 100644 --- a/environment-gpu.yml +++ b/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 diff --git a/hoomd_organics/library/__init__.py b/hoomd_organics/library/__init__.py index a98cc3cb..12d5f818 100644 --- a/hoomd_organics/library/__init__.py +++ b/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 ( diff --git a/hoomd_organics/library/forcefields/custom_forcefields.py b/hoomd_organics/library/forcefields/custom_forcefields.py new file mode 100644 index 00000000..bed5ee5c --- /dev/null +++ b/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) diff --git a/hoomd_organics/library/forcefields.py b/hoomd_organics/library/forcefields/xml_forcefields.py similarity index 55% rename from hoomd_organics/library/forcefields.py rename to hoomd_organics/library/forcefields/xml_forcefields.py index a71b6a95..6bc91a47 100644 --- a/hoomd_organics/library/forcefields.py +++ b/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 diff --git a/hoomd_organics/tests/library/best_model.pth b/hoomd_organics/tests/library/best_model.pth new file mode 100644 index 00000000..c2a965f5 Binary files /dev/null and b/hoomd_organics/tests/library/best_model.pth differ diff --git a/hoomd_organics/tests/library/test_forcefields.py b/hoomd_organics/tests/library/test_forcefields.py index 7537a032..70cb5475 100644 --- a/hoomd_organics/tests/library/test_forcefields.py +++ b/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)