From db33c2b3e8b99291cbf2a4c7cf9fc5dc54fc6e7d Mon Sep 17 00:00:00 2001 From: "Zhu, Jia-Xin" Date: Fri, 10 Jul 2026 13:40:21 +0200 Subject: [PATCH 1/4] Support latest JAX (0.10.x) while keeping jax>=0.4.1 compatibility DMFF called a handful of JAX/optax APIs that have since been removed, so a fresh install (which resolves the newest JAX) broke. All replacements below are valid across the whole supported range, so no version-sniffing is needed. - jax.tree_map -> jax.tree_util.tree_map (removed in jax 0.6.0) - isinstance(..., jnp.ndarray) -> jax.Array (jnp.ndarray is a deprecated alias) - optax._src.base -> public optax.EmptyState / optax.GradientTransformation jax's default PRNG impl changed (jax_threefry_partitionable now defaults to True), which silently shifted the reference energy in test_difftraj. Derive the initial velocities from numpy's seeded Generator instead, whose stream is stable across versions. The resulting energy and gradients agree between jax 0.4.35 and 0.10.2 to ~1e-13 relative. jaxopt is archived upstream and only the QEq module needs it, so it moves from install_requires to an optional [qeq] extra. It does still work on jax 0.10. Also lift the numpy<2 cap and the untrue rdkit==2022.03.2 pin, and extend CI to a matrix covering both ends of the supported JAX range. Co-Authored-By: Claude Opus 4.8 --- .github/workflows/ut.yml | 27 +++++++++++++++++++++------ dmff/api/paramset.py | 2 +- dmff/classical/vsite.py | 3 ++- dmff/eann/eann.py | 2 -- dmff/optimize.py | 15 ++++++++++----- dmff/torch_tools.py | 2 +- dmff/utils.py | 3 ++- docs/user_guide/2.installation.md | 26 ++++++++++++++------------ examples/eann/eann.py | 2 -- requirements.txt | 21 ++++++++++++++------- setup.py | 25 ++++++++++++++++++++++--- tests/test_difftraj/test_difftraj.py | 12 +++++++----- 12 files changed, 94 insertions(+), 46 deletions(-) diff --git a/.github/workflows/ut.yml b/.github/workflows/ut.yml index 45b62bed8..49f3f666a 100644 --- a/.github/workflows/ut.yml +++ b/.github/workflows/ut.yml @@ -8,24 +8,39 @@ jobs: build: runs-on: ubuntu-latest strategy: + fail-fast: false matrix: - python-version: [3.9] + include: + # Oldest JAX we still support. jaxopt is archived and only installable + # here, so the QEq path is exercised on this job alone. + - name: jax-oldest + python-version: "3.11" + jax-spec: "jax==0.4.35 jaxlib==0.4.35 jaxopt 'numpy<2'" + # Latest JAX + numpy 2. + - name: jax-latest + python-version: "3.11" + jax-spec: "jax jaxlib" + name: ${{ matrix.name }} (py${{ matrix.python-version }}) 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 -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 }} 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 + pip install ${{ matrix.jax-spec }} networkx parmed pymbar optax tqdm ase - name: Install DMFF run: | source $CONDA/bin/activate dmff && pip install . + - name: Report resolved versions + run: | + source $CONDA/bin/activate dmff + python -c "import jax, numpy; print('jax', jax.__version__, '| numpy', numpy.__version__)" - name: Run Tests run: | source $CONDA/bin/activate dmff @@ -33,7 +48,7 @@ jobs: 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_mbar/test_* pytest -vs tests/test_sgnn/test_* pytest -vs tests/test_difftraj/test_* pytest -vs tests/test_frontend/test_* 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..bdbe91aa8 100644 --- a/docs/user_guide/2.installation.md +++ b/docs/user_guide/2.installation.md @@ -2,35 +2,37 @@ ## 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 `jax>=0.4.1`, including the current release. If you need to reproduce +results against an older JAX, pin it explicitly (e.g. `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: [jaxopt](https://github.com/google/jaxopt) is required **only** by the QEq module. + It is archived upstream (last release 0.8.5) and is therefore not installed by default: +```bash +pip install "jaxopt>=0.8.0" # or: pip install ".[qeq]" +``` ## 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/requirements.txt b/requirements.txt index 1e2ad7b20..35829a517 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,17 @@ -numpy>=1.18, <2 +# Runtime +numpy>=1.23 +jax>=0.4.1 +jaxlib>=0.4.1 +pymbar>=4.0.0 +tqdm +rdkit +ase + +# Optional: required only by the QEq module (dmff/admp/qeq.py). +# jaxopt is archived upstream; install with `pip install .[qeq]` if needed. +# jaxopt>=0.8.0 + +# Docs mkdocs>=1.3.0 mkdocs-autorefs>=0.4.1 mkdocs-gen-files>=0.3.4 @@ -6,9 +19,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..35412140b 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 = { + # jaxopt is archived upstream (last release 0.8.5, Nov 2023). Only the QEq + # module needs it, and dmff/admp/qeq.py already degrades with a warning when + # it is absent, so it is opt-in rather than a hard dependency. + "qeq": ["jaxopt>=0.8.0"], + "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"]) @@ -44,16 +59,20 @@ def setup(scm=None): 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): From 891a7a3e776db1423ddfb6b2da65fffd0411838d Mon Sep 17 00:00:00 2001 From: "Zhu, Jia-Xin" Date: Fri, 10 Jul 2026 14:23:36 +0200 Subject: [PATCH 2/4] Test and support Python 3.9-3.12 CI previously tested a single Python (3.9 before this branch, 3.11 after), while setup.py claimed 3.9-3.12. Expand the matrix so the metadata is actually backed by test runs. The Python and JAX axes are not independent: upstream JAX drops old Pythons over time, so each Python has a different newest installable JAX (3.9 -> 0.4.30, 3.10 -> 0.6.2, 3.11/3.12 -> current). A "latest JAX" job on 3.9 would silently resolve an old JAX, so the matrix pins each leg explicitly and the install docs spell out the table. Installing on 3.9 was broken outright: setup.py declared setup_requires= ['setuptools_scm'], which resolves via the legacy easy_install path and fetches setuptools_scm as an egg *without* its transitive dependencies. setuptools_scm 10.x moved Configuration into a separate vcs_versioning package, so the egg failed to import and metadata generation died. Declare the build deps in a PEP 517 pyproject.toml instead, which makes pip install them properly in an isolated build env, and drop setup_requires. Verified locally: full suite passes on py3.9/jax0.4.30, py3.10/jax0.6.2, py3.11/jax0.10.2 and py3.12/jax0.10.2 (test_dimer and 3 test_frontend tests fail in every env for lack of ambertools, both before and after this branch). The test_difftraj reference energy now holds across jax 0.4.30, 0.4.35, 0.6.2, 0.10.2. Co-Authored-By: Claude Opus 4.8 --- .github/workflows/ut.yml | 32 ++++++++++++++++++++++++------- docs/user_guide/2.installation.md | 17 ++++++++++++++-- pyproject.toml | 7 +++++++ setup.py | 1 - 4 files changed, 47 insertions(+), 10 deletions(-) create mode 100644 pyproject.toml diff --git a/.github/workflows/ut.yml b/.github/workflows/ut.yml index 49f3f666a..0b5e87fea 100644 --- a/.github/workflows/ut.yml +++ b/.github/workflows/ut.yml @@ -4,6 +4,10 @@ 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 @@ -11,18 +15,32 @@ jobs: fail-fast: false matrix: include: - # Oldest JAX we still support. jaxopt is archived and only installable - # here, so the QEq path is exercised on this job alone. - - name: jax-oldest + # 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 jaxopt 'numpy<2'" + # Oldest JAX we pin against on a modern Python. jaxopt is archived and + # optional; test_qeq runs everywhere, but only the two jobs that + # install jaxopt exercise its jitted jaxopt-backed solver path. + - name: py3.11-jax0.4.35 python-version: "3.11" jax-spec: "jax==0.4.35 jaxlib==0.4.35 jaxopt 'numpy<2'" - # Latest JAX + numpy 2. - - name: jax-latest + # 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: ${{ matrix.name }} (py${{ matrix.python-version }}) + - name: py3.12-jax-latest + python-version: "3.12" + jax-spec: "jax jaxlib" + name: ${{ matrix.name }} steps: - 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@v5 with: @@ -40,7 +58,7 @@ jobs: - name: Report resolved versions run: | source $CONDA/bin/activate dmff - python -c "import jax, numpy; print('jax', jax.__version__, '| numpy', numpy.__version__)" + python -c "import sys, jax, numpy; print('python', sys.version.split()[0], '| jax', jax.__version__, '| numpy', numpy.__version__)" - name: Run Tests run: | source $CONDA/bin/activate dmff diff --git a/docs/user_guide/2.installation.md b/docs/user_guide/2.installation.md index bdbe91aa8..c93cef76d 100644 --- a/docs/user_guide/2.installation.md +++ b/docs/user_guide/2.installation.md @@ -12,8 +12,21 @@ pip install "jax[cpu]" # GPU version pip install "jax[cuda12]" ``` -DMFF supports `jax>=0.4.1`, including the current release. If you need to reproduce -results against an older JAX, pin it explicitly (e.g. `pip install jax==0.4.35 jaxlib==0.4.35 "numpy<2"`). +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 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/setup.py b/setup.py index 35412140b..a8476f7aa 100644 --- a/setup.py +++ b/setup.py @@ -52,7 +52,6 @@ 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", From e00679b07d8a14242d7f5379e1b401e7b6803171 Mon Sep 17 00:00:00 2001 From: "Zhu, Jia-Xin" Date: Fri, 10 Jul 2026 14:40:24 +0200 Subject: [PATCH 3/4] Fix CI: pin swig<4.1 for the OpenMM plugin, drive pip via `python -m pip` Two unrelated CI failures. 1. The OpenMM plugin build fails at `make PythonInstall` with "Error: Unknown directive '%factory'". OpenMM's bundled include/swig/OpenMMSwigHeaders.i does `%include "factory.i"` and `%factory(...)`, but SWIG removed the factory.i library in 4.1.0. CMakeLists.txt locates swig with FIND_PROGRAM, and ubuntu-latest now ships swig 4.2+, so the wrapper generation broke without any change on our side. Pin "swig<4.1" into the conda env and pass -DSWIG_EXECUTABLE explicitly so FIND_PROGRAM cannot pick the system one. Documented in the plugin README, since manual builds hit the same. 2. `python -c "import jax"` failed in a step right after `pip install .` reported success. That only happens when `pip` and `python` resolve to different interpreters -- a bare `pip` can come from the setup-python installation ahead of the conda env on PATH. Use `python -m pip` (and `python -m pytest`) throughout so the installing and importing interpreters are the same by construction, add `pip` explicitly to the conda env, and print `which python` plus a `pip list` on failure so a recurrence is self-diagnosing. Co-Authored-By: Claude Opus 4.8 --- .github/workflows/test_openmm_dmff_plugin.yml | 22 ++++++---- .github/workflows/ut.yml | 40 +++++++++++-------- backend/openmm_dmff_plugin/README.md | 10 ++++- 3 files changed, 47 insertions(+), 25 deletions(-) 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 0b5e87fea..7c9c4e61c 100644 --- a/.github/workflows/ut.yml +++ b/.github/workflows/ut.yml @@ -48,28 +48,36 @@ jobs: - name: Install Dependencies run: | source $CONDA/bin/activate - conda create -n dmff -y python=${{ matrix.python-version }} numpy openmm 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 ${{ matrix.jax-spec }} networkx parmed pymbar 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 + python -m pip install ${{ matrix.jax-spec }} networkx parmed pymbar optax 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 - python -c "import sys, jax, numpy; print('python', sys.version.split()[0], '| jax', jax.__version__, '| numpy', numpy.__version__)" + 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 ``` From 2ac4d29895fd23ddfff16daa06b75e33b7b2d903 Mon Sep 17 00:00:00 2001 From: "Zhu, Jia-Xin" Date: Fri, 10 Jul 2026 18:55:25 +0200 Subject: [PATCH 4/4] Replace archived jaxopt with optax.lbfgs in QEq jaxopt is archived (last release Nov 2023) and was DMFF's only use of it. QEq's inner charge optimizer used jaxopt.LBFGS on a projected gradient. Port it to optax.lbfgs (added in optax 0.2.3), which is maintained and already in the dependency tree, so jaxopt is dropped entirely rather than merely made optional. Adds a small `lbfgs_minimize` helper that mirrors the jaxopt.LBFGS contract used here -- minimize over the first argument, converge on the gradient L2 norm, and stay jittable via lax.while_loop. The projected gradient from get_proj_grad is used both as the descent direction and the stopping criterion; the Armijo backtracking line search only evaluates (never differentiates) the objective, so the charge-constraint projection is preserved. Numerically validated against the jaxopt implementation on a perturbed-charge case that actually exercises the LBFGS branch (test_qeq.py only hits the initial linear-solve branch): both converge to the same projected minimum (81.33959058 vs 81.33959057, final |q| identical to 1e-8). lbfgs_minimize also verified jittable and converging to machine precision on optax 0.2.4/jax 0.4.30 and optax 0.2.8/jax 0.10.2. Full test_admp (23 tests) passes on both. The qeq extra now points at optax>=0.2.3 instead of jaxopt; CI installs optax>=0.2.3 and no longer installs jaxopt. Co-Authored-By: Claude Opus 4.8 --- .github/workflows/ut.yml | 11 ++-- dmff/admp/qeq.py | 90 +++++++++++++++++++++++++------ docs/user_guide/2.installation.md | 7 +-- requirements.txt | 5 +- setup.py | 8 +-- 5 files changed, 88 insertions(+), 33 deletions(-) diff --git a/.github/workflows/ut.yml b/.github/workflows/ut.yml index 7c9c4e61c..27faa3143 100644 --- a/.github/workflows/ut.yml +++ b/.github/workflows/ut.yml @@ -18,13 +18,11 @@ jobs: # 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 jaxopt 'numpy<2'" - # Oldest JAX we pin against on a modern Python. jaxopt is archived and - # optional; test_qeq runs everywhere, but only the two jobs that - # install jaxopt exercise its jitted jaxopt-backed solver path. + 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 jaxopt 'numpy<2'" + 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" @@ -56,7 +54,8 @@ jobs: # interpreter than `python` imports from. which python python -m pip install --upgrade pip - python -m pip install ${{ matrix.jax-spec }} networkx parmed pymbar optax tqdm ase + # 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 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/docs/user_guide/2.installation.md b/docs/user_guide/2.installation.md index c93cef76d..03c4073e4 100644 --- a/docs/user_guide/2.installation.md +++ b/docs/user_guide/2.installation.md @@ -41,10 +41,11 @@ conda install -c conda-forge openmm ```bash conda install -c conda-forge rdkit ``` -+ Optional: [jaxopt](https://github.com/google/jaxopt) is required **only** by the QEq module. - It is archived upstream (last release 0.8.5) and is therefore not installed by default: ++ 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 "jaxopt>=0.8.0" # or: pip install ".[qeq]" +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`. : diff --git a/requirements.txt b/requirements.txt index 35829a517..2d9fdfefd 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,9 +7,8 @@ tqdm rdkit ase -# Optional: required only by the QEq module (dmff/admp/qeq.py). -# jaxopt is archived upstream; install with `pip install .[qeq]` if needed. -# jaxopt>=0.8.0 +# 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 diff --git a/setup.py b/setup.py index a8476f7aa..416aaea9d 100644 --- a/setup.py +++ b/setup.py @@ -30,10 +30,10 @@ ] extras_require = { - # jaxopt is archived upstream (last release 0.8.5, Nov 2023). Only the QEq - # module needs it, and dmff/admp/qeq.py already degrades with a warning when - # it is absent, so it is opt-in rather than a hard dependency. - "qeq": ["jaxopt>=0.8.0"], + # 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",