Add possibility to pass an array of temperatures

bussilab/wham.py

@@ -162,7 +162,8 @@ def wham(bias,
 
        T: float, optional
            The temperature of the system. This number is just used to divide the `bias` array in
-           order to make it adimensional.
+           order to make it adimensional. In case replicas are simulated at different temperatures,
+           it is possible to pass an array here, with ntraj elements.
 
        maxiter: int, optional
            Maximum number of iterations in the minimization procedure.
@@ -209,6 +210,10 @@ def wham(bias,
     nframes = bias.shape[0]
     ntraj = bias.shape[1]
 
+    if isinstance(T,float):
+        T=T*np.ones(ntraj)
+    T = coretools.ensure_np_array(T)
+
     # default values
     if frame_weight is None:
         frame_weight = np.ones(nframes)
@@ -219,7 +224,7 @@ def wham(bias,
     assert len(frame_weight) == nframes
 
     # divide by T once for all
-    shifted_bias = bias/T
+    shifted_bias = bias/T[np.newaxis,:]
     # track shifts
     shifts0 = np.min(shifted_bias, axis=0)
     shifted_bias -= shifts0[np.newaxis,:]

tests/test_wham.py

@@ -27,6 +27,8 @@ def test_wham(self):
         ]))**2), 0.0)
         d = np.exp(wham(np.array([[3, 5], [4, 4]]), frame_weight=(1, 10)).logW)
         self.assertAlmostEqual(np.sum((d-np.array([[0.06421006, 0.93578994]]))**2), 0.0)
+        e = np.exp(wham(np.array([[1, 5, 14], [2, 4.5, 12], [3, 4, 10]]),T=[1.0,0.5,2.0]).logW)
+        self.assertAlmostEqual(np.sum((e-np.array([0.41728571, 0.39684866, 0.18586563]))**2), 0.0)
 
     def test_wham1s(self):
         import numpy as np
@@ -52,6 +54,8 @@ def test_wham1s(self):
         ]))**2), 0.0)
         d = np.exp(wham(np.array([[3, 5], [4, 4]]), frame_weight=(1, 10),method="substitute").logW)
         self.assertAlmostEqual(np.sum((d-np.array([[0.06421006, 0.93578994]]))**2), 0.0)
+        e = np.exp(wham(np.array([[1, 5, 14], [2, 4.5, 12], [3, 4, 10]]),T=[1.0,0.5,2.0],method="substitute").logW)
+        self.assertAlmostEqual(np.sum((e-np.array([0.41728571, 0.39684866, 0.18586563]))**2), 0.0)
 
     def test_wham2(self):
         # test passing lists instead of np arrays