Refactoring of MC moves

.github/workflows/CI.yaml

@@ -9,6 +9,7 @@ on:
   pull_request:
     branches:
       - "main"
+      - "multistage"
   schedule:
     # Weekly tests run on main by default:
     #   Scheduled workflows run on the latest commit on the default or base branch.

Examples/Idealgas.py

@@ -0,0 +1,150 @@
+from openmmtools.testsystems import IdealGas
+from openmm import unit
+
+"""
+This example explore an ideal gas system, where the particles are non-interacting. 
+This will use the MonteCarloBarostatMove to sample the volume of the system and 
+MonteCarloDisplacementMove to sample the particle positions.
+
+This utilizes the IdealGas example from openmmtools to initialize particle positions and topology.
+
+"""
+
+# Use the IdealGas example from openmmtools to initialize particle positions and topology
+# For this example, the topology provides the masses for the particles
+
+n_particles = 216
+temperature = 298 * unit.kelvin
+pressure = 1 * unit.atmosphere
+
+ideal_gas = IdealGas(nparticles=n_particles, temperature=temperature, pressure=pressure)
+
+
+from chiron.potential import IdealGasPotential
+from chiron.utils import PRNG, get_list_of_mass
+import jax.numpy as jnp
+
+# particles are non interacting
+cutoff = 0.0 * unit.nanometer
+ideal_gas_potential = IdealGasPotential(ideal_gas.topology)
+
+from chiron.states import SamplerState, ThermodynamicState
+
+# define the thermodynamic state
+thermodynamic_state = ThermodynamicState(
+    potential=ideal_gas_potential,
+    temperature=temperature,
+    pressure=pressure,
+)
+
+PRNG.set_seed(1234)
+
+
+# define the sampler state
+sampler_state = SamplerState(
+    positions=ideal_gas.positions,
+    current_PRNG_key=PRNG.get_random_key(),
+    box_vectors=ideal_gas.system.getDefaultPeriodicBoxVectors(),
+)
+
+from chiron.neighbors import PairListNsqrd, OrthogonalPeriodicSpace
+
+# define the pair list for an orthogonal periodic space
+# since particles are non-interacting, this will not really do much
+# but will be used to appropriately wrap particles in space
+nbr_list = PairListNsqrd(OrthogonalPeriodicSpace(), cutoff=cutoff)
+nbr_list.build_from_state(sampler_state)
+
+from chiron.reporters import MCReporter
+
+# initialize a reporter to save the simulation data
+filename = "test_mc_ideal_gas.h5"
+import os
+
+if os.path.isfile(filename):
+    os.remove(filename)
+reporter = MCReporter(filename, 100)
+
+
+from chiron.mcmc import (
+    MonteCarloDisplacementMove,
+    MonteCarloBarostatMove,
+    MoveSchedule,
+    MCMCSampler,
+)
+
+# initialize the displacement move
+mc_barostat_move = MonteCarloBarostatMove(
+    volume_max_scale=0.2,
+    number_of_moves=10,
+    reporter=reporter,
+    autotune=True,
+    autotune_interval=100,
+)
+
+# initialize the barostat move and the move schedule
+metropolis_displacement_move = MonteCarloDisplacementMove(
+    displacement_sigma=0.1 * unit.nanometer,
+    number_of_moves=100,
+    autotune=True,
+    autotune_interval=100,
+)
+
+# define the move schedule
+move_set = MoveSchedule(
+    [
+        ("MonteCarloDisplacementMove", metropolis_displacement_move),
+        ("MonteCarloBarostatMove", mc_barostat_move),
+    ]
+)
+
+sampler = MCMCSampler(move_set)
+sampler.run(
+    sampler_state, thermodynamic_state, n_iterations=10, nbr_list=nbr_list
+)  # how many times to repeat
+
+# get the volume from the reporter
+volume = reporter.get_property("volume")
+step = reporter.get_property("elapsed_step")
+
+
+import matplotlib.pyplot as plt
+
+plt.plot(step, volume)
+plt.show()
+
+# get expectations
+ideal_volume = ideal_gas.get_volume_expectation(thermodynamic_state)
+ideal_volume_std = ideal_gas.get_volume_standard_deviation(thermodynamic_state)
+
+print("ideal volume and standard deviation: ", ideal_volume, ideal_volume_std)
+
+
+volume_mean = jnp.mean(jnp.array(volume)) * unit.nanometer**3
+volume_std = jnp.std(jnp.array(volume)) * unit.nanometer**3
+
+
+print("measured volume and standard deviation: ", volume_mean, volume_std)
+
+# get the masses of particles from the topology
+masses = get_list_of_mass(ideal_gas.topology)
+
+sum_of_masses = jnp.sum(jnp.array(masses.value_in_unit(unit.amu))) * unit.amu
+
+ideal_density = sum_of_masses / unit.AVOGADRO_CONSTANT_NA / ideal_volume
+measured_density = sum_of_masses / unit.AVOGADRO_CONSTANT_NA / volume_mean
+
+assert jnp.isclose(
+    ideal_density.value_in_unit(unit.kilogram / unit.meter**3),
+    measured_density.value_in_unit(unit.kilogram / unit.meter**3),
+    atol=1e-1,
+)
+# see if within 5% of ideal volume
+assert (
+    abs(ideal_volume - volume_mean) / ideal_volume < 0.05
+), f"Warning: {abs(ideal_volume - volume_mean) / ideal_volume} exceeds the 5% threshold"
+
+# see if within 10% of the ideal standard deviation of the volume
+assert (
+    abs(ideal_volume_std - volume_std) / ideal_volume_std < 0.1
+), f"Warning: {abs(ideal_volume_std - volume_std) / ideal_volume_std} exceeds the 10% threshold"

