Repository: shirtsgroup/physical_validation

Issue:
Compatibility with pymbar 4
New release of pymbar breaks compatibility. We should
* Fix the erroneous calls
* Be more careful in defining our dependencies to explicitly include supported versions.

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/shirtsgroup/physical_validation/pull/220?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=shirtsgroup) Report
> Merging [#220](https://codecov.io/gh/shirtsgroup/physical_validation/pull/220?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=shirtsgroup) (9bf9c82) into [master](https://codecov.io/gh/shirtsgroup/physical_validation/commit/9af9f760cff8c4fc53eeb5170a43cbdf49eda34e?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=shirtsgroup) (9af9f76) will **increase** coverage by `0.05%`.
> The diff coverage is `100.00%`.


Patch:
diff --git a/.github/workflows/continous_integration.yaml b/.github/workflows/continous_integration.yaml
index 67eeb8b..51a6df9 100644
--- a/.github/workflows/continous_integration.yaml
+++ b/.github/workflows/continous_integration.yaml
@@ -16,6 +16,7 @@ jobs:
       matrix:
         os: ["macOS-latest", "ubuntu-latest"]
         python-version: ["3.7", "3.8", "3.9", "3.10"]
+        pymbar-version: ["pymbar3", "pymbar4"]
 
     steps:
       - uses: actions/checkout@v3
@@ -24,9 +25,9 @@ jobs:
         uses: conda-incubator/setup-miniconda@v2.1.1
         with:
           python-version: ${{ matrix.python-version }}
-          activate-environment: test
+          activate-environment: test-${{ matrix.pymbar-version }}
           channel-priority: true
-          environment-file: devtools/conda-envs/test_env.yaml
+          environment-file: devtools/conda-envs/test-${{ matrix.pymbar-version }}_env.yaml
           auto-activate-base: false
 
       - name: Additional info about the build
diff --git a/physical_validation/util/ensemble.py b/physical_validation/util/ensemble.py
index 570e6d8..d67877c 100644
--- a/physical_validation/util/ensemble.py
+++ b/physical_validation/util/ensemble.py
@@ -20,12 +20,46 @@ from typing import Dict, List, Optional, Tuple, Union
 
 import numpy as np
 import pymbar
+
+try:
+    # pymbar >= 4
+    from pymbar.timeseries import statistical_inefficiency
+except ImportError:
+    # pymbar < 4
+    from pymbar.timeseries import statisticalInefficiency as statistical_inefficiency
+
 import scipy.optimize
 
 from . import error as pv_error
 from . import plot, trajectory
 
 
+def pymbar_bar(
+    work_forward: np.ndarray,
+    work_backward: np.ndarray,
+) -> Dict[str, float]:
+    r"""
+    Wrapper around the pymbar BAR functionality, allowing the remaining code to be
+    agnostic of the pymbar version used.
+
+    Parameters
+    ----------
+    work_forward: Array of forward work values
+    work_backward: Array of backward work values
+
+    Returns
+    -------
+    A dictionary containing the free energy estimate and its error estimate
+
+    """
+    try:
+        # pymbar >= 4
+        return pymbar.other_estimators.bar(work_forward, work_backward)
+    except AttributeError:
+        # pymbar < 4
+        return pymbar.BAR(work_forward, work_backward, return_dict=True)
+
+
 def chemical_potential_energy(
     chemical_potential: np.ndarray,
     number_of_species: np.ndarray,
@@ -979,15 +1013,17 @@ def check_1d(
         )
 
     # calculate inefficiency
-    g1 = pymbar.timeseries.statisticalInefficiency(traj1)
-    g2 = pymbar.timeseries.statisticalInefficiency(traj2)
+    g1 = statistical_inefficiency(traj1)
+    g2 = statistical_inefficiency(traj2)
 
     w_f = -trueslope * traj1
     w_r = trueslope * traj2
 
     if verbosity > 2:
         print("Computing log of partition functions using pymbar.BAR...")
-    df, ddf = pymbar.BAR(w_f, w_r)
+    bar_results = pymbar_bar(w_f, w_r)
+    df = bar_results["Delta_f"]
+    ddf = bar_results["dDelta_f"]
     if verbosity > 2:
         print(
             "Using {:.5f} for log of partition functions as computed from BAR.".format(
@@ -1104,8 +1140,8 @@ def check_1d(
         # use overlap region
         t1, t2, min_ene, max_ene = trajectory.overlap(traj1=t1, traj2=t2)
         # calculate inefficiency
-        g1 = pymbar.timeseries.statisticalInefficiency(t1)
-        g2 = pymbar.timeseries.statisticalInefficiency(t2)
+        g1 = statistical_inefficiency(t1)
+        g2 = statistical_inefficiency(t2)
         # calculate max_likelihood fit
         fv, _ = do_max_likelihood_fit(
             t1,
@@ -1388,14 +1424,14 @@ def check_2d(
     # calculate inefficiency
     g1 = np.array(
         [
-            pymbar.timeseries.statisticalInefficiency(traj1[0]),
-            pymbar.timeseries.statisticalInefficiency(traj1[1]),
+            statistical_inefficiency(traj1[0]),
+            statistical_inefficiency(traj1[1]),
         ]
     )
     g2 = np.array(
         [
-            pymbar.timeseries.statisticalInefficiency(traj2[0]),
-            pymbar.timeseries.statisticalInefficiency(traj2[1]),
+            statistical_inefficiency(traj2[0]),
+            statistical_inefficiency(traj2[1]),
         ]
     )
 
@@ -1404,7 +1440,9 @@ def check_2d(
 
     if verbosity > 2:
         print("Computing log of partition functions using pymbar.BAR...")
-    df, ddf = pymbar.BAR(w_f, w_r)
+    bar_results = pymbar_bar(w_f, w_r)
+    df = bar_results["Delta_f"]
+    ddf = bar_results["dDelta_f"]
     if verbosity > 2:
         print(
             "Using {:.5f} for log of partition functions as computed from BAR.".format(
@@ -1474,14 +1512,14 @@ def check_2d(
         # calculate inefficiency
         g1 = np.array(
             [
-                pymbar.timeseries.statisticalInefficiency(t1[0]),
-                pymbar.timeseries.statisticalInefficiency(t1[1]),
+                statistical_inefficiency(t1[0]),
+                statistical_inefficiency(t1[1]),
             ]
         )
         g2 = np.array(
             [
-                pymbar.timeseries.statisticalInefficiency(t2[0]),
-                pymbar.timeseries.statisticalInefficiency(t2[1]),
+                statistical_inefficiency(t2[0]),
+                statistical_inefficiency(t2[1]),
             ]
         )
         # calculate max_likelihood fit
diff --git a/physical_validation/util/trajectory.py b/physical_validation/util/trajectory.py
index 8dbd1d4..c22117d 100644
--- a/physical_validation/util/trajectory.py
+++ b/physical_validation/util/trajectory.py
@@ -15,7 +15,20 @@ import warnings
 from typing import Iterator, Optional, Tuple
 
 import numpy as np
-from pymbar import timeseries
+
+try:
+    # pymbar >= 4
+    from pymbar.timeseries import (
+        detect_equilibration,
+        statistical_inefficiency,
+        subsample_correlated_data,
+    )
+except ImportError:
+    # pymbar < 4
+    from pymbar.timeseries import detectEquilibration as detect_equilibration
+    from pymbar.timeseries import statisticalInefficiency as statistical_inefficiency
+    from pymbar.timeseries import subsampleCorrelatedData as subsample_correlated_data
+
 from scipy import stats
 
 from . import error as pv_error
@@ -24,22 +37,22 @@ from . import error as pv_error
 def equilibrate(traj: np.ndarray) -> np.ndarray:
     traj = np.array(traj)
     if traj.ndim == 1:
-        t0, g, n_eff = timeseries.detectEquilibration(traj)
+        t0, g, n_eff = detect_equilibration(traj)
         if t0 == 0 and traj.size > 10:
             # See https://github.com/choderalab/pymbar/issues/277
-            t0x, gx, n_effx = timeseries.detectEquilibration(traj[10:])
+            t0x, gx, n_effx = detect_equilibration(traj[10:])
             if t0x != 0:
                 t0 = t0x + 10
         res = traj[t0:]
 
     elif traj.ndim == 2 and traj.shape[0] == 2:
-        t01, g1, n_eff1 = timeseries.detectEquilibration(traj[0])
-        t02, g2, n_eff2 = timeseries.detectEquilibration(traj[1])
+        t01, g1, n_eff1 = detect_equilibration(traj[0])
+        t02, g2, n_eff2 = detect_equilibration(traj[1])
         t0 = max(t01, t02)
         if t0 == 0 and traj.shape[1] > 10:
             # See https://github.com/choderalab/pymbar/issues/277
-            t01x, g1x, n_eff1x = timeseries.detectEquilibration(traj[0, 10:])
-            t02x, g2x, n_eff2x = timeseries.detectEquilibration(traj[1, 10:])
+            t01x, g1x, n_eff1x = detect_equilibration(traj[0, 10:])
+            t02x, g2x, n_eff2x = detect_equilibration(traj[1, 10:])
             t0x = max(t01x, t02x)
             if t0x != 0:
                 t0 = t0x + 10
@@ -61,15 +74,15 @@ def equilibrate(traj: np.ndarray) -> np.ndarray:
 def decorrelate(traj: np.ndarray) -> np.ndarray:
     traj = np.array(traj)
     if traj.ndim == 1:
-        idx = timeseries.subsampleCorrelatedData(traj)
+        idx = subsample_correlated_data(traj)
         res = traj[idx]
     elif traj.ndim == 2:
         # pymbar doesn't offer to decorrelate two samples, so let's do it ourselves
         # and just use the decorrelation of the sample more strongly correlated
         #
         # calculate (maximal) inefficiency
-        g1 = timeseries.statisticalInefficiency(traj[0])
-        g2 = timeseries.statisticalInefficiency(traj[1])
+        g1 = statistical_inefficiency(traj[0])
+        g2 = statistical_inefficiency(traj[1])
         g = np.max([g1, g2])
         # calculate index
         n0 = traj.shape[1]


Test patch:
diff --git a/devtools/conda-envs/test-pymbar3_env.yaml b/devtools/conda-envs/test-pymbar3_env.yaml
new file mode 100644
index 0000000..10647cf
--- /dev/null
+++ b/devtools/conda-envs/test-pymbar3_env.yaml
@@ -0,0 +1,21 @@
+name: test-pymbar3
+channels:
+  - conda-forge
+  - bioconda
+dependencies:
+    # Base depends
+  - python
+  - pip
+
+    # Testing and development
+  - pytest
+  - pytest-cov
+  - pytest-regressions
+  - codecov
+  - gromacs
+
+    # Standard dependencies
+  - numpy
+  - scipy
+  - matplotlib
+  - pymbar >=3,<4  # See issue #216
diff --git a/devtools/conda-envs/test_env.yaml b/devtools/conda-envs/test-pymbar4_env.yaml
similarity index 85%
rename from devtools/conda-envs/test_env.yaml
rename to devtools/conda-envs/test-pymbar4_env.yaml
index 78f4dfe..beda471 100644
--- a/devtools/conda-envs/test_env.yaml
+++ b/devtools/conda-envs/test-pymbar4_env.yaml
@@ -1,4 +1,4 @@
-name: test
+name: test-pymbar4
 channels:
   - conda-forge
   - bioconda
@@ -18,4 +18,4 @@ dependencies:
   - numpy
   - scipy
   - matplotlib
-  - pymbar <4  # See issue #216
+  - pymbar >=4