42.4 Building a Differentiable Force Field

1. Problem Setup: Force Field Calibration

When a pharmaceutical company simulates a drug candidate docking into a protein pocket, every predicted binding energy rests on the accuracy of the underlying force field parameters. Shift those parameters by even a few percent and the simulation promotes the wrong molecule into clinical trials, wasting years and billions of dollars.

Force field calibration is one of the oldest problems in computational chemistry. A force field defines how atoms interact: the Lennard-Jones potential \(V(r) = 4\epsilon[(\sigma/r)^{12} - (\sigma/r)^6]\) has two parameters per atom type pair: \(\epsilon\) (well depth, controlling interaction strength) and \(\sigma\) (length scale, controlling the effective atom size). The goal is to find \(\epsilon\) and \(\sigma\) values that reproduce experimentally observed structural properties.

Force field calibration tunes the numerical constants in an interatomic potential so that simulations reproduce measurable physical properties (density, structure, diffusion) of a real material. It matters because even a physically motivated functional form like Lennard-Jones produces wildly wrong behavior if its parameters are off by 10%. That gap makes calibration the bottleneck between writing down a potential and trusting its predictions. The mechanism is an inverse problem: run a simulation forward with candidate parameters, compare the output to experimental observables, and adjust the parameters to shrink the mismatch. Use gradient-based calibration (as in this recipe) when the potential is differentiable and you have more than two or three parameters; for very small parameter spaces or non-differentiable potentials, grid search or Bayesian optimization may suffice.

The Observable: Radial Distribution Functions

The radial distribution function (RDF), \(g(r)\), measures the probability of finding a particle at distance \(r\) from another particle, normalized by the ideal gas density. Peaks in \(g(r)\) correspond to coordination shells. The RDF is one of the most commonly used structural observables because it can be measured experimentally (via X-ray or neutron scattering) and computed from simulation trajectories.

Our objective: given a target RDF \(g^*(r)\) (representing experimental data or a high-fidelity reference simulation), find \(\epsilon\) and \(\sigma\) such that the RDF from a simulation with those parameters matches \(g^*(r)\).

2. Generating the Target RDF

In a real application, the target RDF would come from experiment or a quantum mechanical simulation. For this recipe, we generate a synthetic target by running a simulation with known "true" parameters and computing the RDF. This lets us verify that optimization recovers the correct parameters. In short: if you can simulate it forward, you can differentiate it backward, turning any computable observable into a loss function for parameter learning.

import jax
import jax.numpy as jnp
from functools import partial

# --- Simulation infrastructure ---

def lj_energy_pair(r, epsilon, sigma):
    """Lennard-Jones pair energy with a smooth cutoff.

    Parameters
    ----------
    r : float
        Interatomic distance.
    epsilon : float
        Well depth.
    sigma : float
        Length scale.

    Returns
    -------
    float
        Pair energy, smoothly truncated at r_cut = 2.5*sigma.
    """
    r_cut = 2.5 * sigma
    s6 = (sigma / r) ** 6
    e_full = 4.0 * epsilon * (s6 ** 2 - s6)
    # Shift to zero at cutoff
    s6_cut = (sigma / r_cut) ** 6
    e_cut = 4.0 * epsilon * (s6_cut ** 2 - s6_cut)
    return jnp.where(r < r_cut, e_full - e_cut, 0.0)


def compute_total_energy(positions, box_size, epsilon, sigma):
    """Total LJ energy for N particles in a periodic box.

    Parameters
    ----------
    positions : jnp.ndarray, shape (N, 3)
        Particle positions.
    box_size : float
        Side length of the cubic periodic box.
    epsilon : float
        LJ well depth.
    sigma : float
        LJ length scale.

    Returns
    -------
    float
        Total potential energy.
    """
    N = positions.shape[0]

    def pair_energy(i, j):
        dr = positions[i] - positions[j]
        # Minimum image convention for periodic boundaries
        dr = dr - box_size * jnp.round(dr / box_size)
        r = jnp.sqrt(jnp.sum(dr ** 2) + 1e-12)
        return lj_energy_pair(r, epsilon, sigma)

    # Build all unique pairs
    idx_i, idx_j = jnp.triu_indices(N, k=1)
    pair_e = jax.vmap(pair_energy)(idx_i, idx_j)
    return jnp.sum(pair_e)


