{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a0000001",
   "metadata": {},
   "source": [
    "# ChemView\n",
    "\n",
    "``ChemView`` is a wrapper around [Chemiscope](https://chemiscope.org/) for visualising\n",
    "ASE trajectories in Jupyter notebooks. It computes geometric and atomic properties from\n",
    "``ase.Atoms`` objects and renders them as a linked scatter plot and 3-D structure viewer.\n",
    "\n",
    "Clicking a point on the scatter plot loads the corresponding frame. The trajectory\n",
    "player at the bottom of the viewer animates through all frames."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "64473a82",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ase.io import read\n",
    "from sparc.src.utils.chemview import ChemView"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0000002",
   "metadata": {},
   "source": [
    "### Load a trajectory\n",
    "\n",
    "Any ASE-readable format works. Here we load a molecular dynamics trajectory stored\n",
    "as an ASE ``.traj`` file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "1d66f812",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total frames: 54\n"
     ]
    }
   ],
   "source": [
    "traj = read('../tests/iter_000000/00.dft/AseMD.traj', index=':')\n",
    "print(f\"Total frames: {len(traj)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0000003",
   "metadata": {},
   "source": [
    "### Energy vs. frame\n",
    "\n",
    "The simplest usage: plot potential energy as a function of frame index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "1eb425d2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aeab4908bfe1411991a6a4defebbc98d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/html": [
       "<ChemiscopeWidget(meta={'name': 'ChemView'}, settings={'target': 'structure', 'structure': [{'unitCell': False, 'space...}, structures=[{'size': 3, 'data': 'structure-0'}, {'size': 3, 'data': 'struct...}, environments=[{'structure': 0, 'center': 0, 'cutoff': 3.5}, {'structure': 0, ...}, properties={'frame': {'target': 'structure', 'values': [0.0, 1.0, 2.0, 3.0,...}, )>"
      ],
      "text/plain": [
       "<ChemiscopeWidget(meta={'name': 'ChemView'}, settings={'target': 'structure', 'structure': [{'unitCell': False…"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ChemView(traj, specs=[\"frame\", \"energy\"], x=\"frame\", y=\"energy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0000004",
   "metadata": {},
   "source": [
    "### Geometric properties with colour mapping\n",
    "\n",
    "Multiple geometric specs can be combined. Here we plot the O–H distance\n",
    "against the H–O–H angle and colour points by energy. Auto-generated property\n",
    "names follow the pattern ``distance_i_j`` and ``angle_i_j_k``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "08f12be1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f21ee9288126460b993b5b075fbaefd8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/html": [
       "<ChemiscopeWidget(meta={'name': 'ChemView'}, settings={'target': 'structure', 'structure': [{'unitCell': False, 'space...}, structures=[{'size': 3, 'data': 'structure-0'}, {'size': 3, 'data': 'struct...}, environments=[{'structure': 0, 'center': 0, 'cutoff': 3.5}, {'structure': 0, ...}, properties={'frame': {'target': 'structure', 'values': [0.0, 1.0, 2.0, 3.0,...}, )>"
      ],
      "text/plain": [
       "<ChemiscopeWidget(meta={'name': 'ChemView'}, settings={'target': 'structure', 'structure': [{'unitCell': False…"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ChemView(traj, specs=[\"frame\", \"energy\", \"distance:0,2\", \"distance:0,1\", \"angle:0,1,2\"], x=\"distance_0_2\", y=\"angle_0_1_2\", color=\"energy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0000005",
   "metadata": {},
   "source": [
    "### Atom-level forces\n",
    "\n",
    "Atom-level specs such as ``force_norm`` produce one value per atom per frame. Setting  ``color_atoms=\"atom_index\"`` colours atoms by their index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "228be46c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d9aded8c75c64f47a1a4793654162863",
       "version_major": 2,
       "version_minor": 0
      },
      "text/html": [
       "<ChemiscopeWidget(meta={'name': 'ChemView'}, settings={'target': 'structure', 'structure': [{'unitCell': False, 'space...}, structures=[{'size': 3, 'data': 'structure-0'}, {'size': 3, 'data': 'struct...}, environments=[{'structure': 0, 'center': 0, 'cutoff': 3.5}, {'structure': 0, ...}, properties={'frame': {'target': 'structure', 'values': [0.0, 1.0, 2.0, 3.0,...}, )>"
      ],
      "text/plain": [
       "<ChemiscopeWidget(meta={'name': 'ChemView'}, settings={'target': 'structure', 'structure': [{'unitCell': False…"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ChemView(traj, specs=[\"frame\", \"energy\", \"force_norm\"], x=\"frame\", y=\"energy\", color_atoms=\"atom_index\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0000006",
   "metadata": {},
   "source": [
    "### Cell parameters\n",
    "\n",
    "For periodic systems, lattice vector lengths and cell volume can be tracked\n",
    "across frames."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "bbe86790",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d11fe37c10934f9795bedfde93ab25da",
       "version_major": 2,
       "version_minor": 0
      },
      "text/html": [
       "<ChemiscopeWidget(meta={'name': 'ChemView'}, settings={'target': 'structure', 'structure': [{'unitCell': False, 'space...}, structures=[{'size': 3, 'data': 'structure-0'}, {'size': 3, 'data': 'struct...}, environments=[{'structure': 0, 'center': 0, 'cutoff': 3.5}, {'structure': 0, ...}, properties={'frame': {'target': 'structure', 'values': [0.0, 1.0, 2.0, 3.0,...}, )>"
      ],
      "text/plain": [
       "<ChemiscopeWidget(meta={'name': 'ChemView'}, settings={'target': 'structure', 'structure': [{'unitCell': False…"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ChemView(traj, specs=[\"frame\", \"volume\", \"cell_a\", \"cell_b\", \"cell_c\"], x=\"frame\", y=\"volume\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0000007",
   "metadata": {},
   "source": [
    "### Single structure\n",
    "\n",
    "A single ``ase.Atoms`` object can be passed directly. Set ``plot=False`` to\n",
    "show only the 3-D viewer without a scatter plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "f53a195b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8cde4e40074c44abb9cc4b55f48000bb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/html": [
       "<StructureWidget(meta={'name': 'ChemView'}, settings={'target': 'structure', 'structure': [{'unitCell': False, 'space...}, structures=[{'size': 3, 'data': 'structure-0'}], environments=[{'structure': 0, 'center': 0, 'cutoff': 3.5}, {'structure': 0, ...}, properties={'x_pos': {'target': 'atom', 'values': [0.0, 0.0, 0.0], 'units':...}, )>"
      ],
      "text/plain": [
       "<StructureWidget(meta={'name': 'ChemView'}, settings={'target': 'structure', 'structure': [{'unitCell': False,…"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from ase.build import molecule\n",
    "water = molecule(\"H2O\")\n",
    "ChemView(water, specs=[\"x_pos\", \"y_pos\", \"z_pos\"], plot=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0000008",
   "metadata": {},
   "source": [
    "### Custom property names\n",
    "\n",
    "When the same spec type appears more than once (e.g. two distances), use\n",
    "``names`` to assign distinct keys. The ``size`` parameter scales scatter-plot\n",
    "points by a property value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "4acdb7e5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "adccdb9996eb4ddbb229d4a5f40996c0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/html": [
       "<ChemiscopeWidget(meta={'name': 'ChemView'}, settings={'target': 'structure', 'structure': [{'unitCell': False, 'space...}, structures=[{'size': 3, 'data': 'structure-0'}, {'size': 3, 'data': 'struct...}, environments=[{'structure': 0, 'center': 0, 'cutoff': 3.5}, {'structure': 0, ...}, properties={'energy': {'target': 'structure', 'values': [-14.31756055, -14....}, )>"
      ],
      "text/plain": [
       "<ChemiscopeWidget(meta={'name': 'ChemView'}, settings={'target': 'structure', 'structure': [{'unitCell': False…"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ChemView(traj, specs=[\"energy\", \"distance:0,1\", \"distance:0,2\"], names=[\"energy\", \"d_OH1\", \"d_OH2\"], x=\"d_OH1\", y=\"d_OH2\", size=\"energy\", color=\"d_OH2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0000009",
   "metadata": {},
   "source": [
    "### Periodic bulk structure\n",
    "\n",
    "``ChemView`` also works with bulk periodic systems. Here we build a 2×2×2\n",
    "FCC copper supercell and view the atomic positions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "6de5e766",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7448ecaa86a04eac96cacb1336244474",
       "version_major": 2,
       "version_minor": 0
      },
      "text/html": [
       "<StructureWidget(meta={'name': 'ChemView'}, settings={'target': 'structure', 'structure': [{'unitCell': False, 'space...}, structures=[{'size': 32, 'data': 'structure-0'}], environments=[{'structure': 0, 'center': 0, 'cutoff': 3.5}, {'structure': 0, ...}, properties={'x_pos': {'target': 'atom', 'values': [0.0, 1.805, 1.805, 0.0, ...}, )>"
      ],
      "text/plain": [
       "<StructureWidget(meta={'name': 'ChemView'}, settings={'target': 'structure', 'structure': [{'unitCell': False,…"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from ase.lattice.cubic import FaceCenteredCubic\n",
    "\n",
    "system = FaceCenteredCubic(directions=[[1,0,0],[0,1,0],[0,0,1]], symbol='Cu', size=(2,2,2), pbc=True)\n",
    "ChemView(system, specs=[\"x_pos\", \"y_pos\", \"z_pos\"], plot=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b43fcf23",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4604c5b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83b1a872",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4404d9a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sparcGNN",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}