JAX-MD Backend¶

Built on the JAX-MD library, this backend adds periodic boundary conditions, neighbor lists, and configurable 1-4 scaling on top of JAX's differentiable energy functions.

Installation¶

The supported CPU lane uses Python 3.14 on Linux/WSL; the current Python 3.15 dependency chain is not available. Install JAX, jaxlib and JAX-MD with one package manager in the same environment, as in the canonical environment. Do not replace a conda-owned jaxlib with pip. This is distinct from the primary Python 3.15 JAX-only lane.

pip install jax jaxlib jax-md

For GPU support:

pip install jax[cuda12] jax-md

Supported energy terms¶

Term	Supported
Bonds (harmonic)	✅
Angles (harmonic)	✅
Torsions (cosine)	✅
Improper torsions	❌
vdW (LJ 12-6)	✅
Electrostatics	Infrastructure only (charges zeroed)
1-4 scaling	✅ Configurable (default: AMBER 0.5)
Periodic boundaries	✅
Neighbor lists	✅ (jax-md native)

Functional forms: Harmonic only.

Preparation gates¶

Populated Urey-Bradley, CMAP, improper parameters, nonzero bond dipoles, and nondefault vdW reduction raise PreparationError before layout or native-state construction. The existing stretch-bend rejection and harmonic-only form restriction remain in place. Reference partial charges do not populate the Coulomb handler. See the shared term-presence rules.

Configuration¶

from q2mm.backends.mm import JaxMdBackend

backend = JaxMdBackend(
    box=(100.0, 100.0, 100.0),   # simulation box dimensions (Å)
    coulomb=None,                 # CoulombHandler; default: CutoffCoulomb(r_cut=12.0)
    nb_options=None,              # NonbondedOptions; default: r_cut=12.0
)

Parameter	Type	Default	Description
`box`	`tuple[float, float, float]`	`(100.0, 100.0, 100.0)`	Periodic box dimensions in Å
`coulomb`	`CoulombHandler \\| None`	`CutoffCoulomb(r_cut=12.0)`	Electrostatic handler
`nb_options`	`NonbondedOptions \\| None`	`NonbondedOptions(r_cut=12.0)`	Non-bonded cutoff and options

Capabilities¶

Prepared-session operation	Supported	Notes
`energy(EnergyRequest)`	✅	—
`minimize(MinimizationRequest)`	✅	JAX gradients + SciPy L-BFGS-B; control limits
`hessian(HessianRequest)`	✅	Analytical via `jax.hessian`
`frequencies(FrequencyRequest)`	✅	From analytical Hessian
`parameter_gradient(ParameterGradientRequest)`	✅	Analytical via `jax.grad`
`batched_energy(BatchedEnergyRequest)`	✅	Vectorized via `jax.vmap`
`Capability.REUSABLE_STATE`	✅	Prepared session reuses compiled JAX functions

Minimization controls¶

Omit MinimizationRequest.tolerance (or use None) to retain the native SciPy stopping defaults. Any explicit value raises EvaluationError before the native minimizer is invoked; it is not silently ignored.

The existing helper passes only maxiter to SciPy L-BFGS-B, with no tol, ftol, or gtol override. SciPy has distinct function-reduction and projected-gradient criteria; this adapter has no established mapping of the generic request's native-unit tolerance to one of them. Rejecting the option avoids inventing that mapping.

max_iterations=None retains the existing 200-iteration limit, and an explicit supported iteration limit is forwarded unchanged. No energy or solver implementation, native stopping default, or result field changes. This control-support boundary does not establish convergence validity: finite returned coordinates are not new convergence evidence, and GEO-01c/D-06 remain separate.

Optax and JaxOpt optimizers

JaxMdBackend exposes analytical parameter gradients through parameter_gradient(ParameterGradientRequest), making it compatible with Optax and JaxOpt optimizers. See the Optimization Guide for workflow recommendations.

GPU support¶

JaxMdBackend runs on whichever device JAX selects. To use a GPU:

Install the CUDA-enabled JAX: pip install jax[cuda12]
Verify: python -c "import jax; print(jax.default_backend())"

The backend name includes the JAX device string (e.g., JAX-MD (OPLSAA, gpu)).

Limitations¶

Harmonic only — MM3 functional form is not yet supported.
Electrostatics zeroed — Coulomb energy is computed with zero charges; charge optimization is not yet supported.
Unsupported populated terms — Urey-Bradley, CMAP, improper parameters, bond dipoles, and nondefault vdW reduction are rejected, not silently omitted or replaced by empty native terms.
64-bit mode forced — importing this module enables jax_enable_x64 globally, which affects all JAX code in the process.

Example¶

from q2mm.backends.contracts import EnergyRequest, ParameterGradientRequest, PreparationRequest
from q2mm.backends.mm import JaxMdBackend
from q2mm.io.xyz import load_xyz
from q2mm.models.forcefield import ForceField

mol = load_xyz("molecule.xyz")
ff = ForceField.create_for_molecule(mol)

backend = JaxMdBackend(box=(50.0, 50.0, 50.0))
session = backend.prepare(PreparationRequest(case_id="example", molecule=mol, force_field=ff))
params = session.layout.vector(ff)

e = session.energy(EnergyRequest(parameters=params)).energy
print(f"JAX-MD energy: {e:.4f} kcal/mol")

# Analytical parameter gradients
grad_result = session.parameter_gradient(ParameterGradientRequest(parameters=params))
print(f"Energy: {grad_result.energy:.4f}, grad shape: {grad_result.gradient.shape}")