docs/source/tutorials/python/md.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/md.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4vcbRxmIhHLL"
   },
   "source": [
    "# Molecular Dynamics"
   ]
  },
  {
   "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 transformers 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": {
    "id": "cgch8VQ--AES"
   },
   "outputs": [],
   "source": [
    "from ase.build import bulk\n",
    "from weas_widget import WeasWidget\n",
    "from ase.io import read\n",
    "from data_tutorials.data import get_data\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from janus_core.calculations.md import NVE, NVT\n",
    "from janus_core.helpers.stats import Stats\n",
    "from janus_core.processing import post_process"
   ]
  },
  {
   "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=[\"precomputed_NaCl-traj.xyz\"],\n",
    "    folder=\"data\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cooling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Build NaCl structure and attach the MACE calculator:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NaCl = bulk(\"NaCl\", \"rocksalt\", a=5.63, cubic=True)\n",
    "NaCl = NaCl * (2, 2, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v=WeasWidget()\n",
    "v.from_ase(NaCl)\n",
    "v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prepare a simulation, cooling the structure from 300K to 200K in stepx of 20K, with 5fs at each temperature:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cooling = NVT(\n",
    "    struct=NaCl.copy(),\n",
    "    arch=\"mace_mp\",\n",
    "    device=\"cpu\",\n",
    "    model=\"small\",\n",
    "    calc_kwargs={\"default_dtype\": \"float64\"},\n",
    "    temp_start=300.0,\n",
    "    temp_end=200.0,\n",
    "    temp_step=20,\n",
    "    temp_time=5,\n",
    "    stats_every=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run cooling:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cooling.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All output files are saved in a results directory, `janus_results`.\n",
    "\n",
    "Within this, the final structures at each temperature are saved in `Cl32Na32-nvt-T300.0-final.xyz`, `Cl32Na32-nvt-T280.0-final.xyz`, ..., `Cl32Na32-nvt-T200.0-final.xyz`.\n",
    "\n",
    "The statiscs from the simulation are also saved every 20 steps in `Cl32Na32-nvt-T300.0-T200.0-stats.dat`. This can then be analysed using the `Stats` module:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = Stats(\"janus_results/Cl32Na32-nvt-T300.0-T200.0-stats.dat\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(data[0], data[5], label=\"Actual\")\n",
    "plt.plot(data[0], data[16], label=\"Target\")\n",
    "plt.legend()\n",
    "plt.xlabel(\"Steps\")\n",
    "plt.ylabel(\"Temperature / K\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Heating, followed by MD"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This will prepare an NVT MD simulation, initially increasing the temperature from 0K to 300K in 20K steps, with 1fs at each temperature, before a further 10 steps (10fs) at 300K.\n",
    "\n",
    "The final structure at each temperature will be saved, e.g. `Cl32Na32-nvt-T0-final.xyz`, `Cl32Na32-nvt-T0-final.xyz`, ..., `Cl32Na32-nvt-T300-final.xyz.`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "heating = NVT(\n",
    "    struct=NaCl.copy(),\n",
    "    arch=\"mace_mp\",\n",
    "    device=\"cpu\",\n",
    "    model=\"small\",\n",
    "    calc_kwargs={\"default_dtype\": \"float64\"},\n",
    "    temp_start=0,  # Start of temperature ramp\n",
    "    temp_end=300.0,  # End of temperature ramp\n",
    "    temp_step=20,  # Temperature ramp increments\n",
    "    temp_time=1,  # Time at each temperature in ramp\n",
    "    temp=300,  # MD temperature\n",
    "    steps=10,  # MD steps at 300K\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "heating.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The same structure can then be used to run an additional NVE MD similation for 50 steps (50fs), with post-processing performed to compute the RDF by setting `post_process_kwargs = {\"rdf_compute\": True}`, with the results saved to `Cl32Na32-nve-rdf.dat`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "md = NVE(\n",
    "    struct=heating.struct,\n",
    "    temp=300,\n",
    "    stats_every=5,\n",
    "    steps=50,\n",
    "    post_process_kwargs={\"rdf_compute\": True, \"rdf_rmax\": 5, \"rdf_bins\": 50},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "md.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# view trajectory\n",
    "\n",
    "v=WeasWidget()\n",
    "traj = read(\"data/precomputed_NaCl-traj.xyz\", index=\":\")\n",
    "v.from_ase(traj)\n",
    "v.avr.model_style = 1\n",
    "v.avr.show_hydrogen_bonds = True\n",
    "v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rdf = np.loadtxt(\"janus_results/Cl32Na32-nve-T300-rdf.dat\")\n",
    "bins, counts = zip(*rdf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(bins, counts)\n",
    "plt.ylabel(\"RDF\")\n",
    "plt.xlabel(\"Distance / Å\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The same post-processing can also be performed separately after completing the simulation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = read(\"data/precomputed_NaCl-traj.xyz\", index=\":\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rdf = post_process.compute_rdf(traj, rmax=5.0, nbins=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(rdf[0], rdf[1])\n",
    "plt.ylabel(\"RDF\")\n",
    "plt.xlabel(\"Distance / Å\")\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": 0
}