Energy-Derivative Training Contract (Forces, Stress, Charge Gradients)

This example demonstrates the recommended way to obtain forces, virial/stress, and charge gradients from PME (and Ewald) for machine-learning interatomic potential (MLIP) training: energy is the only differentiable output, and all derivatives are taken from it with torch.autograd.grad.

The legacy compute_forces / compute_virial / compute_charge_gradients / hybrid_forces flags on particle_mesh_ewald and ewald_summation are deprecated and emit a DeprecationWarning. They remain available for compatibility in v0.4.0, so this example also checks that the autograd results match the legacy direct outputs, which is the migration check.

In this example you will learn:

• Forces from energy autograd: F = -grad(E.sum(), positions)

• Force-loss training (double-backward) with create_graph=True

• Geometry-dependent charges q(R): the full force includes the dE/dq . dq/dR charge-model chain-rule term

• Strain-first virial/stress and stress-loss training

• Charge gradients from energy autograd: dE/dq = grad(E.sum(), charges)

Important

This script is intended as an API demonstration. Do not use this script for performance benchmarking; refer to the benchmarks folder instead.

Setup and Imports

We use particle_mesh_ewald throughout; ewald_summation follows the exact same energy-derivative contract.

from __future__ import annotations

import warnings

import numpy as np
import torch

from nvalchemiops.torch.interactions.electrostatics import particle_mesh_ewald
from nvalchemiops.torch.neighbors import neighbor_list as neighbor_list_fn

Configure Device

if torch.cuda.is_available():
    device = torch.device("cuda:0")
    print("Using CUDA device")
    print(f"  {torch.cuda.get_device_name(0)}")
else:
    device = torch.device("cpu")
    print("Using CPU")

dtype = torch.float64

Using CUDA device
  NVIDIA H100 80GB HBM3

Create a NaCl Crystal System

A small NaCl rock-salt supercell (2x2x2 = 16 atoms) is enough to demonstrate every derivative path. Charges use the same floating dtype as the geometry.

def create_nacl_system(n_cells: int = 2, lattice_constant: float = 5.64):
    """Create a NaCl crystal supercell."""
    base_positions = np.array([[0.0, 0.0, 0.0], [0.5, 0.5, 0.5]])
    base_charges = np.array([1.0, -1.0])

    positions = []
    charges = []
    for i in range(n_cells):
        for j in range(n_cells):
            for k in range(n_cells):
                offset = np.array([i, j, k])
                for pos, charge in zip(base_positions, base_charges):
                    positions.append((pos + offset) * lattice_constant)
                    charges.append(charge)

    positions = torch.tensor(np.array(positions), dtype=dtype, device=device)
    charges = torch.tensor(np.array(charges), dtype=dtype, device=device)
    cell = torch.eye(3, dtype=dtype, device=device) * lattice_constant * n_cells
    cell = cell.unsqueeze(0)
    pbc = torch.tensor([[True, True, True]], dtype=torch.bool, device=device)
    return positions, charges, cell, pbc


positions, charges, cell, pbc = create_nacl_system(n_cells=2)

# Rattle the perfect lattice so the per-atom forces are nonzero -- otherwise the
# centrosymmetric crystal gives zero forces and the autograd-vs-direct force
# check below would compare 0 to 0.
torch.manual_seed(0)
positions = positions + 0.05 * torch.randn_like(positions)

print(f"\nSystem: {len(positions)} atoms NaCl crystal (rattled)")

# Build the real-space neighbor list once; reuse it for every call below.
neighbor_list, neighbor_ptr, neighbor_shifts = neighbor_list_fn(
    positions, 8.0, cell=cell, pbc=pbc, return_neighbor_list=True
)

# Fixed PME parameters keep the example deterministic and fast.
pme_kwargs = dict(
    alpha=0.35,
    mesh_dimensions=(32, 32, 32),
    spline_order=4,
    neighbor_list=neighbor_list,
    neighbor_ptr=neighbor_ptr,
    neighbor_shifts=neighbor_shifts,
)

System: 16 atoms NaCl crystal (rattled)

Forces From Energy Autograd

With no deprecated flag, the call returns the per-atom energy only. The full force is the negative gradient of the total energy w.r.t. positions.

positions_f = positions.detach().requires_grad_(True)
energy = particle_mesh_ewald(positions_f, charges, cell, **pme_kwargs)
print(f"\nenergy shape: {tuple(energy.shape)}  (per-atom)")

forces = -torch.autograd.grad(energy.sum(), positions_f)[0]
print(f"forces shape: {tuple(forces.shape)}")
print(f"max force magnitude: {forces.norm(dim=1).max().item():.6f}")

energy shape: (16,)  (per-atom)
forces shape: (16, 3)
max force magnitude: 0.005254

Migration check: the autograd force equals the legacy compute_forces=True direct output (the deprecated flag still works, with a DeprecationWarning).

with warnings.catch_warnings():
    warnings.filterwarnings(
        "ignore",
        category=DeprecationWarning,
        message=".*compute_forces.*",
    )
    _, forces_direct = particle_mesh_ewald(
        positions, charges, cell, compute_forces=True, **pme_kwargs
    )

