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/metatomic_ase/_neighbors.py:78: UserWarning: `compute_requested_neighbors_from_options` is deprecated and will be removed in a future version. Please use `neighbor_lists_for_model` to get the calculators and call them directly.
  vesin.metatomic.compute_requested_neighbors_from_options(
       Step     Time          Energy          fmax
LBFGS:    0 15:48:54      -45.099880        3.402630
LBFGS:    1 15:48:54      -45.548882        2.344578
LBFGS:    2 15:48:54      -46.199650        0.816876
LBFGS:    3 15:48:54      -46.238434        0.565600
LBFGS:    4 15:48:54      -46.264099        0.571426
LBFGS:    5 15:48:54      -46.289146        0.468049
LBFGS:    6 15:48:54      -46.307289        0.326139
LBFGS:    7 15:48:54      -46.316322        0.298759
LBFGS:    8 15:48:54      -46.325405        0.241907
LBFGS:    9 15:48:54      -46.333160        0.176211
LBFGS:   10 15:48:54      -46.335651        0.084599
LBFGS:   11 15:48:54      -46.336025        0.042398
       Step     Time          Energy          fmax
LBFGS:    0 15:48:54      -46.336025        1.307608
LBFGS:    1 15:48:54      -46.407185        1.285505
LBFGS:    2 15:48:54      -46.979885        0.576497
LBFGS:    3 15:48:54      -47.144047        0.084707
LBFGS:    4 15:48:54      -47.144936        0.064516
LBFGS:    5 15:48:54      -47.145432        0.047223

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.calculator 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.5.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()