Adding support for bootstrapped errors to all expectations.

Computes bootstrap uncertainties and covariances for:

• compute_free_energy_differences
• compute_expectations
• compute_multiple_expectations
• compute_perturbed_free_energies
• compute_entropy_and_enthalpy

These are accessed by using the uncertainty_method = "bootstrap". Note that MBAR needs to be initialized with "n_bootstrap=(some integer)" samples in order to do use this method - an error is omitted otherwise. Possibly could be done differently, but the logic is simpler this way for now.

I could use a little help in rewriting the pytests. Currently, to get bootstrap uncertainties from various functions, as one will now want to check both the default uncertainties and the bootstrap uncertainties, and MBAR needs to be initialized differently each way, which the pytests don't seem to handle.

Note this is built off of #444, so many of the differences are from that branch.

Just putting it up here to show what is going on. Works for free energies so far, still need to get it working for other sections. No merging yet until that is working!

Here's the proposed generalization of the expectations. Not quite as clean as I would have hoped, and took a while to get it to this state, but does eliminate a lot of the non-vectorized loops.

It appears to be slightly slower for larger N (500K to 2M states), and I have not yet tested the performance on large numbers of states (sampled or unsampled) between the codes. It may require a bit of tweaking.

I didn't try to run the PMF code through it as well; that would have meant a number of other conditionals that seemed to complicated the code too much.

It may make sense to split out the computePerturbedFreeEnergies functions; that will simplify the inner loop a bit more.

includes a slightly modified harmonic_oscillator_update.py, which gives identical behavior with the new code to pymbar-examples/harmonic_oscillator.py run with the previous code up to numerical precision.

This new code automatically converts from old-style u_kln matrices internally.

It supports all existing functionality, including updating the weave code and the _pymbar.c code.

The one major difference is in handling computeExpectations. This change hopefully provides a much more general solution which can possibly reduce the total amount of repetitive code as well.

Previously, computeExpectations recognized whether observables were state_dependent or state_independent by whether the observable array was KxN_max, or KxLxN_max. Now, this is done by giving an explicit state_dependent flag. If state_dependent is False, computeExpectations handles the code like the 2D version before (although now the array is 1D, the observable at each state). The big difference is if state_dependent is True. In this case, computeExpectations expects K different observables (the K different observables at each state - for example, the potential energies in each of the K states), as well as K different energies (which defaults to just the u_kn matrix MBAR was generated with). It then calculates the expectations of the first observable and the first state, 2nd observable and the 2nd state, and so forth, which is what the previous code did.

In order to make this happen, I've created a computeGeneralExpectations routine, that takes in I observable arrays and K state observables (see see #89 ).

I'm up for lots of other improvements; but I though it would be easier to start by solving the biggest problem first independently of other issues.

For data sets where we have different numbers of samples at different states, rather than having the energies stored in a KxKxN_each matrix, it would be more efficient to store data as a KxN_tot matrix, where N is the total number of sampled collected, and K is the number of states we are evaluating the energies at. This will take a bit of extra work, however,

