Migrate to hoomd v4

environment-cpu.yml

@@ -6,7 +6,7 @@ dependencies:
   - foyer
   - freud
   - gsd<3.0
-  - hoomd=3.11=*cpu*
+  - hoomd=4.1=*cpu*
   - mbuild
   - numpy
   - openbabel

environment-gpu.yml

@@ -6,7 +6,7 @@ dependencies:
   - foyer
   - freud
   - gsd<3.0
-  - hoomd=3.11=*gpu*
+  - hoomd=4.1=*gpu*
   - mbuild
   - numpy
   - openbabel

hoomd_organics/base/simulation.py

@@ -1,3 +1,4 @@
+import inspect
 import pickle
 import warnings
 
@@ -8,6 +9,7 @@
 import unyt as u
 
 from hoomd_organics.utils import (
+    HOOMDThermostats,
     StdOutLogger,
     UpdateWalls,
     calculate_box_length,
@@ -66,6 +68,7 @@
         gsd_file_name="trajectory.gsd",
         log_write_freq=1e3,
         log_file_name="sim_data.txt",
+        thermostat=HOOMDThermostats.MTTK,
     ):
         super(Simulation, self).__init__(device, seed)
         self.initial_state = initial_state
@@ -95,6 +98,7 @@
         self._create_state(self.initial_state)
         # Add a gsd and thermo props logger to sim operations
         self._add_hoomd_writers()
+        self._thermostat = thermostat
 
     @classmethod
     def from_system(cls, system, **kwargs):
@@ -281,6 +285,20 @@
                 "have been called for the first time."
             )
 
+    @property
+    def thermostat(self):
+        return self._thermostat
+
+    @thermostat.setter
+    def thermostat(self, thermostat):
+        if not issubclass(
+            self._thermostat, hoomd.md.methods.thermostats.Thermostat
+        ):
+            raise ValueError(
+                f"Invalid thermostat. Please choose from: {HOOMDThermostats}"
+            )
+        self._thermostat = thermostat
+
     def add_force(self, hoomd_force):
         """"""
         self._forcefield.append(hoomd_force)
@@ -317,6 +335,17 @@
             elif shift_by:
                 lj_forces.params[k]["sigma"] = sigma + shift_by
 
