Handle undetermined energies in BAR calculations #1098

mcwitt · 2023-07-25T23:18:21Z

Since #1084, we detect numerical overflows during the evaluation of potentials and represent undetermined energies as NaNs. This has exposed some cases where we previously returned invalid energies (for example, overflows appear to occur frequently during the initial evaluation of the BAR df error and overlap between the end states, before the first iteration of bisection).

Now, when we encounter an energy overflow, the resulting NaN in the u_kln matrix causes pymbar.MBAR (used to compute overlap) to fail with a LinAlgError, crashing the simulation. Also, pymbar.BAR warns and returns zeros for the $\Delta f$ and uncertainty estimates when there are NaN work values.

This PR addresses the latter issues by:

Detecting NaNs in u_kln. When we detect a NaN, we raise a warning and replace NaNs with np.inf (representing a configuration with zero probability).
Replacing usage of the pymbar.BAR estimator with pymbar.MBAR (on a 2-state u_kn matrix). The latter interprets np.inf as a configuration with zero weight (as desired), while the former returns (0.0, 0.0) if there are any infs or NaNs.
Switching to a cost function for bisection based on MBAR overlap (rather than bootstrapped $\Delta f$ error). See Handle undetermined energies in BAR calculations #1098 (comment)

This uncovered a related issue where the timeout parameter in bootstrap_bar, intended to cap computational cost, can lead to nondeterminism. This was addressed by

Removing timeout logic and the timeout parameter from bootstrap_bar and bar_with_bootstrapped_uncertainty
Modifying MBAR relative_tolerance and maximum_iterations to reduce cost
Upgrading pymbar from 3.0.5 to 3.1.0 (3.0.6 fixed a bug that prevented maximum_iterations from being respected)

Todo:

Add test for case with many NaNs, e.g. completely non-overlapping states
Is bootstrapped $\Delta f$ error still reliable for bisection?
- ~~seems robust for non-overlapping states in testing: uncertainty estimate produced by MBAR is finite and large~~
- EDIT: point above was wrong; MBAR-estimated $\Delta f$ errors are finite and large for pairs with zero overlap, but bootstrapped errors are zero in this case (since MBAR estimates $\Delta f \approx 0$ for each bootstrap sample). Addressed by switching to bisection cost function based on MBAR overlap.

tests/test_bar.py

timemachine/fe/bar.py

tests/test_bar.py

mcwitt · 2023-07-26T19:33:52Z

Even after preventing BAR calculations from failing when there are undetermined energies, there is a separate issue that bootstrapped $\Delta f$ errors are no longer useful for bisection: in the extreme case of zero overlap between states, pymbar's MBAR returns a $\Delta f$ estimate of 0.0 for every bootstrap sample, so the resulting bootstrapped error is 0. (See the failing test added in 7495063)

In 54e4b5c I modified the bisection cost function to bisect pairs of states with lowest MBAR overlap, rather than maximum bootstrapped $\Delta f$ error. This works in the case of non-overlapping distributions, where the overlap is zero. We've also observed the MBAR overlap estimate to be more robust numerically compared with MBAR's $\Delta f$ error estimate.

This was previously optimized for BAR

df is undefined in this case

Was intended to cap computational effort, but can lead to nondeterminism

Upgrade to pymbar>=3.0.6 which fixed a bug causing maximum_iterations to be ignored: choderalab/pymbar#425

timemachine/fe/bar.py

maxentile · 2023-07-28T13:39:47Z

timemachine/fe/bar.py