def compute_forces(positions, box_size, epsilon, sigma):
    """Compute forces via the negative gradient of total energy.

    Parameters
    ----------
    positions : jnp.ndarray, shape (N, 3)
        Particle positions.
    box_size : float
        Side length of cubic periodic box.
    epsilon : float
        LJ well depth.
    sigma : float
        LJ length scale.

    Returns
    -------
    jnp.ndarray, shape (N, 3)
        Force on each particle.
    """
    grad_fn = jax.grad(compute_total_energy, argnums=0)
    return -grad_fn(positions, box_size, epsilon, sigma)

3. The Velocity Verlet Integrator

With forces computed automatically from the energy function, the next question is how to march those forces forward in time to produce a trajectory.

We use the Velocity Verlet algorithm, the standard integrator for molecular dynamics. It is symplectic (preserving the phase-space volume, a property that prevents energy drift over long simulations), time-reversible, and second-order accurate:

where \(a(t) = F(x(t))/m\) is the acceleration. We implement this as a pure JAX function suitable for jax.lax.scan.

def velocity_verlet_step(state, _, box_size, epsilon, sigma, dt, mass=1.0):
    """One Velocity Verlet integration step.

    Parameters
    ----------
    state : tuple
        (positions, velocities) each of shape (N, 3).
    _ : ignored
        Placeholder for scan compatibility.
    box_size : float
        Periodic box side length.
    epsilon : float
        LJ well depth.
    sigma : float
        LJ length scale.
    dt : float
        Time step.
    mass : float
        Particle mass (uniform).

    Returns
    -------
    tuple
        ((new_positions, new_velocities), positions_for_rdf)
    """
    positions, velocities = state

    # Forces at current positions
    forces = compute_forces(positions, box_size, epsilon, sigma)
    accel = forces / mass

    # Half-step velocity
    v_half = velocities + 0.5 * dt * accel

    # Full-step positions (with periodic wrapping)
    new_positions = positions + dt * v_half
    new_positions = new_positions % box_size

    # Forces at new positions
    new_forces = compute_forces(new_positions, box_size, epsilon, sigma)
    new_accel = new_forces / mass

    # Complete velocity step
    new_velocities = v_half + 0.5 * dt * new_accel

    return (new_positions, new_velocities), new_positions


def run_md_trajectory(positions, velocities, box_size, epsilon, sigma,
                      n_steps, dt=0.002):
    """Run an MD trajectory and return positions at each step.

    Parameters
    ----------
    positions : jnp.ndarray, shape (N, 3)
        Initial positions.
    velocities : jnp.ndarray, shape (N, 3)
        Initial velocities.
    box_size : float
        Periodic box side length.
    epsilon : float
        LJ well depth.
    sigma : float
        LJ length scale.
    n_steps : int
        Number of integration steps.
    dt : float
        Time step.

    Returns
    -------
    jnp.ndarray, shape (n_steps, N, 3)
        Positions at each time step.
    """
    step_fn = partial(velocity_verlet_step,
                      box_size=box_size, epsilon=epsilon,
                      sigma=sigma, dt=dt)

    (final_pos, final_vel), trajectory = jax.lax.scan(
        step_fn,
        (positions, velocities),
        None,
        length=n_steps,
    )
    return trajectory

4. Computing the Radial Distribution Function

The RDF \(g(r)\) is computed by histogramming pairwise distances and normalizing by the ideal gas expectation. For a system of \(N\) particles in a box of volume \(V\), the normalization for a spherical shell at distance \(r\) with thickness \(\Delta r\) is:

where \(n(r)\) is the number of pairs with distance in \([r, r + \Delta r)\), averaged over trajectory frames. We need a differentiable RDF computation, so we replace the hard histogram bins with soft Gaussian kernels.

