docs/source/tutorials/python/geom_opt.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/geom_opt.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4vcbRxmIhHLL"
   },
   "source": [
    "# Geometry Optimization"
   ]
  },
  {
   "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.io import read\n",
    "from ase.optimize import FIRE\n",
    "from weas_widget import WeasWidget\n",
    "\n",
    "from data_tutorials.data import get_data\n",
    "\n",
    "from janus_core.calculations.single_point import SinglePoint\n",
    "from janus_core.calculations.geom_opt import GeomOpt"
   ]
  },
  {
   "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=[\"NaCl-deformed.cif\"],\n",
    "    folder=\"data\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare for optimization of a deformed salt structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NaCl = read(\"data/NaCl-deformed.cif\")\n",
    "NaCl.wrap()\n",
    "v=WeasWidget()\n",
    "v.from_ase(NaCl)\n",
    "v.avr.model_style = 1\n",
    "v.avr.show_hydrogen_bonds = True\n",
    "v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sp_mace = SinglePoint(\n",
    "    struct=NaCl.copy(),\n",
    "    arch=\"mace_mp\",\n",
    "    device=\"cpu\",\n",
    "    model_path=\"small\",\n",
    "    calc_kwargs={\"default_dtype\": \"float64\"},\n",
    "    properties=\"energy\",\n",
    ")\n",
    "\n",
    "init_energy = sp_mace.run()[\"energy\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To optimize only the atomic positions and not the cell, set `filter_func = None`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_NaCl = GeomOpt(\n",
    "    struct=sp_mace.struct,\n",
    "    fmax=0.001,\n",
    "    filter_func=None,\n",
    ")\n",
    "\n",
    "optimized_NaCl.run()\n",
    "v=WeasWidget()\n",
    "v.from_ase(optimized_NaCl.struct)\n",
    "v.avr.model_style = 1\n",
    "v.avr.show_hydrogen_bonds = True\n",
    "v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check energy has been lowered, and cell is unchanged:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Initial cell: {NaCl.cell.cellpar()}\")\n",
    "print(f\"Initial energy: {init_energy}\")\n",
    "\n",
    "print(f\"Final cell: {optimized_NaCl.struct.cell.cellpar()}\")\n",
    "print(f\"Final energy: {optimized_NaCl.struct.get_potential_energy()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimizing cell vectors and atomic positions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setting `filter_kwargs = {\"hydrostatic_strain\": True}` allows the cell lengths to be changed, in addition to atomic positions, but cell angles remain fixed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_NaCl_lengths = GeomOpt(\n",
    "    struct=NaCl.copy(),\n",
    "    arch=\"mace_mp\",\n",
    "    device=\"cpu\",\n",
    "    model_path=\"small\",\n",
    "    calc_kwargs={\"default_dtype\": \"float64\"},\n",
    "    fmax=0.001,\n",
    "    filter_kwargs={\"hydrostatic_strain\": True},\n",
    ")\n",
    "optimized_NaCl_lengths.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check energy has been lowered, and cell lengths have been updated, but angles remain unchanged:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Initial cell: {NaCl.cell.cellpar()}\")\n",
    "print(f\"Initial energy: {init_energy}\")\n",
    "\n",
    "print(f\"Final cell: {optimized_NaCl_lengths.struct.cell.cellpar()}\")\n",
    "print(f\"Final energy: {optimized_NaCl_lengths.struct.get_potential_energy()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimizing at constant pressure and volume"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculations can also be run at a fixed pressure and volume, by setting `filter_kwargs = {\"scalar_pressure\": x, \"constant_volume\": True}`\n",
    "\n",
    "By default, both the cell lengths and angles will be optimized, in addition to the atomic positions.\n",
    "\n",
    "We can also set the optimizer function and filter function used, either by passing the function itself (e.g. `FIRE`) or passing the name of the ASE function (e.g. `\"ExpCellFilter\"`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_NaCl_pressure = GeomOpt(\n",
    "    struct=NaCl.copy(),\n",
    "    arch=\"mace_mp\",\n",
    "    device=\"cpu\",\n",
    "    model_path=\"small\",\n",
    "    calc_kwargs={\"default_dtype\": \"float64\"},\n",
    "    fmax=0.01,\n",
    "    filter_kwargs={\"scalar_pressure\": 0.05, \"constant_volume\": True},\n",
    "    optimizer=FIRE,\n",
    "    filter_func=\"ExpCellFilter\",\n",
    ")\n",
    "optimized_NaCl_pressure.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check cell lengths and angles have both been updated:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Initial cell: {NaCl.cell.cellpar()}\")\n",
    "print(f\"Initial energy: {init_energy}\")\n",
    "\n",
    "print(f\"Final cell: {optimized_NaCl_pressure.struct.cell.cellpar()}\")\n",
    "print(f\"Final energy: {optimized_NaCl_pressure.struct.get_potential_energy()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Comparing MACE to CHGNet and SevenNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_NaCl_mace = GeomOpt(\n",
    "    struct=NaCl.copy(),\n",
    "    arch=\"mace_mp\",\n",
    "    device=\"cpu\",\n",
    "    model_path=\"small\",\n",
    "    calc_kwargs={\"default_dtype\": \"float64\"},\n",
    "    fmax=0.01,\n",
    ")\n",
    "optimized_NaCl_mace.run()\n",
    "\n",
    "optimized_NaCl_chgnet = GeomOpt(\n",
    "    struct=NaCl.copy(),\n",
    "    arch=\"chgnet\",\n",
    "    device=\"cpu\",\n",
    "    calc_kwargs={\"default_dtype\": \"float64\"},\n",
    "    fmax=0.01,\n",
    ")\n",
    "optimized_NaCl_chgnet.run()\n",
    "\n",
    "optimized_NaCl_sevennet = GeomOpt(\n",
    "    struct=NaCl.copy(),\n",
    "    arch=\"sevennet\",\n",
    "    device=\"cpu\",\n",
    "    calc_kwargs={\"default_dtype\": \"float64\"},\n",
    "    fmax=0.01,\n",
    ")\n",
    "optimized_NaCl_sevennet.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Initial energy: {init_energy}\")\n",
    "\n",
    "print(f\"Final energy (MACE): {optimized_NaCl_mace.struct.get_potential_energy()}\")\n",
    "print(f\"Final energy (CHGNet): {optimized_NaCl_chgnet.struct.get_potential_energy()}\")\n",
    "print(f\"Final energy (SevenNet): {optimized_NaCl_sevennet.struct.get_potential_energy()}\")"
   ]
  }
 ],
 "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
}