force_diff = (forces - forces_direct).abs().max().item()
print(f"max |autograd force - direct force|: {force_diff:.2e}")
if force_diff >= 1e-6:
    raise RuntimeError(f"autograd force does not match direct force: {force_diff:.2e}")

max |autograd force - direct force|: 0.00e+00

Force-Loss Training (Double-Backward)

To train on a force loss, build the force with create_graph=True so the subsequent loss.backward() can differentiate through the force. We use a trivial linear charge model q(R) to stand in for a learned charge head.

torch.manual_seed(0)
weight = torch.randn(3, dtype=dtype, device=device, requires_grad=True)


def charge_model(pos: torch.Tensor) -> torch.Tensor:
    """Toy geometry-dependent charge model with enforced neutrality."""
    raw = pos @ weight
    return raw - raw.mean()


positions_q = positions.detach().requires_grad_(True)
charges_qr = charge_model(positions_q)  # kept in the graph -> full q(R) force

energy = particle_mesh_ewald(positions_q, charges_qr, cell, **pme_kwargs)
forces_qr = -torch.autograd.grad(energy.sum(), positions_q, create_graph=True)[0]

target_forces = torch.zeros_like(forces_qr)
force_loss = (forces_qr - target_forces).pow(2).sum()
force_loss.backward()  # differentiates through the force construction

print(f"\nq(R) force-loss backward OK; weight.grad shape: {tuple(weight.grad.shape)}")
if not torch.isfinite(weight.grad).all():
    raise RuntimeError("q(R) force-loss produced non-finite parameter gradients")

q(R) force-loss backward OK; weight.grad shape: (3,)

Because charges = charge_model(positions) stays connected to positions, the autograd force includes both the fixed-charge term and the dE/dq . dq/dR chain-rule term. The legacy compute_forces=True output is only the fixed-charge partial and does not include the charge-model term – this is the central reason direct force output on the full API is deprecated.

Strain-First Virial and Stress

strain is not a PME argument. Build a differentiable strain tensor, deform positions and cell by I + strain, and let autograd map gradients back to strain. The virial is W = -dE/d(strain) and tensile-positive stress is dE/d(strain) / V; see conventions.md for the project-wide sign.

num_systems = cell.shape[0]
positions_s = positions.detach().requires_grad_(True)
strain = torch.zeros(num_systems, 3, 3, device=device, dtype=dtype, requires_grad=True)
eye = torch.eye(3, device=device, dtype=dtype).unsqueeze(0)
deform = eye + strain

# Single system: every atom maps to system 0.
batch_idx = torch.zeros(positions_s.shape[0], dtype=torch.int32, device=device)
positions_def = torch.einsum("ni,nij->nj", positions_s, deform[batch_idx])
cell_def = torch.einsum("bij,bjk->bik", cell, deform)

energy = particle_mesh_ewald(positions_def, charges, cell_def, **pme_kwargs)
grad_strain = torch.autograd.grad(energy.sum(), strain)[0]
virial = -grad_strain
volume = torch.abs(torch.linalg.det(cell_def))
stress = grad_strain / volume[:, None, None]  # tensile-positive Cauchy

print(f"\nvirial shape: {tuple(virial.shape)}")
print(f"stress shape: {tuple(stress.shape)}")

virial shape: (1, 3, 3)
stress shape: (1, 3, 3)

Migration check: the strain-first virial equals the legacy compute_virial=True direct virial (both are -dE/d(strain)).

with warnings.catch_warnings():
    warnings.filterwarnings(
        "ignore",
        category=DeprecationWarning,
        message=".*compute_virial.*",
    )
    _, virial_direct = particle_mesh_ewald(
        positions, charges, cell, compute_virial=True, **pme_kwargs
    )

virial_diff = (virial - virial_direct).abs().max().item()
print(f"max |strain-first virial - direct virial|: {virial_diff:.2e}")
if virial_diff >= 1e-5:
    raise RuntimeError(
        f"strain-first virial does not match direct virial: {virial_diff:.2e}"
    )

max |strain-first virial - direct virial|: 1.11e-16

Stress-Loss Training (Double-Backward)

Stress training uses the same strain-first recipe with create_graph=True, so the stress loss back-propagates to model parameters.

weight_s = torch.randn(3, dtype=dtype, device=device, requires_grad=True)
positions_s = positions.detach().requires_grad_(True)
strain = torch.zeros(num_systems, 3, 3, device=device, dtype=dtype, requires_grad=True)
deform = torch.eye(3, device=device, dtype=dtype).unsqueeze(0) + strain
positions_def = torch.einsum("ni,nij->nj", positions_s, deform[batch_idx])
cell_def = torch.einsum("bij,bjk->bik", cell, deform)

charges_qr = positions_def @ weight_s
charges_qr = charges_qr - charges_qr.mean()

