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

[ChiahsinChu]

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

Motivation

DMFF calls a few JAX and optax APIs that upstream has since removed. Because
setup.py declares only jax>=0.4.1 with no upper bound, a fresh pip install .
resolves the newest JAX and then fails. Anyone installing DMFF today on Python
3.11+ hits this.

This PR makes DMFF work on current JAX without dropping support for older
versions. The declared floor stays at jax>=0.4.1: every replacement below is
valid across the whole range, so no version-sniffing or compat shim is needed.

What was broken (removed / private JAX & optax APIs)

Removed / private API	Replacement	Where
jax.tree_map (removed in jax 0.6.0)	jax.tree_util.tree_map	dmff/optimize.py, dmff/api/paramset.py
isinstance(x, jnp.ndarray) (deprecated alias)	isinstance(x, jax.Array)	dmff/utils.py, dmff/classical/vsite.py, dmff/torch_tools.py
optax._src.base.{EmptyState,GradientTransformation,NO_PARAMS_MSG} (private)	public optax.EmptyState / optax.GradientTransformation; message inlined	dmff/optimize.py

jax.tree_util.tree_map and jax.Array exist as far back as jax 0.4.1, and
optax.EmptyState / optax.GradientTransformation are public in every optax from
0.1.4 onward, so the floor is unchanged. (jax.tree.map was deliberately not
used — it requires jax>=0.4.25 and would have raised the floor.) Two dead
# from jax.config import config comments (referencing a module deleted in jax
0.4.27) were also removed.

test_difftraj and the PRNG default

test_difftraj failed on new JAX with 5401.08336042 — which sits in the file
today as a commented-out alternative expected value, suggesting this has been
papered over before by swapping the constant.

The cause is not DMFF. JAX flipped jax_threefry_partitionable from False to
True by default, so jax.random.normal(PRNGKey(0)) yields different initial
velocities. Rather than pin a legacy flag JAX intends to remove, the test now
draws its initial velocities from numpy.random.default_rng(seed), whose stream
numpy guarantees stable across versions. With that fixture the reference energy
5412.24832246 and gradients reproduce identically on jax 0.4.30, 0.4.35,
0.6.2 and 0.10.2 (~1e-13 relative) — a stronger check than the old constant,
which only pinned one version's RNG output.

Drop jaxopt: port QEq's L-BFGS to optax

jaxopt is archived upstream (last release 0.8.5, Nov 2023) and was DMFF's only
use of it — QEq's inner charge optimizer (dmff/admp/qeq.py) used
jaxopt.LBFGS on a projected gradient. It is ported to optax.lbfgs (added in
optax 0.2.3, already in the dependency tree), so jaxopt is removed entirely
rather than merely made optional.

A small lbfgs_minimize helper mirrors the jaxopt contract used here: minimize
over the first argument, converge on the gradient L2 norm, stay jittable via
lax.while_loop. QEq feeds L-BFGS a projected gradient to keep charges on the
constraint plane; that is preserved by using the projected gradient as both the
descent direction and the stopping criterion, while the Armijo backtracking line
search only evaluates (never differentiates) the objective, so it cannot
reintroduce the unprojected gradient.

Validated against the jaxopt implementation on a perturbed-charge case that
actually exercises the L-BFGS branch — note test_qeq.py only ever hits the
initial linear-solve branch (hessinv_stride=1), so the test suite alone does
not cover this swap:

	jaxopt.LBFGS	optax.lbfgs
converged energy	81.33959057	81.33959058
final |q|	0.19126898	0.19126898

