Speed up the MBAR calculation

[xiki-tempula]

I have adapted @mrshirts 's advice of using BAR results as initial guess into the MBAR, which seems to provide a 5-fold speed up.

from alchemtest.gmx import load_ABFE
from alchemlyb.parsing.gmx import extract_u_nk
import pandas as pd
df_list = [extract_u_nk(data, 310) for data in load_ABFE()['data']['complex']]
df = pd.concat(df_list)
mbar = MBAR(initial_f_k=None)
%timeit mbar.fit(df)
1.66 s ± 7.42 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
mbar = MBAR(initial_f_k='BAR')
%timeit mbar.fit(df)
248 ms ± 2.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

I tried to optimise the convergence detection as well, where using the result from the previous MBAR run as input into the next MBAR run. Compared to using BAR as input, this approach still provides some speed up.

from alchemlyb.convergence import forward_backward_convergence
%timeit forward_backward_convergence(df_list)

Using BAR as initial guess
3.58 s ± 35.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
Using the previous step as initial guess.
2.25 s ± 23.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 98.83%. Comparing base (55870c8) to head (8f94e37).
Report is 11 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master     #357   +/-   ##
=======================================
  Coverage   98.82%   98.83%           
=======================================
  Files          28       28           
  Lines        1875     1890   +15     
  Branches      405      409    +4     
=======================================
+ Hits         1853     1868   +15     
  Misses          2        2           
  Partials       20       20           

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

[orbeckst]

Neat idea and cool that it gives a real performance boost.

Definitely needs documentation!

[orbeckst]

Is this a change due to a merge or did you smuggle it into this PR??

[xiki-tempula]

So you could see that the CI failed for codecov token issues, I tried to see if I could fix it.

[orbeckst]

I didn't fully follow the code path but how is delta_f already initialized with the BAR results?

[xiki-tempula]

So you did a MBAR calculation first result = estimator_fit.fit(sample)
Which populates the result.delta_f_ with for the current MBAR results.
This MBAR results is then propagated back as initial guess for the next calculation.

[xiki-tempula]

Some codecov is still failing for token issue but otherwise this seems fine. When this PR is merged. Do you mind if I do a new 2.2.1 release? I think the speed increase is quite significant.

[orbeckst]

Let's make a 2.3.0 release (see my suggested changes — you can just accept the whole bunch) and document carefully API change (new default).

Please add a test for MBAR(initial_nk=None) for the old behavior, given that default changed.

[orbeckst]

Changing the default is technically speaking an API change so we can't just do it in a patch release.

The conservative approach would be to keep the original default.

However, I think we want to make this speed-up available from release so my suggestion is to make a 2.3 release and make clearer that this changed — see other comments.