.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "generated_examples/2-nvalchemi/run_nva_relaxation.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_generated_examples_2-nvalchemi_run_nva_relaxation.py: Geometry optimization (FIRE) =============================================== :py:class:`~nvalchemi.dynamics.FIRE` drives atomic positions toward a local energy minimum using the Fast Inertial Relaxation Engine algorithm, with :py:class:`~upet.nvalchemi.UPETWrapper` supplying forces and :py:class:`~nvalchemi.dynamics.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]"``. .. GENERATED FROM PYTHON SOURCE LINES 18-49 .. code-block:: Python 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) .. GENERATED FROM PYTHON SOURCE LINES 50-59 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 :py:attr:`~nvalchemi.dynamics.DynamicsStage.AFTER_STEP` is the nvalchemi counterpart of ASE's ``dyn.attach(logger, interval=1)``. :py:func:`~nvalchemi.hooks.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. .. GENERATED FROM PYTHON SOURCE LINES 59-102 .. code-block:: Python 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/Å") .. rst-class:: sphx-glr-script-out .. code-block:: none Relaxed after 52 steps Energy : -47.2383 eV Fmax : 0.0199 eV/Å .. GENERATED FROM PYTHON SOURCE LINES 103-110 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. .. GENERATED FROM PYTHON SOURCE LINES 110-129 .. code-block:: Python 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() .. image-sg:: /generated_examples/2-nvalchemi/images/sphx_glr_run_nva_relaxation_001.png :alt: Energy vs. step, Force convergence :srcset: /generated_examples/2-nvalchemi/images/sphx_glr_run_nva_relaxation_001.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 19.990 seconds) .. _sphx_glr_download_generated_examples_2-nvalchemi_run_nva_relaxation.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: run_nva_relaxation.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: run_nva_relaxation.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: run_nva_relaxation.zip ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_