Tutorial_model_force_check.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "26ed2377",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import jraph\n",
    "import ase\n",
    "import ase.io\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "import optax\n",
    "\n",
    "from phonax.predictors import predict_energy_forces_stress\n",
    "\n",
    "from phonax.trained_models import (\n",
    "    NequIP_JAXMD_uniIAP_model,\n",
    "    NequIP_JAXMD_uniIAP_PBEsol_finetuned_model,\n",
    "    MACE_uniIAP_model,\n",
    "    MACE_uniIAP_PBEsol_finetuned_model,\n",
    "    NequIP_JAX_uniIAP_model,\n",
    ")\n",
    "\n",
    "from phonax.data_utils import (\n",
    "    crystal_struct_to_period_graph,\n",
    "    crystal_atoms_to_period_graph,\n",
    ")\n",
    "\n",
    "jax.config.update(\"jax_debug_nans\", True)\n",
    "jax.config.update(\"jax_debug_infs\", True)\n",
    "np.set_printoptions(precision=3, suppress=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "503a92cb",
   "metadata": {},
   "source": [
    "# Atomic structure force examination\n",
    "\n",
    "One of the conditions in deriving physical phonon or vibrational spectrum of an atomic structure is the (local) equilibrium structure.\n",
    "That is, the atomic forces should be (close to) zero, and the eigenvalues for the Hessians can be interpreted as the vibrational mode energies.\n",
    "Out-of-equilibrium structures can lead to unstable negative modes or spurious mode frequencies in diagonalizing the Hessian matrix.\n",
    "\n",
    "Therefore, to make consistent phonon predictions of a given atomic structure, one has to first check the predicted forces within the structure and make sure they are vanishing and hence close to a local equilibrium.\n",
    "In the interatomic potential models parametrized by equivariant graph neural networks, different training data and model architectures can lead to different predictions of the local equilibrium structure.\n",
    "For example, different exchange-correlation functional used in the DFT such as LDA, PBE, PBEsol, or the more advanced hybrid functionals, can lead to different equilibrium structure predictions.\n",
    "\n",
    "Here in this tutorial, we demonstrate how to check the (residual) atomic forces given an atomic structure and a trained energy model. For out-of-equilibrum structures, we also show how we can perform relaxations of the atomic coordinates to reach a local equilibrum state.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "259bd935",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Create NequIP (JAX-MD version) with parameters {'use_sc': True, 'graph_net_steps': 2, 'hidden_irreps': '128x0e+ 128x0o  + 64x1e +64x1o +64x2e +64x2o', 'nonlinearities': {'e': 'swish', 'o': 'tanh'}, 'r_max': 5.0, 'avg_num_neighbors': 36.712880186304396, 'avg_r_min': None, 'num_species': 100, 'radial_basis': <function bessel_basis at 0x7fd2910ce840>, 'radial_envelope': <function soft_envelope at 0x7fd2910ce7a0>}\n"
     ]
    }
   ],
   "source": [
    "# NequIP model trained with universal-IAP (PBE)\n",
    "model_fn, params, num_message_passing, r_max = NequIP_JAXMD_uniIAP_model(os.path.join(os.getcwd(), 'trained-models'))\n",
    "\n",
    "# NequIP model trained with universal-IAP (PBE) + PBEsol fine-tuning\n",
    "#model_fn, params, num_message_passing, r_max = NequIP_JAXMD_uniIAP_PBEsol_finetuned_model(os.path.join(os.getcwd(), 'trained-models'))\n",
    "\n",
    "# MACE model trained with universal-IAP (PBE)\n",
    "#model_fn, params, num_message_passing, r_max = MACE_uniIAP_model(os.path.join(os.getcwd(), 'trained-models'))\n",
    "\n",
    "# MACE model trained with universal-IAP (PBE) + PBEsol fine-tuning\n",
    "#model_fn, params, num_message_passing, r_max = MACE_uniIAP_PBEsol_finetuned_model(os.path.join(os.getcwd(), 'trained-models'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fecd0d2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "67667c6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "atoms, graph = crystal_struct_to_period_graph('data/mp-149/mp-149.vasp', r_max, num_message_passing)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa8d295a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a5597871",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'energy': Array([-10.781], dtype=float32), 'forces': Array([[ 0., -0., -0.],\n",
      "       [ 0.,  0.,  0.]], dtype=float32), 'pressure': Array([0.014], dtype=float32), 'stress': Array([[[ 0.005,  0.   ,  0.   ],\n",
      "        [-0.   ,  0.005, -0.   ],\n",
      "        [ 0.   , -0.   ,  0.005]]], dtype=float32), 'stress_cell': Array([[[ 0.005,  0.   ,  0.   ],\n",
      "        [-0.   ,  0.005, -0.   ],\n",
      "        [ 0.   , -0.   ,  0.005]]], dtype=float32), 'stress_forces': Array([[[0., 0., 0.],\n",
      "        [0., 0., 0.],\n",
      "        [0., 0., 0.]]], dtype=float32)}\n"
     ]
    }
   ],
   "source": [
    "predictor = jax.jit(\n",
    "    lambda w, g: predict_energy_forces_stress(lambda *x: model_fn(w, *x), g)\n",
    ")\n",
    "pred_out = predictor(params, graph)\n",
    "\n",
    "print(pred_out)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cc4efd1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "8d398357",
   "metadata": {},
   "source": [
    "## Atomic relaxations\n",
    "\n",
    "Given an out-of-equilibrium atomic structure, we would like to perform atomic coordinate updates to approach the local equilibrium with vanishing atomic forces.\n",
    "We can use the computed atomic forces as the effective gradients in guiding the relaxations of the atomic positions.\n",
    "Here we utilize the Optax optimization library for JAX to perform the relaxation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "18379301",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Init stage: {'energy': Array([-21.427], dtype=float32), 'forces': Array([[ 0.   ,  0.   , -0.   ],\n",
      "       [ 0.   , -0.   , -0.004],\n",
      "       [-0.   , -0.   ,  0.004]], dtype=float32), 'pressure': Array([0.002], dtype=float32), 'stress': Array([[[ 0.001, -0.   , -0.   ],\n",
      "        [-0.   ,  0.001,  0.   ],\n",
      "        [-0.   ,  0.   ,  0.   ]]], dtype=float32), 'stress_cell': Array([[[ 0.001,  0.   , -0.   ],\n",
      "        [-0.   ,  0.001, -0.   ],\n",
      "        [ 0.   , -0.   , -0.   ]]], dtype=float32), 'stress_forces': Array([[[ 0., -0., -0.],\n",
      "        [ 0., -0.,  0.],\n",
      "        [-0.,  0.,  0.]]], dtype=float32)}\n",
      "[[ 1.596  1.014  0.   ]\n",
      " [ 0.     2.027 -1.721]\n",
      " [ 0.     2.027  1.721]]\n",
      "[[ 1.595  1.06  -0.   ]\n",
      " [-0.001  1.981 -1.565]\n",
      " [-0.001  1.981  1.565]]\n"
     ]
    }
   ],
   "source": [
    "struct_atoms_init = ase.io.read('data/crystals/MoS2-mono.vasp')\n",
    "struct_graph_init = crystal_atoms_to_period_graph(struct_atoms_init,\n",
    "                                                  r_max, \n",
    "                                                  num_message_passing)\n",
    "struct_pred_init = predictor(params,struct_graph_init)\n",
    "struct_atoms = struct_atoms_init.copy()\n",
    "print('Init stage:',struct_pred_init)\n",
    "\n",
    "struct_atoms = ase.io.read('data/crystals/MoS2-mono.vasp')\n",
    "def struct_relax_grad(pos):\n",
    "    struct_atoms.set_positions(pos)\n",
    "    struct_graph = crystal_atoms_to_period_graph(struct_atoms,\n",
    "                                                 r_max, \n",
    "                                                 num_message_passing)\n",
    "    \n",
    "    struct_pred = predictor(params,struct_graph)\n",
    "    pred_force = struct_pred['forces']\n",
    "    return -pred_force\n",
    "\n",
    "\n",
    "relax_lr = 0.01\n",
    "optimizer = optax.adam(relax_lr)\n",
    "initpos = struct_atoms_init.get_positions()\n",
    "initpos[:,1] *= 1.10\n",
    "initpos[:,2] *= 1.10\n",
    "initpos[:,2] -= np.mean(initpos[:,2])\n",
    "relax_pos = jnp.array(initpos)\n",
    "opt_state = optimizer.init(relax_pos)\n",
    "\n",
    "train_loss = []\n",
    "print(relax_pos)\n",
    "for _ in range(100):\n",
    "    grads = struct_relax_grad(relax_pos)\n",
    "    updates, opt_state = optimizer.update(grads, opt_state)\n",
    "    relax_pos = optax.apply_updates(relax_pos, updates)\n",
    "    \n",
    "    #print(relax_pos[2,2] - relax_pos[0,2])\n",
    "    \n",
    "relax_pos = np.array(relax_pos)\n",
    "relax_pos[:,2] -= np.mean(relax_pos[:,2])\n",
    "print(relax_pos)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ac6fd58",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f798fd19",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03552253",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a25b6b5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}