{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ff2b43f9",
   "metadata": {},
   "source": [
    "# Quick Start\n",
    "\n",
    "## Introduction\n",
    "In this tutorial, you will get a quick view of `weas-widget`. To run this tutorial, you need to install `weas-widget`. Open a terminal and run:\n",
    "\n",
    "```console\n",
    "pip install weas-widget\n",
    "```\n",
    "\n",
    "### \n",
    "\n",
    "## Load structure\n",
    "One can load a structure from ASE or Pymatgen."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e37d2ab4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ase.build import molecule\n",
    "from weas_widget import WeasWidget\n",
    "atoms = molecule(\"C2H6SO\")\n",
    "viewer = WeasWidget()\n",
    "viewer.from_ase(atoms)\n",
    "viewer\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "145edf54",
   "metadata": {},
   "source": [
    "## Edit the structure with mouse and keyboard\n",
    "WEAS supports editing the atoms directly in the GUI and synchronizing with the structure of the Python object.\n",
    "\n",
    "### Select Atoms\n",
    "There are three methods for selecting atoms:\n",
    "- Single Selection: Click directly on an atom to select it.\n",
    "- Box Selection: Hold the `Shift` key and drag to select a group of atoms.\n",
    "- Lasso Selection: Hold `Shift + Alt` and drag to select a group of atoms.\n",
    "\n",
    "\n",
    "### Move, Rotate selected atoms\n",
    "\n",
    "Press the transform shortcut, and move your mouse.\n",
    "\n",
    "|Operation | Shortcut|\n",
    "|----------|---------|\n",
    "| Move     | `g`   |\n",
    "| Rotate   | `r`   |\n",
    "| Duplicate| `d`   |\n",
    "\n",
    "\n",
    "### Delete selected atoms\n",
    "Press the ``Delete`` key to delete the selected atoms\n",
    "\n",
    "\n",
    "### Export edited atoms\n",
    "One can export the edited atoms to ASE or Pymatgen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40d7764b",
   "metadata": {},
   "outputs": [],
   "source": [
    "viewer.to_ase()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31abe86e",
   "metadata": {},
   "source": [
    "\n",
    "## Change the style\n",
    "\n",
    "### GUI\n",
    "On the top left, click `open controls`. One can change the viewer settings directly in the GUI:\n",
    "\n",
    "- model style\n",
    "    - Ball\n",
    "    - Ball and Stick\n",
    "    - Polyhedra\n",
    "    - Stick only\n",
    "- color type\n",
    "    - CPK\n",
    "    - VESTA\n",
    "    - JMOL\n",
    "- material type\n",
    "    - Standard\n",
    "    - Phong\n",
    "    - Basic\n",
    "- atom label\n",
    "    - None\n",
    "    - Atom sybol\n",
    "    - Atom index\n",
    "\n",
    "### Python\n",
    "One can use Python change the viewer settings. For example, change atoms style\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "add547b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ball and stick\n",
    "viewer.avr.model_style = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "159dcdbb",
   "metadata": {},
   "source": [
    "One set the style for individue atom"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2dddf63",
   "metadata": {},
   "outputs": [],
   "source": [
    "# only draw stick (bond) for first four atoms\n",
    "viewer.avr.model_sticks = [1, 1, 1, 1, 0, 0, 0, 0]\n",
    "viewer.avr.draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "895b75ae",
   "metadata": {},
   "source": [
    "Set scale for each atoms:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d53f98f5-5683-438e-9393-04900505f0f8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "viewer.avr.atom_scales = [1, 1, 1, 1, 1, 0.6, 0.6, 0.6, 1.5, 1.5]\n",
    "viewer.avr.draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b12b2131",
   "metadata": {},
   "source": [
    "### Select atoms\n",
    "One can get the indices of the selected atoms by:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4dcca10-f4b5-4495-8767-24f844e9e760",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "viewer.avr.selected_atoms_indices"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0441f64",
   "metadata": {},
   "source": [
    "One can set the indices of the selected atoms, thus highlight them by:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05651abf-6e18-43a6-a653-709f8c72bb27",
   "metadata": {},
   "outputs": [],
   "source": [
    "viewer.avr.selected_atoms_indices = [0, 1, 2]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ad1e44b",
   "metadata": {},
   "source": [
    "## Crystal\n",
    "For a nice visualization of a crystal, one usually shows the polyhedra and the atoms on the unit cell boundary, as well as the bonded atoms outside the cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d362b850-867b-4784-83f1-a92eb1372df0",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from weas_widget import WeasWidget\n",
    "viewer1 = WeasWidget()\n",
    "viewer1.load_example(\"tio2.cif\")\n",
    "# show polyhedra\n",
    "viewer1.avr.model_style = 2\n",
    "# show boundary atoms\n",
    "viewer1.avr.boundary = [[-0.1, 1.1], [-0.1, 1.1], [-0.1, 1.1]]\n",
    "# show bonded atoms outside the cell\n",
    "viewer1.avr.show_bonded_atoms = True\n",
    "# Change color tyoe to \"VESTA\"\n",
    "viewer1.avr.color_type = \"VESTA\"\n",
    "viewer1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9a536fb",
   "metadata": {},
   "source": [
    "## Animation\n",
    "One can read a trajectory file, such as an xyz file, for animation. You can play the animation using timeline bar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d564a07",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ase.build import molecule\n",
    "from weas_widget import WeasWidget\n",
    "atoms = molecule(\"C2H6SO\")\n",
    "# create a list of atoms by rotating it.\n",
    "images = []\n",
    "for i in range(0, 360, 10):\n",
    "    new_atoms = atoms.copy()\n",
    "    new_atoms.rotate(\"z\", i)\n",
    "    images.append(new_atoms)\n",
    "viewer = WeasWidget()\n",
    "viewer.from_ase(images)\n",
    "viewer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d383600c",
   "metadata": {},
   "source": [
    "## Isosurface\n",
    "Here is an example of drawing isosurfaces for HOMO of H2O molecule."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "36503e95",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b7d69dd0dc3a415b8a3be9d6f260b134",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "WeasWidget(atomScales=[1, 1, 1], atoms={'species': {'O': ['O', 8], 'H': ['H', 1]}, 'cell': [10.222218836840568…"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from weas_widget import WeasWidget\n",
    "from ase.io.cube import read_cube_data\n",
    "volume, atoms = read_cube_data(\"../../examples/h2o-homo.cube\")\n",
    "viewer = WeasWidget()\n",
    "viewer.from_ase(atoms)\n",
    "viewer.avr.iso.volumetric_data = {\"values\": volume}\n",
    "viewer.avr.iso.settings = {\"positive\": {\"isovalue\": 0.001},\n",
    "                           \"negative\": {\"isovalue\": -0.001, \"color\": \"yellow\"},}\n",
    "viewer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06c38890",
   "metadata": {},
   "source": [
    "### Vector Field\n",
    "#### Magnetic moment\n",
    "Show the magnetic moment as a vector field.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58259e80",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ase.build import bulk\n",
    "from weas_widget import WeasWidget\n",
    "import numpy as np\n",
    "atoms = bulk(\"Fe\", cubic=True)\n",
    "atoms*=[2, 2, 1]\n",
    "atoms.set_array(\"moment\", np.ones(len(atoms)))\n",
    "viewer = WeasWidget()\n",
    "viewer.from_ase(atoms)\n",
    "viewer.avr.model_style = 1\n",
    "viewer.camera.setting = {\"direction\": [0, -1, 0]}\n",
    "viewer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf942ff8",
   "metadata": {},
   "source": [
    "One can add any vector field by:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34875083",
   "metadata": {},
   "outputs": [],
   "source": [
    "# origin of the vector\n",
    "origins = atoms.positions\n",
    "# the vertor\n",
    "vectors = [[0, 0, 1]]*len(atoms)\n",
    "viewer.avr.vf.settings[\"test\"]= {\"origins\": origins, \"vectors\": vectors, \"color\": \"red\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d2569f1",
   "metadata": {},
   "source": [
    "#### Phonon visualization\n",
    "One can visualize the phonon dispersion via lattice vibrations. One only need to use the eigenstates (calculated with an external software) to generate the trajectory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8d60bac",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from ase.build import bulk\n",
    "from weas_widget import WeasWidget\n",
    "atoms = bulk(\"Fe\", cubic=True)\n",
    "phonon_setting = {\"eigenvectors\": np.array([[[0, 0], [0, 0],[0.5, 0]],\n",
    "                                    [[0, 0], [0, 0], [-0.5, 0]]]\n",
    "                                   ),\n",
    "        \"kpoint\": [0, 0, 0], # optional\n",
    "        \"amplitude\": 5, # scale the motion of the atoms\n",
    "        \"factor\": 1.5, # scale the length of the arrows\n",
    "        \"nframes\": 20,\n",
    "        \"repeat\": [4, 4, 1],\n",
    "        \"color\": \"blue\",\n",
    "        \"radius\": 0.1,\n",
    "        }\n",
    "viewer = WeasWidget()\n",
    "viewer.from_ase(atoms)\n",
    "viewer.avr.phonon_setting = phonon_setting\n",
    "viewer.camera.setting = {\"direction\": [0, -1, 0]}\n",
    "viewer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98817cd6",
   "metadata": {},
   "source": [
    "#### Color by attribute\n",
    "Color by attribute is a powerful tool to visualize the data. For example, color the atoms by their index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec316530",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ase.build import molecule\n",
    "from weas_widget import WeasWidget\n",
    "atoms = molecule(\"C2H6SO\")\n",
    "viewer = WeasWidget()\n",
    "viewer.from_ase(atoms)\n",
    "viewer.avr.model_style = 1\n",
    "viewer.avr.color_by = \"Index\"\n",
    "viewer.avr.color_ramp = [\"red\", \"yellow\", \"blue\"]\n",
    "viewer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69f8b8f3",
   "metadata": {},
   "source": [
    "### Camera setting\n",
    "\n",
    "One can set the direction and zoom of the camera:\n",
    "\n",
    "```python\n",
    "viewer.camera.setting = {\"direction\": [0, 5, 1], \"zoom\": 2}\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ede9f9c6",
   "metadata": {},
   "source": [
    "### Export image\n",
    "\n",
    "Save image to a path by:\n",
    "\n",
    "```python\n",
    "viewer.save_image(\"/homg/xing/filename.png\")\n",
    "```\n",
    "\n",
    "\n",
    "\n",
    "Trigger the download panel by:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b854bffa-3ef6-4dd2-bfc5-73856d7b1076",
   "metadata": {},
   "source": [
    "## Real world example\n",
    "### Adsorption"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa2ade62-65e6-49cc-8721-124f57b1e9ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ase.build import surface, molecule\n",
    "from weas_widget import WeasWidget\n",
    "au111 = surface(\"Au\", (1, 1, 1), 4, vacuum=5.0)\n",
    "co = molecule(\"CO\")\n",
    "au111 += co\n",
    "viewer = WeasWidget()\n",
    "viewer.from_ase(au111)\n",
    "viewer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdb10de0-e02b-4488-a710-d81303a490d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "au111_co = viewer.to_ase()\n",
    "au111_co.write(\"au111_co.xyz\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d482e01-a71a-4f8d-a7e4-070086b72f88",
   "metadata": {},
   "source": [
    "### Large system\n",
    "WEAS can handle thousands of atoms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63afd654-2a38-4488-a3e7-48f339200a02",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from ase.build import surface, molecule\n",
    "from weas_widget import WeasWidget\n",
    "co = molecule(\"CO\")\n",
    "co.center(3.0)\n",
    "co.pbc = [True, True, True]\n",
    "co *= [10, 10, 10]\n",
    "print(\"Number of atoms: \", len(co))\n",
    "viewer2 = WeasWidget()\n",
    "viewer2.from_ase(co)\n",
    "viewer2"
   ]
  }
 ],
 "metadata": {
  "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}