Fix assertions in test_multistate_class and test_sampler_state_inputs

chiron/tests/test_multistate.py

@@ -76,7 +76,7 @@ def ho_multistate_sampler_multiple_ks() -> MultiStateSampler:
     Create a multi-state sampler for a harmonic oscillator system with different spring constants.
     Returns
     -------
-    MultiStateSampler 
+    MultiStateSampler
         The multi-state sampler object.
     """
     from openmm import unit
@@ -131,11 +131,11 @@ def test_multistate_class(ho_multistate_sampler_multiple_minima: MultiStateSampl
 
     Parameters:
     -------
-    ho_multistate_sampler_multiple_minima: MultiStateSampler 
+    ho_multistate_sampler_multiple_minima: MultiStateSampler
         An instance of the MultiStateSampler class.
     Raises:
     -------
-    AssertionError: 
+    AssertionError:
         If any of the assertions fail.
 
     """
@@ -212,6 +212,6 @@ def test_multistate_run(ho_multistate_sampler_multiple_ks: MultiStateSampler):
     print(ho_sampler.delta_f_ij_analytical)
     print(ho_sampler._last_mbar_f_k_offline)
 
-    assert np.isclose(
-        ho_sampler.delta_f_ij_analytical[0], ho_sampler._last_mbar_f_k_offline
+    assert np.allclose(
+        ho_sampler.delta_f_ij_analytical[0], ho_sampler._last_mbar_f_k_offline, atol=0.1
     )

chiron/tests/test_states.py

@@ -124,11 +124,12 @@ def test_sampler_state_inputs():
         x0=unit.Quantity(jnp.array([[1, 2, 3]]), unit.nanometers),
         box_vectors=openmm_box,
     )
-    assert jnp.all(
+    assert jnp.allclose(
         state.box_vectors
         == jnp.array(
             [[4.0311456, 0.0, 0.0], [0.0, 4.0311456, 0.0], [0.0, 0.0, 4.0311456]]
-        )
+        ),
+        atol=1e-4,
     )
 
     # openmm box vectors end up as a list with contents; check to make sure we capture an error if we pass a bad list