Geometry optimization (FIRE)

FIRE drives atomic positions toward a local energy minimum using the Fast Inertial Relaxation Engine algorithm, with UPETWrapper supplying forces and ConvergenceHook stopping the run early once the maximum force norm drops below a threshold. The energy and the maximum force are recorded at every step and plotted, so the approach to the minimum is visible at a glance.

Note

This example requires the optional nvalchemi extra: pip install "upet[nvalchemi]".

from collections import defaultdict

import matplotlib.pyplot as plt
import torch
from ase.build import bulk
from nvalchemi.data import AtomicData, Batch
from nvalchemi.dynamics import FIRE, ConvergenceHook, DynamicsStage
from nvalchemi.hooks import NeighborListHook, extract_dynamics_scalars
from nvalchemi.neighbors import compute_neighbors

from upet.nvalchemi import UPETWrapper


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = UPETWrapper.from_checkpoint(model="pet-mad-xs", version="1.6.0", device=device)

atoms = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond")
atoms.rattle(0.1, seed=0)
data = AtomicData.from_atoms(atoms, device=device)

# The integrators write model outputs back into the batch fields that already
# exist and silently skip the ones that don't, while `AtomicData.from_atoms`
# only fills the fields the `Atoms` object itself carries -- so the output
# buffers have to be allocated up front.
data.energy = torch.zeros(1, 1, device=device)
data.forces = torch.zeros_like(data.positions)

batch = Batch.from_data_list([data], device=device)
compute_neighbors(batch, config=model.model_config.neighbor_config)

Recording the trajectory

nvalchemi runs the whole loop inside run(), so anything to be plotted afterwards has to be collected while the loop is running. A hook registered at AFTER_STEP is the nvalchemi counterpart of ASE’s dyn.attach(logger, interval=1). extract_dynamics_scalars() pulls the standard observables (potential energy, max force norm, temperature) off the batch, so the hook itself only has to accumulate them.

class TrajectoryRecorder:
    """Append the per-step scalar observables to a history dict."""

    def __init__(self):
        self.frequency = 1  # fire on every step
        self.stage = DynamicsStage.AFTER_STEP
        self.history: dict[str, list[float]] = defaultdict(list)

    def __call__(self, ctx, stage):
        self.history["step"].append(ctx.step_count)
        for key, value in extract_dynamics_scalars(ctx).items():
            self.history[key].append(value)


FMAX = 0.02  # eV/Å, the convergence threshold

recorder = TrajectoryRecorder()
fire = FIRE(
    model=model,
    dt=0.1,
    n_steps=300,
    convergence_hook=ConvergenceHook.from_fmax(FMAX),
    hooks=[
        NeighborListHook(
            model.model_config.neighbor_config,
            skin=0.5,
            stage=DynamicsStage.BEFORE_COMPUTE,
        ),
        recorder,
    ],
)

# The first integrator step already reads the forces, so the batch needs one
# model evaluation before the loop starts.
fire.compute(batch)
batch = fire.run(batch)

print(f"Relaxed after {fire.step_count} steps")
print(f"Energy : {batch.energy.item():+.4f} eV")
print(f"Fmax   : {torch.linalg.vector_norm(batch.forces, dim=-1).max():.4f} eV/Å")

Relaxed after 52 steps
Energy : -47.2383 eV
Fmax   : 0.0199 eV/Å

Convergence of the relaxation

The energy decreases toward the local minimum while the maximum force norm falls by more than two orders of magnitude. FIRE accelerates along the downhill direction and resets its velocity whenever it overshoots, which is what produces the steps and plateaus in the force curve rather than a smooth exponential decay.

steps = recorder.history["step"]

fig, (ax_energy, ax_force) = plt.subplots(1, 2, figsize=(9.5, 3.8))

ax_energy.plot(steps, recorder.history["energy"], "o-", ms=3)
ax_energy.set_xlabel("FIRE step")
ax_energy.set_ylabel("potential energy [eV]")
ax_energy.set_title("Energy vs. step")

ax_force.semilogy(steps, recorder.history["fmax"], "o-", ms=3)
ax_force.axhline(FMAX, color="k", ls="--", lw=0.8, label=f"threshold ({FMAX} eV/Å)")
ax_force.set_xlabel("FIRE step")
ax_force.set_ylabel("max |force| [eV/Å]")
ax_force.set_title("Force convergence")
ax_force.legend()

fig.tight_layout()
plt.show()