A branch that supports JAX, including minimize using BFGS and adaptive. It does not pass all the tests yet (because some pathways are mixed JAX/numpy, and those don't get handled well). But it will pass a lot of of the tests locally. I have not tested in GPU's yet. But it at least can be configured to run with JAX now on the conda environment. solver_protocol="jax" will exercise all of the JAX accelerated options.

in 1D, its clear that BAR is between neighboring states. In 2D is is not so clear which path to choose. Is there a way to guess a good path? Or do we just give up and let the user specify one, or just do i->i+1 as given.

Addresses #365

I added a temporary _deprecate function in _deprecate.py to ease the transition and avoid duplication. I have tested it with functions and should also work with methods...

I took master as a base, but most of the changes here won't reach Pymbar 4 (only the name replacements), and we have this excellent table here now:

BAR_zero                                             BARzero
BAR_zero                                             computeBARzero
BAR                                                  computeBAR
order_replicates                                     OrderReplicates
anderson_darling                                     AndersonDarling
QQ_plot                                              QQPlot
generate_confidence_intervals                        generateConfidenceIntervals
EXP_gauss                                            EXPGauss
EXP                                                  computeEXP
EXPGauss                                             computeEXPGauss
weights                                              getWeights
compute_effective_sample_number                      computeEffectiveSampleNumber
compute_overlap                                      computeOverlap
compute_free_energy_differences                      getFreeEnergyDifferences
compute_expectations_inner                           computeExpectationsInner
compute_covariance_of_sums                           computeCovarianceOfSums
compute_expectations                                 computeExpectations
compute_multiple_expectations                        computeMultipleExpectations
compute_perturbed_free_energies                      computePerturbedFreeEnergies
compute_entropy_and_enthalpy                         computeEntropyAndEnthalpy
compute_pmf                                          computePMF
statistical_inefficiency                             statisticalInefficiency
statistical_inefficiency_multiple                    statisticalInefficiencyMultiple
integrated_autocorrelation_time                      integratedAutocorrelationTime
integrated_autocorrelation_time_multiple             integratedAutocorrelationTimeMultiple
normalized_fluctuation_correlation_function          normalizedFluctuationCorrelationFunction
normalized_fluctuation_correlation_function_multiple normalizedFluctuationCorrelationFunctionMultiple
subsample_correlated_data                            subsampleCorrelatedData
detect_equilibration                                 detectEquilibration
statistical_inefficiency_fft                         statisticalInefficiency_fft
detect_equilibration_binary_search                   detectEquilibration_binary_search

I have only modified callables for now. Attributes (if relevant), parameters and free variables are untouched for now.

Lots of stuff going on here, but it's all logically related so a "big" PR is mostly unavoidable.

1. Moved current version of MBAR to old_mbar.py as a reference implementation for testing against.
2. Cleared all MBAR solver code into mbar_solvers.py
3. Added fast NumExpr based logsumexp() function to use as foundation of fast-but-precise MBAR calculations. If NumExpr is unavailable, a pure-numpy implementation is used.
4. Implemented matrix multiplication-based faster-but-imprecise MBAR calculations.
5. Replaced hand-written solvers with a single interface to scipy.optimize.minimize and scipy.optimize.root, which are scipy's unified interfaces to minimization and root-finding algorithms, respectively.
6. Because balancing precision and performance is pretty challenging, we need to make multiple passes through the solver--first get a rough guess with BFGS, then refine it with the MINPACK implementation of NR / HYBR. mbar_solvers.solve_mbar() uses a "protocol" to call BFGS and NR code several times in a way that achieves similar precision while matching the performance of the old MBAR code.
7. Added pymbar/testsystems/pymbar_datasets.py to automatically load the large datasets from compressed HDF5 files, which I will add to the pymbar-datasets repository soon.
8. Added tests/test_mbar_solvers.py to test the accuracy of the new MBAR code on several challenging systems, and also compare to the old MBAR code.
9. Added scripts/benchmark_pymbar.py to do performance and accuracy benchmarks against the old mbar code on several challenging datasets, including the gas and 8proteins datasets.

Overall, the speed is somewhat improved for the challenging gas and 8proteins datasets. More importantly, this restructuring helps us eliminate the weave dependence (needed for Py3 compliance) and the hand-written C code. I also think that mbar_solvers.py is an improvement in readability. If we need more speed in the future, we can plug in faster versions of the gradient / hessian calculation.

Following PEP8 recommendations and using table of replacements as provided in #368.

Note that examples/ is not affected in this PR. They will be "modernized" in a separate PR.

Questions

• [x] Shall I apply snake_case to public attributes and module-level variables?
• [x] Shall I apply snake_case to unexposed variables?

Notes

• This PR introduces breaking changes! Check replacements in #368.
  • Special attention: getFreeEnergyDifferences is now compute_free_energy_differences, for consistency with other methods and functions.
• pymbar.bar and pymbar.exp modules have been merged into pymbar.other_estimators.

Summary of changed names:

AndersonDarling                                    anderson_darling
BARzero                                            bar_zero
computeBAR                                         bar
computeBARzero                                     bar_zero
computeCovarianceOfSums                            compute_covariance_of_sums
computeEffectiveSampleNumber                       compute_effective_sample_number
computeEntropyAndEnthalpy                          compute_entropy_and_enthalpy
computeEXP                                         exp
computeExpectations                                compute_expectations
computeExpectationsInner                           compute_expectations_inner
computeEXPGauss                                    exp_gauss
computeMultipleExpectations                        compute_multiple_expectations
computeOverlap                                     compute_overlap
computePerturbedFreeEnergies                       compute_perturbed_free_energies
computePMF                                         compute_pmf
detectEquilibration                                detect_equilibration
detectEquilibration_binary_search                  detect_equilibration_binary_search
EXPGauss                                           exp_gauss
generateConfidenceIntervals                        generate_confidence_intervals
generatePMF                                        generate_pmf
GetAnalytical                                      get_analytical
getFreeEnergyDifferences                           compute_free_energy_differences
getInformationCriteria                             get_information_criteria
getKDE                                             get_kde
getMBAR                                            get_mbar
getPMF                                             get_pmf
getWeights                                         weights
integratedAutocorrelationTime                      integrated_autocorrelation_time
integratedAutocorrelationTimeMultiple              integrated_autocorrelation_time_multiple
normalizedFluctuationCorrelationFunction           normalized_fluctuation_correlation_function
normalizedFluctuationCorrelationFunctionMultiple   normalized_fluctuation_correlation_function_multiple
OrderReplicates                                    order_replicates
PrintResults                                       print_results
QQPlot                                             qq_plot
statisticalInefficiency                            statistical_inefficiency
statisticalInefficiency_fft                        statistical_inefficiency_fft
statisticalInefficiencyMultiple                    statistical_inefficiency_multiple
subsampleCorrelatedData                            subsample_correlated_data

One issue is that some of the previous "tricks" are ineffective (e.g. fast=True), as the FFT-based approach calculates all lagtimes simultaneously.

For the search the discard region t, I've elected to just do binary search of the logarithmically-spaced lagtime grid.

>>> import numpy
>>> numpy.__version__
'1.24.0'
>>> numpy.bool
<stdin>:1: FutureWarning: In the future `np.bool` will be defined as the corresponding NumPy scalar.  (This may have returned Python scalars in past versions.
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/mattthompson/mambaforge/envs/openff-interchange-env/lib/python3.9/site-packages/numpy/__init__.py", line 284, in __getattr__
    raise AttributeError("module {!r} has no attribute "
AttributeError: module 'numpy' has no attribute 'bool'

Aliases for built-ins were deprecated in February 2021: https://numpy.org/doc/stable/release/1.20.0-notes.html#using-the-aliases-of-builtin-types-like-np-int-is-deprecated

https://github.com/choderalab/pymbar/actions/runs/3736207271/jobs/6340264997

After reading the examples for calculating the PMF based on data by parallel-templering and umbrella-sampling simulations, i found that there is a discrepancy in calculating u_kln(reduced potential energy).

Following the line 222 to 226 in (https://github.com/mrshirts/pymbar/blob/master/examples/parallel-tempering-2dpmf/parallel-tempering-2dpmf.py), reduced potential energy( u_kln[k,l,n] ) of trajectory segment n of temperature k evaluated at temperature l was calculated by following:.

    print("Computing reduced potential energies...")
    u_kln = numpy.zeros([K,K,N_max]) # u_kln[k,l,n] is reduced potential energy of trajectory segment n of temperature k evaluated at temperature l
    for k in range(K):
       for l in range(K):
          u_kln[k,l,0:N_k[k]] = beta_k[l] * U_kn[k,0:N_k[k]]

which means that u_kln[k,l,n] = beta_k[l] * U_kn[k,n] = beta_k[l] * U_kn[k,l,n] ( without bias, U_kn[k,n] = U_kn[k,1,n] = U_kn[k,2,n] = ... = U_kn[k,l,n] = .... ). [equation 1]

While in the examples of analyzing the data by umbrella sampling, as shown in the line 141 to 143 of https://github.com/mrshirts/pymbar/blob/master/examples/umbrella-sampling-pmf/umbrella-sampling.py, it is calculated by:

    # Compute energy of snapshot n from simulation k in umbrella potential l
    u_kln[k,:,n] = u_kn[k,n] + beta_k[k] * (K_k/2.0) * dchi**2

where

    u_kn[k,n] = beta_k[k] * (float(tokens[2]) - float(tokens[1])) # reduced potential energy without umbrella restraint(line 89)

which indicates: u_kln[k,l,n] = beta_k[k] * U0_kn[k,n] + beta_k[k] * (K_k[l] / 2.0) * dchi[l]**2 = beta_k[k] * ( U0_kn[k,n] + (K_k[l] / 2.0) * dchi[l]**2 ) = beta_k[k] * U_kn[k,l,n] [ equation 2]

It seems that the [equation 1] is the correct expression, while the [ equation 2 ] is not.

Generally, based on [ euqation 2 ], for umbrella sampling of constanst temperature(beta_k[k] = beta_k[l]), it is ok, while for different temperature, it will lead to incorrect result.

Here are my suggestions: 1)In line 89 of https://github.com/mrshirts/pymbar/blob/master/examples/umbrella-sampling-pmf/umbrella-sampling.py,

  u_kn[k,n] = beta_k[k] * (float(tokens[2]) - float(tokens[1])) # reduced potential energy without umbrella restraint

should be changed to :

  U_kn[k,n] = float(tokens[2]) - float(tokens[1])  # potential energy without umbrella restraint

2)In line 143,

  u_kln[k,:,n] = u_kn[k,n] + beta_k[k] * (K_k/2.0) * dchi**2 

should be changed to:

  u_kln[k,:,n] = beta_k * ( U_kn[k,n] + (K_k/2.0) * dchi**2 )

I found that in line 544-548 of pymbar/fes.py you wrote

                # how do we label the bins? if N-dimensional:
                # bins[0] + bins[1]*bin_length[1]**1 + bins[2]*bin_length[2]**2
                sample_label[n] = int(
                    np.sum([bin_n[n][d] * bin_length[d] ** d for d in range(dims)])
                )

However, this setting can not give unique bin values. For example, given bin_length=[100,2], bins=(30,1) and bins=(32,0) would have the same bin value of 32 (i.e. 30*100**0+1*2**1=32*100**0+0*2**1=32).

Non-unique bin values would cause under-counting of histogram_data["bin_order"] in line 566-576 of pymbar/fes.py

        if b == 0:
            bin_order = {}
            i = 0
            for bv in bin_label.values():
                if bv not in bin_order:
                    bin_order[bv] = i
                    i += 1
            histogram_data["bin_order"] = bin_order
            histogram_data["bin_label"] = bin_label
        else:
            bin_order = self.histogram_data["bin_order"]

I would suggest to correct line 546-548 of pymbar/fes.py to be

                sample_label[n] = int(
                    np.sum([bin_n[n][d] * np.sum(bin_length[:d+1]) ** d for d in range(dims)])
                )

I want to construct a 2D-PMF using fes.get_fes function. However, I found that the number of bins in each dimension must be the same to avoid an error.

The problem might rooted In line 467 and 468 of pymbar/fes.py

        if len(np.shape(histogram_parameters["bin_edges"])) == 1:
            histogram_parameters["bin_edges"] = [histogram_parameters["bin_edges"]]

The if-statement would pass if the number of bins in each dimension (each element in histogram_parameters["bin_edges"]) are not the same. Say, if histogram_parameters["bin_edges"]=[np.array(1,2,3,4),np.array(1,2)], the if-statement could also pass.