Geometry optimization (LBFGS)

Two-stage geometry optimization of a slightly perturbed silicon unit cell. First, positions are relaxed at fixed cell with ase.optimize.LBFGS. Then the cell itself is relaxed jointly with the atomic positions by wrapping the Atoms in ase.filters.FrechetCellFilter.

The script records the maximum force and total energy at every step and plots them so that the convergence of the two stages is visible at a glance.

/home/runner/work/upet/upet/.tox/docs/lib/python3.13/site-packages/torch/jit/_script.py:1491: FutureWarning: `torch.jit.script` is deprecated. Please switch to `torch.compile` or `torch.export`.
  warnings.warn(
/home/runner/work/upet/upet/.tox/docs/lib/python3.13/site-packages/torch/jit/_script.py:1491: FutureWarning: `torch.jit.script` is deprecated. Please switch to `torch.compile` or `torch.export`.
  warnings.warn(
       Step     Time          Energy          fmax
LBFGS:    0 11:21:20      -45.148903        3.371936
LBFGS:    1 11:21:20      -45.593739        2.373620
LBFGS:    2 11:21:20      -46.253155        0.850057
LBFGS:    3 11:21:20      -46.289375        0.575552
LBFGS:    4 11:21:20      -46.315495        0.521655
LBFGS:    5 11:21:20      -46.337902        0.390082
LBFGS:    6 11:21:20      -46.348885        0.287841
LBFGS:    7 11:21:20      -46.356419        0.264732
LBFGS:    8 11:21:20      -46.365788        0.258870
LBFGS:    9 11:21:20      -46.373112        0.179407
LBFGS:   10 11:21:21      -46.375755        0.094185
LBFGS:   11 11:21:21      -46.376369        0.057819
LBFGS:   12 11:21:21      -46.376656        0.037955
       Step     Time          Energy          fmax
LBFGS:    0 11:21:21      -46.376656        1.457212
LBFGS:    1 11:21:21      -46.462555        1.391114
LBFGS:    2 11:21:21      -47.069180        0.695960
LBFGS:    3 11:21:21      -47.233089        0.137012
LBFGS:    4 11:21:21      -47.238922        0.059607
LBFGS:    5 11:21:21      -47.239189        0.039990

import matplotlib.pyplot as plt
import numpy as np
from ase.build import bulk
from ase.filters import FrechetCellFilter
from ase.optimize import LBFGS

from upet.ase import UPETCalculator


atoms = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond")

# perturb positions and cell so the optimizer has something to do
atoms.rattle(0.1, seed=0)  # ASE's built-in random displacement method
atoms.set_cell(atoms.cell * 1.05, scale_atoms=True)

calculator = UPETCalculator(model="pet-mad-xs", version="1.6.0", device="cpu")
atoms.calc = calculator

history = {"stage": [], "energy": [], "fmax": []}  # type: ignore


def record(stage_name):
    def _cb():
        results = calculator.results
        history["stage"].append(stage_name)
        history["energy"].append(float(results["energy"]))
        history["fmax"].append(float(np.linalg.norm(results["forces"], axis=1).max()))

    return _cb


# stage 1: positions only
opt_pos = LBFGS(atoms)
opt_pos.attach(record("positions"), interval=1)
opt_pos.run(fmax=0.05, steps=30)

# stage 2: joint position + cell relaxation
filtered = FrechetCellFilter(atoms)
opt_cell = LBFGS(filtered)
opt_cell.attach(record("cell"), interval=1)
opt_cell.run(fmax=0.05, steps=30)

steps = np.arange(len(history["energy"]))
stages = np.array(history["stage"])
boundary = (
    int(np.searchsorted(stages == "cell", True))
    if (stages == "cell").any()
    else len(stages)
)

fig, (ax_e, ax_f) = plt.subplots(1, 2, figsize=(9, 3.5))
ax_e.plot(steps, history["energy"], "o-")
ax_e.axvline(boundary - 0.5, color="k", ls="--", lw=0.8)
ax_e.set_xlabel("optimization step")
ax_e.set_ylabel("total energy [eV]")
ax_e.set_title("Energy vs. step")

ax_f.semilogy(steps, history["fmax"], "o-")
ax_f.axvline(boundary - 0.5, color="k", ls="--", lw=0.8)
ax_f.set_xlabel("optimization step")
ax_f.set_ylabel("max |force| [eV/Å]")
ax_f.set_title("Force convergence")

fig.tight_layout()
plt.show()