energy = particle_mesh_ewald(positions_def, charges_qr, cell_def, **pme_kwargs)
grad_strain = torch.autograd.grad(energy.sum(), strain, create_graph=True)[0]
virial = -grad_strain
volume = torch.abs(torch.linalg.det(cell_def))
stress = grad_strain / volume[:, None, None]

stress_loss = stress.pow(2).sum()
stress_loss.backward()

stress_grad_ok = bool(torch.isfinite(weight_s.grad).all())
print(f"\nstress-loss backward OK; weight_s.grad finite: {stress_grad_ok}")
if not stress_grad_ok:
    raise RuntimeError("stress-loss produced non-finite parameter gradients")

stress-loss backward OK; weight_s.grad finite: True

Combined Energy + Force + Stress Loss (Performance Tip)

When a single loss mixes energy, forces, AND stress, take the forces and the virial from one torch.autograd.grad call over both positions and strain – NOT two separate calls. Each create_graph=True grad call builds its own first-derivative graph node, and loss.backward() then runs the (O(K*N)) reciprocal second-derivative once per node. Fusing the two into a single call avoids duplicate reciprocal double-backward work; the gradients are bit-identical to the two-call form.

weight_m = torch.randn(3, dtype=dtype, device=device, requires_grad=True)
positions_m = positions.detach().requires_grad_(True)
strain_m = torch.zeros(
    num_systems, 3, 3, device=device, dtype=dtype, requires_grad=True
)
deform_m = torch.eye(3, device=device, dtype=dtype).unsqueeze(0) + strain_m
positions_md = torch.einsum("ni,nij->nj", positions_m, deform_m[batch_idx])
cell_md = torch.einsum("bij,bjk->bik", cell, deform_m)
charges_m = positions_md @ weight_m
charges_m = charges_m - charges_m.mean()

energy = particle_mesh_ewald(positions_md, charges_m, cell_md, **pme_kwargs)
# ONE combined grad call -> one double-backward (do this instead of separate
# ``grad(E, positions)`` and ``grad(E, strain)`` calls):
grad_pos, grad_strain = torch.autograd.grad(
    energy.sum(), (positions_m, strain_m), create_graph=True
)
# ``positions_m`` are the undeformed reference coordinates, so ``forces_m`` are
# reference-frame forces. Differentiate with respect to ``positions_md`` instead
# when training against deformed-coordinate force targets.
forces_m = -grad_pos
virial_m = -grad_strain
volume_m = torch.abs(torch.linalg.det(cell_md))
stress_m = grad_strain / volume_m[:, None, None]

mixed_loss = energy.sum() + forces_m.pow(2).sum() + stress_m.pow(2).sum()
mixed_loss.backward()
mixed_grad_ok = bool(torch.isfinite(weight_m.grad).all())
print(f"\nmixed E+F+stress backward OK; weight_m.grad finite: {mixed_grad_ok}")
if not mixed_grad_ok:
    raise RuntimeError("mixed-loss produced non-finite parameter gradients")

mixed E+F+stress backward OK; weight_m.grad finite: True

Charge Gradients From Energy Autograd

dE/dq is an ordinary gradient of the energy w.r.t. charges.

charges_g = charges.detach().requires_grad_(True)
energy = particle_mesh_ewald(positions, charges_g, cell, **pme_kwargs)
charge_grad = torch.autograd.grad(energy.sum(), charges_g)[0]
print(f"\ncharge gradient shape: {tuple(charge_grad.shape)}")

with warnings.catch_warnings():
    warnings.filterwarnings(
        "ignore",
        category=DeprecationWarning,
        message=".*compute_forces.*",
    )
    _, _, charge_grad_direct = particle_mesh_ewald(
        positions,
        charges,
        cell,
        compute_forces=True,
        compute_charge_gradients=True,
        **pme_kwargs,
    )

cg_diff = (charge_grad - charge_grad_direct).abs().max().item()
print(f"max |autograd dE/dq - direct dE/dq|: {cg_diff:.2e}")
if cg_diff >= 1e-6:
    raise RuntimeError(f"autograd dE/dq does not match direct dE/dq: {cg_diff:.2e}")

charge gradient shape: (16,)
max |autograd dE/dq - direct dE/dq|: 1.80e-09

Summary

This example demonstrated the energy-derivative training contract:

1. Forces – F = -torch.autograd.grad(E.sum(), positions)[0]; with create_graph=True for force-loss training.

2. q(R) forces – keep charges = charge_model(positions) in the graph so the full dE/dR includes the dE/dq . dq/dR charge-model term.

3. Virial / stress – strain-first: deform positions and cell by I + strain, then virial = -grad(E.sum(), strain) and stress = grad(E.sum(), strain) / volume[:, None, None].

4. Charge gradients – dE/dq = torch.autograd.grad(E.sum(), charges)[0].

5. Combined E + F + stress loss – take forces and virial from a single torch.autograd.grad(E.sum(), (positions, strain), create_graph=True) call, not two separate calls, so the reciprocal double-backward runs once.

Each autograd result matched the corresponding (deprecated) direct kernel output, confirming the migration is numerically exact.

print("\nEnergy-derivative training contract example complete!")

Energy-derivative training contract example complete!