Fix the case where BayesMBAR cannot handle the case of only have two states

src/bayesmbar/bayesmbar.py

@@ -119,7 +119,8 @@ def logdensity(dF):
 
         ## compute the mean, covariance, and precision of dF based on the samples from the likelihood
         self._dF_mean_ll = jnp.mean(self._dF_samples_ll, axis=0)
-        self._dF_cov_ll = jnp.cov(self._dF_samples_ll.T)
+        # Ensure covariance is 2D for Cholesky decomposition (handle scalar case)
+        self._dF_cov_ll = jnp.atleast_2d(jnp.cov(self._dF_samples_ll.T))
 
         L = jnp.linalg.cholesky(self._dF_cov_ll)
         L_inv = jax.scipy.linalg.solve_triangular(L, jnp.eye(L.shape[0]), lower=True)

src/test/test_bayesmbar.py

@@ -1,7 +1,9 @@
+import numpy as np
 import pytest
 from pytest import approx
 from bayesmbar import BayesMBAR
 
+
 @pytest.mark.parametrize("method", ["Newton", "L-BFGS-B"])
 def test_BayesMBAR(setup_mbar_data, method):
     energy, num_conf, F_ref, energy_p, F_ref_p = setup_mbar_data
@@ -11,12 +13,45 @@ def test_BayesMBAR(setup_mbar_data, method):
         verbose=True,
         method=method,
     )
-    assert mbar.F_mode == approx(F_ref, abs = 1e-1)
-    assert mbar.F_mean == approx(F_ref, abs = 1e-1)
+    assert mbar.F_mode == approx(F_ref, abs=1e-1)
+    assert mbar.F_mean == approx(F_ref, abs=1e-1)
+
+
+def test_two_states():
+    M = 2  ## number of states
+    mu = np.linspace(0, 1, M)  ## equilibrium positions
+    k = np.random.uniform(10, 30, M)  ## force constants
+    sigma = np.sqrt(1.0 / k)
+    F_reference = -np.log(sigma)
+    F_reference -= F_reference[0]
+    n = 100
+    x = [np.random.normal(mu[i], sigma[i], (n,)) for i in range(M)]
+    x = np.concatenate(x)
+    u = 0.5 * k.reshape((-1, 1)) * (x - mu.reshape((-1, 1))) ** 2
+    num_conf = np.array([n for i in range(M)])
+    mbar = BayesMBAR(
+        u,
+        num_conf,
+        prior="uniform",
+        mean=None,
+        state_cv=None,
+        kernel=None,
+        sample_size=1000,
+        warmup_steps=100,
+        optimize_steps=0,
+        random_seed=0,
+        verbose=False,
+    )
+    F_reference = F_reference[-1] - F_reference[0]
+    F_mean = mbar.F_mean
+    F_mode = mbar.F_mode
+    assert (F_mean[-1] - F_mean[0]) == approx(F_reference, abs=1)
+    assert (F_mode[-1] - F_mode[0]) == approx(F_reference, abs=1)
+
+    # results = fastmbar.calculate_free_energies_of_perturbed_states(energy_p)
+    # results['F'] = results['F'] - results['F'].mean()
+    # assert results['F'] == approx(F_ref_p, abs = 1e-1)
 
-    #results = fastmbar.calculate_free_energies_of_perturbed_states(energy_p)
-    #results['F'] = results['F'] - results['F'].mean()
-    #assert results['F'] == approx(F_ref_p, abs = 1e-1)
 
 # import pytest
 # from pytest import approx