diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 89d81580e..0ede903f3 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -14,7 +14,7 @@ defaults: jobs: test: - name: Test on ${{ matrix.os }} (py=${{ matrix.python-version }}, gmx=${{ matrix.gmx-version }}) + name: Test on ${{ matrix.os }} (py=${{ matrix.python-version }}, gmx=${{ matrix.gmx-version }}, pymbar=${{ matrix.pymbar-version }}) runs-on: ${{ matrix.os }} timeout-minutes: 120 strategy: @@ -35,6 +35,9 @@ jobs: - "2019" - "2024" - "2025" + pymbar-version: + - "3" + - "4" exclude: # a 2022 build doesn't seem to exist for mac @@ -81,7 +84,7 @@ jobs: with: environment-file: devtools/conda-envs/test_env.yaml create-args: >- # beware the >- instead of |, we don't split on newlines but on spaces - python=${{ matrix.python-version }} gromacs=${{ matrix.gmx-version }}.*=nompi_dblprec* + python=${{ matrix.python-version }} gromacs=${{ matrix.gmx-version }}.*=nompi_dblprec* pymbar=${{ matrix.pymbar-version }}.* - name: Set up conda environment without gromacs if: matrix.gmx-version == '5' || matrix.gmx-version == '4' @@ -90,7 +93,7 @@ jobs: # Use the legacy env for Python 3.9/3.10, full env otherwise. environment-file: ${{ (matrix.python-version == '3.9' || matrix.python-version == '3.10') && 'devtools/conda-envs/test_env_py39.yaml' || 'devtools/conda-envs/test_env.yaml' }} create-args: >- # beware the >- instead of |, we don't split on newlines but on spaces - python=${{ matrix.python-version }} + python=${{ matrix.python-version }} pymbar=${{ matrix.pymbar-version }}.* - name: Set up conda environment with single-precision gromacs (rejected versions) if: matrix.gmx-version == '2022' || matrix.gmx-version == '2023' @@ -101,13 +104,13 @@ jobs: # Use the legacy env for Python 3.9/3.10, full env otherwise. environment-file: ${{ (matrix.python-version == '3.9' || matrix.python-version == '3.10') && 'devtools/conda-envs/test_env_py39.yaml' || 'devtools/conda-envs/test_env.yaml' }} create-args: >- - python=${{ matrix.python-version }} gromacs=${{ matrix.gmx-version }} + python=${{ matrix.python-version }} gromacs=${{ matrix.gmx-version }} pymbar=${{ matrix.pymbar-version }}.* - name: Pin setuptools for Python <= 3.10 if: matrix.python-version == '3.9' || matrix.python-version == '3.10' shell: bash -l {0} run: pip install "setuptools<82" - + - name: Additional info about the build shell: bash run: | diff --git a/devtools/conda-envs/test_env.yaml b/devtools/conda-envs/test_env.yaml index aba44d2e9..0e4a1d811 100644 --- a/devtools/conda-envs/test_env.yaml +++ b/devtools/conda-envs/test_env.yaml @@ -18,7 +18,7 @@ dependencies: - zlib - swig - future - - pymbar =3 + - pymbar - openmm >= 8 - ambertools - ndcctools diff --git a/devtools/conda-envs/test_env_py39.yaml b/devtools/conda-envs/test_env_py39.yaml index 3664f71af..9127956a5 100644 --- a/devtools/conda-envs/test_env_py39.yaml +++ b/devtools/conda-envs/test_env_py39.yaml @@ -18,7 +18,7 @@ dependencies: - zlib - swig - future - - pymbar =3 + - pymbar - openmm >= 8 # ambertools has no Python 3.9 builds on conda-forge - ndcctools diff --git a/src/lipid.py b/src/lipid.py index 015e0000a..3507b2b6f 100644 --- a/src/lipid.py +++ b/src/lipid.py @@ -25,7 +25,7 @@ try: from lxml import etree except: pass -from pymbar import pymbar +from forcebalance.liquid import _mbar_weights import itertools from collections import defaultdict, namedtuple, OrderedDict import csv @@ -663,8 +663,7 @@ def get(self, mvals, AGrad=True, AHess=True): W1 = None if len(BPoints) > 1: logger.info("Running MBAR analysis on %i states...\n" % len(BPoints)) - mbar = pymbar.MBAR(U_kln, N_k, verbose=mbar_verbose, relative_tolerance=5.0e-8) - W1 = mbar.getWeights() + W1 = _mbar_weights(U_kln, N_k, verbose=mbar_verbose) logger.info("Done\n") elif len(BPoints) == 1: W1 = np.ones((BPoints*Shots,BPoints)) diff --git a/src/liquid.py b/src/liquid.py index 2467032ed..37d5ed0b5 100644 --- a/src/liquid.py +++ b/src/liquid.py @@ -24,7 +24,18 @@ try: from lxml import etree except: pass -from pymbar import pymbar +try: + from pymbar import pymbar # pymbar 3: MBAR lives in pymbar.pymbar submodule + _MBAR_SOLVER_KW = {} # v3: self-consistent-iteration default is fine +except ImportError: + import pymbar # pymbar 4: MBAR lives at pymbar top-level + _MBAR_SOLVER_KW = {'solver_protocol': 'robust'} # v4: default hybr diverges on some data + +def _mbar_weights(U_kln, N_k, verbose=False): + """Run MBAR and return weight matrix W[n, m] (shape (sum(N_k), K)).""" + mbar = pymbar.MBAR(U_kln, N_k, verbose=verbose, relative_tolerance=5.0e-8, **_MBAR_SOLVER_KW) + return mbar.weights() if hasattr(mbar, 'weights') else mbar.getWeights() + import itertools from forcebalance.optimizer import Counter from collections import defaultdict, namedtuple, OrderedDict @@ -894,8 +905,7 @@ def get_normal(self, mvals, AGrad=True, AHess=True): W1 = None if len(BPoints) > 1: logger.info("Running MBAR analysis on %i states...\n" % len(BPoints)) - mbar = pymbar.MBAR(U_kln, N_k, verbose=mbar_verbose, relative_tolerance=5.0e-8) - W1 = mbar.getWeights() + W1 = _mbar_weights(U_kln, N_k, verbose=mbar_verbose) logger.info("Done\n") elif len(BPoints) == 1: W1 = np.ones((Shots,1)) @@ -935,8 +945,7 @@ def fill_weights(weights, phase_points, mbar_points, snapshots): mU_kln[k, m, :] = mE[mE_idx] mU_kln[k, m, :] *= beta if np.abs(np.std(mE)) > 1e-6 and mBSims > 1: - mmbar = pymbar.MBAR(mU_kln, mN_k, verbose=False, relative_tolerance=5.0e-8, method='self-consistent-iteration') - mW1 = mmbar.getWeights() + mW1 = _mbar_weights(mU_kln, mN_k) elif len(mBPoints) == 1: mW1 = np.ones((mShots,1)) mW1 /= mShots diff --git a/src/tests/files/test_liquid/N_k.npy b/src/tests/files/test_liquid/N_k.npy new file mode 100644 index 000000000..4d8cb99fa Binary files /dev/null and b/src/tests/files/test_liquid/N_k.npy differ diff --git a/src/tests/files/test_liquid/U_kln.npy b/src/tests/files/test_liquid/U_kln.npy new file mode 100644 index 000000000..ba81a15ae Binary files /dev/null and b/src/tests/files/test_liquid/U_kln.npy differ diff --git a/src/tests/files/test_liquid/mbar_weights_ref.npy b/src/tests/files/test_liquid/mbar_weights_ref.npy new file mode 100644 index 000000000..10785c324 Binary files /dev/null and b/src/tests/files/test_liquid/mbar_weights_ref.npy differ diff --git a/src/tests/test_liquid.py b/src/tests/test_liquid.py index 86308b0d9..974a92c16 100644 --- a/src/tests/test_liquid.py +++ b/src/tests/test_liquid.py @@ -2,13 +2,63 @@ import os import sys import shutil +import numpy as np import pytest from forcebalance.parser import parse_inputs from forcebalance.forcefield import FF from forcebalance.objective import Objective from forcebalance.optimizer import Optimizer +from forcebalance.liquid import _mbar_weights from .__init__ import ForceBalanceTestCase, check_for_openmm +FIXTURE_DIR = os.path.join(os.path.dirname(__file__), 'files', 'test_liquid') + +# U_kln.npy — shape (6, 6, 801), float64 +# Reduced potential energy matrix U_kln[k, m, n] = (E_k[n] + P_m * V_k[n] * pvkj) * beta_m +# where k = source simulation index, m = evaluation state index, n = snapshot index. +# Built from the six npt_result.p pickles in files/test_liquid/single.tmp/Liquid/iter_0000/ +# (water at 249.15 K/1 atm, 273.15 K/1 atm, 298.15 K/1 atm, 373.15 K/1 atm, +# 298.15 K/20 bar, 298.15 K/2000 bar; 801 snapshots each). +# Regenerate with: conda run -n python tools/mbar_mre.py --save-ref +U_KLN_PATH = os.path.join(FIXTURE_DIR, 'U_kln.npy') + +# N_k.npy — shape (6,), int64, all entries = 801 +# Number of uncorrelated snapshots per simulation state; matches the first axis of U_kln. +N_K_PATH = os.path.join(FIXTURE_DIR, 'N_k.npy') + +# mbar_weights_ref.npy — shape (4806, 6), float64 +# MBAR weight matrix W[n, m] produced by pymbar 3.0.5 on U_kln above. +# Used as the reference for cross-version agreement checks (atol=1e-4). +# Regenerate with: conda run -n python tools/mbar_mre.py --save-ref +REF_PATH = os.path.join(FIXTURE_DIR, 'mbar_weights_ref.npy') + + +@pytest.fixture(scope='module') +def mbar_weights(): + if not os.path.exists(U_KLN_PATH) or not os.path.exists(N_K_PATH): + pytest.skip(f"MBAR fixtures not found in {FIXTURE_DIR}; run tools/mbar_mre.py --save-ref") + return _mbar_weights(np.load(U_KLN_PATH), np.load(N_K_PATH)) + + +def test_mbar_weights_normalized(mbar_weights): + """MBAR weight matrix columns must each sum to 1 (pymbar v3 and v4).""" + col_sums = mbar_weights.sum(axis=0) + assert np.allclose(col_sums, 1.0, atol=1e-6), ( + f"MBAR weight columns do not sum to 1: {col_sums}" + ) + + +def test_mbar_weights_match_reference(mbar_weights): + """MBAR weights must agree with the pymbar-v3 reference within 1e-4.""" + if not os.path.exists(REF_PATH): + pytest.skip(f"Reference weights not found at {REF_PATH}; run tools/mbar_mre.py --save-ref") + W_ref = np.load(REF_PATH) + assert mbar_weights.shape == W_ref.shape, ( + f"Weight matrix shape mismatch: got {mbar_weights.shape}, expected {W_ref.shape}" + ) + np.testing.assert_allclose(mbar_weights, W_ref, atol=1e-7, rtol=1e-2, + err_msg="MBAR weights differ from v3 reference beyond rtol=1e-2") + class TestWaterTutorial(ForceBalanceTestCase): def setup_method(self, method): if not check_for_openmm(): pytest.skip("No OpenMM modules found.")