diff --git a/.github/workflows/test_openmm_dmff_plugin.yml b/.github/workflows/test_openmm_dmff_plugin.yml index 597d0acb2..c25b5bf0b 100644 --- a/.github/workflows/test_openmm_dmff_plugin.yml +++ b/.github/workflows/test_openmm_dmff_plugin.yml @@ -11,18 +11,22 @@ jobs: matrix: python-version: [3.9] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: Install Dependencies run: | source $CONDA/bin/activate - conda create -n dmff_omm -y python=${{ matrix.python-version }} numpy=1.* openmm=7.7 -c conda-forge + # swig<4.1 is required: OpenMM's shipped include/swig/OpenMMSwigHeaders.i + # does `%include "factory.i"` / `%factory(...)`, and SWIG 4.1 removed the + # factory.i library, so newer swig fails with "Unknown directive '%factory'". + # The runner image's system swig is now 4.2+, hence we pin one in the env. + conda create -n dmff_omm -y python=${{ matrix.python-version }} pip numpy=1.* openmm=7.7 "swig<4.1" -c conda-forge conda activate dmff_omm conda install -y libtensorflow_cc=2.9.1 -c conda-forge - pip install setuptools==59.5.0 + python -m pip install setuptools==59.5.0 mkdir /tmp/omm_dmff_working_dir cd /tmp/omm_dmff_working_dir wget https://github.com/tensorflow/tensorflow/archive/refs/tags/v2.9.1.tar.gz @@ -41,9 +45,13 @@ jobs: export OPENMM_INSTALLED_DIR=$CONDA_PREFIX export CPPFLOW_INSTALLED_DIR=$CONDA_PREFIX export LIBTENSORFLOW_INSTALLED_DIR=$CONDA_PREFIX - cmake .. -DOPENMM_DIR=${OPENMM_INSTALLED_DIR} -DCPPFLOW_DIR=${CPPFLOW_INSTALLED_DIR} -DTENSORFLOW_DIR=${LIBTENSORFLOW_INSTALLED_DIR} -DUSE_HIGH_PRECISION=OFF - make && make install - make PythonInstall + # CMakeLists.txt does FIND_PROGRAM(SWIG_EXECUTABLE swig), which would + # otherwise pick up the runner's system swig (4.2+) instead of the + # pinned one in the conda env. Setting it here wins over FIND_PROGRAM. + "${CONDA_PREFIX}/bin/swig" -version | head -3 + cmake .. -DOPENMM_DIR=${OPENMM_INSTALLED_DIR} -DCPPFLOW_DIR=${CPPFLOW_INSTALLED_DIR} -DTENSORFLOW_DIR=${LIBTENSORFLOW_INSTALLED_DIR} -DUSE_HIGH_PRECISION=OFF -DSWIG_EXECUTABLE=${CONDA_PREFIX}/bin/swig + make && make install + make PythonInstall - name: Run Tests run: | source $CONDA/bin/activate dmff_omm diff --git a/.github/workflows/ut.yml b/.github/workflows/ut.yml index 45b62bed8..27faa3143 100644 --- a/.github/workflows/ut.yml +++ b/.github/workflows/ut.yml @@ -4,39 +4,79 @@ on: push: pull_request: +# DMFF supports Python 3.9-3.12. The two axes are not independent: upstream JAX +# drops old Pythons over time, so each Python has a different newest usable JAX +# (3.9 -> 0.4.30, 3.10 -> 0.6.2, 3.11/3.12 -> current). Pinning "latest" on 3.9 +# would silently resolve an old JAX and stop testing what the job claims to. jobs: build: runs-on: ubuntu-latest strategy: + fail-fast: false matrix: - python-version: [3.9] + include: + # Floor of both axes: oldest supported Python, newest JAX it can run. + - name: py3.9-jax0.4.30 + python-version: "3.9" + jax-spec: "jax==0.4.30 jaxlib==0.4.30 'numpy<2'" + # Oldest JAX we pin against on a modern Python. + - name: py3.11-jax0.4.35 + python-version: "3.11" + jax-spec: "jax==0.4.35 jaxlib==0.4.35 'numpy<2'" + # Newest JAX each remaining Python can install, with numpy 2. + - name: py3.10-jax-latest + python-version: "3.10" + jax-spec: "jax jaxlib" + - name: py3.11-jax-latest + python-version: "3.11" + jax-spec: "jax jaxlib" + - name: py3.12-jax-latest + python-version: "3.12" + jax-spec: "jax jaxlib" + name: ${{ matrix.name }} steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 + with: + # setuptools_scm needs full history to derive the version. + fetch-depth: 0 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: Install Dependencies run: | source $CONDA/bin/activate - conda create -n dmff -y python=${{ matrix.python-version }} numpy openmm==7.7.0 pytest rdkit openbabel mdtraj ambertools -c conda-forge + conda create -n dmff -y python=${{ matrix.python-version }} pip numpy openmm pytest rdkit openbabel mdtraj ambertools -c conda-forge conda activate dmff - pip install --upgrade pip - pip install jax jaxlib jaxopt networkx parmed pymbar==4.0.1 optax tqdm ase + # Always drive pip through the interpreter we are about to test with. + # A bare `pip` can resolve to the setup-python installation ahead of + # the conda env on PATH, silently installing into a different + # interpreter than `python` imports from. + which python + python -m pip install --upgrade pip + # optax>=0.2.3 gives optax.lbfgs, which the QEq module needs. + python -m pip install ${{ matrix.jax-spec }} networkx parmed pymbar "optax>=0.2.3" tqdm ase - name: Install DMFF run: | - source $CONDA/bin/activate dmff && pip install . + source $CONDA/bin/activate dmff + python -m pip install . + - name: Report resolved versions + run: | + source $CONDA/bin/activate dmff + which python + python -c "import sys, jax, numpy; print('python', sys.version.split()[0], '| jax', jax.__version__, '| numpy', numpy.__version__)" \ + || { echo '--- jax missing; installed packages: ---'; python -m pip list; exit 1; } - name: Run Tests run: | source $CONDA/bin/activate dmff - pytest -vs tests/test_classical/test_* - pytest -vs tests/test_common/test_* - pytest -vs tests/test_admp/test_* - pytest -vs tests/test_utils.py - pytest -vs tests/test_mbar/test_* - pytest -vs tests/test_sgnn/test_* - pytest -vs tests/test_difftraj/test_* - pytest -vs tests/test_frontend/test_* - pytest -vs tests/test_dimer/test_* - pytest -vs tests/test_energy.py - pytest -vs tests/test_ase/test_* + python -m pytest -vs tests/test_classical/test_* + python -m pytest -vs tests/test_common/test_* + python -m pytest -vs tests/test_admp/test_* + python -m pytest -vs tests/test_utils.py + python -m pytest -vs tests/test_mbar/test_* + python -m pytest -vs tests/test_sgnn/test_* + python -m pytest -vs tests/test_difftraj/test_* + python -m pytest -vs tests/test_frontend/test_* + python -m pytest -vs tests/test_dimer/test_* + python -m pytest -vs tests/test_energy.py + python -m pytest -vs tests/test_ase/test_* diff --git a/backend/openmm_dmff_plugin/README.md b/backend/openmm_dmff_plugin/README.md index c059c5eb4..ab80b1435 100644 --- a/backend/openmm_dmff_plugin/README.md +++ b/backend/openmm_dmff_plugin/README.md @@ -11,9 +11,15 @@ Install the python, openmm and cudatoolkit. ```shell mkdir omm_dmff_working_dir && cd omm_dmff_working_dir -conda create -n dmff_omm -c conda-forge python=3.9 openmm cudatoolkit=11.6 +conda create -n dmff_omm -c conda-forge python=3.9 openmm cudatoolkit=11.6 "swig<4.1" conda activate dmff_omm ``` +The `swig<4.1` pin is required. OpenMM's bundled `include/swig/OpenMMSwigHeaders.i` +uses the `%factory` directive, and SWIG removed the `factory.i` library in 4.1.0, +so building the Python wrappers with a newer SWIG fails with +`Error: Unknown directive '%factory'`. Many Linux distributions now ship SWIG 4.2+, +so install a pinned SWIG into the environment and point CMake at it explicitly +with `-DSWIG_EXECUTABLE=${CONDA_PREFIX}/bin/swig` (see below). ### Download `libtensorflow_cc` and install `cppflow` package Install the precompiled libtensorflow_cc library from conda. ```shell @@ -50,7 +56,7 @@ Compile the plugin from the source with the following steps. 2. Run `cmake` command with the required parameters. ```shell - cmake .. -DOPENMM_DIR=${OPENMM_INSTALLED_DIR} -DCPPFLOW_DIR=${CPPFLOW_INSTALLED_DIR} -DTENSORFLOW_DIR=${LIBTENSORFLOW_INSTALLED_DIR} + cmake .. -DOPENMM_DIR=${OPENMM_INSTALLED_DIR} -DCPPFLOW_DIR=${CPPFLOW_INSTALLED_DIR} -DTENSORFLOW_DIR=${LIBTENSORFLOW_INSTALLED_DIR} -DSWIG_EXECUTABLE=${CONDA_PREFIX}/bin/swig make && make install make PythonInstall ``` diff --git a/dmff/admp/qeq.py b/dmff/admp/qeq.py index 4f1ebea78..04bb0de6a 100644 --- a/dmff/admp/qeq.py +++ b/dmff/admp/qeq.py @@ -16,21 +16,20 @@ CONST_1 = jnp.array(1, dtype=jnp.float32) try: - import jaxopt + import optax - try: - from jaxopt import Broyden - - JAXOPT_OLD = False - except ImportError: - JAXOPT_OLD = True + # optax.lbfgs was added in optax 0.2.3; it replaces the archived jaxopt. + OPTAX_LBFGS = hasattr(optax, "lbfgs") + if not OPTAX_LBFGS: import warnings warnings.warn( - "jaxopt is too old. The QEQ potential function cannot be jitted. Please update jaxopt to the latest version for speed concern." + "optax is too old (need >=0.2.3 for optax.lbfgs). QEQ cannot be used." ) except ImportError: + optax = None + OPTAX_LBFGS = False import warnings - warnings.warn("jaxopt not found, QEQ cannot be used.") + warnings.warn("optax not found, QEQ cannot be used.") import jax from jax.scipy.special import erf, erfc @@ -203,6 +202,64 @@ def etainv_piecewise(eta): etainv_piecewise = jax.vmap(etainv_piecewise, in_axes=0) +def lbfgs_minimize( + value_and_grad_fn, value_fn, init_params, tol=1e-2, maxiter=500, memory_size=10 +): + """Minimize with optax's L-BFGS, mirroring the jaxopt.LBFGS contract. + + Replaces ``jaxopt.LBFGS(fun=..., value_and_grad=True, tol=...)``: jaxopt is + archived, and optax.lbfgs is the maintained equivalent. + + ``value_and_grad_fn`` returns ``(value, grad)`` for the parameters being + optimized. The gradient it returns may be a *projected* gradient (as QEq + needs, to stay on the charge-constraint plane); it is used both as the + descent direction and as the convergence criterion, so the projection is + respected. ``value_fn`` returns the objective alone and is only consulted by + the line search, which is backtracking/Armijo -- it never differentiates + ``value_fn``, so it cannot reintroduce the unprojected gradient. + + Stops when the gradient L2 norm drops to ``tol`` or ``maxiter`` is reached. + Runs under ``jit`` via ``lax.while_loop``. + """ + if not OPTAX_LBFGS: + raise ImportError( + "QEQ requires optax>=0.2.3 for optax.lbfgs. Install it with " + "`pip install 'optax>=0.2.3'` or `pip install 'dmff[qeq]'`." + ) + + solver = optax.lbfgs( + memory_size=memory_size, + linesearch=optax.scale_by_backtracking_linesearch( + max_backtracking_steps=20, store_grad=True + ), + ) + + def cond_fn(carry): + _, _, _, _, gnorm, nit = carry + return jnp.logical_and(gnorm > tol, nit < maxiter) + + def body_fn(carry): + params, state, value, gradient, _, nit = carry + updates, state = solver.update( + gradient, state, params, value=value, grad=gradient, value_fn=value_fn + ) + params = optax.apply_updates(params, updates) + value, gradient = value_and_grad_fn(params) + return params, state, value, gradient, jnp.linalg.norm(gradient), nit + 1 + + value, gradient = value_and_grad_fn(init_params) + carry = ( + init_params, + solver.init(init_params), + value, + gradient, + jnp.linalg.norm(gradient), + jnp.asarray(0), + ) + params = jax.lax.while_loop(cond_fn, body_fn, carry)[0] + return params + + class ADMPQeqForce: def __init__( self, @@ -476,11 +533,6 @@ def get_step_energy(positions, box, pairs, mscales, eta, chi, J, aux=None): return energy func = get_proj_grad(E_full,self.const_mat) - solver = jaxopt.LBFGS( - fun=func, - value_and_grad=True, - tol=1e-2, - ) pos = positions buffer_scales = pair_buffer_scales(pairs) if self.has_aux: @@ -490,8 +542,7 @@ def get_step_energy(positions, box, pairs, mscales, eta, chi, J, aux=None): q = self.init_q lagmt = self.init_lagmt - res = solver.run( - q, + rest_args = ( lagmt, chi, J, @@ -502,7 +553,12 @@ def get_step_energy(positions, box, pairs, mscales, eta, chi, J, aux=None): buffer_scales, mscales, ) - q_opt = res.params + q_opt = lbfgs_minimize( + lambda x: func(x, *rest_args), + lambda x: E_full(x, *rest_args), + q, + tol=1e-2, + ) energy = E_full( q_opt, lagmt, diff --git a/dmff/api/paramset.py b/dmff/api/paramset.py index 141adb366..b0ab6177a 100644 --- a/dmff/api/paramset.py +++ b/dmff/api/paramset.py @@ -118,7 +118,7 @@ def __getitem__(self, key: str) -> Union[Dict[str, jnp.ndarray], jnp.ndarray]: return self.parameters[key] def update_mask(self, gradients): - gradients = jax.tree_map( + gradients = jax.tree_util.tree_map( lambda g, m: jnp.where(jnp.abs(m - 1.0) > 1e-5, g, 0.0), gradients, self.mask ) return gradients diff --git a/dmff/classical/vsite.py b/dmff/classical/vsite.py index c34410fe9..0e9039860 100644 --- a/dmff/classical/vsite.py +++ b/dmff/classical/vsite.py @@ -1,5 +1,6 @@ from typing import Tuple, List, Dict, Callable, Optional import numpy as np +import jax import jax.numpy as jnp @@ -61,7 +62,7 @@ def addVirtualSiteToMol(self, rdmol, vtypes=None, vdist=None): newmol: rdkit.Chem.Mol Mol with virtual sites added """ - if isinstance(vtypes, jnp.ndarray) and isinstance(vdist, jnp.ndarray): + if isinstance(vtypes, jax.Array) and isinstance(vdist, jax.Array): func = self.getAddVirtualSiteFunc() # convert between angstrom and nm pos = jnp.array(rdmol.GetConformer(0).GetPositions()) / 10 diff --git a/dmff/eann/eann.py b/dmff/eann/eann.py index a550e0d3d..25b966612 100644 --- a/dmff/eann/eann.py +++ b/dmff/eann/eann.py @@ -10,8 +10,6 @@ from functools import partial import jax.nn.initializers import pickle -# from jax.config import config -# config.update("jax_debug_nans", True) def get_elem_indices(topology): diff --git a/dmff/optimize.py b/dmff/optimize.py index f2bc67ffb..0f8d86090 100644 --- a/dmff/optimize.py +++ b/dmff/optimize.py @@ -4,7 +4,12 @@ from typing import Optional import optax -PeriodicParamsState = optax._src.base.EmptyState +PeriodicParamsState = optax.EmptyState + +NO_PARAMS_MSG = ( + "You are using a transformation that requires the current value of " + "parameters, but you are not passing `params` when calling `update`." +) def periodic_move(pmin, pmax): @@ -14,17 +19,17 @@ def init_fn(params): def update_fn(updates, state, params): if params is None: - raise ValueError(optax._src.base.NO_PARAMS_MSG) + raise ValueError(NO_PARAMS_MSG) - updates = jax.tree_map( + updates = jax.tree_util.tree_map( lambda p, u: jnp.where((p + u) < pmin, u + pmax - pmin, u), params, updates ) - updates = jax.tree_map( + updates = jax.tree_util.tree_map( lambda p, u: jnp.where((p + u) > pmax, u - pmax + pmin, u), params, updates ) return updates, state - return optax._src.base.GradientTransformation(init_fn, update_fn) + return optax.GradientTransformation(init_fn, update_fn) def genOptimizer( diff --git a/dmff/torch_tools.py b/dmff/torch_tools.py index f812a2d94..25f7322ff 100644 --- a/dmff/torch_tools.py +++ b/dmff/torch_tools.py @@ -14,7 +14,7 @@ def t2j_element(e): return e def j2t_element(e): - if isinstance(e, jax.numpy.ndarray): + if isinstance(e, jax.Array): e = j2t(e) e.requires_grad = True return e diff --git a/dmff/utils.py b/dmff/utils.py index 544379cf4..bc721466b 100644 --- a/dmff/utils.py +++ b/dmff/utils.py @@ -1,3 +1,4 @@ +import jax from jax import jit, vmap, tree_util import jax.numpy as jnp from .settings import DO_JIT @@ -86,7 +87,7 @@ def pair_buffer_scales(p): def isinstance_jnp(*args): def _check(arg): - if not isinstance(arg, jnp.ndarray): + if not isinstance(arg, jax.Array): raise TypeError( "all arguments must be jnp.array, \ otherwise they won't be able to take derivatives \ diff --git a/docs/user_guide/2.installation.md b/docs/user_guide/2.installation.md index d2e21cebf..03c4073e4 100644 --- a/docs/user_guide/2.installation.md +++ b/docs/user_guide/2.installation.md @@ -2,35 +2,51 @@ ## 2.1 Install Dependencies + Create conda environment: ``` -conda create -n dmff python=3.9 --yes +conda create -n dmff python=3.11 --yes conda activate dmff ``` + Install [jax](https://github.com/google/jax) (select the correct cuda version, see more details in the Jax installation guide): ```bash # CPU version -pip install "jax[cpu]==0.4.14" +pip install "jax[cpu]" # GPU version -pip install "jax[cuda11_local]==0.4.14" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html -``` -+ Install [mdtraj](https://github.com/mdtraj/mdtraj), [optax](https://github.com/deepmind/optax), [jaxopt](https://github.com/google/jaxopt) and [pymbar](https://github.com/choderalab/pymbar): -```bash -conda install -c conda-forge mdtraj==1.9.7 -pip install optax==0.1.3 pymbar==4.0.1 jaxopt==0.8.1 +pip install "jax[cuda12]" ``` +DMFF supports Python 3.9-3.12 and `jax>=0.4.1`, including the current JAX release. + +Note that JAX itself drops old Python versions over time, so the newest JAX you can +install depends on your Python. `pip install jax` will quietly resolve an older +release rather than fail: + +| Python | newest installable JAX | +| ------ | ---------------------- | +| 3.9 | 0.4.30 | +| 3.10 | 0.6.2 | +| 3.11 | current | +| 3.12 | current | + +Use Python 3.11+ if you want the current JAX. To reproduce results against an older +JAX, pin it explicitly: `pip install jax==0.4.35 jaxlib==0.4.35 "numpy<2"`. + + Install [mdtraj](https://github.com/mdtraj/mdtraj), [optax](https://github.com/deepmind/optax) and [pymbar](https://github.com/choderalab/pymbar): ```bash -conda install -c conda-forge mdtraj==1.9.7 -pip install optax==0.1.3 -pip install pymbar==4.0.1 +conda install -c conda-forge mdtraj +pip install optax pymbar ``` + Install [OpenMM](https://openmm.org/): ```bash -conda install -c conda-forge openmm==7.7.0 +conda install -c conda-forge openmm ``` + Install [RDKit](https://www.rdkit.org/) (required for SMIRKS-based parametrization): ```bash conda install -c conda-forge rdkit ``` ++ Optional: the QEq module needs a newer [optax](https://github.com/google-deepmind/optax) + than the core dependency floor (`optax>=0.2.3`, for `optax.lbfgs`). Install the extra + if you use QEq: +```bash +pip install ".[qeq]" # or: pip install "optax>=0.2.3" +``` ## 2.2 Install DMFF from Source Code One can download the DMFF source code from github and install it using `pip`. : ```bash diff --git a/examples/eann/eann.py b/examples/eann/eann.py index 7a5564dbe..06200706d 100644 --- a/examples/eann/eann.py +++ b/examples/eann/eann.py @@ -11,8 +11,6 @@ from functools import partial import jax.nn.initializers import pickle -# from jax.config import config -# config.update("jax_debug_nans", True) # Make printing parameters a little more readable def parameter_shapes(params): diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 000000000..e98d15526 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,7 @@ +[build-system] +# Declared here rather than via setup.py's `setup_requires`, which uses the +# legacy easy_install path: it fetches setuptools_scm as an egg without its +# transitive dependencies, so setuptools_scm>=10 (which moved Configuration +# into the separate vcs_versioning package) fails to import at build time. +requires = ["setuptools>=64", "setuptools_scm>=8", "wheel"] +build-backend = "setuptools.build_meta" diff --git a/requirements.txt b/requirements.txt index 1e2ad7b20..2d9fdfefd 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,16 @@ -numpy>=1.18, <2 +# Runtime +numpy>=1.23 +jax>=0.4.1 +jaxlib>=0.4.1 +pymbar>=4.0.0 +tqdm +rdkit +ase + +# The QEq module (dmff/admp/qeq.py) additionally needs optax>=0.2.3 for +# optax.lbfgs; install it with `pip install .[qeq]` if you use QEq. + +# Docs mkdocs>=1.3.0 mkdocs-autorefs>=0.4.1 mkdocs-gen-files>=0.3.4 @@ -6,9 +18,3 @@ mkdocs-literate-nav>=0.4.1 mkdocstrings>=0.19.0 mkdocstrings-python>=0.7.0 pygments>=2.12 -jax>=0.4.1 -jaxlib>=0.4.1 -pymbar>=4.0.0 -tqdm -rdkit==2022.03.2 -ase diff --git a/setup.py b/setup.py index 8fa0cca05..416aaea9d 100644 --- a/setup.py +++ b/setup.py @@ -25,11 +25,26 @@ "freud-analysis", "networkx>=3.0", "optax>=0.1.4", - "jaxopt>=0.8.0", "pymbar>=4.0.0", "tqdm" ] +extras_require = { + # The QEq module needs optax.lbfgs (added in optax 0.2.3), which is above the + # core optax floor. It replaces the archived jaxopt. dmff/admp/qeq.py degrades + # with a warning if optax is too old, so this stays an opt-in extra. + "qeq": ["optax>=0.2.3"], + "docs": [ + "mkdocs>=1.3.0", + "mkdocs-autorefs>=0.4.1", + "mkdocs-gen-files>=0.3.4", + "mkdocs-literate-nav>=0.4.1", + "mkdocstrings>=0.19.0", + "mkdocstrings-python>=0.7.0", + "pygments>=2.12", + ], +} + def setup(scm=None): packages = find_packages(exclude=["tests"]) @@ -37,23 +52,26 @@ def setup(scm=None): setuptools.setup( name=NAME, use_scm_version=scm, - setup_requires=['setuptools_scm'], author="DeepModeling", author_email="windwhisper.yu@gmail.com", description="Differentiable Molecular Force Field", long_description=readme, long_description_content_type="text/markdown", url="https://github.com/deepmodeling/DMFF", - python_requires="~=3.8", + python_requires=">=3.9", packages=packages, data_files=[], package_data={}, classifiers=[ - "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", "License :: OSI Approved :: GNU Lesser General Public License v3 (LGPLv3)", ], keywords='DMFF', install_requires=install_requires, + extras_require=extras_require, entry_points={} ) diff --git a/tests/test_difftraj/test_difftraj.py b/tests/test_difftraj/test_difftraj.py index 304c731d2..fdc9fe6c8 100644 --- a/tests/test_difftraj/test_difftraj.py +++ b/tests/test_difftraj/test_difftraj.py @@ -1,6 +1,6 @@ import numpy as np import jax -from jax import jit, value_and_grad, random +from jax import jit, value_and_grad import jax.numpy as jnp from dmff.api import Hamiltonian from openmm import * @@ -24,8 +24,7 @@ @pytest.mark.parametrize( "pdbfile, prm, values", - [("tests/data/water_nvt.pdb", "tests/data/qspc-fw.xml", 5412.57173719)]) - # [("tests/data/water_nvt.pdb", "tests/data/qspc-fw.xml", 5401.08336042)]) + [("tests/data/water_nvt.pdb", "tests/data/qspc-fw.xml", 5412.24832246)]) def test_difftraj(pdbfile, prm, values): pdb = PDBFile(pdbfile) @@ -36,10 +35,13 @@ def test_difftraj(pdbfile, prm, values): cov_map = pots.meta['cov_map'] box = jnp.array(pdb.topology.getPeriodicBoxVectors()._value) - key = random.PRNGKey(seed) state = {} state['pos'] = jnp.array(pdb.getPositions()._value).reshape([1, 648, 3]) - state['vel'] = jnp.einsum('ijk,j->ijk', random.normal(key, shape=state['pos'].shape), jnp.sqrt(kT/mass * 1e3)) / 1e3 + # numpy's seeded Generator, not jax.random: jax's default PRNG impl changed + # (jax_threefry_partitionable flipped to True), which would silently move the + # reference energy below. numpy guarantees a stable stream across versions. + noise = np.random.default_rng(seed).standard_normal(state['pos'].shape) + state['vel'] = jnp.einsum('ijk,j->ijk', jnp.array(noise), jnp.sqrt(kT/mass * 1e3)) / 1e3 @jit def L(traj):