+        bar_result = df_from_u_kln(
+            u_kln_sample,
+            initial_f_k=mbar.f_k,  # warm start


q: would it make sense to add bootstrap_maximum_iterations to signature of bootstrap_bar, then forward it here?

Good call, it seems like useful flexibility to be able to specify max iterations here. Added in f7a1ab4

Unsure whether it would be useful to be able to specify max iterations separately for the point estimate and the bootstrap samples, but this can probably be added later if useful.

Unsure whether it would be useful to be able to specify max iterations separately for the point estimate and the bootstrap samples, but this can probably be added later if useful.

Separately makes sense to me (to control expense), but can be added later if needed.

jkausrelay

Minor comments, overall LGTM.

jkausrelay · 2023-07-28T15:33:31Z

setup.py

@@ -102,7 +102,7 @@ def build_extension(self, ext):
        "jaxlib>0.4.1",
        "networkx",
        "numpy",
-        "pymbar>3.0.4,<4",
+        "pymbar>=3.0.6,<4",


nit: Should this be 3.1.0 or higher? Not sure why this is different from requirements

<3.0.6 definitely won't work because of choderalab/pymbar#425, which was merged in 3.0.6. In general I prefer to use version constraints in setup.py only to exclude known-incompatible versions.

When upgrading, I opted to use the latest release <4.

jkausrelay · 2023-07-28T15:36:54Z

tests/test_bar.py

+    print(f"bootstrap uncertainty = {bootstrap_sigma}, pymbar.MBAR uncertainty = {df_err_ref}")
+    assert df_0 == df_ref
+    assert df_1 == df_ref
+    assert len(bootstrap_samples) == n_bootstrap, "timed out on default problem size!"


nit: Is this still relevant?

Kept assertion but removed misleading message in 28c6d9f

jkausrelay · 2023-07-28T15:39:20Z

tests/test_free_energy.py

+    # should give the same result with inf
+    result_with_inf = estimate_free_energy_bar(np.array([u_kln_with_inf]), DEFAULT_TEMP)
+    assert result_with_nan.dG == result_with_inf.dG
+    assert result_with_nan.dG_err == result_with_inf.dG_err


nit: assert finite

Added in e52fd30

jkausrelay · 2023-07-28T15:42:12Z

timemachine/fe/bar.py

+        # As of pymbar 3.1.0, computation of the covariance matrix can raise an exception on incomplete convergence.
+        # In this case, return the unconverged estimate with NaN as uncertainty.
+        df = mbar.getFreeEnergyDifferences(compute_uncertainty=False)[0]
+        return df[0, 1], np.nan


nit: nan or inf?

Leaning toward NaN as more appropriate here, since we can't be sure of the failure reason. The choice here doesn't affect bisection, since we now use overlap.

Makes sense.

timemachine/fe/bar.py

* Change to MBAR in #1098 produces much more significant differences than BAR

* Change to MBAR in #1098 produces much more significant differences than BAR in cases of poor convergence

mcwitt added 9 commits July 25, 2023 16:02

Simplify: avoid recomputing u_kln

5cc41b2

Clean: use builtin alias for lru_cache(None)

f4eb94d

Compute 2-state delta f using MBAR

b651d53

Add failing test

7e8a9b2

Replace NaN with inf in u_kln

c51e808

Clean: use works_from_ukln

0b72da5

Fix test

da1d363

Update test to use df_from_u_kln

06066fb

Add tests for uniform distributions with partial and zero overlap

352c27d

mcwitt marked this pull request as ready for review July 26, 2023 15:16

mcwitt requested review from maxentile and jkausrelay July 26, 2023 15:16

mcwitt added 3 commits July 26, 2023 08:19

Clean: use df_and_err_from_u_kln

7b9c853

Strengthen tests: add comparison with exact dlogZ

f3f2442

Clean: add file-level nogpu mark

135dad6

maxentile reviewed Jul 26, 2023

View reviewed changes

tests/test_bar.py Outdated Show resolved Hide resolved

timemachine/fe/bar.py Outdated Show resolved Hide resolved

timemachine/fe/bar.py Outdated Show resolved Hide resolved

tests/test_bar.py Show resolved Hide resolved

mcwitt added 8 commits July 26, 2023 09:15

Strengthen partial overlap test to compare with exact result

d3120b8

Fix typo, add docstring note about u_kln convention

0ed9719

Refactor to avoid forwarding kwargs

75da78a

Fix missed update

8da93ba

Fix sign errors

4ab37d8

Strengthen test: also check that error estimate is consistent

1c54793

Add failing test

7495063

Switch to -log(overlap) as cost function for bisection

54e4b5c

mcwitt added 5 commits July 26, 2023 12:46

Fix docstring

7d41f55

Remove relative tolerance setting

5f93cf6

This was previously optimized for BAR

Clean, remove test case with zero overlap

ec4cc2b

df is undefined in this case

Merge branch 'master' into fix/handle-energy-overflow-bisection

35fe889

Add assertion for self-consistent iteration method

dc2c1b2

mcwitt added 9 commits July 27, 2023 14:28

Remove timeout logic from bootstrap_bar

358edc3

Was intended to cap computational effort, but can lead to nondeterminism

Reduce number of bootstrap samples

fd95420

Increase relative tolerance, reduce max iterations for MBAR

11e30ff

Upgrade to pymbar>=3.0.6 which fixed a bug causing maximum_iterations to be ignored: choderalab/pymbar#425

Use pymbar version released on pypi

be4d5c5

Catch pymbar exception on incomplete convergence

92a9f70

Update n_boostrap for consistency

4792f09

Clean: axis=-1 -> axis=2

2e0c9c1

Fix and clean test

f4a4b54

Add assertions for pymbar behavior

f7fe229

mcwitt requested a review from maxentile July 28, 2023 06:04

maxentile approved these changes Jul 28, 2023

View reviewed changes

mcwitt added 3 commits July 28, 2023 08:28

Remove mentions of timeout in docstring

0ddf9a0

Fix name of function in docstring, tweak formatting

f0f3414

Add option to specify max solver iterations for bootstrapping

f7a1ab4

maxentile approved these changes Jul 28, 2023

View reviewed changes

jkausrelay approved these changes Jul 28, 2023

View reviewed changes

badisa approved these changes Jul 28, 2023

View reviewed changes

mcwitt added 4 commits July 28, 2023 09:00

Fix typo

6950d82

Remove obsolete assertion failure message

28c6d9f

Assert finite results with nan and inf inputs

e52fd30

Merge branch 'master' into fix/handle-energy-overflow-bisection

e01c2a3

mcwitt enabled auto-merge (squash) July 28, 2023 16:06

Merge branch 'master' into fix/handle-energy-overflow-bisection

4a900c6

mcwitt merged commit dbababb into master Jul 28, 2023

mcwitt deleted the fix/handle-energy-overflow-bisection branch July 28, 2023 17:24

maxentile mentioned this pull request Jul 31, 2023

[wip] Return inf not nan from jax potentials #1085

Closed

badisa added a commit that referenced this pull request Aug 9, 2023

Updates tolerances of plotting forward and reverse dg

ffe104b

* Change to MBAR in #1098 produces much more significant differences than BAR

badisa mentioned this pull request Aug 9, 2023

Updates tolerances of plotting forward and reverse dg #1114

Merged

badisa added a commit that referenced this pull request Aug 9, 2023

Updates tolerances of plotting forward and reverse dg (#1114)

d7af228

* Change to MBAR in #1098 produces much more significant differences than BAR in cases of poor convergence

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Handle undetermined energies in BAR calculations #1098

Handle undetermined energies in BAR calculations #1098

mcwitt commented Jul 25, 2023 •

edited

Loading

mcwitt commented Jul 26, 2023 •

edited

Loading

maxentile Jul 28, 2023

mcwitt Jul 28, 2023

mcwitt Jul 28, 2023

maxentile Jul 28, 2023

jkausrelay left a comment

jkausrelay Jul 28, 2023

mcwitt Jul 28, 2023

jkausrelay Jul 28, 2023

mcwitt Jul 28, 2023

jkausrelay Jul 28, 2023

mcwitt Jul 28, 2023

jkausrelay Jul 28, 2023

mcwitt Jul 28, 2023

jkausrelay Jul 28, 2023

Handle undetermined energies in BAR calculations #1098

Handle undetermined energies in BAR calculations #1098

Conversation

mcwitt commented Jul 25, 2023 • edited Loading

mcwitt commented Jul 26, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jkausrelay left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mcwitt commented Jul 25, 2023 •

edited

Loading

mcwitt commented Jul 26, 2023 •

edited

Loading