JAX Backend¶

A pure-JAX implementation supporting both harmonic (OPLSAA-style) and MM3 functional forms, including bond, angle, torsion, stretch-bend cross-term, and vdW energy terms. Near-linear torsion terms (central angle >170°) are smoothly suppressed to prevent the well-known dihedral gradient singularity. Best for small-to-medium molecules where periodic boundaries and neighbor lists are not needed. All energy functions are differentiable via jax.grad, enabling analytical gradient computation.

Installation¶

pip install jax jaxlib

For GPU support, install the CUDA-enabled jaxlib:

pip install jax[cuda12]

Verify installation

import jax
print(jax.__version__)
print(jax.default_backend())  # "cpu" or "gpu"

Supported energy terms¶

Term	Supported
Bonds (harmonic + MM3)	✅
Angles (harmonic + MM3)	✅
Torsions (cosine)	✅
Improper torsions	❌
vdW (LJ 12-6 + Buckingham exp-6)	✅
Electrostatics	❌
1-4 scaling	❌ Not implemented

Functional forms: Harmonic and MM3.

Configuration¶

from q2mm.backends.mm import JaxBackend

backend = JaxBackend()

JaxBackend has no constructor parameters. It runs on whichever JAX backend is active (cpu or gpu), detected via jax.default_backend().

Capabilities¶

Method	Supported	Notes
`energy()`	✅	Pure JAX
`minimize()`	✅	JAX gradients + SciPy L-BFGS-B
`hessian()`	✅	Analytical via `jax.hessian`
`frequencies()`	✅	From analytical Hessian
`parameter_gradient()`	✅	Analytical via `jax.grad`
`batched_energy()`	✅	Vectorized via `jax.vmap`
`hessian_parameter_jacobian()`	✅	Analytical Hessian-parameter derivatives

GPU support¶

JaxBackend 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 (harmonic, gpu) or JAX (harmonic, cpu)).

Performance

In the current benchmark set, JaxBackend is one of the fastest in-process backends for harmonic CH₃F optimization and offers analytical gradients for energy-based evaluators. Exact speedups depend on system size, objective, and device, so use the benchmark overview and GPU benchmarks for workload-specific numbers.

Optax optimizers

JaxBackend pairs naturally with optax adaptive optimizers (Adam, AdaGrad, SGD) via OptaxOptimizer. These use JAX's analytical gradients automatically — no finite-difference overhead. On CH₃F MM3, Adam achieves 56.3 cm⁻¹ RMSD (10× better than L-BFGS-B). See Small Molecules for full results.

Analytical-gradient optimization

For multi-molecule transition-state systems, pair JaxBackend with JaxObjectiveExecutor + ScipyOptimizer(method="L-BFGS-B"). That route builds per-case JIT loss fragments and feeds analytical gradients to SciPy's L-BFGS-B implementation without placing all molecules in one XLA program. JaxOptOptimizer remains useful for small single-molecule systems, but it is not the recommended default for the literature-scale TS benchmarks.

Limitations¶

No 1-4 pair scaling — non-bonded energies differ from OpenMM/JAX-MD for molecules with 1-4 interactions. See the compatibility notes.
No periodic boundaries — gas-phase only.

Example¶

from q2mm.backends.mm import JaxBackend
from q2mm.backends.contracts import (
    EnergyRequest,
    ParameterGradientRequest,
    PreparationRequest,
)
from q2mm.io.xyz import load_xyz
from q2mm.models.forcefield import ForceField
from q2mm.models.parameters import ParameterLayout

mol = load_xyz("molecule.xyz")
ff = ForceField.create_for_molecule(mol)
layout = ParameterLayout.from_force_field(ff)
params = layout.vector(ff)

backend = JaxBackend()
session = backend.prepare(PreparationRequest(case_id="0", molecule=mol, force_field=ff))
e = session.energy(EnergyRequest(parameters=params)).energy
print(f"JAX 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}")