Correct TI error estimation II

src/alchemlyb/estimators/ti_.py

@@ -54,12 +54,14 @@ def fit(self, dHdl):
         means = dHdl.mean(level=dHdl.index.names[1:])
         variances = np.square(dHdl.sem(level=dHdl.index.names[1:]))
 
+        # get the lambda names
+        l_types = dHdl.index.names[1:]
+
         # obtain vector of delta lambdas between each state
         dl = means.reset_index()[means.index.names[:]].diff().iloc[1:].values
 
         # apply trapezoid rule to obtain DF between each adjacent state
         deltas = (dl * (means.iloc[:-1].values + means.iloc[1:].values)/2).sum(axis=1)
-        d_deltas = (dl**2 * (variances.iloc[:-1].values + variances.iloc[1:].values)/4).sum(axis=1)
 
         # build matrix of deltas between each state
         adelta = np.zeros((len(deltas)+1, len(deltas)+1))
@@ -70,8 +72,19 @@ def fit(self, dHdl):
             dout = []
             for i in range(len(deltas) - j):
                 out.append(deltas[i] + deltas[i+1:i+j+1].sum())
-                dout.append(d_deltas[i] + d_deltas[i+1:i+j+1].sum())
 
+                # Define additional zero lambda
+                a = [0.0] * len(l_types)
+
+                # Define dl series' with additional zero lambda on the left and right
+                dll = np.insert(dl[i:i + j + 1], 0, [a], axis=0)
+                dlr = np.append(dl[i:i + j + 1], [a], axis=0)
+
+                # Get a series of the form: x1, x1 + x2, ..., x(n-1) + x(n), x(n)
+                dllr = dll + dlr
+
+                # Append deviation of free energy difference between state i and i+j+1
+                dout.append((dllr ** 2 * variances.iloc[i:i + j + 2].values / 4).sum(axis=1).sum())
             adelta += np.diagflat(np.array(out), k=j+1)
             ad_delta += np.diagflat(np.array(dout), k=j+1)
 

src/alchemlyb/tests/test_ti_estimators.py

@@ -52,6 +52,30 @@ def gmx_expanded_ensemble_case_3_dHdl():
 
     return dHdl
 
+def gmx_water_particle_with_total_energy_dHdl():
+    dataset = alchemtest.gmx.load_water_particle_with_total_energy()
+
+    dHdl = pd.concat([gmx.extract_dHdl(filename, T=300)
+                      for filename in dataset['data']['AllStates']])
+
+    return dHdl
+
+def gmx_water_particle_with_potential_energy_dHdl():
+    dataset = alchemtest.gmx.load_water_particle_with_potential_energy()
+
+    dHdl = pd.concat([gmx.extract_dHdl(filename, T=300)
+                      for filename in dataset['data']['AllStates']])
+
+    return dHdl
+
+def gmx_water_particle_without_energy_dHdl():
+    dataset = alchemtest.gmx.load_water_particle_without_energy()
+
+    dHdl = pd.concat([gmx.extract_dHdl(filename, T=300)
+                      for filename in dataset['data']['AllStates']])
+
+    return dHdl
+
 def amber_simplesolvated_charge_dHdl():
     dataset = alchemtest.amber.load_simplesolvated()
 
@@ -71,18 +95,23 @@ def amber_simplesolvated_vdw_dHdl():
 
 class TIestimatorMixin:
 
-    @pytest.mark.parametrize('X_delta_f', ((gmx_benzene_coul_dHdl(), 3.089),
-                                           (gmx_benzene_vdw_dHdl(), -3.056),
-                                           (gmx_expanded_ensemble_case_1_dHdl(), 76.220),
-                                           (gmx_expanded_ensemble_case_2_dHdl(), 76.247),
-                                           (gmx_expanded_ensemble_case_3_dHdl(), 76.387),
-                                           (amber_simplesolvated_charge_dHdl(), -60.114),
-                                           (amber_simplesolvated_vdw_dHdl(), 3.824)))
+    @pytest.mark.parametrize('X_delta_f', ((gmx_benzene_coul_dHdl(), 3.089, 0.02157),
+                                           (gmx_benzene_vdw_dHdl(), -3.056, 0.04863),
+                                           (gmx_expanded_ensemble_case_1_dHdl(), 76.220, 0.15568),
+                                           (gmx_expanded_ensemble_case_2_dHdl(), 76.247, 0.15889),
+                                           (gmx_expanded_ensemble_case_3_dHdl(), 76.387, 0.12532),
+                                           (gmx_water_particle_with_total_energy_dHdl(), -11.696, 0.091775),
+                                           (gmx_water_particle_with_potential_energy_dHdl(), -11.751, 0.091149),
+                                           (gmx_water_particle_without_energy_dHdl(), -11.687, 0.091604),
+                                           (amber_simplesolvated_charge_dHdl(), -60.114, 0.08186),
+                                           (amber_simplesolvated_vdw_dHdl(), 3.824, 0.13254)))
     def test_get_delta_f(self, X_delta_f):
         est = self.cls().fit(X_delta_f[0])
         delta_f = est.delta_f_.iloc[0, -1]
+        d_delta_f = est.d_delta_f_.iloc[0, -1]
 
         assert X_delta_f[1] == pytest.approx(delta_f, rel=1e-3)
+        assert X_delta_f[2] == pytest.approx(d_delta_f, rel=1e-3)
 
 class TestTI(TIestimatorMixin):
     """Tests for TI.