From 6af6ef6812e59e05ce7016cacb6d8f935fe9d512 Mon Sep 17 00:00:00 2001
From: wiederm <marcus.wieder@gmail.com>
Date: Tue, 19 Dec 2023 16:05:43 +0100
Subject: [PATCH] adopt example from openmmtools

---
 chiron/tests/test_minization.py |  3 +--
 chiron/tests/test_multistate.py | 39 +++++++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 2 deletions(-)
 create mode 100644 chiron/tests/test_multistate.py

diff --git a/chiron/tests/test_minization.py b/chiron/tests/test_minization.py
index 157df1e..b5cc290 100644
--- a/chiron/tests/test_minization.py
+++ b/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)    
\ No newline at end of file
diff --git a/chiron/tests/test_multistate.py b/chiron/tests/test_multistate.py
new file mode 100644
index 0000000..62c1b59
--- /dev/null
+++ b/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()