lbfgs_minimize was additionally checked jittable and converging to machine
precision on optax 0.2.4 / jax 0.4.30 and optax 0.2.8 / jax 0.10.2, and
test_qeq.py passes in an environment with no jaxopt installed at all. The
qeq extra now requires optax>=0.2.3 (above the core optax>=0.1.4 floor, so
old-JAX non-QEq users aren't forced up).

Python 3.9–3.12

CI tested a single Python (3.9) while setup.py said python_requires="~=3.8".
The matrix now covers 3.9–3.12 and the classifiers match.

The two axes are not independent: JAX drops old Pythons over time, so each Python
has a different newest installable JAX. Critically, pip install jax on 3.9
silently resolves an old release rather than failing, so a naive
python × [oldest, latest] matrix would contain jobs that test nothing they claim
to. Each leg is therefore pinned explicitly:

Python	newest installable JAX
3.9	0.4.30
3.10	0.6.2
3.11 / 3.12	current

Documented in docs/user_guide/2.installation.md so users on 3.9/3.10 understand
why they cannot get current JAX.

Installing on Python 3.9 was broken outright

Independently of JAX: setup.py declared setup_requires=['setuptools_scm'],
which resolves through 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 fails to
import and metadata generation dies:

File ".eggs/setuptools_scm-10.2.0-py3.9.egg/setuptools_scm/__init__.py", line 8
    from vcs_versioning import Configuration
ModuleNotFoundError: No module named 'vcs_versioning'
error: metadata-generation-failed

Fixed by declaring build dependencies in a PEP 517 pyproject.toml
(requires = ["setuptools>=64", "setuptools_scm>=8", "wheel"]), which makes pip
install them correctly in an isolated build environment, and dropping
setup_requires. CI gains fetch-depth: 0 so setuptools_scm can derive a
version.

Two CI failures fixed along the way

Both are pre-existing breakage caused by the environment moving under us, not by
this PR's code changes.

test_openmm_dmff_plugin.yml fails at make PythonInstall:

/usr/share/miniconda/envs/dmff_omm/include/swig/OpenMMSwigHeaders.i:12:
    Error: Unknown directive '%factory'.

OpenMM's bundled include/swig/OpenMMSwigHeaders.i does %include "factory.i"
then %factory(...). SWIG removed the factory.i library in 4.1.0.
backend/openmm_dmff_plugin/CMakeLists.txt:138 locates the tool with
FIND_PROGRAM(SWIG_EXECUTABLE swig), and the ubuntu-latest runner image now
ships SWIG 4.2+, so wrapper generation broke with no change on our side. Fixed by
pinning "swig<4.1" into the conda env (4.0.2 solves cleanly with openmm=7.7 /
python=3.9 / numpy=1.*) and passing -DSWIG_EXECUTABLE=${CONDA_PREFIX}/bin/swig
so FIND_PROGRAM cannot pick up the system one. Same note added to the plugin
README, since manual builds hit this identically.

ut.yml fails in the version-report step with ModuleNotFoundError: No module named 'jax', immediately after pip install . reported success. That
combination only occurs when pip and python resolve to different
interpreters: a bare pip can come from the actions/setup-python installation
ahead of the conda env on PATH, so JAX is installed into one interpreter and
imported from another. Fixed by driving everything through python -m pip /
python -m pytest, adding pip explicitly to the conda env, and printing
which python plus a pip list on failure so any recurrence is self-diagnosing.

Other dependency cleanups

• requirements.txt: drop the numpy<2 cap (current JAX supports numpy 2; the
  source has no numpy-2-removed API usage). Verified under numpy 2.4.
• requirements.txt: drop the stale rdkit==2022.03.2 pin — the suite passes
  against current rdkit, and setup.py never listed rdkit, so the pin only ever
  applied to people installing from this file.
• requirements.txt: separate runtime deps from the mkdocs tooling; docs extras
  are also exposed as extras_require["docs"].
• docs/user_guide/2.installation.md had a duplicated dependency block; consolidated.

Verification

Full test suite run on five Python × JAX combinations:

Python	JAX	numpy	optax	result
3.9	0.4.30	1.26	0.2.4	pass
3.10	0.6.2	2.2	0.2.8	pass
3.11	0.4.35	1.26	0.2.5	pass (baseline, unchanged)
3.11	0.10.2	2.4	0.2.8	pass
3.12	0.10.2	2.4	0.2.8	pass

tests/test_dimer (6) and 3 of tests/test_frontend fail in every environment,
before and after this branch, because antechamber/sqm (ambertools) are absent.
They are unrelated to this change.

Beyond the suite, examples/water_fullpol/run.py runs end to end on both jax
0.4.35 and 0.10.2; across ~27k printed energies and gradients the max relative
difference is 1.9e-12 — float64 round-off from different XLA codegen, not a
semantic change. This matters because tests/test_admp/test_gradDisp.py enables
jax_debug_nans, so a clean ADMP suite plus a clean water_fullpol run together
are the real signal that autodiff still behaves.

Notes for reviewers

• Nothing here raises the minimum JAX version.
• Python 3.13 is untested — JAX supports it, but openmm/rdkit/mdtraj availability
  on conda-forge for 3.13 was not verified, so the classifiers stop at 3.12.
• pytest tests/ as a whole still fails at collection: tests/test_ase/test_energy.py
  and tests/test_sgnn/test_energy.py share a basename with no __init__.py.
  Pre-existing, not touched here; Makefile and CI run per-module and are fine.
• Also pre-existing: import dmff.admp.qeq as the very first import raises a
  circular-import error; it works via the normal path (package init first), which
  is why the suite is unaffected. Left as-is.