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.

Monopole Ewald/PME/slab entry points accept energy_reduction="atom" (default) or "system". Atom mode returns per-atom energies (N,); system mode returns per-system totals (B,) for batched per-system losses. Direct-output fields (forces, charge gradients, virials) keep their existing shapes. On Torch CUDA, eager atom mode may synchronize when proving a materialized uniform cotangent; energy_reduction="system" is structurally sync-free.

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)

• energy_reduction="system" for batched per-system energy layout (B,)

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 PCIe

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, energy_reduction='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, energy_reduction='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.

6. Batched per-system layout – pass energy_reduction="system" to get (B,) energies for per-system losses without manual scatter_add.

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!