.. 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_basics.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_basics.py: Basics: single-structure evaluation =================================== :py:class:`~upet.nvalchemi.UPETWrapper` wraps a UPET / PET-MAD checkpoint as an `nvalchemi-toolkit `_ ``BaseModelMixin`` model, so it can be driven through nvalchemi's batched :py:class:`~nvalchemi.data.Batch` data pipeline instead of ASE ``Atoms``. This is the entry point for large-scale batched inference (see :doc:`run_nva_batched_eval`) and for nvalchemi's GPU-accelerated MD integrators (see :doc:`run_nva_md_nvt`). This example builds a bulk silicon cell as an ASE ``Atoms`` object, converts it to a single-graph :py:class:`~nvalchemi.data.Batch` with :py:meth:`~nvalchemi.data.AtomicData.from_atoms`, and evaluates energy, forces, and stress with :py:class:`~upet.nvalchemi.UPETWrapper`. Both tensorial outputs are then visualized: the forces as arrows on the projected cell, and the Cauchy stress as an annotated 3x3 map. .. note:: This example requires the optional ``nvalchemi`` extra: ``pip install "upet[nvalchemi]"``. .. GENERATED FROM PYTHON SOURCE LINES 25-37 .. code-block:: Python import matplotlib.pyplot as plt import numpy as np import torch from ase.build import bulk from ase.visualize.plot import plot_atoms from nvalchemi.data import AtomicData, Batch from nvalchemi.neighbors import compute_neighbors from upet.nvalchemi import UPETWrapper .. GENERATED FROM PYTHON SOURCE LINES 38-43 Loading a checkpoint -------------------- :py:meth:`~upet.nvalchemi.UPETWrapper.from_checkpoint` fetches a named UPET model from HuggingFace (here the small ``pet-mad-xs`` model), or loads a local checkpoint file when ``checkpoint_path`` is given instead. .. GENERATED FROM PYTHON SOURCE LINES 43-48 .. code-block:: Python 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) .. GENERATED FROM PYTHON SOURCE LINES 49-56 From ASE ``Atoms`` to a single-graph ``Batch`` ------------------------------------------------ :py:meth:`~nvalchemi.data.AtomicData.from_atoms` builds an :py:class:`~nvalchemi.data.AtomicData` instance directly from an ASE ``Atoms`` object, carrying over positions, atomic numbers, cell and PBC. :py:meth:`~nvalchemi.data.Batch.from_data_list` then promotes it to a single-graph batch. .. GENERATED FROM PYTHON SOURCE LINES 56-66 .. code-block:: Python atoms = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond") # Displace every atom slightly off its equilibrium site, so the predicted # forces are non-zero and of comparable magnitude across the cell. atoms.rattle(0.05, seed=0) data = AtomicData.from_atoms(atoms, device=device) batch = Batch.from_data_list([data], device=device) .. GENERATED FROM PYTHON SOURCE LINES 67-73 Neighbor list and evaluation ------------------------------ :py:func:`~nvalchemi.neighbors.compute_neighbors` is the one-shot convenience function for populating a batch's neighbor list outside a dynamics loop; ``model.model_config.neighbor_config`` already encodes the cutoff and list format the model expects. .. GENERATED FROM PYTHON SOURCE LINES 73-81 .. code-block:: Python compute_neighbors(batch, config=model.model_config.neighbor_config) outputs = model(batch) print(f"Energy : {outputs['energy'].item():+.4f} eV") print(f"Forces :\n{outputs['forces'].detach().cpu().numpy()}") print(f"Stress :\n{outputs['stress'].squeeze(0).detach().cpu().numpy()}") .. rst-class:: sphx-glr-script-out .. code-block:: none Energy : -46.8500 eV Forces : [[-0.6098818 0.79013574 -0.30601698] [-0.55659235 -1.0620064 0.96056247] [ 0.03621662 0.81610477 -0.3269433 ] [-0.22019836 0.2661117 -1.3370566 ] [-0.12413108 0.15760766 -0.23730144] [ 1.1662254 -1.7095635 -0.4695834 ] [ 0.5255206 1.4886923 1.489167 ] [-0.2171591 -0.74708265 0.22717234]] Stress : [[-0.00786457 0.00477575 0.0163688 ] [ 0.00477575 -0.00465882 -0.0008237 ] [ 0.0163688 -0.0008237 -0.00593601]] .. GENERATED FROM PYTHON SOURCE LINES 82-91 Visualizing the forces and the stress tensor ---------------------------------------------- ``AtomicData.from_atoms`` keeps the atom ordering of the ``Atoms`` object, so the force rows line up with ``atoms`` and can be drawn straight onto a projection of the cell. The left panel looks down the ``z`` axis and overlays the in-plane force components as arrows; the right panel shows the full Cauchy stress tensor, whose off-diagonal (shear) components are non-zero because the rattle breaks the cubic symmetry of the ideal diamond cell. .. GENERATED FROM PYTHON SOURCE LINES 91-133 .. code-block:: Python positions = atoms.get_positions() forces = outputs["forces"].detach().cpu().numpy() stress = outputs["stress"].squeeze(0).detach().cpu().numpy() # Scale the arrows so the largest one spans ~1.5 Å on the plot, whatever # the force magnitudes happen to be. arrow_scale = np.linalg.norm(forces[:, :2], axis=1).max() / 1.5 fig, (ax_forces, ax_stress) = plt.subplots(1, 2, figsize=(9.5, 4.2)) plot_atoms(atoms, ax_forces, radii=0.6, show_unit_cell=2) ax_forces.quiver( positions[:, 0], positions[:, 1], forces[:, 0], forces[:, 1], color="tab:red", angles="xy", scale_units="xy", scale=arrow_scale, width=0.007, ) ax_forces.set_title("Forces projected along z") ax_forces.set_xlabel("x [Å]") ax_forces.set_ylabel("y [Å]") limit = np.abs(stress).max() image = ax_stress.imshow(stress, cmap="RdBu_r", vmin=-limit, vmax=limit) for i in range(3): for j in range(3): ax_stress.text( j, i, f"{stress[i, j]:+.1e}", ha="center", va="center", fontsize=9 ) labels = ["x", "y", "z"] ax_stress.set_xticks(range(3), labels) ax_stress.set_yticks(range(3), labels) ax_stress.set_title("Cauchy stress") fig.colorbar(image, ax=ax_stress, label="stress [eV/ų]", fraction=0.046) fig.tight_layout() plt.show() .. image-sg:: /generated_examples/2-nvalchemi/images/sphx_glr_run_nva_basics_001.png :alt: Forces projected along z, Cauchy stress :srcset: /generated_examples/2-nvalchemi/images/sphx_glr_run_nva_basics_001.png :class: sphx-glr-single-img .. _sphx_glr_download_generated_examples_2-nvalchemi_run_nva_basics.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_basics.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: run_nva_basics.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: run_nva_basics.zip ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_