From b313867cf283d8b1ae35d56a55e9432dc322e6a2 Mon Sep 17 00:00:00 2001 From: wiederm Date: Thu, 18 Jan 2024 16:50:54 +0100 Subject: [PATCH 01/43] Fix u_kn shape in test_multistate_run --- chiron/tests/test_multistate.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/chiron/tests/test_multistate.py b/chiron/tests/test_multistate.py index ce74e1d..1a36de4 100644 --- a/chiron/tests/test_multistate.py +++ b/chiron/tests/test_multistate.py @@ -207,6 +207,8 @@ def test_multistate_run(ho_multistate_sampler_multiple_ks: MultiStateSampler): assert ho_sampler.n_replicas == 4 assert ho_sampler.n_states == 4 + u_kn = ho_sampler._reporter.get_property("u_kn") + assert u_kn.shape == (4, 4, n_iteratinos + 1) # check that the free energies are correct print(ho_sampler.analytical_f_i) print(ho_sampler.delta_f_ij_analytical) From a9d9f166984c288d4940addd0caf8bd2286ceecd Mon Sep 17 00:00:00 2001 From: wiederm Date: Thu, 18 Jan 2024 22:36:48 +0100 Subject: [PATCH 02/43] Refactor MCMCMove subclasses and update method names --- chiron/mcmc.py | 124 +++++++++++++++++++------------------------------ 1 file changed, 48 insertions(+), 76 deletions(-) diff --git a/chiron/mcmc.py b/chiron/mcmc.py index 2cb7d75..4f26856 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -4,6 +4,10 @@ import jax.numpy as jnp from chiron.reporters import SimulationReporter + +# MCMCMOve - > MCMove with differenf flavors + + class MCMCMove: def __init__(self, nr_of_moves: int, seed: int): """ @@ -21,6 +25,11 @@ def __init__(self, nr_of_moves: int, seed: int): self.nr_of_moves = nr_of_moves self.key = jrandom.PRNGKey(seed) # 'seed' is an integer seed value + # draw proposal move + # compute probability of acceptance + # TOD: @abc + # def run + class LangevinDynamicsMove(MCMCMove): def __init__( @@ -59,16 +68,17 @@ def __init__( save_traj_in_memory=save_traj_in_memory, ) - def run( + def update( self, sampler_state: SamplerState, thermodynamic_state: ThermodynamicState, ): """ - Run the integrator to perform molecular dynamics simulation. + Update the sampler state in place by running the langevin integrator. Args: state_variables (StateVariablesCollection): State variables of the system. + # NOTE: update """ assert isinstance( @@ -78,6 +88,7 @@ def run( thermodynamic_state, ThermodynamicState ), f"Thermodynamic state must be ThermodynamicState, not {type(thermodynamic_state)}" + # NOTE: should this return the sampler state object? self.integrator.run( thermodynamic_state=thermodynamic_state, sampler_state=sampler_state, @@ -93,29 +104,7 @@ class MCMove(MCMCMove): def __init__(self, nr_of_moves: int, seed: int) -> None: super().__init__(nr_of_moves, seed) - def _check_state_compatiblity( - self, - old_state: SamplerState, - new_state: SamplerState, - ): - """ - Check if the states are compatible. - - Parameters - ---------- - old_state : StateVariablesCollection - The state of the system before the move. - new_state : StateVariablesCollection - The state of the system after the move. - - Raises - ------ - ValueError - If the states are not compatible. - """ - pass - - def apply_move(self): + def step(self): """ Apply a Monte Carlo move to the system. @@ -126,59 +115,27 @@ def apply_move(self): NotImplementedError If the method is not implemented in subclasses. """ - + # needs to be subclassed + # ( + # proposed_sampler_state, + # log_proposal_ratio, + # proposed_thermodynamic_state, + # ) = self._propose(current_sampler_state, current_thermodynamic_state) # log proposal ratio + proposal sampler state + # # current_reduced_pot = current_thermodynamic_state.get_reduced_potential(current_sampler_state) + # # proposed_reduced_pot = proposed_thermodynamic_state.get_reduced_potential(proposed_sampler_state) + # decicion = self._accept_or_reject( + # current_reduced_pot, + # proposed_reduced_pot, + # log_proposal_ratio, + # method="metropolis", # or other flavors + # ) # including the log acceptance ratio + # if decicion: + # self._replace_states(proposed_sampler_state, proposed_thermodynamic_state) raise NotImplementedError("apply_move() must be implemented in subclasses") - def compute_acceptance_probability( - self, - old_state: SamplerState, - new_state: SamplerState, - ): - """ - Compute the acceptance probability for a move from an old state to a new state. - - Parameters - ---------- - old_state : object - The state of the system before the move. - new_state : object - The state of the system after the move. - - Returns - ------- - float - Acceptance probability as a float. - """ - self._check_state_compatiblity(old_state, new_state) - old_system = self.system(old_state) - new_system = self.system(new_state) - - energy_before_state_change = old_system.compute_energy(old_state.position) - energy_after_state_change = new_system.compute_energy(new_state.position) - # Implement the logic to compute the acceptance probability - pass - - def accept_or_reject(self, probability): - """ - Decide whether to accept or reject the move based on the acceptance probability. - - Parameters - ---------- - probability : float - Acceptance probability. - - Returns - ------- - bool - Boolean indicating if the move is accepted. - """ - import jax.numpy as jnp - - return jnp.random.rand() < probability - class RotamerMove(MCMove): - def apply_move(self): + def step(self): """ Implement the logic specific to rotamer changes. """ @@ -186,15 +143,21 @@ def apply_move(self): class ProtonationStateMove(MCMove): - def apply_move(self): + def step(self): """ Implement the logic specific to protonation state changes. """ + + # this becomes more complicated + # proposed_sampler_state, proposed_thermodynamic_state = self._propose() + + # ... + # ... pass class TautomericStateMove(MCMove): - def apply_move(self): + def step(self): """ Implement the logic specific to tautomeric state changes. """ @@ -238,7 +201,16 @@ def _validate_sequence(self): if not isinstance(move_class, MCMCMove): raise ValueError(f"Move {move_name} in the sequence is not available.") + # self.acceptance_statistics = [] # provide probabilites for each MC move + + # def _bias_sequence_based_on_acceptace_statistics() + + # def get_sequence(): + + # def random_sequence() + +# NOTE: update the MultistateSampler class using the MCMCSampler class MCMCSampler(object): """ Basic Markov chain Monte Carlo Gibbs sampler. From 4af534a5431f84d1c5fc518b869b555eba27f8d6 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Sun, 21 Jan 2024 22:17:03 -0800 Subject: [PATCH 03/43] Started MC refactor, merging from multistage branch that includes new reporters and random number scheme. This primarily sketches out the new MCMove class (updates to the actual moves is forthcoming). --- chiron/mcmc.py | 211 ++++++++++++++++++++++++++++++++++++++----------- 1 file changed, 167 insertions(+), 44 deletions(-) diff --git a/chiron/mcmc.py b/chiron/mcmc.py index 8063012..57e22aa 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -4,6 +4,8 @@ import jax.numpy as jnp from chiron.reporters import LangevinDynamicsReporter, _SimulationReporter +from abc import ABC, abstractmethod + class MCMCMove: def __init__( @@ -36,6 +38,28 @@ def __init__( ) assert self.report_frequency is not None + @abstractmethod + def update( + self, + sampler_state: SamplerState, + thermodynamic_state: ThermodynamicState, + nbr_list: Optional[PairsBase] = None, + ): + """ + Update the state of the system. + + Parameters + ---------- + sampler_state : SamplerState + The sampler state to run the integrator on. + thermodynamic_state : ThermodynamicState + The thermodynamic state to run the integrator on. + nbr_list : PairsBase, optional + The neighbor list to use for the simulation. + + """ + pass + class LangevinDynamicsMove(MCMCMove): def __init__( @@ -89,10 +113,11 @@ def __init__( save_traj_in_memory=save_traj_in_memory, ) - def run( + def update( self, sampler_state: SamplerState, thermodynamic_state: ThermodynamicState, + nbr_list: Optional[PairsBase] = None, ): """ Run the integrator to perform molecular dynamics simulation. @@ -103,6 +128,8 @@ def run( The sampler state to run the integrator on. thermodynamic_state : ThermodynamicState The thermodynamic state to run the integrator on. + nbr_list : PairsBase, optional + The neighbor list to use for the simulation. """ assert isinstance( @@ -116,6 +143,7 @@ def run( thermodynamic_state=thermodynamic_state, sampler_state=sampler_state, n_steps=self.nr_of_moves, + nbr_list=nbr_list, ) if self.save_traj_in_memory: @@ -125,74 +153,169 @@ def run( class MCMove(MCMCMove): def __init__( - self, nr_of_moves: int, reporter: Optional[_SimulationReporter] + self, + nr_of_moves: int, + reporter: Optional[_SimulationReporter], + report_frequency: int = 1, + method: str = "metropolis", ) -> None: - super().__init__(nr_of_moves, reporter=reporter) - - def apply_move(self): """ - Apply a Monte Carlo move to the system. - - This method should be overridden by subclasses to define specific types of moves. + Initialize the move. - Raises - ------ - NotImplementedError - If the method is not implemented in subclasses. + Parameters + ---------- + nr_of_moves + Number of moves to be applied in each call to update. + reporter + Reporter object for saving the simulation step data. + report_frequency + Frequency of saving the simulation data. + method + Methodology to use for accepting or rejecting the proposed state. + Default is "metropolis". """ + super().__init__(nr_of_moves, reporter=reporter) + self.method = "metropolis" - raise NotImplementedError("apply_move() must be implemented in subclasses") - - def compute_acceptance_probability( + def update( self, - old_state: SamplerState, - new_state: SamplerState, + sampler_state: SamplerState, + thermodynamic_state: ThermodynamicState, + nbr_list: Optional[PairsBase] = None, ): """ - Compute the acceptance probability for a move from an old state to a new state. + Perform the defined move and update the state. Parameters ---------- - old_state : object - The state of the system before the move. - new_state : object - The state of the system after the move. + sampler_state : SamplerState + The initial state of the simulation, including positions. + thermodynamic_state : ThermodynamicState + The thermodynamic state of the system, including temperature and potential. + nbr_list : PairBase, optional + Neighbor list for the system. + - Returns - ------- - float - Acceptance probability as a float. """ - self._check_state_compatiblity(old_state, new_state) - old_system = self.system(old_state) - new_system = self.system(new_state) + calculate_current_potential = True + for i in range(self.nr_of_moves): + self._step( + sampler_state, + thermodynamic_state, + nbr_list, + calculate_current_potential=calculate_current_potential, + ) + # after the first step, we don't need to recalculate the current potential, it will be stored + calculate_current_potential = False - energy_before_state_change = old_system.compute_energy(old_state.position) - energy_after_state_change = new_system.compute_energy(new_state.position) - # Implement the logic to compute the acceptance probability - pass + def _step( + self, + ): + # if this is the first time we are calling this, + # we will need to recalculate the reduced potential for the current state + if calculate_current_potential: + current_reduced_pot = current_thermodynamic_state.get_reduced_potential( + current_sampler_state + ) + # save the current_reduced_pot so we don't have to recalculate + # it on the next iteration if the move is rejected + self._current_reduced_pot = current_reduced_pot + else: + current_reduced_pot = self._current_reduced_pot + + # propose a new state and calculate the log proposal ratio + # this will be specific to the type of move + # in addition to the sampler_state, this will require/return the thermodynamic state + # for systems that e.g., make changes to particle identity. + ( + proposed_sampler_state, + log_proposal_ratio, + proposed_thermodynamic_state, + ) = self._propose(current_sampler_state, current_thermodynamic_state) + + # calculate the reduced potential for the proposed state + proposed_reduced_pot = proposed_thermodynamic_state.get_reduced_potential( + proposed_sampler_state + ) + + # accept or reject the proposed state + decision = self._accept_or_reject( + current_reduced_pot, + proposed_reduced_pot, + log_proposal_ratio, + method=self.method, + ) - def accept_or_reject(self, probability): + if decision: + # save the reduced potential of the accepted state so + # we don't have to recalculate it the next iteration + self._current_reduced_pot = proposed_reduced_pot + + # replace the current state with the proposed state + # not sure this needs to be a separate function but for simplicity in outlining the code it is fine + # or should this return the new sampler_state and thermodynamic_state? + self._replace_states( + current_sampler_state, + proposed_sampler_state, + current_thermodynamic_state, + proposed_thermodynamic_state, + ) + + # a function that will update the statistics for the move + self._update_statistics(decision) + + @abstractmethod + def _propose(self, current_sampler_state, current_thermodynamic_state): """ - Decide whether to accept or reject the move based on the acceptance probability. + Propose a new state and calculate the log proposal ratio. + + This will need to be defined for each move Parameters ---------- - probability : float - Acceptance probability. + current_sampler_state : SamplerState, required + Current sampler state. + current_thermodynamic_state : ThermodynamicState, required Returns ------- - bool - Boolean indicating if the move is accepted. + proposed_sampler_state : SamplerState + Proposed sampler state. + log_proposal_ratio : float + Log proposal ratio. + proposed_thermodynamic_state : ThermodynamicState + Proposed thermodynamic state. + """ - import jax.numpy as jnp + pass - return jnp.random.rand() < probability + def _replace_states( + self, + current_sampler_state, + proposed_sampler_state, + current_thermodynamic_state, + proposed_thermodynamic_state, + ): + """ + Replace the current state with the proposed state. + """ + # define the code to copy the proposed state to the current state + + def _accept_or_reject( + self, + current_reduced_pot, + proposed_reduced_pot, + log_proposal_ratio, + method=method, + ): + """ + Accept or reject the proposed state with a given methodology. + """ + # define the acceptance probability class RotamerMove(MCMove): - def apply_move(self): + def _step(self): """ Implement the logic specific to rotamer changes. """ @@ -200,7 +323,7 @@ def apply_move(self): class ProtonationStateMove(MCMove): - def apply_move(self): + def _step(self): """ Implement the logic specific to protonation state changes. """ @@ -208,7 +331,7 @@ def apply_move(self): class TautomericStateMove(MCMove): - def apply_move(self): + def _step(self): """ Implement the logic specific to tautomeric state changes. """ From 53f06166d103b55a9bc1962fec3b31a8e3b32219 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Tue, 23 Jan 2024 16:37:48 -0500 Subject: [PATCH 04/43] Small refactor to Langevin integrator. Small changes: run command now takes variable to allow for velocity initialization. The Langevin ntegrator now returns a sampler state rather tha updating. States now has a velocity setter. --- Examples/LJ_langevin.py | 19 ++++++++++++++----- chiron/integrators.py | 24 ++++++++++++++++++++---- chiron/mcmc.py | 14 ++++++++++++-- chiron/states.py | 23 +++++++++++++++++++++++ chiron/tests/test_integrators.py | 10 ++++++++++ 5 files changed, 79 insertions(+), 11 deletions(-) diff --git a/Examples/LJ_langevin.py b/Examples/LJ_langevin.py index d769b1a..55d8e0b 100644 --- a/Examples/LJ_langevin.py +++ b/Examples/LJ_langevin.py @@ -8,6 +8,7 @@ from chiron.potential import LJPotential from openmm import unit +from chiron.utils import PRNG # initialize the LennardJones potential in chiron # @@ -21,9 +22,12 @@ from chiron.states import SamplerState, ThermodynamicState +PRNG.set_seed(1234) # define the sampler state sampler_state = SamplerState( - x0=lj_fluid.positions, box_vectors=lj_fluid.system.getDefaultPeriodicBoxVectors() + x0=lj_fluid.positions, + current_PRNG_key=PRNG.get_random_key(), + box_vectors=lj_fluid.system.getDefaultPeriodicBoxVectors(), ) # define the thermodynamic state @@ -43,7 +47,7 @@ # build the neighbor list from the sampler state nbr_list.build_from_state(sampler_state) -from chiron.reporters import _SimulationReporter +from chiron.reporters import LangevinDynamicsReporter # initialize a reporter to save the simulation data filename = "test_lj.h5" @@ -51,7 +55,11 @@ if os.path.isfile(filename): os.remove(filename) -reporter = _SimulationReporter("test_lj.h5", lj_fluid.topology, 1) +reporter = LangevinDynamicsReporter( + "test_lj.h5", + 1, + lj_fluid.topology, +) from chiron.integrators import LangevinIntegrator @@ -59,19 +67,20 @@ integrator = LangevinIntegrator(reporter=reporter, report_frequency=100) print("init_energy: ", lj_potential.compute_energy(sampler_state.x0, nbr_list)) -integrator.run( +updated_sampler_state = integrator.run( sampler_state, thermodynamic_state, n_steps=5000, nbr_list=nbr_list, progress_bar=True, + initialize_velocities=True, ) import h5py # read the data from the reporter with h5py.File("test_lj.h5", "r") as f: - energies = f["energy"][:] + energies = f["potential_energy"][:] steps = f["step"][:] diff --git a/chiron/integrators.py b/chiron/integrators.py index 0d77452..8018647 100644 --- a/chiron/integrators.py +++ b/chiron/integrators.py @@ -85,6 +85,7 @@ def run( n_steps: int = 5_000, nbr_list: Optional[PairsBase] = None, progress_bar=False, + initialize_velocities: bool = False, ): """ Run the integrator to perform Langevin dynamics molecular dynamics simulation. @@ -102,6 +103,10 @@ def run( progress_bar : bool, optional Flag indicating whether to display a progress bar during integration. + Returns + ------- + sampler_state : SamplerState + The final state of the simulation, including positions, velocities, and current PRNG key. """ from .utils import get_list_of_mass from tqdm import tqdm @@ -137,9 +142,12 @@ def run( b = jnp.sqrt(1 - jnp.exp(-2 * collision_rate_unitless * stepsize_unitless)) # Initialize velocities - if self.velocities is None: + if initialize_velocities: + # we should probably move this to a separate function for unit testing purposes v0 = sigma_v * random.normal(key, x0.shape) else: + if self.velocities is None: + raise ValueError("Velocities must be set before running the integrator") v0 = self.velocities.value_in_unit_system(unit.md_unit_system) x = x0 @@ -180,9 +188,17 @@ def run( self.traj.append(x) log.debug("Finished running Langevin dynamics") - # save the final state of the simulation in the sampler_state object - sampler_state.x0 = x - sampler_state.v0 = v + + # return the final state of the simulation as a sampler_state object + import copy + + updated_sampler_state = copy.deepcopy(sampler_state) + + updated_sampler_state.x0 = x + updated_sampler_state._velocities = v + updated_sampler_state.current_PRNG_key = key + + return updated_sampler_state def _wrap_and_rebuild_neighborlist(self, x: jnp.array, nbr_list: PairsBase): """ diff --git a/chiron/mcmc.py b/chiron/mcmc.py index 57e22aa..29093b1 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -57,6 +57,13 @@ def update( nbr_list : PairsBase, optional The neighbor list to use for the simulation. + Returns + ------- + sampler_state : SamplerState + The updated sampler state. + thermodynamic_state : ThermodynamicState + The updated thermodynamic state. + """ pass @@ -139,7 +146,7 @@ def update( thermodynamic_state, ThermodynamicState ), f"Thermodynamic state must be ThermodynamicState, not {type(thermodynamic_state)}" - self.integrator.run( + updated_sampler_state = self.integrator.run( thermodynamic_state=thermodynamic_state, sampler_state=sampler_state, n_steps=self.nr_of_moves, @@ -150,6 +157,9 @@ def update( self.traj.append(self.integrator.traj) self.integrator.traj = [] + # The thermodynamic_state will not change for the langevin move + return updated_sampler_state, thermodynamic_state + class MCMove(MCMCMove): def __init__( @@ -175,7 +185,7 @@ def __init__( Default is "metropolis". """ super().__init__(nr_of_moves, reporter=reporter) - self.method = "metropolis" + self.method = "metropolis" # I think we should pass a class/function instead of a string, like space. def update( self, diff --git a/chiron/states.py b/chiron/states.py index 99459ae..1dcef10 100644 --- a/chiron/states.py +++ b/chiron/states.py @@ -18,6 +18,18 @@ class SamplerState: box_vectors : unit.Quantity, optional The box vectors defining the simulation's periodic boundary conditions. + Examples + -------- + + from chiron.states import SamplerState + from chiron.utils import PRNG + from openmmtools.testsystems import HarmonicOscillator + + ho = HarmonicOscillator() + PRNG.set_seed(1234) + + sampler_state = SamplerState(x0 = ho.positions, PRNG.get_random_key()) + """ def __init__( @@ -75,6 +87,7 @@ def __init__( self._current_PRNG_key = current_PRNG_key self._box_vectors = box_vectors self._distance_unit = unit.nanometer + self._velocity_unit = unit.nanometer / unit.picosecond @property def n_particles(self) -> int: @@ -103,10 +116,20 @@ def x0(self, x0: Union[jnp.array, unit.Quantity]) -> None: else: self._x0 = unit.Quantity(x0, self._distance_unit) + @velocities.setter + def velocities(self, velocities: Union[jnp.array, unit.Quantity]) -> None: + if isinstance(velocities, unit.Quantity): + self._velocities = velocities + else: + self._velocities = unit.Quantity(velocities, self._velocity_unit) + @property def distance_unit(self) -> unit.Unit: return self._distance_unit + def velocity_unit(self) -> unit.Unit: + return self._velocity_unit + @property def new_PRNG_key(self) -> random.PRNGKey: key, subkey = random.split(self._current_PRNG_key) diff --git a/chiron/tests/test_integrators.py b/chiron/tests/test_integrators.py index 341c987..2e67928 100644 --- a/chiron/tests/test_integrators.py +++ b/chiron/tests/test_integrators.py @@ -43,9 +43,19 @@ def test_langevin_dynamics(prep_temp_dir, provide_testsystems_and_potentials): reporter = LangevinDynamicsReporter() integrator = LangevinIntegrator(reporter=reporter, report_frequency=1) + + with pytest.raises(ValueError): + integrator.run( + sampler_state, + thermodynamic_state, + n_steps=20, + initialize_velocities=False, + ) + integrator.run( sampler_state, thermodynamic_state, n_steps=20, + initialize_velocities=True, ) i = i + 1 From 01176307be90327944326cb3e446f171be470203 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Tue, 23 Jan 2024 17:06:29 -0500 Subject: [PATCH 05/43] created initialize_velocities funtion in utilities; langevin code now uses this if users request velocities to be regenerated each time. The reaosn to split was to make it easy to define velocities outside of the langevin integrator. the option flag to generate each time we call the integrator because in a hybrid workflow, there would effectively be no connection between subsequent langevin moves separated by many MC transformations. furthermore, MC moves that changes thermodynamic states would require regeneration. --- Examples/LJ_langevin.py | 14 ++++++++++++-- chiron/integrators.py | 8 +++++++- chiron/mcmc.py | 3 +++ chiron/utils.py | 39 +++++++++++++++++++++++++++++++++++---- 4 files changed, 57 insertions(+), 7 deletions(-) diff --git a/Examples/LJ_langevin.py b/Examples/LJ_langevin.py index 55d8e0b..42d8e02 100644 --- a/Examples/LJ_langevin.py +++ b/Examples/LJ_langevin.py @@ -8,7 +8,8 @@ from chiron.potential import LJPotential from openmm import unit -from chiron.utils import PRNG +from chiron.utils import PRNG, initialize_velocities + # initialize the LennardJones potential in chiron # @@ -20,6 +21,7 @@ lj_fluid.topology, sigma=sigma, epsilon=epsilon, cutoff=cutoff ) + from chiron.states import SamplerState, ThermodynamicState PRNG.set_seed(1234) @@ -30,11 +32,19 @@ box_vectors=lj_fluid.system.getDefaultPeriodicBoxVectors(), ) +velocities = initialize_velocities( + 300 * unit.kelvin, lj_fluid.topology, PRNG.get_random_key() +) + +print(velocities) + # define the thermodynamic state thermodynamic_state = ThermodynamicState( - potential=lj_potential, temperature=300 * unit.kelvin + potential=lj_potential, + temperature=300 * unit.kelvin, ) + from chiron.neighbors import NeighborListNsqrd, OrthogonalPeriodicSpace # define the neighbor list for an orthogonal periodic space diff --git a/chiron/integrators.py b/chiron/integrators.py index 8018647..6000df1 100644 --- a/chiron/integrators.py +++ b/chiron/integrators.py @@ -144,7 +144,13 @@ def run( # Initialize velocities if initialize_velocities: # we should probably move this to a separate function for unit testing purposes - v0 = sigma_v * random.normal(key, x0.shape) + # v0 = sigma_v * random.normal(key, x0.shape) + from .utils import initialize_velocities + + v0 = initialize_velocities( + temperature, potential.topology, key + ).value_in_unit_system(unit.md_unit_system) + else: if self.velocities is None: raise ValueError("Velocities must be set before running the integrator") diff --git a/chiron/mcmc.py b/chiron/mcmc.py index 29093b1..2506ac4 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -125,6 +125,7 @@ def update( sampler_state: SamplerState, thermodynamic_state: ThermodynamicState, nbr_list: Optional[PairsBase] = None, + initialize_velocities: bool = False, ): """ Run the integrator to perform molecular dynamics simulation. @@ -137,6 +138,8 @@ def update( The thermodynamic state to run the integrator on. nbr_list : PairsBase, optional The neighbor list to use for the simulation. + initialize_velocities : bool, optional + Flag indicating whether to initialize the velocities. """ assert isinstance( diff --git a/chiron/utils.py b/chiron/utils.py index a9a7da0..6b349e1 100644 --- a/chiron/utils.py +++ b/chiron/utils.py @@ -11,20 +11,21 @@ def __init__(self) -> None: """ A PRNG class that can be used to generate random numbers in JAX. The intended use case is to initialize new PRN streams in the `SamplerState` class. - + Example: -------- from chiron.utils import PRNG from chiron.states import SamplerState from openmmtools.testsystems import HarmonicOscillator - + ho = HarmonicOscillator() PRNG.set_seed(1234) sampler_state = [SamplerState(ho.positions, PRNG.get_random_key()) for _ in x0s] - + """ - + pass + @classmethod def set_seed(cls, seed: int) -> None: cls._seed = seed @@ -91,3 +92,33 @@ def get_list_of_mass(topology: Topology) -> unit.Quantity: for atom in topology.atoms(): mass.append(atom.element.mass.value_in_unit(unit.amu)) return mass * unit.amu + + +def initialize_velocities( + temperature: unit.Quantity, topology: Topology, key +) -> unit.Quantity: + """Initialize the velocities from the Maxwell-Boltzmann distribution at the given temperature. + + Parameters + ---------- + temperature : unit.Quantity + The temperature of the system. + topology : Topology + The topology of the system. + key : int + The PRNG key. + + """ + from openmm import unit + import jax.numpy as jnp + + mass = get_list_of_mass(topology) + kB = unit.BOLTZMANN_CONSTANT_kB * unit.AVOGADRO_CONSTANT_NA + + kbT_unitless = (kB * temperature).value_in_unit_system(unit.md_unit_system) + mass_unitless = jnp.array(mass.value_in_unit_system(unit.md_unit_system))[:, None] + sigma_v = jnp.sqrt(kbT_unitless / mass_unitless) + + v0 = sigma_v * random.normal(key, [len(mass), 3]) + + return v0 * unit.nanometer / unit.picosecond From 5e52df2f4d31632a6fc6402a198502d6f7fe496c Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Wed, 24 Jan 2024 21:59:10 -0500 Subject: [PATCH 06/43] Finished fixing up langevin integrator/move --- Examples/LJ_langevin.py | 6 ++--- chiron/integrators.py | 32 ++++++++++-------------- chiron/mcmc.py | 51 +++++++++++++++++++++------------------ chiron/states.py | 4 +++ chiron/tests/test_mcmc.py | 43 +++++++++++++++++++++++++++++++-- 5 files changed, 89 insertions(+), 47 deletions(-) diff --git a/Examples/LJ_langevin.py b/Examples/LJ_langevin.py index 42d8e02..0bf6846 100644 --- a/Examples/LJ_langevin.py +++ b/Examples/LJ_langevin.py @@ -36,7 +36,6 @@ 300 * unit.kelvin, lj_fluid.topology, PRNG.get_random_key() ) -print(velocities) # define the thermodynamic state thermodynamic_state = ThermodynamicState( @@ -74,7 +73,9 @@ from chiron.integrators import LangevinIntegrator # initialize the Langevin integrator -integrator = LangevinIntegrator(reporter=reporter, report_frequency=100) +integrator = LangevinIntegrator( + reporter=reporter, report_frequency=100, reinitialize_velocities=True +) print("init_energy: ", lj_potential.compute_energy(sampler_state.x0, nbr_list)) updated_sampler_state = integrator.run( @@ -83,7 +84,6 @@ n_steps=5000, nbr_list=nbr_list, progress_bar=True, - initialize_velocities=True, ) import h5py diff --git a/chiron/integrators.py b/chiron/integrators.py index 6000df1..922a5a1 100644 --- a/chiron/integrators.py +++ b/chiron/integrators.py @@ -26,6 +26,7 @@ def __init__( self, stepsize=1.0 * unit.femtoseconds, collision_rate=1.0 / unit.picoseconds, + reinitialize_velocities: bool = False, report_frequency: int = 100, reporter: Optional[LangevinDynamicsReporter] = None, save_traj_in_memory: bool = False, @@ -46,6 +47,9 @@ def __init__( save_traj_in_memory: bool Flag indicating whether to save the trajectory in memory. Default is False. NOTE: Only for debugging purposes. + reinitialize_velocities: bool + Whether to reinitialize the velocities each time the run function is called. + Default is False. """ from loguru import logger as log @@ -66,17 +70,7 @@ def __init__( self.velocities = None self.save_traj_in_memory = save_traj_in_memory self.traj = [] - - def set_velocities(self, vel: unit.Quantity) -> None: - """ - Set the initial velocities for the Langevin Integrator. - - Parameters - ---------- - vel : unit.Quantity - Velocities to be set for the integrator. - """ - self.velocities = vel + self.reinitialize_velocities = reinitialize_velocities def run( self, @@ -85,7 +79,6 @@ def run( n_steps: int = 5_000, nbr_list: Optional[PairsBase] = None, progress_bar=False, - initialize_velocities: bool = False, ): """ Run the integrator to perform Langevin dynamics molecular dynamics simulation. @@ -142,19 +135,20 @@ def run( b = jnp.sqrt(1 - jnp.exp(-2 * collision_rate_unitless * stepsize_unitless)) # Initialize velocities - if initialize_velocities: - # we should probably move this to a separate function for unit testing purposes + if self.reinitialize_velocities: # v0 = sigma_v * random.normal(key, x0.shape) from .utils import initialize_velocities - v0 = initialize_velocities( + sampler_state.velocities = initialize_velocities( temperature, potential.topology, key - ).value_in_unit_system(unit.md_unit_system) + ) else: - if self.velocities is None: + if sampler_state._velocities is None: raise ValueError("Velocities must be set before running the integrator") - v0 = self.velocities.value_in_unit_system(unit.md_unit_system) + + # extract the velocities from the sampler state + v0 = sampler_state.velocities x = x0 v = v0 @@ -201,7 +195,7 @@ def run( updated_sampler_state = copy.deepcopy(sampler_state) updated_sampler_state.x0 = x - updated_sampler_state._velocities = v + updated_sampler_state.velocities = v updated_sampler_state.current_PRNG_key = key return updated_sampler_state diff --git a/chiron/mcmc.py b/chiron/mcmc.py index 2506ac4..e63e66c 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -3,6 +3,7 @@ from typing import Tuple, List, Optional import jax.numpy as jnp from chiron.reporters import LangevinDynamicsReporter, _SimulationReporter +from .neighbors import PairsBase from abc import ABC, abstractmethod @@ -13,6 +14,7 @@ def __init__( nr_of_moves: int, reporter: Optional[_SimulationReporter] = None, report_frequency: Optional[int] = 100, + reinitialize_velocities: bool = False, ): """ Initialize a move within the molecular system. @@ -25,6 +27,9 @@ def __init__( Reporter object for saving the simulation data. Default is None. report_frequency : int, optional + Frequency of saving the simulation data in the reporter. + Default is 100. + """ self.nr_of_moves = nr_of_moves @@ -73,6 +78,7 @@ def __init__( self, stepsize=1.0 * unit.femtoseconds, collision_rate=1.0 / unit.picoseconds, + reinitialize_velocities: bool = False, reporter: Optional[LangevinDynamicsReporter] = None, report_frequency: int = 100, nr_of_steps=1_000, @@ -87,6 +93,9 @@ def __init__( Time step size for the integration. collision_rate : unit.Quantity Collision rate for the Langevin dynamics. + reinitialize_velocities : bool, optional + Whether to reinitialize the velocities each time the run function is called. + Default is False. reporter : LangevinDynamicsReporter, optional Reporter object for saving the simulation data. Default is None. @@ -115,6 +124,7 @@ def __init__( self.integrator = LangevinIntegrator( stepsize=self.stepsize, collision_rate=self.collision_rate, + reinitialize_velocities=reinitialize_velocities, report_frequency=report_frequency, reporter=reporter, save_traj_in_memory=save_traj_in_memory, @@ -125,7 +135,6 @@ def update( sampler_state: SamplerState, thermodynamic_state: ThermodynamicState, nbr_list: Optional[PairsBase] = None, - initialize_velocities: bool = False, ): """ Run the integrator to perform molecular dynamics simulation. @@ -138,8 +147,7 @@ def update( The thermodynamic state to run the integrator on. nbr_list : PairsBase, optional The neighbor list to use for the simulation. - initialize_velocities : bool, optional - Flag indicating whether to initialize the velocities. + """ assert isinstance( @@ -209,10 +217,12 @@ def update( Neighbor list for the system. + """ calculate_current_potential = True + for i in range(self.nr_of_moves): - self._step( + sampler_state, thermodynamic_state = self._step( sampler_state, thermodynamic_state, nbr_list, @@ -244,11 +254,9 @@ def _step( proposed_sampler_state, log_proposal_ratio, proposed_thermodynamic_state, - ) = self._propose(current_sampler_state, current_thermodynamic_state) - - # calculate the reduced potential for the proposed state - proposed_reduced_pot = proposed_thermodynamic_state.get_reduced_potential( - proposed_sampler_state + proposed_reduced_pot, + ) = self._propose( + current_sampler_state, current_thermodynamic_state, current_reduced_pot ) # accept or reject the proposed state @@ -258,6 +266,9 @@ def _step( log_proposal_ratio, method=self.method, ) + # a function that will update the statistics for the move + + self._update_statistics(decision) if decision: # save the reduced potential of the accepted state so @@ -267,15 +278,10 @@ def _step( # replace the current state with the proposed state # not sure this needs to be a separate function but for simplicity in outlining the code it is fine # or should this return the new sampler_state and thermodynamic_state? - self._replace_states( - current_sampler_state, - proposed_sampler_state, - current_thermodynamic_state, - proposed_thermodynamic_state, - ) - # a function that will update the statistics for the move - self._update_statistics(decision) + return proposed_sampler_state, proposed_thermodynamic_state + else: + return current_sampler_state, current_thermodynamic_state @abstractmethod def _propose(self, current_sampler_state, current_thermodynamic_state): @@ -319,7 +325,7 @@ def _accept_or_reject( current_reduced_pot, proposed_reduced_pot, log_proposal_ratio, - method=method, + method, ): """ Accept or reject the proposed state with a given methodology. @@ -419,7 +425,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: Optional[PairsBase] = None): """ Run the sampler for a specified number of iterations. @@ -436,7 +442,9 @@ 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) + self.sampler_state, self.thermodynamic_state = move.update( + self.sampler_state, self.thermodynamic_state, nbr_list + ) log.info("Finished running MCMC sampler") log.debug("Closing reporter") @@ -447,9 +455,6 @@ def run(self, n_iterations: int = 1): log.debug(f"Closed reporter {move.reporter.log_file_path}") -from .neighbors import PairsBase - - class MetropolizedMove(MCMove): """A base class for metropolized moves. diff --git a/chiron/states.py b/chiron/states.py index 1dcef10..b6d506d 100644 --- a/chiron/states.py +++ b/chiron/states.py @@ -118,6 +118,10 @@ def x0(self, x0: Union[jnp.array, unit.Quantity]) -> None: @velocities.setter def velocities(self, velocities: Union[jnp.array, unit.Quantity]) -> None: + if velocities.shape != self._x0.shape: + raise ValueError( + f"velocities must have the same shape as x0, got {velocities.shape} and {self._x0.shape} instead." + ) if isinstance(velocities, unit.Quantity): self._velocities = velocities else: diff --git a/chiron/tests/test_mcmc.py b/chiron/tests/test_mcmc.py index 21bf8b6..d746c6f 100644 --- a/chiron/tests/test_mcmc.py +++ b/chiron/tests/test_mcmc.py @@ -53,7 +53,10 @@ def test_sample_from_harmonic_osciallator(prep_temp_dir): reporter = LangevinDynamicsReporter() integrator = LangevinIntegrator( - stepsize=2 * unit.femtosecond, reporter=reporter, report_frequency=1 + stepsize=2 * unit.femtosecond, + reporter=reporter, + report_frequency=1, + reinitialize_velocities=True, ) integrator.run( @@ -122,7 +125,43 @@ def test_sample_from_harmonic_osciallator_with_MCMC_classes_and_LangevinDynamics BaseReporter.set_directory(prep_temp_dir) simulation_reporter = LangevinDynamicsReporter(1) - langevin_move = LangevinDynamicsMove(nr_of_steps=10, reporter=simulation_reporter) + + # this will fail because we have not initialized the velocity + with pytest.raises(ValueError): + langevin_move = LangevinDynamicsMove( + nr_of_steps=10, reinitialize_velocities=False, reporter=simulation_reporter + ) + move_set = MoveSchedule([("LangevinMove", langevin_move)]) + + # Initalize the sampler + sampler = MCMCSampler(move_set, sampler_state, thermodynamic_state) + + # Run the sampler with the thermodynamic state and sampler state and return the sampler state + sampler.run(n_iterations=2) # how many times to repeat + + # the following will reinitialize the velocities for each iteration + langevin_move = LangevinDynamicsMove( + nr_of_steps=10, reinitialize_velocities=True, reporter=simulation_reporter + ) + + move_set = MoveSchedule([("LangevinMove", langevin_move)]) + + # Initalize the sampler + sampler = MCMCSampler(move_set, sampler_state, thermodynamic_state) + + # Run the sampler with the thermodynamic state and sampler state and return the sampler state + sampler.run(n_iterations=2) # how many times to repeat + + # the following will use the initialize velocities function + from chiron.utils import initialize_velocities + + sampler_state.velocities = initialize_velocities( + thermodynamic_state.temperature, ho.topology, sampler_state._current_PRNG_key + ) + + langevin_move = LangevinDynamicsMove( + nr_of_steps=10, reinitialize_velocities=False, reporter=simulation_reporter + ) move_set = MoveSchedule([("LangevinMove", langevin_move)]) From 469071d0bd43b90a8db3aab35ba88e8e420a203a Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Thu, 25 Jan 2024 00:14:46 -0500 Subject: [PATCH 07/43] Implemented Metropolis displacement move in new scheme. --- Examples/LJ_mcmove.py | 22 ++++-- chiron/mcmc.py | 171 +++++++++++++++++++++++++++++++++++------- chiron/states.py | 9 ++- 3 files changed, 167 insertions(+), 35 deletions(-) diff --git a/Examples/LJ_mcmove.py b/Examples/LJ_mcmove.py index bc673f6..f28daaa 100644 --- a/Examples/LJ_mcmove.py +++ b/Examples/LJ_mcmove.py @@ -20,10 +20,15 @@ ) from chiron.states import SamplerState, ThermodynamicState +from chiron.utils import PRNG + +PRNG.set_seed(1234) # define the sampler state sampler_state = SamplerState( - x0=lj_fluid.positions, box_vectors=lj_fluid.system.getDefaultPeriodicBoxVectors() + x0=lj_fluid.positions, + current_PRNG_key=PRNG.get_random_key(), + box_vectors=lj_fluid.system.getDefaultPeriodicBoxVectors(), ) # define the thermodynamic state @@ -53,15 +58,18 @@ if os.path.isfile(filename): os.remove(filename) -reporter = _SimulationReporter("test_mc_lj.h5", lj_fluid.topology, 1) +reporter = _SimulationReporter("test_mc_lj.h5", 1) -from chiron.mcmc import MetropolisDisplacementMove +from chiron.mcmc import MetropolisDispMove -mc_move = MetropolisDisplacementMove( - seed=1234, - displacement_sigma=0.01 * unit.nanometer, +mc_move = MetropolisDispMove( + displacement_sigma=0.001 * unit.nanometer, nr_of_moves=1000, reporter=reporter, + report_frequency=1, ) -mc_move.run(sampler_state, thermodynamic_state, nbr_list, True) +mc_move.update(sampler_state, thermodynamic_state, nbr_list) + +stats = mc_move.statistics +print(stats["n_accepted"] / stats["n_proposed"]) diff --git a/chiron/mcmc.py b/chiron/mcmc.py index e63e66c..560e210 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -14,7 +14,6 @@ def __init__( nr_of_moves: int, reporter: Optional[_SimulationReporter] = None, report_frequency: Optional[int] = 100, - reinitialize_velocities: bool = False, ): """ Initialize a move within the molecular system. @@ -157,7 +156,7 @@ def update( thermodynamic_state, ThermodynamicState ), f"Thermodynamic state must be ThermodynamicState, not {type(thermodynamic_state)}" - updated_sampler_state = self.integrator.run( + sampler_state = self.integrator.run( thermodynamic_state=thermodynamic_state, sampler_state=sampler_state, n_steps=self.nr_of_moves, @@ -169,7 +168,7 @@ def update( self.integrator.traj = [] # The thermodynamic_state will not change for the langevin move - return updated_sampler_state, thermodynamic_state + # return updated_sampler_state, thermodynamic_state class MCMove(MCMCMove): @@ -195,8 +194,14 @@ def __init__( Methodology to use for accepting or rejecting the proposed state. Default is "metropolis". """ - super().__init__(nr_of_moves, reporter=reporter) - self.method = "metropolis" # I think we should pass a class/function instead of a string, like space. + super().__init__( + nr_of_moves, + reporter=reporter, + report_frequency=report_frequency, + ) + self.method = method # I think we should pass a class/function instead of a string, like space. + + self.reset_statistics() def update( self, @@ -233,12 +238,16 @@ def update( def _step( self, + current_sampler_state, + current_thermodynamic_state, + nbr_list, + calculate_current_potential=True, ): # if this is the first time we are calling this, # we will need to recalculate the reduced potential for the current state if calculate_current_potential: current_reduced_pot = current_thermodynamic_state.get_reduced_potential( - current_sampler_state + current_sampler_state, nbr_list ) # save the current_reduced_pot so we don't have to recalculate # it on the next iteration if the move is rejected @@ -256,18 +265,23 @@ def _step( proposed_thermodynamic_state, proposed_reduced_pot, ) = self._propose( - current_sampler_state, current_thermodynamic_state, current_reduced_pot + current_sampler_state, + current_thermodynamic_state, + current_reduced_pot, + nbr_list, ) # accept or reject the proposed state decision = self._accept_or_reject( - current_reduced_pot, - proposed_reduced_pot, log_proposal_ratio, + proposed_sampler_state.new_PRNG_key, method=self.method, ) # a function that will update the statistics for the move + if jnp.isnan(proposed_reduced_pot): + decision = False + self._update_statistics(decision) if decision: @@ -281,10 +295,38 @@ def _step( return proposed_sampler_state, proposed_thermodynamic_state else: + current_sampler_state._current_PRNG_key = ( + proposed_sampler_state._current_PRNG_key + ) return current_sampler_state, current_thermodynamic_state + def _update_statistics(self, decision): + """ + Update the statistics for the move. + """ + if decision: + self.n_accepted += 1 + self.n_proposed += 1 + + @property + def statistics(self): + """The acceptance statistics as a dictionary.""" + return dict(n_accepted=self.n_accepted, n_proposed=self.n_proposed) + + @statistics.setter + def statistics(self, value): + self.n_accepted = value["n_accepted"] + self.n_proposed = value["n_proposed"] + + def reset_statistics(self): + """Reset the acceptance statistics.""" + self.n_accepted = 0 + self.n_proposed = 0 + @abstractmethod - def _propose(self, current_sampler_state, current_thermodynamic_state): + def _propose( + self, current_sampler_state, current_thermodynamic_state, current_reduced_pot + ): """ Propose a new state and calculate the log proposal ratio. @@ -295,6 +337,9 @@ def _propose(self, current_sampler_state, current_thermodynamic_state): current_sampler_state : SamplerState, required Current sampler state. current_thermodynamic_state : ThermodynamicState, required + Current thermodynamic state. + current_reduced_pot : float, required + Current reduced potential. Returns ------- @@ -304,37 +349,111 @@ def _propose(self, current_sampler_state, current_thermodynamic_state): Log proposal ratio. proposed_thermodynamic_state : ThermodynamicState Proposed thermodynamic state. + proposed_reduced_pot : float + Proposed reduced potential. """ pass - def _replace_states( + def _accept_or_reject( self, - current_sampler_state, - proposed_sampler_state, - current_thermodynamic_state, - proposed_thermodynamic_state, + log_proposal_ratio, + key, + method, ): """ - Replace the current state with the proposed state. + Accept or reject the proposed state with a given methodology. """ - # define the code to copy the proposed state to the current state + # define the acceptance probability + if method == "metropolis": + import jax.random as jrandom - def _accept_or_reject( + compare_to = jrandom.uniform(key) + if log_proposal_ratio <= 0.0 or compare_to < jnp.exp(-log_proposal_ratio): + return True + else: + return False + + +class MetropolisDispMove(MCMove): + def __init__( self, + displacement_sigma=1.0 * unit.nanometer, + nr_of_moves: int = 100, + atom_subset: Optional[List[int]] = None, + report_frequency: int = 1, + reporter: Optional[LangevinDynamicsReporter] = None, + ): + """ + Initialize the Displacement Move class. + + Parameters + ---------- + displacement_sigma : float or unit.Quantity, optional + The standard deviation of the displacement for each move. Default is 1.0 nm. + nr_of_moves : int, optional + The number of moves to perform. Default is 100. + atom_subset : list of int, optional + A subset of atom indices to consider for the moves. Default is None. + reporter : SimulationReporter, optional + The reporter to write the data to. Default is None. + Returns + ------- + None + """ + super().__init__( + nr_of_moves=nr_of_moves, + reporter=reporter, + report_frequency=report_frequency, + method="metropolis", + ) + self.displacement_sigma = displacement_sigma + self.atom_subset = atom_subset + + def _propose( + self, + current_sampler_state, + current_thermodynamic_state, current_reduced_pot, - proposed_reduced_pot, - log_proposal_ratio, - method, + nbr_list, ): """ - Accept or reject the proposed state with a given methodology. + Implement the logic specific to displacement changes. """ - # define the acceptance probability + + key = current_sampler_state.new_PRNG_key + + nr_of_atoms = current_sampler_state.n_particles + unitless_displacement_sigma = self.displacement_sigma.value_in_unit_system( + unit.md_unit_system + ) + import jax.random as jrandom + + scaled_displacement_vector = ( + jrandom.normal(key, shape=(nr_of_atoms, 3)) * unitless_displacement_sigma + ) + updated_position = current_sampler_state.x0 + scaled_displacement_vector + import copy + + proposed_sampler_state = copy.deepcopy(current_sampler_state) + proposed_sampler_state.x0 = updated_position + + proposed_reduced_pot = current_thermodynamic_state.get_reduced_potential( + proposed_sampler_state, nbr_list + ) + delta_U = proposed_reduced_pot - current_reduced_pot + + # we do not change the thermodynamic state so we can return 'current_thermodnamic_state' + return ( + proposed_sampler_state, + delta_U, + current_thermodynamic_state, + proposed_reduced_pot, + ) class RotamerMove(MCMove): - def _step(self): + def _propose(self): """ Implement the logic specific to rotamer changes. """ @@ -342,7 +461,7 @@ def _step(self): class ProtonationStateMove(MCMove): - def _step(self): + def _propose(self): """ Implement the logic specific to protonation state changes. """ @@ -350,7 +469,7 @@ def _step(self): class TautomericStateMove(MCMove): - def _step(self): + def _propose(self): """ Implement the logic specific to tautomeric state changes. """ diff --git a/chiron/states.py b/chiron/states.py index b6d506d..74318bb 100644 --- a/chiron/states.py +++ b/chiron/states.py @@ -228,7 +228,7 @@ def __init__( from .utils import get_nr_of_particles self.nr_of_particles = get_nr_of_particles(self.potential.topology) - self._check_completness() + self._check_completeness() def check_variables(self) -> None: """ @@ -242,7 +242,7 @@ def check_variables(self) -> None: set_variables = [var for var in variables if getattr(self, var) is not None] return set_variables - def _check_completness(self): + def _check_completeness(self): # check which variables are set set_variables = self.check_variables() from loguru import logger as log @@ -299,6 +299,11 @@ def get_reduced_potential( ) / unit.AVOGADRO_CONSTANT_NA # log.debug(f"reduced potential: {reduced_potential}") if self.pressure is not 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 return self.beta * reduced_potential From a3ea44d5514c7de0033dba1d2fd82223c7d30a41 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Thu, 25 Jan 2024 17:30:15 -0500 Subject: [PATCH 08/43] Added refactored barostat move and ideal gas test case. --- Examples/Idealgas.py | 117 +++++++++++++++++++++++++++++++++ Examples/LJ_langevin.py | 1 + Examples/LJ_mcmove.py | 8 +-- chiron/mcmc.py | 140 +++++++++++++++++++++++++++++++++++++--- chiron/potential.py | 66 +++++++++++++++++++ chiron/reporters.py | 6 +- chiron/states.py | 7 ++ 7 files changed, 328 insertions(+), 17 deletions(-) create mode 100644 Examples/Idealgas.py diff --git a/Examples/Idealgas.py b/Examples/Idealgas.py new file mode 100644 index 0000000..923038d --- /dev/null +++ b/Examples/Idealgas.py @@ -0,0 +1,117 @@ +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 +from chiron.utils import PRNG +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, +) + +PRNG.set_seed(1234) + + +# define the sampler state +sampler_state = SamplerState( + x0=ideal_gas.positions, + current_PRNG_key=PRNG.get_random_key(), + 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 = "test_mc_ideal_gas.h5" +import os + +if os.path.isfile(filename): + os.remove(filename) +reporter = _SimulationReporter(filename, 1) + + +from chiron.mcmc import ( + MetropolisDisplacementMove, + MonteCarloBarostatMove, + MoveSchedule, + MCMCSampler, +) + +mc_disp_move = MetropolisDisplacementMove( + displacement_sigma=0.1 * unit.nanometer, + nr_of_moves=10, +) + +mc_barostat_move = MonteCarloBarostatMove( + volume_max_scale=0.2, + nr_of_moves=100, + reporter=reporter, +) +move_set = MoveSchedule( + [ + ("MetropolisDisplacementMove", mc_disp_move), + ("MonteCarloBarostatMove", 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) + +print(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 + +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 diff --git a/Examples/LJ_langevin.py b/Examples/LJ_langevin.py index 0bf6846..5b3cdb4 100644 --- a/Examples/LJ_langevin.py +++ b/Examples/LJ_langevin.py @@ -90,6 +90,7 @@ # read the data from the reporter with h5py.File("test_lj.h5", "r") as f: + print(f.keys()) energies = f["potential_energy"][:] steps = f["step"][:] diff --git a/Examples/LJ_mcmove.py b/Examples/LJ_mcmove.py index f28daaa..f838f8e 100644 --- a/Examples/LJ_mcmove.py +++ b/Examples/LJ_mcmove.py @@ -50,7 +50,7 @@ # build the neighbor list from the sampler state nbr_list.build_from_state(sampler_state) -from chiron.reporters import _SimulationReporter +from chiron.reporters import MCReporter # initialize a reporter to save the simulation data filename = "test_lj.h5" @@ -58,11 +58,11 @@ if os.path.isfile(filename): os.remove(filename) -reporter = _SimulationReporter("test_mc_lj.h5", 1) +reporter = MCReporter("test_mc_lj.h5", 1) -from chiron.mcmc import MetropolisDispMove +from chiron.mcmc import MetropolisDisplacementMove -mc_move = MetropolisDispMove( +mc_move = MetropolisDisplacementMove( displacement_sigma=0.001 * unit.nanometer, nr_of_moves=1000, reporter=reporter, diff --git a/chiron/mcmc.py b/chiron/mcmc.py index 560e210..20c86f3 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -375,7 +375,7 @@ def _accept_or_reject( return False -class MetropolisDispMove(MCMove): +class MetropolisDisplacementMove(MCMove): def __init__( self, displacement_sigma=1.0 * unit.nanometer, @@ -432,18 +432,140 @@ def _propose( scaled_displacement_vector = ( jrandom.normal(key, shape=(nr_of_atoms, 3)) * unitless_displacement_sigma ) - updated_position = current_sampler_state.x0 + scaled_displacement_vector import copy proposed_sampler_state = copy.deepcopy(current_sampler_state) - proposed_sampler_state.x0 = updated_position + + proposed_sampler_state.x0 = ( + current_sampler_state.x0 + scaled_displacement_vector + ) + + # after proposing a move we need to wrap particles and see if we need to rebuild + # the neighborlist + if nbr_list is not None: + proposed_sampler_state.x0 = nbr_list.space.wrap(proposed_sampler_state.x0) + + if nbr_list.check(proposed_sampler_state.x0): + nbr_list.build( + proposed_sampler_state.x0, proposed_sampler_state.box_vectors + ) proposed_reduced_pot = current_thermodynamic_state.get_reduced_potential( proposed_sampler_state, nbr_list ) + delta_U = proposed_reduced_pot - current_reduced_pot - # we do not change the thermodynamic state so we can return 'current_thermodnamic_state' + # we do not change the thermodynamic state; thus we can return 'current_thermodynamic_state' + return ( + proposed_sampler_state, + delta_U, + current_thermodynamic_state, + proposed_reduced_pot, + ) + + +class MonteCarloBarostatMove(MCMove): + def __init__( + self, + volume_max_scale=0.01, + nr_of_moves: int = 100, + atom_subset: Optional[List[int]] = None, + report_frequency: int = 1, + reporter: Optional[LangevinDynamicsReporter] = None, + ): + """ + Initialize the Monte Carlo Barostat Move class. + + Parameters + ---------- + displacement_sigma : float or unit.Quantity, optional + The standard deviation of the displacement for each move. Default is 1.0 nm. + nr_of_moves : int, optional + The number of moves to perform. Default is 100. + atom_subset : list of int, optional + A subset of atom indices to consider for the moves. Default is None. + reporter : SimulationReporter, optional + The reporter to write the data to. Default is None. + Returns + ------- + None + """ + super().__init__( + nr_of_moves=nr_of_moves, + reporter=reporter, + report_frequency=report_frequency, + method="metropolis", + ) + self.volume_max_scale = volume_max_scale + self.atom_subset = atom_subset + + def _propose( + self, + current_sampler_state, + current_thermodynamic_state, + current_reduced_pot, + nbr_list, + ): + """ + Implement the logic specific to displacement changes. + """ + from loguru import logger as log + + key = current_sampler_state.new_PRNG_key + + import jax.random as jrandom + + nr_of_atoms = current_sampler_state.n_particles + + initial_volume = ( + current_sampler_state.box_vectors[0][0] + * current_sampler_state.box_vectors[1][1] + * current_sampler_state.box_vectors[2][2] + ) + + # Calculate the maximum amount the volume can change by + delta_volume_max = self.volume_max_scale * initial_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(key, minval=-1, maxval=1) * delta_volume_max + + log.debug(f"Delta volume is {delta_volume}.") + + # calculate the new volume + proposed_volume = initial_volume + delta_volume + log.debug(f"New volume is {proposed_volume}.") + + # calculate the length scale factor for particle positions and box vectors + length_scaling_factor = jnp.power(proposed_volume / initial_volume, 1.0 / 3.0) + + log.debug(f"Length scaling factor is {length_scaling_factor}.") + import copy + + proposed_sampler_state = copy.deepcopy(current_sampler_state) + proposed_sampler_state.x0 = current_sampler_state.x0 * length_scaling_factor + proposed_sampler_state.box_vectors = ( + current_sampler_state.box_vectors * length_scaling_factor + ) + + if nbr_list is not None: + # after scaling the box vectors we should rebuild the neighborlist + nbr_list.build( + proposed_sampler_state.x0, proposed_sampler_state.box_vectors + ) + + proposed_reduced_pot = current_thermodynamic_state.get_reduced_potential( + proposed_sampler_state, nbr_list + ) + + delta_U = ( + proposed_reduced_pot + - current_reduced_pot + - nr_of_atoms * jnp.log(proposed_volume / initial_volume) + ) + + # we do not change the thermodynamic state so we can return 'current_thermodynamic_state' return ( proposed_sampler_state, delta_U, @@ -561,9 +683,7 @@ def run(self, n_iterations: int = 1, nbr_list: Optional[PairsBase] = None): log.info(f"Iteration {iteration + 1}/{n_iterations}") for move_name, move in self.move.move_schedule: log.debug(f"Performing: {move_name}") - self.sampler_state, self.thermodynamic_state = move.update( - self.sampler_state, self.thermodynamic_state, nbr_list - ) + move.update(self.sampler_state, self.thermodynamic_state, nbr_list) log.info("Finished running MCMC sampler") log.debug("Closing reporter") @@ -692,7 +812,7 @@ def apply( delta_energy <= 0.0 or compare_to < jnp.exp(-delta_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}." ) @@ -739,7 +859,7 @@ def _propose_positions(self, positions: jnp.array): ) -class MetropolisDisplacementMove(MetropolizedMove): +class MetropolisDispMove(MetropolizedMove): """A metropolized move that randomly displace a subset of atoms. Parameters @@ -771,7 +891,7 @@ def __init__( displacement_sigma=1.0 * unit.nanometer, nr_of_moves: int = 100, atom_subset: Optional[List[int]] = None, - reporter: Optional[LangevinDynamicsReporter] = None, + reporter: Optional[MCReporter] = None, ): """ Initialize the MCMC class. diff --git a/chiron/potential.py b/chiron/potential.py index 6d9b415..8362ef1 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/reporters.py b/chiron/reporters.py index 27457a6..156e86a 100644 --- a/chiron/reporters.py +++ b/chiron/reporters.py @@ -367,7 +367,7 @@ def _write_to_trajectory(self, positions: np.ndarray) -> None: file_handler=self._write_xtc_file_handle, positions=positions, iteration=self.buffer.get("step"), - box_vecotrs=self.buffer.get("box_vectors"), + box_vectors=self.buffer.get("box_vectors"), ) def read_from_trajectory(self) -> np.ndarray: @@ -409,7 +409,7 @@ def _write_to_xtc( file_handler: md.formats.XTCTrajectoryFile, positions: np.ndarray, iteration: np.ndarray, - box_vecotrs: Optional[np.ndarray] = None, + box_vectors: Optional[np.ndarray] = None, ): """ Write position data to an XTC file. @@ -428,5 +428,5 @@ def _write_to_xtc( file_handler.write( positions, time=iteration, - box=box_vecotrs, + box=box_vectors, ) diff --git a/chiron/states.py b/chiron/states.py index 74318bb..1d9855a 100644 --- a/chiron/states.py +++ b/chiron/states.py @@ -116,6 +116,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) + @velocities.setter def velocities(self, velocities: Union[jnp.array, unit.Quantity]) -> None: if velocities.shape != self._x0.shape: From 4ad677e9d557b9186362547e57065bae59733eb9 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Thu, 25 Jan 2024 17:40:55 -0500 Subject: [PATCH 09/43] Abstractmethod for _reporting wrapper (this will make it easier for each move to have custom/appropriate logging for what is being changed). --- Examples/Idealgas.py | 9 ++++++--- chiron/mcmc.py | 27 ++++++++++++++++++++++++++- 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/Examples/Idealgas.py b/Examples/Idealgas.py index 923038d..5d96a8c 100644 --- a/Examples/Idealgas.py +++ b/Examples/Idealgas.py @@ -88,9 +88,9 @@ import h5py -# with h5py.File(filename, "r") as f: -# volume = f["volume"][:] -# steps = f["step"][:] +with h5py.File(filename, "r") as f: + volume = f["volume"][:] + steps = f["step"][:] # get expectations ideal_volume = ideal_gas.get_volume_expectation(thermodynamic_state) @@ -102,6 +102,9 @@ volume_mean = jnp.mean(jnp.array(volume)) * unit.nanometer**3 volume_std = jnp.std(jnp.array(volume)) * unit.nanometer**3 + +print(volume_mean, volume_std) + ideal_density = mass * n_particles / unit.AVOGADRO_CONSTANT_NA / ideal_volume measured_density = mass * n_particles / unit.AVOGADRO_CONSTANT_NA / volume_mean diff --git a/chiron/mcmc.py b/chiron/mcmc.py index 20c86f3..fcdd542 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -2,7 +2,7 @@ from openmm import unit from typing import Tuple, List, Optional import jax.numpy as jnp -from chiron.reporters import LangevinDynamicsReporter, _SimulationReporter +from chiron.reporters import LangevinDynamicsReporter, _SimulationReporter, MCReporter from .neighbors import PairsBase from abc import ABC, abstractmethod @@ -235,6 +235,20 @@ def update( ) # after the first step, we don't need to recalculate the current potential, it will be stored calculate_current_potential = False + if hasattr(self, "reporter"): + if self.reporter is not None: + if i % self.report_frequency == 0: + self._report(i, sampler_state, thermodynamic_state, nbr_list) + + @abstractmethod + def _report(self, step, sampler_state, thermodynamic_state, nbr_list): + """ + Report the current state of the MC move. + + Since different moves will be modifying different quantities, + this needs to be defined for each move. + """ + pass def _step( self, @@ -500,6 +514,17 @@ def __init__( self.volume_max_scale = volume_max_scale self.atom_subset = atom_subset + def _report(self, step, sampler_state, thermodynamic_state, nbr_list): + potential = thermodynamic_state.get_reduced_potential(sampler_state, nbr_list) + volume = ( + sampler_state.box_vectors[0][0] + * sampler_state.box_vectors[1][1] + * sampler_state.box_vectors[2][2] + ) + self.reporter.report( + {"step": step, "potential_energy": potential, "volume": volume} + ) + def _propose( self, current_sampler_state, From 9c4d47bc680534d417a4964562232e1273ed7d02 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Fri, 26 Jan 2024 14:36:22 -0500 Subject: [PATCH 10/43] Added _reporter to the displacement moves and the ability to only move a subset of atoms: tests need to be updated/implemented still. --- chiron/mcmc.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/chiron/mcmc.py b/chiron/mcmc.py index fcdd542..54c20fa 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -314,6 +314,7 @@ def _step( ) return current_sampler_state, current_thermodynamic_state + def _updated_ def _update_statistics(self, decision): """ Update the statistics for the move. @@ -424,6 +425,10 @@ def __init__( self.displacement_sigma = displacement_sigma self.atom_subset = atom_subset + def _report(self, step, sampler_state, thermodynamic_state, nbr_list): + potential = thermodynamic_state.get_reduced_potential(sampler_state, nbr_list) + self.reporter.report({"step": step, "potential_energy": potential}) + def _propose( self, current_sampler_state, @@ -450,9 +455,14 @@ def _propose( proposed_sampler_state = copy.deepcopy(current_sampler_state) - proposed_sampler_state.x0 = ( - current_sampler_state.x0 + scaled_displacement_vector - ) + if self.atom_subset is not None: + proposed_sampler_state.x0 = ( + current_sampler_state.x0 + scaled_displacement_vector * self.atom_subset + ) + else: + proposed_sampler_state.x0 = ( + current_sampler_state.x0 + scaled_displacement_vector + ) # after proposing a move we need to wrap particles and see if we need to rebuild # the neighborlist From b6d88cc1cb7d3926b6b3a7e59ef405363674383b Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Mon, 29 Jan 2024 11:41:51 -0800 Subject: [PATCH 11/43] Added in stepsize updater to allow move_parameters to be updated on the fly. --- chiron/mcmc.py | 92 +++++++++++++++++++++++++++++++++++++++++++------- 1 file changed, 80 insertions(+), 12 deletions(-) diff --git a/chiron/mcmc.py b/chiron/mcmc.py index 54c20fa..f069284 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -177,6 +177,8 @@ def __init__( nr_of_moves: int, reporter: Optional[_SimulationReporter], report_frequency: int = 1, + update_stepsize: bool = False, + update_stepsize_frequency: int = 100, method: str = "metropolis", ) -> None: """ @@ -190,6 +192,11 @@ def __init__( Reporter object for saving the simulation step data. report_frequency Frequency of saving the simulation data. + update_stepsize + Whether to update the "stepsize" of the move. Stepsize is a generic term for the key move parameters. + For example, for a simple displacement move this would be the displacement_sigma. + update_stepsize_frequency + Frequency of updating the stepsize of the move. method Methodology to use for accepting or rejecting the proposed state. Default is "metropolis". @@ -202,6 +209,8 @@ def __init__( self.method = method # I think we should pass a class/function instead of a string, like space. self.reset_statistics() + self.update_stepsize = update_stepsize + self.update_stepsize_frequency = update_stepsize_frequency def update( self, @@ -222,7 +231,6 @@ def update( Neighbor list for the system. - """ calculate_current_potential = True @@ -240,6 +248,10 @@ def update( if i % self.report_frequency == 0: self._report(i, sampler_state, thermodynamic_state, nbr_list) + if self.update_stepsize: + if i % self.update_stepsize_frequency == 0: + self._update_stepsize() + @abstractmethod def _report(self, step, sampler_state, thermodynamic_state, nbr_list): """ @@ -250,6 +262,17 @@ def _report(self, step, sampler_state, thermodynamic_state, nbr_list): """ pass + @abstractmethod + def _update_stepsize(self): + """ + Update the "stepsize" for a move to reach a target acceptance probability range. + This will be specific to the type of move, e.g., a displacement_sigma for a displacement move + or a maximum volume change factor for a Monte Carlo barostat move. + + Since different moves will be modifying different quantities, this needs to be defined for each move. + """ + pass + def _step( self, current_sampler_state, @@ -259,6 +282,7 @@ def _step( ): # if this is the first time we are calling this, # we will need to recalculate the reduced potential for the current state + # this is toggled by the calculate_current_potential flag if calculate_current_potential: current_reduced_pot = current_thermodynamic_state.get_reduced_potential( current_sampler_state, nbr_list @@ -275,9 +299,9 @@ def _step( # for systems that e.g., make changes to particle identity. ( proposed_sampler_state, - log_proposal_ratio, proposed_thermodynamic_state, proposed_reduced_pot, + log_proposal_ratio, ) = self._propose( current_sampler_state, current_thermodynamic_state, @@ -309,12 +333,14 @@ def _step( return proposed_sampler_state, proposed_thermodynamic_state else: + # if we reject the move, we need to update the current_PRNG key to ensure that + # we are using a different random number for the next iteration + # this is needed because the _step function returns a SamplerState instead of updating it in place current_sampler_state._current_PRNG_key = ( proposed_sampler_state._current_PRNG_key ) return current_sampler_state, current_thermodynamic_state - def _updated_ def _update_statistics(self, decision): """ Update the statistics for the move. @@ -345,7 +371,10 @@ def _propose( """ Propose a new state and calculate the log proposal ratio. - This will need to be defined for each move + This will accept the relevant quantities for the current state, returning the proposed state quantities + and the log proposal ratio. + + This will need to be defined for each new move. Parameters ---------- @@ -360,12 +389,12 @@ def _propose( ------- proposed_sampler_state : SamplerState Proposed sampler state. - log_proposal_ratio : float - Log proposal ratio. proposed_thermodynamic_state : ThermodynamicState Proposed thermodynamic state. proposed_reduced_pot : float Proposed reduced potential. + log_proposal_ratio : float + Log proposal ratio. """ pass @@ -398,6 +427,8 @@ def __init__( atom_subset: Optional[List[int]] = None, report_frequency: int = 1, reporter: Optional[LangevinDynamicsReporter] = None, + update_stepsize: bool = True, + update_stepsize_frequency: int = 100, ): """ Initialize the Displacement Move class. @@ -412,6 +443,10 @@ def __init__( A subset of atom indices to consider for the moves. Default is None. reporter : SimulationReporter, optional The reporter to write the data to. Default is None. + update_stepsize : bool, optional + Whether to update the stepsize of the move. Default is True. + update_stepsize_frequency : int, optional + Frequency of updating the stepsize of the move. Default is 100. Returns ------- None @@ -420,6 +455,8 @@ def __init__( nr_of_moves=nr_of_moves, reporter=reporter, report_frequency=report_frequency, + update_stepsize=update_stepsize, + update_stepsize_frequency=update_stepsize_frequency, method="metropolis", ) self.displacement_sigma = displacement_sigma @@ -429,6 +466,16 @@ def _report(self, step, sampler_state, thermodynamic_state, nbr_list): potential = thermodynamic_state.get_reduced_potential(sampler_state, nbr_list) self.reporter.report({"step": step, "potential_energy": potential}) + def _update_stepsize(self): + """ + Update the displacement_sigma to reach a target acceptance probability of 0.5. + """ + + if self.n_accepted / self.n_proposed > 0.5: + self.displacement_sigma *= 1.1 + else: + self.displacement_sigma /= 1.1 + def _propose( self, current_sampler_state, @@ -443,6 +490,7 @@ def _propose( key = current_sampler_state.new_PRNG_key nr_of_atoms = current_sampler_state.n_particles + unitless_displacement_sigma = self.displacement_sigma.value_in_unit_system( unit.md_unit_system ) @@ -478,14 +526,15 @@ def _propose( proposed_sampler_state, nbr_list ) - delta_U = proposed_reduced_pot - current_reduced_pot + log_proposal_ratio = proposed_reduced_pot - current_reduced_pot - # we do not change the thermodynamic state; thus we can return 'current_thermodynamic_state' + # since do not change the thermodynamic state we can return + # 'current_thermodynamic_state' rather than making a copy return ( proposed_sampler_state, - delta_U, current_thermodynamic_state, proposed_reduced_pot, + log_proposal_ratio, ) @@ -497,6 +546,8 @@ def __init__( atom_subset: Optional[List[int]] = None, report_frequency: int = 1, reporter: Optional[LangevinDynamicsReporter] = None, + update_stepsize: bool = True, + update_stepsize_frequency: int = 100, ): """ Initialize the Monte Carlo Barostat Move class. @@ -511,6 +562,10 @@ def __init__( A subset of atom indices to consider for the moves. Default is None. reporter : SimulationReporter, optional The reporter to write the data to. Default is None. + update_stepsize : bool, optional + Whether to update the stepsize of the move. Default is True. + update_stepsize_frequency : int, optional + Frequency of updating the stepsize of the move. Default is 100. Returns ------- None @@ -519,10 +574,11 @@ def __init__( nr_of_moves=nr_of_moves, reporter=reporter, report_frequency=report_frequency, + update_stepsize=update_stepsize, + update_stepsize_frequency=update_stepsize_frequency, method="metropolis", ) self.volume_max_scale = volume_max_scale - self.atom_subset = atom_subset def _report(self, step, sampler_state, thermodynamic_state, nbr_list): potential = thermodynamic_state.get_reduced_potential(sampler_state, nbr_list) @@ -535,6 +591,17 @@ def _report(self, step, sampler_state, thermodynamic_state, nbr_list): {"step": step, "potential_energy": potential, "volume": volume} ) + def _update_stepsize(self): + """ + Update the volume_max_scale parameter to ensure our acceptance probability is within the range of 0.25 to 0.75. + The maximum volume_max_scale will be capped at 0.3. + """ + acceptance_ratio = self.n_accepted / self.n_proposed + if acceptance_ratio < 0.25: + self.volume_max_scale /= 1.1 + elif acceptance_ratio > 0.75: + self.volume_max_scale = min(self.volume_max_scale * 1.1, 0.3) + def _propose( self, current_sampler_state, @@ -594,7 +661,8 @@ def _propose( proposed_sampler_state, nbr_list ) - delta_U = ( + # + log_proposal_ratio = ( proposed_reduced_pot - current_reduced_pot - nr_of_atoms * jnp.log(proposed_volume / initial_volume) @@ -603,9 +671,9 @@ def _propose( # we do not change the thermodynamic state so we can return 'current_thermodynamic_state' return ( proposed_sampler_state, - delta_U, current_thermodynamic_state, proposed_reduced_pot, + log_proposal_ratio, ) From c5bce7acf62a66c405a43accf62e8ef2d77ac493 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Tue, 30 Jan 2024 00:59:24 -0800 Subject: [PATCH 12/43] Continued reworking to ensure expected behavior of loggers and stepsize adjustments. --- Examples/Idealgas.py | 9 +- Examples/LJ_MCMC.py | 149 +++++++++++++ Examples/LJ_mcmove.py | 39 +++- chiron/mcmc.py | 483 ++++++++++++++---------------------------- chiron/states.py | 3 + 5 files changed, 352 insertions(+), 331 deletions(-) create mode 100644 Examples/LJ_MCMC.py diff --git a/Examples/Idealgas.py b/Examples/Idealgas.py index 5d96a8c..f696474 100644 --- a/Examples/Idealgas.py +++ b/Examples/Idealgas.py @@ -66,25 +66,22 @@ MCMCSampler, ) -mc_disp_move = MetropolisDisplacementMove( - displacement_sigma=0.1 * unit.nanometer, - nr_of_moves=10, -) mc_barostat_move = MonteCarloBarostatMove( volume_max_scale=0.2, nr_of_moves=100, reporter=reporter, + update_stepsize=True, + update_stepsize_frequency=100, ) move_set = MoveSchedule( [ - ("MetropolisDisplacementMove", mc_disp_move), ("MonteCarloBarostatMove", 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 +sampler.run(n_iterations=50, nbr_list=nbr_list) # how many times to repeat import h5py diff --git a/Examples/LJ_MCMC.py b/Examples/LJ_MCMC.py new file mode 100644 index 0000000..8c92d1f --- /dev/null +++ b/Examples/LJ_MCMC.py @@ -0,0 +1,149 @@ +from openmm import unit +from openmm import app + + +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 + +positions = jnp.load("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], + ] +) + +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( + x0=positions, + current_PRNG_key=PRNG.get_random_key(), + box_vectors=box_vectors * unit.nanometer, +) + + +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) + +# CRI: minimizer is not working correctly on my mac +# from chiron.minimze import minimize_energy +# +# results = minimize_energy( +# sampler_state.x0, lj_potential.compute_energy, nbr_list, maxiter=100 +# ) +# +# min_x = results.params +# +# sampler_state.x0 = min_x + +from chiron.reporters import MCReporter, LangevinDynamicsReporter + +# initialize a reporter to save the simulation data +filename_barostat = "test_mc_lj_barostat.h5" +filename_displacement = "test_mc_lj_disp.h5" +filename_langevin = "test_mc_lj_langevin.h5" + +import os + +if os.path.isfile(filename_barostat): + os.remove(filename_barostat) +reporter_barostat = MCReporter(filename_barostat, 1) + +if os.path.isfile(filename_displacement): + os.remove(filename_displacement) +reporter_displacement = MCReporter(filename_displacement, 10) + +if os.path.isfile(filename_langevin): + os.remove(filename_langevin) +reporter_langevin = LangevinDynamicsReporter(filename_langevin, 10) + +from chiron.mcmc import ( + MetropolisDisplacementMove, + MonteCarloBarostatMove, + LangevinDynamicsMove, + MoveSchedule, + MCMCSampler, +) + +mc_displacement_move = MetropolisDisplacementMove( + displacement_sigma=0.001 * unit.nanometer, + nr_of_moves=100, + reporter=reporter_displacement, + report_frequency=10, + update_stepsize=True, + update_stepsize_frequency=100, +) + +mc_barostat_move = MonteCarloBarostatMove( + volume_max_scale=0.1, + nr_of_moves=10, + reporter=reporter_barostat, + report_frequency=1, + update_stepsize=True, + update_stepsize_frequency=50, +) + +langevin_dynamics_move = LangevinDynamicsMove( + stepsize=1.0 * unit.femtoseconds, + collision_rate=1.0 / unit.picoseconds, + nr_of_steps=100, + reporter=reporter_langevin, + report_frequency=10, + reinitialize_velocities=True, +) + + +move_set = MoveSchedule( + [ + # ("LangevinDynamicsMove", langevin_dynamics_move), + ("MetropolisDisplacementMove", mc_displacement_move), + ("MonteCarloBarostatMove", mc_barostat_move), + ] +) + +sampler = MCMCSampler(move_set, sampler_state, thermodynamic_state) +sampler.run(n_iterations=1000, nbr_list=nbr_list) # how many times to repeat diff --git a/Examples/LJ_mcmove.py b/Examples/LJ_mcmove.py index f838f8e..69c2859 100644 --- a/Examples/LJ_mcmove.py +++ b/Examples/LJ_mcmove.py @@ -53,23 +53,54 @@ from chiron.reporters import MCReporter # initialize a reporter to save the simulation data -filename = "test_lj.h5" +filename = "test_mc_lj.h5" import os if os.path.isfile(filename): os.remove(filename) -reporter = MCReporter("test_mc_lj.h5", 1) +reporter = MCReporter(filename, 1) from chiron.mcmc import MetropolisDisplacementMove mc_move = MetropolisDisplacementMove( - displacement_sigma=0.001 * unit.nanometer, - nr_of_moves=1000, + displacement_sigma=0.01 * unit.nanometer, + nr_of_moves=5000, reporter=reporter, report_frequency=1, + update_stepsize=True, + update_stepsize_frequency=100, ) mc_move.update(sampler_state, thermodynamic_state, nbr_list) stats = mc_move.statistics print(stats["n_accepted"] / stats["n_proposed"]) + +import h5py + +with h5py.File(filename, "r") as f: + acceptance_probability = f["acceptance_probability"][:] + displacement_sigma = f["displacement_sigma"][:] + potential_energy = f["potential_energy"][:] + step = f["step"][:] + +# plot the energy +import matplotlib.pyplot as plt + +plt.subplot(3, 1, 1) + +plt.plot(step, displacement_sigma) +plt.ylabel("displacement_sigma (nm)") + +plt.subplot(3, 1, 2) + +plt.plot(step, acceptance_probability) +plt.ylabel("acceptance_probability") + + +plt.subplot(3, 1, 3) + +plt.plot(step, potential_energy) +plt.xlabel("Step") +plt.ylabel("potential_energy (kj/mol)") +plt.show() diff --git a/chiron/mcmc.py b/chiron/mcmc.py index f069284..dea694f 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -34,6 +34,10 @@ def __init__( self.nr_of_moves = nr_of_moves self.reporter = reporter self.report_frequency = report_frequency + + # we need to keep track of which iteration we are on + self.iteration = 0 + from loguru import logger as log if self.reporter is not None: @@ -167,8 +171,10 @@ def update( self.traj.append(self.integrator.traj) self.integrator.traj = [] + self.iteration += 1 + # The thermodynamic_state will not change for the langevin move - # return updated_sampler_state, thermodynamic_state + return sampler_state, thermodynamic_state class MCMove(MCMCMove): @@ -230,7 +236,12 @@ def update( nbr_list : PairBase, optional Neighbor list for the system. - + Returns + ------- + sampler_state : SamplerState + The updated sampler state. + thermodynamic_state : ThermodynamicState + The updated thermodynamic state. """ calculate_current_potential = True @@ -243,22 +254,65 @@ def update( ) # after the first step, we don't need to recalculate the current potential, it will be stored calculate_current_potential = False + + elapsed_step = i + self.iteration * self.nr_of_moves if hasattr(self, "reporter"): if self.reporter is not None: - if i % self.report_frequency == 0: - self._report(i, sampler_state, thermodynamic_state, nbr_list) - + # I think it makes sense to i + self.nr_of_moves*self.iteration as our current "step" + # otherwise, instances where self.report_frequency > self.nr_of_moves would only report on the + # first step which might actually be more frequent than we specify + + if elapsed_step % self.report_frequency == 0: + self._report( + i, + self.iteration, + self.n_accepted / self.n_proposed, + sampler_state, + thermodynamic_state, + nbr_list, + ) if self.update_stepsize: - if i % self.update_stepsize_frequency == 0: + # if we only used i, we might never actually update the parameters if we have a move that is called infrequently + if ( + elapsed_step % self.update_stepsize_frequency == 0 + and elapsed_step > 0 + ): self._update_stepsize() + self.iteration += 1 + + return sampler_state, thermodynamic_state + @abstractmethod - def _report(self, step, sampler_state, thermodynamic_state, nbr_list): + def _report( + self, + step: int, + iteration: int, + acceptance_probability: float, + sampler_state: SamplerState, + thermodynamic_state: ThermodynamicState, + nbr_list: Optional[PairsBase] = None, + ): """ Report the current state of the MC move. Since different moves will be modifying different quantities, this needs to be defined for each move. + + Parameters + ---------- + step : int + The current step of the simulation move. + iteration : int + The current iteration of the move sequence (i.e., how many times has this been called thus far). + acceptance_probability : float + The acceptance probability of the move. + sampler_state : SamplerState + The sampler state of the system. + thermodynamic_state : ThermodynamicState + The thermodynamic state of the system. + nbr_list : Optional[PairBase]=None + The neighbor list or pair list for evaluating interactions in the system, default None """ pass @@ -339,6 +393,12 @@ def _step( current_sampler_state._current_PRNG_key = ( proposed_sampler_state._current_PRNG_key ) + if nbr_list is not None: + if nbr_list.check(current_sampler_state.x0): + nbr_list.build( + current_sampler_state.x0, current_sampler_state.box_vectors + ) + return current_sampler_state, current_thermodynamic_state def _update_statistics(self, decision): @@ -413,7 +473,7 @@ def _accept_or_reject( import jax.random as jrandom compare_to = jrandom.uniform(key) - if log_proposal_ratio <= 0.0 or compare_to < jnp.exp(-log_proposal_ratio): + if -log_proposal_ratio <= 0.0 or compare_to < jnp.exp(log_proposal_ratio): return True else: return False @@ -462,18 +522,57 @@ def __init__( self.displacement_sigma = displacement_sigma self.atom_subset = atom_subset - def _report(self, step, sampler_state, thermodynamic_state, nbr_list): - potential = thermodynamic_state.get_reduced_potential(sampler_state, nbr_list) - self.reporter.report({"step": step, "potential_energy": potential}) + def _report( + self, + step, + iteration, + acceptance_probability, + sampler_state, + thermodynamic_state, + nbr_list, + ): + """ + Report the current state of the MC displacement move. + + Parameters + ---------- + step : int + The current step of the simulation move. + iteration : int + The current iteration of the move sequence (i.e., how many times has this been called thus far). + acceptance_probability : float + The acceptance probability of the move. + sampler_state : SamplerState + The sampler state of the system. + thermodynamic_state : ThermodynamicState + The thermodynamic state of the system. + nbr_list : Optional[PairBase]=None + The neighbor list or pair list for evaluating interactions in the system, default None + + """ + potential = thermodynamic_state.potential.compute_energy( + sampler_state.x0, nbr_list + ) + self.reporter.report( + { + "step": step, + "iteration": iteration, + "potential_energy": potential, + "displacement_sigma": self.displacement_sigma.value_in_unit_system( + unit.md_unit_system + ), + "acceptance_probability": acceptance_probability, + } + ) def _update_stepsize(self): """ Update the displacement_sigma to reach a target acceptance probability of 0.5. """ - - if self.n_accepted / self.n_proposed > 0.5: + acceptance_ratio = self.n_accepted / self.n_proposed + if acceptance_ratio > 0.6: self.displacement_sigma *= 1.1 - else: + elif acceptance_ratio < 0.4: self.displacement_sigma /= 1.1 def _propose( @@ -505,11 +604,12 @@ def _propose( if self.atom_subset is not None: proposed_sampler_state.x0 = ( - current_sampler_state.x0 + scaled_displacement_vector * self.atom_subset + proposed_sampler_state.x0 + + scaled_displacement_vector * self.atom_subset ) else: proposed_sampler_state.x0 = ( - current_sampler_state.x0 + scaled_displacement_vector + proposed_sampler_state.x0 + scaled_displacement_vector ) # after proposing a move we need to wrap particles and see if we need to rebuild @@ -526,7 +626,7 @@ def _propose( proposed_sampler_state, nbr_list ) - log_proposal_ratio = proposed_reduced_pot - current_reduced_pot + log_proposal_ratio = -proposed_reduced_pot + current_reduced_pot # since do not change the thermodynamic state we can return # 'current_thermodynamic_state' rather than making a copy @@ -580,15 +680,49 @@ def __init__( ) self.volume_max_scale = volume_max_scale - def _report(self, step, sampler_state, thermodynamic_state, nbr_list): - potential = thermodynamic_state.get_reduced_potential(sampler_state, nbr_list) + def _report( + self, + step, + iteration, + acceptance_probability, + sampler_state, + thermodynamic_state, + nbr_list, + ): + """ + + Parameters + ---------- + step : int + The current step of the simulation move. + iteration : int + The current iteration of the move sequence (i.e., how many times has this been called thus far). + acceptance_probability : float + The acceptance probability of the move. + sampler_state : SamplerState + The sampler state of the system. + thermodynamic_state : ThermodynamicState + The thermodynamic state of the system. + nbr_list : Optional[PairBase]=None + The neighbor list or pair list for evaluating interactions in the system, default None + """ + potential = thermodynamic_state.potential.compute_energy( + sampler_state.x0, nbr_list + ) volume = ( sampler_state.box_vectors[0][0] * sampler_state.box_vectors[1][1] * sampler_state.box_vectors[2][2] ) self.reporter.report( - {"step": step, "potential_energy": potential, "volume": volume} + { + "step": step, + "iteration": iteration, + "potential_energy": potential, + "volume": volume, + "max_volume_scale": self.volume_max_scale, + "acceptance_probability": acceptance_probability, + } ) def _update_stepsize(self): @@ -632,21 +766,17 @@ def _propose( # 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(key, minval=-1, maxval=1) * delta_volume_max - - log.debug(f"Delta volume is {delta_volume}.") - # calculate the new volume proposed_volume = initial_volume + delta_volume - log.debug(f"New volume is {proposed_volume}.") # calculate the length scale factor for particle positions and box vectors length_scaling_factor = jnp.power(proposed_volume / initial_volume, 1.0 / 3.0) - log.debug(f"Length scaling factor is {length_scaling_factor}.") import copy proposed_sampler_state = copy.deepcopy(current_sampler_state) proposed_sampler_state.x0 = current_sampler_state.x0 * length_scaling_factor + proposed_sampler_state.box_vectors = ( current_sampler_state.box_vectors * length_scaling_factor ) @@ -661,11 +791,11 @@ def _propose( proposed_sampler_state, nbr_list ) - # + # χ = exp ⎡−β (ΔU + PΔV ) + N ln(V new /V old )⎤ log_proposal_ratio = ( - proposed_reduced_pot - - current_reduced_pot - - nr_of_atoms * jnp.log(proposed_volume / initial_volume) + -proposed_reduced_pot + + current_reduced_pot + + nr_of_atoms * jnp.log(proposed_volume / initial_volume) ) # we do not change the thermodynamic state so we can return 'current_thermodynamic_state' @@ -786,7 +916,9 @@ def run(self, n_iterations: int = 1, nbr_list: Optional[PairsBase] = None): log.info(f"Iteration {iteration + 1}/{n_iterations}") for move_name, move in self.move.move_schedule: log.debug(f"Performing: {move_name}") - move.update(self.sampler_state, self.thermodynamic_state, nbr_list) + self.sampler_state, self.thermodynamic_state = move.update( + self.sampler_state, self.thermodynamic_state, nbr_list + ) log.info("Finished running MCMC sampler") log.debug("Closing reporter") @@ -795,294 +927,3 @@ def run(self, n_iterations: int = 1, nbr_list: Optional[PairsBase] = None): move.reporter.flush_buffer() # TODO: flush reporter log.debug(f"Closed reporter {move.reporter.log_file_path}") - - -class MetropolizedMove(MCMove): - """A base class for metropolized moves. - - Only the proposal needs to be specified by subclasses through the method - _propose_positions(). - - Parameters - ---------- - 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. - atom_subset - - Examples - -------- - TBC - """ - - def __init__( - self, - atom_subset: Optional[List[int]] = None, - nr_of_moves: int = 100, - reporter: Optional[_SimulationReporter] = None, - report_frequency: int = 1, - ): - self.n_accepted = 0 - self.n_proposed = 0 - self.atom_subset = atom_subset - super().__init__(nr_of_moves=nr_of_moves, reporter=reporter) - from loguru import logger as log - - self.report_frequency = report_frequency - log.debug(f"Atom subset is {atom_subset}.") - - @property - def statistics(self): - """The acceptance statistics as a dictionary.""" - return dict(n_accepted=self.n_accepted, n_proposed=self.n_proposed) - - @statistics.setter - def statistics(self, value): - self.n_accepted = value["n_accepted"] - self.n_proposed = value["n_proposed"] - - def apply( - self, - thermodynamic_state: ThermodynamicState, - sampler_state: SamplerState, - nbr_list=Optional[PairsBase], - ): - """Apply a metropolized move to the sampler state. - - Total number of acceptances and proposed move are updated. - - Parameters - ---------- - thermodynamic_state : ThermodynamicState - The thermodynamic state to use to apply the move. - sampler_state : SamplerState - The initial sampler state to apply the move to. This is modified. - 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 - from loguru import logger as log - - # 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. - x0 = sampler_state.x0 - atom_subset = self.atom_subset - if atom_subset is None: - 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.") - # Propose perturbed positions. Modifying the reference changes the sampler state. - proposed_positions = self._propose_positions(initial_positions) - - 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 - else: - sampler_state.x0 = sampler_state.x0.at[jnp.array(atom_subset)].set( - proposed_positions - ) - if nbr_list is not None: - 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 - # Accept or reject with Metropolis criteria. - delta_energy = proposed_energy - initial_energy - log.debug(f"Delta energy is {delta_energy} kT.") - import jax.random as jrandom - - self.key, subkey = jrandom.split(self.key) - - compare_to = jrandom.uniform(subkey) - if not jnp.isnan(proposed_energy) and ( - delta_energy <= 0.0 or compare_to < jnp.exp(-delta_energy) - ): - self.n_accepted += 1 - 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}." - ) - if self.n_proposed % self.report_frequency == 0: - self.reporter.report( - { - "energy": proposed_energy, # in kT - "step": self.n_proposed, - "traj": sampler_state.x0, - } - ) - else: - # Restore original positions. - if atom_subset is None: - sampler_state.x0 = initial_positions - else: - sampler_state.x0 = sampler_state.x0.at[jnp.array([atom_subset])].set( - initial_positions - ) - 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. - - These method must be implemented in subclasses. - - Parameters - ---------- - positions : nx3 jnp.ndarray - The original positions of the subset of atoms that these move - applied to. - - Returns - ------- - proposed_positions : nx3 jnp.ndarray - The new proposed positions. - - """ - raise NotImplementedError( - "This MetropolizedMove does not know how to propose new positions." - ) - - -class MetropolisDispMove(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, - displacement_sigma=1.0 * unit.nanometer, - nr_of_moves: int = 100, - atom_subset: Optional[List[int]] = None, - reporter: Optional[MCReporter] = None, - ): - """ - Initialize the MCMC class. - - Parameters - ---------- - seed : int, optional - The seed for the random number generator. Default is 1234. - displacement_sigma : float or unit.Quantity, optional - The standard deviation of the displacement for each move. Default is 1.0 nm. - nr_of_moves : int, optional - The number of moves to perform. Default is 100. - atom_subset : list of int, optional - A subset of atom indices to consider for the moves. Default is None. - reporter : SimulationReporter, optional - The reporter to write the data to. Default is None. - Returns - ------- - None - """ - super().__init__(nr_of_moves=nr_of_moves, reporter=reporter) - self.displacement_sigma = displacement_sigma - self.atom_subset = atom_subset - self.key = None - - def displace_positions( - self, positions: jnp.array, displacement_sigma=1.0 * unit.nanometer - ): - """Return the positions after applying a random displacement to them. - - Parameters - ---------- - positions : nx3 jnp.array unit.Quantity - The positions to displace. - 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). - - Returns - ------- - rotated_positions : nx3 numpy.ndarray openmm.unit.Quantity - The displaced positions. - - """ - import jax.random as jrandom - - self.key, subkey = jrandom.split(self.key) - nr_of_atoms = positions.shape[0] - unitless_displacement_sigma = displacement_sigma.value_in_unit_system( - unit.md_unit_system - ) - displacement_vector = ( - jrandom.normal(subkey, shape=(nr_of_atoms, 3)) * 0.1 - ) # NOTE: convert from Angstrom to nm - scaled_displacement_vector = displacement_vector * unitless_displacement_sigma - updated_position = positions + scaled_displacement_vector - - return updated_position - - def _propose_positions(self, initial_positions: jnp.array) -> jnp.array: - """Implement MetropolizedMove._propose_positions for apply().""" - return self.displace_positions(initial_positions, self.displacement_sigma) - - def run( - self, - sampler_state: SamplerState, - thermodynamic_state: ThermodynamicState, - nbr_list=None, - progress_bar=True, - ): - from tqdm import tqdm - from loguru import logger as log - from jax import random - - self.key = sampler_state.new_PRNG_key - - for trials in ( - tqdm(range(self.nr_of_moves)) if progress_bar else range(self.nr_of_moves) - ): - self.apply(thermodynamic_state, sampler_state, nbr_list) - if trials % 100 == 0: - log.debug(f"Acceptance rate: {self.n_accepted / self.n_proposed}") - if self.reporter is not None: - self.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}") diff --git a/chiron/states.py b/chiron/states.py index 1d9855a..0c498c4 100644 --- a/chiron/states.py +++ b/chiron/states.py @@ -311,6 +311,9 @@ def get_reduced_potential( * sampler_state.box_vectors[1][1] * sampler_state.box_vectors[2][2] ) * unit.nanometer**3 + + from loguru import logger as log + reduced_potential += self.pressure * self.volume return self.beta * reduced_potential From df930039eb13ade6259e6b0b109b550c9897c325 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Tue, 30 Jan 2024 11:34:50 -0800 Subject: [PATCH 13/43] Fixed test_integrator and atom_subsetting in displacement move. --- chiron/mcmc.py | 10 +++++++++- chiron/tests/conftest.py | 6 ++---- chiron/tests/test_integrators.py | 12 +++++++----- chiron/utils.py | 1 + 4 files changed, 19 insertions(+), 10 deletions(-) diff --git a/chiron/mcmc.py b/chiron/mcmc.py index dea694f..47971d8 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -520,7 +520,9 @@ def __init__( method="metropolis", ) self.displacement_sigma = displacement_sigma + self.atom_subset = atom_subset + self.atom_subset_mask = None def _report( self, @@ -585,6 +587,12 @@ def _propose( """ Implement the logic specific to displacement changes. """ + if self.atom_subset is not None and self.atom_subset_mask is None: + import jax.numpy as jnp + + self.atom_subset_mask = jnp.zeros(current_sampler_state.n_particles) + for atom in self.atom_subset: + self.atom_subset_mask = self.atom_subset_mask.at[atom].set(1) key = current_sampler_state.new_PRNG_key @@ -605,7 +613,7 @@ def _propose( if self.atom_subset is not None: proposed_sampler_state.x0 = ( proposed_sampler_state.x0 - + scaled_displacement_vector * self.atom_subset + + scaled_displacement_vector * self.atom_subset_mask ) else: proposed_sampler_state.x0 = ( diff --git a/chiron/tests/conftest.py b/chiron/tests/conftest.py index 742d5ab..9f1cca6 100644 --- a/chiron/tests/conftest.py +++ b/chiron/tests/conftest.py @@ -29,8 +29,8 @@ def provide_testsystems_and_potentials(): import jax.numpy as jnp hoa_potential = HarmonicOscillatorPotential( - ho.topology, - ho.K, + hoa.topology, + hoa.K, x0=unit.Quantity( jnp.array( [ @@ -54,5 +54,3 @@ def provide_testsystems_and_potentials(): (hoa, hoa_potential), ] return TESTSYSTEM_AND_POTENTIAL - - diff --git a/chiron/tests/test_integrators.py b/chiron/tests/test_integrators.py index 2e67928..a68678a 100644 --- a/chiron/tests/test_integrators.py +++ b/chiron/tests/test_integrators.py @@ -42,20 +42,22 @@ def test_langevin_dynamics(prep_temp_dir, provide_testsystems_and_potentials): reporter = LangevinDynamicsReporter() - integrator = LangevinIntegrator(reporter=reporter, report_frequency=1) - with pytest.raises(ValueError): + integrator = LangevinIntegrator( + reporter=reporter, report_frequency=1, reinitialize_velocities=False + ) + integrator.run( sampler_state, thermodynamic_state, n_steps=20, - initialize_velocities=False, ) - + integrator = LangevinIntegrator( + reporter=reporter, report_frequency=1, reinitialize_velocities=True + ) integrator.run( sampler_state, thermodynamic_state, n_steps=20, - initialize_velocities=True, ) i = i + 1 diff --git a/chiron/utils.py b/chiron/utils.py index 6b349e1..07927e0 100644 --- a/chiron/utils.py +++ b/chiron/utils.py @@ -113,6 +113,7 @@ def initialize_velocities( import jax.numpy as jnp mass = get_list_of_mass(topology) + kB = unit.BOLTZMANN_CONSTANT_kB * unit.AVOGADRO_CONSTANT_NA kbT_unitless = (kB * temperature).value_in_unit_system(unit.md_unit_system) From acfe5e23eb4a25ece1166ff40e4eec2726cdcc58 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Tue, 30 Jan 2024 11:54:57 -0800 Subject: [PATCH 14/43] Added in flag to initialize velocities the first time langevin is called. This is different than reinitialize which will call reinit every time run is called. The velocity init function was spun off into the utils file. velocities can be set in the sampler state; code will throw and error if initialize_velocities or reinitialize_velocities are false, and velocities are not set in the sampler state. --- chiron/integrators.py | 15 +++++++++++++++ chiron/mcmc.py | 5 +++++ chiron/multistate.py | 4 ++-- chiron/tests/test_multistate.py | 4 +++- chiron/tests/test_utils.py | 4 +++- 5 files changed, 28 insertions(+), 4 deletions(-) diff --git a/chiron/integrators.py b/chiron/integrators.py index 922a5a1..d3a5c3f 100644 --- a/chiron/integrators.py +++ b/chiron/integrators.py @@ -26,6 +26,7 @@ def __init__( self, stepsize=1.0 * unit.femtoseconds, collision_rate=1.0 / unit.picoseconds, + initialize_velocities: bool = False, reinitialize_velocities: bool = False, report_frequency: int = 100, reporter: Optional[LangevinDynamicsReporter] = None, @@ -40,6 +41,10 @@ def __init__( Time step of integration with units of time. Default is 1.0 * unit.femtoseconds. collision_rate : unit.Quantity, optional Collision rate for the Langevin dynamics, with units 1/time. Default is 1.0 / unit.picoseconds. + initialize_velocities : bool, optional + Flag indicating whether to initialize the velocities the first time the run function is called. Default is False. + reinitialize_velocities : bool, optional + Flag indicating whether to reinitialize the velocities each time the run function is called. Default is False. report_frequency : int, optional Frequency of saving the simulation data. Default is 100. reporter : SimulationReporter, optional @@ -71,6 +76,7 @@ def __init__( self.save_traj_in_memory = save_traj_in_memory self.traj = [] self.reinitialize_velocities = reinitialize_velocities + self.initialize_velocities = initialize_velocities def run( self, @@ -142,7 +148,16 @@ def run( sampler_state.velocities = initialize_velocities( temperature, potential.topology, key ) + self.initialize_velocities = False + elif self.initialize_velocities: + # v0 = sigma_v * random.normal(key, x0.shape) + from .utils import initialize_velocities + + sampler_state.velocities = initialize_velocities( + temperature, potential.topology, key + ) + self.initialize_velocities = False else: if sampler_state._velocities is None: raise ValueError("Velocities must be set before running the integrator") diff --git a/chiron/mcmc.py b/chiron/mcmc.py index 47971d8..0c8fd5a 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -81,6 +81,7 @@ def __init__( self, stepsize=1.0 * unit.femtoseconds, collision_rate=1.0 / unit.picoseconds, + initialize_velocities: bool = False, reinitialize_velocities: bool = False, reporter: Optional[LangevinDynamicsReporter] = None, report_frequency: int = 100, @@ -96,6 +97,9 @@ def __init__( Time step size for the integration. collision_rate : unit.Quantity Collision rate for the Langevin dynamics. + initialize_velocities: bool, optional + Whether to initialize the velocities the first time the run function is called. + Default is False. reinitialize_velocities : bool, optional Whether to reinitialize the velocities each time the run function is called. Default is False. @@ -127,6 +131,7 @@ def __init__( self.integrator = LangevinIntegrator( stepsize=self.stepsize, collision_rate=self.collision_rate, + initialize_velocities=initialize_velocities, reinitialize_velocities=reinitialize_velocities, report_frequency=report_frequency, reporter=reporter, diff --git a/chiron/multistate.py b/chiron/multistate.py index b9100ec..2d17766 100644 --- a/chiron/multistate.py +++ b/chiron/multistate.py @@ -75,7 +75,7 @@ def __init__( self._neighborhoods = None self._n_accepted_matrix = None self._n_proposed_matrix = None - self._reporter = reporter # NOTE: reporter needs to be putlic, API change ahead + self._reporter = reporter # NOTE: reporter needs to be putlic, API change ahead self._metadata = None self._mcmc_moves = copy.deepcopy(mcmc_moves) self._online_estimator = None @@ -396,7 +396,7 @@ def _propagate_replica(self, replica_id: int): thermodynamic_state = self._thermodynamic_states[thermodynamic_state_id] mcmc_move = self._mcmc_moves[thermodynamic_state_id] # Apply the MCMC move to the replica. - mcmc_move.run(sampler_state, thermodynamic_state) + mcmc_move.update(sampler_state, thermodynamic_state) # Append the new state to the trajectory for analysis. self._traj[replica_id].append(sampler_state.x0) diff --git a/chiron/tests/test_multistate.py b/chiron/tests/test_multistate.py index ef31a84..bbfdd9c 100644 --- a/chiron/tests/test_multistate.py +++ b/chiron/tests/test_multistate.py @@ -24,7 +24,9 @@ def setup_sampler() -> Tuple[NeighborListNsqrd, MultiStateSampler]: OrthogonalPeriodicSpace(), cutoff=cutoff, skin=skin, n_max_neighbors=180 ) - move = LangevinDynamicsMove(stepsize=1.0 * unit.femtoseconds, nr_of_steps=100) + move = LangevinDynamicsMove( + stepsize=1.0 * unit.femtoseconds, nr_of_steps=100, initialize_velocities=True + ) BaseReporter.set_directory("multistate_test") reporter = MultistateReporter() reporter.reset_reporter_file() diff --git a/chiron/tests/test_utils.py b/chiron/tests/test_utils.py index 59f6d10..119a7f5 100644 --- a/chiron/tests/test_utils.py +++ b/chiron/tests/test_utils.py @@ -64,7 +64,9 @@ def test_reporter(prep_temp_dir, ho_multistate_sampler_multiple_ks): reporter = LangevinDynamicsReporter("langevin_test") reporter.reset_reporter_file() - integrator = LangevinIntegrator(reporter=reporter, report_frequency=1) + integrator = LangevinIntegrator( + reporter=reporter, report_frequency=1, initialize_velocities=True + ) integrator.run( sampler_state, thermodynamic_state, From beb0bdc305532f10c4f8d6a8ca7248d3f9aa332e Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Tue, 30 Jan 2024 15:29:49 -0800 Subject: [PATCH 15/43] test_multistate still is failing an assertion, but I think that is fine, since this PR was focused on revamping the MC moves. That can be tackled separately in the multistate PR. --- chiron/multistate.py | 2 +- chiron/tests/test_multistate.py | 16 ++++++++++------ 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/chiron/multistate.py b/chiron/multistate.py index 2d17766..3306ee1 100644 --- a/chiron/multistate.py +++ b/chiron/multistate.py @@ -75,7 +75,7 @@ def __init__( self._neighborhoods = None self._n_accepted_matrix = None self._n_proposed_matrix = None - self._reporter = reporter # NOTE: reporter needs to be putlic, API change ahead + self._reporter = reporter # NOTE: reporter needs to be public, API change ahead self._metadata = None self._mcmc_moves = copy.deepcopy(mcmc_moves) self._online_estimator = None diff --git a/chiron/tests/test_multistate.py b/chiron/tests/test_multistate.py index bbfdd9c..ff8feb6 100644 --- a/chiron/tests/test_multistate.py +++ b/chiron/tests/test_multistate.py @@ -214,17 +214,21 @@ def test_multistate_run(ho_multistate_sampler_multiple_ks: MultiStateSampler): print(f"Analytical free energy difference: {ho_sampler.delta_f_ij_analytical[0]}") - n_iteratinos = 250 - ho_sampler.run(n_iteratinos) + n_iterations = 250 + ho_sampler.run(n_iterations) # check that we have the correct number of iterations, replicas and states - assert ho_sampler.iteration == n_iteratinos - assert ho_sampler._iteration == n_iteratinos + assert ho_sampler.iteration == n_iterations + assert ho_sampler._iteration == n_iterations assert ho_sampler.n_replicas == 4 assert ho_sampler.n_states == 4 - u_kn = ho_sampler.reporter.get_property("u_kn") - assert u_kn.shape == (n_iteratinos, 4, 4) + u_kn = ho_sampler._reporter.get_property("u_kn") + + # this appears to need to be n+1 iterations; I assume because of the initial logging of the value + assert u_kn.shape == (4, 4, n_iterations + 1) + + # assert u_kn.shape == (n_iterations, 4, 4) # check that the free energies are correct print(ho_sampler.analytical_f_i) # [ 0. , -0.28593054, -0.54696467, -0.78709279] From 0b57c2babb7f980be81000677261cef4dfbd7502 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Wed, 31 Jan 2024 11:08:53 -0800 Subject: [PATCH 16/43] Modified the routines to ensure that passing the nbr_list fits in more with the functional programming model (passing and return a neighborlist). --- Examples/LJ_langevin.py | 4 +- chiron/integrators.py | 27 ++- chiron/mcmc.py | 274 ++++++++++++++++++++++--------- chiron/tests/test_integrators.py | 2 +- 4 files changed, 222 insertions(+), 85 deletions(-) diff --git a/Examples/LJ_langevin.py b/Examples/LJ_langevin.py index 5b3cdb4..b9c5f78 100644 --- a/Examples/LJ_langevin.py +++ b/Examples/LJ_langevin.py @@ -78,10 +78,10 @@ ) print("init_energy: ", lj_potential.compute_energy(sampler_state.x0, nbr_list)) -updated_sampler_state = integrator.run( +updated_sampler_state, updated_nbr_list = integrator.run( sampler_state, thermodynamic_state, - n_steps=5000, + n_steps=1000, nbr_list=nbr_list, progress_bar=True, ) diff --git a/chiron/integrators.py b/chiron/integrators.py index d3a5c3f..da0021d 100644 --- a/chiron/integrators.py +++ b/chiron/integrators.py @@ -5,7 +5,7 @@ from openmm import unit from .states import SamplerState, ThermodynamicState from .reporters import LangevinDynamicsReporter -from typing import Optional +from typing import Optional, Tuple from .potential import NeuralNetworkPotential from .neighbors import PairsBase @@ -85,7 +85,7 @@ def run( n_steps: int = 5_000, nbr_list: Optional[PairsBase] = None, progress_bar=False, - ): + ) -> Tuple[SamplerState, PairsBase]: """ Run the integrator to perform Langevin dynamics molecular dynamics simulation. @@ -106,6 +106,8 @@ def run( ------- sampler_state : SamplerState The final state of the simulation, including positions, velocities, and current PRNG key. + nbr_list : PairBase + The neighbor list for the final state of the simulation. If the NeighborList object is None, the function returns None. """ from .utils import get_list_of_mass from tqdm import tqdm @@ -181,15 +183,19 @@ def run( # r x += (stepsize_unitless * 0.5) * v - if nbr_list is not None: - x = self._wrap_and_rebuild_neighborlist(x, nbr_list) + # we can actually skip this for now, and just wrap/check/rebuild + # right before we call the force + + # if nbr_list is not None: + # x = self._wrap_and_rebuild_neighborlist(x, nbr_list) # o random_noise_v = random.normal(subkey, x.shape) v = (a * v) + (b * sigma_v * random_noise_v) x += (stepsize_unitless * 0.5) * v + if nbr_list is not None: - x = self._wrap_and_rebuild_neighborlist(x, nbr_list) + x, nbr_list = self._wrap_and_rebuild_neighborlist(x, nbr_list) F = potential.compute_force(x, nbr_list) # v @@ -213,7 +219,7 @@ def run( updated_sampler_state.velocities = v updated_sampler_state.current_PRNG_key = key - return updated_sampler_state + return updated_sampler_state, nbr_list def _wrap_and_rebuild_neighborlist(self, x: jnp.array, nbr_list: PairsBase): """ @@ -225,13 +231,20 @@ def _wrap_and_rebuild_neighborlist(self, x: jnp.array, nbr_list: PairsBase): The coordinates of the particles. nbr_list: PairsBsse The neighborlist object. + + Returns + ------- + x: jnp.array + The wrapped coordinates. + nbr_list: PairsBase + The neighborlist object; this may or may not have been rebuilt. """ x = nbr_list.space.wrap(x) # check if we need to rebuild the neighborlist after moving the particles if nbr_list.check(x): nbr_list.build(x, self.box_vectors) - return x + return x, nbr_list def _report( self, diff --git a/chiron/mcmc.py b/chiron/mcmc.py index 0c8fd5a..2afb50e 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -36,7 +36,7 @@ def __init__( self.report_frequency = report_frequency # we need to keep track of which iteration we are on - self.iteration = 0 + self._move_iteration = 0 from loguru import logger as log @@ -52,7 +52,7 @@ def update( sampler_state: SamplerState, thermodynamic_state: ThermodynamicState, nbr_list: Optional[PairsBase] = None, - ): + ) -> Tuple[SamplerState, ThermodynamicState, Optional[PairsBase]]: """ Update the state of the system. @@ -71,6 +71,8 @@ def update( The updated sampler state. thermodynamic_state : ThermodynamicState The updated thermodynamic state. + nbr_list: PairsBase + The updated neighbor/pair list. If no nbr_list is passed, this will be None. """ pass @@ -79,13 +81,13 @@ def update( class LangevinDynamicsMove(MCMCMove): def __init__( self, - stepsize=1.0 * unit.femtoseconds, - collision_rate=1.0 / unit.picoseconds, + stepsize: unit.Quantity = 1.0 * unit.femtoseconds, + collision_rate: unit.Quantity = 1.0 / unit.picoseconds, initialize_velocities: bool = False, reinitialize_velocities: bool = False, reporter: Optional[LangevinDynamicsReporter] = None, report_frequency: int = 100, - nr_of_steps=1_000, + nr_of_steps: int = 1_000, save_traj_in_memory: bool = False, ): """ @@ -143,7 +145,7 @@ def update( sampler_state: SamplerState, thermodynamic_state: ThermodynamicState, nbr_list: Optional[PairsBase] = None, - ): + ) -> Tuple[SamplerState, ThermodynamicState, Optional[PairsBase]]: """ Run the integrator to perform molecular dynamics simulation. @@ -156,6 +158,14 @@ def update( nbr_list : PairsBase, optional The neighbor list to use for the simulation. + Returns + ------- + sampler_state : SamplerState + The updated sampler state. + thermodynamic_state : ThermodynamicState + The thermodynamic state; note this is not modified by the Langevin dynamics algorithm. + nbr_list: PairsBase + The updated neighbor/pair list. If a nbr_list is not set, this will be None. """ assert isinstance( @@ -165,7 +175,7 @@ def update( thermodynamic_state, ThermodynamicState ), f"Thermodynamic state must be ThermodynamicState, not {type(thermodynamic_state)}" - sampler_state = self.integrator.run( + updated_sampler_state, updated_nbr_list = self.integrator.run( thermodynamic_state=thermodynamic_state, sampler_state=sampler_state, n_steps=self.nr_of_moves, @@ -176,10 +186,10 @@ def update( self.traj.append(self.integrator.traj) self.integrator.traj = [] - self.iteration += 1 + self._move_iteration += 1 # The thermodynamic_state will not change for the langevin move - return sampler_state, thermodynamic_state + return updated_sampler_state, thermodynamic_state, updated_nbr_list class MCMove(MCMCMove): @@ -228,7 +238,7 @@ def update( sampler_state: SamplerState, thermodynamic_state: ThermodynamicState, nbr_list: Optional[PairsBase] = None, - ): + ) -> Tuple[SamplerState, ThermodynamicState, Optional[PairsBase]]: """ Perform the defined move and update the state. @@ -247,11 +257,13 @@ def update( The updated sampler state. thermodynamic_state : ThermodynamicState The updated thermodynamic state. + nbr_list: PairsBase + The updated neighbor/pair list. If a nbr_list is not set, this will be None. """ calculate_current_potential = True for i in range(self.nr_of_moves): - sampler_state, thermodynamic_state = self._step( + sampler_state, thermodynamic_state, nbr_list = self._step( sampler_state, thermodynamic_state, nbr_list, @@ -260,33 +272,33 @@ def update( # after the first step, we don't need to recalculate the current potential, it will be stored calculate_current_potential = False - elapsed_step = i + self.iteration * self.nr_of_moves + elapsed_step = i + self._move_iteration * self.nr_of_moves if hasattr(self, "reporter"): if self.reporter is not None: - # I think it makes sense to i + self.nr_of_moves*self.iteration as our current "step" - # otherwise, instances where self.report_frequency > self.nr_of_moves would only report on the - # first step which might actually be more frequent than we specify + # I think it makes sense to use i + self.nr_of_moves*self._move_iteration as our current "step" + # otherwise, if we just used i, instances where self.report_frequency > self.nr_of_moves would only report on the + # first step, which might actually be more frequent than we specify if elapsed_step % self.report_frequency == 0: self._report( i, - self.iteration, + self._move_iteration, self.n_accepted / self.n_proposed, sampler_state, thermodynamic_state, nbr_list, ) if self.update_stepsize: - # if we only used i, we might never actually update the parameters if we have a move that is called infrequently + # if we only used i, we might never actually update the parameters if we have a move that is called infrequently if ( elapsed_step % self.update_stepsize_frequency == 0 and elapsed_step > 0 ): self._update_stepsize() + # keep track of how many times this function has been called + self._move_iteration += 1 - self.iteration += 1 - - return sampler_state, thermodynamic_state + return sampler_state, thermodynamic_state, nbr_list @abstractmethod def _report( @@ -329,22 +341,56 @@ def _update_stepsize(self): or a maximum volume change factor for a Monte Carlo barostat move. Since different moves will be modifying different quantities, this needs to be defined for each move. + + Note this will modify the "stepsize" in place. """ pass def _step( self, - current_sampler_state, - current_thermodynamic_state, - nbr_list, - calculate_current_potential=True, - ): - # if this is the first time we are calling this, - # we will need to recalculate the reduced potential for the current state + current_sampler_state: SamplerState, + current_thermodynamic_state: ThermodynamicState, + current_nbr_list: Optional[PairsBase] = None, + calculate_current_potential: bool = True, + ) -> Tuple[SamplerState, ThermodynamicState, Optional[PairsBase]]: + """ + Performs an individual MC step. + + This will call the _propose function which will be specific to the type of move. + + Parameters + ---------- + current_sampler_state : SamplerState, required + Current sampler state. + current_thermodynamic_state : ThermodynamicState, required + Current thermodynamic state. + current_nbr_list : Optional[PairsBase] + Neighbor list associated with the current state. + calculate_current_potential : bool, optional + Whether to calculate the current reduced potential. Default is True. + + Returns + ------- + sampler_state : SamplerState + The updated sampler state; if a move is rejected this will be unchanged. + Note, if the proposed move is rejected, the current PRNG key will be updated to ensure + that we are using a different random number for the next iteration. + thermodynamic_state : ThermodynamicState + The updated thermodynamic state; if a move is rejected this will be unchanged. + Note, many MC moves will not modify the thermodynamic state regardless of acceptance of the move. + nbr_list: PairsBase, optional + The updated neighbor/pair list. If a nbr_list is not set, this will be None. + If the move is rejected, this will correspond to the neighbor + + """ + + # if this is the first time we are calling this function during this iteration + # we will need to calculate the reduced potential for the current state # this is toggled by the calculate_current_potential flag + # otherwise, we can use the one that was saved from the last step, for efficiency if calculate_current_potential: current_reduced_pot = current_thermodynamic_state.get_reduced_potential( - current_sampler_state, nbr_list + current_sampler_state, current_nbr_list ) # save the current_reduced_pot so we don't have to recalculate # it on the next iteration if the move is rejected @@ -356,16 +402,20 @@ def _step( # this will be specific to the type of move # in addition to the sampler_state, this will require/return the thermodynamic state # for systems that e.g., make changes to particle identity. + # For efficiency, we will also return a copy of the nbr_list associated with the proposed state + # because if the move is rejected, we can move back the original state without having to rebuild the nbr_list + # if it were modified due to the proposed state. ( proposed_sampler_state, proposed_thermodynamic_state, proposed_reduced_pot, log_proposal_ratio, + proposed_nbr_list, ) = self._propose( current_sampler_state, current_thermodynamic_state, current_reduced_pot, - nbr_list, + current_nbr_list, ) # accept or reject the proposed state @@ -390,7 +440,11 @@ def _step( # not sure this needs to be a separate function but for simplicity in outlining the code it is fine # or should this return the new sampler_state and thermodynamic_state? - return proposed_sampler_state, proposed_thermodynamic_state + return ( + proposed_sampler_state, + proposed_thermodynamic_state, + proposed_nbr_list, + ) else: # if we reject the move, we need to update the current_PRNG key to ensure that # we are using a different random number for the next iteration @@ -398,13 +452,8 @@ def _step( current_sampler_state._current_PRNG_key = ( proposed_sampler_state._current_PRNG_key ) - if nbr_list is not None: - if nbr_list.check(current_sampler_state.x0): - nbr_list.build( - current_sampler_state.x0, current_sampler_state.box_vectors - ) - return current_sampler_state, current_thermodynamic_state + return current_sampler_state, current_thermodynamic_state, current_nbr_list def _update_statistics(self, decision): """ @@ -431,8 +480,12 @@ def reset_statistics(self): @abstractmethod def _propose( - self, current_sampler_state, current_thermodynamic_state, current_reduced_pot - ): + self, + current_sampler_state: SamplerState, + current_thermodynamic_state: ThermodynamicState, + current_reduced_pot: float, + current_nbr_list: Optional[PairsBase] = None, + ) -> Tuple[SamplerState, ThermodynamicState, float, float, Optional[PairsBase]]: """ Propose a new state and calculate the log proposal ratio. @@ -449,6 +502,8 @@ def _propose( Current thermodynamic state. current_reduced_pot : float, required Current reduced potential. + current_nbr_list : PairsBase, required + Neighbor list associated with the current state. Returns ------- @@ -460,6 +515,8 @@ def _propose( Proposed reduced potential. log_proposal_ratio : float Log proposal ratio. + proposed_nbr_list : PairsBase + Proposed neighbor list. If not defined, this will be None. """ pass @@ -531,12 +588,12 @@ def __init__( def _report( self, - step, - iteration, - acceptance_probability, - sampler_state, - thermodynamic_state, - nbr_list, + step: int, + iteration: int, + acceptance_probability: float, + sampler_state: SamplerState, + thermodynamic_state: ThermodynamicState, + nbr_list: Optional[PairsBase] = None, ): """ Report the current state of the MC displacement move. @@ -584,14 +641,41 @@ def _update_stepsize(self): def _propose( self, - current_sampler_state, - current_thermodynamic_state, - current_reduced_pot, - nbr_list, - ): + current_sampler_state: SamplerState, + current_thermodynamic_state: ThermodynamicState, + current_reduced_pot: float, + current_nbr_list: Optional[PairsBase] = None, + ) -> Tuple[SamplerState, ThermodynamicState, float, float, Optional[PairsBase]]: """ Implement the logic specific to displacement changes. + + Parameters + ---------- + current_sampler_state : SamplerState, required + Current sampler state. + current_thermodynamic_state : ThermodynamicState, required + Current thermodynamic state. + current_reduced_pot : float, required + Current reduced potential. + current_nbr_list : Optional[PairsBase] + Neighbor list associated with the current state. + + Returns + ------- + proposed_sampler_state : SamplerState + Proposed sampler state. + proposed_thermodynamic_state : ThermodynamicState + Proposed thermodynamic state. + proposed_reduced_pot : float + Proposed reduced potential. + log_proposal_ratio : float + Log proposal ratio. + proposed_nbr_list : PairsBase + Proposed neighbor list. If not defined, this will be None. """ + + # create a mask for the atom subset: if a value of the mask is 0 + # the particle won't move; if 1 the particle will be moved if self.atom_subset is not None and self.atom_subset_mask is None: import jax.numpy as jnp @@ -625,18 +709,30 @@ def _propose( proposed_sampler_state.x0 + scaled_displacement_vector ) - # after proposing a move we need to wrap particles and see if we need to rebuild - # the neighborlist - if nbr_list is not None: - proposed_sampler_state.x0 = nbr_list.space.wrap(proposed_sampler_state.x0) + # after proposing a move we need to wrap particles and see if we need to rebuild the neighborlist + if current_nbr_list is not None: + proposed_sampler_state.x0 = current_nbr_list.space.wrap( + proposed_sampler_state.x0 + ) + + # if we need to rebuild the neighbor the neighborlist + # we will make a copy and then build + if current_nbr_list.check(proposed_sampler_state.x0): + import copy - if nbr_list.check(proposed_sampler_state.x0): - nbr_list.build( + proposed_nbr_list = copy.deepcopy(current_nbr_list) + + proposed_nbr_list.build( proposed_sampler_state.x0, proposed_sampler_state.box_vectors ) + # if we don't need to update the neighborlist, just make a new variable that refers to the original + else: + proposed_nbr_list = current_nbr_list + else: + proposed_nbr_list = None proposed_reduced_pot = current_thermodynamic_state.get_reduced_potential( - proposed_sampler_state, nbr_list + proposed_sampler_state, proposed_nbr_list ) log_proposal_ratio = -proposed_reduced_pot + current_reduced_pot @@ -648,6 +744,7 @@ def _propose( current_thermodynamic_state, proposed_reduced_pot, log_proposal_ratio, + proposed_nbr_list, ) @@ -656,7 +753,6 @@ def __init__( self, volume_max_scale=0.01, nr_of_moves: int = 100, - atom_subset: Optional[List[int]] = None, report_frequency: int = 1, reporter: Optional[LangevinDynamicsReporter] = None, update_stepsize: bool = True, @@ -671,8 +767,6 @@ def __init__( The standard deviation of the displacement for each move. Default is 1.0 nm. nr_of_moves : int, optional The number of moves to perform. Default is 100. - atom_subset : list of int, optional - A subset of atom indices to consider for the moves. Default is None. reporter : SimulationReporter, optional The reporter to write the data to. Default is None. update_stepsize : bool, optional @@ -695,12 +789,12 @@ def __init__( def _report( self, - step, - iteration, - acceptance_probability, - sampler_state, - thermodynamic_state, - nbr_list, + step: int, + iteration: int, + acceptance_probability: float, + sampler_state: SamplerState, + thermodynamic_state: ThermodynamicState, + nbr_list: Optional[PairsBase] = None, ): """ @@ -719,6 +813,7 @@ def _report( nbr_list : Optional[PairBase]=None The neighbor list or pair list for evaluating interactions in the system, default None """ + potential = thermodynamic_state.potential.compute_energy( sampler_state.x0, nbr_list ) @@ -751,13 +846,38 @@ def _update_stepsize(self): def _propose( self, - current_sampler_state, - current_thermodynamic_state, - current_reduced_pot, - nbr_list, - ): + current_sampler_state: SamplerState, + current_thermodynamic_state: ThermodynamicState, + current_reduced_pot: float, + current_nbr_list: Optional[PairsBase] = None, + ) -> Tuple[SamplerState, ThermodynamicState, float, float, Optional[PairsBase]]: """ Implement the logic specific to displacement changes. + + Parameters + ---------- + current_sampler_state : SamplerState, required + Current sampler state. + current_thermodynamic_state : ThermodynamicState, required + Current thermodynamic state. + current_reduced_pot : float, required + Current reduced potential. + current_nbr_list : PairsBase, optional + Neighbor list associated with the current state. + + Returns + ------- + proposed_sampler_state : SamplerState + Proposed sampler state. + proposed_thermodynamic_state : ThermodynamicState + Proposed thermodynamic state. + proposed_reduced_pot : float + Proposed reduced potential. + log_proposal_ratio : float + Log proposal ratio. + proposed_nbr_list : PairsBase + Proposed neighbor list. If not defined, this will be None. + """ from loguru import logger as log @@ -794,14 +914,15 @@ def _propose( current_sampler_state.box_vectors * length_scaling_factor ) - if nbr_list is not None: - # after scaling the box vectors we should rebuild the neighborlist - nbr_list.build( + if current_nbr_list is not None: + proposed_nbr_list = copy.deepcopy(current_nbr_list) + # after scaling the box vectors and coordinates we should always rebuild the neighborlist + proposed_nbr_list.build( proposed_sampler_state.x0, proposed_sampler_state.box_vectors ) proposed_reduced_pot = current_thermodynamic_state.get_reduced_potential( - proposed_sampler_state, nbr_list + proposed_sampler_state, proposed_nbr_list ) # χ = exp ⎡−β (ΔU + PΔV ) + N ln(V new /V old )⎤ @@ -817,6 +938,7 @@ def _propose( current_thermodynamic_state, proposed_reduced_pot, log_proposal_ratio, + proposed_nbr_list, ) @@ -929,7 +1051,7 @@ def run(self, n_iterations: int = 1, nbr_list: Optional[PairsBase] = None): log.info(f"Iteration {iteration + 1}/{n_iterations}") for move_name, move in self.move.move_schedule: log.debug(f"Performing: {move_name}") - self.sampler_state, self.thermodynamic_state = move.update( + self.sampler_state, self.thermodynamic_state, nbr_list = move.update( self.sampler_state, self.thermodynamic_state, nbr_list ) @@ -940,3 +1062,5 @@ def run(self, n_iterations: int = 1, nbr_list: Optional[PairsBase] = None): move.reporter.flush_buffer() # TODO: flush reporter log.debug(f"Closed reporter {move.reporter.log_file_path}") + + # I think we should return the sampler/thermo state to be consistent diff --git a/chiron/tests/test_integrators.py b/chiron/tests/test_integrators.py index a68678a..ed6f01f 100644 --- a/chiron/tests/test_integrators.py +++ b/chiron/tests/test_integrators.py @@ -55,7 +55,7 @@ def test_langevin_dynamics(prep_temp_dir, provide_testsystems_and_potentials): integrator = LangevinIntegrator( reporter=reporter, report_frequency=1, reinitialize_velocities=True ) - integrator.run( + updated_sampler_state, updated_nbr_list = integrator.run( sampler_state, thermodynamic_state, n_steps=20, From 85a9ae4904424f7fd965299c9ab26b6d82f85efd Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Wed, 31 Jan 2024 11:58:47 -0800 Subject: [PATCH 17/43] Added ideal gas test written in the other PR. updated logic in langevin to take into account the iteration not just the current step; reporters now also log the elapsed_step. --- Examples/Idealgas.py | 17 +++-- Examples/LJ_langevin.py | 2 + chiron/integrators.py | 20 ++++-- chiron/mcmc.py | 10 +++ chiron/tests/test_testsystems.py | 117 +++++++++++++++++++++++++++++++ 5 files changed, 154 insertions(+), 12 deletions(-) diff --git a/Examples/Idealgas.py b/Examples/Idealgas.py index f696474..19f6427 100644 --- a/Examples/Idealgas.py +++ b/Examples/Idealgas.py @@ -48,7 +48,7 @@ nbr_list = PairList(OrthogonalPeriodicSpace(), cutoff=cutoff) nbr_list.build_from_state(sampler_state) -from chiron.reporters import _SimulationReporter +from chiron.reporters import MCReporter # initialize a reporter to save the simulation data filename = "test_mc_ideal_gas.h5" @@ -56,7 +56,7 @@ if os.path.isfile(filename): os.remove(filename) -reporter = _SimulationReporter(filename, 1) +reporter = MCReporter(filename, 100) from chiron.mcmc import ( @@ -81,13 +81,16 @@ ) sampler = MCMCSampler(move_set, sampler_state, thermodynamic_state) -sampler.run(n_iterations=50, nbr_list=nbr_list) # how many times to repeat +sampler.run(n_iterations=10, nbr_list=nbr_list) # how many times to repeat -import h5py -with h5py.File(filename, "r") as f: - volume = f["volume"][:] - steps = f["step"][:] +volume = reporter.get_property("volume") +step = reporter.get_property("step") + +import matplotlib.pyplot as plt + +plt.plot(step, volume) +plt.show() # get expectations ideal_volume = ideal_gas.get_volume_expectation(thermodynamic_state) diff --git a/Examples/LJ_langevin.py b/Examples/LJ_langevin.py index b9c5f78..6a34f0d 100644 --- a/Examples/LJ_langevin.py +++ b/Examples/LJ_langevin.py @@ -94,6 +94,8 @@ energies = f["potential_energy"][:] steps = f["step"][:] +energies = reporter.get_property("potential_energy") +steps = reporter.get_property("step") # plot the energy import matplotlib.pyplot as plt diff --git a/chiron/integrators.py b/chiron/integrators.py index da0021d..6b40721 100644 --- a/chiron/integrators.py +++ b/chiron/integrators.py @@ -77,6 +77,7 @@ def __init__( self.traj = [] self.reinitialize_velocities = reinitialize_velocities self.initialize_velocities = initialize_velocities + self._move_iteration = 0 def run( self, @@ -183,8 +184,8 @@ def run( # r x += (stepsize_unitless * 0.5) * v - # we can actually skip this for now, and just wrap/check/rebuild - # right before we call the force + # we can actually skip this wrapping, and just wrap/check/rebuild + # right before we call the force again. # if nbr_list is not None: # x = self._wrap_and_rebuild_neighborlist(x, nbr_list) @@ -201,9 +202,9 @@ def run( # v v += (stepsize_unitless * 0.5) * F / mass_unitless - if step % self.report_frequency == 0: + if (step + self._move_iteration * n_steps) % self.report_frequency == 0: if hasattr(self, "reporter") and self.reporter is not None: - self._report(x, potential, nbr_list, step) + self._report(x, potential, nbr_list, step, self._move_iteration) if self.save_traj_in_memory: self.traj.append(x) @@ -252,6 +253,8 @@ def _report( potential: NeuralNetworkPotential, nbr_list: PairsBase, step: int, + iteration: int, + elapsed_step: int, ): """ Reports the trajectory, energy, step, and box vectors (if available) to the reporter. @@ -265,7 +268,12 @@ def _report( nbr_list: PairsBase The neighbor list step: int - The current time step. + The current step in the move; this resets each iteration. + iteration: int + The number iterations the move has been called. + elapsed_step: int, + The total number of steps that have been taken in the simulation move. + Returns: None @@ -274,6 +282,8 @@ def _report( "positions": x, "potential_energy": potential.compute_energy(x, nbr_list), "step": step, + "iteration": iteration, + "elapsed_step": elapsed_step, } if nbr_list is not None: d["box_vectors"] = nbr_list.space.box_vectors diff --git a/chiron/mcmc.py b/chiron/mcmc.py index 2afb50e..04be25e 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -283,6 +283,7 @@ def update( self._report( i, self._move_iteration, + elapsed_step, self.n_accepted / self.n_proposed, sampler_state, thermodynamic_state, @@ -305,6 +306,7 @@ def _report( self, step: int, iteration: int, + elapsed_step: int, acceptance_probability: float, sampler_state: SamplerState, thermodynamic_state: ThermodynamicState, @@ -322,6 +324,8 @@ def _report( The current step of the simulation move. iteration : int The current iteration of the move sequence (i.e., how many times has this been called thus far). + elapsed_step : int + The total number of steps that have been taken in the simulation move. step+ nr_moves*iteration acceptance_probability : float The acceptance probability of the move. sampler_state : SamplerState @@ -590,6 +594,7 @@ def _report( self, step: int, iteration: int, + elapsed_step: int, acceptance_probability: float, sampler_state: SamplerState, thermodynamic_state: ThermodynamicState, @@ -604,6 +609,8 @@ def _report( The current step of the simulation move. iteration : int The current iteration of the move sequence (i.e., how many times has this been called thus far). + elapsed_step : int + The total number of steps that have been taken in the simulation move. step+ nr_moves*iteration acceptance_probability : float The acceptance probability of the move. sampler_state : SamplerState @@ -791,6 +798,7 @@ def _report( self, step: int, iteration: int, + elapsed_step: int, acceptance_probability: float, sampler_state: SamplerState, thermodynamic_state: ThermodynamicState, @@ -804,6 +812,8 @@ def _report( The current step of the simulation move. iteration : int The current iteration of the move sequence (i.e., how many times has this been called thus far). + elapsed_step : int + The total number of steps that have been taken in the simulation move. step+ nr_moves*iteration acceptance_probability : float The acceptance probability of the move. sampler_state : SamplerState diff --git a/chiron/tests/test_testsystems.py b/chiron/tests/test_testsystems.py index 3506b77..c1f9401 100644 --- a/chiron/tests/test_testsystems.py +++ b/chiron/tests/test_testsystems.py @@ -207,3 +207,120 @@ 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 + + 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 + from chiron.utils import PRNG + 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, + ) + + PRNG.set_seed(1234) + + # define the sampler state + sampler_state = SamplerState( + x0=ideal_gas.positions, + current_PRNG_key=PRNG.get_random_key(), + 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 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, 1) + + from chiron.mcmc import ( + MetropolisDisplacementMove, + MonteCarloBarostatMove, + MoveSchedule, + MCMCSampler, + ) + + mc_displacement_move = MetropolisDisplacementMove( + stepsize=0.1, + n_steps=10, + reporter=reporter, + update_stepsize=True, + update_stepsize_frequency=100, + ) + + mc_barostat_move = MonteCarloBarostatMove( + volume_max_scale=0.2, + nr_of_moves=100, + reporter=reporter, + update_stepsize=True, + update_stepsize_frequency=100, + ) + move_set = MoveSchedule( + [ + ("MetropolisDisplacementMove", mc_displacement_move), + ("MonteCarloBarostatMove", mc_barostat_move), + ] + ) + + sampler = MCMCSampler(move_set, sampler_state, thermodynamic_state) + sampler.run(n_iterations=10, nbr_list=nbr_list) # how many times to repeat + + volume = reporter.get_property("volume") + + # 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, ideal_volume_std) + + volume_mean = jnp.mean(jnp.array(volume)) * unit.nanometer**3 + volume_std = jnp.std(jnp.array(volume)) * unit.nanometer**3 + + print(volume_mean, volume_std) + + 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 From 67a98313e6ee91e7d95a64f90568e8e416a8a85b Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Wed, 31 Jan 2024 12:36:59 -0800 Subject: [PATCH 18/43] Added ideal gas test written in the other PR. updated logic in langevin to take into account the iteration not just the current step; reporters now also log the elapsed_step. Updated the examples --- Examples/Idealgas.py | 27 +++++++++++++++++++-------- Examples/LJ_MCMC.py | 5 ++++- Examples/LJ_mcmove.py | 10 ++++------ Examples/methane_coords.npy | Bin 0 -> 26528 bytes chiron/mcmc.py | 2 ++ 5 files changed, 29 insertions(+), 15 deletions(-) create mode 100644 Examples/methane_coords.npy diff --git a/Examples/Idealgas.py b/Examples/Idealgas.py index 19f6427..9e307e5 100644 --- a/Examples/Idealgas.py +++ b/Examples/Idealgas.py @@ -2,9 +2,8 @@ from openmm import unit -# Use the LennardJonesFluid 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 -# The default LennardJonesFluid example considers the system to be Argon with 39.9 amu n_particles = 216 temperature = 298 * unit.kelvin @@ -17,7 +16,7 @@ from chiron.utils import PRNG import jax.numpy as jnp -# +# particles are non interacting cutoff = 0.0 * unit.nanometer ideal_gas_potential = IdealGasPotential(ideal_gas.topology) @@ -44,7 +43,7 @@ # 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 +# but will be used to appropriately wrap particles in space nbr_list = PairList(OrthogonalPeriodicSpace(), cutoff=cutoff) nbr_list.build_from_state(sampler_state) @@ -66,16 +65,27 @@ MCMCSampler, ) - +# initialize the displacement move mc_barostat_move = MonteCarloBarostatMove( volume_max_scale=0.2, - nr_of_moves=100, + nr_of_moves=10, reporter=reporter, update_stepsize=True, update_stepsize_frequency=100, ) + +# initialize the barostat move and the move schedule +metropolis_displacement_move = MetropolisDisplacementMove( + displacement_sigma=0.1 * unit.nanometer, + nr_of_moves=100, + update_stepsize=True, + update_stepsize_frequency=100, +) + +# define the move schedule move_set = MoveSchedule( [ + ("MetropolisDisplacementMove", metropolis_displacement_move), ("MonteCarloBarostatMove", mc_barostat_move), ] ) @@ -83,9 +93,10 @@ sampler = MCMCSampler(move_set, sampler_state, thermodynamic_state) sampler.run(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("step") +step = reporter.get_property("elapsed_step") + import matplotlib.pyplot as plt diff --git a/Examples/LJ_MCMC.py b/Examples/LJ_MCMC.py index 8c92d1f..c173352 100644 --- a/Examples/LJ_MCMC.py +++ b/Examples/LJ_MCMC.py @@ -17,6 +17,9 @@ 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. + positions = jnp.load("Examples/methane_coords.npy") * unit.nanometer box_vectors = jnp.array( @@ -139,7 +142,7 @@ move_set = MoveSchedule( [ - # ("LangevinDynamicsMove", langevin_dynamics_move), + ("LangevinDynamicsMove", langevin_dynamics_move), ("MetropolisDisplacementMove", mc_displacement_move), ("MonteCarloBarostatMove", mc_barostat_move), ] diff --git a/Examples/LJ_mcmove.py b/Examples/LJ_mcmove.py index 69c2859..c4ac913 100644 --- a/Examples/LJ_mcmove.py +++ b/Examples/LJ_mcmove.py @@ -76,13 +76,11 @@ stats = mc_move.statistics print(stats["n_accepted"] / stats["n_proposed"]) -import h5py -with h5py.File(filename, "r") as f: - acceptance_probability = f["acceptance_probability"][:] - displacement_sigma = f["displacement_sigma"][:] - potential_energy = f["potential_energy"][:] - step = f["step"][:] +acceptance_probability = reporter.get_property("acceptance_probability") +displacement_sigma = reporter.get_property("displacement_sigma") +potential_energy = reporter.get_property("potential_energy") +step = reporter.get_property("step") # plot the energy import matplotlib.pyplot as plt diff --git a/Examples/methane_coords.npy b/Examples/methane_coords.npy new file mode 100644 index 0000000000000000000000000000000000000000..b769491b98b3900c02c416ef91ee3d2a4d04a594 GIT binary patch literal 26528 zcmbTec|4Tw_dh-}X3Q97NQ<;+BZ{KZLh2k_$x^Zul{OWLNGg<+Xc3Vmm91#8R4Pf5 zQi;f(vSi=)eP6$q*ZY%xfB*dB;dak`U&}e?I_LR(p7GsosH1C4rJSXlmppRxw5hG6 z;zCKKV;dx87D^towzac8bn<|;?UAFC*S8$9Jbjd0KW%=<<|z4FYSk(^IhloPmdPx1 zTKNC(hpqCQD=~4Gg0^zFRrfOo=dEBmMp$t%E0g?a3t`7>^XqooYN2=_z?YZ9n!2u= zxV<=(&BXGwv!?aus5tKj!+u_8A6!o>&PiH42%^u!qxcp#fyA%;^(nDTyhMLvT70+% zEc|M%R~ho4uAgH3r(%A*e{b?99qhGilV3cc_0KcNxZjP_>Ql*c z_eb~aKe3j9{rxKY{NB_+=KNM3XKp`?zs&V^<6tQjCz7|0g_WY;xo*od;ICbr=tpTedM_mB8$^W) z)9_frw8Xt1YvJ#@1#b(CIxk)gmfD+T+&%l@o)yi^RhT(7{=fHCZm8IN=o1}vlBOFy zChN7q&7Zw|c|NFrZm6Ez$J*$h>lavcLqW~+$m(Yt91}DQc8nYV-Iax{K3^(9_Q;%W zyXloMJLYTkW#K-MxyMSb4dlT!GTzN@tt@={K_G5lYa6TyO;Y34r%tWczOluxjlEcS zDIQ&_cXKgVyL9<=986R0k zIw7zs`V0SV9t{2}}7*;pB<|>=f(BCsDE04@$lcRacsmJW8`ObDQp3Xm70jkYs z`~A&2A;-vTmr%$AtY24om%XbAb{KB?Ugz2g-D}?~XgPPoOTVMnO*i#Jf%nru#~LPf zUJIQN`;Z9{j5^EC6b?pxO9HS8Cuv0wjsOD9+PYXw%~n}fjb-=9OrHO^|T+-D2a}wP4^{C#qiA0sE{3 zJ+kihg5s!iebNY-SEy@^Dp|h)n$Zoh6)ZG)u#M@xj)Eto1IsOZD99X14wW9E;DYDh zXz$(I;M0Lw;%SmC5In+rxpUtT$gIlI+;fPDvIFTmFEv(A+0{9Y+QI%(VsG0dsv79s z(Eng0Ye5JDXUhF(ye-5+qpne-bE|n!Tx7sePq`J2hl&1Lv%Lyrf7V?wz(3%Wz@RoB z8;6F0t;<%6Pe50<-3zJ5JUCbxr!{w+gXf0MXS*=isOx1VYSYg`HfvypCWV9c^Aimh z?cm4OffPlFfs>_W2!#Z6<2JYBn(z8RIfidw?+!3P=Xv|Yxa!a}o&IP}QzJe~8_eqP+ z*v7^Wx-ptXSwmCv5BYjRC-M|8+SPwG(mKnF=BAfIGs!wPc%>}$l%Zl+YrZu5U?#lp z%(|398HLf1F!?%HF8;hh*Uuv3Neg^vvV0>4nLD4`t#~;Bn)_9vCa>R3|Is(?iXfUh zWvnt5;78U1%{R4+SolC{PMm;CD-@mcaw_R=h1I1Q`<5@{BE2~)VyJ`<-*-wWt`{By zFGYW|gxyrseP=Gr8^e#A-1bG9Rk2axz!8)6HVo``4&%(#6+r7H1_!j8Mxee+@a>wX z4Uo1p#7Mh}hL4x>N4C0EfU`rqd36dyjV7KqnQ~dfR_Z zTQvr1hcllR>G9xhtxlz7d-*_Z)Z6(}|y>1(;z>DQ~hmNn0>w-Bi?+1^pB5}Jo zfBf*Y0Z{DCyj&1e1sm2sk(B9TqMxv2kjv#hFkE)+;IZ@pXudIc*LWrsk7_;0&rfCG zw?fYq@A|;JGTKT;zq@0bL4UT-AT(@gimRYyHpkl&ek74JR4Iq-B-Tj?h zmpgwk$&9ZJnDS{6i?A22J3T%5)3p_N`2w{PU54Nc>)e+o2{b%?Cw!XIJ6`OJ@F+2& zGx4ay>;C+Xe{ucBktHY0#fISe;I$Q-$UYB0Z-`sMI}8gJZlCphF%Rm^61CrEM&h%@ z;3=<#-N2FC6g?~K57@j40hlAr8lA;pg|;f${{9`K`l?Hc(ZDnFj=%-Lmql8Kg8=1!L$ zu`p`$!^R(Eo(ze^>yOV3g64gA5*OJ9$A1Nv35gbg^Us5yChh-P-|BTbes7mHeOoY2w2QPYPatTw1jC%P{;fukCR+>xF$i z_ripA82Bvvo8n7J20l3E6IKzyGZoK+HF+;BZ?6Y_-yT8BExb4;neH07dl(qe&2{I< zx|uO@6>fx4(Y>T~+>EDp>Ks$uRFq^xs{w-<21Ba${jiYb{m^th6))&UtEK}B$7Jp- z9?B5FFiokVk~;zz8!kEbn^YwzD!F$3d^ZGorVYZesx80_xTTO&)Cy;+MR>fOM?mgj z*@cU1+Q3jJF7X+66q48#uWk|l|M_}C%alGz5YRJq?y!Yl9r+t7_h&BT%?&YtxElyi@DBK2v(N zzXKiLxTcz9|E6R9%xM(L%~wx9V~kj0rX#K#Q0Xi#X6m0 zwa*Qm;4*#grb)ZCG&Vh>dy9@0S1z0x-ZBm|eKMQp-{PBkPxo%&<~3baP&_xaUAvot z`qEGDmJjv9ozvO9@yl4~`C{G)@9F$Fr}^00=_(WyxhL`DWdxLFlP$5GKMJ=>R;gOy)j;0;lJlnNBdq+h3FP;Bq+ zi?Y3N;K!LG{a0AXifZZey;BQ}a34i#5+4qJITCfSfP(4zBcfMg>4@TaPKym!Gk;e=yzOZlZ@})v(XEC2Y`6wL)9cG$F zW$>eYrJk$%!CDY}k&&}|4juW#6z@i)Q*cdqioX0@8gAS^cfq-h{b1-L9UZ!&38Ykq zeO71Z!wNmyyAB?_SU0m~H9o?#d1)_5MlUWs@`j57DZ49kkr#VA zZZ+Im#*a6H6(76}sD%R*>AdjM&7eDB+*2bch>rT}c6K_9f}7}(QnBnVu$DgT6w$=N zrd1Mc4>xqeC9~hphdf$9J#*dVHH#^D`f1MbQxCf!H30l?>>LCc#WvS&5}#sJ*Gd=% z@Z-Ei5ubRj_kg5nszebV4JUTG3mnqo$9?6;Hg+Zs!*IbX%_efcR@B?T$`&f_v6f9y zSKwmvqEAb5=Jdm!OwmBIU0o0#w|wQAX*3LXE>pX{g@Uqf%&ZrGxmJk5sy(DTq=|(&O4g!F`9>hE zO7gg|Fa^KNa2LJ$W&pH3*vhfGT*NZGe6E@g_d1<1eDIr!DG!BLw_NH0={czzN|b1L z;NZx~1J8KSdYMPQ=95my^}b`hIERBHrf~)e=0gxH{xo?yy=uxHZgDm|aXCZ9HCM~8 zNfH0}TmKk*SNt4aO~>Fj@9RS?+Mv~U{oJef`7sEn)}@;`=v!>D@aa%LWN8T+6iN?6 zY4sT+rOpZHU%py#VR|nx_FkIiWi5z$fn^hxHY1>spdO{rJP0$qF{>m#9pqXaV?}Oq zaMj&UD)oE#(8f!1SMzojK4-eTy-D)WPTHl9*SmV*U1`+Xa>-7ZdGUboeETM-2>HF_ z>%o6@N9(i6fJy4Yo#~oo2S(tUMu0}_*ESe!J{T*!tOm|FIi|V#_CT##RC%uZD2NUQ z-kL|um&nBB0+*gT9}vVi z>KeuKq`n(}|1((f;t*8wxZPa3h=Kuri8)<|sMtl)(MkWLubQK(nbZz3_b$X{l~Hm2 zIP{daF>rl0wXFMtARegu)bGL=19k(Sq7aLLvZe;6`qw$g>5y_w`9MWMn*k|dlK1I# zH_zkAXX95|OEjki-+P*x(TtS1;>*$$OED!i-x7ewL?${GH*im1Ml- z+^=Kn?(<;XF9qHs&NOu1W>;B1J_0i$9aW{P1h7U>SNe_zsT=plvpo8G;qN+FM7cbp zbq+wol4VVYzK_GNCl79a%wglD=TypD+a4H=b*%P#GY*6MSF6#5|JjN9yoZMKNc|&~ z;rA^zgO0a~HNN*hAa&>w-`y(t?I8Bm6mA}+;nl@uIm&jqXsMMzs;)h#n`WW|au&~uRL&eQ* z6s{?`%w6_}4^MIz4K91v3HcdG$~2 zUBSk9u^fBLGkkb_J5RJ%YbAK=-ZB~U9fo;@a~?k#8>;?_!tlAm4dx_lhA z#=ds?zL$l8iP}7U^O~SAbjKt%tJWdtcD-lYO+l`(gp~;JB;_ zrd2}019P!tJuY5v9^ihD83nH$!u$Fe0vMGdJt3sbhkMqV%B+m!py$|Ji%A@zU?QtN z^A$h#^zL;lZ>@&N18dm6H4GHn&bX|&hk@_H+Y@LHSa|>B_rAwY9F&mrh?+^pmC?8^ zAa6g}5B6qTArjyABz!0j+fV};vs0dZAkS~Mata$=IWcw4ykQ4YR}TK$e_e}$Mn~qq zbLj86qB^52rQ~=6bguro#hAn&kyz^un-T<2TIkx_=j)rGO0D-qi&-UzKfBV4&w8tAW38}W*R7T>k{^$RXi0wG&WqokiKr^%{X563yD5nfD}a|7zoM-s zsJQ!rYE>+$j~-nC&N4F3AX~R{2^zyN{XFN`yQ@R6Oz_xA(UvY~$%|nbuVP^6gZ+vF zEFLU7RWNrVhk>!luUG7EU?MEcz4zvLJ-pA(y|JCt2d6P!>HPF65FZwb*)pvGv{NW6 zt;*U#ZMPGjjo~1~Co&s9DhXh3)ah4TnHI2owD;(~r{!>UQ(I>t*%uG$_fKv=h9E*p zM`0!z=L28AsBmJ>C4(l5r8LT+Nia((VG9ituKV#TE@h&1!>f6P{v14+~j2=ILBP*dQiYbyPwB`u4?dN z3actFb>{>GG+whjU(ZF^vCg|o7jf`+&BtGguX>@R=+C_1Vh*l8y;Qg>aRLO18x4oMX!ql(}4Vxcc6BvaXHTov}J4 zj|U6q3SBVs>VVx%zLz6NJz03AeZ8|e4WoR7Gpj2IUh?vd%GS#?eB*oUt}_{z{jAfe z_2l=%<2U^CSOefKHLfsAT>zgs8tlGj$3XrMzB|00SHeyA4Z%l#&~Rwwq^;}W8JfeCxs7JI8>d^3~7SH z``2|EXK`^3yGb#G;2qYmTyWkYF8Y7p(aR?FY<-m2FT=A;6in=xEq1;Vg5;tw=WMbG&r3snVVemclqiWmbDyUxM5t=$S z4BG|Snt7T$xJR^cZC-UhFlKnky%Ot(m<%T=M||cjiccSr*$IYE}gJY>zB;cMZ-|cN@*E(%ZHpjN2w=o4};Pvh5eW9 z*+@6k2=P`K1b5-hkNpm{!13YRY4#3-V7~B9);ymP=<=`(J{d~KFdx@jW&~#oYv=JU z@#DwLR6pbF36k$k7hZ77bpRq<0)NMsvCu^={r=WB6YwRrr!}Io8EjtgX6aXVL(n7D zF1x`Tc*oE>L-p#1j0PLl%Nl;H-?@=<_!Jcjuif`lJkSiwr!g zKF^5^v-e-y4%R)X+YX%+z$@l@F_aTLxI-iL+cTd&*uATEQ84l6|8QB7kNoy{_3lvz zc~5OYwp?63ls^=Clrw(-toYqNM%`c{(|(@rC1C-)(EU>P1i9b*x@hj_NIFI>z8qb} z6u?7~D=Uxk61x^45v-VjPu@kfCfYJEd3}J)22Wm8bAMt~lFhBbz+uR@a9ZiBCNs%e9a4JYNMne|*cHw7>ZkQeSDJ z96WQP=3do7Djr@Q_5&4p(9a_DRvP)g%^pR$43j<>J2`W~Fu{-h`u#)AEHQI8CT86& zYe$0HKi%y8;q*HSRs{Nozcn3%EXs!{qn-RXH(fUCBYDolx38O?pJk!b_0gujBo1Ea zm?6^IG6WOq(t00AzPpg`K~kt7!S|xVRhUC%o84@+1d}hjp0MbS<4hMJs7z6*2Cw1)-0UwaQ;N%BRXCf^kfeZ`;`lP%y(&4 z3AkoyKDsE?14k=nexBIHMwu6tR({dV;Jz^XE1yRLM1^F+a6cW>)c54X74RaOoohSy zoDcm4m47wH@#CWlK`G}W`=;`QJ5Ial6MMTs(u)@A=tai{E6cNw_7FV2@hmP`+z9(* zMvk0cB8Vpiv~(tMsr@UfH?AW2$EKg1`P~s5Jm?hgTz8@x8pXcLl{t>U^)Z7V9xEmw z=|G}|!7Kr6aq!8xD$)o+YwZvC|Dd3plzq$u$-^C}fwR|Z((%P$)%lVSxnQF3d!Z!3 zkL3BszkjUbL6^G~DJ{RcV1A8kzT6@zx*WJjyE4JXm5#Iz2Y=+j?0ls^;}4iP0pVezJ-`=|}2@|Ks2ES~;fY`YGsI;Ik&dgdYculMhH8_6$zSv*Aj2WfN8kWIwsZ&8_>eff_Ku;Gz9&EK4xHFMU!fm_3S2d+&J4rFNVSV$ zA`P&8bykjfYQxm~{=PTnzKTl61_~1lRY3tNFdu`y;JiBTf%EWmu`x9Pb)+FD+ zsDe=M;P_Xf6zrgn#GKZzhPkuCU&QSmf)aDh(ZEFl_-A_hsK#3gUK43*UPEvSBfBjf zF%ExV#nbIO!C?R_qsoq+QRc^n3uYGIFY=?lVweXly9QckkA_8t|I4q^_iMgunJt6^ z6iJEgM#E4ZL~%H}g6Ds@2C3_GwkSp%e?~#0#w$lHXY-++^I<+AixFVHTe$bmgaF?D zBzX9XDji3bHGNpb;>EHhfzkEGqfn-Q#Zhe+$^ZV3kNi~CTdV8FMvVCxPB*HD!?xp> zC(lJevx{q)Muc~4{^OfxOYmT+mi5!W(os$$bVTI&7$~Lvd2F6s2hSF-7fXw32S4ul zEoqL0@OW!*M(hVZ%&T4%@h605>OB&wO~*<^TVPYiiS!)@1aS1*G@%V!nRw6EPV3Mc z!cR3B#|dqtp&ivB`Rb}hSb|Dao`HHe-yF7D(4q=#t(UO#Pql+F<;AB{-VgF8vX41tvoW8~#{2xoI_5?&hcqJ-2% zkAxqj?w1)bvFxjaX)j+r;k+hwpnJUQ>MLw~e&Wr%#)sLE^-MEOd8Ghus(EF&;T=Ef z&N5{jai!o6zU$lGc{7kA9IzrhYyiT{Vw+DN<;N99y?1@iw!o+3hEIBwDuH$6X!DUB zRCLrTf9xPI0p~AmFn-g_gWhEad7PX^!J_<>`>on4csgH0Dour;CT zc(8Q)-X$)tn3(eHP_F(-77ACrPnuZX2fERXf|GW!_LuL_)`?*lxFoyLQH6>_HsaEA z-maXRsXX5U%V$SUwT|j?HomM@Ei4kHc%d)j;=rECDRlKDhX8iFz zYrL9)6^0f&hqI{op<8fTn{gi;RdKni7}5h5=Jx-d+c^R*i)XEBy+pyJyPr>MlKA}B zZ(o!ib8~va!WnGKW#^Cef$Wm#X)m=qz&%yms`t+@Y^_RpI9V?w+77~C&|)QSbIQXL1H35y3d%Q z+fDMYzx(2>r{Q@-k%B1;hbLUJho|Dj-*N3pCj`S*5|^IMDl&b|z=-7+gX_Y_VTJM* zbPMN4@$=b|O0WOrrGGy+GdlX|$S*GTOLiqpm*K(OS5mi|;)bAM?}FcRAM<0n$G*~* zdsHOZ(?`jXCYVzp5)eh^*BEy(bQZ~uN>1XViObwaLyLIODr07O zoeBeezSoc+jS=w z6$jH%X`G)kd4eW0N-N9)`BFXh5RFVlU<;SAqj~;jT zw?W96!X;vg-C%bz{owg|KVW9#$<+*ko5n58w6IYY!oT@e%;v9n!Ht2_c;3A^Bi#)v z_U$^f2HU}l2EBjF8JWY4be<-<=IE%~{+d9ZoT-KB5IycQOyvyUF4PvuSL+%_Af zKH*1G)l%i#`V8D~tk^gCSP>-k;+Z6_0D3QJ;65|x1Wn`IZ=EhY&(F0AeU!?KByAZdxci7{YdmFTXI&|ed_~y($YN}WaA&60SKaJP-hu5f|$1K z@6xvnT=L?1(@m}rT4xF56@DQ)k^?APnn3cAd2S^V3tJ&P{Ft$Fa5c2cm0E|JF|gLX zeqWs$AL2B*13b*Ssrca$8Fe*wCl5Y7IuVdc_~xh6nyeNtYk~Y{mJ7O5iGHWxxoU|P z3m5qQUiN^jlMdxgkfM4YIL|%&p`OHv|Kd$G$tQcKXP-=BVu2*jPUVhK@No;dGXD&z z-@6s&(PS8S?Qm5IFh)V@^}=%3QU*pPlpc&QW#hgo+8Jx2FY-$W@uY6-(EGRw`CQ zR$cUsD55(E)=Q|JZq)_JZqmEH&Je`WG!;A6yk6*fZ4)kRM47VBAa#Z6l08k};~w|o znjyjeE5&V($Z`!Q z1FyRd-c%Ptqno{3lPq~qWMl47#U&nmV7K_Q!J0yNewXq(6}T8-^0ak_Q!`xG`<+qo zj)QnSVy8xT^VB{XsGe20WZVO$H+^4u?-&H{#TOLD94WXvZ^3+9aG=!W)+Xn}iqbo}(=)uC=_D!RS=bi0n^mvX22o-ZuyfW0j@T_^pSz3b-YUjt-c zLNbo;xWR`R5vx5CyC|54GJy(@dSG@G^OU>(7%)C=4O222h5zQawGS@$Y`Qpk&U^8B z=%$f%crET;vURWpd^Fa*zIu$*`*PQOo4dwA?{)pOwnrS?y>e2iF$(^n1|xQ(OgtE* zFaCZp$pa_(x(75IYd8CXRFW5jRx&+^9nNftc|9$B6fTK2WQM=O3ZabK)| zz$l3m50u`W4!3TqDLT?jMvl) znA%r)@yhxMcCUC92uo`e?YrL&Y4+uXKA}wX;Y_F{6Fm3KUH8V-HiUmEYKxtmCls={ z4jdsAdiD9sHZs1~FT0{<6FgVFBsDa`@&|OyiJBek(gW_#N=nS`_d{fa;SS~FbUf{l zq<-oh4<;{7^?ZAD9LCr+?}F2Ep@xw*cL#~nYtOD+wNaObnpf8Ot#zvcfdrO{`ROiD z?^e84NZuoo`!XqUISu1|%pwAb-X-(+j&ubo(S_vbu8$`3cFq~^m!3}IL{Ngvq@KWR zYgEUmFc+=Y*aXs-voJyQWOBoL#*{zkF3&AHvWJCgTIQv7mSgZo(f2mjy%!=rZW3LQ z*8=B8o~_nj&Vv_1z{-Hcm073!UDAzs@pbsAhzbcVE|(4{+)HQU8PV`LfqI1R+&|}x z=}TTLa?Bm8_xi_k8p&<^b%OAax#61wy-hr(=Y$!#2vR3Cf#Ys*bRbZw13bY^lN zmW_TC-!lNms$UIXJURrwxoeXXTUeM~vFyg}IeaL2uii{PoQuk1^YIM9n<_tYgVxZP z$dzg3PT~mKbcVteO+lQql+XTpUnfjg3~}q3px_?0gyNLVY_zT_IINYy#2C98m!fhS z+Lk}!1}G20HHYTbPmx3L?0A;?YBCPP$17fRiI#(wU((~C15D(HvbssUL2jMtwT6uZ zS1Woe6-(`b9+HlVYA|qgg_hfwt~!`W(m;83ADDe$z0l!u5UzzNcoQDm(&zZaEwLqV z+9!y%!@CO($%$tgAL)fX;bLKV^Qn0A?2i=FwN#WcSlY5wk`HeQ6;x0Q`e2&2GrpSW z1+V>Izf81_LC5u&xqIz8VbPA#jc<-K@W6NZ+9pzuu5_(;og+xamxWQ&ucirN?yBEy zVPfrY_xr(xrs>15<8~u^ZxP`Wq^uI>u!kTh|7CUqjTdbM_Bn{(W1!j$-}T8Jy^tO^ zI#=-aZ+IHq=X$BI6@IFz{r+@`hUd5cx_g4mf1ae9gV*$KaQ*PrM}+VIBubX6Z|{UT zvXbPa;GpVnJ~98UE*N}wsp^S$CCum*aZDK$#2>zK0T-`O(LZx!TWT~D57Sizmdg#Mzxu`2xd z{o5R@{rS9cp(q#2zC?bzqsB!U{vT>@@1{e7g{pUljNFmcNx{tegRQ5DzN*g4 zNcqcEE~YszHJhBzKvcz@;#W9VjX|kr zuThhD9mG8TV`baJ!FUDN%BS934EUH@`!%c_N^h(8ts=bai>I$<=aRTyP*phYI!$P5 z-i*3w>uDsOH_i*l6d-t*nu=za4t)sJc38JecN>F8+h%_V5o02Y8e$wzbdn2JDzy6% zeqH*Ywi=%rsT*Wfz9@T)!ZS7}U zdw=!`;yS;Qo95EIIPu^Hf1nZ*drsz^yeGwjZW3a;>*@xk{Q1J~wR^a48X!Q{LC5$h z!7;L&$5*NLz$5+Ln}5WP0(00f<3d+AsQS4t@_XF{9SdHH#C#lr$8$fhU1t+rlcuSl zM(`+HF1XA*M9y7Xw|SIVBoEBZ3ct4H11~cE0~o2eVeZvtWe+COQ$uxPB}T!3#xzds zCVcIm#OSZ&+_L-gaqHOHVOS|FPTAH?$EQ5w{)(IE=x+MXKf#rTTesVc$SQYDvyjNROo3B`A)X3HXp4nT2poS`2N&QRAV%jOc? zcAuYo$+sSWClmcw_p#ALD(d^{TQy*n+Z}L<%Y*;br_0;nYmWG?mdGwJyV`i7>k*0L z13z49%?VHRW$6XVybU*~vx^KIQ-Y)onLc;|7$eIprC8^c6+@OFC7@(XD^*xN{mo?Z@)-XI0?VFA4PG4r#N7$3F}n$mjX07#N! z)J>%Y>~mY+$Ey&YNi3;qvk&3v{x2V~sHr)5VoM*~%>5YQOZ4w@r3rHRH_IUKW22JY zc^dY~JeTGlVWYu)so?N*UZm^_+u9I61Q{y^rdP~qhZ^~&y*G&-!`#j3j>^a{$dr`6 z8Xe2QjjPt`pO9_@0X_e{6Fr@O_CMBt40Je*g~NcXHTg&AY&R{{+2dbI%8=5`EU>G&N~B>@rJOZqG0j3~ANJ zZso;UgR*g&zj&wowc)+uH12K=zP$W#v+HXPnnhm@s-IpDpW{PSjxYz{lgDhCn}gl( zYuCD@U-AS`lz%z=UASW^56$>IZQH$}Ziqjhu;W(P7`)xM=O)8|iZVr8l#1swQOM?b z;DFiy{D=4N`?oG0(U(?DBmDwIZ7MjuM*s(A#;MnkbMs@0(USvVG*o*3c3)U%D|lBv zymBLhjY8KKU0JGM3@R!LXZwiW$HSzRd3Dn$R1Q7(VrfEji>01Bge{3~ba1IV`&)qx zS_D`8WIBIMpK}+y^V{yc)1(4oEeAfx5ZwDO9u%{+%5la=1`f7F%1zU30*2F?^Ud8X z{GsWqmuo^r%UiRn$E;`=JJYLv3B3-&w8vdGG_!Hj_G@&N&qG9407_4Y&fmQG%*|xN zL;rnlQ24D|(YGo1y;4MHK3V_ezYiQ;!aEMh_4AablfO-Zf?pLlbFi;g_sqpgesq*d zK4u-t!SjE*elW;>D3PpGZw)Wf5{?)7*^K>ff$W8HE&f`~`6mH?`M^S-z4M!pvtKoJpN* z^8Niub#X+eWxNsViz#Tb*3RfE!F|k1WrvRv9n$^+-;i4u7&yOedSElbZ-@4h!j9nA z-ugLZeAB8xWvAB?k27>!?5#D=;XysbiG9?bQAKnO%vOtT`Cceq6dqAz(glzFh{7a( zWNIE3YqqNz9QwyU__zz14?iJ#$3xjpS1vU{GUOc4Ko$OC5Md+_yM)gCBn!DXZOQP4s8kfdy6h6kIh~_pR@TWo?{~a%6ut|SJ@xZMq9egO~fND~}rR+E<_8K3t6 zW5li_cGRKKqJ; zPlShV?0+~2vwg~Doh9{T(w=Qs&-qbMH~;k`A5tf7<5_u;?*|JD{9Z&woEQV;!jG5s zZxF)ZUr#?5>t#X6iw*TxUk<{KuIYzwEUt#DXKKoB-=?GQ_OX&N zSs}}wJC}%lyrT6R@8`+*m21h_yR8c* zgtfgbo9Gyx$$Rezod;(gy<`ZKKG;}TA2*rjL_SWK`GwsJmM8C&+FH}GZ^Zc8_>nUB zx>@VtrtdUl8Xqj46-xBDDuWVDK@TkX?&!&wF)XH7JqfyEr9>aGuTt- zbnX5f%X8-wouT*&t>ri9n2Tqn43wKd==8Q7`-)j`v=009aS?@z1zi) z=H8O+^XD{y#-7S5nYjY^;6t&TLB|j*$z|NPJJr>%Tki>w|r4jgbYp4E#N( zW6e($;_q|cu9);2LCv1UCvS{G+Nq6Rb;Pdrd>8gmSV_SnerEf*BTSreH}X-W0s{qz zT;aIqzx7&v{z-S92?Lkf1*8~LnYiV%O3)_WL3sRFM{dB5i$%M}_Dc|or_M3clf*@`}yM)uY!UV1s#tUcz!19@?W11 zl^<8$)i;^@tqyj`8@ELho;GGB?c+LLeyos=OjmK|$3&kS6zz+hAaGCmipAG4Xkq4R z#?7H%)W}|w$klXI4Lmon_jfnwrYycXe+%iuNK&>)C-3vwklPrmF%By7^l3f)tyBIr z!1Hx~)^8eK$TfQ4B1Ol2C(0$q)(lVmpIH#Nv*#!aX}u%kSwC7KenObW<5LCiuQ*la zNU-ts!E3)|vp87BBl=`-U;)%dhW*@L#6;Wb==tB8h>mRFjYIxm9t4~k+_vY=I9xk) zmT{ZJbMNTZS*7lTuaOCwc)hC`WDLh1Kj0;LDSZivyXq9IYgL(D+w@+Y8sUA=zToF= zgdZoed-*$5Bx)@|!O|g^SLEhsc(@;EW!aqsrNfA%PwIr!Wolk^?vIDPwcgo5Afv!WF@3ZlVe!_s3p>{q1U<*yy2c{MZ%)2o3otQI0p`iQ6O%98|+a6?&BaIoqx(ZQ2_I9Y_Du+>Vy3IZ5@&r3=}^r{8mlB0aQNli7p{JvpUbR9fe1# zU`4!Ofh*BRq-g~1_(J+&eoja|0ACv3{A@llx~~^TEmvx&5#HnP`~&7}sb5F*Z3)!S z1zXJd@R3xNzUDSMQYIS&PchIfUqmJQ1Q#_l+J|RtC3yCO2Q!}U>jLw9>Z9BDDfr1u zTxE$p7aPT_&RHmMaPOhv{i+le&b;rugGus26VI*m&BTvI($;6Z#3~qA?%Cs$K|>wu zk_HdcUQl}RPJs52{U{>=6D{#f)ERx|DX`*SUoOzZ;oAOD*l_)~`#NSD zh`(Dml&{N!MQ797-vy*Y?w6P_U81wENqAoK!fpr*c#kjINBrBK#@k-^lOKKdg}YxX z?SK>0c50W~5Is^tM!F;6r@|f)VByFFXv-*u4~P)GXR6lv8<%UQ>}J(#zF)TccrlOK zE~cJO$5w`+afW#nyft~fv3=&h@%dT5O<1+54d#6EzrCQJ)bTrSy?gtX^S}OGYCYuj zm{4-#Nq@nE-sOc4IOse$)bTlxhURIzwckdwP{eN42aUpEFp}|dewoFOTd}hA{o7WU z^L*U&f^Ro0=(}K>K=`Xg%bYF*3w8tZZAbq%O%CdnJE_sM1#t5A{)5HfAURy9)60u) z-gng}^B{w4k@K<_8$s^&yx|S56R^e^xvIM5$dR*;T=$*d{6d;KDLi$)I}dne^UJ6R|* z)4*U%lLwUpX38d#aTjTSEUrH`1iqry(n)6tK98>yQTq*AQLOW}X}J_7Gd zzM5^hUk;KhCmW;*|8T20rb3QBL1+bO5F$G7|JG5z3yPX{ z(w-6i=ESpDp`T>^)XZw~RaAL#_nYz=Bgl_!^JbLG+O&hfWFMw575~GrGrNHmvdZOx zGpQ5yiiSk|CjG+O7OTG?^D>UV>6pKUg6|3ZQBwA=PbwpdwO-V)1 zxa8d*IM`7Q2501BH>Q_CuK-DaZ*uYQ#_T0)^$8z4c7NyHQYOx0^4QV+8$tE7if4n- zFK|t|u}FdVO}@V=$M!W18|9B#z33|gnN2r7p3)$FWe4}(e@prar+t^WrIgzRO69er z)3^@~2M#%!a!8+3Tl9yOS%mjmJRthPfKBRzU={s8gH!%@SK(cb8>xf;>qiwCfp|++ zldDz_1m1f5CftdI2AA8H>XUwL&(KBZEWHSS_sUsSC!K>b@+Vx2C;s_CyrhQZxttNO zbp9f5_KuD!!aJV^^o+uP?eho6z}(tDbJuA$R`06Uqdp_`NBOX!#4I6PN7^8oR4c*r zOk#v2>7$!JyKZy25EHMBlfqY?=s539X2$98`8HnBar!rIb+L04M>%4r{mD*0qY5F{ zOZxulZyY=^(`c^%AE`%kR#psq6Mm+B-{~;>F|hW~HytE>(HaA@cRH&Oy{@9;?z`li zdVaDmsJDFRf4X*O3O;wz(!81S2aM%|g9To6LuwxXiifM{Xj$QtzxdQY+~M!J`Si!( z!#$Td__ch2rXbmGuQlCb0tBDv+GG_XpH0J-kMW469hPBe`lx0Ui%VhPf+cMoS zJEU8c@?!*EdCzE_Gn|LM|_v#tr~E+GByhL5kC zrF27_vF>S|pn6z)McTUWj7Jh0Vjl7qBLss+rlh@YqU}vQy_d!J)Jah6d zd|buFH{DC;Ob|PKd|e>gKD-5X`&>ODO!j~Fx*69S4|T(=y#wllZB+bT>oD^pk04$* zkvOh#mLK0VXP><8(FuCS`*|!$KfA(Y!8J_yZy~kJKaySl`~Htb#9ngqqVC9ZvHZZ@ zr`aE?L07k+hV^v-Y6lCKoM{3`zTLozm++EC6>Qr>w@6<%C#IxDpXgaficU_} z=Uy8hU+IhCqEST2^hy2F2mX$%jNAY4U?o&O-Qe2^D^|HB!5h*)8^8Nz@40`vN!aw|{c79>u(Iu&i@*G=%oir}Kd4L~ncZKqr{V0a@-Le}LRudiZ zM7w?a>@FCI$aPOxTMbHQqgD139rm|Hb2UUro>-tB6Mb@+^v9mkHS=YWzEXG9R+@nz zddw^-pMH?&cng(`g#7q$ex#|Q=w&9hz6k4kLU?FrtFhkDUkv;XMcF}Qy#qgbUZF}3xi4t9+!aiMB zhc~2dExq-=;Xe8PiwE&v7#5u{u;3mycstPxr#{`Ar#uE8scY|P3-!ZjwVwZ5()VJK z;PhZ=F9nrE9lvm%zOciI=&IdfV$P8Mqv(FM)!;+;-LYq?&WqVN_OheEij32B zTUQs~JwALEf5OAryIHxj^S%@3brt)`ON*YGrO#tOiXN5>U(GxhY@tF08 z^z}_&QP8)bZ|Zy|Rpjk{dlQjyXNKdE$~p1 zeJ{6|j_c2lWV;g{+<&e5ACLA4@cX=zbO3e$B~MbPyoDFT>0z1Cf;{-8)gWftRxS?r zsan3&9fmg);qLl4KK!pwB94J0lGBAJ=RMxw*lJJu$eo-Oj^}O}h1cJY`fTl|p}fkC z%JkV>RGYu+U0ux>Ov;lJ;=5t%6Hj?s^gp}Xt&`e7=OOVtpxDEn@YjAs_H+OAzrHm? zKkBZrmpmw4@_9vy7x}&ix9fG4;Y}c~Z54Ut8R_HXW~isyjX>C{_CxOE`+#gG0WE@m zeE4XQe4d=IMYi^OHvg}vGmnS5`{Fob8D=b%ib|zLv`MH$b#6sU3ZbNstXYy|DQl6G zJVZ#@Qt^bchDxhEBt=r;AxrjM3^N#hXL^3}r`KyRbKQHs=YG!T^VUo!U-;*{QomD( zfdFGiyaFc6b8{njPEyE~Ox5rnjb6y*Q`z@0MS#S5F3mPv%1^GKtSW=&YMV~-wr9Hh zB+h^R=bR}fG#-|3TvQ`WMyzks^H48vJE3{nq^%LWSNe{4n;^%>VD}YEi6;1^*_+@H zIszLOG1LA*e^S|)tarOYFTgC@iYB`m{?Mz}RGN7243~uZA`F}}*|7D;m!QxS{A9>7 zTZvnz8Tt(>j?&RDfGO$X>v;c-V-|6>r1KKC?gqU+tQ)pD^n!g@R|Hw=j|qNV@cPb5 ztt-4_)~DpZ3{*skl9*8Y!c{{M@8^EH+a38HRSR!k7!V=CD95gIL?7f16_-}UpKyJ? zdd^$?{e`n1rP!UVVf2!&O&?eV2qFPwGgM>Nh9k zuPmk{KKB`d!>Z>sJm<|KbzhdIKCm5z$DDZQ13sc3%%fkwx9`vXT|Mw{G~JquY>7@@ zoqw8_)G13OtmqhoT^ygrS>zNAt9ky`7o{dst2>(^g?p>`ramUTUdHOTN)#pJ{g=O8 zJ!k6S`E36F)Jh(rYHbj~@m(AV*BVm9bGPIBp*!IU{N#!1eF-VdZ%jWIA8(jbdsdE{ zXpdapV})q>x#s`@xcr)XSyHGYTa+5HzGnbf(185;m>bU^>A+t79bBE z>Vq_%yR=8urmpGyWLKMGyY(wB@|)AOi3NL8jVg6N^~1jBpNl=^xruR2%6q0{FWAZ3 zcs^675ZY2M8IJDdpU$GZlj9>71A6PfF2nUZbW$OsbLRa{zYis$Q*jPcKjBqR$NN1S zry%b`i7hYkvP$0MhO82pf&5pw_z6eKP4JGj*jf=Xv(STYri2a@4c;mx*Y-F9Q(MiSlsk0;y87ovG^u z$=2D1Hy!q4Jw{xBvwS~1v~D}?Ij0p=E#}U9o{Z=5wYi*WEZ(#KfA2{cquvl=bwZ=i zULh@aKJq$xsC_f;v!i}7F?&q0&Z0|{%oQP_`}TZ#yQCjJHr_NLH+x~1S*od8ZWlzE zZ7H6s(GBw=Y*xlz=OH;)(pK7_?pER8E-!(3fS#;_`%juq!Pih@krTOnP_+Nz{No(o ziKhP3+w+CVhTHycx*rZfz5pt@#=T%-Ip;R#a|X+X>FHE;Zrm4bezme>x{E!f|vKfbn#f6F)s z^F$smT0WyYEfZDjn)Me0q!JFVtLYhnR<-A})jmHV@aTE5;09q*c8(=A4|AJq^0R3p zrRe8-`L;d>{YJn2N_fsn3a_j9ucZ=WGkVY+%m+-bXOo4ildIJz+|#@Jp;evfqen1E&n#_{pT{!nNO8xk)p>r?&(*AFtr+ zDmmAI_xbOR7frnnagkDIucacD6JQjQ`MyG&kNh~iG^Y^P!|X`kiKT`(Cuo{d>_`1E zhj#(@Jds*R@9m9jIf1_L-fDBNIAJ0fS@Vq#^^!4#wCx%jc*)=AGA7H-z5wOLhlEY| zxkgIlm2VzHPJ!K!kpK;FlScu#2xVgM;n@V_C}Q9q^B@k~ zcNqh%&G010B7U}O51g#@4?8_TBPz-F)-(G={>$0B;Qg#Kt z^9QjG6<-kDW`6${aMd@xW+?PQ?DfV?aV9(@Vpu*$g@wKoj!;+6OBO%KYcZ~+lQ|I+ zJ1TR>q4m?C>*@=`ux<1f8fqi`NTn_yq-X>h6WL zVveI?1J~}Nevn$#v}*~@ImW1)^*H|U5@|)hSZX<}Kjo3aA}q*jLIWVbAYoXYFI@gy zlqg?~Szx%l?@ygsk1>UxILB+_vNGq?)#>%6{3^Z8<{$d=I!UyDe?LSeLPAIDyPhIH z?d7)zdpgFzg>Un&+xEgF|I+bWs{80fEabwDr*+Mc@AKNE0J$6|elhIs-fn}2CRZCn zNdeNXs4l$_U-y0R;BY$jwQdH7ui9Lo5vK)p9>!RAN?r6y?JUQjTHRGMEt5tJTE9)* zyxaf+r}t$q!TfVWc=Ha9{$-vk7&F@zxd7-c1=MMJ4hX&Hocp@Z`?&VHjRI$Er8G<@ zPS+II`o{78(alc24ofxELSNOyJMLl3S+AL^D0Iq1khs<8U$V7fLYdm^LwSl{;TiRN zn3}XO*&6iGOBUy*T)v0y?TH$NuxHH8hK0y?3}%&ne1krL`3UZF#@ww5sajq<0_g>? z)l#qzp4cQwwOmA>+UY~<=EyNY#wNT_hK+d+OUKpibo2otHt6;mVg2%lVd}COZf1!U8d=xj z+$>yB@MKCioCuTt)VgyVayKN1_ibasG^cX9&Is?TtgVTefE#%248-Tqqu7ty^PTBr z3t75oOX@f@6kPHiSUUw>7{hFD@BY(=yyE|B-kx_ex>bTphqmiq6aV>d(Wkz9h2mMK zbo4We-7Wrs{=Q^OBi*gX`(Z_A%i%-lANMts))+wE!}Z46mi<2`!R3hSn?{--X|vxi zH|wt<(8xWQr=f}S1hjd@Q0V0I?vkqfIqhJS-LO&|bHQc7icyv|RI>M|i;%Z28{+aZ zYi@M$66N4~vH@6sHkh*>HhvQ&Zh=14mG_vCj5a;JAKj2~;q^8PTLx@bd3E*^o-ft> zHDG$G27dR=Cy+M1bHlWV60j3xMOf29Hi5@~*?AO!W&7SzvL| z@|Y(RGCu4-+Vq(P_Lq2`IHN!53dY#|Gmr-vcj5jUhZ+9D^m?zp?sh}(0p{bgEjr`Q z41+KCrx<7#Cbx`tw(Pja0^fxKVxb8%5;^bu-uAPDP;0f|7RR6bNQM8EO;-bCa_rL{ zBk+@Ltmb=*hp6l>^A+i%6YliOqW>cQZ~wu%#<9Z`viI%FShcxaBnGikoLu0lPR6Aa zrCI0$AJB2c{Yv?`(N@4+h!pVU8I=vs=pQ|T$8DH8{6xgp_*V;Z$grxs39g#aS3Q-2 z_#1BX{;3ZG<>OCM_po3`pkB<;osGcU%&rNUTLlXh3cKUWCSmNXK}gpc%t4v0x>JDu zq54Z)=Z0+95Tle;;hfO``dOpf)jI{rigBqAmyo-czQrSRXXw z#O4YS4UTR79Ro~b&B@5Z0a&%&+UgC*f4(&GtU2;5Y(3|dhw#x!l~br)oCzDkbuA^_ zc$#7M!&BlM-ISt{tXJraoHLo{(Vt~dhYk*OUV{3vGXLXkVd&r6eoXgfL@1T;C~?tR zWQXDYS+OTd$hi}7p*-tnpZ)Z4pSDmQ??V|_2i+N3fon3sV?-U=&sBaGN zuNr|A>zF39O?Z!Ql3G^Uj@&N0y9-+#$3dNzewzO&ofJvY*Ztef0twOntr@E*xfG&LfD|w{HalMP>59hi1c<(8wG)b z#gD?(s(>lJ(|py94!FNQwrFRy5J}2CAJ|=woPsxOU+M@C@pxvMvO^Mc-2d#xFMsH!_OHDoeA&ZBciw!I za2)+M>)uCvEuW8{=l`75$$*wucdoLadV$Jk!RsPqf`4RXYk&dZG5*J`{|N1O|7Wgh4Z*>I=^NV z)HwC1e+&kc)ius`;@oDwcw!u$e?z)PH&2TQdBtRpRN&XspOZl|@L2wI67pWuhe@KY za{TPYGbu(wB-nhVhdAPEo|D2DUtM!4(8j~Q|no_vvXD7U#Fn#5eEkLA~PXr|(PlWo` zwEiuA-?p5Z#Q7~0QaSN$sCFaz9JeeDX88ybzg8-VhK_)wo?*XMm~!9T#81U}vT!HsYnK0Ik zvN4byQ=pDUkHTeDvoAM)U_SF(vmO`v<$1egPudG{leIcQZ#aEwX@ZYTC1nb(O1Jho z3-OY3I_|GtpXDNRuOJB8ScK^8+f|&cGX=bUQvyq#^?}X@(c`V>`H7aHtK@C;JKVVK zx{=dA40c(Ca_)5*D3etuGa%7FI@gD99M(7w1$Ru2z(9pnkdaP3+-AIJil#F_!`rYl zI2P-^t7F2VO)aoF8HKPsUPAw-eCIF1gc4)_)ujyNgGCB2iyasNYl}^kNbUw;Z7H%C&)Hm0@KFAMyYbK&0NQCxCyHIeSPWN%9+m$>^7@mPMyvW1$SE~&ef+v+mSB@^h=P-xM<;PDN5|w+*Rrm=*?nlVK zm^%$|Dk{=t(aCHl=QRtYD5PQ4mR3vjqlJ{Xt#(;WC;C_phl=@$4&RnXj>z$oHAx?H z!1E-G*n@YIjQ{Q{$X)FheL&TI%S}{AZJz7d&2WSC-#>^vG%*398%>5%J_(VpZ!)s% zW8Q=PdYOqwdwI!04la^13|Bl{X*{@J7pOK#D@X7X^@{DGW@7^|r>}wkAN2DZzFZt7 z6EMR;`r4$BKWdBT=^bs(W-8`o7Buzl+}rgZSG<>(#K+%C{kk5{XLeDoKk^TQ26ak} zwl>2G{xhe#H_vb$T2|$Z1U64X=mSw>omUK)994DX@GO3FDV>o2AuGIzh8)K|9tEKy z+D+hp4?@yLX{36qcqp|%h}^0VzjC>hkK`VD_(dUiZ5Kx_R+GmFAJ=Gsb`-=X-^u5V~o3H+Wx8BiPPw#4kuLu1M zd*<@hg??X@#E1+aujksY|r!8$GV|l?~!=36C;24 zvT4-^HAhc%!uop}yJj!rAtthXmmZW0kQ&4{J=}x2=FIWV+xOXUq7V&=xmb5*y-K^i z6MaxFee2|FMMy>pPmxkJ&L4aIuDmmBf<#vfT`!?J_*?$i5vL7o;Kag>eWu{IPt2G` z5*kmc3-gbHp*IHC=#$Vq5%t8xYY;A;_2v7zpNpK?p1ydaBj)eAO0Ca(kAtOpp@c@a z2)Q3>akxLD6?%di97T;=pvYa;@{y+qdF$^w#ERu7BjWN~IC-w3)HH`{#f=bO>E<9^ zOCeF)cGss_jlqgsz6c85CvGbFXxZuUkf^(j68XDm#Et&`G<_}RGteF@{tpZEGB>8L zkVE~S!%sdx4icx+pKqw20PF&__R4b+*Q+7JBA&?QV3mx1``iPT`^SUSnG?{Ky-l?5 zG0s1mOO+V&8o-yCelmGdm_&zT!928 ztZ%_lVn=87^OC@GSCv$e6EMB6r|TvMJ9lkz4A$ixHH){L0(Vc3?+F_iP>uG_e9Zrk zZ6E(Py$t76H^Su-5B0*+cKzU|xu~B@L{tgx9r}M>gNu9%Iq++R9Off27RQr>TxX+h z`42o3pnb)3Tx`Aw(Xx;a&BXpx?9?KkMC9_dvED9RyoaA8+tukFifjbk!a250!EJEb z!X_YYauA*sP1IG~WW(3Oruo}ZH&1d7$yLRE#z>%`zV7M>WLcM6+E$8?q{u_tzR3TA z-#qUp$fN#n&vl{N#GgK8%Nc6dP0S}9F*Tm|_k$U~ir4C>Wu2HK51F+g`kW;bGJpI% zFUheTE-5HD+#ei=n6>@d2Qe4#n1lf|E*2Ol*9)^nFt0k#;8i8|<1Q9=XvwPR$87!( z&(Q-+V_&jt(%5jv@9=!piVk?|;xY9V>$|V0$@83Lt+3raY{6Xg3;)(_C(-{MOpCEc z-mSFC8Rv3sJ`xcl=3#@}p*+dBxiJ}oP%mPm_D)~`O6Y%Wi90t0CFiG>%Gm#c-}e@r z7ph#C_>qKjg#$&)T)AnO(>`)iUjYB!F8`K08)%9Un>9=8&Cwrr>!#B#4$o~h+U*+_ z`~+)G{^SnU&4pPXxQ}E0?f?9g>q{4yk)kP>cQGMs_17xckY^h_itm5LPiicue^lA+ zRD5R}FWE6RZ`oSRBL~n<<(ka!Nk``$UG?!KgUaGnQEDw(cP5R-NFQu%{X7#za^P!$hxkCKPFu@>m@ z`>dPoiQI`j;Xxab!?n(Io9xFAZSZl<;~S?V`k=+FZa%pEb4$FB{$@H9IF~+*O z=;-H;xtsqJ^9J88@!1bM%YK!xZtsFhN9p{uOTJ7g^g~%-Xr0X30WVFflaSQgvd^>L_yJ*jFa5$LL zU;BiKQc$Y2`y_H2IwP`#GW$TXNcLr@S~dK(`6Tla={J_9g~D{=fAq-nI#&j)bawUR z_?Sy~N9bg@v_iG7fKzKv7y4iCwk*Z`MyPO~gbn%&Y_!~hwf4{nWp;CEO)BQLjI|W& zQGa-?AiejP9r6K%I`@|{o1iQ)BPg-B6Q9Sn^9vJjKCaiUJBa)t&ynLvt1gO=BQn;q zedrJ9l7m$O!fkNt=b4?e`SH25h$SA=gZ{iFA>|8^ck-K;Q#uBV-JjG&BiHJ8o(*&C z{5alH!8)kHhFj9e@3fUcWBN4#Vs2+|>3$M<_YyhXTdQV#!X5pZTjVfjJ6(4+Eq`%f z<@PBUiXfUj&jiQ|3C9W^-Zm(&ons<{`#4A5;6ffxI~Ztbwz2(3z+646r_V~@Pk-xs zOj4zJ9*unJc^JgmSB!x1l1KY0z>5k_TptBVo$+DuB+Gu_$?~Ld`dB}y5vDWN*(e9C z;W$wd5;8^E=^G)$4ou> zw|et8Q9FJjTgi^OgxtFVu|$O&F7&gQ+;~`qz6Z(6E3+@iAg|hAMPXZ-rwtkI!5JO?j4wuZm2etJHbnOGvvg&;)IBgos8zLOG0E%p2-)#T^-n%`C)HS|tWx>8)OU8urIFv%YG6$GlrvR6$aj z-54;#JpZNS&-k4xR885PI$z-k{#OcFP2?2KK`KfXn8qZS!#s7Ol+VP{4qjF5%8E@^L Date: Wed, 31 Jan 2024 15:03:16 -0800 Subject: [PATCH 19/43] Added in additional tests for barostat --- chiron/integrators.py | 7 ++- chiron/mcmc.py | 1 + chiron/tests/test_mcmc.py | 126 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 132 insertions(+), 2 deletions(-) diff --git a/chiron/integrators.py b/chiron/integrators.py index 6b40721..f818214 100644 --- a/chiron/integrators.py +++ b/chiron/integrators.py @@ -202,9 +202,12 @@ def run( # v v += (stepsize_unitless * 0.5) * F / mass_unitless - if (step + self._move_iteration * n_steps) % self.report_frequency == 0: + elapsed_step = step + self._move_iteration * n_steps + if (elapsed_step) % self.report_frequency == 0: if hasattr(self, "reporter") and self.reporter is not None: - self._report(x, potential, nbr_list, step, self._move_iteration) + self._report( + x, potential, nbr_list, step, self._move_iteration, elapsed_step + ) if self.save_traj_in_memory: self.traj.append(x) diff --git a/chiron/mcmc.py b/chiron/mcmc.py index dfd6bfb..3d8bce2 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -840,6 +840,7 @@ def _report( "elapsed_step": elapsed_step, "potential_energy": potential, "volume": volume, + "box_vectors": sampler_state.box_vectors, "max_volume_scale": self.volume_max_scale, "acceptance_probability": acceptance_probability, } diff --git a/chiron/tests/test_mcmc.py b/chiron/tests/test_mcmc.py index d746c6f..abecea0 100644 --- a/chiron/tests/test_mcmc.py +++ b/chiron/tests/test_mcmc.py @@ -329,6 +329,132 @@ 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 MonteCarloBarostatMove + + barostat_move = MonteCarloBarostatMove( + volume_max_scale=0.1, + nr_of_moves=1, + ) + + assert barostat_move.volume_max_scale == 0.1 + assert barostat_move.nr_of_moves == 1 + + +def test_mc_barostat(prep_temp_dir): + import jax.numpy as jnp + + from chiron.reporters import MCReporter, BaseReporter + + wd = prep_temp_dir.join(f"_test_{uuid.uuid4()}") + BaseReporter.set_directory(wd) + simulation_reporter = MCReporter(1) + + from chiron.mcmc import MonteCarloBarostatMove + + barostat_move = MonteCarloBarostatMove( + volume_max_scale=0.1, + nr_of_moves=10, + reporter=simulation_reporter, + report_frequency=1, + ) + + from chiron.potential import IdealGasPotential + from openmm import unit + + 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) + + ideal_gas_potential = IdealGasPotential(topology) + + from chiron.states import SamplerState, ThermodynamicState + from chiron.utils import PRNG + + PRNG.set_seed(1234) + + # define the sampler state + sampler_state = SamplerState( + x0=positions, box_vectors=box_vectors, current_PRNG_key=PRNG.get_random_key() + ) + + # define the thermodynamic state + thermodynamic_state = ThermodynamicState( + potential=ideal_gas_potential, + temperature=300 * unit.kelvin, + pressure=1.0 * unit.atmosphere, + ) + + from chiron.neighbors import PairList, OrthogonalPeriodicSpace + + # since particles are non-interacting and we will not displacece them, the pair list basically + # does nothing in this case. + nbr_list = PairList(OrthogonalPeriodicSpace(), cutoff=0 * unit.nanometer) + + sampler_state, thermodynamic_state, nbr_list = barostat_move.update( + sampler_state, thermodynamic_state, nbr_list + ) + potential_energies = simulation_reporter.get_property("potential_energy") + volumes = simulation_reporter.get_property("volume") + + # ideal gas treatment, so stored energy will only be a + # consequence of pressure, volume, and temperature + from loguru import logger as log + + log.debug(f"PE {potential_energies * unit.kilojoules_per_mole}") + log.debug(thermodynamic_state.pressure) + log.debug(thermodynamic_state.beta) + log.debug(volumes) + log.debug(volumes * unit.nanometer**3) + + # assert that the PE is always zero + assert potential_energies[0] == 0 + assert potential_energies[-1] == 0 + + # the reduced potential will only be a consequence of the pressure, volume, and temperature + + assert jnp.isclose( + thermodynamic_state.get_reduced_potential(sampler_state), + ( + thermodynamic_state.pressure + * thermodynamic_state.beta + * (volumes[-1] * unit.nanometer**3) + ), + 1e-3, + ) + + print(barostat_move.statistics["n_accepted"]) + assert barostat_move.statistics["n_proposed"] == 10 + assert barostat_move.statistics["n_accepted"] == 8 + + 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 From c992b047351f4ca1724bfd372ff3691b5ff80edc Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Wed, 31 Jan 2024 16:59:40 -0800 Subject: [PATCH 20/43] Fixed convergence test syntax: these were missed on my part because they were marked to be skipped because they take too long. --- chiron/tests/test_convergence_tests.py | 26 +++++++++++++++++++------- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/chiron/tests/test_convergence_tests.py b/chiron/tests/test_convergence_tests.py index 16cfad3..95df83d 100644 --- a/chiron/tests/test_convergence_tests.py +++ b/chiron/tests/test_convergence_tests.py @@ -16,8 +16,8 @@ def prep_temp_dir(tmpdir_factory): IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true" -@pytest.mark.skip(reason="Tests takes too long") -@pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Test takes too long.") +# @pytest.mark.skip(reason="Tests takes too long") +# @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Test takes too long.") def test_convergence_of_MC_estimator(prep_temp_dir): from openmm import unit @@ -44,9 +44,17 @@ 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), + ) + from chiron.utils import PRNG + + PRNG.set_seed(1234) + + sampler_state = SamplerState( + x0=ho.positions, current_PRNG_key=PRNG.get_random_key() ) - sampler_state = SamplerState(ho.positions) from chiron.reporters import _SimulationReporter @@ -113,8 +121,8 @@ def __init__(self, temperature): ) -@pytest.mark.skip(reason="Tests takes too long") -@pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Test 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 from chiron.states import SamplerState, ThermodynamicState @@ -133,10 +141,14 @@ def test_langevin_dynamics_with_LJ_fluid(prep_temp_dir): ) print(lj_fluid.system.getDefaultPeriodicBoxVectors()) + from chiron.utils import PRNG + + PRNG.set_seed(1234) sampler_state = SamplerState( x0=lj_fluid.positions, box_vectors=lj_fluid.system.getDefaultPeriodicBoxVectors(), + current_PRNG_key=PRNG.get_random_key(), ) print(sampler_state.x0.shape) print(sampler_state.box_vectors) @@ -162,7 +174,7 @@ def test_langevin_dynamics_with_LJ_fluid(prep_temp_dir): integrator.run( sampler_state, thermodynamic_state, - n_steps=2000, + n_steps=1000, nbr_list=nbr_list, progress_bar=True, ) From 0f049b24e26fea2b4113ec25cdbf9356f854635f Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Thu, 1 Feb 2024 12:24:49 -0800 Subject: [PATCH 21/43] Fixed convergence test syntax: these were missed on my part because they were marked to be skipped because they take too long. for Harmonic oscillator we do not seem to need to run as long as set initially to pass and get convergence. --- chiron/tests/test_convergence_tests.py | 22 +++++++++++++--------- chiron/utils.py | 2 +- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/chiron/tests/test_convergence_tests.py b/chiron/tests/test_convergence_tests.py index 95df83d..ebd31e6 100644 --- a/chiron/tests/test_convergence_tests.py +++ b/chiron/tests/test_convergence_tests.py @@ -16,8 +16,8 @@ def prep_temp_dir(tmpdir_factory): IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true" -# @pytest.mark.skip(reason="Tests takes too long") -# @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Test takes too long.") +@pytest.mark.skip(reason="Tests takes too long") +@pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Test takes too long.") def test_convergence_of_MC_estimator(prep_temp_dir): from openmm import unit @@ -66,7 +66,7 @@ def test_convergence_of_MC_estimator(prep_temp_dir): from chiron.mcmc import MetropolisDisplacementMove, MoveSchedule, MCMCSampler mc_displacement_move = MetropolisDisplacementMove( - nr_of_moves=100_000, + nr_of_moves=1_000, displacement_sigma=0.5 * unit.angstrom, atom_subset=[0], reporter=simulation_reporter, @@ -90,7 +90,9 @@ def test_convergence_of_MC_estimator(prep_temp_dir): plt.plot(chiron_energy) print("Expectation values generated with chiron") - es = chiron_energy + import jax.numpy as jnp + + es = jnp.array(chiron_energy) print(es.mean(), es.std()) print("Expectation values from openmmtools") @@ -121,8 +123,8 @@ def __init__(self, temperature): ) -# @pytest.mark.skip(reason="Tests takes too long") -# @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Test 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 from chiron.states import SamplerState, ThermodynamicState @@ -165,12 +167,14 @@ def test_langevin_dynamics_with_LJ_fluid(prep_temp_dir): potential=lj_potential, temperature=300 * unit.kelvin ) - from chiron.reporters import _SimulationReporter + from chiron.reporters import LangevinDynamicsReporter id = uuid.uuid4() - reporter = _SimulationReporter(f"{prep_temp_dir}/test_{id}.h5") + reporter = LangevinDynamicsReporter(f"{prep_temp_dir}/test_{id}.h5") - integrator = LangevinIntegrator(reporter=reporter, report_frequency=100) + integrator = LangevinIntegrator( + reporter=reporter, report_frequency=100, initialize_velocities=True + ) integrator.run( sampler_state, thermodynamic_state, diff --git a/chiron/utils.py b/chiron/utils.py index 07927e0..6f8fbde 100644 --- a/chiron/utils.py +++ b/chiron/utils.py @@ -86,7 +86,7 @@ def get_nr_of_particles(topology: Topology) -> int: def get_list_of_mass(topology: Topology) -> unit.Quantity: """Get the mass of the system from the topology.""" - from simtk import unit + from openmm import unit mass = [] for atom in topology.atoms(): From 21ff068f786631b201d9b735244fa6f6a6ec45e6 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Fri, 2 Feb 2024 13:02:08 -0800 Subject: [PATCH 22/43] Updated CI.yaml to enabled CI to run for branches commiting to the multistage branch. Thanks Mike Henry! --- .github/workflows/CI.yaml | 1 + Examples/LJ_MCMC.py | 2 +- chiron/mcmc.py | 10 ++++------ 3 files changed, 6 insertions(+), 7 deletions(-) diff --git a/.github/workflows/CI.yaml b/.github/workflows/CI.yaml index 0e307f5..2917564 100644 --- a/.github/workflows/CI.yaml +++ b/.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. diff --git a/Examples/LJ_MCMC.py b/Examples/LJ_MCMC.py index c173352..be3bb59 100644 --- a/Examples/LJ_MCMC.py +++ b/Examples/LJ_MCMC.py @@ -149,4 +149,4 @@ ) sampler = MCMCSampler(move_set, sampler_state, thermodynamic_state) -sampler.run(n_iterations=1000, nbr_list=nbr_list) # how many times to repeat +sampler.run(n_iterations=10, nbr_list=nbr_list) # how many times to repeat diff --git a/chiron/mcmc.py b/chiron/mcmc.py index 3d8bce2..b2d86e3 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -938,12 +938,10 @@ def _propose( proposed_sampler_state, proposed_nbr_list ) - # χ = exp ⎡−β (ΔU + PΔV ) + N ln(V new /V old )⎤ - log_proposal_ratio = ( - -proposed_reduced_pot - + current_reduced_pot - + nr_of_atoms * jnp.log(proposed_volume / initial_volume) - ) + # ⎡−β (ΔU + PΔV ) + N ln(V new /V old )⎤ + log_proposal_ratio = -( + proposed_reduced_pot - current_reduced_pot + ) + nr_of_atoms * jnp.log(proposed_volume / initial_volume) # we do not change the thermodynamic state so we can return 'current_thermodynamic_state' return ( From 2074665be540729a8124db48eeb244847bb45aa9 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Fri, 2 Feb 2024 13:16:01 -0800 Subject: [PATCH 23/43] Added in skip to the multistate testing, as this still needs to be worked on in the multistate branch this is being commited to" --- chiron/states.py | 2 +- chiron/tests/test_multistate.py | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/chiron/states.py b/chiron/states.py index 0c498c4..ded4842 100644 --- a/chiron/states.py +++ b/chiron/states.py @@ -315,7 +315,7 @@ def get_reduced_potential( from loguru import logger as log reduced_potential += self.pressure * self.volume - + # add chemical potential return self.beta * reduced_potential def kT_to_kJ_per_mol(self, energy): diff --git a/chiron/tests/test_multistate.py b/chiron/tests/test_multistate.py index ff8feb6..29f3815 100644 --- a/chiron/tests/test_multistate.py +++ b/chiron/tests/test_multistate.py @@ -195,6 +195,9 @@ def test_multistate_minimize(ho_multistate_sampler_multiple_minima: MultiStateSa ) +@pytest.mark.skip( + reason="Multistate code still needs to be modified in the multistage branch" +) def test_multistate_run(ho_multistate_sampler_multiple_ks: MultiStateSampler): """ Test function for running the multistate sampler. From e339e4498d68c91cbd6e448157ded410eb3b6432 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Fri, 2 Feb 2024 13:32:01 -0800 Subject: [PATCH 24/43] match/case statements don't exist in python 3.9. I commented this out and just added in if/elif statements. --- chiron/multistate.py | 34 +++++++++++++++++++++++----------- 1 file changed, 23 insertions(+), 11 deletions(-) diff --git a/chiron/multistate.py b/chiron/multistate.py index 3306ee1..269d32c 100644 --- a/chiron/multistate.py +++ b/chiron/multistate.py @@ -598,17 +598,29 @@ def _report(self, property: str) -> None: from loguru import logger as log log.debug(f"Reporting {property}...") - match property: - case "positions": - return self._report_positions() - case "states": - pass - case "u_kn": - return self._report_energy_matrix() - case "trajectory": - return - case "mixing_statistics": - return + if property == "positions": + return self._report_positions() + elif property == "states": + pass + elif property == "u_kn": + return self._report_energy_matrix() + elif property == "trajectory": + return + elif "mixing_statistics": + return + + # match isn't in python 3.9; we can discuss if we want to drop python 3.0 support or just keep the if/else structure + # match property: + # case "positions": + # return self._report_positions() + # case "states": + # pass + # case "u_kn": + # return self._report_energy_matrix() + # case "trajectory": + # return + # case "mixing_statistics": + # return def _report_iteration(self): """ From 231f830e1fea53dc405f15f75779dd9be7465f60 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Fri, 2 Feb 2024 13:42:25 -0800 Subject: [PATCH 25/43] fixture was missing in test_testsystems --- chiron/tests/test_testsystems.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/chiron/tests/test_testsystems.py b/chiron/tests/test_testsystems.py index c1f9401..7cadcac 100644 --- a/chiron/tests/test_testsystems.py +++ b/chiron/tests/test_testsystems.py @@ -14,6 +14,13 @@ def compute_openmm_reference_energy(testsystem, positions): return e +@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 test_HO(): """ Test the harmonic oscillator system using a Langevin integrator. From 9085a438be3f0dd070807fd800b60d8f5c4aef40 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Fri, 2 Feb 2024 13:49:03 -0800 Subject: [PATCH 26/43] fixture was missing in test_testsystems --- chiron/tests/test_testsystems.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/chiron/tests/test_testsystems.py b/chiron/tests/test_testsystems.py index 7cadcac..c2f60a1 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 @@ -14,13 +24,6 @@ def compute_openmm_reference_energy(testsystem, positions): return e -@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 test_HO(): """ Test the harmonic oscillator system using a Langevin integrator. From 88ee9ba164d6fb329fff6d02d0fea7b5c1c33458 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Fri, 2 Feb 2024 14:20:38 -0800 Subject: [PATCH 27/43] weirdly wrong syntax in the test ideal gas test. --- chiron/tests/test_testsystems.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/chiron/tests/test_testsystems.py b/chiron/tests/test_testsystems.py index c2f60a1..a3edabe 100644 --- a/chiron/tests/test_testsystems.py +++ b/chiron/tests/test_testsystems.py @@ -284,8 +284,8 @@ def test_ideal_gas(prep_temp_dir): ) mc_displacement_move = MetropolisDisplacementMove( - stepsize=0.1, - n_steps=10, + displacement_sigma=0.1 * unit.nanometer, + nr_of_moves=10, reporter=reporter, update_stepsize=True, update_stepsize_frequency=100, From 030435c83720be9ced255152a2748b71fbd7cda6 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Fri, 9 Feb 2024 13:16:38 -0800 Subject: [PATCH 28/43] Updating ideal gas example. --- Examples/Idealgas.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/Examples/Idealgas.py b/Examples/Idealgas.py index 9e307e5..3ce3286 100644 --- a/Examples/Idealgas.py +++ b/Examples/Idealgas.py @@ -8,12 +8,14 @@ n_particles = 216 temperature = 298 * unit.kelvin pressure = 1 * unit.atmosphere -mass = unit.Quantity(39.9, unit.gram / unit.mole) + +# mass = unit.Quantity(39.9, unit.gram / unit.mole) ideal_gas = IdealGas(nparticles=n_particles, temperature=temperature, pressure=pressure) + from chiron.potential import IdealGasPotential -from chiron.utils import PRNG +from chiron.utils import PRNG, get_list_of_mass import jax.numpy as jnp # particles are non interacting @@ -107,17 +109,22 @@ ideal_volume = ideal_gas.get_volume_expectation(thermodynamic_state) ideal_volume_std = ideal_gas.get_volume_standard_deviation(thermodynamic_state) -print(ideal_volume, ideal_volume_std) +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(volume_mean, volume_std) +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 = mass * n_particles / unit.AVOGADRO_CONSTANT_NA / ideal_volume -measured_density = mass * n_particles / unit.AVOGADRO_CONSTANT_NA / volume_mean +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), From 3ec2aa062001bc04b6919c35fcd787170901b8e4 Mon Sep 17 00:00:00 2001 From: Chris Iacovella Date: Fri, 9 Feb 2024 13:17:29 -0800 Subject: [PATCH 29/43] Update Examples/Idealgas.py with descriptive assert statement Co-authored-by: Marcus Wieder <31651017+wiederm@users.noreply.github.com> --- Examples/Idealgas.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/Examples/Idealgas.py b/Examples/Idealgas.py index 3ce3286..32b996e 100644 --- a/Examples/Idealgas.py +++ b/Examples/Idealgas.py @@ -132,7 +132,9 @@ atol=1e-1, ) # see if within 5% of ideal volume -assert abs(ideal_volume - volume_mean) / ideal_volume < 0.05 +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 From 0987cf2434ded9b43bf62a4b55aceef639f56dca Mon Sep 17 00:00:00 2001 From: Chris Iacovella Date: Fri, 9 Feb 2024 13:17:46 -0800 Subject: [PATCH 30/43] Update Examples/Idealgas.py with descriptive assert statement Co-authored-by: Marcus Wieder <31651017+wiederm@users.noreply.github.com> --- Examples/Idealgas.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/Examples/Idealgas.py b/Examples/Idealgas.py index 32b996e..29380f3 100644 --- a/Examples/Idealgas.py +++ b/Examples/Idealgas.py @@ -137,4 +137,6 @@ ), 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 +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" From fc9ede3989fb56ab697fd5ea44cc06db547b8eb8 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Fri, 9 Feb 2024 13:19:29 -0800 Subject: [PATCH 31/43] Updating ideal gas example. --- Examples/Idealgas.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/Examples/Idealgas.py b/Examples/Idealgas.py index 29380f3..e512ada 100644 --- a/Examples/Idealgas.py +++ b/Examples/Idealgas.py @@ -9,8 +9,6 @@ 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 1daab8a7eb023f46f51c769bba3bca010d06ff28 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Fri, 9 Feb 2024 14:02:49 -0800 Subject: [PATCH 32/43] Updating ideal gas example. --- Examples/Idealgas.py | 8 ++++++++ Examples/LJ_MCMC.py | 15 +++++++++++---- chiron/integrators.py | 20 +++++++------------- chiron/tests/test_convergence_tests.py | 4 +--- chiron/tests/test_integrators.py | 10 ---------- chiron/tests/test_multistate.py | 4 +--- chiron/utils.py | 19 +++++++++++++++++++ 7 files changed, 47 insertions(+), 33 deletions(-) diff --git a/Examples/Idealgas.py b/Examples/Idealgas.py index e512ada..6eed6fa 100644 --- a/Examples/Idealgas.py +++ b/Examples/Idealgas.py @@ -1,6 +1,14 @@ 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 +MetropolisDisplacementMove 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 diff --git a/Examples/LJ_MCMC.py b/Examples/LJ_MCMC.py index be3bb59..4b069c6 100644 --- a/Examples/LJ_MCMC.py +++ b/Examples/LJ_MCMC.py @@ -1,7 +1,12 @@ 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 @@ -20,7 +25,9 @@ # these were generated in Mbuild using fill_box which wraps packmol # a minimum spacing of 0.4 nm was used during construction. -positions = jnp.load("Examples/methane_coords.npy") * unit.nanometer +from chiron.utils import get_full_path + +positions = jnp.load(get_full_path("Examples/methane_coords.npy")) * unit.nanometer box_vectors = jnp.array( [ @@ -136,17 +143,17 @@ nr_of_steps=100, reporter=reporter_langevin, report_frequency=10, - reinitialize_velocities=True, + initialize_velocities=True, ) move_set = MoveSchedule( [ ("LangevinDynamicsMove", langevin_dynamics_move), - ("MetropolisDisplacementMove", mc_displacement_move), + # ("MetropolisDisplacementMove", mc_displacement_move), ("MonteCarloBarostatMove", mc_barostat_move), ] ) sampler = MCMCSampler(move_set, sampler_state, thermodynamic_state) -sampler.run(n_iterations=10, nbr_list=nbr_list) # how many times to repeat +sampler.run(n_iterations=100, nbr_list=nbr_list) # how many times to repeat diff --git a/chiron/integrators.py b/chiron/integrators.py index f818214..3760e38 100644 --- a/chiron/integrators.py +++ b/chiron/integrators.py @@ -26,7 +26,6 @@ def __init__( self, stepsize=1.0 * unit.femtoseconds, collision_rate=1.0 / unit.picoseconds, - initialize_velocities: bool = False, reinitialize_velocities: bool = False, report_frequency: int = 100, reporter: Optional[LangevinDynamicsReporter] = None, @@ -41,8 +40,6 @@ def __init__( Time step of integration with units of time. Default is 1.0 * unit.femtoseconds. collision_rate : unit.Quantity, optional Collision rate for the Langevin dynamics, with units 1/time. Default is 1.0 / unit.picoseconds. - initialize_velocities : bool, optional - Flag indicating whether to initialize the velocities the first time the run function is called. Default is False. reinitialize_velocities : bool, optional Flag indicating whether to reinitialize the velocities each time the run function is called. Default is False. report_frequency : int, optional @@ -52,9 +49,6 @@ def __init__( save_traj_in_memory: bool Flag indicating whether to save the trajectory in memory. Default is False. NOTE: Only for debugging purposes. - reinitialize_velocities: bool - Whether to reinitialize the velocities each time the run function is called. - Default is False. """ from loguru import logger as log @@ -76,7 +70,6 @@ def __init__( self.save_traj_in_memory = save_traj_in_memory self.traj = [] self.reinitialize_velocities = reinitialize_velocities - self.initialize_velocities = initialize_velocities self._move_iteration = 0 def run( @@ -151,19 +144,20 @@ def run( sampler_state.velocities = initialize_velocities( temperature, potential.topology, key ) - self.initialize_velocities = False - elif self.initialize_velocities: + elif sampler_state._velocities is None: # v0 = sigma_v * random.normal(key, x0.shape) from .utils import initialize_velocities sampler_state.velocities = initialize_velocities( temperature, potential.topology, key ) - self.initialize_velocities = False - else: - if sampler_state._velocities is None: - raise ValueError("Velocities must be set before running the integrator") + elif sampler_state._velocities.shape[0] != sampler_state.x0.shape[0]: + from .utils import initialize_velocities + + sampler_state.velocities = initialize_velocities( + temperature, potential.topology, key + ) # extract the velocities from the sampler state v0 = sampler_state.velocities diff --git a/chiron/tests/test_convergence_tests.py b/chiron/tests/test_convergence_tests.py index ebd31e6..c05c072 100644 --- a/chiron/tests/test_convergence_tests.py +++ b/chiron/tests/test_convergence_tests.py @@ -172,9 +172,7 @@ def test_langevin_dynamics_with_LJ_fluid(prep_temp_dir): id = uuid.uuid4() reporter = LangevinDynamicsReporter(f"{prep_temp_dir}/test_{id}.h5") - integrator = LangevinIntegrator( - reporter=reporter, report_frequency=100, initialize_velocities=True - ) + integrator = LangevinIntegrator(reporter=reporter, report_frequency=100) integrator.run( sampler_state, thermodynamic_state, diff --git a/chiron/tests/test_integrators.py b/chiron/tests/test_integrators.py index ed6f01f..e8121de 100644 --- a/chiron/tests/test_integrators.py +++ b/chiron/tests/test_integrators.py @@ -42,16 +42,6 @@ def test_langevin_dynamics(prep_temp_dir, provide_testsystems_and_potentials): reporter = LangevinDynamicsReporter() - with pytest.raises(ValueError): - integrator = LangevinIntegrator( - reporter=reporter, report_frequency=1, reinitialize_velocities=False - ) - - integrator.run( - sampler_state, - thermodynamic_state, - n_steps=20, - ) integrator = LangevinIntegrator( reporter=reporter, report_frequency=1, reinitialize_velocities=True ) diff --git a/chiron/tests/test_multistate.py b/chiron/tests/test_multistate.py index 29f3815..7da5f7e 100644 --- a/chiron/tests/test_multistate.py +++ b/chiron/tests/test_multistate.py @@ -24,9 +24,7 @@ def setup_sampler() -> Tuple[NeighborListNsqrd, MultiStateSampler]: OrthogonalPeriodicSpace(), cutoff=cutoff, skin=skin, n_max_neighbors=180 ) - move = LangevinDynamicsMove( - stepsize=1.0 * unit.femtoseconds, nr_of_steps=100, initialize_velocities=True - ) + move = LangevinDynamicsMove(stepsize=1.0 * unit.femtoseconds, nr_of_steps=100) BaseReporter.set_directory("multistate_test") reporter = MultistateReporter() reporter.reset_reporter_file() diff --git a/chiron/utils.py b/chiron/utils.py index 6f8fbde..41daf33 100644 --- a/chiron/utils.py +++ b/chiron/utils.py @@ -38,6 +38,25 @@ def get_random_key(cls) -> int: return subkey +def get_full_path(relative_path: str) -> str: + """Get the fill path of a file that is defined relative to the chiron module root directory. + + Parameters + ---------- + relative_path : str + The relative path of the file. + + Returns + ------- + str + The full path of the file. + """ + from importlib.resources import files + + _MODULE_ROOT = files("chiron") + return f"{_MODULE_ROOT}/../{relative_path}" + + def get_data_file_path(relative_path: str) -> str: """Get the full path to one of the reference files in testsystems. In the source distribution, these files are in ``chiron/data/``, From e419169b9d9bab005a3c10c085e2c9897b7d2c49 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Fri, 9 Feb 2024 14:04:51 -0800 Subject: [PATCH 33/43] removed velocity initialization flag in langevin --- chiron/mcmc.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/chiron/mcmc.py b/chiron/mcmc.py index b2d86e3..7477de8 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -83,7 +83,6 @@ def __init__( self, stepsize: unit.Quantity = 1.0 * unit.femtoseconds, collision_rate: unit.Quantity = 1.0 / unit.picoseconds, - initialize_velocities: bool = False, reinitialize_velocities: bool = False, reporter: Optional[LangevinDynamicsReporter] = None, report_frequency: int = 100, @@ -99,9 +98,6 @@ def __init__( Time step size for the integration. collision_rate : unit.Quantity Collision rate for the Langevin dynamics. - initialize_velocities: bool, optional - Whether to initialize the velocities the first time the run function is called. - Default is False. reinitialize_velocities : bool, optional Whether to reinitialize the velocities each time the run function is called. Default is False. From 9ad8a23f4fe1f1627a0d58aeabef1b925e5156a0 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Fri, 9 Feb 2024 15:07:13 -0800 Subject: [PATCH 34/43] updated various functions in reponse to marcus' comments. --- chiron/mcmc.py | 44 ++++++++++++++++++++++---------------------- chiron/potential.py | 10 ++++------ chiron/states.py | 8 +++++--- 3 files changed, 31 insertions(+), 31 deletions(-) diff --git a/chiron/mcmc.py b/chiron/mcmc.py index 7477de8..f912408 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -256,25 +256,24 @@ def update( nbr_list: PairsBase The updated neighbor/pair list. If a nbr_list is not set, this will be None. """ - calculate_current_potential = True + calculate_current_reduced_potential = True for i in range(self.nr_of_moves): sampler_state, thermodynamic_state, nbr_list = self._step( sampler_state, thermodynamic_state, nbr_list, - calculate_current_potential=calculate_current_potential, + calculate_current_reduced_potential=calculate_current_reduced_potential, ) - # after the first step, we don't need to recalculate the current potential, it will be stored - calculate_current_potential = False + # after the first step, we don't need to recalculate the current reduced_potential, it will be stored + calculate_current_reduced_potential = False + # I think it makes sense to use i + self.nr_of_moves*self._move_iteration as our current "step" + # otherwise, if we just used i, instances where self.report_frequency > self.nr_of_moves would only report on the + # first step, which might actually be more frequent than we specify elapsed_step = i + self._move_iteration * self.nr_of_moves if hasattr(self, "reporter"): if self.reporter is not None: - # I think it makes sense to use i + self.nr_of_moves*self._move_iteration as our current "step" - # otherwise, if we just used i, instances where self.report_frequency > self.nr_of_moves would only report on the - # first step, which might actually be more frequent than we specify - if elapsed_step % self.report_frequency == 0: self._report( i, @@ -351,7 +350,7 @@ def _step( current_sampler_state: SamplerState, current_thermodynamic_state: ThermodynamicState, current_nbr_list: Optional[PairsBase] = None, - calculate_current_potential: bool = True, + calculate_current_reduced_potential: bool = True, ) -> Tuple[SamplerState, ThermodynamicState, Optional[PairsBase]]: """ Performs an individual MC step. @@ -366,7 +365,7 @@ def _step( Current thermodynamic state. current_nbr_list : Optional[PairsBase] Neighbor list associated with the current state. - calculate_current_potential : bool, optional + calculate_current_reduced_potential : bool, optional Whether to calculate the current reduced potential. Default is True. Returns @@ -386,9 +385,9 @@ def _step( # if this is the first time we are calling this function during this iteration # we will need to calculate the reduced potential for the current state - # this is toggled by the calculate_current_potential flag + # this is toggled by the calculate_current_reduced_potential flag # otherwise, we can use the one that was saved from the last step, for efficiency - if calculate_current_potential: + if calculate_current_reduced_potential: current_reduced_pot = current_thermodynamic_state.get_reduced_potential( current_sampler_state, current_nbr_list ) @@ -418,16 +417,16 @@ def _step( current_nbr_list, ) - # accept or reject the proposed state - decision = self._accept_or_reject( - log_proposal_ratio, - proposed_sampler_state.new_PRNG_key, - method=self.method, - ) - # a function that will update the statistics for the move - if jnp.isnan(proposed_reduced_pot): decision = False + else: + # accept or reject the proposed state + decision = self._accept_or_reject( + log_proposal_ratio, + proposed_sampler_state.new_PRNG_key, + method=self.method, + ) + # a function that will update the statistics for the move self._update_statistics(decision) @@ -933,8 +932,9 @@ def _propose( proposed_reduced_pot = current_thermodynamic_state.get_reduced_potential( proposed_sampler_state, proposed_nbr_list ) - - # ⎡−β (ΔU + PΔV ) + N ln(V new /V old )⎤ + # NPT acceptance criteria was originally defined in McDonald 1972, https://doi.org/10.1080/00268977200100031 + # (see equation 9). The acceptance probability is given by: + # ⎡−β (ΔU + PΔV ) + N ln(V new /V old )⎤ log_proposal_ratio = -( proposed_reduced_pot - current_reduced_pot ) + nr_of_atoms * jnp.log(proposed_volume / initial_volume) diff --git a/chiron/potential.py b/chiron/potential.py index 8362ef1..097fea5 100644 --- a/chiron/potential.py +++ b/chiron/potential.py @@ -78,12 +78,10 @@ def __init__( """ - 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)}" - ) + if not isinstance(topology, (Topology, property)) and topology is not None: + raise TypeError( + f"Topology must be a Topology object, a property, or None, got type(topology) = {type(topology)}" + ) self.topology = topology diff --git a/chiron/states.py b/chiron/states.py index ded4842..a9ad46a 100644 --- a/chiron/states.py +++ b/chiron/states.py @@ -87,7 +87,7 @@ def __init__( self._current_PRNG_key = current_PRNG_key self._box_vectors = box_vectors self._distance_unit = unit.nanometer - self._velocity_unit = unit.nanometer / unit.picosecond + self._time_unit = unit.picosecond @property def n_particles(self) -> int: @@ -132,14 +132,16 @@ def velocities(self, velocities: Union[jnp.array, unit.Quantity]) -> None: if isinstance(velocities, unit.Quantity): self._velocities = velocities else: - self._velocities = unit.Quantity(velocities, self._velocity_unit) + self._velocities = unit.Quantity( + velocities, self._distance_unit / self._time_unit + ) @property def distance_unit(self) -> unit.Unit: return self._distance_unit def velocity_unit(self) -> unit.Unit: - return self._velocity_unit + return self._distance_unit / self._time_unit @property def new_PRNG_key(self) -> random.PRNGKey: From 60e0814644b4e6cbe69508ece4203244151efbfd Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Fri, 9 Feb 2024 22:10:44 -0800 Subject: [PATCH 35/43] Merge failed to correctly merge, and broke MCMCSampler; fixed now. multistate reporter giving an error. --- Examples/Idealgas.py | 6 ++++-- Examples/LJ_MCMC.py | 7 ++++--- chiron/mcmc.py | 29 +++++++++++++++++++++++------ chiron/tests/test_mcmc.py | 19 ++++--------------- chiron/tests/test_testsystems.py | 6 ++++-- chiron/tests/test_utils.py | 3 ++- 6 files changed, 41 insertions(+), 29 deletions(-) diff --git a/Examples/Idealgas.py b/Examples/Idealgas.py index 6eed6fa..34ed686 100644 --- a/Examples/Idealgas.py +++ b/Examples/Idealgas.py @@ -98,8 +98,10 @@ ] ) -sampler = MCMCSampler(move_set, sampler_state, thermodynamic_state) -sampler.run(n_iterations=10, nbr_list=nbr_list) # how many times to repeat +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") diff --git a/Examples/LJ_MCMC.py b/Examples/LJ_MCMC.py index 4b069c6..aa30c13 100644 --- a/Examples/LJ_MCMC.py +++ b/Examples/LJ_MCMC.py @@ -143,7 +143,6 @@ nr_of_steps=100, reporter=reporter_langevin, report_frequency=10, - initialize_velocities=True, ) @@ -155,5 +154,7 @@ ] ) -sampler = MCMCSampler(move_set, sampler_state, thermodynamic_state) -sampler.run(n_iterations=100, nbr_list=nbr_list) # how many times to repeat +sampler = MCMCSampler(move_set) +sampler.run( + sampler_state, thermodynamic_state, n_iterations=100, nbr_list=nbr_list +) # how many times to repeat diff --git a/chiron/mcmc.py b/chiron/mcmc.py index 0ca2a92..dea4677 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -129,7 +129,6 @@ def __init__( self.integrator = LangevinIntegrator( stepsize=self.stepsize, collision_rate=self.collision_rate, - initialize_velocities=initialize_velocities, reinitialize_velocities=reinitialize_velocities, report_frequency=report_frequency, reporter=reporter, @@ -1036,26 +1035,44 @@ def __init__( log.info("Initializing MCMC sampler") self.move = move_set - def run( self, sampler_state: SamplerState, thermodynamic_state: ThermodynamicState, n_iterations: int = 1, + nbr_list: Optional[PairsBase] = None, ): """ Run the sampler for a specified number of iterations. Parameters ---------- + sampler_state : SamplerState + The initial state of the sampler. + thermodynamic_state : ThermodynamicState + The thermodynamic state of the system. n_iterations : int, optional Number of iterations of the sampler to run. + Default is 1. + nbr_list : PairsBase, optional + The neighbor list to use for the simulation. + + Returns + ------- + sampler_state : SamplerState + The updated sampler state. + thermodynamic_state : ThermodynamicState + The updated thermodynamic state. + nbr_list: PairsBase + The updated neighbor/pair list. If a nbr_list is not set, this will be None. + """ from loguru import logger as log from copy import deepcopy sampler_state = deepcopy(sampler_state) thermodynamic_state = deepcopy(thermodynamic_state) + nbr_list = deepcopy(nbr_list) log.info("Running MCMC sampler") log.info(f"move_schedule = {self.move.move_schedule}") @@ -1064,7 +1081,9 @@ def run( for move_name, move in self.move.move_schedule: log.debug(f"Performing: {move_name}") - move.run(sampler_state, thermodynamic_state) + sampler_state, thermodynamic_state, nbr_list = move.update( + sampler_state, thermodynamic_state, nbr_list + ) log.info("Finished running MCMC sampler") log.debug("Closing reporter") @@ -1072,6 +1091,4 @@ def run( if move.reporter is not None: move.reporter.flush_buffer() log.debug(f"Closed reporter {move.reporter.log_file_path}") - return sampler_state - - # I think we should return the sampler/thermo state to be consistent + return sampler_state, thermodynamic_state, nbr_list diff --git a/chiron/tests/test_mcmc.py b/chiron/tests/test_mcmc.py index ce707f1..3c182c3 100644 --- a/chiron/tests/test_mcmc.py +++ b/chiron/tests/test_mcmc.py @@ -126,19 +126,6 @@ def test_sample_from_harmonic_osciallator_with_MCMC_classes_and_LangevinDynamics simulation_reporter = LangevinDynamicsReporter(1) - # this will fail because we have not initialized the velocity - with pytest.raises(ValueError): - langevin_move = LangevinDynamicsMove( - nr_of_steps=10, reinitialize_velocities=False, reporter=simulation_reporter - ) - move_set = MoveSchedule([("LangevinMove", langevin_move)]) - - # Initalize the sampler - sampler = MCMCSampler(move_set, sampler_state, thermodynamic_state) - - # Run the sampler with the thermodynamic state and sampler state and return the sampler state - sampler.run(n_iterations=2) # how many times to repeat - # the following will reinitialize the velocities for each iteration langevin_move = LangevinDynamicsMove( nr_of_steps=10, reinitialize_velocities=True, reporter=simulation_reporter @@ -147,10 +134,12 @@ def test_sample_from_harmonic_osciallator_with_MCMC_classes_and_LangevinDynamics move_set = MoveSchedule([("LangevinMove", langevin_move)]) # Initalize the sampler - sampler = MCMCSampler(move_set, sampler_state, thermodynamic_state) + sampler = MCMCSampler(move_set) # Run the sampler with the thermodynamic state and sampler state and return the sampler state - sampler.run(n_iterations=2) # how many times to repeat + sampler.run( + sampler_state, thermodynamic_state, n_iterations=2 + ) # how many times to repeat # the following will use the initialize velocities function from chiron.utils import initialize_velocities diff --git a/chiron/tests/test_testsystems.py b/chiron/tests/test_testsystems.py index a3edabe..953effd 100644 --- a/chiron/tests/test_testsystems.py +++ b/chiron/tests/test_testsystems.py @@ -305,8 +305,10 @@ def test_ideal_gas(prep_temp_dir): ] ) - sampler = MCMCSampler(move_set, sampler_state, thermodynamic_state) - sampler.run(n_iterations=10, nbr_list=nbr_list) # how many times to repeat + sampler = MCMCSampler(move_set) + sampler.run( + sampler_state, thermodynamic_state, n_iterations=10, nbr_list=nbr_list + ) # how many times to repeat volume = reporter.get_property("volume") diff --git a/chiron/tests/test_utils.py b/chiron/tests/test_utils.py index 119a7f5..5477fbd 100644 --- a/chiron/tests/test_utils.py +++ b/chiron/tests/test_utils.py @@ -65,7 +65,8 @@ def test_reporter(prep_temp_dir, ho_multistate_sampler_multiple_ks): reporter.reset_reporter_file() integrator = LangevinIntegrator( - reporter=reporter, report_frequency=1, initialize_velocities=True + reporter=reporter, + report_frequency=1, ) integrator.run( sampler_state, From e78992db038d53b566a590e1afb1a0a4752726de Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Mon, 12 Feb 2024 13:00:58 -0800 Subject: [PATCH 36/43] Working through comments from Jchodera. --- Examples/Idealgas.py | 12 +-- Examples/LJ_MCMC.py | 81 +++++++++--------- Examples/LJ_langevin.py | 10 +-- Examples/LJ_mcmove.py | 6 +- chiron/mcmc.py | 113 ++++++++++++------------- chiron/tests/test_convergence_tests.py | 2 +- chiron/tests/test_mcmc.py | 10 +-- chiron/tests/test_testsystems.py | 12 +-- 8 files changed, 121 insertions(+), 125 deletions(-) diff --git a/Examples/Idealgas.py b/Examples/Idealgas.py index 34ed686..bbeffe8 100644 --- a/Examples/Idealgas.py +++ b/Examples/Idealgas.py @@ -76,18 +76,18 @@ # initialize the displacement move mc_barostat_move = MonteCarloBarostatMove( volume_max_scale=0.2, - nr_of_moves=10, + number_of_moves=10, reporter=reporter, - update_stepsize=True, - update_stepsize_frequency=100, + autotune=True, + autotune_interval=100, ) # initialize the barostat move and the move schedule metropolis_displacement_move = MetropolisDisplacementMove( displacement_sigma=0.1 * unit.nanometer, - nr_of_moves=100, - update_stepsize=True, - update_stepsize_frequency=100, + number_of_moves=100, + autotune=True, + autotune_interval=100, ) # define the move schedule diff --git a/Examples/LJ_MCMC.py b/Examples/LJ_MCMC.py index aa30c13..1ba26dc 100644 --- a/Examples/LJ_MCMC.py +++ b/Examples/LJ_MCMC.py @@ -29,12 +29,15 @@ 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], - ] +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 @@ -65,9 +68,7 @@ # define the sampler state sampler_state = SamplerState( - x0=positions, - current_PRNG_key=PRNG.get_random_key(), - box_vectors=box_vectors * unit.nanometer, + x0=positions, current_PRNG_key=PRNG.get_random_key(), box_vectors=box_vectors ) @@ -90,53 +91,54 @@ # # sampler_state.x0 = min_x -from chiron.reporters import MCReporter, LangevinDynamicsReporter +from chiron.reporters import MCReporter # initialize a reporter to save the simulation data -filename_barostat = "test_mc_lj_barostat.h5" -filename_displacement = "test_mc_lj_disp.h5" -filename_langevin = "test_mc_lj_langevin.h5" - import os +filename_barostat = "test_mc_lj_barostat.h5" if os.path.isfile(filename_barostat): os.remove(filename_barostat) reporter_barostat = MCReporter(filename_barostat, 1) -if os.path.isfile(filename_displacement): - os.remove(filename_displacement) -reporter_displacement = MCReporter(filename_displacement, 10) - -if os.path.isfile(filename_langevin): - os.remove(filename_langevin) -reporter_langevin = LangevinDynamicsReporter(filename_langevin, 10) - -from chiron.mcmc import ( - MetropolisDisplacementMove, - MonteCarloBarostatMove, - LangevinDynamicsMove, - MoveSchedule, - MCMCSampler, -) +from chiron.mcmc import MetropolisDisplacementMove mc_displacement_move = MetropolisDisplacementMove( displacement_sigma=0.001 * unit.nanometer, - nr_of_moves=100, + number_of_moves=100, reporter=reporter_displacement, report_frequency=10, - update_stepsize=True, - update_stepsize_frequency=100, + autotune=True, + autotune_interval=100, ) +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 MonteCarloBarostatMove + mc_barostat_move = MonteCarloBarostatMove( volume_max_scale=0.1, - nr_of_moves=10, + number_of_moves=10, reporter=reporter_barostat, report_frequency=1, - update_stepsize=True, - update_stepsize_frequency=50, + 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( stepsize=1.0 * unit.femtoseconds, collision_rate=1.0 / unit.picoseconds, @@ -145,16 +147,17 @@ report_frequency=10, ) +from chiron.mcmc import MoveSchedule move_set = MoveSchedule( [ ("LangevinDynamicsMove", langevin_dynamics_move), - # ("MetropolisDisplacementMove", mc_displacement_move), + ("MetropolisDisplacementMove", 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 -) # how many times to repeat +sampler.run(sampler_state, thermodynamic_state, n_iterations=100, nbr_list=nbr_list) diff --git a/Examples/LJ_langevin.py b/Examples/LJ_langevin.py index 6a34f0d..442ba37 100644 --- a/Examples/LJ_langevin.py +++ b/Examples/LJ_langevin.py @@ -8,7 +8,7 @@ from chiron.potential import LJPotential from openmm import unit -from chiron.utils import PRNG, initialize_velocities +from chiron.utils import PRNG # initialize the LennardJones potential in chiron @@ -32,10 +32,6 @@ box_vectors=lj_fluid.system.getDefaultPeriodicBoxVectors(), ) -velocities = initialize_velocities( - 300 * unit.kelvin, lj_fluid.topology, PRNG.get_random_key() -) - # define the thermodynamic state thermodynamic_state = ThermodynamicState( @@ -73,9 +69,7 @@ from chiron.integrators import LangevinIntegrator # initialize the Langevin integrator -integrator = LangevinIntegrator( - reporter=reporter, report_frequency=100, reinitialize_velocities=True -) +integrator = LangevinIntegrator(reporter=reporter, report_frequency=100) print("init_energy: ", lj_potential.compute_energy(sampler_state.x0, nbr_list)) updated_sampler_state, updated_nbr_list = integrator.run( diff --git a/Examples/LJ_mcmove.py b/Examples/LJ_mcmove.py index c4ac913..6b75e5b 100644 --- a/Examples/LJ_mcmove.py +++ b/Examples/LJ_mcmove.py @@ -64,11 +64,11 @@ mc_move = MetropolisDisplacementMove( displacement_sigma=0.01 * unit.nanometer, - nr_of_moves=5000, + number_of_moves=5000, reporter=reporter, report_frequency=1, - update_stepsize=True, - update_stepsize_frequency=100, + autotune=True, + autotune_interval=100, ) mc_move.update(sampler_state, thermodynamic_state, nbr_list) diff --git a/chiron/mcmc.py b/chiron/mcmc.py index dea4677..231cbc2 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -11,7 +11,7 @@ class MCMCMove: def __init__( self, - nr_of_moves: int, + number_of_moves: int, reporter: Optional[_SimulationReporter] = None, report_frequency: Optional[int] = 100, ): @@ -20,7 +20,7 @@ def __init__( Parameters ---------- - nr_of_moves : int + number_of_moves : int Number of moves to be applied. reporter : _SimulationReporter, optional Reporter object for saving the simulation data. @@ -31,7 +31,7 @@ def __init__( """ - self.nr_of_moves = nr_of_moves + self.number_of_moves = number_of_moves self.reporter = reporter self.report_frequency = report_frequency @@ -115,7 +115,7 @@ def __init__( Default is False. NOTE: Only for debugging purposes. """ super().__init__( - nr_of_moves=nr_of_steps, + number_of_moves=number_of_moves, reporter=reporter, report_frequency=report_frequency, ) @@ -173,7 +173,7 @@ def update( updated_sampler_state, updated_nbr_list = self.integrator.run( thermodynamic_state=thermodynamic_state, sampler_state=sampler_state, - n_steps=self.nr_of_moves, + n_steps=self.number_of_moves, nbr_list=nbr_list, ) @@ -190,11 +190,11 @@ def update( class MCMove(MCMCMove): def __init__( self, - nr_of_moves: int, + number_of_moves: int, reporter: Optional[_SimulationReporter], report_frequency: int = 1, - update_stepsize: bool = False, - update_stepsize_frequency: int = 100, + autotune: bool = False, + autotune_interval: int = 100, method: str = "metropolis", ) -> None: """ @@ -202,31 +202,31 @@ def __init__( Parameters ---------- - nr_of_moves + number_of_moves Number of moves to be applied in each call to update. reporter Reporter object for saving the simulation step data. report_frequency Frequency of saving the simulation data. - update_stepsize - Whether to update the "stepsize" of the move. Stepsize is a generic term for the key move parameters. - For example, for a simple displacement move this would be the displacement_sigma. - update_stepsize_frequency - Frequency of updating the stepsize of the move. + autotune + Whether to automatically tune the parameters of the MC move to achieve a target acceptance ratio. + For example, for a simple displacement move this would update the displacement_sigma. + autotune_interval + Frequency of autotuning the MC move parameters to achieve a target acceptance ratio. method Methodology to use for accepting or rejecting the proposed state. Default is "metropolis". """ super().__init__( - nr_of_moves, + number_of_moves, reporter=reporter, report_frequency=report_frequency, ) self.method = method # I think we should pass a class/function instead of a string, like space. self.reset_statistics() - self.update_stepsize = update_stepsize - self.update_stepsize_frequency = update_stepsize_frequency + self.autotune = autotune + self.autotune_interval = autotune_interval def update( self, @@ -257,7 +257,7 @@ def update( """ calculate_current_reduced_potential = True - for i in range(self.nr_of_moves): + for i in range(self.number_of_moves): sampler_state, thermodynamic_state, nbr_list = self._step( sampler_state, thermodynamic_state, @@ -267,10 +267,10 @@ def update( # after the first step, we don't need to recalculate the current reduced_potential, it will be stored calculate_current_reduced_potential = False - # I think it makes sense to use i + self.nr_of_moves*self._move_iteration as our current "step" - # otherwise, if we just used i, instances where self.report_frequency > self.nr_of_moves would only report on the + # I think it makes sense to use i + self.number_of_moves*self._move_iteration as our current "step" + # otherwise, if we just used i, instances where self.report_frequency > self.number_of_moves would only report on the # first step, which might actually be more frequent than we specify - elapsed_step = i + self._move_iteration * self.nr_of_moves + elapsed_step = i + self._move_iteration * self.number_of_moves if hasattr(self, "reporter"): if self.reporter is not None: if elapsed_step % self.report_frequency == 0: @@ -283,13 +283,10 @@ def update( thermodynamic_state, nbr_list, ) - if self.update_stepsize: + if self.autotune: # if we only used i, we might never actually update the parameters if we have a move that is called infrequently - if ( - elapsed_step % self.update_stepsize_frequency == 0 - and elapsed_step > 0 - ): - self._update_stepsize() + if elapsed_step % self.autotune_interval == 0 and elapsed_step > 0: + self._autotune() # keep track of how many times this function has been called self._move_iteration += 1 @@ -332,15 +329,15 @@ def _report( pass @abstractmethod - def _update_stepsize(self): + def _autotune(self): """ - Update the "stepsize" for a move to reach a target acceptance probability range. + This will autotune the move parameters to reach a target acceptance probability. This will be specific to the type of move, e.g., a displacement_sigma for a displacement move or a maximum volume change factor for a Monte Carlo barostat move. Since different moves will be modifying different quantities, this needs to be defined for each move. - Note this will modify the "stepsize" in place. + Note this will modify the class parameters in place. """ pass @@ -543,12 +540,12 @@ class MetropolisDisplacementMove(MCMove): def __init__( self, displacement_sigma=1.0 * unit.nanometer, - nr_of_moves: int = 100, + number_of_moves: int = 100, atom_subset: Optional[List[int]] = None, report_frequency: int = 1, reporter: Optional[LangevinDynamicsReporter] = None, - update_stepsize: bool = True, - update_stepsize_frequency: int = 100, + autotune: bool = False, + autotune_interval: int = 100, ): """ Initialize the Displacement Move class. @@ -557,26 +554,27 @@ def __init__( ---------- displacement_sigma : float or unit.Quantity, optional The standard deviation of the displacement for each move. Default is 1.0 nm. - nr_of_moves : int, optional + number_of_moves : int, optional The number of moves to perform. Default is 100. atom_subset : list of int, optional A subset of atom indices to consider for the moves. Default is None. reporter : SimulationReporter, optional The reporter to write the data to. Default is None. - update_stepsize : bool, optional - Whether to update the stepsize of the move. Default is True. - update_stepsize_frequency : int, optional - Frequency of updating the stepsize of the move. Default is 100. + autotune : bool, optional + Whether to autotune the displacement_sigma of the move to achieve an acceptance ratio between 0.4 and 0.6. + Default is False. + autotune_interval : int, optional + Frequency of autotuning displacement_sigma of the move. Default is 100. Returns ------- None """ super().__init__( - nr_of_moves=nr_of_moves, + number_of_moves=number_of_moves, reporter=reporter, report_frequency=report_frequency, - update_stepsize=update_stepsize, - update_stepsize_frequency=update_stepsize_frequency, + autotune=autotune, + autotune_interval=autotune_interval, method="metropolis", ) self.displacement_sigma = displacement_sigma @@ -631,9 +629,9 @@ def _report( } ) - def _update_stepsize(self): + def _autotune(self): """ - Update the displacement_sigma to reach a target acceptance probability of 0.5. + Update the displacement_sigma to reach a target acceptance probability between 0.4 and 0.6. """ acceptance_ratio = self.n_accepted / self.n_proposed if acceptance_ratio > 0.6: @@ -754,37 +752,38 @@ class MonteCarloBarostatMove(MCMove): def __init__( self, volume_max_scale=0.01, - nr_of_moves: int = 100, + number_of_moves: int = 100, report_frequency: int = 1, reporter: Optional[LangevinDynamicsReporter] = None, - update_stepsize: bool = True, - update_stepsize_frequency: int = 100, + autotune: bool = False, + autotune_interval: int = 100, ): """ Initialize the Monte Carlo Barostat Move class. Parameters ---------- - displacement_sigma : float or unit.Quantity, optional - The standard deviation of the displacement for each move. Default is 1.0 nm. - nr_of_moves : int, optional + volume_max_scale : float, optional + The scaling factor multiplied by volume to set the maximum volume change allowed. + number_of_moves : int, optional The number of moves to perform. Default is 100. reporter : SimulationReporter, optional The reporter to write the data to. Default is None. - update_stepsize : bool, optional - Whether to update the stepsize of the move. Default is True. - update_stepsize_frequency : int, optional - Frequency of updating the stepsize of the move. Default is 100. + autotune : bool, optional + Whether to autotune the volume_max_scale value of the move to achieve a target probability + between 0.25 and 0.75. Default is False. volume_max_scale is capped at 0.3 + autotune_interval : int, optional + Frequency of autotuning the volume_max_scale of the move. Default is 100. Returns ------- None """ super().__init__( - nr_of_moves=nr_of_moves, + number_of_moves=number_of_moves, reporter=reporter, report_frequency=report_frequency, - update_stepsize=update_stepsize, - update_stepsize_frequency=update_stepsize_frequency, + autotune=autotune, + autotune_interval=autotune_interval, method="metropolis", ) self.volume_max_scale = volume_max_scale @@ -840,7 +839,7 @@ def _report( } ) - def _update_stepsize(self): + def _autotune(self): """ Update the volume_max_scale parameter to ensure our acceptance probability is within the range of 0.25 to 0.75. The maximum volume_max_scale will be capped at 0.3. diff --git a/chiron/tests/test_convergence_tests.py b/chiron/tests/test_convergence_tests.py index c05c072..1bd1ae9 100644 --- a/chiron/tests/test_convergence_tests.py +++ b/chiron/tests/test_convergence_tests.py @@ -66,7 +66,7 @@ def test_convergence_of_MC_estimator(prep_temp_dir): from chiron.mcmc import MetropolisDisplacementMove, MoveSchedule, MCMCSampler mc_displacement_move = MetropolisDisplacementMove( - nr_of_moves=1_000, + number_of_moves=1_000, displacement_sigma=0.5 * unit.angstrom, atom_subset=[0], reporter=simulation_reporter, diff --git a/chiron/tests/test_mcmc.py b/chiron/tests/test_mcmc.py index 3c182c3..b63b3d4 100644 --- a/chiron/tests/test_mcmc.py +++ b/chiron/tests/test_mcmc.py @@ -207,7 +207,7 @@ def test_sample_from_harmonic_osciallator_with_MCMC_classes_and_MetropolisDispla simulation_reporter = MCReporter(1) mc_displacement_move = MetropolisDisplacementMove( - nr_of_moves=10, + number_of_moves=10, displacement_sigma=0.1 * unit.angstrom, atom_subset=[0], reporter=simulation_reporter, @@ -269,7 +269,7 @@ def test_sample_from_harmonic_osciallator_array_with_MCMC_classes_and_Metropolis simulation_reporter = MCReporter(1) mc_displacement_move = MetropolisDisplacementMove( - nr_of_moves=10, + number_of_moves=10, displacement_sigma=0.1 * unit.angstrom, atom_subset=None, reporter=simulation_reporter, @@ -330,11 +330,11 @@ def test_mc_barostat_parameter_setting(): barostat_move = MonteCarloBarostatMove( volume_max_scale=0.1, - nr_of_moves=1, + number_of_moves=1, ) assert barostat_move.volume_max_scale == 0.1 - assert barostat_move.nr_of_moves == 1 + assert barostat_move.number_of_moves == 1 def test_mc_barostat(prep_temp_dir): @@ -350,7 +350,7 @@ def test_mc_barostat(prep_temp_dir): barostat_move = MonteCarloBarostatMove( volume_max_scale=0.1, - nr_of_moves=10, + number_of_moves=10, reporter=simulation_reporter, report_frequency=1, ) diff --git a/chiron/tests/test_testsystems.py b/chiron/tests/test_testsystems.py index 953effd..4f7b758 100644 --- a/chiron/tests/test_testsystems.py +++ b/chiron/tests/test_testsystems.py @@ -285,18 +285,18 @@ def test_ideal_gas(prep_temp_dir): mc_displacement_move = MetropolisDisplacementMove( displacement_sigma=0.1 * unit.nanometer, - nr_of_moves=10, + number_of_moves=10, reporter=reporter, - update_stepsize=True, - update_stepsize_frequency=100, + autotune=True, + autotune_interval=100, ) mc_barostat_move = MonteCarloBarostatMove( volume_max_scale=0.2, - nr_of_moves=100, + number_of_moves=100, reporter=reporter, - update_stepsize=True, - update_stepsize_frequency=100, + autotune=True, + autotune_interval=100, ) move_set = MoveSchedule( [ From 07fe3682d048cb935ae6a72bf5c49d5c54041a3f Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Tue, 20 Feb 2024 09:33:24 -0800 Subject: [PATCH 37/43] Fixed error with multistate systems. Various other updates based on comments. --- Examples/Idealgas.py | 10 +- Examples/LJ_MCMC.py | 22 ++-- Examples/LJ_langevin.py | 8 +- Examples/LJ_mcmove.py | 8 +- chiron/integrators.py | 66 ++++++------ chiron/mcmc.py | 134 ++++++++++++++----------- chiron/multistate.py | 17 ++-- chiron/neighbors.py | 14 +-- chiron/potential.py | 20 ++-- chiron/states.py | 46 +++++---- chiron/tests/test_convergence_tests.py | 16 +-- chiron/tests/test_integrators.py | 4 +- chiron/tests/test_mcmc.py | 32 +++--- chiron/tests/test_minization.py | 14 +-- chiron/tests/test_multistate.py | 11 +- chiron/tests/test_pairs.py | 14 +-- chiron/tests/test_potential.py | 2 +- chiron/tests/test_states.py | 26 ++--- chiron/tests/test_testsystems.py | 12 +-- chiron/tests/test_utils.py | 4 +- 20 files changed, 260 insertions(+), 220 deletions(-) diff --git a/Examples/Idealgas.py b/Examples/Idealgas.py index bbeffe8..d6830ce 100644 --- a/Examples/Idealgas.py +++ b/Examples/Idealgas.py @@ -4,7 +4,7 @@ """ 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 -MetropolisDisplacementMove to sample the particle positions. +MonteCarloDisplacementMove to sample the particle positions. This utilizes the IdealGas example from openmmtools to initialize particle positions and topology. @@ -42,7 +42,7 @@ # define the sampler state sampler_state = SamplerState( - x0=ideal_gas.positions, + positions=ideal_gas.positions, current_PRNG_key=PRNG.get_random_key(), box_vectors=ideal_gas.system.getDefaultPeriodicBoxVectors(), ) @@ -67,7 +67,7 @@ from chiron.mcmc import ( - MetropolisDisplacementMove, + MonteCarloDisplacementMove, MonteCarloBarostatMove, MoveSchedule, MCMCSampler, @@ -83,7 +83,7 @@ ) # initialize the barostat move and the move schedule -metropolis_displacement_move = MetropolisDisplacementMove( +metropolis_displacement_move = MonteCarloDisplacementMove( displacement_sigma=0.1 * unit.nanometer, number_of_moves=100, autotune=True, @@ -93,7 +93,7 @@ # define the move schedule move_set = MoveSchedule( [ - ("MetropolisDisplacementMove", metropolis_displacement_move), + ("MonteCarloDisplacementMove", metropolis_displacement_move), ("MonteCarloBarostatMove", mc_barostat_move), ] ) diff --git a/Examples/LJ_MCMC.py b/Examples/LJ_MCMC.py index 1ba26dc..b2e8684 100644 --- a/Examples/LJ_MCMC.py +++ b/Examples/LJ_MCMC.py @@ -68,7 +68,7 @@ # define the sampler state sampler_state = SamplerState( - x0=positions, current_PRNG_key=PRNG.get_random_key(), box_vectors=box_vectors + positions=positions, current_PRNG_key=PRNG.get_random_key(), box_vectors=box_vectors ) @@ -84,12 +84,12 @@ # from chiron.minimze import minimize_energy # # results = minimize_energy( -# sampler_state.x0, lj_potential.compute_energy, nbr_list, maxiter=100 +# sampler_state.positions, lj_potential.compute_energy, nbr_list, maxiter=100 # ) # # min_x = results.params # -# sampler_state.x0 = min_x +# sampler_state.positions = min_x from chiron.reporters import MCReporter @@ -101,13 +101,13 @@ os.remove(filename_barostat) reporter_barostat = MCReporter(filename_barostat, 1) -from chiron.mcmc import MetropolisDisplacementMove +from chiron.mcmc import MonteCarloDisplacementMove -mc_displacement_move = MetropolisDisplacementMove( +mc_displacement_move = MonteCarloDisplacementMove( displacement_sigma=0.001 * unit.nanometer, number_of_moves=100, reporter=reporter_displacement, - report_frequency=10, + report_interval=10, autotune=True, autotune_interval=100, ) @@ -124,7 +124,7 @@ volume_max_scale=0.1, number_of_moves=10, reporter=reporter_barostat, - report_frequency=1, + report_interval=1, autotune=True, autotune_interval=50, ) @@ -140,11 +140,11 @@ from chiron.mcmc import LangevinDynamicsMove langevin_dynamics_move = LangevinDynamicsMove( - stepsize=1.0 * unit.femtoseconds, + timestep=1.0 * unit.femtoseconds, collision_rate=1.0 / unit.picoseconds, - nr_of_steps=100, + number_of_steps=100, reporter=reporter_langevin, - report_frequency=10, + report_interval=10, ) from chiron.mcmc import MoveSchedule @@ -152,7 +152,7 @@ move_set = MoveSchedule( [ ("LangevinDynamicsMove", langevin_dynamics_move), - ("MetropolisDisplacementMove", mc_displacement_move), + ("MonteCarloDisplacementMove", mc_displacement_move), ("MonteCarloBarostatMove", mc_barostat_move), ] ) diff --git a/Examples/LJ_langevin.py b/Examples/LJ_langevin.py index 442ba37..1086cdc 100644 --- a/Examples/LJ_langevin.py +++ b/Examples/LJ_langevin.py @@ -27,7 +27,7 @@ PRNG.set_seed(1234) # define the sampler state sampler_state = SamplerState( - x0=lj_fluid.positions, + positions=lj_fluid.positions, current_PRNG_key=PRNG.get_random_key(), box_vectors=lj_fluid.system.getDefaultPeriodicBoxVectors(), ) @@ -69,13 +69,13 @@ from chiron.integrators import LangevinIntegrator # initialize the Langevin integrator -integrator = LangevinIntegrator(reporter=reporter, report_frequency=100) -print("init_energy: ", lj_potential.compute_energy(sampler_state.x0, nbr_list)) +integrator = LangevinIntegrator(reporter=reporter, report_interval=100) +print("init_energy: ", lj_potential.compute_energy(sampler_state.positions, nbr_list)) updated_sampler_state, updated_nbr_list = integrator.run( sampler_state, thermodynamic_state, - n_steps=1000, + number_of_steps=1000, nbr_list=nbr_list, progress_bar=True, ) diff --git a/Examples/LJ_mcmove.py b/Examples/LJ_mcmove.py index 6b75e5b..fac45d4 100644 --- a/Examples/LJ_mcmove.py +++ b/Examples/LJ_mcmove.py @@ -26,7 +26,7 @@ # define the sampler state sampler_state = SamplerState( - x0=lj_fluid.positions, + positions=lj_fluid.positions, current_PRNG_key=PRNG.get_random_key(), box_vectors=lj_fluid.system.getDefaultPeriodicBoxVectors(), ) @@ -60,13 +60,13 @@ os.remove(filename) reporter = MCReporter(filename, 1) -from chiron.mcmc import MetropolisDisplacementMove +from chiron.mcmc import MonteCarloDisplacementMove -mc_move = MetropolisDisplacementMove( +mc_move = MonteCarloDisplacementMove( displacement_sigma=0.01 * unit.nanometer, number_of_moves=5000, reporter=reporter, - report_frequency=1, + report_interval=1, autotune=True, autotune_interval=100, ) diff --git a/chiron/integrators.py b/chiron/integrators.py index 3760e38..9228529 100644 --- a/chiron/integrators.py +++ b/chiron/integrators.py @@ -24,10 +24,10 @@ class LangevinIntegrator: def __init__( self, - stepsize=1.0 * unit.femtoseconds, + timestep=1.0 * unit.femtoseconds, collision_rate=1.0 / unit.picoseconds, - reinitialize_velocities: bool = False, - report_frequency: int = 100, + refresh_velocities: bool = False, + report_interval: int = 100, reporter: Optional[LangevinDynamicsReporter] = None, save_traj_in_memory: bool = False, ) -> None: @@ -36,14 +36,14 @@ def __init__( Parameters ---------- - stepsize : unit.Quantity, optional + timestep : unit.Quantity, optional Time step of integration with units of time. Default is 1.0 * unit.femtoseconds. collision_rate : unit.Quantity, optional Collision rate for the Langevin dynamics, with units 1/time. Default is 1.0 / unit.picoseconds. - reinitialize_velocities : bool, optional + refresh_velocities : bool, optional Flag indicating whether to reinitialize the velocities each time the run function is called. Default is False. - report_frequency : int, optional - Frequency of saving the simulation data. Default is 100. + report_interval : int, optional + Interval between saving the simulation data. Default is 100. reporter : SimulationReporter, optional Reporter object for saving the simulation data. Default is None. save_traj_in_memory: bool @@ -53,11 +53,11 @@ def __init__( from loguru import logger as log self.kB = unit.BOLTZMANN_CONSTANT_kB * unit.AVOGADRO_CONSTANT_NA - log.info(f"stepsize = {stepsize}") + log.info(f"timestep = {timestep}") log.info(f"collision_rate = {collision_rate}") - log.info(f"report_frequency = {report_frequency}") + log.info(f"report_interval = {report_interval}") - self.stepsize = stepsize + self.timestep = timestep self.collision_rate = collision_rate if reporter: log.info( @@ -65,18 +65,18 @@ def __init__( ) log.info(f"and logging to {reporter.log_file_path}") self.reporter = reporter - self.report_frequency = report_frequency + self.report_interval = report_interval self.velocities = None self.save_traj_in_memory = save_traj_in_memory self.traj = [] - self.reinitialize_velocities = reinitialize_velocities + self.refresh_velocities = refresh_velocities self._move_iteration = 0 def run( self, sampler_state: SamplerState, thermodynamic_state: ThermodynamicState, - n_steps: int = 5_000, + number_of_steps: int = 5_000, nbr_list: Optional[PairsBase] = None, progress_bar=False, ) -> Tuple[SamplerState, PairsBase]: @@ -89,7 +89,7 @@ def run( The initial state of the simulation, including positions. thermodynamic_state : ThermodynamicState The thermodynamic state of the system, including temperature and potential. - n_steps : int, optional + number_of_steps : int, optional Number of simulation steps to perform. nbr_list : PairBase, optional Neighbor list for the system. @@ -114,10 +114,10 @@ def run( self.box_vectors = sampler_state.box_vectors self.progress_bar = progress_bar temperature = thermodynamic_state.temperature - x0 = sampler_state.x0 + x0 = sampler_state.positions log.debug("Running Langevin dynamics") - log.debug(f"n_steps = {n_steps}") + log.debug(f"number_of_steps = {number_of_steps}") log.debug(f"temperature = {temperature}") # Initialize the random number generator @@ -129,16 +129,16 @@ def run( :, None ] sigma_v = jnp.sqrt(kbT_unitless / mass_unitless) - stepsize_unitless = self.stepsize.value_in_unit_system(unit.md_unit_system) + timestep_unitless = self.timestep.value_in_unit_system(unit.md_unit_system) collision_rate_unitless = self.collision_rate.value_in_unit_system( unit.md_unit_system ) - a = jnp.exp((-collision_rate_unitless * stepsize_unitless)) - b = jnp.sqrt(1 - jnp.exp(-2 * collision_rate_unitless * stepsize_unitless)) + a = jnp.exp((-collision_rate_unitless * timestep_unitless)) + b = jnp.sqrt(1 - jnp.exp(-2 * collision_rate_unitless * timestep_unitless)) # Initialize velocities - if self.reinitialize_velocities: - # v0 = sigma_v * random.normal(key, x0.shape) + if self.refresh_velocities: + # v0 = sigma_v * random.normal(key, positions.shape) from .utils import initialize_velocities sampler_state.velocities = initialize_velocities( @@ -146,13 +146,13 @@ def run( ) elif sampler_state._velocities is None: - # v0 = sigma_v * random.normal(key, x0.shape) + # v0 = sigma_v * random.normal(key, positions.shape) from .utils import initialize_velocities sampler_state.velocities = initialize_velocities( temperature, potential.topology, key ) - elif sampler_state._velocities.shape[0] != sampler_state.x0.shape[0]: + elif sampler_state._velocities.shape[0] != sampler_state.positions.shape[0]: from .utils import initialize_velocities sampler_state.velocities = initialize_velocities( @@ -171,12 +171,16 @@ def run( F = potential.compute_force(x, nbr_list) # propagation loop - for step in tqdm(range(n_steps)) if self.progress_bar else range(n_steps): + for step in ( + tqdm(range(number_of_steps)) + if self.progress_bar + else range(number_of_steps) + ): key, subkey = random.split(key) # v - v += (stepsize_unitless * 0.5) * F / mass_unitless + v += (timestep_unitless * 0.5) * F / mass_unitless # r - x += (stepsize_unitless * 0.5) * v + x += (timestep_unitless * 0.5) * v # we can actually skip this wrapping, and just wrap/check/rebuild # right before we call the force again. @@ -187,17 +191,17 @@ def run( random_noise_v = random.normal(subkey, x.shape) v = (a * v) + (b * sigma_v * random_noise_v) - x += (stepsize_unitless * 0.5) * v + x += (timestep_unitless * 0.5) * v if nbr_list is not None: x, nbr_list = self._wrap_and_rebuild_neighborlist(x, nbr_list) F = potential.compute_force(x, nbr_list) # v - v += (stepsize_unitless * 0.5) * F / mass_unitless + v += (timestep_unitless * 0.5) * F / mass_unitless - elapsed_step = step + self._move_iteration * n_steps - if (elapsed_step) % self.report_frequency == 0: + elapsed_step = step + self._move_iteration * number_of_steps + if (elapsed_step) % self.report_interval == 0: if hasattr(self, "reporter") and self.reporter is not None: self._report( x, potential, nbr_list, step, self._move_iteration, elapsed_step @@ -213,7 +217,7 @@ def run( updated_sampler_state = copy.deepcopy(sampler_state) - updated_sampler_state.x0 = x + updated_sampler_state.positions = x updated_sampler_state.velocities = v updated_sampler_state.current_PRNG_key = key diff --git a/chiron/mcmc.py b/chiron/mcmc.py index 231cbc2..cbeb38a 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -13,7 +13,7 @@ def __init__( self, number_of_moves: int, reporter: Optional[_SimulationReporter] = None, - report_frequency: Optional[int] = 100, + report_interval: Optional[int] = 100, ): """ Initialize a move within the molecular system. @@ -25,15 +25,15 @@ def __init__( reporter : _SimulationReporter, optional Reporter object for saving the simulation data. Default is None. - report_frequency : int, optional - Frequency of saving the simulation data in the reporter. + report_interval : int, optional + Interval for saving the simulation data in the reporter. Default is 100. """ self.number_of_moves = number_of_moves self.reporter = reporter - self.report_frequency = report_frequency + self.report_interval = report_interval # we need to keep track of which iteration we are on self._move_iteration = 0 @@ -44,7 +44,7 @@ def __init__( log.info( f"Using reporter {self.reporter} saving to {self.reporter.workdir}" ) - assert self.report_frequency is not None + assert self.report_interval is not None @abstractmethod def update( @@ -81,12 +81,12 @@ def update( class LangevinDynamicsMove(MCMCMove): def __init__( self, - stepsize: unit.Quantity = 1.0 * unit.femtoseconds, + timestep: unit.Quantity = 1.0 * unit.femtoseconds, collision_rate: unit.Quantity = 1.0 / unit.picoseconds, - reinitialize_velocities: bool = False, + refresh_velocities: bool = False, reporter: Optional[LangevinDynamicsReporter] = None, - report_frequency: int = 100, - nr_of_steps: int = 1_000, + report_interval: int = 100, + number_of_steps: int = 1_000, save_traj_in_memory: bool = False, ): """ @@ -94,20 +94,20 @@ def __init__( Parameters ---------- - stepsize : unit.Quantity + timestep : unit.Quantity Time step size for the integration. collision_rate : unit.Quantity Collision rate for the Langevin dynamics. - reinitialize_velocities : bool, optional + refresh_velocities : bool, optional Whether to reinitialize the velocities each time the run function is called. Default is False. reporter : LangevinDynamicsReporter, optional Reporter object for saving the simulation data. Default is None. - report_frequency : int - Frequency of saving the simulation data. + report_interval : int + Interval for saving the simulation data. Default is 100. - nr_of_steps : int, optional + number_of_steps : int, optional Number of steps to run the integrator for. Default is 1_000. save_traj_in_memory: bool @@ -115,22 +115,22 @@ def __init__( Default is False. NOTE: Only for debugging purposes. """ super().__init__( - number_of_moves=number_of_moves, + number_of_moves=number_of_steps, reporter=reporter, - report_frequency=report_frequency, + report_interval=report_interval, ) - self.stepsize = stepsize + self.timestep = timestep self.collision_rate = collision_rate self.save_traj_in_memory = save_traj_in_memory self.traj = [] from chiron.integrators import LangevinIntegrator self.integrator = LangevinIntegrator( - stepsize=self.stepsize, + timestep=self.timestep, collision_rate=self.collision_rate, - reinitialize_velocities=reinitialize_velocities, - report_frequency=report_frequency, + refresh_velocities=refresh_velocities, + report_interval=report_interval, reporter=reporter, save_traj_in_memory=save_traj_in_memory, ) @@ -173,7 +173,7 @@ def update( updated_sampler_state, updated_nbr_list = self.integrator.run( thermodynamic_state=thermodynamic_state, sampler_state=sampler_state, - n_steps=self.number_of_moves, + number_of_steps=self.number_of_moves, nbr_list=nbr_list, ) @@ -192,10 +192,10 @@ def __init__( self, number_of_moves: int, reporter: Optional[_SimulationReporter], - report_frequency: int = 1, + report_interval: int = 1, autotune: bool = False, autotune_interval: int = 100, - method: str = "metropolis", + acceptance_method: str = "Metropolis-Hastings", ) -> None: """ Initialize the move. @@ -203,26 +203,26 @@ def __init__( Parameters ---------- number_of_moves - Number of moves to be applied in each call to update. + Number of moves to be attempted in each call to update. reporter Reporter object for saving the simulation step data. - report_frequency - Frequency of saving the simulation data. + report_interval + Interval for saving the simulation data. autotune Whether to automatically tune the parameters of the MC move to achieve a target acceptance ratio. For example, for a simple displacement move this would update the displacement_sigma. autotune_interval Frequency of autotuning the MC move parameters to achieve a target acceptance ratio. - method + acceptance_method Methodology to use for accepting or rejecting the proposed state. - Default is "metropolis". + Default is "Metropolis-Hastings". """ super().__init__( - number_of_moves, + number_of_moves=number_of_moves, reporter=reporter, - report_frequency=report_frequency, + report_interval=report_interval, ) - self.method = method # I think we should pass a class/function instead of a string, like space. + self.acceptance_method = acceptance_method # I think we should pass a class/function instead of a string, like space. self.reset_statistics() self.autotune = autotune @@ -268,12 +268,12 @@ def update( calculate_current_reduced_potential = False # I think it makes sense to use i + self.number_of_moves*self._move_iteration as our current "step" - # otherwise, if we just used i, instances where self.report_frequency > self.number_of_moves would only report on the + # otherwise, if we just used i, instances where self.report_interval > self.number_of_moves would only report on the # first step, which might actually be more frequent than we specify elapsed_step = i + self._move_iteration * self.number_of_moves if hasattr(self, "reporter"): if self.reporter is not None: - if elapsed_step % self.report_frequency == 0: + if elapsed_step % self.report_interval == 0: self._report( i, self._move_iteration, @@ -420,7 +420,7 @@ def _step( decision = self._accept_or_reject( log_proposal_ratio, proposed_sampler_state.new_PRNG_key, - method=self.method, + acceptance_method=self.acceptance_method, ) # a function that will update the statistics for the move @@ -520,13 +520,13 @@ def _accept_or_reject( self, log_proposal_ratio, key, - method, + acceptance_method, ): """ Accept or reject the proposed state with a given methodology. """ # define the acceptance probability - if method == "metropolis": + if acceptance_method == "Metropolis-Hastings": import jax.random as jrandom compare_to = jrandom.uniform(key) @@ -536,16 +536,17 @@ def _accept_or_reject( return False -class MetropolisDisplacementMove(MCMove): +class MonteCarloDisplacementMove(MCMove): def __init__( self, displacement_sigma=1.0 * unit.nanometer, number_of_moves: int = 100, atom_subset: Optional[List[int]] = None, - report_frequency: int = 1, - reporter: Optional[LangevinDynamicsReporter] = None, + report_interval: int = 1, + reporter: Optional[MCReporter] = None, autotune: bool = False, autotune_interval: int = 100, + acceptance_method="Metropolis-Hastings", ): """ Initialize the Displacement Move class. @@ -555,9 +556,13 @@ def __init__( displacement_sigma : float or unit.Quantity, optional The standard deviation of the displacement for each move. Default is 1.0 nm. number_of_moves : int, optional - The number of moves to perform. Default is 100. + The number of move attempts to perform. Default is 100. + For a given move, all particles will be randomly displaced at once (unless atom_subset is), + rather than moving each particle one at a time. atom_subset : list of int, optional - A subset of atom indices to consider for the moves. Default is None. + A list of particle indices that represent a subset of all particles. + If defined, only those particles in the list will have their positions random displaced. + Default is None. reporter : SimulationReporter, optional The reporter to write the data to. Default is None. autotune : bool, optional @@ -565,6 +570,10 @@ def __init__( Default is False. autotune_interval : int, optional Frequency of autotuning displacement_sigma of the move. Default is 100. + acceptance_method : str, optional + Methodology to use for accepting or rejecting the proposed state. + Default is "Metropolis-Hastings". + Returns ------- None @@ -572,10 +581,10 @@ def __init__( super().__init__( number_of_moves=number_of_moves, reporter=reporter, - report_frequency=report_frequency, + report_interval=report_interval, autotune=autotune, autotune_interval=autotune_interval, - method="metropolis", + acceptance_method=acceptance_method, ) self.displacement_sigma = displacement_sigma @@ -614,7 +623,7 @@ def _report( """ potential = thermodynamic_state.potential.compute_energy( - sampler_state.x0, nbr_list + sampler_state.positions, nbr_list ) self.reporter.report( { @@ -700,30 +709,30 @@ def _propose( proposed_sampler_state = copy.deepcopy(current_sampler_state) if self.atom_subset is not None: - proposed_sampler_state.x0 = ( - proposed_sampler_state.x0 + proposed_sampler_state.positions = ( + proposed_sampler_state.positions + scaled_displacement_vector * self.atom_subset_mask ) else: - proposed_sampler_state.x0 = ( - proposed_sampler_state.x0 + scaled_displacement_vector + proposed_sampler_state.positions = ( + proposed_sampler_state.positions + scaled_displacement_vector ) # after proposing a move we need to wrap particles and see if we need to rebuild the neighborlist if current_nbr_list is not None: - proposed_sampler_state.x0 = current_nbr_list.space.wrap( - proposed_sampler_state.x0 + proposed_sampler_state.positions = current_nbr_list.space.wrap( + proposed_sampler_state.positions ) # if we need to rebuild the neighbor the neighborlist # we will make a copy and then build - if current_nbr_list.check(proposed_sampler_state.x0): + if current_nbr_list.check(proposed_sampler_state.positions): import copy proposed_nbr_list = copy.deepcopy(current_nbr_list) proposed_nbr_list.build( - proposed_sampler_state.x0, proposed_sampler_state.box_vectors + proposed_sampler_state.positions, proposed_sampler_state.box_vectors ) # if we don't need to update the neighborlist, just make a new variable that refers to the original else: @@ -753,10 +762,11 @@ def __init__( self, volume_max_scale=0.01, number_of_moves: int = 100, - report_frequency: int = 1, + report_interval: int = 1, reporter: Optional[LangevinDynamicsReporter] = None, autotune: bool = False, autotune_interval: int = 100, + acceptance_method="Metropolis-Hastings", ): """ Initialize the Monte Carlo Barostat Move class. @@ -766,7 +776,7 @@ def __init__( volume_max_scale : float, optional The scaling factor multiplied by volume to set the maximum volume change allowed. number_of_moves : int, optional - The number of moves to perform. Default is 100. + The number of volume update moves attempts to perform. Default is 100. reporter : SimulationReporter, optional The reporter to write the data to. Default is None. autotune : bool, optional @@ -774,6 +784,10 @@ def __init__( between 0.25 and 0.75. Default is False. volume_max_scale is capped at 0.3 autotune_interval : int, optional Frequency of autotuning the volume_max_scale of the move. Default is 100. + acceptance_method : str, optional + Methodology to use for accepting or rejecting the proposed state. + Default is "Metropolis-Hastings". + Returns ------- None @@ -781,10 +795,10 @@ def __init__( super().__init__( number_of_moves=number_of_moves, reporter=reporter, - report_frequency=report_frequency, + report_interval=report_interval, autotune=autotune, autotune_interval=autotune_interval, - method="metropolis", + acceptance_method=acceptance_method, ) self.volume_max_scale = volume_max_scale @@ -819,7 +833,7 @@ def _report( """ potential = thermodynamic_state.potential.compute_energy( - sampler_state.x0, nbr_list + sampler_state.positions, nbr_list ) volume = ( sampler_state.box_vectors[0][0] @@ -914,7 +928,9 @@ def _propose( import copy proposed_sampler_state = copy.deepcopy(current_sampler_state) - proposed_sampler_state.x0 = current_sampler_state.x0 * length_scaling_factor + proposed_sampler_state.positions = ( + current_sampler_state.positions * length_scaling_factor + ) proposed_sampler_state.box_vectors = ( current_sampler_state.box_vectors * length_scaling_factor @@ -924,7 +940,7 @@ def _propose( proposed_nbr_list = copy.deepcopy(current_nbr_list) # after scaling the box vectors and coordinates we should always rebuild the neighborlist proposed_nbr_list.build( - proposed_sampler_state.x0, proposed_sampler_state.box_vectors + proposed_sampler_state.positions, proposed_sampler_state.box_vectors ) proposed_reduced_pot = current_thermodynamic_state.get_reduced_potential( diff --git a/chiron/multistate.py b/chiron/multistate.py index d3c782a..e67d7bb 100644 --- a/chiron/multistate.py +++ b/chiron/multistate.py @@ -323,14 +323,14 @@ def _minimize_replica( # Perform minimization minimized_state = minimize_energy( - sampler_state.x0, + sampler_state.positions, thermodynamic_state.potential.compute_energy, self.nbr_list, maxiter=max_iterations, ) # Update the sampler state - self._sampler_states[replica_id].x0 = minimized_state.params + self._sampler_states[replica_id].positions = minimized_state.params # Compute and log final energy final_energy = thermodynamic_state.get_reduced_potential(sampler_state) @@ -398,9 +398,14 @@ def _propagate_replica(self, replica_id: int): mcmc_sampler = self._mcmc_sampler[thermodynamic_state_id] # Propagate using the mcmc sampler - self._sampler_states[replica_id] = mcmc_sampler.run(sampler_state, thermodynamic_state) + # NOTE this needs to be updated to support neighborlists + ( + self._sampler_states[replica_id], + self._thermodynamic_states[thermodynamic_state_id], + nbr_list, + ) = mcmc_sampler.run(sampler_state, thermodynamic_state) # Append the new state to the trajectory for analysis. - self._traj[replica_id].append(self._sampler_states[replica_id].x0) + self._traj[replica_id].append(self._sampler_states[replica_id].positions) def _perform_swap_proposals(self): """ @@ -579,9 +584,9 @@ def _report_positions(self): log.debug("Reporting positions...") # numpy array with shape (n_replicas, n_atoms, 3) - xyz = np.zeros((self.n_replicas, self._sampler_states[0].x0.shape[0], 3)) + xyz = np.zeros((self.n_replicas, self._sampler_states[0].positions.shape[0], 3)) for replica_id in range(self.n_replicas): - xyz[replica_id] = self._sampler_states[replica_id].x0 + xyz[replica_id] = self._sampler_states[replica_id].positions return {"positions": xyz} def _report(self, property: str) -> None: diff --git a/chiron/neighbors.py b/chiron/neighbors.py index 9941e44..64b897d 100644 --- a/chiron/neighbors.py +++ b/chiron/neighbors.py @@ -207,13 +207,13 @@ class PairsBase(ABC): >>> import jax.numpy as jnp >>> >>> space = OrthogonalPeriodicSpace() # define the simulation space, in this case an orthogonal periodic space - >>> sampler_state = SamplerState(x0=jnp.array([[0.0, 0.0, 0.0], [2, 0.0, 0.0], [0.0, 2, 0.0]]), + >>> sampler_state = SamplerState(positions=jnp.array([[0.0, 0.0, 0.0], [2, 0.0, 0.0], [0.0, 2, 0.0]]), >>> box_vectors=jnp.array([[10, 0.0, 0.0], [0.0, 10, 0.0], [0.0, 0.0, 10]])) >>> >>> pair_list = PairsBase(space, cutoff=2.5*unit.nanometer) # initialize the pair list >>> pair_list.build_from_state(sampler_state) # build the pair list from the sampler state >>> - >>> coordinates = sampler_state.x0 # get the coordinates from the sampler state, without units attached + >>> coordinates = sampler_state.positions # get the coordinates from the sampler state, without units attached >>> >>> # the calculate function will produce information used to calculate the energy >>> n_neighbors, padding_mask, dist, r_ij = pair_list.calculate(coordinates) @@ -309,7 +309,7 @@ def build_from_state(self, sampler_state: SamplerState): if not isinstance(sampler_state, SamplerState): raise TypeError(f"Expected SamplerState, got {type(sampler_state)} instead") - coordinates = sampler_state.x0 + coordinates = sampler_state.positions if sampler_state.box_vectors is None: raise ValueError(f"SamplerState does not contain box vectors") box_vectors = sampler_state.box_vectors @@ -526,7 +526,7 @@ def build( """ # set our reference coordinates - # the call to x0 and box_vectors automatically convert these to jnp arrays in the correct unit system + # the call to positions and box_vectors automatically convert these to jnp arrays in the correct unit system if isinstance(coordinates, unit.Quantity): if not coordinates.unit.is_compatible(unit.nanometer): raise ValueError( @@ -669,7 +669,7 @@ def calculate(self, coordinates: jnp.array): r_ij: jnp.array Array of displacement vectors between each particle and its neighbors. Shape (n_particles, n_max_neighbors, 3) """ - # coordinates = sampler_state.x0 + # coordinates = sampler_state.positions # note, we assume the box vectors do not change between building and calculating the neighbor list # changes to the box vectors require rebuilding the neighbor list @@ -761,7 +761,7 @@ class PairList(PairsBase): >>> >>> space = OrthogonalPeriodicSpace() >>> pair_list = PairList(space, cutoff=2.5) - >>> sampler_state = SamplerState(x0=jnp.array([[0.0, 0.0, 0.0], [2, 0.0, 0.0], [0.0, 2, 0.0]]), + >>> sampler_state = SamplerState(positions=jnp.array([[0.0, 0.0, 0.0], [2, 0.0, 0.0], [0.0, 2, 0.0]]), >>> box_vectors=jnp.array([[10, 0.0, 0.0], [0.0, 10, 0.0], [0.0, 0.0, 10]])) >>> pair_list.build_from_state(sampler_state) >>> @@ -769,7 +769,7 @@ class PairList(PairsBase): >>> displacement_vectors of shape (n_particles, n_particles-1, 3) >>> # mask, is a bool array that is True if the particle is within the cutoff distance, False if it is not >>> # n_pairs is of shape (n_particles) and is per row sum of the mask. The mask ensure we also do not double count pairs - >>> n_pairs, mask, distances, displacement_vectors = pair_list.calculate(sampler_state.x0) + >>> n_pairs, mask, distances, displacement_vectors = pair_list.calculate(sampler_state.positions) """ def __init__( diff --git a/chiron/potential.py b/chiron/potential.py index 097fea5..d567a4c 100644 --- a/chiron/potential.py +++ b/chiron/potential.py @@ -349,7 +349,7 @@ def __init__( The topology object representing the molecular system. k : unit.Quantity, optional The spring constant of the harmonic potential. Default is 1.0 kcal/mol/Å^2. - x0 : unit.Quantity, optional + positions : unit.Quantity, optional The equilibrium position of the harmonic potential. Default is [0.0,0.0,0.0] Å. U0 : unit.Quantity, optional The offset potential energy of the harmonic potential. Default is 0.0 kcal/mol. @@ -366,7 +366,9 @@ def __init__( if not isinstance(k, unit.Quantity): raise TypeError(f"k must be a unit.Quantity, type(k) = {type(k)}") if not isinstance(x0, unit.Quantity): - raise TypeError(f"x0 must be a unit.Quantity, type(x0) = {type(x0)}") + raise TypeError( + f"positions must be a unit.Quantity, type(positions) = {type(x0)}" + ) if not isinstance(U0, unit.Quantity): raise TypeError(f"U0 must be a unit.Quantity, type(U0) = {type(U0)}") @@ -376,9 +378,11 @@ def __init__( ) if not x0.unit.is_compatible(unit.angstrom): raise ValueError( - f"x0 must be a unit.Quantity with units of distance, x0.unit = {x0.unit}" + f"positions must be a unit.Quantity with units of distance, positions.unit = {x0.unit}" ) - assert x0.shape[1] == 3, f"x0 must be a NX3 vector, x0.shape = {x0.shape}" + assert ( + x0.shape[1] == 3 + ), f"positions must be a NX3 vector, positions.shape = {x0.shape}" if not U0.unit.is_compatible(unit.kilocalories_per_mole): raise ValueError( f"U0 must be a unit.Quantity with units of energy, U0.unit = {U0.unit}" @@ -388,9 +392,11 @@ def __init__( log.debug("Initializing HarmonicOscillatorPotential") log.debug(f"k = {k}") - log.debug(f"x0 = {x0}") + log.debug(f"positions = {x0}") log.debug(f"U0 = {U0}") - log.debug("Energy is calculate: U(x) = (K/2) * ( (x-x0)^2 + y^2 + z^2 ) + U0") + log.debug( + "Energy is calculate: U(x) = (K/2) * ( (x-positions)^2 + y^2 + z^2 ) + U0" + ) self.k = jnp.array( k.value_in_unit_system(unit.md_unit_system) ) # spring constant @@ -403,7 +409,7 @@ def __init__( self.topology = topology def compute_energy(self, positions: jnp.array, nbr_list=None): - # the functional form is given by U(x) = (K/2) * ( (x-x0)^2 + y^2 + z^2 ) + U0 + # the functional form is given by U(x) = (K/2) * ( (x-positions)^2 + y^2 + z^2 ) + U0 # https://github.com/choderalab/openmmtools/blob/main/openmmtools/testsystems.py#L695 # compute the displacement vectors diff --git a/chiron/states.py b/chiron/states.py index a9ad46a..c76613b 100644 --- a/chiron/states.py +++ b/chiron/states.py @@ -11,7 +11,7 @@ class SamplerState: Parameters ---------- - x0 : unit.Quantity + positions : unit.Quantity The current positions of the particles in the simulation. velocities : unit.Quantity, optional The velocities of the particles in the simulation. @@ -28,21 +28,23 @@ class SamplerState: ho = HarmonicOscillator() PRNG.set_seed(1234) - sampler_state = SamplerState(x0 = ho.positions, PRNG.get_random_key()) + sampler_state = SamplerState(positions = ho.positions, PRNG.get_random_key()) """ def __init__( self, - x0: unit.Quantity, + positions: unit.Quantity, current_PRNG_key: random.PRNGKey, velocities: Optional[unit.Quantity] = None, box_vectors: Optional[unit.Quantity] = None, ) -> None: # NOTE: all units are internally in the openMM units system as documented here: # http://docs.openmm.org/latest/userguide/theory/01_introduction.html#units - if not isinstance(x0, unit.Quantity): - raise TypeError(f"x0 must be a unit.Quantity, got {type(x0)} instead.") + if not isinstance(positions, unit.Quantity): + raise TypeError( + f"positions must be a unit.Quantity, got {type(positions)} instead." + ) if velocities is not None and not isinstance(velocities, unit.Quantity): raise TypeError( f"velocities must be a unit.Quantity, got {type(velocities)} instead." @@ -57,8 +59,10 @@ def __init__( raise TypeError( f"box_vectors must be a unit.Quantity or openMM box, got {type(box_vectors)} instead." ) - if not x0.unit.is_compatible(unit.nanometer): - raise ValueError(f"x0 must have units of distance, got {x0.unit} instead.") + if not positions.unit.is_compatible(unit.nanometer): + raise ValueError( + f"positions must have units of distance, got {positions.unit} instead." + ) if velocities is not None and not velocities.unit.is_compatible( unit.nanometer / unit.picosecond ): @@ -75,14 +79,14 @@ def __init__( raise ValueError( f"box_vectors must be a 3x3 array, got {box_vectors.shape} instead." ) - if velocities is not None and x0.shape != velocities.shape: + if velocities is not None and positions.shape != velocities.shape: raise ValueError( - f"x0 and velocities must have the same shape, got {x0.shape} and {velocities.shape} instead." + f"positions and velocities must have the same shape, got {positions.shape} and {velocities.shape} instead." ) if current_PRNG_key is None: raise ValueError(f"random_seed must be set.") - self._x0 = x0 + self._positions = positions self._velocities = velocities self._current_PRNG_key = current_PRNG_key self._box_vectors = box_vectors @@ -91,11 +95,11 @@ def __init__( @property def n_particles(self) -> int: - return self._x0.shape[0] + return self._positions.shape[0] @property - def x0(self) -> jnp.array: - return self._convert_to_jnp(self._x0) + def positions(self) -> jnp.array: + return self._convert_to_jnp(self._positions) @property def velocities(self) -> jnp.array: @@ -109,12 +113,12 @@ def box_vectors(self) -> jnp.array: return None return self._convert_to_jnp(self._box_vectors) - @x0.setter - def x0(self, x0: Union[jnp.array, unit.Quantity]) -> None: + @positions.setter + def positions(self, x0: Union[jnp.array, unit.Quantity]) -> None: if isinstance(x0, unit.Quantity): - self._x0 = x0 + self._positions = x0 else: - self._x0 = unit.Quantity(x0, self._distance_unit) + self._positions = unit.Quantity(x0, self._distance_unit) @box_vectors.setter def box_vectors(self, box_vectors: Union[jnp.array, unit.Quantity]) -> None: @@ -125,9 +129,9 @@ def box_vectors(self, box_vectors: Union[jnp.array, unit.Quantity]) -> None: @velocities.setter def velocities(self, velocities: Union[jnp.array, unit.Quantity]) -> None: - if velocities.shape != self._x0.shape: + if velocities.shape != self._positions.shape: raise ValueError( - f"velocities must have the same shape as x0, got {velocities.shape} and {self._x0.shape} instead." + f"velocities must have the same shape as positions, got {velocities.shape} and {self._positions.shape} instead." ) if isinstance(velocities, unit.Quantity): self._velocities = velocities @@ -299,10 +303,10 @@ def get_reduced_potential( self.beta = 1.0 / ( unit.BOLTZMANN_CONSTANT_kB * (self.temperature * unit.kelvin) ) - # log.debug(f"sample state: {sampler_state.x0}") + # log.debug(f"sample state: {sampler_state.positions}") reduced_potential = ( unit.Quantity( - self.potential.compute_energy(sampler_state.x0, nbr_list), + self.potential.compute_energy(sampler_state.positions, nbr_list), unit.kilojoule_per_mole, ) ) / unit.AVOGADRO_CONSTANT_NA diff --git a/chiron/tests/test_convergence_tests.py b/chiron/tests/test_convergence_tests.py index 1bd1ae9..5a3ad69 100644 --- a/chiron/tests/test_convergence_tests.py +++ b/chiron/tests/test_convergence_tests.py @@ -53,7 +53,7 @@ def test_convergence_of_MC_estimator(prep_temp_dir): PRNG.set_seed(1234) sampler_state = SamplerState( - x0=ho.positions, current_PRNG_key=PRNG.get_random_key() + positions=ho.positions, current_PRNG_key=PRNG.get_random_key() ) from chiron.reporters import _SimulationReporter @@ -63,16 +63,16 @@ def test_convergence_of_MC_estimator(prep_temp_dir): simulation_reporter = _SimulationReporter(f"{prep_temp_dir}/test_{id}.h5") # Initalize the move set (here only LangevinDynamicsMove) - from chiron.mcmc import MetropolisDisplacementMove, MoveSchedule, MCMCSampler + from chiron.mcmc import MonteCarloDisplacementMove, MoveSchedule, MCMCSampler - mc_displacement_move = MetropolisDisplacementMove( + mc_displacement_move = MonteCarloDisplacementMove( number_of_moves=1_000, displacement_sigma=0.5 * unit.angstrom, atom_subset=[0], reporter=simulation_reporter, ) - move_set = MoveSchedule([("MetropolisDisplacementMove", mc_displacement_move)]) + move_set = MoveSchedule([("MonteCarloDisplacementMove", mc_displacement_move)]) # Initalize the sampler sampler = MCMCSampler(move_set, sampler_state, thermodynamic_state) @@ -148,11 +148,11 @@ def test_langevin_dynamics_with_LJ_fluid(prep_temp_dir): PRNG.set_seed(1234) sampler_state = SamplerState( - x0=lj_fluid.positions, + positions=lj_fluid.positions, box_vectors=lj_fluid.system.getDefaultPeriodicBoxVectors(), current_PRNG_key=PRNG.get_random_key(), ) - print(sampler_state.x0.shape) + print(sampler_state.positions.shape) print(sampler_state.box_vectors) nbr_list = NeighborListNsqrd( @@ -172,11 +172,11 @@ def test_langevin_dynamics_with_LJ_fluid(prep_temp_dir): id = uuid.uuid4() reporter = LangevinDynamicsReporter(f"{prep_temp_dir}/test_{id}.h5") - integrator = LangevinIntegrator(reporter=reporter, report_frequency=100) + integrator = LangevinIntegrator(reporter=reporter, report_interval=100) integrator.run( sampler_state, thermodynamic_state, - n_steps=1000, + number_of_steps=1000, nbr_list=nbr_list, progress_bar=True, ) diff --git a/chiron/tests/test_integrators.py b/chiron/tests/test_integrators.py index e8121de..9ed5b4e 100644 --- a/chiron/tests/test_integrators.py +++ b/chiron/tests/test_integrators.py @@ -43,11 +43,11 @@ def test_langevin_dynamics(prep_temp_dir, provide_testsystems_and_potentials): reporter = LangevinDynamicsReporter() integrator = LangevinIntegrator( - reporter=reporter, report_frequency=1, reinitialize_velocities=True + reporter=reporter, report_interval=1, refresh_velocities=True ) updated_sampler_state, updated_nbr_list = integrator.run( sampler_state, thermodynamic_state, - n_steps=20, + number_of_steps=20, ) i = i + 1 diff --git a/chiron/tests/test_mcmc.py b/chiron/tests/test_mcmc.py index b63b3d4..77e3a24 100644 --- a/chiron/tests/test_mcmc.py +++ b/chiron/tests/test_mcmc.py @@ -41,7 +41,7 @@ def test_sample_from_harmonic_osciallator(prep_temp_dir): PRNG.set_seed(1234) sampler_state = SamplerState( - x0=ho.positions, current_PRNG_key=PRNG.get_random_key() + positions=ho.positions, current_PRNG_key=PRNG.get_random_key() ) from chiron.integrators import LangevinIntegrator @@ -53,16 +53,16 @@ def test_sample_from_harmonic_osciallator(prep_temp_dir): reporter = LangevinDynamicsReporter() integrator = LangevinIntegrator( - stepsize=2 * unit.femtosecond, + timestep=2 * unit.femtosecond, reporter=reporter, - report_frequency=1, - reinitialize_velocities=True, + report_interval=1, + refresh_velocities=True, ) integrator.run( sampler_state, thermodynamic_state, - n_steps=5, + number_of_steps=5, ) integrator.reporter.flush_buffer() import jax.numpy as jnp @@ -128,7 +128,7 @@ def test_sample_from_harmonic_osciallator_with_MCMC_classes_and_LangevinDynamics # the following will reinitialize the velocities for each iteration langevin_move = LangevinDynamicsMove( - nr_of_steps=10, reinitialize_velocities=True, reporter=simulation_reporter + number_of_steps=10, refresh_velocities=True, reporter=simulation_reporter ) move_set = MoveSchedule([("LangevinMove", langevin_move)]) @@ -149,7 +149,7 @@ def test_sample_from_harmonic_osciallator_with_MCMC_classes_and_LangevinDynamics ) langevin_move = LangevinDynamicsMove( - nr_of_steps=10, reinitialize_velocities=False, reporter=simulation_reporter + number_of_steps=10, refresh_velocities=False, reporter=simulation_reporter ) move_set = MoveSchedule([("LangevinMove", langevin_move)]) @@ -174,7 +174,7 @@ def test_sample_from_harmonic_osciallator_with_MCMC_classes_and_MetropolisDispla """ from openmm import unit from chiron.potential import HarmonicOscillatorPotential - from chiron.mcmc import MetropolisDisplacementMove, MoveSchedule, MCMCSampler + from chiron.mcmc import MonteCarloDisplacementMove, MoveSchedule, MCMCSampler # Initalize the testsystem from openmmtools.testsystems import HarmonicOscillator @@ -206,14 +206,14 @@ def test_sample_from_harmonic_osciallator_with_MCMC_classes_and_MetropolisDispla BaseReporter.set_directory(wd) simulation_reporter = MCReporter(1) - mc_displacement_move = MetropolisDisplacementMove( + mc_displacement_move = MonteCarloDisplacementMove( number_of_moves=10, displacement_sigma=0.1 * unit.angstrom, atom_subset=[0], reporter=simulation_reporter, ) - move_set = MoveSchedule([("MetropolisDisplacementMove", mc_displacement_move)]) + move_set = MoveSchedule([("MonteCarloDisplacementMove", mc_displacement_move)]) # Initalize the sampler sampler = MCMCSampler(move_set) @@ -234,7 +234,7 @@ def test_sample_from_harmonic_osciallator_array_with_MCMC_classes_and_Metropolis sampler states, and uses the Metropolis displacement move in an MCMC sampling scheme. """ from openmm import unit - from chiron.mcmc import MetropolisDisplacementMove, MoveSchedule, MCMCSampler + from chiron.mcmc import MonteCarloDisplacementMove, MoveSchedule, MCMCSampler # Initalize the testsystem from openmmtools.testsystems import HarmonicOscillatorArray @@ -268,14 +268,14 @@ def test_sample_from_harmonic_osciallator_array_with_MCMC_classes_and_Metropolis simulation_reporter = MCReporter(1) - mc_displacement_move = MetropolisDisplacementMove( + mc_displacement_move = MonteCarloDisplacementMove( number_of_moves=10, displacement_sigma=0.1 * unit.angstrom, atom_subset=None, reporter=simulation_reporter, ) - move_set = MoveSchedule([("MetropolisDisplacementMove", mc_displacement_move)]) + move_set = MoveSchedule([("MonteCarloDisplacementMove", mc_displacement_move)]) # Initalize the sampler sampler = MCMCSampler(move_set) @@ -352,7 +352,7 @@ def test_mc_barostat(prep_temp_dir): volume_max_scale=0.1, number_of_moves=10, reporter=simulation_reporter, - report_frequency=1, + report_interval=1, ) from chiron.potential import IdealGasPotential @@ -397,7 +397,9 @@ def test_mc_barostat(prep_temp_dir): # define the sampler state sampler_state = SamplerState( - x0=positions, box_vectors=box_vectors, current_PRNG_key=PRNG.get_random_key() + positions=positions, + box_vectors=box_vectors, + current_PRNG_key=PRNG.get_random_key(), ) # define the thermodynamic state diff --git a/chiron/tests/test_minization.py b/chiron/tests/test_minization.py index cf0cf4e..e5a7b8c 100644 --- a/chiron/tests/test_minization.py +++ b/chiron/tests/test_minization.py @@ -29,8 +29,10 @@ def test_minimization(): nbr_list.build_from_state(sampler_state) # compute intial energy with and without pairlist - initial_e_with_nbr_list = lj_potential.compute_energy(sampler_state.x0, nbr_list) - initial_e_without_nbr_list = lj_potential.compute_energy(sampler_state.x0) + initial_e_with_nbr_list = lj_potential.compute_energy( + sampler_state.positions, nbr_list + ) + initial_e_without_nbr_list = lj_potential.compute_energy(sampler_state.positions) print(f"initial_e_with_nbr_list: {initial_e_with_nbr_list}") print(f"initial_e_without_nbr_list: {initial_e_without_nbr_list}") assert not jnp.isclose( @@ -38,7 +40,7 @@ def test_minimization(): ), "initial_e_with_nbr_list and initial_e_without_nbr_list should not be close" # minimize energy for 0 steps results = minimize_energy( - sampler_state.x0, lj_potential.compute_energy, nbr_list, maxiter=0 + sampler_state.positions, lj_potential.compute_energy, nbr_list, maxiter=0 ) # check that the minimization did not change the energy @@ -48,7 +50,7 @@ def test_minimization(): min_x, nbr_list ) after_0_steps_minimization_e_without_nbr_list = lj_potential.compute_energy( - sampler_state.x0 + sampler_state.positions ) print( f"after_0_steps_minimization_e_with_nbr_list: {after_0_steps_minimization_e_with_nbr_list}" @@ -67,7 +69,7 @@ def test_minimization(): # after 100 steps of minimization steps = 100 results = minimize_energy( - sampler_state.x0, lj_potential.compute_energy, nbr_list, maxiter=steps + sampler_state.positions, lj_potential.compute_energy, nbr_list, maxiter=steps ) min_x = results.params e_min = lj_potential.compute_energy(min_x, nbr_list) @@ -103,7 +105,7 @@ def test_minimize_two_particles(): # define the sampler state sampler_state = SamplerState( - x0=coordinates * unit.nanometer, + positions=coordinates * unit.nanometer, current_PRNG_key=PRNG.get_random_key(), box_vectors=jnp.array([[10.0, 0.0, 0.0], [0.0, 10.0, 0.0], [0.0, 0.0, 10.0]]) * unit.nanometer, diff --git a/chiron/tests/test_multistate.py b/chiron/tests/test_multistate.py index 8ed4449..533d879 100644 --- a/chiron/tests/test_multistate.py +++ b/chiron/tests/test_multistate.py @@ -25,7 +25,9 @@ def setup_sampler() -> Tuple[NeighborListNsqrd, MultiStateSampler]: OrthogonalPeriodicSpace(), cutoff=cutoff, skin=skin, n_max_neighbors=180 ) - lang_move = LangevinDynamicsMove(stepsize=1.0 * unit.femtoseconds, nr_of_steps=100) + lang_move = LangevinDynamicsMove( + timestep=1.0 * unit.femtoseconds, number_of_steps=100 + ) BaseReporter.set_directory("multistate_test") reporter = MultistateReporter() reporter.reset_reporter_file() @@ -183,16 +185,16 @@ def test_multistate_minimize(ho_multistate_sampler_multiple_minima: MultiStateSa ho_multistate_sampler_multiple_minima.minimize() assert np.allclose( - ho_multistate_sampler_multiple_minima.sampler_states[0].x0, + ho_multistate_sampler_multiple_minima.sampler_states[0].positions, np.array([[0.0, 0.0, 0.0]]), ) assert np.allclose( - ho_multistate_sampler_multiple_minima.sampler_states[1].x0, + ho_multistate_sampler_multiple_minima.sampler_states[1].positions, np.array([[0.05, 0.0, 0.0]]), atol=1e-2, ) assert np.allclose( - ho_multistate_sampler_multiple_minima.sampler_states[2].x0, + ho_multistate_sampler_multiple_minima.sampler_states[2].positions, np.array([[0.1, 0.0, 0.0]]), atol=1e-2, ) @@ -220,7 +222,6 @@ def test_multistate_run(ho_multistate_sampler_multiple_ks: MultiStateSampler): print(f"Analytical free energy difference: {ho_sampler.delta_f_ij_analytical[0]}") - n_iteratinos = 25 ho_sampler.run(n_iteratinos) diff --git a/chiron/tests/test_pairs.py b/chiron/tests/test_pairs.py index fb2bf2c..237675e 100644 --- a/chiron/tests/test_pairs.py +++ b/chiron/tests/test_pairs.py @@ -100,7 +100,7 @@ def test_neighborlist_pair(): PRNG.set_seed(1234) state = SamplerState( - x0=unit.Quantity(coordinates, unit.nanometer), + positions=unit.Quantity(coordinates, unit.nanometer), current_PRNG_key=PRNG.get_random_key(), box_vectors=unit.Quantity(box_vectors, unit.nanometer), ) @@ -125,7 +125,7 @@ def test_neighborlist_pair(): assert jnp.all(nbr_list.box_vectors == box_vectors) assert nbr_list.is_built == True - nbr_list.build(state.x0, state.box_vectors) + nbr_list.build(state.positions, state.box_vectors) assert jnp.all(nbr_list.ref_coordinates == coordinates) assert jnp.all(nbr_list.box_vectors == box_vectors) @@ -213,7 +213,7 @@ def test_inputs(): PRNG.set_seed(1234) state = SamplerState( - x0=unit.Quantity(coordinates, unit.nanometer), + positions=unit.Quantity(coordinates, unit.nanometer), current_PRNG_key=PRNG.get_random_key(), box_vectors=None, ) @@ -287,7 +287,7 @@ def test_neighborlist_pair_multiple_particles(): PRNG.set_seed(1234) state = SamplerState( - x0=unit.Quantity(coordinates, unit.nanometer), + positions=unit.Quantity(coordinates, unit.nanometer), current_PRNG_key=PRNG.get_random_key(), box_vectors=unit.Quantity(box_vectors, unit.nanometer), ) @@ -342,7 +342,7 @@ def test_neighborlist_pair_multiple_particles(): ) ) # test passing coordinates and box vectors directly - nbr_list.build(state.x0, state.box_vectors) + nbr_list.build(state.positions, state.box_vectors) assert jnp.all(nbr_list.n_neighbors == jnp.array([7, 6, 5, 4, 3, 2, 1, 0])) @@ -362,7 +362,7 @@ def test_pairlist_pair(): PRNG.set_seed(1234) state = SamplerState( - x0=unit.Quantity(coordinates, unit.nanometer), + positions=unit.Quantity(coordinates, unit.nanometer), current_PRNG_key=PRNG.get_random_key(), box_vectors=unit.Quantity(box_vectors, unit.nanometer), ) @@ -416,7 +416,7 @@ def test_pair_list_multiple_particles(): PRNG.set_seed(1234) state = SamplerState( - x0=unit.Quantity(coordinates, unit.nanometer), + positions=unit.Quantity(coordinates, unit.nanometer), current_PRNG_key=PRNG.get_random_key(), box_vectors=unit.Quantity(box_vectors, unit.nanometer), ) diff --git a/chiron/tests/test_potential.py b/chiron/tests/test_potential.py index 230fa1c..3013df1 100644 --- a/chiron/tests/test_potential.py +++ b/chiron/tests/test_potential.py @@ -178,7 +178,7 @@ def test_lennard_jones(): positions = jnp.array([[0, 0, 0], [i * 0.25 * 2 ** (1 / 6), 0, 0]]) state = SamplerState( - x0=unit.Quantity(positions, unit.nanometer), + positions=unit.Quantity(positions, unit.nanometer), current_PRNG_key=PRNG.get_random_key(), box_vectors=unit.Quantity(box_vectors, unit.nanometer), ) diff --git a/chiron/tests/test_states.py b/chiron/tests/test_states.py index 4499f18..e437fad 100644 --- a/chiron/tests/test_states.py +++ b/chiron/tests/test_states.py @@ -31,14 +31,14 @@ def test_initialize_state(): sampler_state = SamplerState(ho.positions, current_PRNG_key=PRNG.get_random_key()) assert jnp.allclose( - sampler_state.x0, + sampler_state.positions, jnp.array([[0.0, 0.0, 0.0]]), ) def test_sampler_state_conversion(): """Test converting a sampler state to jnp arrays. - Note, testing the conversion of x0, where internal unit length is nanometers + Note, testing the conversion of positions, where internal unit length is nanometers and thus output jnp.arrays (with units dropped) should reflect this. """ from chiron.states import SamplerState @@ -54,7 +54,7 @@ def test_sampler_state_conversion(): ) assert jnp.allclose( - sampler_state.x0, + sampler_state.positions, jnp.array([[10.0, 10.0, 10.0]]), ) @@ -64,7 +64,7 @@ def test_sampler_state_conversion(): ) assert jnp.allclose( - sampler_state.x0, + sampler_state.positions, jnp.array([[1.0, 1.0, 1.0]]), ) @@ -81,11 +81,11 @@ def test_sampler_state_inputs(): # test input of positions # should have units with pytest.raises(TypeError): - SamplerState(x0=jnp.array([1, 2, 3])) + SamplerState(positions=jnp.array([1, 2, 3])) # throw and error because of incompatible units with pytest.raises(ValueError): SamplerState( - x0=unit.Quantity(jnp.array([[1, 2, 3]]), unit.radians), + positions=unit.Quantity(jnp.array([[1, 2, 3]]), unit.radians), current_PRNG_key=PRNG.get_random_key(), ) @@ -93,14 +93,14 @@ def test_sampler_state_inputs(): # velocities should have units with pytest.raises(TypeError): SamplerState( - x0=unit.Quantity(jnp.array([[1, 2, 3]]), unit.nanometers), + positions=unit.Quantity(jnp.array([[1, 2, 3]]), unit.nanometers), current_PRNG_key=PRNG.get_random_key(), velocities=jnp.array([1, 2, 3]), ) # velocities should have units of distance/time with pytest.raises(ValueError): SamplerState( - x0=unit.Quantity(jnp.array([[1, 2, 3]]), unit.nanometers), + positions=unit.Quantity(jnp.array([[1, 2, 3]]), unit.nanometers), current_PRNG_key=PRNG.get_random_key(), velocities=unit.Quantity(jnp.array([1, 2, 3]), unit.nanometers), ) @@ -109,14 +109,14 @@ def test_sampler_state_inputs(): # box_vectors should have units with pytest.raises(TypeError): SamplerState( - x0=unit.Quantity(jnp.array([[1, 2, 3]]), unit.nanometers), + positions=unit.Quantity(jnp.array([[1, 2, 3]]), unit.nanometers), current_PRNG_key=PRNG.get_random_key(), box_vectors=jnp.array([1, 2, 3]), ) # box_vectors should have units of distance with pytest.raises(ValueError): SamplerState( - x0=unit.Quantity(jnp.array([[1, 2, 3]]), unit.nanometers), + positions=unit.Quantity(jnp.array([[1, 2, 3]]), unit.nanometers), current_PRNG_key=PRNG.get_random_key(), box_vectors=unit.Quantity( jnp.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]), unit.radians @@ -125,7 +125,7 @@ def test_sampler_state_inputs(): # check to see that the size of the box vectors are correct with pytest.raises(ValueError): SamplerState( - x0=unit.Quantity(jnp.array([[1, 2, 3]]), unit.nanometers), + positions=unit.Quantity(jnp.array([[1, 2, 3]]), unit.nanometers), current_PRNG_key=PRNG.get_random_key(), box_vectors=unit.Quantity( jnp.array([[1, 0, 0], [0, 1, 0]]), unit.nanometers @@ -140,7 +140,7 @@ def test_sampler_state_inputs(): # check openmm_box conversion: state = SamplerState( - x0=unit.Quantity(jnp.array([[1, 2, 3]]), unit.nanometers), + positions=unit.Quantity(jnp.array([[1, 2, 3]]), unit.nanometers), current_PRNG_key=PRNG.get_random_key(), box_vectors=openmm_box, ) @@ -155,7 +155,7 @@ def test_sampler_state_inputs(): # openmm box vectors end up as a list with contents; check to make sure we capture an error if we pass a bad list with pytest.raises(TypeError): SamplerState( - x0=unit.Quantity(jnp.array([[1, 2, 3]]), unit.nanometers), + positions=unit.Quantity(jnp.array([[1, 2, 3]]), unit.nanometers), current_PRNG_key=PRNG.get_random_key(), box_vectors=[123], ) diff --git a/chiron/tests/test_testsystems.py b/chiron/tests/test_testsystems.py index 4f7b758..ce9bb19 100644 --- a/chiron/tests/test_testsystems.py +++ b/chiron/tests/test_testsystems.py @@ -196,7 +196,7 @@ def test_LJ_fluid(): PRNG.set_seed(1234) state = SamplerState( - x0=lj_openmm.positions, + positions=lj_openmm.positions, current_PRNG_key=PRNG.get_random_key(), box_vectors=lj_openmm.system.getDefaultPeriodicBoxVectors(), ) @@ -210,7 +210,7 @@ def test_LJ_fluid(): lj_openmm.topology, sigma=sigma, epsilon=epsilon, cutoff=cutoff ) - e_chiron_energy = lj_chiron.compute_energy(state.x0, nbr_list) + e_chiron_energy = lj_chiron.compute_energy(state.positions, nbr_list) e_openmm_energy = compute_openmm_reference_energy( lj_openmm, lj_openmm.positions ) @@ -253,7 +253,7 @@ def test_ideal_gas(prep_temp_dir): # define the sampler state sampler_state = SamplerState( - x0=ideal_gas.positions, + positions=ideal_gas.positions, current_PRNG_key=PRNG.get_random_key(), box_vectors=ideal_gas.system.getDefaultPeriodicBoxVectors(), ) @@ -277,13 +277,13 @@ def test_ideal_gas(prep_temp_dir): reporter = MCReporter(filename, 1) from chiron.mcmc import ( - MetropolisDisplacementMove, + MonteCarloDisplacementMove, MonteCarloBarostatMove, MoveSchedule, MCMCSampler, ) - mc_displacement_move = MetropolisDisplacementMove( + mc_displacement_move = MonteCarloDisplacementMove( displacement_sigma=0.1 * unit.nanometer, number_of_moves=10, reporter=reporter, @@ -300,7 +300,7 @@ def test_ideal_gas(prep_temp_dir): ) move_set = MoveSchedule( [ - ("MetropolisDisplacementMove", mc_displacement_move), + ("MonteCarloDisplacementMove", mc_displacement_move), ("MonteCarloBarostatMove", mc_barostat_move), ] ) diff --git a/chiron/tests/test_utils.py b/chiron/tests/test_utils.py index 5477fbd..7ce4474 100644 --- a/chiron/tests/test_utils.py +++ b/chiron/tests/test_utils.py @@ -66,12 +66,12 @@ def test_reporter(prep_temp_dir, ho_multistate_sampler_multiple_ks): integrator = LangevinIntegrator( reporter=reporter, - report_frequency=1, + report_interval=1, ) integrator.run( sampler_state, thermodynamic_state, - n_steps=20, + number_of_steps=20, ) import numpy as np From 6a2733b5d086cc04ffcebf53d6fe05d1566f3809 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Tue, 20 Feb 2024 13:49:31 -0800 Subject: [PATCH 38/43] Changed `coordinates` to `positions` in neighbor/pairlist routines to be consistent with samplerstate variable name change. --- chiron/integrators.py | 6 +- chiron/mcmc.py | 2 +- chiron/neighbors.py | 255 +++++++++++++------------ chiron/tests/test_convergence_tests.py | 119 ++++++++++++ chiron/tests/test_pairs.py | 4 +- chiron/tests/test_testsystems.py | 119 ------------ 6 files changed, 263 insertions(+), 242 deletions(-) diff --git a/chiron/integrators.py b/chiron/integrators.py index 9228529..cbcb3a6 100644 --- a/chiron/integrators.py +++ b/chiron/integrators.py @@ -225,19 +225,19 @@ def run( def _wrap_and_rebuild_neighborlist(self, x: jnp.array, nbr_list: PairsBase): """ - Wrap the coordinates and rebuild the neighborlist if necessary. + Wrap the positions and rebuild the neighborlist if necessary. Parameters ---------- x: jnp.array - The coordinates of the particles. + The positions of the particles. nbr_list: PairsBsse The neighborlist object. Returns ------- x: jnp.array - The wrapped coordinates. + The wrapped positions. nbr_list: PairsBase The neighborlist object; this may or may not have been rebuilt. """ diff --git a/chiron/mcmc.py b/chiron/mcmc.py index cbeb38a..2904590 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -938,7 +938,7 @@ def _propose( if current_nbr_list is not None: proposed_nbr_list = copy.deepcopy(current_nbr_list) - # after scaling the box vectors and coordinates we should always rebuild the neighborlist + # after scaling the box vectors and positions we should always rebuild the neighborlist proposed_nbr_list.build( proposed_sampler_state.positions, proposed_sampler_state.box_vectors ) diff --git a/chiron/neighbors.py b/chiron/neighbors.py index 64b897d..baa5674 100644 --- a/chiron/neighbors.py +++ b/chiron/neighbors.py @@ -89,9 +89,9 @@ def displacement( Parameters ---------- xyz_1: jnp.array - Coordinates of the first point + Positions of the first point xyz_2: jnp.array - Coordinates of the second point + Positions of the second point Returns ------- @@ -117,17 +117,17 @@ def displacement( @partial(jax.jit, static_argnums=(0,)) def wrap(self, xyz: jnp.array) -> jnp.array: """ - Wrap the coordinates of the system. + Wrap the positions of the system. Parameters ---------- xyz: jnp.array - Coordinates of the system + Positions of the system Returns ------- jnp.array - Wrapped coordinates of the system + Wrapped positions of the system """ xyz = xyz - jnp.floor(xyz / self._box_lengths) * self._box_lengths @@ -148,9 +148,9 @@ def displacement( Parameters ---------- xyz_1: jnp.array - Coordinates of the first point + Positions of the first point xyz_2: jnp.array - Coordinates of the second point + Positions of the second point Returns ------- @@ -171,18 +171,18 @@ def displacement( @partial(jax.jit, static_argnums=(0,)) def wrap(self, xyz: jnp.array) -> jnp.array: """ - Wrap the coordinates of the system. - For the Non-periodic system, this does not alter the coordinates + Wrap the positions of the system. + For the Non-periodic system, this does not alter the positions Parameters ---------- xyz: jnp.array - Coordinates of the system + Positions of the system Returns ------- jnp.array - Wrapped coordinates of the system + Wrapped positions of the system """ return xyz @@ -190,7 +190,7 @@ def wrap(self, xyz: jnp.array) -> jnp.array: class PairsBase(ABC): """ - Abstract Base Class for different algorithms that determine which particles are interacting. + Abstract Base Class for different algorithms that determine which particle pairs are interacting. Parameters ---------- @@ -213,10 +213,10 @@ class PairsBase(ABC): >>> pair_list = PairsBase(space, cutoff=2.5*unit.nanometer) # initialize the pair list >>> pair_list.build_from_state(sampler_state) # build the pair list from the sampler state >>> - >>> coordinates = sampler_state.positions # get the coordinates from the sampler state, without units attached + >>> positions = sampler_state.positions # get the positions from the sampler state, without units attached >>> >>> # the calculate function will produce information used to calculate the energy - >>> n_neighbors, padding_mask, dist, r_ij = pair_list.calculate(coordinates) + >>> n_neighbors, padding_mask, dist, r_ij = pair_list.calculate(positions) >>> """ @@ -237,16 +237,16 @@ def __init__( @abstractmethod def build( self, - coordinates: Union[jnp.array, unit.Quantity], + positions: Union[jnp.array, unit.Quantity], box_vectors: Union[jnp.array, unit.Quantity], ): """ - Build list from an array of coordinates and array of box vectors. + Build list from an array of positions and array of box vectors. Parameters ---------- - coordinates: jnp.array or unit.Quantity - Shape[n_particles,3] array of particle coordinates, either with or without units attached. + positions: jnp.array or unit.Quantity + Shape[n_particles,3] array of particle positions, either with or without units attached. If the array is passed as a unit.Quantity, the units must be distances and will be converted to nanometers. box_vectors: jnp.array or unit.Quantity Shape[3,3] array of box vectors for the system, either with or without units attached. @@ -261,24 +261,36 @@ def build( def _validate_build_inputs( self, - coordinates: Union[jnp.array, unit.Quantity], + positions: Union[jnp.array, unit.Quantity], box_vectors: Union[jnp.array, unit.Quantity], ): """ Validate the inputs to the build function. + + This will raise ValueErrors if the inputs are not of the correct type or shape or compatible units + + Parameters + ---------- + positions: jnp.array or unit.Quantity + Shape[n_particles,3] array of particle positions, either with or without units attached. + If the array is passed as a unit.Quantity, the units must be distances and will be converted to nanometers. + box_vectors: jnp.array or unit.Quantity + Shape[3,3] array of box vectors for the system, either with or without units attached. + If the array is passed as a unit.Quantity, the units must be distances and will be converted to nanometers. + """ - if isinstance(coordinates, unit.Quantity): - if not coordinates.unit.is_compatible(unit.nanometer): + if isinstance(positions, unit.Quantity): + if not positions.unit.is_compatible(unit.nanometer): raise ValueError( - f"Coordinates require distance units, not {coordinates.unit}" + f"Positions require distance units, not {positions.unit}" ) - self.ref_coordinates = coordinates.value_in_unit_system(unit.md_unit_system) - if isinstance(coordinates, jnp.ndarray): - if coordinates.shape[1] != 3: + self.ref_positions = positions.value_in_unit_system(unit.md_unit_system) + if isinstance(positions, jnp.ndarray): + if positions.shape[1] != 3: raise ValueError( - f"coordinates should be a Nx3 array, shape provided: {coordinates.shape}" + f"positions should be a Nx3 array, shape provided: {positions.shape}" ) - self.ref_coordinates = coordinates + self.ref_positions = positions if isinstance(box_vectors, unit.Quantity): if not box_vectors.unit.is_compatible(unit.nanometer): raise ValueError( @@ -300,7 +312,7 @@ def build_from_state(self, sampler_state: SamplerState): Parameters ---------- sampler_state: SamplerState - SamplerState object containing the coordinates and box vectors + SamplerState object containing the positions and box vectors Returns ------- @@ -309,22 +321,22 @@ def build_from_state(self, sampler_state: SamplerState): if not isinstance(sampler_state, SamplerState): raise TypeError(f"Expected SamplerState, got {type(sampler_state)} instead") - coordinates = sampler_state.positions + positions = sampler_state.positions if sampler_state.box_vectors is None: raise ValueError(f"SamplerState does not contain box vectors") box_vectors = sampler_state.box_vectors - self.build(coordinates, box_vectors) + self.build(positions, box_vectors) @abstractmethod - def calculate(self, coordinates: jnp.array): + def calculate(self, positions: jnp.array): """ - Calculate the neighbor list for the current state + Calculate the list of interacting particles for the current state Parameters ---------- - coordinates: jnp.array - Shape[N,3] array of particle coordinates + positions: jnp.array + Shape[N,3] array of particle positions Returns ------- @@ -343,15 +355,16 @@ def calculate(self, coordinates: jnp.array): pass @abstractmethod - def check(self, coordinates: jnp.array) -> bool: + def check(self, positions: jnp.array) -> bool: """ - Check if the internal variables need to be reset. E.g., rebuilding a neighborlist - Should do nothing for a simple pairlist. + Check if the internal variables need to be reset. E.g., rebuilding a neighborlist if particles moved to far, + or rebuilding if number of particles changes. + Parameters ---------- - coordinates: jnp.array - Array of particle coordinates + positions: jnp.array + Array of particle positions Returns ------- bool @@ -362,19 +375,26 @@ def check(self, coordinates: jnp.array) -> bool: class NeighborListNsqrd(PairsBase): """ - N^2 neighborlist implementation that returns the particle pair ids, displacement vectors, and distances. + Neighbor list implementation that returns the particle pair ids, displacement vectors, + and distances between pairs within the specified cutoff range. + The particle neighbor lists (i.e., lists of particles within cutoff+skin) are generated using an N^2 calculation. + + This will pad the neighbor list to a fixed size, n_max_neighbors, for efficiency purposes for JITTED functions. + The code will automatically increase n_max_neighbors if the number of neighbors exceeds the current value. Parameters ---------- space: Space Class that defines how to calculate the displacement between two points and apply the boundary conditions - cutoff: float, default = 2.5 + cutoff: unit.Quantity, default = 1.2 unit.nanometer Cutoff distance for the neighborlist - skin: float, default = 0.4 - Skin distance for the neighborlist + skin: unit.Quantity, default = 0.4 unit.nanometer + Skin distance, i.e., buffer, for the neighborlist + Larger values of the skin will reduce the frequency of rebuilding the neighbor list, + but will increase the number of neighbors to consider. n_max_neighbors: int, default=200 - Maximum number of neighbors for each particle. Used for padding arrays for efficient jax computations - This will be checked and dynamically updated during the build stage + Maximum number of neighbors for each particle. This is used for padding arrays for efficient jax computations + n_max_neighbors will be dynamically updated (in increments of 10) as part of the build function. Examples -------- @@ -446,7 +466,7 @@ def _pairs_mask(self, particle_ids: jnp.array): @partial(jax.jit, static_argnums=(0, 5)) def _build_neighborlist( - self, particle_i, reduction_mask, pid, coordinates, n_max_neighbors + self, particle_i, reduction_mask, pid, positions, n_max_neighbors ): """ Jitted function to build the neighbor list for a single particle @@ -454,11 +474,11 @@ def _build_neighborlist( Parameters ---------- particle_i: jnp.array - X,Y,Z coordinates of particle i + X,Y,Z positions of particle i reduction_mask: jnp.array Mask to exclude self-interactions and double counting of pairs - coordinates: jnp.array - X,Y,Z coordinates of all particles + positions: jnp.array + X,Y,Z positions of all particles n_max_neighbors: int Maximum number of neighbors for each particle. Used for padding arrays for efficient jax computations @@ -473,9 +493,9 @@ def _build_neighborlist( """ # calculate the displacement between particle i and all other particles - r_ij, dist = self.space.displacement(particle_i, coordinates) + r_ij, dist = self.space.displacement(particle_i, positions) - # neighbor_mask will be an array of length n_particles (i.e., length of coordinates) + # neighbor_mask will be an array of length n_particles (i.e., length of positions) # where each element is True if the particle is a neighbor, False if it is not # subject to both the cutoff+skin and the reduction mask that eliminates double counting and self-interactions neighbor_mask = jnp.where( @@ -506,16 +526,16 @@ def _build_neighborlist( def build( self, - coordinates: Union[jnp.array, unit.Quantity], + positions: Union[jnp.array, unit.Quantity], box_vectors: Union[jnp.array, unit.Quantity], ): """ - Build the neighborlist from an array of coordinates and box vectors. + Build the neighborlist from an array of positions and box vectors. Parameters ---------- - coordinates: jnp.array - Shape[N,3] array of particle coordinates + positions: jnp.array + Shape[N,3] array of particle positions box_vectors: jnp.array Shape[3,3] array of box vectors @@ -525,14 +545,14 @@ def build( """ - # set our reference coordinates + # set our reference positions # the call to positions and box_vectors automatically convert these to jnp arrays in the correct unit system - if isinstance(coordinates, unit.Quantity): - if not coordinates.unit.is_compatible(unit.nanometer): + if isinstance(positions, unit.Quantity): + if not positions.unit.is_compatible(unit.nanometer): raise ValueError( - f"Coordinates require distance units, not {coordinates.unit}" + f"Positions require distance units, not {positions.unit}" ) - coordinates = coordinates.value_in_unit_system(unit.md_unit_system) + positions = positions.value_in_unit_system(unit.md_unit_system) if isinstance(box_vectors, unit.Quantity): if not box_vectors.unit.is_compatible(unit.nanometer): @@ -546,7 +566,7 @@ def build( f"box_vectors should be a 3x3 array, shape provided: {box_vectors.shape}" ) - self.ref_coordinates = coordinates + self.ref_positions = positions self.box_vectors = box_vectors # the neighborlist assumes that the box vectors do not change between building and calculating the neighbor list @@ -555,7 +575,7 @@ def build( # store the ids of all the particles self.particle_ids = jnp.array( - range(0, self.ref_coordinates.shape[0]), dtype=jnp.uint32 + range(0, self.ref_positions.shape[0]), dtype=jnp.uint32 ) # calculate which pairs to exclude @@ -571,10 +591,10 @@ def build( self.neighbor_mask, self.neighbor_list, self.n_neighbors = jax.vmap( self._build_neighborlist, in_axes=(0, 0, 0, None, None) )( - self.ref_coordinates, + self.ref_positions, reduction_mask, self.particle_ids, - self.ref_coordinates, + self.ref_positions, self.n_max_neighbors, ) @@ -590,10 +610,10 @@ def build( self.neighbor_mask, self.neighbor_list, self.n_neighbors = jax.vmap( self._build_neighborlist, in_axes=(0, 0, 0, None, None) )( - self.ref_coordinates, + self.ref_positions, reduction_mask, self.particle_ids, - self.ref_coordinates, + self.ref_positions, self.n_max_neighbors, ) @@ -603,7 +623,7 @@ def build( @partial(jax.jit, static_argnums=(0,)) def _calc_distance_per_particle( - self, particle1, neighbors, neighbor_mask, coordinates + self, particle1, neighbors, neighbor_mask, positions ): """ Jitted function to calculate the distance between a particle and its neighbors @@ -616,8 +636,8 @@ def _calc_distance_per_particle( Array of particle ids for the neighbors of particle1 neighbor_mask: jnp.array Mask to exclude padding from the neighbor list of particle1 - coordinates: jnp.array - X,Y,Z coordinates of all particles + positions: jnp.array + X,Y,Z positions of all particles Returns ------- @@ -636,7 +656,7 @@ def _calc_distance_per_particle( # calculate the displacement between particle i and all neighbors r_ij, dist = self.space.displacement( - coordinates[particles1], coordinates[neighbors] + positions[particles1], positions[neighbors] ) # calculate the mask to determine if the particle is a neighbor # this will be done based on the interaction cutoff and using the neighbor_mask to exclude padding @@ -647,14 +667,14 @@ def _calc_distance_per_particle( return n_pairs, mask, dist, r_ij - def calculate(self, coordinates: jnp.array): + def calculate(self, positions: jnp.array): """ Calculate the neighbor list for the current state Parameters ---------- - coordinates: jnp.array - Shape[N,3] array of particle coordinates + positions: jnp.array + Shape[N,3] array of particle positions Returns ------- @@ -669,20 +689,20 @@ def calculate(self, coordinates: jnp.array): r_ij: jnp.array Array of displacement vectors between each particle and its neighbors. Shape (n_particles, n_max_neighbors, 3) """ - # coordinates = sampler_state.positions + # positions = sampler_state.positions # note, we assume the box vectors do not change between building and calculating the neighbor list # changes to the box vectors require rebuilding the neighbor list n_neighbors, padding_mask, dist, r_ij = jax.vmap( self._calc_distance_per_particle, in_axes=(0, 0, 0, None) - )(self.particle_ids, self.neighbor_list, self.neighbor_mask, coordinates) + )(self.particle_ids, self.neighbor_list, self.neighbor_mask, positions) # mask = mask.reshape(-1, self.n_max_neighbors) return n_neighbors, self.neighbor_list, padding_mask, dist, r_ij @partial(jax.jit, static_argnums=(0,)) - def _calculate_particle_displacement(self, particle, coordinates, ref_coordinates): + def _calculate_particle_displacement(self, particle, positions, ref_positions): """ - Calculate the displacement of a particle from the reference coordinates. + Calculate the displacement of a particle from the reference positions. If the displacement exceeds the half the skin distance, return True, otherwise return False. This function is designed to allow it to be jitted and vmapped over particle indices. @@ -691,50 +711,50 @@ def _calculate_particle_displacement(self, particle, coordinates, ref_coordinate ---------- particle: int Particle id - coordinates: jnp.array - Array of particle coordinates - ref_coordinates: jnp.array - Array of reference particle coordinates + positions: jnp.array + Array of particle positions + ref_positions: jnp.array + Array of reference particle positions Returns ------- bool True if the particle is outside the skin distance, False if it is not. """ - # calculate the displacement of a particle from the initial coordinates + # calculate the displacement of a particle from the initial positions r_ij, displacement = self.space.displacement( - coordinates[particle], ref_coordinates[particle] + positions[particle], ref_positions[particle] ) status = jnp.where(displacement >= self.skin / 2.0, True, False) del displacement return status - def check(self, coordinates: jnp.array) -> bool: + def check(self, positions: jnp.array) -> bool: """ - Check if the neighbor list needs to be rebuilt based on displacement of the particles from the reference coordinates. + Check if the neighbor list needs to be rebuilt based on displacement of the particles from the reference positions. If a particle moves more than 0.5 skin distance, the neighborlist will be rebuilt. - Will also return True if the size of the coordinates array changes. + Will also return True if the size of the positions array changes. Note, this could also accept a user defined criteria for distance, but this is not implemented yet. Parameters ---------- - coordinates: jnp.array - Array of particle coordinates + positions: jnp.array + Array of particle positions Returns ------- bool True if the neighbor list needs to be rebuilt, False if it does not. """ - if self.ref_coordinates.shape[0] != coordinates.shape[0]: + if self.ref_positions.shape[0] != positions.shape[0]: return True status = jax.vmap( self._calculate_particle_displacement, in_axes=(0, None, None) - )(self.particle_ids, coordinates, self.ref_coordinates) + )(self.particle_ids, positions, self.ref_positions) if jnp.any(status): del status return True @@ -826,12 +846,14 @@ def _pairs_and_mask(self, particle_ids: jnp.array): particles_i = jnp.reshape(particle_ids, (particle_ids.shape[0], 1)) # create a mask to exclude self interactions and double counting temp_mask = particles_i != particles_j + # remove self interactions all_pairs = jax.vmap(self._remove_self_interactions, in_axes=(0, 0))( particles_j, temp_mask ) del temp_mask all_pairs = jnp.array(all_pairs[0], dtype=jnp.uint32) + # create the mask that will remove any double counting of pairs reduction_mask = jnp.where(particles_i < all_pairs, True, False) return all_pairs, reduction_mask @@ -844,16 +866,16 @@ def _remove_self_interactions(self, particles, temp_mask): def build( self, - coordinates: Union[jnp.array, unit.Quantity], + positions: Union[jnp.array, unit.Quantity], box_vectors: Union[jnp.array, unit.Quantity], ): """ - Build the neighborlist from an array of coordinates and box vectors. + Build the list from an array of positions and box vectors. Parameters ---------- - coordinates: jnp.array - Shape[n_particles,3] array of particle coordinates + positions: jnp.array + Shape[n_particles,3] array of particle positions box_vectors: jnp.array Shape[3,3] array of box vectors @@ -863,18 +885,17 @@ def build( """ - # set our reference coordinates - # this will set self.ref_coordinates=coordinates and self.box_vectors - self._validate_build_inputs(coordinates, box_vectors) + # validate the positions and box vectors + self._validate_build_inputs(positions, box_vectors) - self.n_particles = self.ref_coordinates.shape[0] + self.n_particles = self.ref_positions.shape[0] - # 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 + # the PairsList assumes that the box vectors do not change between building and calculating the neighbor list + # changes to the box vectors require rebuilding the list self.space.box_vectors = self.box_vectors # store the ids of all the particles - self.particle_ids = jnp.array(range(0, coordinates.shape[0]), dtype=jnp.uint32) + self.particle_ids = jnp.array(range(0, positions.shape[0]), dtype=jnp.uint32) # calculate which pairs to exclude self.all_pairs, self.reduction_mask = self._pairs_and_mask(self.particle_ids) @@ -883,7 +904,7 @@ def build( @partial(jax.jit, static_argnums=(0,)) def _calc_distance_per_particle( - self, particle1, neighbors, neighbor_mask, coordinates + self, particle1, neighbors, neighbor_mask, positions ): """ Jitted function to calculate the distance between a particle and all possible neighbors @@ -896,8 +917,8 @@ def _calc_distance_per_particle( Array of particle ids for the possible particle pairs of particle1 neighbor_mask: jnp.array Mask to exclude double particles to prevent double counting - coordinates: jnp.array - X,Y,Z coordinates of all particles, shaped (n_particles, 3) + positions: jnp.array + X,Y,Z positions of all particles, shaped (n_particles, 3) Returns ------- @@ -920,7 +941,7 @@ def _calc_distance_per_particle( # calculate the displacement between particle i and all neighbors r_ij, dist = self.space.displacement( - coordinates[particles1], coordinates[neighbors] + positions[particles1], positions[neighbors] ) # calculate the mask to determine if the particle is a neighbor # this will be done based on the interaction cutoff and using the neighbor_mask to exclude padding @@ -931,14 +952,14 @@ def _calc_distance_per_particle( return n_pairs, mask, dist, r_ij - def calculate(self, coordinates: jnp.array): + def calculate(self, positions: jnp.array): """ - Calculate the neighbor list for the current state + Calculate the list of neighbor pairs for the current state Parameters ---------- - coordinates: jnp.array - Shape[n_particles,3] array of particle coordinates + positions: jnp.array + Shape[n_particles,3] array of particle positions Returns ------- @@ -953,35 +974,35 @@ def calculate(self, coordinates: jnp.array): r_ij: jnp.array Array of displacement vectors between particle pairs. Shape: (n_particles, n_particles-1, 3). """ - if coordinates.shape[0] != self.n_particles: + if positions.shape[0] != self.n_particles: raise ValueError( f"Number of particles cannot changes without rebuilding. " - f"Coordinates must have shape ({self.n_particles}, 3), found {coordinates.shape}" + f"Positions must have shape ({self.n_particles}, 3), found {positions.shape}" ) - # coordinates = self.space.wrap(coordinates) + # positions = self.space.wrap(positions) 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) + )(self.particle_ids, self.all_pairs, self.reduction_mask, positions) return n_neighbors, self.all_pairs, padding_mask, dist, r_ij - def check(self, coordinates: jnp.array) -> bool: + def check(self, positions: jnp.array) -> bool: """ Check if we need to reconstruct internal arrays. For a simple pairlist this will always return False, unless the number of particles change. Parameters ---------- - coordinates: jnp.array - Array of particle coordinates + positions: jnp.array + Array of particle positions Returns ------- bool True if we need to rebuild the neighbor list, False if we do not. """ - if coordinates.shape[0] != self.n_particles: + if positions.shape[0] != self.n_particles: return True else: return False diff --git a/chiron/tests/test_convergence_tests.py b/chiron/tests/test_convergence_tests.py index 5a3ad69..deeb3e5 100644 --- a/chiron/tests/test_convergence_tests.py +++ b/chiron/tests/test_convergence_tests.py @@ -180,3 +180,122 @@ def test_langevin_dynamics_with_LJ_fluid(prep_temp_dir): nbr_list=nbr_list, progress_bar=True, ) + + +def test_ideal_gas(prep_temp_dir): + from openmmtools.testsystems import IdealGas + from openmm import unit + + 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 + from chiron.utils import PRNG + 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, + ) + + 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 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 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, 1) + + from chiron.mcmc import ( + MonteCarloDisplacementMove, + MonteCarloBarostatMove, + MoveSchedule, + MCMCSampler, + ) + + mc_displacement_move = MonteCarloDisplacementMove( + displacement_sigma=0.1 * unit.nanometer, + number_of_moves=10, + reporter=reporter, + autotune=True, + autotune_interval=100, + ) + + mc_barostat_move = MonteCarloBarostatMove( + volume_max_scale=0.2, + number_of_moves=100, + reporter=reporter, + autotune=True, + autotune_interval=100, + ) + move_set = MoveSchedule( + [ + ("MonteCarloDisplacementMove", mc_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 + + volume = reporter.get_property("volume") + + # 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, ideal_volume_std) + + volume_mean = jnp.mean(jnp.array(volume)) * unit.nanometer**3 + volume_std = jnp.std(jnp.array(volume)) * unit.nanometer**3 + + print(volume_mean, volume_std) + + 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 diff --git a/chiron/tests/test_pairs.py b/chiron/tests/test_pairs.py index 237675e..e30c008 100644 --- a/chiron/tests/test_pairs.py +++ b/chiron/tests/test_pairs.py @@ -121,13 +121,13 @@ def test_neighborlist_pair(): nbr_list.build_from_state(state) - assert jnp.all(nbr_list.ref_coordinates == coordinates) + assert jnp.all(nbr_list.ref_positions == coordinates) assert jnp.all(nbr_list.box_vectors == box_vectors) assert nbr_list.is_built == True nbr_list.build(state.positions, state.box_vectors) - assert jnp.all(nbr_list.ref_coordinates == coordinates) + assert jnp.all(nbr_list.ref_positions == coordinates) assert jnp.all(nbr_list.box_vectors == box_vectors) assert nbr_list.is_built == True diff --git a/chiron/tests/test_testsystems.py b/chiron/tests/test_testsystems.py index ce9bb19..5b16ef8 100644 --- a/chiron/tests/test_testsystems.py +++ b/chiron/tests/test_testsystems.py @@ -217,122 +217,3 @@ 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 - - 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 - from chiron.utils import PRNG - 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, - ) - - 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 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 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, 1) - - from chiron.mcmc import ( - MonteCarloDisplacementMove, - MonteCarloBarostatMove, - MoveSchedule, - MCMCSampler, - ) - - mc_displacement_move = MonteCarloDisplacementMove( - displacement_sigma=0.1 * unit.nanometer, - number_of_moves=10, - reporter=reporter, - autotune=True, - autotune_interval=100, - ) - - mc_barostat_move = MonteCarloBarostatMove( - volume_max_scale=0.2, - number_of_moves=100, - reporter=reporter, - autotune=True, - autotune_interval=100, - ) - move_set = MoveSchedule( - [ - ("MonteCarloDisplacementMove", mc_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 - - volume = reporter.get_property("volume") - - # 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, ideal_volume_std) - - volume_mean = jnp.mean(jnp.array(volume)) * unit.nanometer**3 - volume_std = jnp.std(jnp.array(volume)) * unit.nanometer**3 - - print(volume_mean, volume_std) - - 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 From 6337ce59fc79182d21de771e152ff2e702095c81 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Wed, 21 Feb 2024 15:55:31 -0800 Subject: [PATCH 39/43] Changed PairListNsqrd to allow setting the cutoff to None which will use no cutoff (i.e. all pairs interact). Revamped some internal tooling regarding how units are treated internally. Space was revampped such that it takes box_vectors as a argument rather than storing them internally (fewer copies of this floating around), and will help to ensure treating the class as static will not mean using the wrong box vectors accidentally. --- Examples/Idealgas.py | 4 +- Examples/LJ_MCMC.py | 4 +- Examples/LJ_mcmove.py | 2 +- chiron/integrators.py | 2 +- chiron/mcmc.py | 3 +- chiron/neighbors.py | 462 ++++++++++++++++++------- chiron/potential.py | 2 +- chiron/states.py | 4 +- chiron/tests/test_convergence_tests.py | 6 +- chiron/tests/test_mcmc.py | 4 +- chiron/tests/test_minization.py | 8 +- chiron/tests/test_pairs.py | 162 ++++----- 12 files changed, 433 insertions(+), 230 deletions(-) diff --git a/Examples/Idealgas.py b/Examples/Idealgas.py index d6830ce..32ef6f0 100644 --- a/Examples/Idealgas.py +++ b/Examples/Idealgas.py @@ -47,12 +47,12 @@ box_vectors=ideal_gas.system.getDefaultPeriodicBoxVectors(), ) -from chiron.neighbors import PairList, OrthogonalPeriodicSpace +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 = PairList(OrthogonalPeriodicSpace(), cutoff=cutoff) +nbr_list = PairListNsqrd(OrthogonalPeriodicSpace(), cutoff=cutoff) nbr_list.build_from_state(sampler_state) from chiron.reporters import MCReporter diff --git a/Examples/LJ_MCMC.py b/Examples/LJ_MCMC.py index b2e8684..63d2412 100644 --- a/Examples/LJ_MCMC.py +++ b/Examples/LJ_MCMC.py @@ -72,12 +72,12 @@ ) -from chiron.neighbors import PairList, OrthogonalPeriodicSpace +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 = PairList(OrthogonalPeriodicSpace(), cutoff=cutoff) +nbr_list = PairListNsqrd(OrthogonalPeriodicSpace(), cutoff=cutoff) nbr_list.build_from_state(sampler_state) # CRI: minimizer is not working correctly on my mac diff --git a/Examples/LJ_mcmove.py b/Examples/LJ_mcmove.py index fac45d4..09fa6fa 100644 --- a/Examples/LJ_mcmove.py +++ b/Examples/LJ_mcmove.py @@ -44,7 +44,7 @@ nbr_list = NeighborListNsqrd( OrthogonalPeriodicSpace(), cutoff=cutoff, skin=skin, n_max_neighbors=180 ) -from chiron.neighbors import PairList +from chiron.neighbors import PairListNsqrd # build the neighbor list from the sampler state diff --git a/chiron/integrators.py b/chiron/integrators.py index cbcb3a6..50e6406 100644 --- a/chiron/integrators.py +++ b/chiron/integrators.py @@ -242,7 +242,7 @@ def _wrap_and_rebuild_neighborlist(self, x: jnp.array, nbr_list: PairsBase): The neighborlist object; this may or may not have been rebuilt. """ - x = nbr_list.space.wrap(x) + x = nbr_list.space.wrap(x, self.box_vectors) # check if we need to rebuild the neighborlist after moving the particles if nbr_list.check(x): nbr_list.build(x, self.box_vectors) diff --git a/chiron/mcmc.py b/chiron/mcmc.py index 2904590..c95ce46 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -721,7 +721,8 @@ def _propose( # after proposing a move we need to wrap particles and see if we need to rebuild the neighborlist if current_nbr_list is not None: proposed_sampler_state.positions = current_nbr_list.space.wrap( - proposed_sampler_state.positions + proposed_sampler_state.positions, + proposed_sampler_state.box_vectors, ) # if we need to rebuild the neighbor the neighborlist diff --git a/chiron/neighbors.py b/chiron/neighbors.py index baa5674..ccdc706 100644 --- a/chiron/neighbors.py +++ b/chiron/neighbors.py @@ -3,62 +3,36 @@ import jax import jax.numpy as jnp from functools import partial -from typing import Tuple, Union +from typing import Tuple, Union, Optional from .states import SamplerState from openmm import unit -# split out the displacement calculation from the neighborlist for flexibility +# split out the displacement calculation from the neighbor list and pair list for flexibility from abc import ABC, abstractmethod class Space(ABC): - def __init__( - self, box_vectors: Union[jnp.array, unit.Quantity, None] = None - ) -> None: - """ - Abstract base class for defining the simulation space. + """ + Abstract Base Class for different simulation spaces. - Parameters - ---------- - box_vectors: jnp.array, optional - Box vectors for the system. - """ - if box_vectors is not None: - if isinstance(box_vectors, unit.Quantity): - if not box_vectors.unit.is_compatible(unit.nanometer): - raise ValueError( - f"Box vectors require distance unit, not {box_vectors.unit}" - ) - self.box_vectors = box_vectors.value_in_unit_system(unit.md_unit_system) - elif isinstance(box_vectors, jnp.ndarray): - if box_vectors.shape != (3, 3): - raise ValueError( - f"box_vectors should be a 3x3 array, shape provided: {box_vectors.shape}" - ) - - self.box_vectors = box_vectors - else: - raise TypeError( - f"box_vectors must be a jnp.array or unit.Quantity, not {type(box_vectors)}" - ) + This class will define two functions: + - displacement, i.e., how to calculate the displacement vector and distance between two points + - wrap, i.e., how to wrap a particle in the box (i.e., apply boundary conditions). + + Note, this class does not store the box_vectors; they will need to be passed to each function. - @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 + """ @abstractmethod def displacement( - self, xyz_1: jnp.array, xyz_2: jnp.array + self, xyz_1: jnp.array, xyz_2: jnp.array, box_vectors: jnp.array ) -> Tuple[jnp.array, jnp.array]: pass @abstractmethod - def wrap(self, xyz: jnp.array) -> jnp.array: + def wrap(self, xyz: jnp.array, box_vectors: jnp.array) -> jnp.array: pass @@ -68,20 +42,9 @@ class OrthogonalPeriodicSpace(Space): """ - @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._box_lengths = jnp.array( - [box_vectors[0][0], box_vectors[1][1], box_vectors[2][2]] - ) - @partial(jax.jit, static_argnums=(0,)) def displacement( - self, xyz_1: jnp.array, xyz_2: jnp.array + self, xyz_1: jnp.array, xyz_2: jnp.array, box_vectors: jnp.array ) -> Tuple[jnp.array, jnp.array]: """ Calculate the periodic distance between two points. @@ -92,6 +55,7 @@ def displacement( Positions of the first point xyz_2: jnp.array Positions of the second point + box_vectors: jnp.array Returns ------- @@ -101,21 +65,22 @@ def displacement( Distance between the two points """ - # calculate uncorrect r_ij + # calculate uncorrected r_ij r_ij = xyz_1 - xyz_2 - # calculated corrected displacement vector - r_ij = ( - jnp.mod(r_ij + self._box_lengths * 0.5, self._box_lengths) - - self._box_lengths * 0.5 + box_lengths = jnp.array( + [box_vectors[0][0], box_vectors[1][1], box_vectors[2][2]] ) + # calculated corrected displacement vector + # using modulus seems faster in JAX + r_ij = jnp.mod(r_ij + box_lengths * 0.5, box_lengths) - box_lengths * 0.5 # calculate the scalar distance dist = jnp.linalg.norm(r_ij, axis=-1) return r_ij, dist @partial(jax.jit, static_argnums=(0,)) - def wrap(self, xyz: jnp.array) -> jnp.array: + def wrap(self, xyz: jnp.array, box_vectors: jnp.array) -> jnp.array: """ Wrap the positions of the system. @@ -123,6 +88,8 @@ def wrap(self, xyz: jnp.array) -> jnp.array: ---------- xyz: jnp.array Positions of the system + box_vectors: jnp.array + Box vectors for the system Returns ------- @@ -130,7 +97,11 @@ def wrap(self, xyz: jnp.array) -> jnp.array: Wrapped positions of the system """ - xyz = xyz - jnp.floor(xyz / self._box_lengths) * self._box_lengths + box_lengths = jnp.array( + [box_vectors[0][0], box_vectors[1][1], box_vectors[2][2]] + ) + + xyz = xyz - jnp.floor(xyz / box_lengths) * box_lengths return xyz @@ -141,6 +112,7 @@ def displacement( self, xyz_1: jnp.array, xyz_2: jnp.array, + box_vectors: jnp.array, ) -> Tuple[jnp.array, jnp.array]: """ Calculate the periodic distance between two points. @@ -151,6 +123,8 @@ def displacement( Positions of the first point xyz_2: jnp.array Positions of the second point + box_vectors: jnp.array + Box vectors for the system. Returns ------- @@ -169,7 +143,7 @@ def displacement( return r_ij, dist @partial(jax.jit, static_argnums=(0,)) - def wrap(self, xyz: jnp.array) -> jnp.array: + def wrap(self, xyz: jnp.array, box_vectors: jnp.array) -> jnp.array: """ Wrap the positions of the system. For the Non-periodic system, this does not alter the positions @@ -178,6 +152,8 @@ def wrap(self, xyz: jnp.array) -> jnp.array: ---------- xyz: jnp.array Positions of the system + box_vectors: jnp.array + Box vectors for the system Returns ------- @@ -231,7 +207,7 @@ def __init__( raise ValueError( f"cutoff must be a unit.Quantity with units of distance, cutoff.unit = {cutoff.unit}" ) - self.cutoff = cutoff.value_in_unit_system(unit.md_unit_system) + self.cutoff = cutoff self.space = space @abstractmethod @@ -375,29 +351,68 @@ def check(self, positions: jnp.array) -> bool: class NeighborListNsqrd(PairsBase): """ - Neighbor list implementation that returns the particle pair ids, displacement vectors, - and distances between pairs within the specified cutoff range. - The particle neighbor lists (i.e., lists of particles within cutoff+skin) are generated using an N^2 calculation. + Neighbor list implementation used to determine which particles are interacting (i.e., within the cutoff). - This will pad the neighbor list to a fixed size, n_max_neighbors, for efficiency purposes for JITTED functions. - The code will automatically increase n_max_neighbors if the number of neighbors exceeds the current value. + This `calculate` function of this class returns the particle pair ids, displacement vectors, and distances + between pairs within the specified cutoff range, subject to the boundary conditions defined by the simulation + Space class passed to the constructor. + + The neighbor list (i.e., list of particles within cutoff+skin of a given particle) + is generated using an N^2 calculation rather than using a spatial partitioning scheme (e.g., cell-list). - Parameters - ---------- - space: Space - Class that defines how to calculate the displacement between two points and apply the boundary conditions - cutoff: unit.Quantity, default = 1.2 unit.nanometer - Cutoff distance for the neighborlist - skin: unit.Quantity, default = 0.4 unit.nanometer - Skin distance, i.e., buffer, for the neighborlist - Larger values of the skin will reduce the frequency of rebuilding the neighbor list, - but will increase the number of neighbors to consider. - n_max_neighbors: int, default=200 - Maximum number of neighbors for each particle. This is used for padding arrays for efficient jax computations - n_max_neighbors will be dynamically updated (in increments of 10) as part of the build function. - Examples - -------- + Notes: + This does not include self-interactions and only includes unique pairs (i.e., no double-counting). + This is sometimes referred to as a "half" neighbor list. E.g. consider the pair of neighboring particles (A, B): + in the "half" neighbor list approach, B is in the neighbor list of A, but A is not in the neighbor list of B + as that pair is already accounted for. + . + The output of the `calculate` function is padded to a fixed size, n_max_neighbors, to allow for efficient + jitted computations in JAX. As such, values need to be masked using the `padding_mask` array returned by the + `calculate` function. The padding mask is simple an array of 1s and 0s, where 1 indicates a valid neighbor and + 0 indicates padding. The code will automatically increase n_max_neighbors by 10 if the number of neighbors + exceeds the current value. + + The check function, which indicates if the neighbor list should be rebuilt, will return True if: + - the number of particles changes + - any of the particles have moved more than half the skin distance from their positions at the time of last + building the neighbor list. + + + Parameters + ---------- + space: Space + Class that defines how to calculate the displacement between two points and apply the boundary conditions + cutoff: unit.Quantity, default = 1.2 unit.nanometer + Cutoff distance for the neighborlist + skin: unit.Quantity, default = 0.4 unit.nanometer + Skin distance, i.e., buffer, for the neighborlist + Larger values of the skin will reduce the frequency of rebuilding the neighbor list, + but will increase the number of neighbors to consider. + n_max_neighbors: int, default=200 + Maximum number of neighbors for each particle. This is used for padding arrays for efficient jax computations + n_max_neighbors will be dynamically updated (in increments of 10) as part of the build function. + Examples + -------- + >>> from openmm import unit + >>> import jax.numpy as jnp + >>> + >>> from chiron.states import SamplerState + >>> sampler_state = SamplerState(positions=jnp.array([[0.0, 0.0, 0.0], [2, 0.0, 0.0], [0.0, 2, 0.0]])*unit.nanometer, + >>> box_vectors=jnp.array([[10, 0.0, 0.0], [0.0, 10, 0.0], [0.0, 0.0, 10]])*unit.nanometer) + >>> + >>> from chiron.neighbors import NeighborListNsqrd, OrthogonalPeriodicSpace + >>> nbr_list = NeighborListNsqrd(OrthogonalPeriodicSpace(), cutoff=1.2*unit.nanometer, skin=0.4*unit.nanometer) + >>> + >>> # build the neighborlist + >>> nbr_list.build_from_state(sampler_state) # build the pair list from the sampler state + >>> + >>> # calculate which particles are interacting along with their distances and displacement vectors + >>> n_neighbors, neighbor_list, padding_mask, dist, r_ij = nbr_list.calculate(sampler_state.positions) + >>> + >>> # check the neighborlist + >>> if nbr_list.check(sampler_state.positions): + >>> nbr_list.build_from_state(sampler_state) # rebuild the pair list from the sampler state """ @@ -410,18 +425,14 @@ def __init__( ): if not isinstance(space, Space): raise TypeError(f"space must be of type Space, found {type(space)}") - if not cutoff.unit.is_compatible(unit.angstrom): - raise ValueError( - f"cutoff must be a unit.Quantity with units of distance, cutoff.unit = {cutoff.unit}" - ) + if not skin.unit.is_compatible(unit.angstrom): raise ValueError( f"cutoff must be a unit.Quantity with units of distance, skin.unit = {skin.unit}" ) - self.cutoff = cutoff.value_in_unit_system(unit.md_unit_system) - self.skin = skin.value_in_unit_system(unit.md_unit_system) - self.cutoff_and_skin = self.cutoff + self.skin + self.cutoff = cutoff + self.skin = skin self.n_max_neighbors = n_max_neighbors self.space = space @@ -429,8 +440,44 @@ def __init__( # this does not imply that the neighborlist is up to date self.is_built = False - # note, we need to use the partial decorator in order to use the jit decorate - # so that it knows to ignore the `self` argument + @property + def cutoff(self) -> unit.Quantity: + return self._cutoff + + @cutoff.setter + def cutoff(self, cutoff: unit.Quantity) -> None: + if not cutoff.unit.is_compatible(unit.nanometer): + raise ValueError( + f"cutoff must be a unit.Quantity with units of distance, cutoff.unit = {cutoff.unit}" + ) + self._cutoff = cutoff + + # if we change the cutoff or skin we need to rebuild + self.is_built = False + + @property + def skin(self) -> unit.Quantity: + return self._skin + + @skin.setter + def skin(self, skin: unit.Quantity) -> None: + if not skin.unit.is_compatible(unit.nanometer): + raise ValueError( + f"skin must be a unit.Quantity with units of distance, skin.unit = {skin.unit}" + ) + self._skin = skin + + # if we change the cutoff or skin we need to rebuild + self.is_built = False + + # Note, we need to use the partial decorator in order to jit the method of a class + # so that it knows to ignore the `self` argument. However, this treats it as static. + # This means that any changes to the internal class values within the class will not be updated in this function. + # Hence it is important that we pass the values as arguments to the function so we are getting updated values + # and not reference anything via self.variable_name. + # Alternatively we could create a custom pytree instead of declaring the class static in this function, + # but I don't think that is necessary if we just pass the values as arguments. + @partial(jax.jit, static_argnums=(0,)) def _pairs_mask(self, particle_ids: jnp.array): """ @@ -464,9 +511,18 @@ def _pairs_mask(self, particle_ids: jnp.array): return temp_mask + # note: since n_max_neighbors dictates the output size, we need to define it as a static argument + # to allow us to jit this function @partial(jax.jit, static_argnums=(0, 5)) def _build_neighborlist( - self, particle_i, reduction_mask, pid, positions, n_max_neighbors + self, + particle_i, + reduction_mask, + pid, + positions, + n_max_neighbors, + cutoff_and_skin, + box_vectors, ): """ Jitted function to build the neighbor list for a single particle @@ -481,6 +537,10 @@ def _build_neighborlist( X,Y,Z positions of all particles n_max_neighbors: int Maximum number of neighbors for each particle. Used for padding arrays for efficient jax computations + cutoff_and_skin: float + Cutoff distance for the neighborlist plus the skin distance, in nanometers. + box_vectors: jnp.array + Box vectors for the system Returns ------- @@ -493,18 +553,22 @@ def _build_neighborlist( """ # calculate the displacement between particle i and all other particles - r_ij, dist = self.space.displacement(particle_i, positions) + # we could pass this as a function instead of referencing self.space, but I ran into issues + # doing that with vmap, and I haven't been able to figure out how to resolve that yet -- CRI + r_ij, dist = self.space.displacement(particle_i, positions, box_vectors) # neighbor_mask will be an array of length n_particles (i.e., length of positions) # where each element is True if the particle is a neighbor, False if it is not # subject to both the cutoff+skin and the reduction mask that eliminates double counting and self-interactions neighbor_mask = jnp.where( - (dist < self.cutoff_and_skin) & (reduction_mask), True, False + (dist < cutoff_and_skin) & (reduction_mask), True, False ) # when we pad the neighbor list, we will use last particle id in the neighbor list - # this choice was made such that when we use the neighbor list in the masked energy calculat + # this choice was made such that when we use the neighbor list in the masked energy calculation # the padded values will result in reasonably well defined values fill_value = jnp.argmax(neighbor_mask) + # if the max value is the same as the particle of interest, which can occur if particle 0 has no neighbors + # we will just increment by 1 to avoid calculating a self interaction fill_value = jnp.where(fill_value == pid, fill_value + 1, fill_value) # count up the number of neighbors @@ -530,7 +594,7 @@ def build( box_vectors: Union[jnp.array, unit.Quantity], ): """ - Build the neighborlist from an array of positions and box vectors. + Build the neighbor list from an array of positions and box vectors. Parameters ---------- @@ -569,9 +633,10 @@ def build( self.ref_positions = positions self.box_vectors = box_vectors + cutoff_and_skin = self.cutoff + self.skin + # 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 # store the ids of all the particles self.particle_ids = jnp.array( @@ -589,13 +654,15 @@ def build( # n_neighbors: an array of shape (n_particles) where each element is the number of neighbors for that particle self.neighbor_mask, self.neighbor_list, self.n_neighbors = jax.vmap( - self._build_neighborlist, in_axes=(0, 0, 0, None, None) + self._build_neighborlist, in_axes=(0, 0, 0, None, None, None, None) )( self.ref_positions, reduction_mask, self.particle_ids, self.ref_positions, self.n_max_neighbors, + cutoff_and_skin.value_in_unit_system(unit.md_unit_system), + self.box_vectors, ) self.neighbor_list = self.neighbor_list.reshape(-1, self.n_max_neighbors) @@ -608,13 +675,15 @@ def build( self.n_max_neighbors = int(jnp.max(self.n_neighbors) + 10) self.neighbor_mask, self.neighbor_list, self.n_neighbors = jax.vmap( - self._build_neighborlist, in_axes=(0, 0, 0, None, None) + self._build_neighborlist, in_axes=(0, 0, 0, None, None, None, None) )( self.ref_positions, reduction_mask, self.particle_ids, self.ref_positions, self.n_max_neighbors, + cutoff_and_skin.value_in_unit_system(unit.md_unit_system), + self.box_vectors, ) self.neighbor_list = self.neighbor_list.reshape(-1, self.n_max_neighbors) @@ -623,7 +692,7 @@ def build( @partial(jax.jit, static_argnums=(0,)) def _calc_distance_per_particle( - self, particle1, neighbors, neighbor_mask, positions + self, particle1, neighbors, neighbor_mask, positions, cutoff, box_vectors ): """ Jitted function to calculate the distance between a particle and its neighbors @@ -638,6 +707,10 @@ def _calc_distance_per_particle( Mask to exclude padding from the neighbor list of particle1 positions: jnp.array X,Y,Z positions of all particles + cutoff: float + Cutoff distance for the neighborlist, in nanometers + box_vectors: jnp.array + Box vectors for the system Returns ------- @@ -656,11 +729,11 @@ def _calc_distance_per_particle( # calculate the displacement between particle i and all neighbors r_ij, dist = self.space.displacement( - positions[particles1], positions[neighbors] + positions[particles1], positions[neighbors], box_vectors ) # calculate the mask to determine if the particle is a neighbor # this will be done based on the interaction cutoff and using the neighbor_mask to exclude padding - mask = jnp.where((dist < self.cutoff) & (neighbor_mask), 1, 0) + mask = jnp.where((dist < cutoff) & (neighbor_mask), 1, 0) # calculate the number of pairs n_pairs = mask.sum() @@ -694,13 +767,27 @@ def calculate(self, positions: jnp.array): # changes to the box vectors require rebuilding the neighbor list n_neighbors, padding_mask, dist, r_ij = jax.vmap( - self._calc_distance_per_particle, in_axes=(0, 0, 0, None) - )(self.particle_ids, self.neighbor_list, self.neighbor_mask, positions) + self._calc_distance_per_particle, in_axes=(0, 0, 0, None, None, None) + )( + self.particle_ids, + self.neighbor_list, + self.neighbor_mask, + positions, + self.cutoff.value_in_unit_system(unit.md_unit_system), + self.box_vectors, + ) # mask = mask.reshape(-1, self.n_max_neighbors) return n_neighbors, self.neighbor_list, padding_mask, dist, r_ij @partial(jax.jit, static_argnums=(0,)) - def _calculate_particle_displacement(self, particle, positions, ref_positions): + def _calculate_particle_displacement( + self, + particle: int, + positions: jnp.array, + ref_positions: jnp.array, + skin: float, + box_vectors: jnp.array, + ): """ Calculate the displacement of a particle from the reference positions. If the displacement exceeds the half the skin distance, return True, otherwise return False. @@ -715,6 +802,11 @@ def _calculate_particle_displacement(self, particle, positions, ref_positions): Array of particle positions ref_positions: jnp.array Array of reference particle positions + skin: float + Skin distance for the neighborlist, in nanometers + box_vectors: jnp.array + Box vectors for the system + Returns ------- @@ -724,10 +816,10 @@ def _calculate_particle_displacement(self, particle, positions, ref_positions): # calculate the displacement of a particle from the initial positions r_ij, displacement = self.space.displacement( - positions[particle], ref_positions[particle] + positions[particle], ref_positions[particle], box_vectors ) - status = jnp.where(displacement >= self.skin / 2.0, True, False) + status = jnp.where(displacement >= skin / 2.0, True, False) del displacement return status @@ -753,8 +845,14 @@ def check(self, positions: jnp.array) -> bool: return True status = jax.vmap( - self._calculate_particle_displacement, in_axes=(0, None, None) - )(self.particle_ids, positions, self.ref_positions) + self._calculate_particle_displacement, in_axes=(0, None, None, None, None) + )( + self.particle_ids, + positions, + self.ref_positions, + self.skin.value_in_unit_system(unit.md_unit_system), + self.box_vectors, + ) if jnp.any(status): del status return True @@ -763,30 +861,39 @@ def check(self, positions: jnp.array) -> bool: return False -class PairList(PairsBase): +class PairListNsqrd(PairsBase): """ - N^2 pairlist implementation that returns the particle pair ids, displacement vectors, and distances. + Pair list implementation to determine which particles are interacting, subject to the simulation + Space class. This performs an N^2 calculation to determine the distances between particles which will be + inefficient for all but very small system sizes. This routine can be defined either with or without a cutoff. + + The `calculate` function of this class returns the particle pair ids, displacement vectors, and distances. + For efficiency of the jitted functions, the `calculate` function array sizes are fixed. For example, distance + has shape (n_particles, n_particles-1); note self-interactions are removed, hence n_particles-1). + The masking vector contains values of 1 for interacting particles and 0 otherwise. Parameters ---------- space: Space Class that defines how to calculate the displacement between two points and apply the boundary conditions - cutoff: float, default = 2.5 - Cutoff distance for the pair list calculation + cutoff: Optional[unit.Quantity], default = None + Cutoff distance for the pair list calculation. If None, the pair list will be calculated without a cutoff, + applying the boundary conditions as defined in space. Examples -------- - >>> from chiron.neighbors import PairList, OrthogonalPeriodicSpace - >>> from chiron.states import SamplerState >>> import jax.numpy as jnp + >>> import openmm.unit as unit >>> - >>> space = OrthogonalPeriodicSpace() - >>> pair_list = PairList(space, cutoff=2.5) + >>> from chiron.states import SamplerState >>> sampler_state = SamplerState(positions=jnp.array([[0.0, 0.0, 0.0], [2, 0.0, 0.0], [0.0, 2, 0.0]]), >>> box_vectors=jnp.array([[10, 0.0, 0.0], [0.0, 10, 0.0], [0.0, 0.0, 10]])) + >>> + >>> from chiron.neighbors import PairListNsqrd, OrthogonalPeriodicSpace + >>> pair_list = PairListNsqrd(OrthogonalPeriodicSpace(), cutoff=1.2*unit.nanometer) >>> pair_list.build_from_state(sampler_state) >>> >>> # mask and distances are of shape (n_particles, n_particles-1), - >>> displacement_vectors of shape (n_particles, n_particles-1, 3) + >>> # displacement_vectors of shape (n_particles, n_particles-1, 3) >>> # mask, is a bool array that is True if the particle is within the cutoff distance, False if it is not >>> # n_pairs is of shape (n_particles) and is per row sum of the mask. The mask ensure we also do not double count pairs >>> n_pairs, mask, distances, displacement_vectors = pair_list.calculate(sampler_state.positions) @@ -795,22 +902,35 @@ class PairList(PairsBase): def __init__( self, space: Space, - cutoff: unit.Quantity = unit.Quantity(1.2, unit.nanometer), + cutoff: Optional[unit.Quantity] = None, ): if not isinstance(space, Space): raise TypeError(f"space must be of type Space, found {type(space)}") - if not cutoff.unit.is_compatible(unit.angstrom): - raise ValueError( - f"cutoff must be a unit.Quantity with units of distance, cutoff.unit = {cutoff.unit}" - ) - self.cutoff = cutoff.value_in_unit_system(unit.md_unit_system) + # keeping this public in case we want to change it later + # validation is performed in the setter + self.cutoff = cutoff + self.space = space # set a a simple variable to know if this has at least been built once as opposed to just initialized # this does not imply that the neighborlist is up to date self.is_built = False + @property + def cutoff(self): + return self._cutoff + + @cutoff.setter + def cutoff(self, cutoff): + if cutoff is not None: + if not cutoff.unit.is_compatible(unit.angstrom): + raise ValueError( + f"cutoff must be a unit.Quantity with units of distance, cutoff.unit = {cutoff.unit}" + ) + self._cutoff = cutoff + # Since this is just a simple pair list, we do not need to rebuild by changing a cutoff + # note, we need to use the partial decorator in order to use the jit decorate # so that it knows to ignore the `self` argument @partial(jax.jit, static_argnums=(0,)) @@ -891,8 +1011,6 @@ def build( self.n_particles = self.ref_positions.shape[0] # the PairsList assumes that the box vectors do not change between building and calculating the neighbor list - # changes to the box vectors require rebuilding the list - self.space.box_vectors = self.box_vectors # store the ids of all the particles self.particle_ids = jnp.array(range(0, positions.shape[0]), dtype=jnp.uint32) @@ -902,9 +1020,64 @@ def build( self.is_built = True - @partial(jax.jit, static_argnums=(0,)) - def _calc_distance_per_particle( - self, particle1, neighbors, neighbor_mask, positions + @partial(jax.jit, static_argnums=(0)) + def _calc_distance_per_particle_with_cutoff( + self, particle1, neighbors, neighbor_mask, positions, cutoff, box_vectors + ): + """ + Jitted function to calculate the distance between a particle and all possible neighbors + + Parameters + ---------- + particle1: int + Particle id + neighbors: jnp.array + Array of particle ids for the possible particle pairs of particle1 + neighbor_mask: jnp.array + Mask to exclude double particles to prevent double counting + positions: jnp.array + X,Y,Z positions of all particles, shaped (n_particles, 3) + cutoff: float + Cutoff distance for the interaction. + box_vectors: jnp.array + Box vectors for the system + + Returns + ------- + n_pairs: int + Number of interacting pairs for the particle + mask: jnp.array + Mask to exclude padding particles not within the cutoff particle1. + If a particle is within the interaction cutoff, the mask is 1, otherwise it is 0 + Array has shape (n_particles, n_particles-1) as it excludes self interactions + dist: jnp.array + Array of distances between the particle and all other particles in the system. + Array has shape (n_particles, n_particles-1) as it excludes self interactions + r_ij: jnp.array + Array of displacement vectors between the particle and all other particles in the system. + Array has shape (n_particles, n_particles-1, 3) as it excludes self interactions + . + + """ + # repeat the particle id for each neighbor + particles1 = jnp.repeat(particle1, neighbors.shape[0]) + + # calculate the displacement between particle i and all neighbors + r_ij, dist = self.space.displacement( + positions[particles1], positions[neighbors], box_vectors + ) + # calculate the mask to determine if the particle is a neighbor + # this will be done based on the interaction cutoff and using the neighbor_mask to exclude padding + mask = jnp.where((dist < cutoff) & (neighbor_mask), 1, 0) + + # calculate the number of pairs + n_pairs = mask.sum() + + return n_pairs, mask, dist, r_ij + + @partial(jax.jit, static_argnums=(0)) + def _calc_distance_per_particle_no_cutoff( + self, particle1, neighbors, neighbor_mask, positions, box_vectors ): """ Jitted function to calculate the distance between a particle and all possible neighbors @@ -919,6 +1092,8 @@ def _calc_distance_per_particle( Mask to exclude double particles to prevent double counting positions: jnp.array X,Y,Z positions of all particles, shaped (n_particles, 3) + box_vectors: jnp.array + Box vectors of the system Returns ------- @@ -935,17 +1110,18 @@ def _calc_distance_per_particle( Array of displacement vectors between the particle and all other particles in the system. Array has shape (n_particles, n_particles-1, 3) as it excludes self interactions + """ # repeat the particle id for each neighbor particles1 = jnp.repeat(particle1, neighbors.shape[0]) # calculate the displacement between particle i and all neighbors r_ij, dist = self.space.displacement( - positions[particles1], positions[neighbors] + positions[particles1], positions[neighbors], box_vectors ) # calculate the mask to determine if the particle is a neighbor # this will be done based on the interaction cutoff and using the neighbor_mask to exclude padding - mask = jnp.where((dist < self.cutoff) & (neighbor_mask), 1, 0) + mask = jnp.where(neighbor_mask, 1, 0) # calculate the number of pairs n_pairs = mask.sum() @@ -980,12 +1156,30 @@ def calculate(self, positions: jnp.array): f"Positions must have shape ({self.n_particles}, 3), found {positions.shape}" ) - # positions = self.space.wrap(positions) - - 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, positions) - + # if we did not define a cutoff, we will + if self.cutoff is None: + n_neighbors, padding_mask, dist, r_ij = jax.vmap( + self._calc_distance_per_particle_no_cutoff, + in_axes=(0, 0, 0, None, None), + )( + self.particle_ids, + self.all_pairs, + self.reduction_mask, + positions, + self.box_vectors, + ) + else: + n_neighbors, padding_mask, dist, r_ij = jax.vmap( + self._calc_distance_per_particle_with_cutoff, + in_axes=(0, 0, 0, None, None, None), + )( + self.particle_ids, + self.all_pairs, + self.reduction_mask, + positions, + self.cutoff.value_in_unit_system(unit.md_unit_system), + self.box_vectors, + ) return n_neighbors, self.all_pairs, padding_mask, dist, r_ij def check(self, positions: jnp.array) -> bool: diff --git a/chiron/potential.py b/chiron/potential.py index d567a4c..1d2d340 100644 --- a/chiron/potential.py +++ b/chiron/potential.py @@ -264,7 +264,7 @@ def compute_energy(self, positions: jnp.array, nbr_list=None, debug_mode=False): raise ValueError("Neighborlist must be built before use") # ensure that the cutoff in the neighbor list is the same as the cutoff in the potential - if nbr_list.cutoff != self.cutoff: + if nbr_list.cutoff.value_in_unit_system(unit.md_unit_system) != self.cutoff: raise ValueError( f"Neighborlist cutoff ({nbr_list.cutoff}) must be the same as the potential cutoff ({self.cutoff})" ) diff --git a/chiron/states.py b/chiron/states.py index c76613b..6ebdca1 100644 --- a/chiron/states.py +++ b/chiron/states.py @@ -282,7 +282,7 @@ def get_reduced_potential( ---------- sampler_state : SamplerState The sampler state for which to compute the reduced potential. - nbr_list : NeighborList or PairList, optional + nbr_list : NeighborList or PairListNsqrd, optional The neighbor list or pair list routine to use for calculating the reduced potential. Returns @@ -343,7 +343,7 @@ def calculate_reduced_potential_at_states( The sampler state for which to compute the reduced potential. thermodynamic_states : list of ThermodynamicState The thermodynamic states for which to compute the reduced potential. - nbr_list : NeighborList or PairList, optional + nbr_list : NeighborList or PairListNsqrd, optional Returns ------- list of float diff --git a/chiron/tests/test_convergence_tests.py b/chiron/tests/test_convergence_tests.py index deeb3e5..df30515 100644 --- a/chiron/tests/test_convergence_tests.py +++ b/chiron/tests/test_convergence_tests.py @@ -182,6 +182,8 @@ def test_langevin_dynamics_with_LJ_fluid(prep_temp_dir): ) +@pytest.mark.skip(reason="Tests takes too long") +@pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Test takes too long.") def test_ideal_gas(prep_temp_dir): from openmmtools.testsystems import IdealGas from openmm import unit @@ -221,12 +223,12 @@ def test_ideal_gas(prep_temp_dir): box_vectors=ideal_gas.system.getDefaultPeriodicBoxVectors(), ) - from chiron.neighbors import PairList, OrthogonalPeriodicSpace + 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 = PairList(OrthogonalPeriodicSpace(), cutoff=cutoff) + nbr_list = PairListNsqrd(OrthogonalPeriodicSpace(), cutoff=cutoff) nbr_list.build_from_state(sampler_state) from chiron.reporters import MCReporter diff --git a/chiron/tests/test_mcmc.py b/chiron/tests/test_mcmc.py index 77e3a24..bef08ba 100644 --- a/chiron/tests/test_mcmc.py +++ b/chiron/tests/test_mcmc.py @@ -409,11 +409,11 @@ def test_mc_barostat(prep_temp_dir): pressure=1.0 * unit.atmosphere, ) - from chiron.neighbors import PairList, OrthogonalPeriodicSpace + from chiron.neighbors import PairListNsqrd, OrthogonalPeriodicSpace # since particles are non-interacting and we will not displacece them, the pair list basically # does nothing in this case. - nbr_list = PairList(OrthogonalPeriodicSpace(), cutoff=0 * unit.nanometer) + nbr_list = PairListNsqrd(OrthogonalPeriodicSpace(), cutoff=0 * unit.nanometer) sampler_state, thermodynamic_state, nbr_list = barostat_move.update( sampler_state, thermodynamic_state, nbr_list diff --git a/chiron/tests/test_minization.py b/chiron/tests/test_minization.py index e5a7b8c..4faa41d 100644 --- a/chiron/tests/test_minization.py +++ b/chiron/tests/test_minization.py @@ -3,7 +3,7 @@ def test_minimization(): import jax.numpy as jnp from chiron.states import SamplerState - from chiron.neighbors import PairList, OrthogonalPeriodicSpace + from chiron.neighbors import PairListNsqrd, OrthogonalPeriodicSpace from openmm import unit # initialize testystem @@ -25,7 +25,7 @@ def test_minimization(): box_vectors=lj_fluid.system.getDefaultPeriodicBoxVectors(), ) # use parilist - nbr_list = PairList(OrthogonalPeriodicSpace(), cutoff=cutoff) + nbr_list = PairListNsqrd(OrthogonalPeriodicSpace(), cutoff=cutoff) nbr_list.build_from_state(sampler_state) # compute intial energy with and without pairlist @@ -88,7 +88,7 @@ def test_minimize_two_particles(): import jax.numpy as jnp from chiron.states import SamplerState - from chiron.neighbors import PairList, OrthogonalPeriodicSpace + from chiron.neighbors import PairListNsqrd, OrthogonalPeriodicSpace from openmm import unit from chiron.potential import LJPotential @@ -111,7 +111,7 @@ def test_minimize_two_particles(): * unit.nanometer, ) - pair_list = PairList(OrthogonalPeriodicSpace(), cutoff=cutoff) + pair_list = PairListNsqrd(OrthogonalPeriodicSpace(), cutoff=cutoff) pair_list.build_from_state(sampler_state) e_start = lj_potential.compute_energy(coordinates, pair_list) diff --git a/chiron/tests/test_pairs.py b/chiron/tests/test_pairs.py index e30c008..60df5ad 100644 --- a/chiron/tests/test_pairs.py +++ b/chiron/tests/test_pairs.py @@ -2,7 +2,7 @@ import pytest from chiron.neighbors import ( NeighborListNsqrd, - PairList, + PairListNsqrd, OrthogonalPeriodicSpace, OrthogonalNonperiodicSpace, ) @@ -13,78 +13,47 @@ def test_orthogonal_periodic_displacement(): # test that the incorrect box shapes throw an exception - with pytest.raises(ValueError): - space = OrthogonalPeriodicSpace(jnp.array([10.0, 10.0, 10.0])) - # test that incorrect units throw an exception - with pytest.raises(ValueError): - space = OrthogonalPeriodicSpace( - unit.Quantity( - jnp.array([[10.0, 0.0, 0.0], [0.0, 10.0, 0.0], [0.0, 0.0, 10.0]]), - unit.radians, - ) - ) - - space = OrthogonalPeriodicSpace( - jnp.array([[10.0, 0.0, 0.0], [0.0, 10.0, 0.0], [0.0, 0.0, 10.0]]) - ) - # test that the box vectors are set correctly - assert jnp.all( - space.box_vectors - == jnp.array([[10.0, 0.0, 0.0], [0.0, 10.0, 0.0], [0.0, 0.0, 10.0]]) - ) - - # test that the box lengths for an orthogonal box are set correctly - assert jnp.all(space._box_lengths == jnp.array([10.0, 10.0, 10.0])) + space = OrthogonalPeriodicSpace() + box_vectors = jnp.array([[10.0, 0.0, 0.0], [0.0, 10.0, 0.0], [0.0, 0.0, 10.0]]) # test calculation of the displacement_vector and distance between two points p1 = jnp.array([[0, 0, 0], [0, 0, 0]]) p2 = jnp.array([[1, 0, 0], [6, 0, 0]]) - r_ij, distance = space.displacement(p1, p2) + r_ij, distance = space.displacement(p1, p2, box_vectors) assert jnp.all(r_ij == jnp.array([[-1.0, 0.0, 0.0], [4.0, 0.0, 0.0]])) assert jnp.all(distance == jnp.array([1, 4])) # test that the periodic wrapping works as expected - wrapped_x = space.wrap(jnp.array([11, 0, 0])) + wrapped_x = space.wrap(jnp.array([11, 0, 0]), box_vectors) assert jnp.all(wrapped_x == jnp.array([1, 0, 0])) - wrapped_x = space.wrap(jnp.array([-1, 0, 0])) + wrapped_x = space.wrap(jnp.array([-1, 0, 0]), box_vectors) assert jnp.all(wrapped_x == jnp.array([9, 0, 0])) - wrapped_x = space.wrap(jnp.array([5, 0, 0])) + wrapped_x = space.wrap(jnp.array([5, 0, 0]), box_vectors) assert jnp.all(wrapped_x == jnp.array([5, 0, 0])) - wrapped_x = space.wrap(jnp.array([5, 12, -1])) + wrapped_x = space.wrap(jnp.array([5, 12, -1]), box_vectors) assert jnp.all(wrapped_x == jnp.array([5, 2, 9])) - # test the setter for the box vectors - space.box_vectors = jnp.array( - [[10.0, 0.0, 0.0], [0.0, 20.0, 0.0], [0.0, 0.0, 30.0]] - ) - assert jnp.all( - space._box_vectors - == jnp.array([[10.0, 0.0, 0.0], [0.0, 20.0, 0.0], [0.0, 0.0, 30.0]]) - ) - assert jnp.all(space._box_lengths == jnp.array([10.0, 20.0, 30.0])) - def test_orthogonal_nonperiodic_displacement(): - space = OrthogonalNonperiodicSpace( - jnp.array([[10.0, 0.0, 0.0], [0.0, 10.0, 0.0], [0.0, 0.0, 10.0]]) - ) + space = OrthogonalNonperiodicSpace() + box_vectors = jnp.array([[10.0, 0.0, 0.0], [0.0, 10.0, 0.0], [0.0, 0.0, 10.0]]) p1 = jnp.array([[0, 0, 0], [0, 0, 0]]) p2 = jnp.array([[1, 0, 0], [6, 0, 0]]) - r_ij, distance = space.displacement(p1, p2) + r_ij, distance = space.displacement(p1, p2, box_vectors) assert jnp.all(r_ij == jnp.array([[-1.0, 0.0, 0.0], [-6.0, 0.0, 0.0]])) assert jnp.all(distance == jnp.array([1, 6])) - wrapped_x = space.wrap(jnp.array([11, -1, 2])) + wrapped_x = space.wrap(jnp.array([11, -1, 2]), box_vectors) assert jnp.all(wrapped_x == jnp.array([11, -1, 2])) @@ -106,17 +75,16 @@ def test_neighborlist_pair(): ) space = OrthogonalPeriodicSpace() - cutoff = 1.1 - skin = 0.1 + cutoff = 1.1 * unit.nanometer + skin = 0.1 * unit.nanometer nbr_list = NeighborListNsqrd( space, - cutoff=unit.Quantity(cutoff, unit.nanometer), - skin=unit.Quantity(skin, unit.nanometer), + cutoff=cutoff, + skin=skin, n_max_neighbors=5, ) assert nbr_list.cutoff == cutoff assert nbr_list.skin == skin - assert nbr_list.cutoff_and_skin == cutoff + skin assert nbr_list.n_max_neighbors == 5 nbr_list.build_from_state(state) @@ -195,12 +163,12 @@ def test_neighborlist_pair(): def test_inputs(): space = OrthogonalPeriodicSpace() # every particle should interact with every other particle - cutoff = 2.1 - skin = 0.1 + cutoff = 2.1 * unit.nanometer + skin = 0.1 * unit.nanometer nbr_list = NeighborListNsqrd( space, - cutoff=unit.Quantity(cutoff, unit.nanometer), - skin=unit.Quantity(skin, unit.nanometer), + cutoff=cutoff, + skin=skin, n_max_neighbors=5, ) # check that the state is of the correct type @@ -247,24 +215,24 @@ def test_inputs(): with pytest.raises(TypeError): NeighborListNsqrd( 123, - cutoff=unit.Quantity(cutoff, unit.nanometer), - skin=unit.Quantity(skin, unit.nanometer), + cutoff=cutoff, + skin=skin, n_max_neighbors=5, ) # check units of cutoff with pytest.raises(ValueError): NeighborListNsqrd( space, - cutoff=unit.Quantity(cutoff, unit.radian), - skin=unit.Quantity(skin, unit.nanometer), + cutoff=unit.Quantity(123, unit.radian), + skin=unit.Quantity(123, unit.nanometer), n_max_neighbors=5, ) # check units of skin with pytest.raises(ValueError): NeighborListNsqrd( space, - cutoff=unit.Quantity(cutoff, unit.nanometer), - skin=unit.Quantity(skin, unit.radian), + cutoff=unit.Quantity(123, unit.nanometer), + skin=unit.Quantity(123, unit.radian), n_max_neighbors=5, ) @@ -294,12 +262,12 @@ def test_neighborlist_pair_multiple_particles(): space = OrthogonalPeriodicSpace() # every particle should interact with every other particle - cutoff = 2.1 - skin = 0.1 + cutoff = 2.1 * unit.nanometer + skin = 0.1 * unit.nanometer nbr_list = NeighborListNsqrd( space, - cutoff=unit.Quantity(cutoff, unit.nanometer), - skin=unit.Quantity(skin, unit.nanometer), + cutoff=cutoff, + skin=skin, n_max_neighbors=5, ) nbr_list.build_from_state(state) @@ -310,12 +278,12 @@ def test_neighborlist_pair_multiple_particles(): assert jnp.all(n_interacting == jnp.array([7, 6, 5, 4, 3, 2, 1, 0])) # every particle should be in the nieghbor list, but only a subset in the interacting range - cutoff = 1.1 - skin = 1.1 + cutoff = 1.1 * unit.nanometer + skin = 1.1 * unit.nanometer nbr_list = NeighborListNsqrd( space, - cutoff=unit.Quantity(cutoff, unit.nanometer), - skin=unit.Quantity(skin, unit.nanometer), + cutoff=cutoff, + skin=skin, n_max_neighbors=5, ) nbr_list.build_from_state(state) @@ -368,11 +336,10 @@ def test_pairlist_pair(): ) space = OrthogonalPeriodicSpace() - cutoff = 1.1 - skin = 0.1 - pair_list = PairList( + cutoff = 1.1 * unit.nanometer + pair_list = PairListNsqrd( space, - cutoff=unit.Quantity(cutoff, unit.nanometer), + cutoff=cutoff, ) assert pair_list.cutoff == cutoff @@ -382,7 +349,7 @@ def test_pairlist_pair(): assert jnp.all(pair_list.reduction_mask == jnp.array([[True], [False]])) assert pair_list.is_built == True - n_pairs, all_pairs, mask, dist, displacement = pair_list.calculate(coordinates) + n_pairs, all_pairs, mask, dist, displacement = pair_list.calculate(state.positions) assert jnp.all(n_pairs == jnp.array([1, 0])) assert jnp.all(all_pairs.shape == (2, 1)) @@ -394,10 +361,49 @@ def test_pairlist_pair(): assert pair_list.check(coordinates) == False - coordinates = coordinates = jnp.array([[0, 0, 0], [1, 0, 0], [1, 1, 0]]) + coordinates = jnp.array([[0, 0, 0], [1, 0, 0], [1, 1, 0]]) # we changed number of particles, and thus should rebuild assert pair_list.check(coordinates) == True + # test without using a cutoff + # this will be exactly the same as with a cutoff, given it is just two particles + cutoff = None + pair_list = PairListNsqrd( + space, + cutoff=None, + ) + pair_list.build_from_state(state) + + assert pair_list.cutoff == cutoff + n_pairs, all_pairs, mask, dist, displacement = pair_list.calculate(state.positions) + assert jnp.all(n_pairs == jnp.array([1, 0])) + assert jnp.all(all_pairs.shape == (2, 1)) + assert jnp.all(all_pairs == jnp.array([[1], [0]])) + assert jnp.all(mask == jnp.array([[1], [0]])) + assert jnp.all(dist == jnp.array([[1.0], [1.0]])) + assert displacement.shape == (2, 1, 3) + assert jnp.all(displacement == jnp.array([[[-1.0, 0.0, 0.0]], [[1.0, 0.0, 0.0]]])) + + # test the difference between a short cutoff with no interactions and the same + # system with no cutoff. + + # this test ultimately have no particles in the neighbor list + # because the cutoff is really short + cutoff = 0.5 * unit.nanometer + pair_list = PairListNsqrd(space, cutoff=cutoff) + + assert pair_list.cutoff == cutoff + pair_list.build_from_state(state) + n_pairs, all_pairs, mask, dist, displacement = pair_list.calculate(state.positions) + # the mask will all be false because the cutoff is too short + assert jnp.all(mask == jnp.array([[0], [0]])) + + # set the cutoff to None, and calculate all pairs in the box + pair_list.cutoff = None + n_pairs, all_pairs, mask, dist, displacement = pair_list.calculate(state.positions) + # the mask will have the single pair in the box be true + assert jnp.all(mask == jnp.array([[1], [0]])) + def test_pair_list_multiple_particles(): # test the pair list for multiple particles @@ -423,11 +429,11 @@ def test_pair_list_multiple_particles(): space = OrthogonalPeriodicSpace() # every particle should interact with every other particle - cutoff = 2.1 - skin = 0.1 - pair_list = PairList( + cutoff = 2.1 * unit.nanometer + skin = 0.1 * unit.nanometer + pair_list = PairListNsqrd( space, - cutoff=unit.Quantity(cutoff, unit.nanometer), + cutoff=cutoff, ) pair_list.build_from_state(state) @@ -454,8 +460,8 @@ def test_pair_list_multiple_particles(): # compare to nbr_list nbr_list = NeighborListNsqrd( space, - cutoff=unit.Quantity(cutoff, unit.nanometer), - skin=unit.Quantity(skin, unit.nanometer), + cutoff=cutoff, + skin=skin, n_max_neighbors=20, ) nbr_list.build_from_state(state) From 12d3b1001c860e365f743aa371759f4adcdb8b3f Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Wed, 21 Feb 2024 16:02:50 -0800 Subject: [PATCH 40/43] mised an instance where the reporter called space.box_vectors. fixed. --- Examples/LJ_MCMC.py | 18 ++++++++++-------- chiron/integrators.py | 2 +- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/Examples/LJ_MCMC.py b/Examples/LJ_MCMC.py index 63d2412..95f7956 100644 --- a/Examples/LJ_MCMC.py +++ b/Examples/LJ_MCMC.py @@ -96,10 +96,12 @@ # initialize a reporter to save the simulation data import os -filename_barostat = "test_mc_lj_barostat.h5" -if os.path.isfile(filename_barostat): - os.remove(filename_barostat) -reporter_barostat = MCReporter(filename_barostat, 1) + +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 @@ -112,11 +114,11 @@ autotune_interval=100, ) -filename_displacement = "test_mc_lj_disp.h5" +filename_barostat = "test_mc_lj_barostat.h5" +if os.path.isfile(filename_barostat): + os.remove(filename_barostat) +reporter_barostat = MCReporter(filename_barostat, 1) -if os.path.isfile(filename_displacement): - os.remove(filename_displacement) -reporter_displacement = MCReporter(filename_displacement, 10) from chiron.mcmc import MonteCarloBarostatMove diff --git a/chiron/integrators.py b/chiron/integrators.py index 50e6406..51993df 100644 --- a/chiron/integrators.py +++ b/chiron/integrators.py @@ -287,6 +287,6 @@ def _report( "elapsed_step": elapsed_step, } if nbr_list is not None: - d["box_vectors"] = nbr_list.space.box_vectors + d["box_vectors"] = nbr_list.box_vectors self.reporter.report(d) From 47ab2ccda18c1dfed650fd381f2425cde06f00f2 Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Wed, 21 Feb 2024 22:30:36 -0800 Subject: [PATCH 41/43] name refactoring in MCMC --- Examples/LJ_MCMC.py | 2 +- chiron/mcmc.py | 60 +++++++++++++++++++++------------------------ 2 files changed, 29 insertions(+), 33 deletions(-) diff --git a/Examples/LJ_MCMC.py b/Examples/LJ_MCMC.py index 95f7956..b3cb2e6 100644 --- a/Examples/LJ_MCMC.py +++ b/Examples/LJ_MCMC.py @@ -144,7 +144,7 @@ langevin_dynamics_move = LangevinDynamicsMove( timestep=1.0 * unit.femtoseconds, collision_rate=1.0 / unit.picoseconds, - number_of_steps=100, + number_of_steps=1000, reporter=reporter_langevin, report_interval=10, ) diff --git a/chiron/mcmc.py b/chiron/mcmc.py index c95ce46..d7d1b75 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -255,17 +255,14 @@ def update( nbr_list: PairsBase The updated neighbor/pair list. If a nbr_list is not set, this will be None. """ - calculate_current_reduced_potential = True + self._current_reduced_potential = None for i in range(self.number_of_moves): sampler_state, thermodynamic_state, nbr_list = self._step( sampler_state, thermodynamic_state, nbr_list, - calculate_current_reduced_potential=calculate_current_reduced_potential, ) - # after the first step, we don't need to recalculate the current reduced_potential, it will be stored - calculate_current_reduced_potential = False # I think it makes sense to use i + self.number_of_moves*self._move_iteration as our current "step" # otherwise, if we just used i, instances where self.report_interval > self.number_of_moves would only report on the @@ -346,7 +343,6 @@ def _step( current_sampler_state: SamplerState, current_thermodynamic_state: ThermodynamicState, current_nbr_list: Optional[PairsBase] = None, - calculate_current_reduced_potential: bool = True, ) -> Tuple[SamplerState, ThermodynamicState, Optional[PairsBase]]: """ Performs an individual MC step. @@ -361,8 +357,6 @@ def _step( Current thermodynamic state. current_nbr_list : Optional[PairsBase] Neighbor list associated with the current state. - calculate_current_reduced_potential : bool, optional - Whether to calculate the current reduced potential. Default is True. Returns ------- @@ -383,15 +377,17 @@ def _step( # we will need to calculate the reduced potential for the current state # this is toggled by the calculate_current_reduced_potential flag # otherwise, we can use the one that was saved from the last step, for efficiency - if calculate_current_reduced_potential: - current_reduced_pot = current_thermodynamic_state.get_reduced_potential( - current_sampler_state, current_nbr_list + if self._current_reduced_potential is None: + current_reduced_potential = ( + current_thermodynamic_state.get_reduced_potential( + current_sampler_state, current_nbr_list + ) ) - # save the current_reduced_pot so we don't have to recalculate + # save the current_reduced_potential so we don't have to recalculate # it on the next iteration if the move is rejected - self._current_reduced_pot = current_reduced_pot + self._current_reduced_potential = current_reduced_potential else: - current_reduced_pot = self._current_reduced_pot + current_reduced_potential = self._current_reduced_potential # propose a new state and calculate the log proposal ratio # this will be specific to the type of move @@ -403,17 +399,17 @@ def _step( ( proposed_sampler_state, proposed_thermodynamic_state, - proposed_reduced_pot, + proposed_reduced_potential, log_proposal_ratio, proposed_nbr_list, ) = self._propose( current_sampler_state, current_thermodynamic_state, - current_reduced_pot, + current_reduced_potential, current_nbr_list, ) - if jnp.isnan(proposed_reduced_pot): + if jnp.isnan(proposed_reduced_potential): decision = False else: # accept or reject the proposed state @@ -429,7 +425,7 @@ def _step( if decision: # save the reduced potential of the accepted state so # we don't have to recalculate it the next iteration - self._current_reduced_pot = proposed_reduced_pot + self._current_reduced_potential = proposed_reduced_potential # replace the current state with the proposed state # not sure this needs to be a separate function but for simplicity in outlining the code it is fine @@ -478,7 +474,7 @@ def _propose( self, current_sampler_state: SamplerState, current_thermodynamic_state: ThermodynamicState, - current_reduced_pot: float, + current_reduced_potential: float, current_nbr_list: Optional[PairsBase] = None, ) -> Tuple[SamplerState, ThermodynamicState, float, float, Optional[PairsBase]]: """ @@ -495,7 +491,7 @@ def _propose( Current sampler state. current_thermodynamic_state : ThermodynamicState, required Current thermodynamic state. - current_reduced_pot : float, required + current_reduced_potential : float, required Current reduced potential. current_nbr_list : PairsBase, required Neighbor list associated with the current state. @@ -506,7 +502,7 @@ def _propose( Proposed sampler state. proposed_thermodynamic_state : ThermodynamicState Proposed thermodynamic state. - proposed_reduced_pot : float + proposed_reduced_potential : float Proposed reduced potential. log_proposal_ratio : float Log proposal ratio. @@ -652,7 +648,7 @@ def _propose( self, current_sampler_state: SamplerState, current_thermodynamic_state: ThermodynamicState, - current_reduced_pot: float, + current_reduced_potential: float, current_nbr_list: Optional[PairsBase] = None, ) -> Tuple[SamplerState, ThermodynamicState, float, float, Optional[PairsBase]]: """ @@ -664,7 +660,7 @@ def _propose( Current sampler state. current_thermodynamic_state : ThermodynamicState, required Current thermodynamic state. - current_reduced_pot : float, required + current_reduced_potential : float, required Current reduced potential. current_nbr_list : Optional[PairsBase] Neighbor list associated with the current state. @@ -675,7 +671,7 @@ def _propose( Proposed sampler state. proposed_thermodynamic_state : ThermodynamicState Proposed thermodynamic state. - proposed_reduced_pot : float + proposed_reduced_potential : float Proposed reduced potential. log_proposal_ratio : float Log proposal ratio. @@ -741,18 +737,18 @@ def _propose( else: proposed_nbr_list = None - proposed_reduced_pot = current_thermodynamic_state.get_reduced_potential( + proposed_reduced_potential = current_thermodynamic_state.get_reduced_potential( proposed_sampler_state, proposed_nbr_list ) - log_proposal_ratio = -proposed_reduced_pot + current_reduced_pot + log_proposal_ratio = -proposed_reduced_potential + current_reduced_potential # since do not change the thermodynamic state we can return # 'current_thermodynamic_state' rather than making a copy return ( proposed_sampler_state, current_thermodynamic_state, - proposed_reduced_pot, + proposed_reduced_potential, log_proposal_ratio, proposed_nbr_list, ) @@ -869,7 +865,7 @@ def _propose( self, current_sampler_state: SamplerState, current_thermodynamic_state: ThermodynamicState, - current_reduced_pot: float, + current_reduced_potential: float, current_nbr_list: Optional[PairsBase] = None, ) -> Tuple[SamplerState, ThermodynamicState, float, float, Optional[PairsBase]]: """ @@ -881,7 +877,7 @@ def _propose( Current sampler state. current_thermodynamic_state : ThermodynamicState, required Current thermodynamic state. - current_reduced_pot : float, required + current_reduced_potential : float, required Current reduced potential. current_nbr_list : PairsBase, optional Neighbor list associated with the current state. @@ -892,7 +888,7 @@ def _propose( Proposed sampler state. proposed_thermodynamic_state : ThermodynamicState Proposed thermodynamic state. - proposed_reduced_pot : float + proposed_reduced_potential : float Proposed reduced potential. log_proposal_ratio : float Log proposal ratio. @@ -944,21 +940,21 @@ def _propose( proposed_sampler_state.positions, proposed_sampler_state.box_vectors ) - proposed_reduced_pot = current_thermodynamic_state.get_reduced_potential( + proposed_reduced_potential = current_thermodynamic_state.get_reduced_potential( proposed_sampler_state, proposed_nbr_list ) # NPT acceptance criteria was originally defined in McDonald 1972, https://doi.org/10.1080/00268977200100031 # (see equation 9). The acceptance probability is given by: # ⎡−β (ΔU + PΔV ) + N ln(V new /V old )⎤ log_proposal_ratio = -( - proposed_reduced_pot - current_reduced_pot + proposed_reduced_potential - current_reduced_potential ) + nr_of_atoms * jnp.log(proposed_volume / initial_volume) # we do not change the thermodynamic state so we can return 'current_thermodynamic_state' return ( proposed_sampler_state, current_thermodynamic_state, - proposed_reduced_pot, + proposed_reduced_potential, log_proposal_ratio, proposed_nbr_list, ) From 192604e5bea3ba7a0fcadc0b477396a4bcd1760d Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Fri, 23 Feb 2024 09:51:02 -0800 Subject: [PATCH 42/43] further addressing comments. --- Examples/LJ_langevin.py | 23 ++++-- chiron/integrators.py | 6 -- chiron/mcmc.py | 35 ++++++++- chiron/neighbors.py | 163 +++++++++++++++++++++++++++------------- 4 files changed, 162 insertions(+), 65 deletions(-) diff --git a/Examples/LJ_langevin.py b/Examples/LJ_langevin.py index 1086cdc..8bb38d1 100644 --- a/Examples/LJ_langevin.py +++ b/Examples/LJ_langevin.py @@ -8,8 +8,6 @@ from chiron.potential import LJPotential from openmm import unit -from chiron.utils import PRNG - # initialize the LennardJones potential in chiron # @@ -22,9 +20,12 @@ ) -from chiron.states import SamplerState, ThermodynamicState +from chiron.utils import PRNG PRNG.set_seed(1234) + +from chiron.states import SamplerState, ThermodynamicState + # define the sampler state sampler_state = SamplerState( positions=lj_fluid.positions, @@ -32,7 +33,6 @@ box_vectors=lj_fluid.system.getDefaultPeriodicBoxVectors(), ) - # define the thermodynamic state thermodynamic_state = ThermodynamicState( potential=lj_potential, @@ -42,14 +42,21 @@ from chiron.neighbors import NeighborListNsqrd, OrthogonalPeriodicSpace -# define the neighbor list for an orthogonal periodic space +# Set up a neighbor list for an orthogonal periodic box with a cutoff of 3.0 * sigma and skin of 0.5 * sigma, +# where sigma = 0.34 nm. +# The class we instantiate, NeighborListNsqrd, uses an O(N^2) calculation to build the neighbor list, +# but uses a buffer (i.e., the skin) to avoid needing to perform the O(N^2) calculation at every step. +# With this routine, the calculation at each step between builds is O(N*n_max_neighbors). +# For the conditions considered here, n_max_neighbors is set to 180 (note this will increase if necessary) +# and thus there is ~5 reduction in computational cost compared to a brute force approach (i.e., PairListNsqrd). + skin = 0.5 * unit.nanometer nbr_list = NeighborListNsqrd( OrthogonalPeriodicSpace(), cutoff=cutoff, skin=skin, n_max_neighbors=180 ) -# build the neighbor list from the sampler state +# perform the initial build of the neighbor list from the sampler state nbr_list.build_from_state(sampler_state) from chiron.reporters import LangevinDynamicsReporter @@ -72,6 +79,9 @@ integrator = LangevinIntegrator(reporter=reporter, report_interval=100) print("init_energy: ", lj_potential.compute_energy(sampler_state.positions, nbr_list)) +# run the simulation +# note, typically we will not be calling the integrator directly, +# but instead using the LangevinDynamics Move in the MCMC Sampler. updated_sampler_state, updated_nbr_list = integrator.run( sampler_state, thermodynamic_state, @@ -84,7 +94,6 @@ # read the data from the reporter with h5py.File("test_lj.h5", "r") as f: - print(f.keys()) energies = f["potential_energy"][:] steps = f["step"][:] diff --git a/chiron/integrators.py b/chiron/integrators.py index 51993df..fa04063 100644 --- a/chiron/integrators.py +++ b/chiron/integrators.py @@ -182,12 +182,6 @@ def run( # r x += (timestep_unitless * 0.5) * v - # we can actually skip this wrapping, and just wrap/check/rebuild - # right before we call the force again. - - # if nbr_list is not None: - # x = self._wrap_and_rebuild_neighborlist(x, nbr_list) - # o random_noise_v = random.normal(subkey, x.shape) v = (a * v) + (b * sigma_v * random_noise_v) diff --git a/chiron/mcmc.py b/chiron/mcmc.py index d7d1b75..3aa3aac 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -533,6 +533,23 @@ def _accept_or_reject( class MonteCarloDisplacementMove(MCMove): + """ + A Monte Carlo move that randomly displaces particles in the system. + + For each move, all particles will be randomly displaced at once, where the random displacement is drawn from + a normal distribution. The standard deviation of the distribution is defined by the `displacement_sigma` parameter. + + Displacements can be restricted to a subset of particles by defining the `atom_subset` parameter, which is a list of + particle indices that will be allowed to move. If `atom_subset` is not defined, all particles will be displaced. + + Note, the displacement moves are applied on a per-particle basis; this does not support collective moves. + + The value of the `displacement_sigma` can be autotuned to achieve a target acceptance ratio between 0.4 and 0.6, + by setting the autotune parameter to True. The frequency of autotuning is defined by setting `autotune_interval`. + + + """ + def __init__( self, displacement_sigma=1.0 * unit.nanometer, @@ -652,7 +669,7 @@ def _propose( current_nbr_list: Optional[PairsBase] = None, ) -> Tuple[SamplerState, ThermodynamicState, float, float, Optional[PairsBase]]: """ - Implement the logic specific to displacement changes. + Implements the logic specific to displacement moves. Parameters ---------- @@ -755,6 +772,22 @@ def _propose( class MonteCarloBarostatMove(MCMove): + """ + A Monte Carlo move that randomly changes the volume of the system. + + The volume change is drawn from a normal distribution with a mean of 0 and a standard deviation defined + by the product of the `volume_max_scale` parameter and the current volume. Particle positions are scaled + proportionately with the change in volume. This routine operates on a per-particle basis and does not support + collective moves (i.e., it is an "atomic" barostat move where particle center-of-mass positions are scaled; + it is not aware of "molecules" which would be scaled by the molecule center-of-mass). + + The `volume_max_scale` parameter can be autotuned to achieve a target acceptance ratio between 0.25 and 0.75, + by setting the autotune parameter to True. The frequency of autotuning is defined by setting `autotune_interval`. + Note, the maximum value of `volume_max_scale` is capped at 0.3 in the auto-tuning process. + + + """ + def __init__( self, volume_max_scale=0.01, diff --git a/chiron/neighbors.py b/chiron/neighbors.py index ccdc706..be80cf0 100644 --- a/chiron/neighbors.py +++ b/chiron/neighbors.py @@ -199,8 +199,20 @@ class PairsBase(ABC): def __init__( self, space: Space, - cutoff: unit.Quantity = unit.Quantity(1.2, unit.nanometer), + cutoff: Optional[unit.Quantity] = unit.Quantity(1.2, unit.nanometer), ): + """ + Initialize the PairsBase class + + Parameters + ---------- + space: Space + Class that defines how to calculate the displacement between two points and apply the boundary conditions + This should not be changed after initialization. + cutoff: unit.Quantity, default = 1.2 unit.nanometer + Cutoff distance for the neighborlist + + """ if not isinstance(space, Space): raise TypeError(f"space must be of type Space, found {type(space)}") if not cutoff.unit.is_compatible(unit.angstrom): @@ -351,38 +363,39 @@ def check(self, positions: jnp.array) -> bool: class NeighborListNsqrd(PairsBase): """ - Neighbor list implementation used to determine which particles are interacting (i.e., within the cutoff). - - This `calculate` function of this class returns the particle pair ids, displacement vectors, and distances - between pairs within the specified cutoff range, subject to the boundary conditions defined by the simulation - Space class passed to the constructor. - - The neighbor list (i.e., list of particles within cutoff+skin of a given particle) - is generated using an N^2 calculation rather than using a spatial partitioning scheme (e.g., cell-list). + A JAX based neighbor list implementation used to determine which pairs of particles are interacting + (i.e., those particles that fall within the specified cutoff). + The neighbor list (i.e., list of particles within a distance of cutoff+skin of a given particle) is generated + within the `build` function using an O(N^2) calculation rather than using a spatial partitioning scheme + (e.g., cell-list). The `calculate` function that uses the neighbor list to determine which particle pairs are + interacting and determine the distances and displacement vectors between interacting pairs of particles for + use in the calculation of the interaction energies/forces. The routines are subject to the boundary conditions + specified by the Space class. Notes: - This does not include self-interactions and only includes unique pairs (i.e., no double-counting). + This neighbor list not include self-interactions and only includes unique pairs (i.e., no double-counting). This is sometimes referred to as a "half" neighbor list. E.g. consider the pair of neighboring particles (A, B): in the "half" neighbor list approach, B is in the neighbor list of A, but A is not in the neighbor list of B as that pair is already accounted for. . - The output of the `calculate` function is padded to a fixed size, n_max_neighbors, to allow for efficient - jitted computations in JAX. As such, values need to be masked using the `padding_mask` array returned by the - `calculate` function. The padding mask is simple an array of 1s and 0s, where 1 indicates a valid neighbor and - 0 indicates padding. The code will automatically increase n_max_neighbors by 10 if the number of neighbors - exceeds the current value. + The output of the `calculate` function is padded to a fixed size, `n_max_neighbors` (default=100), + to allow for efficient jitted computations in JAX. As such, values need to be masked using the `padding_mask` + array returned by the `calculate` function. The padding mask is an array of 1s and 0s, where 1 indicates an + interacting neighbor and 0 indicates the pair is either non-interacting or simply a padded value. + The `build` function will iteratively increase `n_max_neighbors` by 10 until we can store all neighbors. - The check function, which indicates if the neighbor list should be rebuilt, will return True if: + The `check` function, which indicates if the neighbor list should be rebuilt, will return True if: - the number of particles changes - - any of the particles have moved more than half the skin distance from their positions at the time of last - building the neighbor list. + - any of the particles have moved more than half the skin distance from their reference positions (i.e., the + positions of particles when the neighbor list was last built). Parameters ---------- space: Space - Class that defines how to calculate the displacement between two points and apply the boundary conditions + Class that defines how to calculate the displacement between two points and apply the boundary conditions. + This should not be changed after initialization. cutoff: unit.Quantity, default = 1.2 unit.nanometer Cutoff distance for the neighborlist skin: unit.Quantity, default = 0.4 unit.nanometer @@ -436,8 +449,8 @@ def __init__( self.n_max_neighbors = n_max_neighbors self.space = space - # set a a simple variable to know if this has at least been built once as opposed to just initialized - # this does not imply that the neighborlist is up to date + # this variable will ensure that `calculate` will fail if we try to call it before building + # note: self.is_built=True does not imply that the neighborlist is up-to-date self.is_built = False @property @@ -453,6 +466,8 @@ def cutoff(self, cutoff: unit.Quantity) -> None: self._cutoff = cutoff # if we change the cutoff or skin we need to rebuild + # we will set the variable to ensure that attempts to call the calculate function will fail if + # we have not rebuilt the neighbor list self.is_built = False @property @@ -468,15 +483,18 @@ def skin(self, skin: unit.Quantity) -> None: self._skin = skin # if we change the cutoff or skin we need to rebuild + # we will set the variable to ensure that attempts to call the calculate function will fail if + # we have not rebuilt the neighbor list self.is_built = False - # Note, we need to use the partial decorator in order to jit the method of a class - # so that it knows to ignore the `self` argument. However, this treats it as static. - # This means that any changes to the internal class values within the class will not be updated in this function. - # Hence it is important that we pass the values as arguments to the function so we are getting updated values - # and not reference anything via self.variable_name. - # Alternatively we could create a custom pytree instead of declaring the class static in this function, - # but I don't think that is necessary if we just pass the values as arguments. + # Note, we need to use the partial decorator and declare self as static in order to JIT a function within a class. + # This approach treats internal variables of the class as static within this function; e.g., if set self.cutoff = 2, + # called the function, then changed it to 3, the value of self.cutoff in this function would still be 2. + # Thus, we need to pass any variables that may change as arguments, rather than referencing self.variable_name. + # While we could create a custom pytree instead of declaring the class as static (allowing us to reference class + # variables directly within the JITTED function), any changes to those internal variables, say self.cutoff, + # would mean a change to the hash of any JITTEd function that depends on the variable, requiring JAX to recompile + # the function, which is a slow operation. As such, it is also more efficient to just pass variables as arguments. @partial(jax.jit, static_argnums=(0,)) def _pairs_mask(self, particle_ids: jnp.array): @@ -511,7 +529,7 @@ def _pairs_mask(self, particle_ids: jnp.array): return temp_mask - # note: since n_max_neighbors dictates the output size, we need to define it as a static argument + # note: since n_max_neighbors dictates the output size, we will define it as a static argument # to allow us to jit this function @partial(jax.jit, static_argnums=(0, 5)) def _build_neighborlist( @@ -552,9 +570,13 @@ def _build_neighborlist( Number of neighbors for the particle """ - # calculate the displacement between particle i and all other particles - # we could pass this as a function instead of referencing self.space, but I ran into issues - # doing that with vmap, and I haven't been able to figure out how to resolve that yet -- CRI + # Calculate the displacement between particle i and all other particles + # NOTE: It would be safer to pass the displacement calculate as a callable function, instead of referencing + # self.space. If someone changes the boundary conditions (i.e., changes space in the class), + # self.space.displacement will not change since the self is marked as status. + # However, I ran into issues passing a function through vmap, and I haven't been able to figure out how to + # resolve it yet. I do not want to remove vmap, as that would require substantially changing the flow of + # the code. For now, I've noted in the docstring that space should not change after initialization -- CRI r_ij, dist = self.space.displacement(particle_i, positions, box_vectors) # neighbor_mask will be an array of length n_particles (i.e., length of positions) @@ -728,6 +750,7 @@ def _calc_distance_per_particle( particles1 = jnp.repeat(particle1, neighbors.shape[0]) # calculate the displacement between particle i and all neighbors + # See note above: if self.space changes, it will not show up here because self is static. r_ij, dist = self.space.displacement( positions[particles1], positions[neighbors], box_vectors ) @@ -814,7 +837,7 @@ def _calculate_particle_displacement( True if the particle is outside the skin distance, False if it is not. """ # calculate the displacement of a particle from the initial positions - + # again, note that if self.space changes, it will not show up here because self is static. r_ij, displacement = self.space.displacement( positions[particle], ref_positions[particle], box_vectors ) @@ -863,14 +886,22 @@ def check(self, positions: jnp.array) -> bool: class PairListNsqrd(PairsBase): """ - Pair list implementation to determine which particles are interacting, subject to the simulation - Space class. This performs an N^2 calculation to determine the distances between particles which will be - inefficient for all but very small system sizes. This routine can be defined either with or without a cutoff. - - The `calculate` function of this class returns the particle pair ids, displacement vectors, and distances. - For efficiency of the jitted functions, the `calculate` function array sizes are fixed. For example, distance - has shape (n_particles, n_particles-1); note self-interactions are removed, hence n_particles-1). - The masking vector contains values of 1 for interacting particles and 0 otherwise. + A class that implements a simple pair list using JAX that determine which pairs of particles are interacting. + This class can be defined with cutoff (i.e., only returning information about pairs separated by distances + less than the cutoff) or without a cutoff (i.e., information about all possible pairs are returned). + Note, in both cases, distances are calculated using the boundary conditions defined by the simulation Space class + and only unique pairs are returned (i.e., no double counting and no self-interactions). + + This performs an O(N^2) calculation each time the `calculate` function is called and thus will be inefficient + for all but very small system sizes. + + The calculate function will return various pieces of information about the interacting pairs + (e.g., number of neighbors, neighbor ids, distances, displacement vectors) that can be used to calculate the + interaction potential/force. For efficiency of the jitted functions, the `calculate` function array + sizes are fixed. For example, distance has shape (n_particles, n_particles-1), regardless of the number of particles + that are actually neighbors (note: self interactions are removed hence n_particles-1). The `padding_mask` array + returned by `calculate` is used to exclude those pairs that are not interacting. The `padding_mask` contains values + of 1 for interacting particles and 0 for non-interacting. Parameters ---------- @@ -879,6 +910,7 @@ class PairListNsqrd(PairsBase): cutoff: Optional[unit.Quantity], default = None Cutoff distance for the pair list calculation. If None, the pair list will be calculated without a cutoff, applying the boundary conditions as defined in space. + Examples -------- >>> import jax.numpy as jnp @@ -892,11 +924,11 @@ class PairListNsqrd(PairsBase): >>> pair_list = PairListNsqrd(OrthogonalPeriodicSpace(), cutoff=1.2*unit.nanometer) >>> pair_list.build_from_state(sampler_state) >>> - >>> # mask and distances are of shape (n_particles, n_particles-1), - >>> # displacement_vectors of shape (n_particles, n_particles-1, 3) - >>> # mask, is a bool array that is True if the particle is within the cutoff distance, False if it is not - >>> # n_pairs is of shape (n_particles) and is per row sum of the mask. The mask ensure we also do not double count pairs - >>> n_pairs, mask, distances, displacement_vectors = pair_list.calculate(sampler_state.positions) + >>> # n_pairs is of shape (n_particles) and is per row sum of the padding_mask. + >>> # pairs, padding mask and distances are of shape (n_particles, n_particles-1), + >>> # displacement_vectors are of shape (n_particles, n_particles-1, 3) + >>> # padding_mask, is a bool array that is True if the particle is within the cutoff distance, False if it is not + >>> n_pairs, pairs, padding_mask, distances, displacement_vectors = pair_list.calculate(sampler_state.positions) """ def __init__( @@ -904,6 +936,17 @@ def __init__( space: Space, cutoff: Optional[unit.Quantity] = None, ): + """ + Initialize the PairListNsqrd class + + Parameters + ---------- + space: Space + Class that defines how to calculate the displacement between two points and apply the boundary conditions. + This should not change after initialization. + cutoff: Optional[unit.Quantity], default = None + Cutoff distance for the pair list calculation. If None, the pair list will be calculated without a cutoff. + """ if not isinstance(space, Space): raise TypeError(f"space must be of type Space, found {type(space)}") @@ -913,12 +956,21 @@ def __init__( self.space = space - # set a a simple variable to know if this has at least been built once as opposed to just initialized - # this does not imply that the neighborlist is up to date + # the init function does not setup the internal arrays we need to use calculate + # this is handled in the `build` function + # this variable can be used to check that the pair list has been built before trying to use it self.is_built = False @property def cutoff(self): + """ + Cutoff distance for the pair list calculation. If None, the pair list will be calculated without a cutoff. + + Returns + ------- + cutoff: unit.Quantity + Cutoff distance for the pair list calculation. If None, the pair list will be calculated without a cutoff. + """ return self._cutoff @cutoff.setter @@ -928,11 +980,18 @@ def cutoff(self, cutoff): raise ValueError( f"cutoff must be a unit.Quantity with units of distance, cutoff.unit = {cutoff.unit}" ) + # Note, since this is just a simple pair list, we do not need to rebuild by changing the cutoff self._cutoff = cutoff - # Since this is just a simple pair list, we do not need to rebuild by changing a cutoff - # note, we need to use the partial decorator in order to use the jit decorate - # so that it knows to ignore the `self` argument + # Note, we need to use the partial decorator and declare self as static in order to JIT a function within a class. + # As mentioned in a comment above in the NeighborListNsqrd class, this approach treats internal variables of the + # class as static within this function; e.g., if set self.cutoff = 2, called the function, then changed it to 3, + # the value of self.cutoff in this function would still be 2. Thus, we need to pass any variables that may change + # as arguments, rather than referencing self.variable_name. While we could create a custom pytree instead of + # declaring the class as static (allowing us to reference class variables directly within the JITTED function), + # any changes to those internal variables, say self.cutoff, would mean a change to the hash of any JITTEd function + # that depends on the variable, requiring JAX to recompile the function, which is a slow operation. + # As such, it is also more efficient to just pass variables as arguments. @partial(jax.jit, static_argnums=(0,)) def _pairs_and_mask(self, particle_ids: jnp.array): """ @@ -1063,6 +1122,7 @@ def _calc_distance_per_particle_with_cutoff( particles1 = jnp.repeat(particle1, neighbors.shape[0]) # calculate the displacement between particle i and all neighbors + # See note above: if self.space changes, it will not show up here because self is static. r_ij, dist = self.space.displacement( positions[particles1], positions[neighbors], box_vectors ) @@ -1116,6 +1176,7 @@ def _calc_distance_per_particle_no_cutoff( particles1 = jnp.repeat(particle1, neighbors.shape[0]) # calculate the displacement between particle i and all neighbors + # See note above: if self.space changes, it will not show up here because self is static. r_ij, dist = self.space.displacement( positions[particles1], positions[neighbors], box_vectors ) From ba3817dbb24787aee3bd7384c098ecb3c1c53e1b Mon Sep 17 00:00:00 2001 From: chrisiacovella Date: Tue, 27 Feb 2024 08:51:28 -0800 Subject: [PATCH 43/43] Added accumulated for number_of_attemps_made instead of elapsed_steps. (I missed this comment in the PR comments) --- chiron/mcmc.py | 46 +++++++++++++++++++++++++++++++--------------- 1 file changed, 31 insertions(+), 15 deletions(-) diff --git a/chiron/mcmc.py b/chiron/mcmc.py index 3aa3aac..64d71d2 100644 --- a/chiron/mcmc.py +++ b/chiron/mcmc.py @@ -38,6 +38,9 @@ def __init__( # we need to keep track of which iteration we are on self._move_iteration = 0 + # we also need to keep track of attempts made (i.e., total elapsed steps), in case the number_of_moves is changed + self._number_of_attempts_made = 0 + from loguru import logger as log if self.reporter is not None: @@ -77,6 +80,13 @@ def update( """ pass + @property + def number_of_attemps_made(self): + """ + Return the total number of steps that have been attempted in the move. + """ + return self._number_of_attempts_made + class LangevinDynamicsMove(MCMCMove): def __init__( @@ -176,6 +186,8 @@ def update( number_of_steps=self.number_of_moves, nbr_list=nbr_list, ) + # update the elapsed steps + self._number_of_attempts_made += self.number_of_moves if self.save_traj_in_memory: self.traj.append(self.integrator.traj) @@ -263,18 +275,19 @@ def update( thermodynamic_state, nbr_list, ) + self._number_of_attempts_made += 1 + + # We should use self._number_of_attempts_made as the "step" otherwise, if we just used i, instances where + # self.report_interval > self.number_of_moves would only report on the + # first step, which might actually be more frequent than we specify - # I think it makes sense to use i + self.number_of_moves*self._move_iteration as our current "step" - # otherwise, if we just used i, instances where self.report_interval > self.number_of_moves would only report on the - # first step, which might actually be more frequent than we specify - elapsed_step = i + self._move_iteration * self.number_of_moves if hasattr(self, "reporter"): if self.reporter is not None: - if elapsed_step % self.report_interval == 0: + if self._number_of_attempts_made % self.report_interval == 0: self._report( i, self._move_iteration, - elapsed_step, + self._number_of_attempts_made, self.n_accepted / self.n_proposed, sampler_state, thermodynamic_state, @@ -282,7 +295,10 @@ def update( ) if self.autotune: # if we only used i, we might never actually update the parameters if we have a move that is called infrequently - if elapsed_step % self.autotune_interval == 0 and elapsed_step > 0: + if ( + self._number_of_attempts_made % self.autotune_interval == 0 + and self._number_of_attempts_made > 0 + ): self._autotune() # keep track of how many times this function has been called self._move_iteration += 1 @@ -294,7 +310,7 @@ def _report( self, step: int, iteration: int, - elapsed_step: int, + number_of_attempts_made: int, acceptance_probability: float, sampler_state: SamplerState, thermodynamic_state: ThermodynamicState, @@ -312,7 +328,7 @@ def _report( The current step of the simulation move. iteration : int The current iteration of the move sequence (i.e., how many times has this been called thus far). - elapsed_step : int + number_of_attempts_made : int The total number of steps that have been taken in the simulation move. step+ nr_moves*iteration acceptance_probability : float The acceptance probability of the move. @@ -608,7 +624,7 @@ def _report( self, step: int, iteration: int, - elapsed_step: int, + number_of_attempts_made: int, acceptance_probability: float, sampler_state: SamplerState, thermodynamic_state: ThermodynamicState, @@ -623,7 +639,7 @@ def _report( The current step of the simulation move. iteration : int The current iteration of the move sequence (i.e., how many times has this been called thus far). - elapsed_step : int + number_of_attempts_made : int The total number of steps that have been taken in the simulation move. step+ nr_moves*iteration acceptance_probability : float The acceptance probability of the move. @@ -642,7 +658,7 @@ def _report( { "step": step, "iteration": iteration, - "elapsed_step": elapsed_step, + "number_of_attempts_made": number_of_attempts_made, "potential_energy": potential, "displacement_sigma": self.displacement_sigma.value_in_unit_system( unit.md_unit_system @@ -836,7 +852,7 @@ def _report( self, step: int, iteration: int, - elapsed_step: int, + number_of_attempts_made: int, acceptance_probability: float, sampler_state: SamplerState, thermodynamic_state: ThermodynamicState, @@ -850,7 +866,7 @@ def _report( The current step of the simulation move. iteration : int The current iteration of the move sequence (i.e., how many times has this been called thus far). - elapsed_step : int + number_of_attempts_made : int The total number of steps that have been taken in the simulation move. step+ nr_moves*iteration acceptance_probability : float The acceptance probability of the move. @@ -874,7 +890,7 @@ def _report( { "step": step, "iteration": iteration, - "elapsed_step": elapsed_step, + "number_of_attempts_made": number_of_attempts_made, "potential_energy": potential, "volume": volume, "box_vectors": sampler_state.box_vectors,