Day_3/icomse-3.4.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lohEE2bT1-qB"
   },
   "source": [
    "This notebook serves as an example of how to analyze a simulation trajectory using unsupervised techniques. Here, specifically, we'll be analyzing a simulation of cyclohexane conformations, simulated using quantum-espresso."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "w9yZxuem1-qC"
   },
   "source": [
    "Before running this notebook, you will need to install:\n",
    "    \n",
    "- [ase](https://wiki.fysik.dtu.dk/ase/index.html)\n",
    "- [scikit-learn](https://scikit-learn.org/)\n",
    "- [scikit-matter](https://github.com/scikit-learn-contrib/scikit-matter)\n",
    "- [chemiscope](https://chemiscope.org)\n",
    "\n",
    "in addition to standard packages [numpy](https://numpy.org/) and [matplotlib](https://matplotlib.org/)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VYTxx33-1-qC"
   },
   "source": [
    "## Loading Chemiscope widgets in Jupyter\n",
    "\n",
    "Please make sure you have jupyter extensions enabled.\n",
    "\n",
    "If at *any time* you are unable to load the chemiscope widgets in Jupyter, you can replace `chemiscope.show(` with `chemiscope.write_input('filename.json', ...` and upload the resulting file to [chemiscope.org](chemiscope.org)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "RYXoLmHs2DTm",
    "outputId": "1b03ee7f-2738-4ff2-8deb-e4b9bcab4258"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: ase in /usr/local/lib/python3.10/dist-packages (3.22.1)\n",
      "Requirement already satisfied: skmatter in /usr/local/lib/python3.10/dist-packages (0.1.4)\n",
      "Requirement already satisfied: chemiscope in /usr/local/lib/python3.10/dist-packages (0.5.2)\n",
      "Requirement already satisfied: dataframe-image in /usr/local/lib/python3.10/dist-packages (0.1.11)\n",
      "Requirement already satisfied: dscribe in /usr/local/lib/python3.10/dist-packages (2.0.0)\n",
      "Requirement already satisfied: matplotlib>=3.1.0 in /usr/local/lib/python3.10/dist-packages (from ase) (3.7.1)\n",
      "Requirement already satisfied: numpy>=1.15.0 in /usr/local/lib/python3.10/dist-packages (from ase) (1.22.4)\n",
      "Requirement already satisfied: scipy>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from ase) (1.10.1)\n",
      "Requirement already satisfied: scikit-learn>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from skmatter) (1.2.2)\n",
      "Requirement already satisfied: ipywidgets<9,>=7 in /usr/local/lib/python3.10/dist-packages (from chemiscope) (7.7.1)\n",
      "Requirement already satisfied: pandas>=0.24 in /usr/local/lib/python3.10/dist-packages (from dataframe-image) (1.5.3)\n",
      "Requirement already satisfied: nbconvert>=5 in /usr/local/lib/python3.10/dist-packages (from dataframe-image) (6.5.4)\n",
      "Requirement already satisfied: aiohttp in /usr/local/lib/python3.10/dist-packages (from dataframe-image) (3.8.4)\n",
      "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from dataframe-image) (2.27.1)\n",
      "Requirement already satisfied: pillow in /usr/local/lib/python3.10/dist-packages (from dataframe-image) (8.4.0)\n",
      "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from dataframe-image) (23.1)\n",
      "Requirement already satisfied: mistune in /usr/local/lib/python3.10/dist-packages (from dataframe-image) (0.8.4)\n",
      "Requirement already satisfied: lxml in /usr/local/lib/python3.10/dist-packages (from dataframe-image) (4.9.3)\n",
      "Requirement already satisfied: beautifulsoup4 in /usr/local/lib/python3.10/dist-packages (from dataframe-image) (4.11.2)\n",
      "Requirement already satisfied: cssutils in /usr/local/lib/python3.10/dist-packages (from dataframe-image) (2.7.1)\n",
      "Requirement already satisfied: html2image in /usr/local/lib/python3.10/dist-packages (from dataframe-image) (2.0.3)\n",
      "Requirement already satisfied: pybind11>=2.4 in /usr/local/lib/python3.10/dist-packages (from dscribe) (2.10.4)\n",
      "Requirement already satisfied: joblib>=1.0.0 in /usr/local/lib/python3.10/dist-packages (from dscribe) (1.3.1)\n",
      "Requirement already satisfied: sparse in /usr/local/lib/python3.10/dist-packages (from dscribe) (0.14.0)\n",
      "Requirement already satisfied: ipykernel>=4.5.1 in /usr/local/lib/python3.10/dist-packages (from ipywidgets<9,>=7->chemiscope) (5.5.6)\n",
      "Requirement already satisfied: ipython-genutils~=0.2.0 in /usr/local/lib/python3.10/dist-packages (from ipywidgets<9,>=7->chemiscope) (0.2.0)\n",
      "Requirement already satisfied: traitlets>=4.3.1 in /usr/local/lib/python3.10/dist-packages (from ipywidgets<9,>=7->chemiscope) (5.7.1)\n",
      "Requirement already satisfied: widgetsnbextension~=3.6.0 in /usr/local/lib/python3.10/dist-packages (from ipywidgets<9,>=7->chemiscope) (3.6.4)\n",
      "Requirement already satisfied: ipython>=4.0.0 in /usr/local/lib/python3.10/dist-packages (from ipywidgets<9,>=7->chemiscope) (7.34.0)\n",
      "Requirement already satisfied: jupyterlab-widgets>=1.0.0 in /usr/local/lib/python3.10/dist-packages (from ipywidgets<9,>=7->chemiscope) (3.0.8)\n",
      "Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3.1.0->ase) (1.1.0)\n",
      "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3.1.0->ase) (0.11.0)\n",
      "Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3.1.0->ase) (4.40.0)\n",
      "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3.1.0->ase) (1.4.4)\n",
      "Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3.1.0->ase) (3.1.0)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3.1.0->ase) (2.8.2)\n",
      "Requirement already satisfied: bleach in /usr/local/lib/python3.10/dist-packages (from nbconvert>=5->dataframe-image) (6.0.0)\n",
      "Requirement already satisfied: defusedxml in /usr/local/lib/python3.10/dist-packages (from nbconvert>=5->dataframe-image) (0.7.1)\n",
      "Requirement already satisfied: entrypoints>=0.2.2 in /usr/local/lib/python3.10/dist-packages (from nbconvert>=5->dataframe-image) (0.4)\n",
      "Requirement already satisfied: jinja2>=3.0 in /usr/local/lib/python3.10/dist-packages (from nbconvert>=5->dataframe-image) (3.1.2)\n",
      "Requirement already satisfied: jupyter-core>=4.7 in /usr/local/lib/python3.10/dist-packages (from nbconvert>=5->dataframe-image) (5.3.1)\n",
      "Requirement already satisfied: jupyterlab-pygments in /usr/local/lib/python3.10/dist-packages (from nbconvert>=5->dataframe-image) (0.2.2)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from nbconvert>=5->dataframe-image) (2.1.3)\n",
      "Requirement already satisfied: nbclient>=0.5.0 in /usr/local/lib/python3.10/dist-packages (from nbconvert>=5->dataframe-image) (0.8.0)\n",
      "Requirement already satisfied: nbformat>=5.1 in /usr/local/lib/python3.10/dist-packages (from nbconvert>=5->dataframe-image) (5.9.0)\n",
      "Requirement already satisfied: pandocfilters>=1.4.1 in /usr/local/lib/python3.10/dist-packages (from nbconvert>=5->dataframe-image) (1.5.0)\n",
      "Requirement already satisfied: pygments>=2.4.1 in /usr/local/lib/python3.10/dist-packages (from nbconvert>=5->dataframe-image) (2.14.0)\n",
      "Requirement already satisfied: tinycss2 in /usr/local/lib/python3.10/dist-packages (from nbconvert>=5->dataframe-image) (1.2.1)\n",
      "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas>=0.24->dataframe-image) (2022.7.1)\n",
      "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from scikit-learn>=1.1.0->skmatter) (3.1.0)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->dataframe-image) (23.1.0)\n",
      "Requirement already satisfied: charset-normalizer<4.0,>=2.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->dataframe-image) (2.0.12)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp->dataframe-image) (6.0.4)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /usr/local/lib/python3.10/dist-packages (from aiohttp->dataframe-image) (4.0.2)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->dataframe-image) (1.9.2)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp->dataframe-image) (1.3.3)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp->dataframe-image) (1.3.1)\n",
      "Requirement already satisfied: soupsieve>1.2 in /usr/local/lib/python3.10/dist-packages (from beautifulsoup4->dataframe-image) (2.4.1)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->dataframe-image) (1.26.16)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->dataframe-image) (2023.5.7)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->dataframe-image) (3.4)\n",
      "Requirement already satisfied: numba>=0.49 in /usr/local/lib/python3.10/dist-packages (from sparse->dscribe) (0.56.4)\n",
      "Requirement already satisfied: jupyter-client in /usr/local/lib/python3.10/dist-packages (from ipykernel>=4.5.1->ipywidgets<9,>=7->chemiscope) (6.1.12)\n",
      "Requirement already satisfied: tornado>=4.2 in /usr/local/lib/python3.10/dist-packages (from ipykernel>=4.5.1->ipywidgets<9,>=7->chemiscope) (6.3.1)\n",
      "Requirement already satisfied: setuptools>=18.5 in /usr/local/lib/python3.10/dist-packages (from ipython>=4.0.0->ipywidgets<9,>=7->chemiscope) (67.7.2)\n",
      "Requirement already satisfied: jedi>=0.16 in /usr/local/lib/python3.10/dist-packages (from ipython>=4.0.0->ipywidgets<9,>=7->chemiscope) (0.18.2)\n",
      "Requirement already satisfied: decorator in /usr/local/lib/python3.10/dist-packages (from ipython>=4.0.0->ipywidgets<9,>=7->chemiscope) (4.4.2)\n",
      "Requirement already satisfied: pickleshare in /usr/local/lib/python3.10/dist-packages (from ipython>=4.0.0->ipywidgets<9,>=7->chemiscope) (0.7.5)\n",
      "Requirement already satisfied: prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from ipython>=4.0.0->ipywidgets<9,>=7->chemiscope) (3.0.39)\n",
      "Requirement already satisfied: backcall in /usr/local/lib/python3.10/dist-packages (from ipython>=4.0.0->ipywidgets<9,>=7->chemiscope) (0.2.0)\n",
      "Requirement already satisfied: matplotlib-inline in /usr/local/lib/python3.10/dist-packages (from ipython>=4.0.0->ipywidgets<9,>=7->chemiscope) (0.1.6)\n",
      "Requirement already satisfied: pexpect>4.3 in /usr/local/lib/python3.10/dist-packages (from ipython>=4.0.0->ipywidgets<9,>=7->chemiscope) (4.8.0)\n",
      "Requirement already satisfied: platformdirs>=2.5 in /usr/local/lib/python3.10/dist-packages (from jupyter-core>=4.7->nbconvert>=5->dataframe-image) (3.8.1)\n",
      "Requirement already satisfied: fastjsonschema in /usr/local/lib/python3.10/dist-packages (from nbformat>=5.1->nbconvert>=5->dataframe-image) (2.17.1)\n",
      "Requirement already satisfied: jsonschema>=2.6 in /usr/local/lib/python3.10/dist-packages (from nbformat>=5.1->nbconvert>=5->dataframe-image) (4.3.3)\n",
      "Requirement already satisfied: llvmlite<0.40,>=0.39.0dev0 in /usr/local/lib/python3.10/dist-packages (from numba>=0.49->sparse->dscribe) (0.39.1)\n",
      "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.7->matplotlib>=3.1.0->ase) (1.16.0)\n",
      "Requirement already satisfied: notebook>=4.4.1 in /usr/local/lib/python3.10/dist-packages (from widgetsnbextension~=3.6.0->ipywidgets<9,>=7->chemiscope) (6.4.8)\n",
      "Requirement already satisfied: webencodings in /usr/local/lib/python3.10/dist-packages (from bleach->nbconvert>=5->dataframe-image) (0.5.1)\n",
      "Requirement already satisfied: parso<0.9.0,>=0.8.0 in /usr/local/lib/python3.10/dist-packages (from jedi>=0.16->ipython>=4.0.0->ipywidgets<9,>=7->chemiscope) (0.8.3)\n",
      "Requirement already satisfied: pyrsistent!=0.17.0,!=0.17.1,!=0.17.2,>=0.14.0 in /usr/local/lib/python3.10/dist-packages (from jsonschema>=2.6->nbformat>=5.1->nbconvert>=5->dataframe-image) (0.19.3)\n",
      "Requirement already satisfied: pyzmq>=13 in /usr/local/lib/python3.10/dist-packages (from jupyter-client->ipykernel>=4.5.1->ipywidgets<9,>=7->chemiscope) (23.2.1)\n",
      "Requirement already satisfied: argon2-cffi in /usr/local/lib/python3.10/dist-packages (from notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets<9,>=7->chemiscope) (21.3.0)\n",
      "Requirement already satisfied: nest-asyncio>=1.5 in /usr/local/lib/python3.10/dist-packages (from notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets<9,>=7->chemiscope) (1.5.6)\n",
      "Requirement already satisfied: Send2Trash>=1.8.0 in /usr/local/lib/python3.10/dist-packages (from notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets<9,>=7->chemiscope) (1.8.2)\n",
      "Requirement already satisfied: terminado>=0.8.3 in /usr/local/lib/python3.10/dist-packages (from notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets<9,>=7->chemiscope) (0.17.1)\n",
      "Requirement already satisfied: prometheus-client in /usr/local/lib/python3.10/dist-packages (from notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets<9,>=7->chemiscope) (0.17.0)\n",
      "Requirement already satisfied: ptyprocess>=0.5 in /usr/local/lib/python3.10/dist-packages (from pexpect>4.3->ipython>=4.0.0->ipywidgets<9,>=7->chemiscope) (0.7.0)\n",
      "Requirement already satisfied: wcwidth in /usr/local/lib/python3.10/dist-packages (from prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0->ipython>=4.0.0->ipywidgets<9,>=7->chemiscope) (0.2.6)\n",
      "Requirement already satisfied: argon2-cffi-bindings in /usr/local/lib/python3.10/dist-packages (from argon2-cffi->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets<9,>=7->chemiscope) (21.2.0)\n",
      "Requirement already satisfied: cffi>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from argon2-cffi-bindings->argon2-cffi->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets<9,>=7->chemiscope) (1.15.1)\n",
      "Requirement already satisfied: pycparser in /usr/local/lib/python3.10/dist-packages (from cffi>=1.0.1->argon2-cffi-bindings->argon2-cffi->notebook>=4.4.1->widgetsnbextension~=3.6.0->ipywidgets<9,>=7->chemiscope) (2.21)\n",
      "Reading package lists... Done\n",
      "Building dependency tree       \n",
      "Reading state information... Done\n",
      "cargo is already the newest version (0.67.1+ds0ubuntu0.libgit2-0ubuntu0.20.04.2).\n",
      "rustc is already the newest version (1.66.1+dfsg0ubuntu1~llvm-0ubuntu0.20.04).\n",
      "0 upgraded, 0 newly installed, 0 to remove and 15 not upgraded.\n",
      "Collecting git+https://github.com/Luthaf/rascaline.git\n",
      "  Cloning https://github.com/Luthaf/rascaline.git to /tmp/pip-req-build-wi5u0qxm\n",
      "  Running command git clone --filter=blob:none --quiet https://github.com/Luthaf/rascaline.git /tmp/pip-req-build-wi5u0qxm\n",
      "  Resolved https://github.com/Luthaf/rascaline.git to commit 881de9b0bb6c71a83b1867421c5cff0b2c997b42\n",
      "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
      "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
      "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
      "Collecting equistore-core@ https://github.com/lab-cosmo/equistore/archive/c022fde.zip#subdirectory=python/equistore-core (from rascaline==0.1.0.dev410)\n",
      "  Using cached https://github.com/lab-cosmo/equistore/archive/c022fde.zip\n",
      "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
      "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
      "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
      "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from equistore-core@ https://github.com/lab-cosmo/equistore/archive/c022fde.zip#subdirectory=python/equistore-core->rascaline==0.1.0.dev410) (1.22.4)\n",
      "Cloning into '5th_workshop_MachineLearning'...\n",
      "remote: Enumerating objects: 199, done.\u001b[K\n",
      "remote: Counting objects: 100% (43/43), done.\u001b[K\n",
      "remote: Compressing objects: 100% (29/29), done.\u001b[K\n",
      "remote: Total 199 (delta 30), reused 26 (delta 14), pack-reused 156\u001b[K\n",
      "Receiving objects: 100% (199/199), 99.50 MiB | 18.21 MiB/s, done.\n",
      "Resolving deltas: 100% (87/87), done.\n"
     ]
    }
   ],
   "source": [
    "!pip install ase skmatter chemiscope dataframe-image dscribe\n",
    "!apt install rustc cargo\n",
    "!pip install git+https://github.com/Luthaf/rascaline.git\n",
    "!git clone https://github.com/icomse/5th_workshop_MachineLearning.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "_uRwVMgI1-qD"
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "import numpy as np\n",
    "from ase.io import read\n",
    "from matplotlib import pyplot as plt\n",
    "import chemiscope\n",
    "import scipy\n",
    "\n",
    "import os\n",
    "\n",
    "from matplotlib.gridspec import GridSpec\n",
    "import pandas as pd\n",
    "from pandas.plotting import table\n",
    "from ase.geometry.analysis import Analysis as asis\n",
    "from ase.neighborlist import NeighborList as NL\n",
    "from ase.neighborlist import natural_cutoffs\n",
    "import dataframe_image as dfi\n",
    "\n",
    "names = {\"C\": \"Carbon\", \"H\": \"Hydrogen\"}\n",
    "colors = {\"H\": (0.6, 0.6, 0.6), \"C\": (0.2, 0.2, 0.2)}\n",
    "os.chdir('5th_workshop_MachineLearning/Day_3')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MdtiAjU61-qD"
   },
   "source": [
    "## Preparing the Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BqfRCiSn1-qE"
   },
   "source": [
    "### Read Data\n",
    "\n",
    "Here we read in 5 MD trajectories and place them in a concatenated list `traj`.\n",
    "\n",
    "`ranges` is storing the range of `traj` corresponding to each original file.\n",
    "`conf_idx` is storing the location of the initial conformations.\n",
    "\n",
    "`rgb_colors` is the set of colors used for each conformer, stored in rgba format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JQr3UBKH1-qE"
   },
   "outputs": [],
   "source": [
    "# read in the frames from each MD simulation\n",
    "traj = []\n",
    "names = [\"chair\", \"twist-boat\", \"boat\", \"half-chair\", \"planar\"]\n",
    "rgb_colors = [\n",
    "    (0.13333333333333333, 0.47058823529411764, 0.7098039215686275),\n",
    "    (0.4588235294117647, 0.7568627450980392, 0.34901960784313724),\n",
    "    (0.803921568627451, 0.6078431372549019, 0.16862745098039217),\n",
    "    (0.803921568627451, 0.13725490196078433, 0.15294117647058825),\n",
    "    (0.4392156862745098, 0.2784313725490196, 0.611764705882353),\n",
    "]\n",
    "\n",
    "ranges = np.zeros((len(names), 2), dtype=int)\n",
    "conf_idx = np.zeros(len(names), dtype=int)\n",
    "\n",
    "for i, n in enumerate(names):\n",
    "    frames = read(f\"./datasets/cyclohexane/{n}.xyz\",\":\",)\n",
    "\n",
    "    ranges[i] = (len(traj), len(traj) + len(frames))\n",
    "    conf_idx[i] = len(traj)\n",
    "    traj = [*traj, *frames]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XIkbQqZ81-qE"
   },
   "outputs": [],
   "source": [
    "# energies of the simulation frames, relative to the chair conformation\n",
    "energy = np.array([a.info[\"relative_energy_eV\"] for a in traj])\n",
    "\n",
    "# energies of the known conformers, relative to the chair conformation\n",
    "c_energy = np.array([traj[c].info[\"relative_energy_eV\"] for c in conf_idx])\n",
    "\n",
    "# extrema for the energies\n",
    "max_e = max(energy)\n",
    "min_e = min(energy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "R9ry8oSM1-qE"
   },
   "source": [
    "Here we can confirm what our analysis will tell us:\n",
    "\n",
    "- the simulation starting in the planar conformation transitions to the chair conformation\n",
    "- the simulations starting in the twist-boat, boat, and half-chair conformations ultimately get stuck in the twist formation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cgdtyz8r1-qE"
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, figsize=(6, 4))\n",
    "\n",
    "for n, c, r, rgb in zip(names, c_energy, ranges, rgb_colors):\n",
    "    ax.plot(\n",
    "        range(0, r[1] - r[0]), energy[r[0] : r[1]] - min_e, label=n, c=rgb, zorder=-1\n",
    "    )\n",
    "\n",
    "ax.legend()\n",
    "ax.set_xlabel(\"Simulation Timestep\")\n",
    "ax.set_ylabel(\"Energy\")\n",
    "\n",
    "ax.set_xlim([0, len(energy) // 5])\n",
    "ax.set_ylim([-0.1, 1.25 * (max_e - min_e)])\n",
    "ax.set_yticklabels([])\n",
    "\n",
    "plt.tight_layout()\n",
    "# plt.savefig('figures/Figure5/energy.png')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ezh6-Nva1-qE"
   },
   "source": [
    "# STOP! Up until now we've just loaded some pre-computed descriptors. But now... let's do this ourself!\n",
    "\n",
    "## Let's start with the first frame:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NAVcwLe61-qE"
   },
   "outputs": [],
   "source": [
    "frame = traj[0]\n",
    "\n",
    "positions = frame.positions\n",
    "positions -= positions.min(axis=0)\n",
    "order = list(sorted(range(len(frame)), key=lambda i: (-frame.numbers[i], *positions[i])))\n",
    "\n",
    "positions = np.round(positions, 3)[order].tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UmDG87IH1-qE"
   },
   "outputs": [],
   "source": [
    "nl = NL(\n",
    "    cutoffs=np.multiply(1, natural_cutoffs(frame)),\n",
    "    bothways=True,\n",
    "    self_interaction=False,\n",
    ")\n",
    "nl.update(frame)\n",
    "a = asis(frame, nl=nl)\n",
    "angles = a.get_angles(\"C\", \"C\", \"C\", unique=True)\n",
    "angle_vals = a.get_values(angles)\n",
    "\n",
    "bonds = a.get_bonds(\"C\",\"C\", unique=True)\n",
    "bond_vals = a.get_values(bonds)\n",
    "\n",
    "# positions\n",
    "positions_table = pd.DataFrame(\n",
    "    positions,\n",
    "    index=[\n",
    "        f\"{frame.symbols[i]}{frame.numbers[:i].tolist().count(n)+1}\"\n",
    "        for i, n in enumerate(frame.numbers)\n",
    "    ],\n",
    "    columns=[r\"$x$\", r\"$y$\", r\"$z$\"],\n",
    ")\n",
    "positions_table"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wRScGwYm1-qF"
   },
   "source": [
    "### These are our _internal coordinates_\n",
    "\n",
    "### What about our bond angles?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "G-49ndx01-qF"
   },
   "outputs": [],
   "source": [
    "# bond angles\n",
    "bond_table = pd.DataFrame(\n",
    "    np.transpose(np.round(angle_vals, 3)),\n",
    "    index=[\"C{}-C{}-C{}\".format(*np.sort(np.add(a, 1))) for a in angles[0]],\n",
    "    columns=[r\"Bond Angle ($^\\circ$)\"],\n",
    ")\n",
    "bond_table"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qJ8Eg6so1-qF"
   },
   "source": [
    "### Let's take a look at the distance matrix (known as the Z-matrix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Db3DEDyA1-qF"
   },
   "outputs": [],
   "source": [
    "distances = np.array(\n",
    "    [[np.linalg.norm(np.subtract(x, y)) for x in positions] for y in positions]\n",
    ")\n",
    "\n",
    "fig, (ax, cax) = plt.subplots(\n",
    "    2, 1, figsize=(4,6), gridspec_kw=dict(height_ratios=(1.0, 0.1))\n",
    ")\n",
    "p = ax.imshow(distances, cmap=\"bwr\", vmax=4.0)\n",
    "\n",
    "ax.axvline(5.5, c=\"k\", lw=2)\n",
    "ax.axhline(5.5, c=\"k\", lw=2)\n",
    "\n",
    "ax.annotate(\"C\", xy=(2.5, -0.5), xytext=(2.5, -1.0), ha=\"center\", va=\"center\")\n",
    "ax.annotate(\"C\", xy=(-0.5, 2.5), xytext=(-1.0, 2.5), ha=\"center\", va=\"center\")\n",
    "ax.annotate(\"H\", xy=(11.5, -0.5), xytext=(11.5, -1.5), ha=\"center\", va=\"center\")\n",
    "ax.annotate(\"H\", xy=(-0.5, 11.5), xytext=(-1.5, 11.5), ha=\"center\", va=\"center\")\n",
    "\n",
    "ax.set_xticks([])\n",
    "ax.set_yticks([])\n",
    "\n",
    "plt.colorbar(\n",
    "    p, label=r\"$d, (\\AA)$\", ax=ax, cax=cax, orientation=\"horizontal\", fraction=0.3\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "77_Mx70l1-qF"
   },
   "source": [
    "### If we know our neighborhood cutoffs, then we can make an _adjacency matrix_."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kpa-Yjgz1-qF"
   },
   "outputs": [],
   "source": [
    "# adjacency-matrix\n",
    "adjacency = np.array(\n",
    "    [[np.linalg.norm(np.subtract(x, y)) for x in positions] for y in positions]\n",
    ")\n",
    "adjacency[adjacency==0] = np.inf\n",
    "\n",
    "fig, ax = plt.subplots(\n",
    "    1, 1, figsize=(4,4)\n",
    ")\n",
    "\n",
    "# Here we are selecting a cutoff of 1.6 Angstrom\n",
    "p = ax.imshow(adjacency<1.6, cmap=\"Greys\", vmax=1.0)\n",
    "\n",
    "ax.axvline(5.5, c=\"k\", lw=2)\n",
    "ax.axhline(5.5, c=\"k\", lw=2)\n",
    "\n",
    "ax.annotate(\"C\", xy=(2.5, -0.5), xytext=(2.5, -1.0), ha=\"center\", va=\"center\")\n",
    "ax.annotate(\"C\", xy=(-0.5, 2.5), xytext=(-1.0, 2.5), ha=\"center\", va=\"center\")\n",
    "ax.annotate(\"H\", xy=(11.5, -0.5), xytext=(11.5, -1.5), ha=\"center\", va=\"center\")\n",
    "ax.annotate(\"H\", xy=(-0.5, 11.5), xytext=(-1.5, 11.5), ha=\"center\", va=\"center\")\n",
    "\n",
    "ax.set_xticks([])\n",
    "ax.set_yticks([])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Wb5JEg4x1-qF"
   },
   "source": [
    "Next, the coulomb matrix!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4Tg7jDnJ1-qF"
   },
   "outputs": [],
   "source": [
    "from dscribe.descriptors import CoulombMatrix\n",
    "\n",
    "cm = CoulombMatrix(n_atoms_max=18)\n",
    "pd.DataFrame(cm.create(frame).reshape(18,18), columns = frame.symbols, index=frame.symbols)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0U_1HWpG1-qG"
   },
   "source": [
    "### Question: How would you combine these descriptors from atom-level to become molecule-level?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lW8LLpSF1-qG"
   },
   "source": [
    "### Now go back to cell 5 (STOP!) and see how your answers change for other configurations. The chair configuration is in traj[0], twist-boat in traj[400], boat in traj[800], half-chair in traj[1200], and planar in traj[1600]."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MWYlguLW1-qG"
   },
   "source": [
    "# A small dive into SOAP"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rTKTYL6R1-qG"
   },
   "source": [
    "### For computing symmetrized representations, there are many, *many* different software packages that offer different benefits.\n",
    "\n",
    "- [dscribe](https://singroup.github.io/dscribe/) is very user-friendly if you want to use off-the-shelf SOAP PowerSpectra (but that's about all that's available here).\n",
    "- [librascal](https://github.com/lab-cosmo/librascal) is a power package for every variation of SOAP, Behler-Parinello symmetry functions, and other atom-centered representations. Unfortunately, its development has been discontinued since 2022, so proceed with caution.\n",
    "- [rascaline](https://github.com/luthaf/rascaline) is another power package for every variation of SOAP, Behler-Parinello symmetry functions, and other atom-centered representations, but requires knowledge of the [equistore](https://github.com/lab-cosmo/equistore) format and has a steep learning curve."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AdKOYYhb1-qH"
   },
   "source": [
    "## We're going to proceed cautiously with rascaline, but I encourage everyone to ask ample questions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "o5A6T8Vk1-qH"
   },
   "source": [
    "In symmetrized descriptors, we have _hyperparameters_ that we can tune to mimic the physics of our system. In a 2-body SOAP descriptor, we tune:\n",
    "\n",
    "- the `atomic_gaussian_width`, how large of a Gaussian we'd like to place on each atom, in Angstrom\n",
    "- the `cutoff`, how far away from each atom we'd like to integrate\n",
    "- the `radial_basis`, what set of [basis functions](https://chem.libretexts.org/Courses/Pacific_Union_College/Quantum_Chemistry/11%3A_Computational_Quantum_Chemistry/11.02%3A_Gaussian_Basis_Sets) to use for radial expansion\n",
    "- the `max_radial` number of these basis functions to use\n",
    "- the `cutoff_function` that we use to weight neighbors (this makes it so that atoms far away from our \"central\" one are weighted less)\n",
    "\n",
    "Here are some good starting parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sq064IQE1-qH"
   },
   "outputs": [],
   "source": [
    "from rascaline.calculators import SoapRadialSpectrum\n",
    "representation = SoapRadialSpectrum(\n",
    "    **{\n",
    "        \"atomic_gaussian_width\": 0.3,\n",
    "        \"max_radial\": 6,\n",
    "        \"cutoff\": 3.5,\n",
    "        \"radial_basis\": {\"Gto\": {}},\n",
    "        \"cutoff_function\": {\"ShiftedCosine\": {\"width\": 0.8}},\n",
    "        \"center_atom_weight\": 1.0,\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Uzcj8I741-qH"
   },
   "outputs": [],
   "source": [
    "rep = representation.compute(frame)\n",
    "\n",
    "# These two steps help with the equistore format\n",
    "rep = rep.keys_to_samples('species_center')\n",
    "rep = rep.keys_to_properties([\"species_neighbor\"])\n",
    "x = rep.block().values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wwE-MwHH1-qH"
   },
   "outputs": [],
   "source": [
    "pd.DataFrame(x, index=frame.symbols, columns=['{}-neighbor, n={}'.format(\"CH\"[i//6], i%6) for i in range(12)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oJQH-z8M1-qH"
   },
   "source": [
    "For every neighboring species, we have a histogram across radial bases of where they fall!\n",
    "\n",
    "But this is *VERY* hard to read in chart form, let's look at a plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "i9-at7Jd1-qH"
   },
   "outputs": [],
   "source": [
    "plt.imshow(x)\n",
    "plt.gca().set_yticks(np.arange(len(frame)))\n",
    "plt.gca().set_yticklabels(frame.symbols)\n",
    "plt.gca().set_ylabel(\"Focal atom\")\n",
    "\n",
    "plt.gca().set_xticks(np.arange(len(x[0])))\n",
    "plt.gca().set_xticklabels(['{}'.format(i%6) for i in range(12)])\n",
    "plt.gca().set_xlabel(\"n\")\n",
    "\n",
    "plt.annotate(\"C-Neighbors\",\n",
    "            xy=(0.25, 1), xycoords='axes fraction',\n",
    "            size=12, va=\"bottom\", ha=\"center\",)\n",
    "\n",
    "plt.annotate(\"H-Neighbors\",\n",
    "            xy=(0.75, 1), xycoords='axes fraction',\n",
    "            size=12, va=\"bottom\", ha=\"center\",)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YONTYoYQ1-qH"
   },
   "source": [
    "When you get to this point, turn to your neighbors to discuss. If you need something to do while you wait, this is a good, albeit lengthy read:\n",
    "    \n",
    "https://doi.org/10.1021/acs.chemrev.1c00021"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xc-Oxwz41-qH"
   },
   "source": [
    "## But physics rarely relies on solely pair terms (at least good physics...)\n",
    "\n",
    "When we go to higher body-orders, we also histogram over angular components through spherical harmonics. So we specify\n",
    "\n",
    "- the `max_angular` number of spherical harmonics we expand over\n",
    "\n",
    "The PowerSpectrum is a three-body descriptor, where we check two neighbors for each focal atom"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Uot9xpJp1-qI"
   },
   "outputs": [],
   "source": [
    "from rascaline.calculators import SoapPowerSpectrum\n",
    "representation = SoapPowerSpectrum(\n",
    "    **{\n",
    "        \"atomic_gaussian_width\": 0.3,\n",
    "        \"max_angular\": 4,\n",
    "        \"max_radial\": 6,\n",
    "        \"cutoff\": 3.5,\n",
    "        \"radial_basis\": {\"Gto\": {}},\n",
    "        \"cutoff_function\": {\"ShiftedCosine\": {\"width\": 0.8}},\n",
    "        \"center_atom_weight\": 1.0,\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fPiN8Hw01-qI"
   },
   "outputs": [],
   "source": [
    "rep = representation.compute(frame)\n",
    "\n",
    "rep = rep.keys_to_samples('species_center')\n",
    "\n",
    "# There are two neighbor species now!\n",
    "rep = rep.keys_to_properties([\"species_neighbor_1\", \"species_neighbor_2\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PRHOHQl81-qI"
   },
   "outputs": [],
   "source": [
    "plt.plot(rep.block().values.T)\n",
    "\n",
    "plt.gca().set_xticks([])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YHZQCil-1-qI"
   },
   "source": [
    "From a PCA, we can see that carbons are all very similar, but we have two types of hydrogens: in-plane and out-of-plane."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oeNb9jOY1-qI"
   },
   "outputs": [],
   "source": [
    "from sklearn.decomposition import PCA\n",
    "\n",
    "pca = PCA(n_components=2).fit(x)\n",
    "t = pca.transform(x)\n",
    "\n",
    "chemiscope.show([frame], properties={\"t\": t}, environments=[(0, i, 3.5) for i in range(len(frame))])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FyXU21ct1-qI"
   },
   "source": [
    "# Okay! I want to make SOAP descriptors like you did for the activities before. How do I do that?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YKvancpw1-qI"
   },
   "source": [
    "First, we're only going to select carbon atoms as our focal atoms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_KvW8qn41-qI"
   },
   "outputs": [],
   "source": [
    "from equistore.core import Labels\n",
    "\n",
    "values = []\n",
    "for i, _ in enumerate(traj):\n",
    "    for j in range(6): # carbons are indices 0-6\n",
    "        values.append([i,j])\n",
    "\n",
    "selection = Labels(\n",
    "    names=[\"structure\", \"center\"],\n",
    "    values=np.array(values),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EFEbaeiT1-qJ"
   },
   "outputs": [],
   "source": [
    "rep = representation.compute(traj, selected_samples=selection)\n",
    "\n",
    "rep = rep.keys_to_samples('species_center')\n",
    "rep = rep.keys_to_properties([\"species_neighbor_1\", \"species_neighbor_2\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EIKC4hG81-qJ"
   },
   "source": [
    "Next, we need to [standardize](https://scikit-matter.readthedocs.io/en/latest/examples/pcovr/PCovR_Scaling.html#sphx-glr-examples-pcovr-pcovr-scaling-py) our descriptors. Here we scale them all together, because their relative magnitude matters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "olIAYY3y1-qJ"
   },
   "outputs": [],
   "source": [
    "from skmatter.preprocessing import StandardFlexibleScaler\n",
    "x = StandardFlexibleScaler(column_wise=False).fit_transform(rep.block().values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ItIDAtFb1-qJ"
   },
   "source": [
    "Last, we'll split them by frame to help with later operations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "i6PSssxt1-qJ"
   },
   "outputs": [],
   "source": [
    "split_soaps = np.array(np.split(x, len(traj)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FRhdidk-1-qJ"
   },
   "source": [
    "# Et voila!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xYbXJRsy1-qJ"
   },
   "outputs": [],
   "source": [
    "np.linalg.norm(split_soaps - np.load('./datasets/cyclohexane/cyclohexane_descriptors.npy'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PvRHO4vO1-qJ"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "hide_input": false,
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.6"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "186.797px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}