Montecarlo barostat

Examples/LJ_MCMC.py

@@ -0,0 +1,143 @@
+from openmmtools.testsystems import LennardJonesFluid
+
+# Use the LennardJonesFluid example from openmmtools to initialize particle positions and topology
+# For this example, the topology provides the masses for the particles
+# The default LennardJonesFluid example considers the system to be Argon with 39.9 amu
+lj_fluid = LennardJonesFluid(reduced_density=0.5, nparticles=1100)
+
+
+from chiron.potential import LJPotential
+from openmm import unit
+
+# initialize the LennardJones potential for UA-TraPPE methane
+#
+sigma = 0.373 * unit.nanometer
+epsilon = 0.2941 * unit.kilocalories_per_mole
+cutoff = 1.4 * unit.nanometer
+
+lj_potential = LJPotential(
+    lj_fluid.topology, sigma=sigma, epsilon=epsilon, cutoff=cutoff
+)
+
+from chiron.states import SamplerState, ThermodynamicState
+
+# define the sampler state
+sampler_state = SamplerState(
+    x0=lj_fluid.positions, box_vectors=lj_fluid.system.getDefaultPeriodicBoxVectors()
+)
+
+# define the thermodynamic state
+thermodynamic_state = ThermodynamicState(
+    potential=lj_potential,
+    temperature=140 * unit.kelvin,
+    pressure=13.00765 * unit.atmosphere,
+)
+
+from chiron.neighbors import PairList, OrthogonalPeriodicSpace
+
+# define the neighbor list for an orthogonal periodic space
+skin = 0.5 * unit.nanometer
+
+nbr_list = PairList(OrthogonalPeriodicSpace(), cutoff=cutoff)
+
+
+# build the neighbor list from the sampler state
+nbr_list.build_from_state(sampler_state)
+
+from chiron.minimze import minimize_energy
+
+results = minimize_energy(
+    sampler_state.x0, lj_potential.compute_energy, nbr_list, maxiter=100
+)
+
+min_x = results.params
+
+
+from chiron.reporters import SimulationReporter
+
+# initialize a reporter to save the simulation data
+filename1 = "test_lj_nvt.h5"
+filename2 = "test_lj_npt.h5"
+filename3 = "test_lj_langevin.h5"
+import os
+
+if os.path.isfile(filename1):
+    os.remove(filename1)
+if os.path.isfile(filename2):
+    os.remove(filename2)
+if os.path.isfile(filename3):
+    os.remove(filename3)
+reporter1 = SimulationReporter(filename1, lj_fluid.topology, 10)
+reporter2 = SimulationReporter(filename2, lj_fluid.topology, 10)
+reporter3 = SimulationReporter(filename3, lj_fluid.topology, 10)
+
+from chiron.mcmc import (
+    MetropolisDisplacementMove,
+    MoveSet,
+    MCMCSampler,
+    MCBarostatMove,
+    LangevinDynamicsMove,
+)
+
+langevin_move = LangevinDynamicsMove(
+    stepsize=1.0 * unit.femtoseconds,
+    collision_rate=1.0 / unit.picoseconds,
+    nr_of_steps=1000,
+    simulation_reporter=reporter3,
+    seed=1234,
+)
+
+mc_disp_move = MetropolisDisplacementMove(
+    seed=1234,
+    displacement_sigma=0.1 * unit.nanometer,
+    nr_of_moves=90,
+    simulation_reporter=reporter1,
+)
+
+mc_barostat_move = MCBarostatMove(
+    seed=1234,
+    volume_max_scale=0.1,
+    nr_of_moves=10,
+    adjust_box_scaling=True,
+    adjust_frequency=100,
+    simulation_reporter=reporter2,
+)
+move_set = MoveSet(
+    [
+        ("MetropolisDisplacementMove", mc_disp_move),
+        ("MCBarostatMove", mc_barostat_move),
+    ]
+)
+
+sampler = MCMCSampler(move_set, sampler_state, thermodynamic_state)
+
+mass = unit.Quantity(16.04, unit.gram / unit.mole)
+volume = (
+    sampler_state.box_vectors[0][0]
+    * sampler_state.box_vectors[1][1]
+    * sampler_state.box_vectors[2][2]
+)
+initial_density = (
+    mass
+    * sampler.sampler_state.x0.shape[0]
+    / unit.AVOGADRO_CONSTANT_NA
+    / (unit.Quantity(volume, unit.nanometer**3))
+)
+densities = []
+densities.append(initial_density)
+sampler.run(n_iterations=100, nbr_list=nbr_list)  # how many times to repeat
+
+final_density = (
+    mass
+    * sampler.sampler_state.x0.shape[0]
+    / unit.AVOGADRO_CONSTANT_NA
+    / (sampler.thermodynamic_state.volume)
+)
+densities.append(final_density)
+
+print(f"initial density: {initial_density}\nfinal density: {final_density}")
+print(
+    f"initial density: {initial_density.value_in_unit(unit.kilogram/unit.meter**3)}\nfinal density: {final_density.value_in_unit(unit.kilogram/unit.meter**3)}"
+)
+print(mc_barostat_move.statistics)
+print(mc_disp_move.statistics)

