{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PAmUc0owFthe"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/ddmms/camml-tutorials/blob/main/notebooks/05-generative/energy-calc.ipynb)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rW1usWUgFthg"
      },
      "source": [
        "# Compare structure energies\n",
        "\n",
        "This notebook gives an example of how to load a `.cif` file and calculate the relaxed geometry energy. This can be used to compare the distributions of the energies of the materials from `Chemeloen` and `CrystaLLM`"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "execute next cell only if you are on google colab"
      ],
      "metadata": {
        "id": "1qkh11ghG4xT"
      }
    },
    {
      "cell_type": "code",
      "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)"
      ],
      "metadata": {
        "id": "T-u9NK35F-Lo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ByoDZHWBFthg"
      },
      "outputs": [],
      "source": [
        "from janus_core.calculations.geom_opt import GeomOpt\n",
        "import ase"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Use data_tutorials to get the data required for this tutorial:"
      ],
      "metadata": {
        "id": "HjWN6p9qGz-T"
      }
    },
    {
      "cell_type": "code",
      "source": [
        " from data_tutorials.data import get_data\n",
        "\n",
        "get_data(\n",
        "    url=\"https://gitlab.com/cam-ml/tutorials/-/raw/main/notebooks/05-generative/chemeleon/\",\n",
        "    filename=[\"generated_structure_0.cif\"],\n",
        "    folder=\"chemeleon\",\n",
        ")"
      ],
      "metadata": {
        "id": "6RceyRhuGI8r"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zRYc4MBlFthh"
      },
      "source": [
        "## Optimise\n",
        "\n",
        "Read in a generated strucrtue and run a BFGS optimisation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3ZYz_FOHFthi"
      },
      "outputs": [],
      "source": [
        "device = 'cuda'\n",
        "#or\n",
        "#device = \"cpu\"\n",
        "\n",
        "mol = ase.io.read('chemeleon/generated_structure_0.cif')\n",
        "\n",
        "# periodic\n",
        "opt = GeomOpt(\n",
        "    struct=mol,\n",
        "    arch=\"mace_mp\",\n",
        "    device=device,\n",
        "    model_path=\"medium-omat-0\",\n",
        "    calc_kwargs={\"default_dtype\": \"float64\"},\n",
        "    fmax=0.001,\n",
        "    filter_kwargs={\"hydrostatic_strain\": True},\n",
        ")\n",
        "\n",
        "opt.run()\n",
        "\n",
        "print(opt.struct.get_potential_energy())"
      ]
    }
  ],
  "metadata": {
    "language_info": {
      "name": "python"
    },
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}