+    def _initialize_thermostat(self, thermostat_kwargs):
+        """Initializes the thermostat used by the integrator."""
+        required_thermostat_kwargs = {}
+        for k in inspect.signature(self.thermostat).parameters:
+            if k not in thermostat_kwargs.keys():
+                raise ValueError(
+                    f"Missing required parameter {k} for thermostat."
+                )
+            required_thermostat_kwargs[k] = thermostat_kwargs[k]
+        return self.thermostat(**required_thermostat_kwargs)
+
     def set_integrator_method(self, integrator_method, method_kwargs):
         """Creates an initial (or updates the existing) method used by
         Hoomd's integrator. This doesn't need to be called directly;
@@ -458,11 +487,12 @@
         )
         self.operations.updaters.append(box_resizer)
         self.set_integrator_method(
-            integrator_method=hoomd.md.methods.NVT,
+            integrator_method=hoomd.md.methods.ConstantVolume,
             method_kwargs={
-                "tau": tau_kt,
+                "thermostat": self._initialize_thermostat(
+                    {"kT": kT, "tau": tau_kt}
+                ),
                 "filter": self.integrate_group,
-                "kT": kT,
             },
         )
         if thermalize_particles:
@@ -487,7 +517,6 @@
         self,
         n_steps,
         kT,
-        alpha,
         tally_reservoir_energy=False,
         default_gamma=1.0,
         default_gamma_r=(1.0, 1.0, 1.0),
@@ -500,7 +529,6 @@
             method_kwargs={
                 "filter": self.integrate_group,
                 "kT": kT,
-                "alpha": alpha,
                 "tally_reservoir_energy": tally_reservoir_energy,
                 "default_gamma": default_gamma,
                 "default_gamma_r": default_gamma_r,
@@ -533,18 +561,18 @@
     ):
         """"""
         self.set_integrator_method(
-            integrator_method=hoomd.md.methods.NPT,
+            integrator_method=hoomd.md.methods.ConstantPressure,
             method_kwargs={
-                "kT": kT,
                 "S": pressure,
-                "tau": tau_kt,
                 "tauS": tau_pressure,
                 "couple": couple,
                 "box_dof": box_dof,
                 "rescale_all": rescale_all,
                 "gamma": gamma,
                 "filter": self.integrate_group,
-                "kT": kT,
+                "thermostat": self._initialize_thermostat(
+                    {"kT": kT, "tau": tau_kt}
+                ),
             },
         )
         if thermalize_particles:
@@ -568,11 +596,12 @@
     ):
         """"""
         self.set_integrator_method(
-            integrator_method=hoomd.md.methods.NVT,
+            integrator_method=hoomd.md.methods.ConstantVolume,
             method_kwargs={
-                "tau": tau_kt,
+                "thermostat": self._initialize_thermostat(
+                    {"kT": kT, "tau": tau_kt}
+                ),
                 "filter": self.integrate_group,
-                "kT": kT,
             },
         )
         if thermalize_particles:
@@ -589,7 +618,7 @@
     def run_NVE(self, n_steps, write_at_start=True):
         """"""
         self.set_integrator_method(
-            integrator_method=hoomd.md.methods.NVE,
+            integrator_method=hoomd.md.methods.ConstantVolume,
             method_kwargs={"filter": self.integrate_group},
         )
         std_out_logger = StdOutLogger(n_steps=n_steps, sim=self)
@@ -711,7 +740,8 @@
             filename=self.gsd_file_name,
             trigger=hoomd.trigger.Periodic(int(self.gsd_write_freq)),
             mode="wb",
-            dynamic=["momentum"],
+            dynamic=["momentum", "property"],
+            filter=hoomd.filter.All(),
             logger=gsd_logger,
         )
 

hoomd_organics/library/simulations/tensile.py

@@ -91,7 +91,7 @@ def run_tensile(self, strain, kT, n_steps, period):
         self.operations.updaters.append(box_resizer)
         self.operations.updaters.append(particle_updater)
         self.set_integrator_method(
-            integrator_method=hoomd.md.methods.NVE,
+            integrator_method=hoomd.md.methods.ConstantVolume,
             method_kwargs={"filter": self.integrate_group},
         )
         self.run(n_steps + 1)

hoomd_organics/tests/base/test_simulation.py

@@ -99,7 +99,7 @@ def test_set_ref_mass(self, benzene_system):
     def test_NVT(self, benzene_system):
         sim = Simulation.from_system(benzene_system)
         sim.run_NVT(kT=1.0, tau_kt=0.01, n_steps=500)
-        assert isinstance(sim.method, hoomd.md.methods.NVT)
+        assert isinstance(sim.method, hoomd.md.methods.ConstantVolume)
 
     def test_NPT(self, benzene_system):
         sim = Simulation.from_system(benzene_system)
@@ -110,17 +110,17 @@ def test_NPT(self, benzene_system):
             tau_kt=0.001,
             tau_pressure=0.01,
         )
-        assert isinstance(sim.method, hoomd.md.methods.NPT)
+        assert isinstance(sim.method, hoomd.md.methods.ConstantPressure)
 
     def test_langevin(self, benzene_system):
         sim = Simulation.from_system(benzene_system)
-        sim.run_langevin(n_steps=500, kT=1.0, alpha=0.5)
+        sim.run_langevin(n_steps=500, kT=1.0)
         assert isinstance(sim.method, hoomd.md.methods.Langevin)
 
     def test_NVE(self, benzene_system):
         sim = Simulation.from_system(benzene_system)
         sim.run_NVE(n_steps=500)
-        assert isinstance(sim.method, hoomd.md.methods.NVE)
+        assert isinstance(sim.method, hoomd.md.methods.ConstantVolume)
 
     def test_displacement_cap(self, benzene_system):
         sim = Simulation.from_system(benzene_system)
@@ -207,11 +207,11 @@ def test_update_volume_two_values(self, benzene_system):
     def test_change_methods(self, benzene_system):
         sim = Simulation.from_system(benzene_system)
         sim.run_NVT(kT=1.0, tau_kt=0.01, n_steps=0)
-        assert isinstance(sim.method, hoomd.md.methods.NVT)
+        assert isinstance(sim.method, hoomd.md.methods.ConstantVolume)
         sim.run_NPT(
             kT=1.0, tau_kt=0.01, tau_pressure=0.1, pressure=0.001, n_steps=0
         )
-        assert isinstance(sim.method, hoomd.md.methods.NPT)
+        assert isinstance(sim.method, hoomd.md.methods.ConstantPressure)
 
     def test_change_dt(self, benzene_system):
         sim = Simulation.from_system(benzene_system)
@@ -309,6 +309,7 @@ def test_save_restart_gsd(self, benzene_system):
     def test_gsd_logger(self, benzene_system):
         sim = Simulation.from_system(benzene_system, gsd_write_freq=1)
         sim.run_NVT(n_steps=5, kT=1.0, tau_kt=0.001)
+        sim.operations.writers[-2].flush()
         expected_gsd_quantities = [
             "hoomd_organics/base/simulation/Simulation/timestep",
             "hoomd_organics/base/simulation/Simulation/tps",

hoomd_organics/utils/__init__.py

@@ -1,5 +1,5 @@
 from .actions import *
-from .base_types import FF_Types
+from .base_types import FF_Types, HOOMDThermostats
 from .ff_utils import xml_to_gmso_ff
 from .utils import (
     calculate_box_length,

hoomd_organics/utils/base_types.py

@@ -1,7 +1,16 @@
+import hoomd
+
+
 class FF_Types:
     opls = "opls"
     pps_opls = "pps_opls"
     oplsaa = "oplsaa"
     gaff = "gaff"
     custom = "custom"
     Hoomd = "Hoomd"
+
+
+class HOOMDThermostats:
+    BERENDSEN = hoomd.md.methods.thermostats.Berendsen
+    BUSSI = hoomd.md.methods.thermostats.Bussi
+    MTTK = hoomd.md.methods.thermostats.MTTK