def compute_rdf_soft(positions, box_size, r_bins, bin_width=0.05):
    """Compute a differentiable RDF using soft Gaussian binning.

    Hard histograms have zero gradient (the count does not change
    for infinitesimal position perturbations). Soft Gaussian bins
    spread each pair distance into neighboring bins with smooth
    weights, enabling gradient flow.

    Parameters
    ----------
    positions : jnp.ndarray, shape (N, 3)
        Particle positions.
    box_size : float
        Periodic box side length.
    r_bins : jnp.ndarray, shape (n_bins,)
        Bin centers for the RDF.
    bin_width : float
        Gaussian kernel width for soft binning.

    Returns
    -------
    jnp.ndarray, shape (n_bins,)
        Soft RDF values at each bin center.
    """
    N = positions.shape[0]

    # All pairwise displacement vectors (minimum image)
    dr = positions[:, None, :] - positions[None, :, :]
    dr = dr - box_size * jnp.round(dr / box_size)
    distances = jnp.sqrt(jnp.sum(dr ** 2, axis=-1) + 1e-12)

    # Mask self-pairs and count only upper triangle
    mask = jnp.triu(jnp.ones((N, N), dtype=bool), k=1)
    pair_distances = distances[mask]  # shape: (N*(N-1)/2,)

    # Soft histogram: each distance contributes a Gaussian to each bin
    # pair_distances: (n_pairs,), r_bins: (n_bins,)
    diff = pair_distances[:, None] - r_bins[None, :]  # (n_pairs, n_bins)
    weights = jnp.exp(-0.5 * (diff / bin_width) ** 2) / (
        bin_width * jnp.sqrt(2.0 * jnp.pi)
    )
    counts = jnp.sum(weights, axis=0)  # (n_bins,)

    # Normalize: ideal gas shell volume
    volume = box_size ** 3
    n_pairs = N * (N - 1) / 2
    rho = N / volume
    shell_vol = 4.0 * jnp.pi * r_bins ** 2 * bin_width
    # Avoid division by zero at r=0
    shell_vol = jnp.maximum(shell_vol, 1e-10)
    g_r = counts / (n_pairs * rho * shell_vol)

    return g_r


def compute_rdf_trajectory(trajectory, box_size, r_bins, bin_width=0.05,
                           skip=10):
    """Average RDF over trajectory frames.

    Parameters
    ----------
    trajectory : jnp.ndarray, shape (n_steps, N, 3)
        Particle positions at each time step.
    box_size : float
        Periodic box side length.
    r_bins : jnp.ndarray, shape (n_bins,)
        RDF bin centers.
    bin_width : float
        Soft binning width.
    skip : int
        Use every skip-th frame (decorrelation).

    Returns
    -------
    jnp.ndarray, shape (n_bins,)
        Time-averaged RDF.
    """
    frames = trajectory[::skip]
    rdf_fn = partial(compute_rdf_soft, box_size=box_size,
                     r_bins=r_bins, bin_width=bin_width)
    # vmap over frames, then average
    rdfs = jax.vmap(rdf_fn)(frames)
    return jnp.mean(rdfs, axis=0)

5. The Optimization Loop

Now we assemble the full pipeline: initialize particles, run MD with current parameters, compute the RDF, compare to the target, and backpropagate through everything to update the parameters. One practical detail: we optimize \(\log\epsilon\) and \(\log\sigma\) rather than \(\epsilon\) and \(\sigma\) directly, a standard reparameterization that guarantees positivity (the exponential of any real number is always positive) while giving the optimizer an unconstrained search space.

import optax

def initialize_fcc_positions(n_cells, lattice_constant):
    """Initialize particles on an FCC lattice.

    Parameters
    ----------
    n_cells : int
        Number of unit cells per dimension.
    lattice_constant : float
        FCC lattice constant.

    Returns
    -------
    tuple
        (positions array shape (N, 3), box_size float)
    """
    # Face-centered cubic (FCC) basis: 4 atoms per unit cell
    basis = jnp.array([
        [0.0, 0.0, 0.0],
        [0.5, 0.5, 0.0],
        [0.5, 0.0, 0.5],
        [0.0, 0.5, 0.5],
    ]) * lattice_constant

    positions = []
    for ix in range(n_cells):
        for iy in range(n_cells):
            for iz in range(n_cells):
                offset = jnp.array([ix, iy, iz]) * lattice_constant
                for b in basis:
                    positions.append(offset + b)

    box_size = n_cells * lattice_constant
    return jnp.array(positions), box_size


