diff --git a/Examples/LJ_MCMC.py b/Examples/LJ_MCMC.py new file mode 100644 index 0000000..3035029 --- /dev/null +++ b/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) diff --git a/Examples/LJ_mcbarostat.py b/Examples/LJ_mcbarostat.py new file mode 100644 index 0000000..72de170 --- /dev/null +++ b/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) diff --git a/LICENSE b/LICENSE index cc8c29e..937ab36 100644 --- a/LICENSE +++ b/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 diff --git a/chiron/integrators.py b/chiron/integrators.py index e94aed7..ad51f8b 100644 --- a/chiron/integrators.py +++ b/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() diff --git a/chiron/mcmc.py b/chiron/mcmc.py index 20efebd..7466336 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -102,6 +102,7 @@ def run( self, sampler_state: SamplerState, thermodynamic_state: ThermodynamicState, + nbr_list=None, ): """ Run the integrator to perform molecular dynamics simulation. @@ -115,7 +116,9 @@ def run( sampler_state=sampler_state, n_steps=self.nr_of_moves, key=self.key, + nbr_list=nbr_list, ) + self.key = self.integrator.key class MCMove(StateUpdateMove): @@ -297,7 +300,7 @@ def __init__( self.sampler_state = deepcopy(sampler_state) self.thermodynamic_state = deepcopy(thermodynamic_state) - def run(self, n_iterations: int = 1): + def run(self, n_iterations: int = 1, nbr_list=None): """ Run the sampler for a specified number of iterations. @@ -305,6 +308,9 @@ def run(self, n_iterations: int = 1): ---------- n_iterations : int, optional Number of iterations of the sampler to run. + nbr_list: Neighbor List or Pair List routine, + The routine to use to calculate the interacting atoms. + Default is None and will use an unoptimized pairlist without PBC """ log.info("Running Gibbs sampler") log.info(f"move_schedule = {self.move.move_schedule}") @@ -313,8 +319,7 @@ def run(self, n_iterations: int = 1): log.info(f"Iteration {iteration + 1}/{n_iterations}") for move_name, move in self.move.move_schedule: log.debug(f"Performing: {move_name}") - move.run(self.sampler_state, self.thermodynamic_state) - + move.run(self.sampler_state, self.thermodynamic_state, nbr_list) log.info("Finished running Gibbs sampler") log.debug("Closing reporter") for _, move in self.move.move_schedule: @@ -375,6 +380,7 @@ def apply( thermodynamic_state: ThermodynamicState, sampler_state: SamplerState, reporter: SimulationReporter, + trial: int, nbr_list=None, ): """Apply a metropolized move to the sampler state. @@ -389,16 +395,25 @@ def apply( The initial sampler state to apply the move to. This is modified. reporter: SimulationReporter The reporter to write the data to. + trial : int + The trial number of the move. This is used to determine if we have a stashed + value of the reduced_energy of the prior state. nbr_list: Neighbor List or Pair List routine, The routine to use to calculate the interacting atoms. Default is None and will use an unoptimized pairlist without PBC """ import jax.numpy as jnp + if trial == 0: + # we need to compute the initial energy + # Compute initial energy + initial_energy = thermodynamic_state.get_reduced_potential( + sampler_state, nbr_list + ) # NOTE: in kT + else: + initial_energy = self.stored_energy # Compute initial energy - initial_energy = thermodynamic_state.get_reduced_potential( - sampler_state, nbr_list - ) # NOTE: in kT + log.debug(f"Initial energy is {initial_energy} kT.") # Store initial positions of the atoms that are moved. # We'll use this also to recover in case the move is rejected. @@ -409,11 +424,38 @@ def apply( initial_positions = jnp.copy(x0) else: initial_positions = jnp.copy(sampler_state.x0[jnp.array(atom_subset)]) - log.debug(f"Initial positions are {initial_positions} nm.") + if sampler_state.box_vectors is not None: + initial_box_vectors = jnp.copy(sampler_state.box_vectors) + else: + initial_box_vectors = None + + # if we are running in NPT, we need to store the initial volume as well + if thermodynamic_state.pressure is not None: + if sampler_state.box_vectors is None: + raise ValueError( + "Cannot run NPT without a volume or box vectors specified" + ) + # units are nm**3 + initial_volume = ( + initial_box_vectors[0][0] + * initial_box_vectors[1][1] + * initial_box_vectors[2][2] + ) + # set the volume in the thermodynamic state + thermodynamic_state.volume = initial_volume * unit.nanometer**3 + + initial_volume = jnp.copy( + thermodynamic_state.volume.value_in_unit_system(unit.md_unit_system) + ) + + # log.debug(f"Initial positions are {initial_positions} nm.") # Propose perturbed positions. Modifying the reference changes the sampler state. - proposed_positions = self._propose_positions(initial_positions) + # This will also return box vectors for methodologies that change both position and box + proposed_positions, proposed_box_vectors = self._propose_positions( + initial_positions, initial_box_vectors + ) - log.debug(f"Proposed positions are {proposed_positions} nm.") + # log.debug(f"Proposed positions are {proposed_positions} nm.") # Compute the energy of the proposed positions. if atom_subset is None: sampler_state.x0 = proposed_positions @@ -421,16 +463,57 @@ def apply( sampler_state.x0 = sampler_state.x0.at[jnp.array(atom_subset)].set( proposed_positions ) + # set the box vectors in the sampler state if they changed + sampler_state.box_vectors = proposed_box_vectors + if not thermodynamic_state.pressure is None: + thermodynamic_state.volume = ( + proposed_box_vectors[0][0] + * proposed_box_vectors[1][1] + * proposed_box_vectors[2][2] + ) * unit.nanometer**3 + if nbr_list is not None: - if nbr_list.check(sampler_state.x0): + # nbr_list box_vectors setter also sets space.box_vectors + nbr_list.box_vectors = proposed_box_vectors + + # this will call the space.wrap function + sampler_state.x0 = nbr_list.wrap(sampler_state.x0) + + # if the box vectors changed, we need to rebuild the neighbor list + if jnp.any(initial_box_vectors != proposed_box_vectors): nbr_list.build(sampler_state.x0, sampler_state.box_vectors) + else: + # check to see if any particles have moved too far + if nbr_list.check(sampler_state.x0): + nbr_list.build(sampler_state.x0, sampler_state.box_vectors) proposed_energy = thermodynamic_state.get_reduced_potential( sampler_state, nbr_list ) # NOTE: in kT + log.debug(f"Proposed energy is {proposed_energy} kT.") + log.debug(f"Initial energy is {initial_energy} kT.") # Accept or reject with Metropolis criteria. delta_energy = proposed_energy - initial_energy log.debug(f"Delta energy is {delta_energy} kT.") + # if we are running NPT, we need to add in the N*ln(V/V0) term to delta energy + if thermodynamic_state.pressure is not None: + # units are nm**3 + proposed_volume = ( + proposed_box_vectors[0][0] + * proposed_box_vectors[1][1] + * proposed_box_vectors[2][2] + ) + thermodynamic_state.volume = proposed_volume * unit.nanometer**3 + + log.debug( + f"Proposed volume is {proposed_volume} initial volume {initial_volume}" + ) + log.debug(f"shape: {sampler_state.x0.shape[0]}") + volume_term = sampler_state.x0.shape[0] * jnp.log( + proposed_volume / initial_volume + ) + delta_energy -= volume_term + log.debug(f"Delta energy for NPT is {delta_energy} kT.") import jax.random as jrandom self.key, subkey = jrandom.split(self.key) @@ -439,20 +522,25 @@ def apply( if not jnp.isnan(proposed_energy) and ( delta_energy <= 0.0 or compare_to < jnp.exp(-delta_energy) ): + self.stored_energy = proposed_energy self.n_accepted += 1 - log.debug(f"Check suceeded: {compare_to=} < {jnp.exp(-delta_energy)}") + log.debug(f"Check succeeded: {compare_to=} < {jnp.exp(-delta_energy)}") log.debug( f"Move accepted. Energy change: {delta_energy:.3f} kT. Number of accepted moves: {self.n_accepted}." ) - reporter.report( - { - "energy": thermodynamic_state.kT_to_kJ_per_mol( - proposed_energy - ).value_in_unit_system(unit.md_unit_system), - "step": self.n_proposed, - "traj": sampler_state.x0, - } - ) + if reporter is not None: + reporter.report( + { + "energy": thermodynamic_state.kT_to_kJ_per_mol( + proposed_energy + ).value_in_unit_system(unit.md_unit_system), + "volume": thermodynamic_state.volume.value_in_unit_system( + unit.md_unit_system + ), + "step": self.n_proposed, + "traj": sampler_state.x0, + } + ) else: # Restore original positions. if atom_subset is None: @@ -461,13 +549,21 @@ def apply( sampler_state.x0 = sampler_state.x0.at[jnp.array([atom_subset])].set( initial_positions ) + if thermodynamic_state.pressure is not None: + thermodynamic_state.volume = initial_volume * unit.nanometer**3 + sampler_state.box_vectors = initial_box_vectors + if nbr_list is not None: + nbr_list.build(sampler_state.x0, sampler_state.box_vectors) + self.stored_energy = initial_energy log.debug( f"Move rejected. Energy change: {delta_energy:.3f} kT. Number of rejected moves: {self.n_proposed - self.n_accepted}." ) self.n_proposed += 1 - def _propose_positions(self, positions: jnp.array): - """Return new proposed positions. + def _propose_positions( + self, positions: jnp.array, box_vectors: jnp.array + ) -> Tuple[jnp.array, jnp.array]: + """Return new proposed positions and changes to box vectors. These method must be implemented in subclasses. @@ -476,11 +572,14 @@ def _propose_positions(self, positions: jnp.array): positions : nx3 jnp.ndarray The original positions of the subset of atoms that these move applied to. + box_vectors : 3x3 jnp.ndarray + The box vectors of the system. Returns ------- - proposed_positions : nx3 jnp.ndarray - The new proposed positions. + Tuple[jnp.array, jnp.array] + proposed_positions : nx3 jnp.ndarray, box_vectors : 3x3 jnp.ndarray + The new proposed positions and box vectors. """ raise NotImplementedError( @@ -590,9 +689,183 @@ def displace_positions( return updated_position - def _propose_positions(self, initial_positions: jnp.array) -> jnp.array: + def _propose_positions( + self, initial_positions: jnp.array, box_vectors: jnp.array + ) -> Tuple[jnp.array, jnp.array]: + """Implement MetropolizedMove._propose_positions for apply().""" + return ( + self.displace_positions(initial_positions, self.displacement_sigma), + box_vectors, + ) + + def run( + self, + sampler_state: SamplerState, + thermodynamic_state: ThermodynamicState, + nbr_list=None, + progress_bar=True, + ): + from tqdm import tqdm + + if nbr_list is not None: + if not nbr_list.is_built: + nbr_list.build_from_state(sampler_state) + + for trials in ( + tqdm(range(self.nr_of_moves)) if progress_bar else range(self.nr_of_moves) + ): + self.apply( + thermodynamic_state, + sampler_state, + self.simulation_reporter, + trials, + nbr_list, + ) + if trials % 100 == 0: + log.debug(f"Acceptance rate: {self.n_accepted / self.n_proposed}") + # if self.simulation_reporter is not None: + # self.simulation_reporter.report( + # { + # "Acceptance rate": self.n_accepted / self.n_proposed, + # "step": self.n_proposed, + # } + # ) + + log.info(f"Acceptance rate: {self.n_accepted / self.n_proposed}") + + +class MCBarostatMove(MetropolizedMove): + """A metropolized move that randomly displace a subset of atoms. + + Parameters + ---------- + displacement_sigma : openmm.unit.Quantity + The standard deviation of the normal distribution used to propose the + random displacement (units of length, default is 1.0*nanometer). + atom_subset : slice or list of int, optional + If specified, the move is applied only to those atoms specified by these + indices. If None, the move is applied to all atoms (default is None). + + Attributes + ---------- + n_accepted : int + The number of proposals accepted. + n_proposed : int + The total number of attempted moves. + displacement_sigma + atom_subset + + See Also + -------- + MetropolizedMove + + """ + + def __init__( + self, + seed: int = 1234, + volume_max_scale=0.01, + nr_of_moves: int = 100, + adjust_box_scaling: bool = False, + adjust_frequency: int = 100, + atom_subset: Optional[List[int]] = None, + simulation_reporter: Optional[SimulationReporter] = None, + ): + """ + Initialize the MCMC class. + + Parameters + ---------- + seed : int, optional + The seed for the random number generator. Default is 1234. + volume_max_scale : float + The length scale to apply to each box length of the system. Default is 1.1. + nr_of_moves : int, optional + The number of moves to perform. Default is 100. + adjust_box_scaling : bool, optional + Whether to adjust the box scaling based on the acceptance rate. Default is False. + adjust_frequency : int, optional + How often to adjust the box scaling of adjust_box_scaling is True. Default is 100. + atom_subset : list of int, optional + A subset of atom indices to consider for the moves. Default is None. + simulation_reporter : SimulationReporter, optional + The reporter to write the data to. Default is None. + Returns + ------- + None + """ + + super().__init__(nr_of_moves=nr_of_moves, seed=seed) + self.volume_max_scale = volume_max_scale + + self.atom_subset = atom_subset + self.adjust_box_scaling = adjust_box_scaling + self.adjust_frequency = adjust_frequency + + self.simulation_reporter = simulation_reporter + if self.simulation_reporter is not None: + log.info( + f"Using reporter {self.simulation_reporter} saving to {self.simulation_reporter.filename}" + ) + + def scale_positions_and_box( + self, positions: jnp.array, volume_max_scale: float, box_vectors: jnp.array + ): + """Return the positions after applying a random displacement to them. + + Parameters + ---------- + positions : nx3 jnp.array unit.Quantity + The positions to displace. + volume_max_scale : float + The maximum magnitude of the volume scaling. + box_vectors : 3x3 jnp.ndarray + The box vectors of the system. + + Returns + ------- + Tuple[jnp.array, jnp.array] + updated_positions : nx3 numpy.ndarray openmm.unit.Quantity of scaled positions + scaled_box : 3x3 jnp.ndarray of scaled box vectors + + """ + import jax.random as jrandom + + self.key, subkey = jrandom.split(self.key) + nr_of_atoms = positions.shape[0] + + volume = box_vectors[0][0] * box_vectors[1][1] * box_vectors[2][2] + + # Calculate the maximum amount the volume can change by + delta_volume_max = volume_max_scale * volume + + # Calculate the volume change by generating a random number between -1 and 1 + # and multiplying by the maximum allowed volume change, delta_volume_max + delta_volume = jrandom.uniform(self.key, minval=-1, maxval=1) * delta_volume_max + + log.debug(f"Delta volume is {delta_volume}.") + + # calculate the new volume + new_volume = volume + delta_volume + log.debug(f"New volume is {new_volume}.") + + # calculate the length scale factor for particle positions and box vectors + length_scaling_factor = jnp.power(new_volume / volume, 1.0 / 3.0) + + log.debug(f"Length scaling factor is {length_scaling_factor}.") + + scaled_positions = positions * length_scaling_factor + scaled_box = box_vectors * length_scaling_factor + + return scaled_positions, scaled_box + + def _propose_positions( + self, initial_positions: jnp.array, box_vectors: jnp.array + ) -> Tuple[jnp.array, jnp.array]: """Implement MetropolizedMove._propose_positions for apply().""" - return self.displace_positions(initial_positions, self.displacement_sigma) + return self.scale_positions_and_box( + initial_positions, self.volume_max_scale, box_vectors + ) def run( self, @@ -603,20 +876,45 @@ def run( ): from tqdm import tqdm + if thermodynamic_state.pressure is None: + raise ValueError( + "Cannot run MCBarostatMove without a pressure specified in the thermodynamic state" + ) + + if nbr_list is not None: + if not nbr_list.is_built: + nbr_list.build_from_state(sampler_state) + for trials in ( tqdm(range(self.nr_of_moves)) if progress_bar else range(self.nr_of_moves) ): self.apply( - thermodynamic_state, sampler_state, self.simulation_reporter, nbr_list + thermodynamic_state, + sampler_state, + self.simulation_reporter, + trials, + nbr_list, ) if trials % 100 == 0: log.debug(f"Acceptance rate: {self.n_accepted / self.n_proposed}") - if self.simulation_reporter is not None: - self.simulation_reporter.report( - { - "Acceptance rate": self.n_accepted / self.n_proposed, - "step": self.n_proposed, - } - ) + log.debug(f"{self.n_proposed} {self.n_accepted}") + # adjust the volume scaling if the acceptance rate is too high or too low + if self.adjust_box_scaling: + if ( + self.n_proposed > self.adjust_frequency + and self.n_proposed % self.adjust_frequency == 0 + ): + if self.n_accepted < 0.25 * self.n_proposed: + self.volume_max_scale /= 1.1 + log.debug( + f"Acceptance rate is too low, reducing volume_max_scale to {self.volume_max_scale}" + ) + elif self.n_accepted > 0.75 * self.n_proposed: + self.volume_max_scale = min(1.1 * self.volume_max_scale, 0.3) + + log.debug( + f"Acceptance rate is too high, increasing volumeScale to {self.volume_max_scale}" + ) + log.debug(f"volume max scaling: {self.volume_max_scale}") log.info(f"Acceptance rate: {self.n_accepted / self.n_proposed}") diff --git a/chiron/neighbors.py b/chiron/neighbors.py index 69a9502..145dc31 100644 --- a/chiron/neighbors.py +++ b/chiron/neighbors.py @@ -360,6 +360,32 @@ def check(self, coordinates: jnp.array) -> bool: """ pass + @property + def box_vectors(self) -> jnp.array: + return self._box_vectors + + @box_vectors.setter + def box_vectors(self, box_vectors: jnp.array) -> None: + self._box_vectors = box_vectors + self.space.box_vectors = box_vectors + + def wrap(self, coordinates: jnp.array) -> jnp.array: + """ + Wrap the coordinates of the system. + + Parameters + ---------- + coordinates: jnp.array + Coordinates of the system + + Returns + ------- + jnp.array + Wrapped coordinates of the system + + """ + return self.space.wrap(coordinates) + class NeighborListNsqrd(PairsBase): """ @@ -547,12 +573,13 @@ def build( f"box_vectors should be a 3x3 array, shape provided: {box_vectors.shape}" ) - self.ref_coordinates = coordinates + # will also set the box vectors in the space object due to the setter in the ABC self.box_vectors = box_vectors + self.ref_coordinates = coordinates # the neighborlist assumes that the box vectors do not change between building and calculating the neighbor list # changes to the box vectors require rebuilding the neighbor list - self.space.box_vectors = self.box_vectors + # self.space.box_vectors = self.box_vectors # store the ids of all the particles self.particle_ids = jnp.array( @@ -959,8 +986,6 @@ def calculate(self, coordinates: jnp.array): f"Coordinates must have shape ({self.n_particles}, 3), found {coordinates.shape}" ) - # coordinates = self.space.wrap(coordinates) - n_neighbors, padding_mask, dist, r_ij = jax.vmap( self._calc_distance_per_particle, in_axes=(0, 0, 0, None) )(self.particle_ids, self.all_pairs, self.reduction_mask, coordinates) diff --git a/chiron/potential.py b/chiron/potential.py index b5982f0..db97e47 100644 --- a/chiron/potential.py +++ b/chiron/potential.py @@ -63,6 +63,72 @@ def compute_pairlist(self, positions, cutoff) -> jnp.array: return distance[interacting_mask], displacement_vectors[interacting_mask], pairs +class IdealGasPotential(NeuralNetworkPotential): + def __init__( + self, + topology: Topology, + ): + """ + Initialize the Ideal Gas potential. + + Parameters + ---------- + topology : Topology + The topology of the system + + """ + + if not isinstance(topology, Topology): + if not isinstance(topology, property): + if topology is not None: + raise TypeError( + f"Topology must be a Topology object or None, type(topology) = {type(topology)}" + ) + + self.topology = topology + + def compute_energy(self, positions: jnp.array, nbr_list=None, debug_mode=False): + """ + Compute the energy for an ideal gas, which is always 0. + + Parameters + ---------- + positions : jnp.array + The positions of the particles in the system + nbr_list : NeighborList, default=None + Instance of a neighbor list or pair list class to use. + If None, an unoptimized N^2 pairlist will be used without PBC conditions. + Returns + ------- + potential_energy : float + The total potential energy of the system. + + """ + # Compute the pair distances and displacement vectors + + return 0.0 + + def compute_force(self, positions: jnp.array, nbr_list=None) -> jnp.array: + """ + Compute the force for ideal gas particles, which is always 0. + + Parameters + ---------- + positions : jnp.array + The positions of the particles in the system + nbr_list : NeighborList, optional + Instance of the neighborlist class to use. By default, set to None, which will use an N^2 pairlist + + Returns + ------- + force : jnp.array + The forces on the particles in the system + + """ + + return 0.0 + + class LJPotential(NeuralNetworkPotential): def __init__( self, diff --git a/chiron/states.py b/chiron/states.py index edfae43..99d9e3f 100644 --- a/chiron/states.py +++ b/chiron/states.py @@ -92,6 +92,13 @@ def x0(self, x0: Union[jnp.array, unit.Quantity]) -> None: else: self._x0 = unit.Quantity(x0, self._distance_unit) + @box_vectors.setter + def box_vectors(self, box_vectors: Union[jnp.array, unit.Quantity]) -> None: + if isinstance(box_vectors, unit.Quantity): + self._box_vectors = box_vectors + else: + self._box_vectors = unit.Quantity(box_vectors, self._distance_unit) + @property def distance_unit(self) -> unit.Unit: return self._distance_unit @@ -179,6 +186,7 @@ def __init__( self.beta = None self.volume = volume + self.pressure = pressure from .utils import get_nr_of_particles @@ -264,20 +272,26 @@ def get_reduced_potential( and N(x) is the number of particles. """ if self.beta is None: - self.beta = 1.0 / ( - unit.BOLTZMANN_CONSTANT_kB * (self.temperature * unit.kelvin) - ) - log.debug(f"sample state: {sampler_state.x0}") + self.beta = 1.0 / (unit.BOLTZMANN_CONSTANT_kB * (self.temperature)) + reduced_potential = ( unit.Quantity( self.potential.compute_energy(sampler_state.x0, nbr_list), unit.kilojoule_per_mole, ) ) / unit.AVOGADRO_CONSTANT_NA - log.debug(f"reduced potential: {reduced_potential}") if self.pressure is not None: + # in case volume is not set, calculate from the box vectors + if self.volume is None: + self.volume = ( + sampler_state.box_vectors[0][0] + * sampler_state.box_vectors[1][1] + * sampler_state.box_vectors[2][2] + ) * unit.nanometer**3 reduced_potential += self.pressure * self.volume + log.debug(f"reduced potential energy: {reduced_potential}") + return self.beta * reduced_potential def kT_to_kJ_per_mol(self, energy): diff --git a/chiron/tests/test_convergence_tests.py b/chiron/tests/test_convergence_tests.py index 3deff74..5a3a185 100644 --- a/chiron/tests/test_convergence_tests.py +++ b/chiron/tests/test_convergence_tests.py @@ -44,7 +44,9 @@ def test_convergence_of_MC_estimator(prep_temp_dir): from chiron.states import ThermodynamicState, SamplerState thermodynamic_state = ThermodynamicState( - harmonic_potential, temperature=300, volume=30 * (unit.angstrom**3) + harmonic_potential, + temperature=300 * unit.kelvin, + volume=30 * (unit.angstrom**3), ) sampler_state = SamplerState(ho.positions) @@ -52,7 +54,9 @@ def test_convergence_of_MC_estimator(prep_temp_dir): id = uuid.uuid4() - simulation_reporter = SimulationReporter(f"{prep_temp_dir}/test_{id}.h5") + simulation_reporter = SimulationReporter( + f"{prep_temp_dir}/test_{id}.h5", ho.topology + ) # Initalize the move set (here only LangevinDynamicsMove) from chiron.mcmc import MetropolisDisplacementMove, MoveSet, MCMCSampler @@ -113,7 +117,7 @@ def __init__(self, temperature): ) -@pytest.mark.skip(reason="Tests takes too long") +# @pytest.mark.skip(reason="Tests takes too long") @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Test takes too long.") def test_langevin_dynamics_with_LJ_fluid(prep_temp_dir): from chiron.integrators import LangevinIntegrator @@ -156,7 +160,7 @@ def test_langevin_dynamics_with_LJ_fluid(prep_temp_dir): from chiron.reporters import SimulationReporter id = uuid.uuid4() - reporter = SimulationReporter(f"{prep_temp_dir}/test_{id}.h5") + reporter = SimulationReporter(f"{prep_temp_dir}/test_{id}.h5", lj_fluid.topology) integrator = LangevinIntegrator(reporter=reporter, save_frequency=100) integrator.run( diff --git a/chiron/tests/test_mcmc.py b/chiron/tests/test_mcmc.py index 950b530..92087c1 100644 --- a/chiron/tests/test_mcmc.py +++ b/chiron/tests/test_mcmc.py @@ -275,6 +275,113 @@ def test_thermodynamic_state_inputs(): ThermodynamicState(potential=harmonic_potential, pressure=100 * unit.atmosphere) +def test_mc_barostat_parameter_setting(): + import jax.numpy as jnp + from chiron.mcmc import MCBarostatMove + + barostat_move = MCBarostatMove( + seed=1234, + volume_max_scale=0.1, + nr_of_moves=1, + ) + + assert barostat_move.volume_max_scale == 0.1 + assert barostat_move.nr_of_moves == 1 + + # when a key is generated from a seed for the first time, we get a jnp.array([0, seed]) + assert jnp.all(barostat_move.key == jnp.array([0, 1234])) + + +def test_mc_barostat(): + import jax.numpy as jnp + from chiron.mcmc import MCBarostatMove + + barostat_move = MCBarostatMove( + seed=1234, + volume_max_scale=0.1, + nr_of_moves=1, + ) + + from chiron.potential import LJPotential + from openmm import unit + + sigma = 0.34 * unit.nanometer + epsilon = 0.0 * unit.kilocalories_per_mole + cutoff = sigma + + positions = ( + jnp.array( + [ + [0, 0, 0], + [1, 0, 0], + [0, 1, 0], + [0, 0, 1], + [1, 1, 0], + [1, 0, 1], + [0, 1, 1], + [1, 1, 1], + ] + ) + * unit.nanometer + ) + box_vectors = ( + jnp.array([[10.0, 0.0, 0.0], [0.0, 10.0, 0.0], [0.0, 0.0, 10.0]]) + * unit.nanometer + ) + volume = box_vectors[0][0] * box_vectors[1][1] * box_vectors[2][2] + + from openmm.app import Topology, Element + + topology = Topology() + element = Element.getBySymbol("Ar") + chain = topology.addChain() + residue = topology.addResidue("system", chain) + for i in range(positions.shape[0]): + topology.addAtom("Ar", element, residue) + + lj_potential = LJPotential(topology, sigma=sigma, epsilon=epsilon, cutoff=cutoff) + + from chiron.states import SamplerState, ThermodynamicState + + # define the sampler state + sampler_state = SamplerState( + x0=positions, + box_vectors=box_vectors, + ) + + # define the thermodynamic state + thermodynamic_state = ThermodynamicState( + potential=lj_potential, + temperature=300 * unit.kelvin, + pressure=1.0 * unit.atmosphere, + ) + + from chiron.neighbors import PairList, OrthogonalPeriodicSpace + + nbr_list = PairList(OrthogonalPeriodicSpace(), cutoff=cutoff) + + barostat_move.run(sampler_state, thermodynamic_state, nbr_list, True) + + # ideal gas treatment, so stored energy will only be a consequence of pressure, volume, and temperature + assert ( + barostat_move.stored_energy + == thermodynamic_state.pressure + * thermodynamic_state.volume + * thermodynamic_state.beta + ) + assert barostat_move.statistics["n_proposed"] == 1 + assert barostat_move.statistics["n_accepted"] == 0 + assert thermodynamic_state.volume == volume + + barostat_move.run(sampler_state, thermodynamic_state, nbr_list, True) + assert barostat_move.statistics["n_proposed"] == 2 + assert barostat_move.statistics["n_accepted"] == 1 + + assert jnp.isclose( + thermodynamic_state.volume.value_in_unit(unit.nanometer**3), 917.576978 + ) + + def test_sample_from_joint_distribution_of_two_HO_with_local_moves_and_MC_updates(): # define two harmonic oscillators with different spring constants and equilibrium positions # sample from the joint distribution of the two HO using local langevin moves diff --git a/chiron/tests/test_states.py b/chiron/tests/test_states.py index c9640a2..9a4b833 100644 --- a/chiron/tests/test_states.py +++ b/chiron/tests/test_states.py @@ -32,6 +32,14 @@ def test_initialize_state(): jnp.array([[0.0, 0.0, 0.0]]), ) + state = ThermodynamicState( + potential, + temperature=300 * unit.kelvin, + pressure=1 * unit.atmosphere, + volume=30 * (unit.angstrom**3), + ) + assert state.pressure == 1 * unit.atmosphere + def test_sampler_state_conversion(): """Test converting a sampler state to jnp arrays. @@ -148,6 +156,7 @@ def test_reduced_potential(): ho = HarmonicOscillator() potential = HarmonicOscillatorPotential(topology=ho.topology, k=ho.K, U0=ho.U0) + potential = HarmonicOscillatorPotential(topology=ho.topology, k=ho.K, U0=ho.U0) state = ThermodynamicState( potential, temperature=300 * unit.kelvin, volume=30 * (unit.angstrom**3) @@ -156,3 +165,20 @@ def test_reduced_potential(): reduced_e = state.get_reduced_potential(sampler_state) assert reduced_e == 0.0 + + # test the reduced potential for a system with pressure defined + box_vectors = ( + jnp.array([[10.0, 0.0, 0.0], [0.0, 10.0, 0.0], [0.0, 0.0, 10.0]]) + * unit.nanometers + ) + volume = box_vectors[0][0] * box_vectors[1][1] * box_vectors[2][2] + sampler_state.box_vectors = box_vectors + + pressure = 1.0 * unit.atmosphere + + state = ThermodynamicState( + potential, temperature=300 * unit.kelvin, volume=volume, pressure=pressure + ) + reduced_e = state.get_reduced_potential(sampler_state) + # since the energy will be zero for the ho at 0, the reduced potential will just be pressure * volume * beta + assert reduced_e == pressure * volume * state.beta diff --git a/chiron/tests/test_testsystems.py b/chiron/tests/test_testsystems.py index 5f14269..d139e13 100644 --- a/chiron/tests/test_testsystems.py +++ b/chiron/tests/test_testsystems.py @@ -1,3 +1,13 @@ +import pytest + + +@pytest.fixture(scope="session") +def prep_temp_dir(tmpdir_factory): + """Create a temporary directory for the test.""" + tmpdir = tmpdir_factory.mktemp("test_testsystems") + return tmpdir + + def compute_openmm_reference_energy(testsystem, positions): from openmm import unit from openmm.app import Simulation @@ -203,3 +213,115 @@ def test_LJ_fluid(): assert jnp.isclose( e_chiron_energy, e_openmm_energy.value_in_unit_system(unit.md_unit_system) ), "Chiron LJ fluid energy does not match openmm" + + +def test_ideal_gas(prep_temp_dir): + from openmmtools.testsystems import IdealGas + from openmm import unit + + # 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 + n_particles = 216 + temperature = 298 * unit.kelvin + pressure = 1 * unit.atmosphere + mass = unit.Quantity(39.9, unit.gram / unit.mole) + + ideal_gas = IdealGas( + nparticles=n_particles, temperature=temperature, pressure=pressure + ) + + from chiron.potential import IdealGasPotential + import jax.numpy as jnp + + # + 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, + ) + + # define the sampler state + sampler_state = SamplerState( + x0=ideal_gas.positions, + box_vectors=ideal_gas.system.getDefaultPeriodicBoxVectors(), + ) + + from chiron.neighbors import PairList, 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 = PairList(OrthogonalPeriodicSpace(), cutoff=cutoff) + nbr_list.build_from_state(sampler_state) + + from chiron.reporters import SimulationReporter + + # initialize a reporter to save the simulation data + filename = f"{prep_temp_dir}/test_ideal_gas_vol.h5" + + reporter1 = SimulationReporter(filename, ideal_gas.topology, 100) + + from chiron.mcmc import ( + MetropolisDisplacementMove, + MoveSet, + MCMCSampler, + MCBarostatMove, + ) + + mc_disp_move = MetropolisDisplacementMove( + seed=1234, + displacement_sigma=0.1 * unit.nanometer, + nr_of_moves=10, + ) + + mc_barostat_move = MCBarostatMove( + seed=1234, + volume_max_scale=0.2, + nr_of_moves=100, + adjust_box_scaling=True, + adjust_frequency=50, + simulation_reporter=reporter1, + ) + move_set = MoveSet( + [ + ("MetropolisDisplacementMove", mc_disp_move), + ("MCBarostatMove", mc_barostat_move), + ] + ) + + sampler = MCMCSampler(move_set, sampler_state, thermodynamic_state) + sampler.run(n_iterations=30, nbr_list=nbr_list) # how many times to repeat + + import h5py + + with h5py.File(filename, "r") as f: + volume = f["volume"][:] + steps = f["step"][:] + + # get expectations + ideal_volume = ideal_gas.get_volume_expectation(thermodynamic_state) + ideal_volume_std = ideal_gas.get_volume_standard_deviation(thermodynamic_state) + + volume_mean = jnp.mean(jnp.array(volume)) * unit.nanometer**3 + volume_std = jnp.std(jnp.array(volume)) * unit.nanometer**3 + + ideal_density = mass * n_particles / unit.AVOGADRO_CONSTANT_NA / ideal_volume + measured_density = mass * n_particles / 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 + + # see if within 10% of the ideal standard deviation of the volume + assert abs(ideal_volume_std - volume_std) / ideal_volume_std < 0.1