return partial convergence analysis if mbar fails for individual fraction by aqemia-benedict-tan · Pull Request #1984 · OpenFreeEnergy/openfe

aqemia-benedict-tan · 2026-06-04T07:29:08Z

When running ABFE calculations in OpenFE, MBAR sometimes fails to produce an estimate at low fractions of uncorrelated samples (~0.1-0.4). It would be helpful to still return the convergence trace where MBAR can generate a meaningful result (e.g. >0.7) This PR modifies get_forward_and_reverse_analysis under openmm_utils so that NaN values are appended to the forwards/reverse estimate lists where MBAR fails and the successful estimates are still returned.

Checklist

All new code is appropriately documented (user-facing code must have complete docstrings).
Ran pre-commit: you can run pre-commit locally or comment on this PR with pre-commit.ci autofix.

Developers certificate of origin

I certify that this contribution is covered by the MIT License here and the Developer Certificate of Origin at https://developercertificate.org/.

IAlibay

Thanks @aqemia-benedict-tan, this looks great - just the one comment on the code as-is.

The main thing we would need (and why I have to request changes) is that we need to add tests to cover the new behaviour

Ideally we would want to:

Test that the check the array of forward and backwards energies returned make sense (similar to

openfe/src/openfe/tests/protocols/test_openmmutils.py

Line 360 in 5347082

np.array(

)
Something in tests/analysis/test_plotting.py that checks that passing in an array with NaN's works fine and probably a manual check that the plot looks ok.

IAlibay · 2026-06-04T08:55:38Z

+                # MBAR can fail to converge at low fractions of uncorrelated
+                # samples. If this happens, append NaN and carry on so that
+                # estimates at higher fractions are still reported.
+                try:


[nit] Given the similar-ish duplicated code, it might be good to just move this to a helper method that takes in the sub-sampled u_ln, fraction, and either "reverse" or "forward" for the warning string.

codecov · 2026-06-04T09:21:05Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 90.16%. Comparing base (75cb2e8) to head (d3b12f7).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1984      +/-   ##
==========================================
- Coverage   94.95%   90.16%   -4.79%     
==========================================
  Files         216      216              
  Lines       20517    20578      +61     
==========================================
- Hits        19481    18554     -927     
- Misses       1036     2024     +988

Flag	Coverage Δ
fast-tests	`90.16% <100.00%> (?)`
slow-tests	`?`

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

…e analysis

aqemia-benedict-tan · 2026-06-04T13:44:19Z

Thanks for the feedback @IAlibay, I have refactored to include the helper function _get_fraction_free_energy for the purpose of appending NaN in the event of low fraction failure. I have also added unit tests to verify that clean arrays are returned in the event of the lowest fraction in the forwards estimate failing and to assert that the convergence plot is returned successfully if an array containing NaN estimates is passed to it. The resulting plot looks like this:

IAlibay · 2026-06-04T15:54:45Z

Thanks @aqemia-benedict-tan ! I'lll have another look at this early next week.

IAlibay

Sorry - I somewhat changed my mind on what the behaviour should look like on failure, that will require some code changes.

Could you also:

Add a news entry under openfe/news.
Add yourself to the CITATION.cff author list?

Thanks!

IAlibay · 2026-06-10T12:54:47Z

+                fraction = chunk / N_l[0]

                # Forward
-                DG, dDG = self._get_free_energy(


I've been thinking about this a little bit more and I've come to the realization that we probably should NaN both the forward and the reverse chunk if either fails.

My reasoning is that whilst a failure is directly caused by a lack of convergence in the estimator, it's more broadly caused by too few data points (generally too high a decorrelation time). So whilst one direction could pass (e.g. the reduced potentials happen to be easier to deal with), it's probably not reasonable to include either.

Yes this makes sense - commits 7b7dba3 and dbf9b71 refactor the _get_fraction_free_energy helper to return NaN for both directions, if either the forwards or reverse estimate fails.
Commit 42ee838 updates the tests, the test at src/openfe/tests/protocols/test_openmmutils.py is adapted to ensure NaN is returned for both directions in the case that the lowest fraction of samples failes in either the forwards or reverse direction.
The new graph from the plotting test looks like this:

Amazing, thanks!

IAlibay · 2026-06-10T12:58:58Z

+    uncorrelated samples) should be rendered as gaps without raising.
+    """
+    forward_reverse = _make_forward_reverse(
+        forward=[np.nan, np.nan, -9.2, -9.4, -9.5, -9.6, -9.65, -9.7, -9.72, -9.75],


What if the nan happens towards the end? It's unlikely but could happen.

My understanding is that the fill_between call would just do nothing. In that case - it might be better to just return None for the whole thing.

Good point! Commit ce66327 fixes this by returningNone if the final element of the reverse and forwards estimate (i.e. fraction 1.0) isNaN. A unit test is also added for this scenario.

IAlibay · 2026-06-10T13:01:56Z

            rtol=5e-01,
        )  # fmt: skip

+    def test_forward_and_reverse_nan_on_mbar_failure(self, analyzer):


Thanks for doing this - this test looks great!

…n where MBAR fails

aqemia-benedict-tan · 2026-06-11T12:18:06Z

Sorry - I somewhat changed my mind on what the behaviour should look like on failure, that will require some code changes.

Could you also:

Add a news entry under openfe/news.

Add yourself to the CITATION.cff author list?

Thanks!

Added in commits 4ce6e8a and e5d52da

IAlibay

A few more things and then I think we're there!

IAlibay · 2026-06-12T08:28:14Z


-                fractions.append(chunk / N_l[0])
+                fractions.append(fraction)
        except ParameterError:


I think this outer try/except is now dead. Can you remove it (alongisde the associated ParameterError import)?

Fixed in commit 8cd3275

IAlibay · 2026-06-15T12:45:45Z

+                "of uncorrelated samples (fraction 1.0); discarding the forward "
+                "and reverse convergence analysis."
+            )
+            warnings.warn(wmsg)


Can you set the stacklevel parameter to 2 here? We're not great at being consistent about it elsewhere, but it's a small QoL thing that will be useful for users.

Fixed in commit 9232453

github-actions · 2026-06-16T12:53:40Z

No API break detected ✅

store partial convergence analysis if mbar fails for individual fraction

59ae19f

IAlibay requested changes Jun 4, 2026

View reviewed changes

aqemia-benedict-tan added 2 commits June 4, 2026 13:37

Refactor forwards/reverse estimates with helper function

4aa683f

Add tests for MBAR failure returning individual NaN during convergenc…

2e00d1d

…e analysis

aqemia-benedict-tan added 2 commits June 5, 2026 10:20

mypy can't follow openff.units/pint's Quantity typing

fbbe6f7

ruff needs double spacing

0cced56

jameseastwood assigned IAlibay Jun 9, 2026

IAlibay requested changes Jun 10, 2026

View reviewed changes

aqemia-benedict-tan added 6 commits June 11, 2026 10:48

Update to raise NaN for both directions even if only one fails

7b7dba3

update tests for returning NaN in both directions for a given fractio…

42ee838

…n where MBAR fails

add units for helper function

dbf9b71

return None if final MBAR estimate for fraction=1.0 errors

ce66327

add Benedict to authors

4ce6e8a

updated forwards and reverse convergence analysis

e5d52da

Merge branch 'main' into main

d3b12f7

IAlibay mentioned this pull request Jun 12, 2026

Update protocol docs to document new forward/reverse behaviour #1998

Open

IAlibay requested changes Jun 15, 2026

View reviewed changes

aqemia-benedict-tan added 2 commits June 16, 2026 12:47

remove dead try/except from get_forward_and_reverse_analysis

8cd3275

set stacklevel=2 in warning

9232453

Conversation

aqemia-benedict-tan commented Jun 4, 2026

Developers certificate of origin

Uh oh!

IAlibay left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Jun 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

aqemia-benedict-tan commented Jun 4, 2026

Uh oh!

IAlibay commented Jun 4, 2026

Uh oh!

IAlibay left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

aqemia-benedict-tan Jun 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

aqemia-benedict-tan commented Jun 11, 2026

Uh oh!

IAlibay left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Jun 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

codecov Bot commented Jun 4, 2026 •

edited

Loading

aqemia-benedict-tan Jun 11, 2026 •

edited

Loading