{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n\n# Visualizing Clusters in Bi2Se3 Slab\n\nIn this example, we'll explore the atomic clusters within a slab of Bi2Se3. The slab structure is formed by quintuple layers (QL) where each QL can be regarded as a cluster. We'll:\n\n1. Parse the POSCAR file to obtain atomic positions.\n2. Identify and label the atomic clusters within the slab.\n3. Visualize the atomic clusters using `pyvista`.\n\nLet's get started!\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import os\nfrom itertools import product\n\nimport numpy as np\nimport pyvista as pv\n\nimport pyprocar.pyposcar as p\nfrom pyprocar.utils import DATA_DIR\n\ndata_dir = os.path.join(DATA_DIR, \"examples\", \"PyPoscar\", \"05-clusters\")\n# You do not need this. This is to ensure an image is rendered off screen when generating exmaple gallery.\npv.OFF_SCREEN = True"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Utility function for creating GIF visualizations\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def create_gif(atoms, labels, unit_cell, scalars, save_file):\n    plotter = pv.Plotter()\n    title = save_file.split(os.sep)[-1].split(\".\")[0]\n    plotter.add_title(title)\n    plotter.add_mesh(\n        unit_cell.delaunay_3d().extract_feature_edges(),\n        color=\"black\",\n        line_width=5,\n        render_lines_as_tubes=True,\n    )\n    plotter.add_point_labels(\n        points=atoms.points, labels=labels, show_points=False, always_visible=True\n    )\n    plotter.add_mesh(\n        atoms,\n        scalars=scalars,\n        point_size=30,\n        render_points_as_spheres=True,\n        show_scalar_bar=False,\n    )\n    path = plotter.generate_orbital_path(n_points=36)\n    plotter.open_gif(os.path.join(data_dir, save_file))\n    plotter.orbit_on_path(path, write_frames=True, viewup=[0, 0, 1], step=0.05)\n    plotter.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Parsing the POSCAR and Identifying Clusters\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print(\"Loading a slab of Bi2Se3\")\na = p.poscar.Poscar(os.path.join(data_dir, \"Bi2Se3.vasp\"), verbose=False)\na.parse()\n\nprint(\n    \"It structure is formed by the so-called Quintuple layers (QL).\"\n    \" Each QL can be regarded as a `cluster`\"\n)\nprint(\"The Quintuple layers are\")\nc = p.Clusters(a)\nprint(c.clusters)\n\nclusters_mapping = {}\nfor i, cluster in enumerate(c.clusters):\n    for i_atom in cluster:\n        clusters_mapping.update({i_atom: i})\nclusters_array = np.array([clusters_mapping[i] for i in range(len(a.elm))])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Visualizing the Clusters\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Convert positions to Cartesian coordinates for visualization\natoms_before = pv.PolyData(np.dot(a.dpos, a.lat))\natoms_before[\"atoms\"] = a.elm\n\natoms_before[\"clusters\"] = clusters_array\nlabels_before = [elm for elm, point in zip(a.elm, a.dpos)]\nlabels_after = clusters_array\n# Define the unit cell using lattice vectors\nunit_cell_comb = list(product([0, 1], repeat=3))\nunit_cell = np.array(\n    [\n        comb[0] * a.lat[0] + comb[1] * a.lat[1] + comb[2] * a.lat[2]\n        for comb in unit_cell_comb\n    ]\n)\nunit_cell_before = pv.PolyData(unit_cell)\n\n\ncreate_gif(\n    atoms=atoms_before,\n    labels=labels_before,\n    unit_cell=unit_cell_before,\n    scalars=\"atoms\",\n    save_file=\"atoms_before_cluster.gif\",\n)\ncreate_gif(\n    atoms=atoms_before,\n    labels=labels_after,\n    unit_cell=unit_cell_before,\n    scalars=\"clusters\",\n    save_file=\"atoms_after_cluster.gif\",\n)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "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.9.20"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}