Examples/LJ_mcbarostat.py

@@ -0,0 +1,70 @@
+from openmmtools.testsystems import LennardJonesFluid
+
+# Use the LennardJonesFluid example from openmmtools to initialize particle positions and topology
+# For this example, the topology provides the masses for the particles
+# The default LennardJonesFluid example considers the system to be Argon with 39.9 amu
+lj_fluid = LennardJonesFluid(reduced_density=0.1, nparticles=1000)
+
+
+from chiron.potential import LJPotential
+from openmm import unit
+
+# initialize the LennardJones potential in chiron
+#
+sigma = 0.34 * unit.nanometer
+epsilon = 0.238 * unit.kilocalories_per_mole
+cutoff = 3.0 * sigma
+
+lj_potential = LJPotential(
+    lj_fluid.topology, sigma=sigma, epsilon=epsilon, cutoff=cutoff
+)
+
+from chiron.states import SamplerState, ThermodynamicState
+
+# define the sampler state
+sampler_state = SamplerState(
+    x0=lj_fluid.positions, box_vectors=lj_fluid.system.getDefaultPeriodicBoxVectors()
+)
+
+# define the thermodynamic state
+thermodynamic_state = ThermodynamicState(
+    potential=lj_potential,
+    temperature=300 * unit.kelvin,
+    pressure=1.0 * unit.atmosphere,
+)
+
+from chiron.neighbors import NeighborListNsqrd, OrthogonalPeriodicSpace
+
+# define the neighbor list for an orthogonal periodic space
+skin = 0.5 * unit.nanometer
+
+nbr_list = NeighborListNsqrd(
+    OrthogonalPeriodicSpace(), cutoff=cutoff, skin=skin, n_max_neighbors=180
+)
+from chiron.neighbors import PairList
+
+
+# build the neighbor list from the sampler state
+nbr_list.build_from_state(sampler_state)
+
+from chiron.reporters import SimulationReporter
+
+# initialize a reporter to save the simulation data
+filename = "test_lj.h5"
+import os
+
+if os.path.isfile(filename):
+    os.remove(filename)
+reporter = SimulationReporter("test_mc_lj.h5", lj_fluid.topology, 1)
+
+print(thermodynamic_state.pressure)
+from chiron.mcmc import MCBarostatMove
+
+barostat_move = MCBarostatMove(
+    seed=1234,
+    volume_max_scale=0.01,
+    nr_of_moves=1000,
+    simulation_reporter=reporter,
+)
+
+barostat_move.run(sampler_state, thermodynamic_state, nbr_list, True)

LICENSE

@@ -1,7 +1,7 @@
 
 MIT License
 
-Copyright (c) 2023 Chris Iacovella
+Copyright (c) 2023 Chodera Lab
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

chiron/integrators.py

@@ -58,6 +58,7 @@ def __init__(
         self.save_frequency = save_frequency
 
         self.velocities = None
+
     def set_velocities(self, vel: unit.Quantity) -> None:
         """
         Set the initial velocities for the Langevin Integrator.
@@ -112,7 +113,7 @@ def run(
         log.info(f"n_steps = {n_steps}")
         log.info(f"temperature = {temperature}")
         log.info(f"Using seed: {key}")
-
+        self.key = key
         kbT_unitless = (self.kB * temperature).value_in_unit_system(unit.md_unit_system)
         mass_unitless = jnp.array(mass.value_in_unit_system(unit.md_unit_system))[
             :, None
@@ -125,7 +126,7 @@ def run(
 
         # Initialize velocities
         if self.velocities is None:
-            v0 = sigma_v * random.normal(key, x0.shape)
+            v0 = sigma_v * random.normal(self.key, x0.shape)
         else:
             v0 = self.velocities.value_in_unit_system(unit.md_unit_system)
         # Convert to dimensionless quantities
@@ -139,7 +140,7 @@ def run(
 
         F = potential.compute_force(x, nbr_list)
         for step in tqdm(range(n_steps)) if self.progress_bar else range(n_steps):
-            key, subkey = random.split(key)
+            self.key, subkey = random.split(self.key)
             # v
             v += (stepsize_unitless * 0.5) * F / mass_unitless
             # r
@@ -178,6 +179,6 @@ def run(
 
                     # log.debug(d)
                     self.reporter.report(d)
-
+        sampler_state.x0 = x
         log.debug("Finished running Langevin dynamics")
         # self.reporter.close()