diff --git a/src/alchemlyb/estimators/ti_.py b/src/alchemlyb/estimators/ti_.py index c11e9d29..7f25dbe3 100644 --- a/src/alchemlyb/estimators/ti_.py +++ b/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) diff --git a/src/alchemlyb/tests/test_ti_estimators.py b/src/alchemlyb/tests/test_ti_estimators.py index 60061ef5..602a4c3c 100644 --- a/src/alchemlyb/tests/test_ti_estimators.py +++ b/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.