adopt example from openmmtools

chiron/tests/test_minization.py

@@ -1,6 +1,5 @@
 def test_minimization():
     from chiron.minimze import minimize_energy
-    import jax
     import jax.numpy as jnp
 
     from chiron.states import SamplerState
@@ -32,4 +31,4 @@ def test_minimization():
 
     min_x = minimize_energy(sampler_state.x0, lj_potential.compute_energy, nbr_list)
     e = lj_potential.compute_energy(min_x, nbr_list)
-    assert jnp.isclose(e, -12506.332)
+    assert jnp.isclose(e, -12506.332)    

chiron/tests/test_multistate.py

@@ -0,0 +1,39 @@
+def test_HO():
+    import math
+    from openmm import unit
+    from openmmtools import testsystems
+    from chiron.multistate import MultiStateSampler
+    from chiron.mcmc import LangevinDynamicsMove
+    from chiron.states import ThermodynamicState
+
+    testsystem = testsystems.AlanineDipeptideImplicit()
+
+    n_replicas = 3
+    T_min = 298.0 * unit.kelvin  # Minimum temperature.
+    T_max = 600.0 * unit.kelvin  # Maximum temperature.
+    temperatures = [
+        T_min
+        + (T_max - T_min)
+        * (math.exp(float(i) / float(n_replicas - 1)) - 1.0)
+        / (math.e - 1.0)
+        for i in range(n_replicas)
+    ]
+    temperatures = [
+        T_min
+        + (T_max - T_min)
+        * (math.exp(float(i) / float(n_replicas - 1)) - 1.0)
+        / (math.e - 1.0)
+        for i in range(n_replicas)
+    ]
+    thermodynamic_states = [
+        ThermodynamicState(system=testsystem.system, temperature=T)
+        for T in temperatures
+    ]
+
+    # Initialize simulation object with options. Run with a langevin integrator.
+
+    move = LangevinDynamicsMove(timestep=2.0 * unit.femtoseconds, n_steps=50)
+    simulation = MultiStateSampler(mcmc_moves=move, number_of_iterations=2)
+
+    # Run the simulation
+    simulation.run()