examples/tutorials/geometry-optimisation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Geometry optimisation calculation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Run geometry optimisation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dLs_a7wLt98_"
   },
   "source": [
    "To run a geometry optimisation using aiida-mlip you need to define some inputs as AiiDA data types, to then pass to the calculation.\n",
    "\n",
    "To start, you will need a structure to optimise. \n",
    "\n",
    "The structure in the example is NaCl, generated using ase.build, or alternatively one can choose one of the structures in the folder `Structures`.\n",
    "The input structure in aiida-mlip needs to be saved as a StructureData type.\n",
    "We can print some properties of the structure, for example the cell or the atoms sites."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "NmNkuWExt8RE",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial cell parameters: [[0.0, 2.815, 2.815], [2.815, 0.0, 2.815], [2.815, 2.815, 0.0]]\n",
      "Structure's atoms sites: [<Site: kind name 'Na' @ 0.0,0.0,0.0>, <Site: kind name 'Cl' @ 2.815,0.0,0.0>]\n"
     ]
    }
   ],
   "source": [
    "from aiida import load_profile\n",
    "load_profile()\n",
    "\n",
    "from aiida.orm import StructureData\n",
    "from ase.io import read\n",
    "from ase.build import bulk\n",
    "\n",
    "\n",
    "#structure = StructureData(ase=read(\"structures/qmof-ffeef76.cif\"))\n",
    "structure = StructureData(ase=bulk(\"NaCl\", \"rocksalt\", 5.63))\n",
    "\n",
    "print(f\"Initial cell parameters: {structure.cell}\")\n",
    "print(f\"Structure's atoms sites: {structure.sites}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YxOTQEg9uGnN"
   },
   "source": [
    "Then we need to choose a model and architecture to be used for the calculation and save it as ModelData type, a specific data type of this plugin.\n",
    "In this example we use mace_mp with a model that we download from this URL: \"https://github.com/stfc/janus-core/raw/main/tests/models/mace_mp_small.model\", and we save the file in the cache folder (default=\"~/.cache/mlips/\"):\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "3iGSzzNithOk",
    "tags": []
   },
   "outputs": [],
   "source": [
    "from aiida_mlip.data.model import ModelData\n",
    "uri = \"https://github.com/stfc/janus-core/raw/main/tests/models/mace_mp_small.model\"\n",
    "model = ModelData.from_uri(uri, architecture=\"mace_mp\", cache_dir=\"mlips\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Cf28b-C1uM_I"
   },
   "source": [
    "If we already have the model saved in some folder we can save it as:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sqoVIB2guSKM",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# from aiida_mlip.data.model import ModelData\n",
    "# model = ModelData.from_local(\"mlips/mace_mp/mace_mp_small.model\", architecture=\"mace_mp\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xvEMrU5zuVn_"
   },
   "source": [
    "Another parameter that we need to define as AiiDA type is the code. Assuming the code is saved as `janus` in the `localhost` computer, the code info that are needed can be loaded as follow:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "mH5E3MtPtyj-",
    "tags": []
   },
   "outputs": [],
   "source": [
    "from aiida.orm import load_code\n",
    "code = load_code(\"janus@localhost\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "u-nJa_viuZMv"
   },
   "source": [
    "The other inputs can be set up as AiiDA Str, Float or Bool. There is a default for every input except the structure and code. This is a list of possible inputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "_XkQgKPMtyhf",
    "tags": []
   },
   "outputs": [],
   "source": [
    "from aiida.orm import Str, Float, Bool\n",
    "inputs = {\n",
    "        \"code\": code,\n",
    "        \"model\": model,\n",
    "        \"struct\": structure,\n",
    "        \"arch\": Str(model.architecture),\n",
    "        \"precision\": Str(\"float64\"),\n",
    "        \"device\": Str(\"cpu\"),\n",
    "        \"fmax\": Float(0.1), \n",
    "        \"opt_cell_lengths\": Bool(False), \n",
    "        \"opt_cell_fully\": Bool(True), \n",
    "        \"metadata\": {\"options\": {\"resources\": {\"num_machines\": 1}}},\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RCPc7-Z5uhoK"
   },
   "source": [
    "It's worth noting that the architecture is already defined within the model, accessible through the architecture property in the ModelData. Even if not explicitly provided as input, it will be automatically retrieved from the model.\n",
    "\n",
    "The calculation must be set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "XOaSjxT8tyek",
    "tags": []
   },
   "outputs": [],
   "source": [
    "from aiida.plugins import CalculationFactory\n",
    "geomoptCalc = CalculationFactory(\"mlip.opt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "aiida_mlip.calculations.geomopt.GeomOpt"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "geomoptCalc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BhWwGwm2uvYR"
   },
   "source": [
    "In this case, since we are running a geometry optimisation the entry point is `mlip.opt`\n",
    "Finally, run the calculation:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "agibHmT-tyU2",
    "tags": []
   },
   "outputs": [],
   "source": [
    "from aiida.engine import run_get_node\n",
    "result, node = run_get_node(geomoptCalc, **inputs)\n",
    "print(\"CALCULATION FINISHED\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analyse results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OSt6cQMeu5w9"
   },
   "source": [
    "`result` is a dictionary of the available results obtained from the calculation, while node contains the infor on the node where the calculation is run:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZwVf9a-Fu5aW",
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(f\"Printing output nodes dictionary: {result}\")\n",
    "print(\" \")\n",
    "print(f\"Printing calcjob node info: {node}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "L5oMOnQPvID2"
   },
   "source": [
    "In this case there are more outputs than the single point calculation, such as the output structure and the trajectory of the optimisation.\n",
    "We can see how many steps it took to optimise:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xyJnGWqlIQNT",
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(f\"The number of optimisation steps is: {result['traj_output'].numsteps}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The calculation can also be interacted with through verdi cli. Use `verdi process list` to show the list of calculations.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "! verdi process list -a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last calculation in the list is the last thing that was run. The PK number should correspond to the PK printed when you printed the node.\n",
    "You can interact with the data through verdi commands, by writing that PK number of the calculation of interest.\n",
    "Every calcjob node has a results dictionary which is outputted when running:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "! verdi calcjob res PK"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the node show command we can see the inputs and outputs of the calculation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! verdi node show PK"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the geometry optimisation we are most likely interested in the final structure and the trajectory of the geometry optimisation. Let's compare the initial and final cell parameters and see if they changed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from aiida.orm import load_node\n",
    "\n",
    "print(f\"Initial cell parameters: {structure.cell}\")\n",
    "final_structure = load_node(PK)\n",
    "print(f\"Final cell parameters: {final_structure.cell}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot energies and visualise provenance graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's analyse the steps of the optimisation. We'll run singlepoint calculation on every steps to see how the energy changed. We'll also use this to visualise a complex provenance graph with more calculations connected.\n",
    "Note the outputs of the calculation can be called either by using the load_node function when the Pk is known or directly with the outputs attribute of the calcjobe node.\n",
    "Note also that we use the `calcfunction` decorator get the single structures from the `TrajectoryData`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from aiida.orm import load_node\n",
    "from aiida.engine import calcfunction\n",
    "\n",
    "\n",
    "#traj = node.outputs.traj_output\n",
    "traj = load_node(PK) \n",
    "    \n",
    "@calcfunction\n",
    "def prepare_struct_inputs(traj, index):\n",
    "    return traj.get_step_structure(index.value)\n",
    "\n",
    "uri = \"https://github.com/stfc/janus-core/raw/main/tests/models/mace_mp_small.model\"\n",
    "model = ModelData.from_uri(uri, architecture=\"mace_mp\", cache_dir=\"mlips\")\n",
    "list_of_nodes = []\n",
    "\n",
    "\n",
    "inputs = {\n",
    "    \"code\": code,\n",
    "    \"model\": model,\n",
    "    \"precision\": Str(\"float64\"),\n",
    "    \"device\": Str(\"cpu\"),\n",
    "    \"metadata\": {\"options\": {\"resources\": {\"num_machines\": 1}}},\n",
    "}\n",
    "    \n",
    "for index in range(traj.numsteps):\n",
    "    print(index)\n",
    "    singlepointCalc = CalculationFactory(\"mlip.sp\")\n",
    "    struc = prepare_struct_inputs(traj, index)\n",
    "    inputs['struct']=struc\n",
    "    result, node = run_get_node(singlepointCalc, **inputs)\n",
    "    list_of_nodes.append(node)\n",
    "print(\"calculations ended\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's print the list of alcjob nodes that we just created"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "list_of_nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can use it for getting the energies in every step and plotting them.\n",
    "(a better alternative to the list of nodes might be to use a AiiDA group, see high-throughput-screening tutorial)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "steps = []\n",
    "energies = []\n",
    "\n",
    "# Loop through each node to extract step number and energy level\n",
    "for step, node in enumerate(list_of_nodes):\n",
    "    energy = node.outputs.results_dict.get_dict()['info']['mace_mp_energy']\n",
    "    steps.append(step)\n",
    "    energies.append(energy)\n",
    "\n",
    "# Plotting the energy levels over steps\n",
    "plt.figure(figsize=(10, 6))\n",
    "plt.plot(steps, energies, marker='o', linestyle='-', color='g')\n",
    "plt.title('Energy Levels over Steps')\n",
    "plt.xlabel('Step')\n",
    "plt.ylabel('Energy')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the energy decreased, which is what we want in a geometry optimisation process.\n",
    "\n",
    "Now let's generate the provenance graph. (Insert PK number of the TrajectoryData in the code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "! verdi node graph generate PK"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The provenance graph shows both the calculation that created the `TrajectoryData` and the calculations that we run using the structures in the `TrajectoryData`.\n",
    "This is made possible by the use of the `calcfunction` decorator that we used. If we had not used it, the graph would stop a the `TrajectoryData` and the other `Singlepoint` calculations would be independent."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "private_outputs": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "venv_aiida_11",
   "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.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}