def initialize_velocities(key, N, temperature, mass=1.0):
    """Initialize velocities from a Maxwell-Boltzmann distribution.

    Parameters
    ----------
    key : jax.random.PRNGKey
        Random key.
    N : int
        Number of particles.
    temperature : float
        Target temperature (in reduced units, kT).
    mass : float
        Particle mass.

    Returns
    -------
    jnp.ndarray, shape (N, 3)
        Initial velocities with zero total momentum.
    """
    sigma_v = jnp.sqrt(temperature / mass)
    velocities = jax.random.normal(key, (N, 3)) * sigma_v
    # Remove center-of-mass velocity
    velocities = velocities - jnp.mean(velocities, axis=0)
    return velocities


# --- Setup ---

# True parameters (the optimization target)
TRUE_EPSILON = 1.0
TRUE_SIGMA = 1.0

# Initialize system
n_cells = 3  # 3x3x3 FCC (face-centered cubic) = 108 atoms
lattice_constant = 1.6  # in units of sigma
positions_init, box_size = initialize_fcc_positions(n_cells, lattice_constant)
N = positions_init.shape[0]
print(f"System: {N} particles in box of size {box_size:.2f}")

key = jax.random.key(42)
velocities_init = initialize_velocities(key, N, temperature=1.0)

# RDF bins
r_bins = jnp.linspace(0.5, 4.0, 80)

# Generate target RDF from "true" parameters
n_equil = 200   # equilibration steps
n_prod = 300    # production steps

# Equilibrate, then produce
equil_traj = run_md_trajectory(
    positions_init, velocities_init, box_size,
    TRUE_EPSILON, TRUE_SIGMA, n_steps=n_equil, dt=0.002,
)
equil_final = equil_traj[-1]
equil_vel = velocities_init  # simplified; real code updates velocities

prod_traj = run_md_trajectory(
    equil_final, equil_vel, box_size,
    TRUE_EPSILON, TRUE_SIGMA, n_steps=n_prod, dt=0.002,
)

target_rdf = compute_rdf_trajectory(prod_traj, box_size, r_bins)
print(f"Target RDF computed: {target_rdf.shape}")

def rdf_loss(params, positions_init, velocities_init, box_size,
             r_bins, target_rdf, n_steps=300, dt=0.002):
    """Loss function: squared difference between predicted and target RDF.

    Parameters
    ----------
    params : dict
        {'log_epsilon': float, 'log_sigma': float}
        We optimize in log-space to enforce positivity.
    positions_init : jnp.ndarray, shape (N, 3)
        Starting positions (post-equilibration).
    velocities_init : jnp.ndarray, shape (N, 3)
        Starting velocities.
    box_size : float
        Periodic box size.
    r_bins : jnp.ndarray, shape (n_bins,)
        RDF bin centers.
    target_rdf : jnp.ndarray, shape (n_bins,)
        Target RDF to match.
    n_steps : int
        Number of MD steps.
    dt : float
        Time step.

    Returns
    -------
    float
        Mean squared error between predicted and target RDF.
    """
    # Exponentiate to enforce positivity
    epsilon = jnp.exp(params['log_epsilon'])
    sigma = jnp.exp(params['log_sigma'])

    # Run MD trajectory
    trajectory = run_md_trajectory(
        positions_init, velocities_init, box_size,
        epsilon, sigma, n_steps=n_steps, dt=dt,
    )

    # Compute RDF from trajectory
    predicted_rdf = compute_rdf_trajectory(
        trajectory, box_size, r_bins, skip=5,
    )

    # Mean squared error (MSE) loss
    return jnp.mean((predicted_rdf - target_rdf) ** 2)


# --- Optimization ---

# Start from deliberately wrong parameters
init_params = {
    'log_epsilon': jnp.log(1.5),   # true: log(1.0) = 0.0
    'log_sigma': jnp.log(1.3),     # true: log(1.0) = 0.0
}

optimizer = optax.adam(learning_rate=0.01)  # Adam: adaptive moment estimation
opt_state = optimizer.init(init_params)

# JIT-compile the loss and gradient computation
# This compiles the entire pipeline: MD + RDF + loss + backprop
loss_and_grad_fn = jax.jit(jax.value_and_grad(rdf_loss))

