{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n\n# Substituting Atoms in a POSCAR File\n\nIn this example, we'll demonstrate how to substitute atoms in a POSCAR file using the `pyprocar` package. \nSpecifically, we will:\n\n1. Read a POSCAR file containing atomic positions and lattice vectors.\n2. Visualize the initial atomic positions.\n3. Substitute a boron (B) atom with a nitrogen (N) atom.\n4. Visualize the atomic positions after the substitution.\n5. Create GIFs to visualize the atomic structures from different angles.\n\nBefore diving in, make sure to have `pyvista`, `numpy`, and `pyprocar` installed.\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\n# Define the data directory\ndata_dir = os.path.join(DATA_DIR, \"examples\", \"PyPoscar\", \"00-Poscar\")\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": [
        "## Reading and Parsing the POSCAR File\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Read the POSCAR file\na = p.Poscar(os.path.join(data_dir, \"POSCAR-YB6.vasp\"))\na.parse()\n\n# Display lattice, elements, and positions\nprint(\"Lattice:\")\nprint(a.lat)\nprint(\"\\nElements:\")\nprint(a.elm)\nprint(\"\\nPositions in Direct coordinates\")\nprint(a.dpos)\n\n\n# Convert positions to Cartesian coordinates for visualization\natoms_before = pv.PolyData(np.dot(a.dpos, a.lat))\natoms_before[\"atoms\"] = a.elm\nlabels_before = [elm + f\":{point}\" for elm, point in zip(a.elm, a.dpos)]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Atom Substitution\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Substitute the second B atom with an N atom\nprint(\"\\n\\nChanging the the second atom from B to N\\n\" + \"-\" * 40)\nnew_pos = a.dpos[1]\na.remove(atoms=1)\na.add(position=new_pos, element=\"N\", direct=True)\nprint(\"\\nElements after substitution:\")\nprint(a.elm)\nprint(\"\\nPositions in Direct coordinates after substitution:\")\nprint(a.dpos)\n\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 = pv.PolyData(unit_cell)\n\n# Convert positions to Cartesian coordinates for visualization\natoms_after = pv.PolyData(np.dot(a.dpos, a.lat))\natoms_after[\"atoms\"] = a.elm\nlabels_after = [elm + f\":{point}\" for elm, point in zip(a.elm, a.dpos)]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Visualization of Atomic Structures\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Visualize the atomic structures side by side\nplotter = pv.Plotter(shape=(1, 2), border=False)\n# Before substitution\nplotter.subplot(0, 0)\nplotter.add_mesh(\n    unit_cell.delaunay_3d().extract_feature_edges(),\n    color=\"black\",\n    line_width=5,\n    render_lines_as_tubes=True,\n)\nplotter.add_point_labels(\n    points=atoms_before.points,\n    labels=labels_before,\n    show_points=False,\n    always_visible=True,\n)\nplotter.add_mesh(\n    atoms_before,\n    scalars=\"atoms\",\n    point_size=30,\n    render_points_as_spheres=True,\n    show_scalar_bar=False,\n)\nplotter.add_title(\"Before substitution\")\n# After substitution\nplotter.subplot(0, 1)\nplotter.add_mesh(\n    unit_cell.delaunay_3d().extract_feature_edges(),\n    color=\"black\",\n    line_width=5,\n    render_lines_as_tubes=True,\n)\nplotter.add_point_labels(\n    points=atoms_after.points,\n    labels=labels_after,\n    show_points=False,\n    always_visible=True,\n)\nplotter.add_mesh(\n    atoms_after,\n    scalars=\"atoms\",\n    point_size=30,\n    render_points_as_spheres=True,\n    show_scalar_bar=False,\n)\nplotter.add_title(\"After substitution\")\nplotter.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Creating GIFs for Visualization\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Define a function to create a GIF visualization of the atomic structure\ndef create_gif(atoms, labels, unit_cell, 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=\"atoms\",\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()\n\n\n# Create GIFs for atomic structures before and after substitution\ncreate_gif(\n    atoms=atoms_before,\n    labels=labels_before,\n    unit_cell=unit_cell,\n    save_file=\"atoms_before.gif\",\n)\ncreate_gif(\n    atoms=atoms_after,\n    labels=labels_after,\n    unit_cell=unit_cell,\n    save_file=\"atoms_after.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
}