diff --git a/bussilab/wham.py b/bussilab/wham.py index 5c518bf..3b98d7f 100644 --- a/bussilab/wham.py +++ b/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,:] diff --git a/tests/test_wham.py b/tests/test_wham.py index 134c940..c8feb5e 100644 --- a/tests/test_wham.py +++ b/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