use openmmtools testsystems directly

chiron/integrator.py

@@ -17,6 +17,20 @@ def __init__(
         box_vectors=None,
         progress_bar=False,
     ):
+        """
+        Initialize the LangevinIntegrator object.
+
+        Parameters
+        ----------
+        potential : NeuralNetworkPotential
+            Object representing the potential energy function.
+        topology : Topology
+            Object representing the molecular system.
+        box_vectors : array_like, optional
+            Box vectors for periodic boundary conditions.
+        progress_bar : bool, optional
+            Flag indicating whether to display a progress bar during integration.
+        """
         from .utils import get_list_of_mass
 
         self.box_vectors = box_vectors
@@ -35,6 +49,30 @@ def run(
         collision_rate=1.0 / unit.picoseconds,
         key=random.PRNGKey(0),
     ):
+        """
+        Run the integrator to perform Langevin dynamics molecular dynamics simulation.
+
+        Parameters
+        ----------
+        x0 : array_like
+            Initial positions of the particles.
+        temperature : unit.Quantity
+            Temperature of the system.
+        n_steps : int, optional
+            Number of simulation steps to perform.
+        stepsize : unit.Quantity, optional
+            Time step size for the integration.
+        collision_rate : unit.Quantity, optional
+            Collision rate for the Langevin dynamics.
+        key : jax.random.PRNGKey
+            Random key for generating random numbers.
+
+        Returns
+        -------
+        list of array_like
+            Trajectory of particle positions at each simulation step.
+        """
+
         kbT_unitless = (self.kB * temperature).value_in_unit_system(unit.md_unit_system)
         mass_unitless = jnp.array(self.mass.value_in_unit_system(unit.md_unit_system))
         sigma_v = jnp.sqrt(kbT_unitless / mass_unitless)

chiron/mcmc.py

@@ -37,6 +37,10 @@
 >>> sampler = MCMCSampler(thermodynamic_state, sampler_state, move=sequence_move)
 
 """
+from chiron.states import StateVariablesCollection
+from chiron.potential import NeuralNetworkPotential
+from openmm import unit
+from loguru import logger as log
 
 
 class GibbsSampler(object):
@@ -53,14 +57,16 @@ class GibbsSampler(object):
 
     """
 
-    def __init__(self, state_variables, move_set):
+    def __init__(self, state_variables: StateVariablesCollection, move_set: MoveSet):
         from copy import deepcopy
 
+        log.info("Initializing Gibbs sampler")
+
         # Make a deep copy of the state so that initial state is unchanged.
         self.state_variables = deepcopy(state_variables)
         self.move = move_set
 
-    def run(self, n_iterations=1):
+    def run(self, n_iterations: int = 1):
         """
         Run the sampler for a specified number of iterations.
 
@@ -71,3 +77,194 @@ def run(self, n_iterations=1):
 
         """
         # Apply move for n_iterations.
+
+
+from typing import Optional
+
+
+class LangevinDynamicsMove:
+    def __init__(
+        self,
+        n_steps: int,
+        NeuralNetworkPotential: NeuralNetworkPotential,
+        stepsize=1.0 * unit.femtoseconds,
+        collision_rate=1.0 / unit.picoseconds,
+    ):
+        """
+        Initialize the LangevinDynamicsMove with a molecular system.
+
+        Parameters
+        ----------
+        n_steps : int
+            Number of simulation steps to perform.
+        NeuralNetworkPotential : object
+            A representation of the molecular system (neural network potential and topology).
+        stepsize : unit.Quantity
+            Time step size for the integration.
+        collision_rate : unit.Quantity
+            Collision rate for the Langevin dynamics.
+        """
+
+        self.n_steps = n_steps
+        self.NeuralNetworkPotential = NeuralNetworkPotential
+        self.stepsize = stepsize
+        self.collision_rate = collision_rate
+        # system represents the potential energy function and topology
+        self.system: Optional[NeuralNetworkPotential] = None
+
+    def run(
+        self,
+        state_variables: StateVariablesCollection,
+    ):
+        """
+        Run the integrator to perform molecular dynamics simulation.
+
+        Args:
+            state_variables (StateVariablesCollection): State variables of the system.
+        """
+        from chiron.integrator import LangevinIntegrator
+
+        if self.system is None:
+            self.system = self.NeuralNetworkPotential(state_variables)
+
+        dynamics = LangevinIntegrator(
+            self.system, self.system.topology, state_variables.box_vectors
+        )
+
+        dynamics.run(
+            x0=state_variables.positions,
+            temperature=state_variables.temperature,
+            n_steps=self.n_steps,
+            stepsize=self.stepsize,
+            collision_rate=self.collision_rate,
+        )
+
+
+class MoveSet:
+    def __init__(self) -> None:
+        pass
+
+
+class MCMove:
+    def __init__(self, NeuralNetworkPotential: NeuralNetworkPotential):
+        """
+        Initialize the MCMove with a molecular system.
+
+        Parameters
+        ----------
+        system : object
+            A representation of the molecular system (e.g., coordinates, topology).
+        """
+        self.system = NeuralNetworkPotential
+
+    def _check_state_compatiblity(
+        self, old_state: StateVariablesCollection, new_state: StateVariablesCollection
+    ):
+        """
+        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):
+        """
+        Apply a Monte Carlo move to the system.
+
+        This method should be overridden by subclasses to define specific types of moves.
+
+        Raises
+        ------
+        NotImplementedError
+            If the method is not implemented in subclasses.
+        """
+        raise NotImplementedError("apply_move() must be implemented in subclasses")
+
+    def compute_acceptance_probability(
+        self, old_state: StateVariablesCollection, new_state: StateVariablesCollection
+    ):
+        """
+        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.NeuronNetworkPotential(old_state)
+        new_system = self.NeuronNetworkPotential(new_state).compute_energy(
+            new_state.position
+        )
+
+        energy_before_state_change = old_system.compute_energy(old_state.position)
+        enegy_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):
+        """
+        Implement the logic specific to rotamer changes.
+        """
+        pass
+
+
+class ProtonationStateMove(MCMove):
+    def apply_move(self):
+        """
+        Implement the logic specific to protonation state changes.
+        """
+        pass
+
+
+class TautomericStateMove(MCMove):
+    def apply_move(self):
+        """
+        Implement the logic specific to tautomeric state changes.
+        """
+        pass
+
+
+class ProposedPositionMove(MCMove):
+    def apply_move(self):
+        """
+        Implement the logic specific to a MC position change.
+        """
+        pass

chiron/tests/test_testsystems.py

@@ -1,10 +1,12 @@
-from chiron.testsystems import HarmonicOscillator
 from chiron.integrator import LangevinIntegrator
 
 
 def test_HO():
     from openmm.unit import kelvin
+    from chiron.potential import HarmonicOscillatorPotential
+    from openmmtools.testsystems import HarmonicOscillator
 
     ho = HarmonicOscillator()
-    integrator = LangevinIntegrator(ho.harmonic_potential, ho.topology)
-    integrator.run(ho.x0, temperature=300 * kelvin, n_steps=1000)
+    harmonic_potential = HarmonicOscillatorPotential(ho.K, ho.positions, ho.U0)
+    integrator = LangevinIntegrator(harmonic_potential, ho.topology)
+    integrator.run(ho.positions, temperature=300 * kelvin, n_steps=1000)