Examples/LJ_MCMC.py

@@ -0,0 +1,165 @@
+from openmm import unit
+from openmm import app
+
+"""
+This example explore a Lennard-Jones system, where a single bead represents a united atom methane molecule, 
+modeled with the UA-TraPPE force field.
+
+
+"""
+n_particles = 1100
+temperature = 140 * unit.kelvin
+pressure = 13.00765 * unit.atmosphere
+mass = unit.Quantity(16.04, unit.gram / unit.mole)
+
+# create the topology
+lj_topology = app.Topology()
+element = app.Element(1000, "CH4", "CH4", mass)
+chain = lj_topology.addChain()
+for i in range(n_particles):
+    residue = lj_topology.addResidue("CH4", chain)
+    lj_topology.addAtom("CH4", element, residue)
+
+import jax.numpy as jnp
+
+# these were generated in Mbuild using fill_box which wraps packmol
+# a minimum spacing of 0.4 nm was used during construction.
+
+from chiron.utils import get_full_path
+
+positions = jnp.load(get_full_path("Examples/methane_coords.npy")) * unit.nanometer
+
+box_vectors = (
+    jnp.array(
+        [
+            [4.275021399280942, 0.0, 0.0],
+            [0.0, 4.275021399280942, 0.0],
+            [0.0, 0.0, 4.275021399280942],
+        ]
+    )
+    * unit.nanometer
+)
+
+from chiron.potential import LJPotential
+from chiron.utils import PRNG
+import jax.numpy as jnp
+
+#
+
+# 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_topology, sigma=sigma, epsilon=epsilon, cutoff=cutoff)
+
+from chiron.states import SamplerState, ThermodynamicState
+
+# define the thermodynamic state
+thermodynamic_state = ThermodynamicState(
+    potential=lj_potential,
+    temperature=temperature,
+    pressure=pressure,
+)
+
+PRNG.set_seed(1234)
+
+
+# define the sampler state
+sampler_state = SamplerState(
+    positions=positions, current_PRNG_key=PRNG.get_random_key(), box_vectors=box_vectors
+)
+
+
+from chiron.neighbors import PairListNsqrd, OrthogonalPeriodicSpace
+
+# define the pair list for an orthogonal periodic space
+# since particles are non-interacting, this will not really do much
+# but will appropriately wrap particles in space
+nbr_list = PairListNsqrd(OrthogonalPeriodicSpace(), cutoff=cutoff)
+nbr_list.build_from_state(sampler_state)
+
+# CRI: minimizer is not working correctly on my mac
+# from chiron.minimze import minimize_energy
+#
+# results = minimize_energy(
+#     sampler_state.positions, lj_potential.compute_energy, nbr_list, maxiter=100
+# )
+#
+# min_x = results.params
+#
+# sampler_state.positions = min_x
+
+from chiron.reporters import MCReporter
+
+# initialize a reporter to save the simulation data
+import os
+
+
+filename_displacement = "test_mc_lj_disp.h5"
+
+if os.path.isfile(filename_displacement):
+    os.remove(filename_displacement)
+reporter_displacement = MCReporter(filename_displacement, 10)
+
+from chiron.mcmc import MonteCarloDisplacementMove
+
+mc_displacement_move = MonteCarloDisplacementMove(
+    displacement_sigma=0.001 * unit.nanometer,
+    number_of_moves=100,
+    reporter=reporter_displacement,
+    report_interval=10,
+    autotune=True,
+    autotune_interval=100,
+)
+
+filename_barostat = "test_mc_lj_barostat.h5"
+if os.path.isfile(filename_barostat):
+    os.remove(filename_barostat)
+reporter_barostat = MCReporter(filename_barostat, 1)
+
+
+from chiron.mcmc import MonteCarloBarostatMove
+
+mc_barostat_move = MonteCarloBarostatMove(
+    volume_max_scale=0.1,
+    number_of_moves=10,
+    reporter=reporter_barostat,
+    report_interval=1,
+    autotune=True,
+    autotune_interval=50,
+)
+
+from chiron.reporters import LangevinDynamicsReporter
+
+filename_langevin = "test_mc_lj_langevin.h5"
+
+if os.path.isfile(filename_langevin):
+    os.remove(filename_langevin)
+reporter_langevin = LangevinDynamicsReporter(filename_langevin, 10)
+
+from chiron.mcmc import LangevinDynamicsMove
+
+langevin_dynamics_move = LangevinDynamicsMove(
+    timestep=1.0 * unit.femtoseconds,
+    collision_rate=1.0 / unit.picoseconds,
+    number_of_steps=1000,
+    reporter=reporter_langevin,
+    report_interval=10,
+)
+
+from chiron.mcmc import MoveSchedule
+
+move_set = MoveSchedule(
+    [
+        ("LangevinDynamicsMove", langevin_dynamics_move),
+        ("MonteCarloDisplacementMove", mc_displacement_move),
+        ("MonteCarloBarostatMove", mc_barostat_move),
+    ]
+)
+
+from chiron.mcmc import MCMCSampler
+
+sampler = MCMCSampler(move_set)
+sampler.run(sampler_state, thermodynamic_state, n_iterations=100, nbr_list=nbr_list)