docs/source/tutorials/python/single_point.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/stfc/janus-core/blob/main/docs/source/tutorials/python/single_point.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Single Point"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4vcbRxmIhHLL"
   },
   "source": [
    "`janus-core` contains various machine learnt interatomic potentials (MLIPs), including MACE based models (MACE-MP, MACE-OFF), CHGNet, SevenNet and more, full list on https://github.com/stfc/janus-core.\n",
    "\n",
    "Other will be added as their utility is proven beyond a specific material."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up environment (optional)\n",
    "\n",
    "These steps are required for Google Colab, but may work on other systems too:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TF-EiWxyuMc7"
   },
   "outputs": [],
   "source": [
    "# import locale\n",
    "# locale.getpreferredencoding = lambda: \"UTF-8\"\n",
    "\n",
    "# ! pip uninstall torch torchaudio torchvision numpy -y\n",
    "# ! uv pip install janus-core[all] data-tutorials torch==2.5.1 --system\n",
    "# get_ipython().kernel.do_shutdown(restart=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use `data_tutorials` to get the data required for this tutorial:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from data_tutorials.data import get_data\n",
    "\n",
    "get_data(\n",
    "    url=\"https://raw.githubusercontent.com/stfc/janus-tutorials/main/data/\",\n",
    "    filename=[\"sucrose.xyz\", \"NaCl-set.xyz\"],\n",
    "    folder=\"data\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "857K7R9Cenca"
   },
   "source": [
    "## Single point calculations for periodic system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cgch8VQ--AES"
   },
   "outputs": [],
   "source": [
    "from ase.build import bulk\n",
    "\n",
    "from janus_core.calculations.single_point import SinglePoint\n",
    "\n",
    "# Change to cuda if you have a gpu or mps for apple silicon\n",
    "device = \"cpu\"\n",
    "\n",
    "NaCl = bulk(\"NaCl\", \"rocksalt\", a=5.63, cubic=True)\n",
    "sp = SinglePoint(\n",
    "    struct=NaCl,\n",
    "    arch=\"mace_mp\",\n",
    "    device=device,\n",
    "    model=\"small\",\n",
    "    calc_kwargs={\"default_dtype\": \"float64\"},\n",
    ")\n",
    "res_mace = sp.run()\n",
    "\n",
    "NaCl = bulk(\"NaCl\", \"rocksalt\", a=5.63, cubic=True)\n",
    "sp = SinglePoint(\n",
    "    struct=NaCl,\n",
    "    arch=\"sevennet\",\n",
    "    device=device,\n",
    ")\n",
    "res_sevennet = sp.run()\n",
    "\n",
    "NaCl = bulk(\"NaCl\", \"rocksalt\", a=5.63, cubic=True)\n",
    "sp = SinglePoint(\n",
    "    struct=NaCl,\n",
    "    arch=\"chgnet\",\n",
    "    device=device,\n",
    ")\n",
    "res_chgnet = sp.run()\n",
    "\n",
    "print(f\"    MACE[eV]: {res_mace['energy']}\")\n",
    "print(f\"SevenNet[eV]: {res_sevennet['energy']}\")\n",
    "print(f\"  CHGNeT[eV]: {res_chgnet['energy']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fDt01muwC5fl"
   },
   "source": [
    "## Simple Molecules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cbvNchA4_dEq"
   },
   "outputs": [],
   "source": [
    "from ase.build import molecule\n",
    "from weas_widget import WeasWidget\n",
    "\n",
    "from janus_core.calculations.single_point import SinglePoint\n",
    "\n",
    "sp = SinglePoint(\n",
    "    struct=molecule(\"H2O\"),\n",
    "    arch=\"mace_off\",\n",
    "    device=device,\n",
    "    model=\"medium\",\n",
    ")\n",
    "res = sp.run()\n",
    "print(res)\n",
    "v=WeasWidget()\n",
    "v.from_ase(sp.struct)\n",
    "v.avr.model_style = 1\n",
    "v.avr.show_hydrogen_bonds = True\n",
    "v\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lazs3WCKC3fY"
   },
   "source": [
    "## Sugar on salt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WHHUSn7gCz-l"
   },
   "outputs": [],
   "source": [
    "from ase.build import add_adsorbate, bulk\n",
    "from ase.io import read, write\n",
    "\n",
    "a = 5.63\n",
    "NaCl = bulk(\n",
    "    \"NaCl\", crystalstructure=\"rocksalt\", cubic=True, orthorhombic=True, a=5.63\n",
    ") * (6, 6, 3)\n",
    "\n",
    "NaCl.center(vacuum=20.0, axis=2)\n",
    "sugar = read(\"data/sucrose.xyz\")\n",
    "add_adsorbate(slab=NaCl, adsorbate=sugar, height=4.5, position=(10, 10))\n",
    "write(\"slab.xyz\", NaCl)\n",
    "\n",
    "sp = SinglePoint(\n",
    "    struct=\"slab.xyz\",\n",
    "    arch=\"mace_mp\",\n",
    "    device=device,\n",
    "    model=\"small\",\n",
    ")\n",
    "res = sp.run()\n",
    "print(res)\n",
    "\n",
    "v = WeasWidget()\n",
    "v.avr.model_style = 1\n",
    "v.avr.show_hydrogen_bonds = True\n",
    "v.from_ase(NaCl)\n",
    "v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7BDbwfmDE-BO"
   },
   "source": [
    "## Calculate an entire collection of data frames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ase.io import read\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "frames = read(\"data/NaCl-set.xyz\", format=\"extxyz\", index=\":\")\n",
    "dft_energy = np.array([s.info[\"dft_energy\"] / len(s) for s in frames])\n",
    "\n",
    "sp = SinglePoint(\n",
    "    struct=\"data/NaCl-set.xyz\",\n",
    "    arch=\"mace_mp\",\n",
    "    device=device,\n",
    "    model=\"medium-0b3\",\n",
    ")\n",
    "sp.run()\n",
    "\n",
    "mace_mp_energy = np.array([s.info[\"mace_mp_energy\"] / len(s) for s in sp.struct])\n",
    "rmse_mace = np.linalg.norm(mace_mp_energy - dft_energy) / np.sqrt(len(dft_energy))\n",
    "\n",
    "sp = SinglePoint(struct=\"data/NaCl-set.xyz\", arch=\"chgnet\", device=device)\n",
    "sp.run()\n",
    "\n",
    "chgnet_energy = np.array([s.info[\"chgnet_energy\"] / len(s) for s in sp.struct])\n",
    "rmse_chgnet = np.linalg.norm(chgnet_energy - dft_energy) / np.sqrt(len(dft_energy))\n",
    "\n",
    "sp = SinglePoint(struct=\"data/NaCl-set.xyz\", arch=\"sevennet\", device=device)\n",
    "sp.run()\n",
    "\n",
    "sevennet_energy = np.array([s.info[\"sevennet_energy\"] / len(s) for s in sp.struct])\n",
    "rmse_sevennet = np.linalg.norm(sevennet_energy - dft_energy) / np.sqrt(len(dft_energy))\n",
    "\n",
    "print(\n",
    "    f\"rmse: mace_mp = {rmse_mace}, chgnet = {rmse_chgnet}, \"\n",
    "    f\"sevennet = {rmse_sevennet} eV/atom\"\n",
    ")\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.scatter(dft_energy, mace_mp_energy, marker=\"o\", label=\"mace-mp-0\")\n",
    "ax.scatter(dft_energy, sevennet_energy, marker=\"x\", label=\"sevennet\")\n",
    "ax.scatter(dft_energy, chgnet_energy, marker=\"+\", label=\"chgnet\")\n",
    "ax.legend()\n",
    "plt.xlabel(\"MLIP [eV/atom]\")\n",
    "plt.ylabel(\"DFT [eV/atom]\")\n",
    "\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "authorship_tag": "ABX9TyNvtIsPHVgkt0NvUv51T6ZG",
   "gpuType": "T4",
   "include_colab_link": true,
   "private_outputs": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "janus",
   "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}