params = init_params
print(f"\nStarting optimization:")
print(f"  Initial epsilon = {jnp.exp(params['log_epsilon']):.4f} "
      f"(target: {TRUE_EPSILON})")
print(f"  Initial sigma   = {jnp.exp(params['log_sigma']):.4f} "
      f"(target: {TRUE_SIGMA})")

for step in range(80):
    loss, grads = loss_and_grad_fn(
        params, equil_final, equil_vel, box_size,
        r_bins, target_rdf, n_steps=300,
    )
    updates, opt_state = optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)

    if step % 10 == 0:
        eps = jnp.exp(params['log_epsilon'])
        sig = jnp.exp(params['log_sigma'])
        print(f"  Step {step:3d}: loss={loss:.6f}, "
              f"epsilon={eps:.4f}, sigma={sig:.4f}")

# Final parameters
final_eps = jnp.exp(params['log_epsilon'])
final_sig = jnp.exp(params['log_sigma'])
print(f"\nOptimized: epsilon={final_eps:.4f}, sigma={final_sig:.4f}")
print(f"Target:    epsilon={TRUE_EPSILON:.4f}, sigma={TRUE_SIGMA:.4f}")
print(f"Error:     epsilon={abs(final_eps - TRUE_EPSILON):.4f}, "
      f"sigma={abs(final_sig - TRUE_SIGMA):.4f}")
# Typical output:
#   Step   0: loss=0.142356, epsilon=1.5000, sigma=1.3000
#   Step  10: loss=0.058921, epsilon=1.3214, sigma=1.2103
#   Step  20: loss=0.018743, epsilon=1.1876, sigma=1.1245
#   Step  30: loss=0.005621, epsilon=1.0923, sigma=1.0634
#   Step  40: loss=0.001843, epsilon=1.0412, sigma=1.0298
#   Step  50: loss=0.000621, epsilon=1.0187, sigma=1.0132
#   Step  60: loss=0.000198, epsilon=1.0078, sigma=1.0054
#   Step  70: loss=0.000067, epsilon=1.0031, sigma=1.0021
#
# Optimized: epsilon=1.0014, sigma=1.0009
# Target:    epsilon=1.0000, sigma=1.0000
# Error:     epsilon=0.0014, sigma=0.0009

6. Anatomy of the Gradient Flow

The gradient computation inside loss_and_grad_fn walks backward through the chain shown in Figure 42.4:

1. MSE loss \(\to\) gradient w.r.t. the predicted RDF bins.
2. Soft RDF computation \(\to\) gradient w.r.t. pairwise distances at each frame, via the Gaussian kernel derivatives.
3. Distance computation \(\to\) gradient w.r.t. particle positions at each frame.
4. 300 Velocity Verlet steps (backwards) \(\to\) gradient w.r.t. the forces at each step, then w.r.t. the energy function's parameters.
5. LJ energy function \(\to\) gradient w.r.t. \(\epsilon\) and \(\sigma\), accumulated over all pairs and all time steps.

The forward pass evaluates roughly \(300 \times \binom{108}{2} \approx 1.7 \times 10^6\) pairs. The backward pass performs the same number of vector-Jacobian product (VJP) operations. JAX compiles this entire chain into a single Accelerated Linear Algebra (XLA) computation that runs on the GPU without returning to Python between operations. The first call compiles for tens of seconds, but subsequent calls execute in seconds. (That is 1.7 million coupled derivative operations, from loss to force field parameters, completing in a single GPU kernel launch.)

7. Discovery Workbench Integration

The differentiable force field pipeline integrates naturally with the Discovery Workbench architecture introduced in Chapter 6. The force field calibration task becomes a Workbench experiment with tracked parameters, metrics, and artifacts:

from discovery_workbench import Experiment, Parameter, Metric

# Register the force field calibration as a Workbench experiment
experiment = Experiment(
    name="lj_force_field_calibration",
    description="Calibrate LJ parameters to match target RDF "
                "using differentiable MD",
    tags=["differentiable-programming", "force-field", "molecular-dynamics"],
)

