{ "cells": [ { "cell_type": "raw", "metadata": { "raw_mimetype": "text/restructuredtext", "vscode": { "languageId": "raw" } }, "source": [ "######\n", "Theory\n", "######\n", "\n", "An (uncharged, ground-state) atomic structure containing :math:`N` atoms is completely defined by:\n", "\n", "* the positions of its atoms, :math:`\\vec{R} \\in \\mathbb{R}^{N \\times 3}`\n", "* their atomic numbers, :math:`{Z} \\in \\mathbb{Z}^N`\n", "* the unit cell, :math:`C \\in \\mathbb{R}^{3 \\times 3}` (if the structure is periodic)\n", "\n", "Energy\n", "======\n", "\n", "Since these three properties fully define the structure, it must be that the total energy, :math:`E \\in \\mathbb{R}`,\n", "can be expressed solely as a function of these three properties:\n", "\n", ".. math::\n", "\n", " E = f\\left(\\vec{R}, Z, C\\right)\n", "\n", "Forces\n", "======\n", "\n", "The force on atom :math:`i`, :math:`\\vec{F}_i \\in \\mathbb{R}^3`, is given by the\n", "negative gradient of the energy with respect to that atom's position:\n", "\n", ".. math::\n", "\n", " \\vec{F}_i = -\\frac{\\partial E}{\\partial \\vec{R}_i}\n", "\n", "By using :class:`torch.Tensor` representations of :math:`\\vec{R}`, :math:`Z`, and :math:`C`, and \n", "ensuring that the energy function, :math:`f`, makes use of ``torch`` operations, we can leverage\n", "automatic differentiation, supplied by :func:`torch.autograd.grad`, to calculate the forces on the atoms in a structure\n", "\"for free\" from any energy prediction.\n", "\n", "Stress\n", "======\n", "\n", "Consider scaling a structure, that is \"stretching\" both the atomic positions and unit cell, by some amount,\n", ":math:`1 + \\lambda`, along the :math:`x` direction. This operation, :math:`\\hat{O}_{\\lambda}`, acts \n", "on the atomic positions, :math:`\\vec{R}`, according to:\n", "\n", ".. math::\n", "\n", " \\hat{O}_{\\lambda} \\left(\\vec{R}\\right) = \\vec{R} \\times \\begin{pmatrix}\n", " 1 + \\lambda & 0 & 0 \\\\\n", " 0 & 1 & 0 \\\\\n", " 0 & 0 & 1\n", " \\end{pmatrix} = \\vec{R} + \\lambda R_x \\begin{pmatrix}\n", " 1 \\\\\n", " 0 \\\\\n", " 0\n", " \\end{pmatrix}\n", "\n", "and analogously for the structure's unit cell, :math:`C`.\n", "The response of the energy to this transformation gives an indication as to the stress acting on the structure. If \n", "\n", ".. math::\n", " \\frac{\\partial E\\left[\\hat{O}_{\\lambda}(\\vec{R}), Z, \\hat{O}_{\\lambda}(C)\\right]}{\\partial \\lambda} \\bigg|_{\\lambda=0} =\n", " \\frac{\\partial E}{\\partial \\lambda} \\bigg|_{\\lambda=0} < 0 \n", "\n", "then the energy decreases as the unit cell expands from its current state. \n", "This would indicate that the system is under *compressive* stress (along the :math:`x` direction) and \"wants\" to expand.\n", "\n", "Now consider a more general scaling operation, :math:`\\hat{\\mathbf{O}}_{\\mathbf{\\lambda}}`, that symmetrically scales both the atomic positions and unit cell as:\n", "\n", ".. math::\n", "\n", " \\begin{aligned}\n", " \\hat{\\mathbf{O}}_{\\mathbf{\\lambda}} \\left(\\vec{R}\\right) &= \\vec{R} \\times \\begin{pmatrix}\n", " 1 + \\lambda_{xx} & \\lambda_{xy} & \\lambda_{xz} \\\\\n", " \\lambda_{yx} & 1 + \\lambda_{yy} & \\lambda_{yz} \\\\\n", " \\lambda_{zx} & \\lambda_{zy} & 1 + \\lambda_{zz}\n", " \\end{pmatrix} \\\\\n", " &= \\vec{R} + \\vec{R} \\times \\begin{pmatrix}\n", " \\lambda_{xx} & \\lambda_{xy} & \\lambda_{xz} \\\\\n", " \\lambda_{yx} & \\lambda_{yy} & \\lambda_{yz} \\\\\n", " \\lambda_{zx} & \\lambda_{zy} & \\lambda_{zz}\n", " \\end{pmatrix}\n", " \\end{aligned}\n", "\n", "where, due to the symmetry of the expansion, :math:`\\lambda_{ij} = \\lambda_{ji} \\quad \\forall \\; i \\neq j \\in \\{x,y,z\\}`.\n", "\n", "The diagonal terms of this matrix again correspond to the compressive/dilative stress along each of the Cartesian axes.\n", "\n", "The **off-diagonal terms** describe the shear stress, *i.e.* the tendency of the structure to slide in one plane relative to another.\n", "\n", "In ``graph-pes``, we follow the common definition of the **stress tensor**, :math:`\\mathbf{\\sigma} \\in \\mathbb{R}^{3 \\times 3}`, as the derivative\n", "of the total energy with respect to these stretching coefficients, as normalised by the cell's volume, :math:`V = \\det(\\mathbf{C})`: [1]_\n", "\n", ".. math::\n", "\n", " \\mathbf{\\sigma} = \\frac{1}{V} \\frac{\\partial E}{\\partial \\mathbf{\\lambda}} \\bigg|_{\\mathbf{\\lambda} = 0} \n", " \\quad \\quad \n", " \\sigma_{ij} = \\frac{1}{V} \\frac{\\partial E}{\\partial \\lambda_{ij}} \\bigg|_{\\lambda_{ij} = 0}\n", "\n", "We can again make use of automatic differentiation to calculate these stress tensors. To do this, we:\n", "\n", "1. define a symmetrized :math:`\\mathbf{\\lambda} = 0^{3 \\times 3}`. This is all zeros since we \n", " - don't want to actually change the atomic positions or unit cell for the energy calculation\n", " - want to evaluate the derivative at :math:`\\mathbf{\\lambda} = 0`\n", "2. apply the scaling operation, :math:`\\hat{\\mathbf{O}}_{\\mathbf{\\lambda}}`, to the atomic positions and unit cell\n", " - again this is a no-op due to evaluating the scaling operation at :math:`\\mathbf{\\lambda} = 0`, but introduces the scaling coefficients into the computational graph\n", "3. evaluate the energy\n", "4. calculate the derivative of the energy with respect to :math:`\\mathbf{\\lambda}` and normalise by the cell's volume\n", "\n", "Interpretation\n", "--------------\n", "\n", "The sign of :math:`\\sigma_{xx}` contains useful information:\n", "\n", "If :math:`\\sigma_{xx} < 0`, then the total energy decreases (i.e. the structure gets more stable) as the unit cell expands along the :math:`x` direction. This indicates that the system is under *compressive* stress (along the :math:`x` direction) and \"wants\" to expand.\n", "\n", "Similarly, if :math:`\\sigma_{xx} > 0`, then the system is under *tensile* stress (along the :math:`x` direction) and \"wants\" to contract.\n", "\n", "Virial\n", "======\n", "\n", "In ``graph-pes``, we follow the common definition of the **virial stress tensor**, :math:`\\mathbf{W} \\in \\mathbb{R}^{3 \\times 3}`, as:\n", "\n", ".. math::\n", "\n", " \\begin{aligned}\n", " \\mathbf{W} &= - \\frac{\\partial E}{\\partial \\mathbf{\\lambda}} \\bigg|_{\\mathbf{\\lambda} = 0} \\\\\n", " & = - \\text{stress} \\times V\n", " \\end{aligned}\n", "\n", "i.e. as the negative of the stress tensor scaled by the cell's volume. Hence, the virial stress tensor is an extensive property, while the stress tensor is an intensive property.\n", "\n", ".. [1] F. Knuth et al. `All-electron formalism for total energy strain\n", " derivatives and stress tensor components for numeric atom-centered\n", " orbitals `__.\n", " Computer Physics Communications 190, 33–50 (2015).\n" ] }, { "cell_type": "raw", "metadata": { "raw_mimetype": "text/restructuredtext", "vscode": { "languageId": "raw" } }, "source": [ "Correctness\n", "=============\n", "\n", "Below we (empirically) show the correctness of the :meth:`~graph_pes.GraphPESModel.predict` implementation on :class:`~graph_pes.GraphPESModel` instances by comparing:\n", "\n", "1. predictions from the inbuilt :class:`graph_pes.models.LennardJones` model against those from ASE's :class:`~ase.calculators.lj.LennardJones` calculator.\n", "2. ``\"force\"`` and ``\"stress\"`` predictions from arbitrary :class:`~graph_pes.GraphPESModel` instances against finite-difference approximations of those quantities.\n", "\n", "Lennard-Jones\n", "-------------" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "0.643428299130453" ] }, "execution_count": 2, "metadata": {}, "output_type": "execute_result" } ], "source": [ "from __future__ import annotations\n", "\n", "import ase\n", "from ase.calculators.lj import LennardJones as ASE_LennardJones\n", "from graph_pes.models import LennardJones as GraphPES_LennardJones\n", "from graph_pes.utils.calculator import GraphPESCalculator\n", "\n", "lj_properties = {\n", " \"rc\": 2.0,\n", " \"sigma\": 1.0,\n", " \"epsilon\": 1.0,\n", " \"smooth\": False,\n", "}\n", "\n", "ase_lj_calc = ASE_LennardJones(**lj_properties)\n", "graph_pes_lj_calc = GraphPESCalculator(\n", " model=GraphPES_LennardJones.from_ase(**lj_properties)\n", ")\n", "\n", "dimer = ase.Atoms(\n", " positions=[\n", " [0, 0, 0],\n", " [0, 0, 0.98],\n", " ]\n", ")\n", "\n", "ase_lj_calc.get_potential_energy(dimer)" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "0.643427848815918" ] }, "execution_count": 3, "metadata": {}, "output_type": "execute_result" } ], "source": [ "graph_pes_lj_calc.get_potential_energy(dimer)" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "array([[ 0. , 0. , -34.77113701],\n", " [ 0. , 0. , 34.77113701]])" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "ase_lj_calc.get_forces(dimer)" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "array([[ -0. , -0. , -34.771133],\n", " [ -0. , -0. , 34.771133]], dtype=float32)" ] }, "execution_count": 6, "metadata": {}, "output_type": "execute_result" } ], "source": [ "graph_pes_lj_calc.get_forces(dimer)" ] }, { "cell_type": "code", "execution_count": 20, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "array([[-18.039, -0. , 0. ],\n", " [ -0. , -18.039, 0. ],\n", " [ 0. , 0. , -18.039]])" ] }, "execution_count": 20, "metadata": {}, "output_type": "execute_result" } ], "source": [ "import numpy as np\n", "from ase.stress import voigt_6_to_full_3x3_stress\n", "\n", "\n", "def get_stress_matrix(calculator, structure):\n", " stress = calculator.get_stress(structure)\n", " if stress.shape != (3, 3):\n", " stress = voigt_6_to_full_3x3_stress(stress)\n", " return np.round(stress, 3)\n", "\n", "\n", "cell_size = 1\n", "periodic_structure = ase.Atoms(\n", " positions=[\n", " [cell_size / 2, cell_size / 2, cell_size / 2],\n", " ],\n", " cell=[cell_size, cell_size, cell_size],\n", " pbc=True,\n", ")\n", "\n", "get_stress_matrix(ase_lj_calc, periodic_structure)" ] }, { "cell_type": "code", "execution_count": 21, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "array([[-18.039, 0. , 0. ],\n", " [ 0. , -18.039, -0. ],\n", " [ 0. , -0. , -18.039]])" ] }, "execution_count": 21, "metadata": {}, "output_type": "execute_result" } ], "source": [ "get_stress_matrix(graph_pes_lj_calc, periodic_structure)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Finite-differences\n", "\n", "The force on atom $i$ along direction $d$ is given by:\n", "\n", "$$\n", "F^{(d)}_i = -\\frac{\\partial E}{\\partial R^{(d)}_i}\n", "$$\n", "\n", "which we can estimate by perturbing $R^{(d)}_i$ and calculating the\n", "finite-difference:\n", "\n", "$$\n", "F^{(d)}_i \\approx \\frac{E\\left(R^{(d)}_i + \\epsilon\\right) - E\\left(R^{(d)}_i - \\epsilon\\right)}{2\\epsilon}\n", "$$\n", "\n", "We can also estimate the stress tensor in a simlar fashion.\n" ] }, { "cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [], "source": [ "from typing import Callable, Literal\n", "\n", "\n", "def finite_difference_force_estimate(\n", " energy_function: Callable[[ase.Atoms], float],\n", " atoms: ase.Atoms,\n", " atom_index: int,\n", " direction: Literal[\"x\", \"y\", \"z\"],\n", " epsilon: float = 1e-4,\n", ") -> float:\n", " direction_index = {\"x\": 0, \"y\": 1, \"z\": 2}[direction]\n", "\n", " copy1 = atoms.copy()\n", " copy1.positions[atom_index][direction_index] += epsilon\n", " copy2 = atoms.copy()\n", " copy2.positions[atom_index][direction_index] -= epsilon\n", "\n", " f1 = energy_function(copy1)\n", " f2 = energy_function(copy2)\n", " return (f2 - f1) / (2 * epsilon)\n", "\n", "\n", "def finite_difference_stress_component(\n", " energy_function: Callable[[ase.Atoms], float],\n", " atoms: ase.Atoms,\n", " component: tuple[int, int],\n", " epsilon: float = 1e-6,\n", ") -> float:\n", " i, j = component\n", "\n", " scaling = np.eye(3)\n", " if i == j:\n", " scaling[i, j] += epsilon\n", " else:\n", " scaling[i, j] = scaling[j, i] = epsilon / 2\n", "\n", " copy = atoms.copy()\n", " copy.positions = atoms.positions @ scaling\n", " copy.cell = atoms.cell @ scaling\n", "\n", " f1 = energy_function(copy)\n", " f2 = energy_function(atoms)\n", " return (f1 - f2) / epsilon / atoms.get_volume()\n", "\n", "\n", "def finite_difference_stress_estimate(\n", " energy_function: Callable[[ase.Atoms], float],\n", " atoms: ase.Atoms,\n", " epsilon: float = 1e-6,\n", ") -> np.ndarray:\n", " stress = np.zeros((3, 3))\n", " for i in range(3):\n", " for j in range(3):\n", " stress[i, j] = finite_difference_stress_component(\n", " energy_function, atoms, (i, j), epsilon\n", " )\n", " return stress" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Lets check that our autograd-based implementation matches the finite-difference estimates:" ] }, { "cell_type": "code", "execution_count": 22, "metadata": {}, "outputs": [], "source": [ "from graph_pes.models import SchNet\n", "\n", "model = SchNet(cutoff=3.0)\n", "schnet_calc = GraphPESCalculator(model)\n", "\n", "cell_size = 4\n", "n_atoms = 10\n", "random_structure = ase.Atoms(\n", " positions=np.random.RandomState(42).rand(n_atoms, 3) * cell_size,\n", " cell=[cell_size, cell_size, cell_size],\n", " pbc=True,\n", ")" ] }, { "cell_type": "code", "execution_count": 23, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "(-0.034377, -0.034358)" ] }, "execution_count": 23, "metadata": {}, "output_type": "execute_result" } ], "source": [ "atom_index = 0\n", "direction = \"x\"\n", "\n", "fd_estimate = finite_difference_force_estimate(\n", " schnet_calc.get_potential_energy,\n", " random_structure,\n", " atom_index,\n", " direction,\n", " epsilon=1e-3,\n", ")\n", "autograd_value = schnet_calc.get_forces(random_structure)[atom_index][\n", " {\"x\": 0, \"y\": 1, \"z\": 2}[direction]\n", "]\n", "\n", "np.round(fd_estimate, 6), np.round(autograd_value, 6)" ] }, { "cell_type": "code", "execution_count": 24, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "array([[ 0.001911, 0.002092, -0.003447],\n", " [ 0.002092, 0.003045, -0.00231 ],\n", " [-0.003447, -0.00231 , 0.005248]])" ] }, "execution_count": 24, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# finite-difference stress\n", "fd_stress = finite_difference_stress_estimate(\n", " schnet_calc.get_potential_energy, random_structure, epsilon=1e-3\n", ")\n", "np.round(fd_stress, 6)" ] }, { "cell_type": "code", "execution_count": 25, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "array([[ 0.001928, 0.002072, -0.003426],\n", " [ 0.002072, 0.003042, -0.002297],\n", " [-0.003426, -0.002297, 0.005268]])" ] }, "execution_count": 25, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# autograd stress\n", "autograd_stress = voigt_6_to_full_3x3_stress(\n", " schnet_calc.get_stress(random_structure)\n", ")\n", "np.round(autograd_stress, 6)" ] } ], "metadata": { "kernelspec": { "display_name": "graph-pes", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.9.21" } }, "nbformat": 4, "nbformat_minor": 2 }