Feat: Support the current JAX release (0.10.x), test Python 3.9–3.12, and drop jaxopt

.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

.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_*

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
    ```

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,

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

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

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):

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(

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

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 \