# Log hyperparameters
experiment.log_params({
    "n_atoms": N,
    "n_md_steps": 300,
    "n_optimization_steps": 80,
    "learning_rate": 0.01,
    "optimizer": "adam",
    "initial_epsilon": 1.5,
    "initial_sigma": 1.3,
    "target_epsilon": TRUE_EPSILON,
    "target_sigma": TRUE_SIGMA,
})

# Training loop with Workbench logging
params = init_params
for step in range(80):
    loss, grads = loss_and_grad_fn(
        params, equil_final, equil_vel, box_size,
        r_bins, target_rdf,
    )
    updates, opt_state = optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)

    # Log metrics at each step
    experiment.log_metrics({
        "loss": float(loss),
        "epsilon": float(jnp.exp(params['log_epsilon'])),
        "sigma": float(jnp.exp(params['log_sigma'])),
        "grad_norm_epsilon": float(jnp.abs(grads['log_epsilon'])),
        "grad_norm_sigma": float(jnp.abs(grads['log_sigma'])),
    }, step=step)

# Log final RDF comparison as an artifact
experiment.log_artifact("final_rdf_comparison", {
    "r_bins": r_bins.tolist(),
    "target_rdf": target_rdf.tolist(),
    "predicted_rdf": compute_rdf_trajectory(
        run_md_trajectory(equil_final, equil_vel, box_size,
                          jnp.exp(params['log_epsilon']),
                          jnp.exp(params['log_sigma']),
                          n_steps=300),
        box_size, r_bins,
    ).tolist(),
})

experiment.finish()

8. Scaling Up: From Toy to Production

Logging every parameter update and RDF comparison ensures reproducibility, but the 108-particle system we have been tracking is far smaller than the thousands of atoms typical of real materials simulations.

The recipe above works for a small system (108 atoms, 300 steps). Scaling to production-grade force field fitting requires several enhancements:

• Neighbor lists: our \(O(N^2)\) pair enumeration must be replaced with neighbor lists for \(O(N)\) scaling. JAX-MD provides differentiable neighbor lists that handle periodic boundaries and dynamic allocation.
• Gradient checkpointing: for \(> 10^3\) MD steps, use jax.checkpoint on the scan body (as shown in Listing 42.18) to bound memory usage.
• Multiple observables: real force field fitting matches not just the RDF but also the mean potential energy, pressure, diffusion coefficient, and possibly forces from quantum mechanical reference calculations. Each observable adds a term to the loss; gradients through all terms come from a single backward pass.
• Multiple state points: a robust force field must reproduce properties at multiple temperatures and densities. Use vmap across state points to run simulations in parallel and average gradients.
• Neural network potentials: replace the two-parameter LJ function with a neural network (as in Listing 42.11) that takes atomic environment descriptors as input. The same optimization pipeline works; only the energy function and its parameter count change.

9. Extensions and Research Frontiers

Once the scaling bottlenecks are addressed with neighbor lists and gradient checkpointing, the same differentiable pipeline opens the door to problems well beyond two-parameter Lennard-Jones fitting.

Differentiable force field calibration is an active area of research. Several directions extend the ideas in this recipe:

• Differentiable free energy: matching free energies (not just structural observables) requires differentiating through thermodynamic integration or alchemical transformations (computational techniques that gradually morph one molecule into another to estimate free energy differences). Recent work (Wirnsberger et al., 2022) shows this is feasible with normalizing flows.
• Multi-fidelity optimization: combine cheap short simulations (for gradient direction) with expensive long simulations (for accurate observables) in a multi-fidelity optimization scheme. The short simulations provide biased but low-variance gradients; the long simulations correct the bias periodically.
• Equivariant neural potentials: models like NequIP, MACE, and Allegro use \(E(3)\)-equivariant (invariant to rotations, reflections, and translations in 3D space) neural networks to learn potentials that respect rotational symmetry by construction. Training these models through differentiable MD is the current state of the art for machine-learned interatomic potentials.
• Coarse-grained model discovery: fit coarse-grained (simplified representations where groups of atoms are merged into single interaction sites) force fields to reproduce structural and dynamical properties of atomistic simulations, automating the coarse-graining process that traditionally requires expert physical intuition.

These extensions connect to the broader themes of scientific machine learning (Chapter 33), scientific simulation (Chapter 43), and chemistry and materials discovery (Chapter 49).