diff --git a/attic/rabfe/estimator.py b/attic/rabfe/estimator.py index 005703411..3261e4aa9 100644 --- a/attic/rabfe/estimator.py +++ b/attic/rabfe/estimator.py @@ -345,7 +345,7 @@ def deltaG_from_results( fwd_delta_u = model.beta * delta_Us[lambda_idx][0] rev_delta_u = model.beta * delta_Us[lambda_idx][1] - dG_exact, exact_bar_err = bar_with_bootstrapped_uncertainty(fwd_delta_u, rev_delta_u) + dG_exact, exact_bar_err = bar_with_pessimistic_uncertainty(fwd_delta_u, rev_delta_u) bar_dG += dG_exact / model.beta exact_bar_overlap = endpoint_correction.overlap_from_cdf(fwd_delta_u, rev_delta_u) @@ -407,7 +407,7 @@ def deltaG_from_results( rhs_xs=results[-1].xs, seed=2021, ) - dG_endpoint, endpoint_err = bar_with_bootstrapped_uncertainty( + dG_endpoint, endpoint_err = bar_with_pessimistic_uncertainty( model.beta * lhs_du, model.beta * np.array(rhs_du) ) dG_endpoint = dG_endpoint / model.beta diff --git a/tests/data/zero_overlap_ukln.npz b/tests/data/zero_overlap_ukln.npz new file mode 100644 index 000000000..fe8f3de67 Binary files /dev/null and b/tests/data/zero_overlap_ukln.npz differ diff --git a/tests/test_bar.py b/tests/test_bar.py index 417c7184d..eb174ff16 100644 --- a/tests/test_bar.py +++ b/tests/test_bar.py @@ -1,5 +1,7 @@ from functools import partial +from pathlib import Path from typing import Tuple +from unittest.mock import patch import numpy as np import pymbar @@ -8,7 +10,7 @@ from pymbar.testsystems import ExponentialTestCase from timemachine.fe.bar import ( - bar_with_bootstrapped_uncertainty, + bar_with_pessimistic_uncertainty, bootstrap_bar, compute_fwd_and_reverse_df_over_time, df_and_err_from_u_kln, @@ -83,8 +85,14 @@ def test_bootstrap_bar(sigma): # estimate 3 times df_ref, df_err_ref = df_and_err_from_u_kln(u_kln) - df_0, bootstrap_samples = bootstrap_bar(u_kln, n_bootstrap=n_bootstrap) - df_1, bootstrap_sigma = bar_with_bootstrapped_uncertainty(u_kln) + df_0, ddf_0, bootstrap_samples = bootstrap_bar(u_kln, n_bootstrap=n_bootstrap) + df_1, bootstrap_sigma = bar_with_pessimistic_uncertainty(u_kln) + + # Full errors should match exactly + assert df_err_ref == ddf_0 + + # The bootstrapped error should be as large or larger than the full error + assert bootstrap_sigma >= df_err_ref # assert estimates identical, uncertainties comparable print(f"bootstrap uncertainty = {bootstrap_sigma}, pymbar.MBAR uncertainty = {df_err_ref}") @@ -216,3 +224,31 @@ def test_compute_fwd_and_reverse_df_over_time(frames_per_step): # The values at the end should be nearly identical since they contain all the samples assert np.allclose(fwd[-1], rev[-1]) assert np.allclose(fwd_err[-1], rev_err[-1]) + + +def test_bootstrap_bar_and_regular_bar_match(): + """In cases where the u_kln has effectively no overlap, bootstrapping returns 0.0 + since the MBAR estimate is always zero. Checks that `bar_with_pessimistic_uncertainty` returns the (non-zero) error + from the MBAR estimate computed on all samples in these cases. + """ + test_ukln = Path(__file__).parent / "data" / "zero_overlap_ukln.npz" + u_kln = np.load(open(test_ukln, "rb"))["u_kln"] + + # The overlap should be closer to zero + overlap = pair_overlap_from_ukln(u_kln) + np.testing.assert_allclose(overlap, 0.0, atol=1e-12) + + boot_df, boot_df_err = bar_with_pessimistic_uncertainty(u_kln) + df, df_err = df_and_err_from_u_kln(u_kln) + assert boot_df == df + assert boot_df_err == df_err + + +@patch("timemachine.fe.bar.pymbar.MBAR.getFreeEnergyDifferences") +def test_nan_bar_error(mock_energy_diff): + df = np.zeros(shape=(1, 2)) + df_err = np.ones(shape=(1, 2)) * np.nan + mock_energy_diff.return_value = (df, df_err) + dummy_ukln = np.ones(shape=(2, 2, 100)) + _, boot_df_err = bar_with_pessimistic_uncertainty(dummy_ukln) + assert boot_df_err == 0.0 diff --git a/tests/test_fe_absolute_hydration.py b/tests/test_fe_absolute_hydration.py index 0d8e49fe5..f02640c32 100644 --- a/tests/test_fe_absolute_hydration.py +++ b/tests/test_fe_absolute_hydration.py @@ -19,7 +19,7 @@ def test_run_solvent(): assert res.plots.overlap_summary_png is not None assert res.plots.overlap_detail_png is not None - assert np.linalg.norm(res.final_result.dG_errs) < 10 + assert np.linalg.norm(res.final_result.dG_errs) < 20.0 assert len(res.frames) == n_windows assert len(res.boxes) == n_windows assert len(res.frames[0]) == n_frames diff --git a/timemachine/fe/bar.py b/timemachine/fe/bar.py index dfc2d2d12..4f33712cd 100644 --- a/timemachine/fe/bar.py +++ b/timemachine/fe/bar.py @@ -151,7 +151,7 @@ def df_from_u_kln( def bootstrap_bar( u_kln: NDArray, n_bootstrap: int = 100, maximum_iterations: int = DEFAULT_MAXIMUM_ITERATIONS -) -> Tuple[float, NDArray]: +) -> Tuple[float, float, NDArray]: """Given a 2-state u_kln matrix, subsample u_kln with replacement and re-run df_from_u_kln many times Parameters @@ -166,8 +166,12 @@ def bootstrap_bar( Returns ------- best_estimate : float - BAR(w_F, w_R, computeUncertainty=False) - bootstrap_samples: array + BAR(w_F, w_R) + + best_estimate_err : float + MBAR(w_F, w_R) error estimate, using all samples + + bootstrap_samples : array shape (n_bootstrap,) Notes @@ -178,7 +182,9 @@ def bootstrap_bar( u_kn, N_k = ukln_to_ukn(u_kln) mbar = pymbar.MBAR(u_kn, N_k, maximum_iterations=maximum_iterations) - full_bar_result = mbar.getFreeEnergyDifferences(compute_uncertainty=False)[0][0, 1] + df, ddf = mbar.getFreeEnergyDifferences() + full_bar_result = df[0, 1] + full_bar_err = ddf[0, 1] _, _, n = u_kln.shape @@ -196,15 +202,34 @@ def bootstrap_bar( ) bootstrap_samples.append(bar_result) - return full_bar_result, np.array(bootstrap_samples) + return full_bar_result, full_bar_err, np.array(bootstrap_samples) -def bar_with_bootstrapped_uncertainty( +def bar_with_pessimistic_uncertainty( u_kln: NDArray, n_bootstrap=100, maximum_iterations: int = DEFAULT_MAXIMUM_ITERATIONS ) -> Tuple[float, float]: - """Given 2-state u_kln, returns free energy difference and uncertainty computed by bootstrapping.""" + """Given 2-state u_kln, returns free energy difference and the uncertainty. The uncertainty can be produced either by + BAR using all samples or the bootstrapped error, whichever is greater. + + Parameters + ---------- + u_kln : array + 2-state u_kln matrix + n_bootstrap : int + number of bootstrap samples + maximum_iterations : int + maximum number of solver iterations to use for each sample + + Returns + ------- + best_estimate : float + BAR(w_F, w_R) + + uncertainty : float + `max(error_estimates)` where `error_estimates = [bootstrapped_bar_stddev, two_state_mbar_uncertainty]` + """ - df, bootstrap_dfs = bootstrap_bar(u_kln, n_bootstrap=n_bootstrap, maximum_iterations=maximum_iterations) + df, ddf, bootstrap_dfs = bootstrap_bar(u_kln, n_bootstrap=n_bootstrap, maximum_iterations=maximum_iterations) # warn if bootstrap distribution deviates significantly from normality normaltest_result = normaltest(bootstrap_dfs) @@ -212,8 +237,12 @@ def bar_with_bootstrapped_uncertainty( if normaltest_result.pvalue < pvalue_threshold: logger.warning(f"bootstrapped errors non-normal: {normaltest_result}") - # regardless, summarize as if normal - ddf = np.std(bootstrap_dfs) + # Take the max of the BAR error estimate using all samples and the bootstrapped error. Summarize as if normal regardless + # Use np.maximum to always return the NaN + if not np.isfinite(ddf): + logger.warning(f"BAR error estimate is not finite, setting to zero: {ddf}") + ddf = 0.0 + ddf = np.maximum(ddf, np.std(bootstrap_dfs)) return df, ddf diff --git a/timemachine/fe/free_energy.py b/timemachine/fe/free_energy.py index 967a28987..4c5d742ba 100644 --- a/timemachine/fe/free_energy.py +++ b/timemachine/fe/free_energy.py @@ -10,7 +10,7 @@ from timemachine.constants import BOLTZ from timemachine.fe import model_utils, topology from timemachine.fe.bar import ( - bar_with_bootstrapped_uncertainty, + bar_with_pessimistic_uncertainty, df_and_err_from_u_kln, pair_overlap_from_ukln, works_from_ukln, @@ -696,7 +696,7 @@ def estimate_free_energy_bar(u_kln_by_component: NDArray, temperature: float) -> u_kln = u_kln_by_component.sum(0) - df, df_err = bar_with_bootstrapped_uncertainty(u_kln) # reduced units + df, df_err = bar_with_pessimistic_uncertainty(u_kln) # reduced units kBT = BOLTZ * temperature dG, dG_err = df * kBT, df_err * kBT # kJ/mol