[Feature Request] BAR style calculation

[akalpokas]

Given the high impact of updateParametersInContext() code on the performance of FEP calculations with SOMD2 (benchmark openmm/openmm#4610) it would be a very useful feature to allow the user to calculate energy differences to a restricted number of λ-windows, for example only neighbouring ones, thus enabling a "BAR style" calculation. Even with currently proposed OpenMM optimizations, going from a 3 λ-window to a 99 λ-window calculation using RF and 1ps of sampling frequency results with a performance drop from ~900ns/day to ~220ns/day. For PME the performance drop is from ~287ns/day to ~125ns/day.

With a BAR style of a calculation, we could do a calculation involving 99 λ-windows at the computational cost of only 3 (MBAR is meant to be more accurate, although I do believe that in vast majority of cases this is negligible, while the performance gain here is substantial). This would help in use cases where higher number of λ-windows are consistently needed, for example ABFE or ring-breaking FEP. The way GROMACS implements this is by using a calc-lambda-neighbors parameter which controls how many neighbouring λ-windows the energy sampling calculation will be carried out with, with 1 meaning the only the adjacent windows, and -1 meaning every window. I don't think this is too hard to implement in SOMD2, although it will probably involve additional modifications to BioSimSpace.FreeEnergy.Relative.analyse code to support BAR type energy output files (alchemlyb has a BAR estimator, although it cannot produce overlap matrices, like the MBAR one).

[jmichel80]

pymbar has a BAR estimator and can do overlap analyses https://pymbar.readthedocs.io/en/master/

could there be precision benefits in computing perturbed energies for 5-9 neighbouring windows ? The MBAR theory allows variable number of samples per thermodynamic state so it may be possible to padd square matrices and fake a full MBAR analysis.

[lohedges]

This should now be solved on the development branch of Sire and the feature_repex branch of somd2. It's now possible to specify the --num-energy-neighbours option to somd2, which specifies the number of neighbouring lambda windows that are included in the energy trajectory calculation. (Windows that aren't computed take a large energy value.) Initial testing shows that this works well and can greatly improve performance when a large number of windows are used.

I'll close this issue following further testing once the feature_repex PR is ready.