docs/source/tutorials/python/eos.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/eos.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4vcbRxmIhHLL"
   },
   "source": [
    "# Equation of State"
   ]
  },
  {
   "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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ase import Atoms\n",
    "from ase.io import read\n",
    "from weas_widget import WeasWidget\n",
    "\n",
    "from data_tutorials.data import get_data\n",
    "\n",
    "from janus_core.calculations.eos import EoS"
   ]
  },
  {
   "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": [
    "get_data(\n",
    "    url=\"https://raw.githubusercontent.com/stfc/janus-tutorials/main/data/\",\n",
    "    filename=[\"beta_quartz.cif\"],\n",
    "    folder=\"data\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Equation of state for α-quartz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Build the structure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "α_quartz = Atoms(\n",
    "    symbols=(*[\"Si\"] * 3, *[\"O\"] * 6),\n",
    "    scaled_positions=[\n",
    "        [0.469700, 0.000000, 0.000000],\n",
    "        [0.000000, 0.469700, 0.666667],\n",
    "        [0.530300, 0.530300, 0.333333],\n",
    "        [0.413500, 0.266900, 0.119100],\n",
    "        [0.266900, 0.413500, 0.547567],\n",
    "        [0.733100, 0.146600, 0.785767],\n",
    "        [0.586500, 0.853400, 0.214233],\n",
    "        [0.853400, 0.586500, 0.452433],\n",
    "        [0.146600, 0.733100, 0.880900],\n",
    "    ],\n",
    "    cell=[\n",
    "        [4.916000, 0.000000, 0.000000],\n",
    "        [-2.45800, 4.257381, 0.000000],\n",
    "        [0.000000, 0.000000, 5.405400],\n",
    "    ],\n",
    "    pbc=True,\n",
    ")\n",
    "\n",
    "v=WeasWidget()\n",
    "v.from_ase(α_quartz)\n",
    "v.avr.model_style = 1\n",
    "v.avr.show_hydrogen_bonds = True\n",
    "v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate the equation of state using the MACE-MP potential:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mace_eos = EoS(\n",
    "    struct=α_quartz.copy(),\n",
    "    arch=\"mace_mp\",\n",
    "    device=\"cpu\",\n",
    "    model=\"small\",\n",
    "    calc_kwargs={\"default_dtype\": \"float64\"},\n",
    "    minimize_kwargs={\"filter_func\": None},\n",
    "    min_volume=0.75,\n",
    "    max_volume=1.25,\n",
    "    n_volumes=20,\n",
    ").run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mace_eos[\"eos\"].plot(show=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Equation of state for β-quartz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perform the same calculation for β-quartz:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "β_quartz = read(\"data/beta_quartz.cif\")\n",
    "\n",
    "v=WeasWidget()\n",
    "v.from_ase(β_quartz)\n",
    "v.avr.model_style = 1\n",
    "v.avr.show_hydrogen_bonds = True\n",
    "v\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mace_eos_beta = EoS(\n",
    "    struct=β_quartz.copy(),\n",
    "    arch=\"mace_mp\",\n",
    "    device=\"cpu\",\n",
    "    model=\"small\",\n",
    "    calc_kwargs={\"default_dtype\": \"float64\"},\n",
    "    minimize_kwargs={\"filter_func\": None},\n",
    "    min_volume=0.75,\n",
    "    max_volume=1.25,\n",
    "    n_volumes=20,\n",
    ").run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mace_eos_beta[\"eos\"].plot(show=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Combining plots for α-quartz and β-quartz:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "ax = plt.gca()\n",
    "\n",
    "data_alpha = mace_eos[\"eos\"].getplotdata()\n",
    "data_beta = mace_eos_beta[\"eos\"].getplotdata()\n",
    "\n",
    "ax.plot(data_alpha[4], data_alpha[5], ls=\"-\", color=\"C3\", label=\"α-quartz\")\n",
    "ax.plot(data_alpha[6], data_alpha[7], ls=\"\", marker=\"x\", color=\"C4\", mfc=\"C4\")\n",
    "\n",
    "ax.plot(data_beta[4], data_beta[5], ls=\"-\", color=\"C0\", label=\"β-quartz\")\n",
    "ax.plot(data_beta[6], data_beta[7], ls=\"\", marker=\"x\", color=\"C2\", mfc=\"C2\")\n",
    "\n",
    "ax.set_xlabel(\"volume [Å$^3$]\")\n",
    "ax.set_ylabel(\"energy [eV]\")\n",
    "ax.legend()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparing MACE to CHGNeT and SevenNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sevennet_eos = EoS(\n",
    "    struct=α_quartz.copy(),\n",
    "    arch=\"sevennet\",\n",
    "    device=\"cpu\",\n",
    "    minimize_kwargs={\"filter_func\": None},\n",
    ").run()\n",
    "\n",
    "sevennet_eos[\"eos\"].plot(show=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chgnet_eos = EoS(\n",
    "    struct=α_quartz.copy(),\n",
    "    arch=\"chgnet\",\n",
    "    device=\"cpu\",\n",
    "    minimize_kwargs={\"filter_func\": None},\n",
    ").run()\n",
    "\n",
    "chgnet_eos[\"eos\"].plot(show=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"MACE energy [eV]: {mace_eos['e_0']}\")\n",
    "print(f\"SevenNet energy [eV]: {sevennet_eos['e_0']}\")\n",
    "print(f\"CHGNeT energy [eV]: {chgnet_eos['e_0']}\")\n",
    "\n",
    "print()\n",
    "\n",
    "print(f\"MACE volume [Å^3]: {mace_eos['v_0']}\")\n",
    "print(f\"SevenNet volume [Å^3]: {sevennet_eos['v_0']}\")\n",
    "print(f\"CHGNeT volume [Å^3]: {chgnet_eos['v_0']}\")\n",
    "\n",
    "print()\n",
    "\n",
    "print(f\"MACE bulk_modulus [GPa]: {mace_eos['bulk_modulus']}\")\n",
    "print(f\"SevenNet bulk_modulus [GPa]: {sevennet_eos['bulk_modulus']}\")\n",
    "print(f\"CHGNeT bulk_modulus [GPa]: {chgnet_eos['bulk_modulus']}\")"
   ]
  }
 ],
 "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": 0
}