{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n\n# Modifying a POSCAR File: Scaling, Supercells, and Defects\n\nIn this example, we'll demonstrate several modifications on a POSCAR file using the `pyprocar` package:\n\n1. Scaling the lattice vectors to reduce vacuum space.\n2. Generating a supercell.\n3. Introducing defects by changing atom types.\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\", \"01-poscar_utils\")\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, 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()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Scaling Vacuum Space in the Lattice\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "a = p.poscarUtils.poscar_modify(\n    os.path.join(data_dir, \"POSCAR-9AGNR.vasp\"), verbose=False\n)\nprint(\"the lattice has too much vacuum space\\n\", a.p.lat)\nprint(\"I will shrink these vector by 1/3\")\n\nscaling = np.array([1, 1 / 3, 1 / 3])\na.scale_lattice(factor=scaling, keep_cartesian=True)\na.write(os.path.join(data_dir, \"POSCAR-9AGNR-smaller.vasp\"))\nprint(\"New lattice\\n\", a.p.lat)\n\ntmp_a = p.Poscar(os.path.join(data_dir, \"POSCAR-9AGNR.vasp\"))\ntmp_a.parse()\n\n# Convert positions to Cartesian coordinates for visualization\natoms_before = pv.PolyData(np.dot(tmp_a.dpos, tmp_a.lat))\natoms_before[\"atoms\"] = tmp_a.elm\nlabels_before = [elm for elm, point in zip(tmp_a.elm, tmp_a.dpos)]\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] * tmp_a.lat[0] + comb[1] * tmp_a.lat[1] + comb[2] * tmp_a.lat[2]\n        for comb in unit_cell_comb\n    ]\n)\nunit_cell_before = pv.PolyData(unit_cell)\n\ntmp_a = p.Poscar(os.path.join(data_dir, \"POSCAR-9AGNR-smaller.vasp\"))\ntmp_a.parse()\n\n# Convert positions to Cartesian coordinates for visualization\natoms_after = pv.PolyData(np.dot(tmp_a.dpos, tmp_a.lat))\natoms_after[\"atoms\"] = tmp_a.elm\nlabels_after = [elm for elm, point in zip(tmp_a.elm, tmp_a.dpos)]\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] * tmp_a.lat[0] + comb[1] * tmp_a.lat[1] + comb[2] * tmp_a.lat[2]\n        for comb in unit_cell_comb\n    ]\n)\nunit_cell_after = pv.PolyData(unit_cell)\n\ncreate_gif(\n    atoms=atoms_before,\n    labels=labels_before,\n    unit_cell=unit_cell_before,\n    save_file=\"atoms_before_scaling.gif\",\n)\ncreate_gif(\n    atoms=atoms_after,\n    labels=labels_after,\n    unit_cell=unit_cell_after,\n    save_file=\"atoms_after_scaling.gif\",\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Creating a Supercell\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print(\"\\n\\nNow I will make an supercell 3x1x1\")\nb = p.poscarUtils.poscar_supercell(\n    os.path.join(data_dir, \"POSCAR-9AGNR-smaller.vasp\"), verbose=False\n)\nsize = np.array([[3, 0, 0], [0, 1, 0], [0, 0, 1]])\nb.supercell(size=size)\nb.write(os.path.join(data_dir, \"POSCAR-9AGNR-311.vasp\"))\nprint(\"It was saved as POSCAR-9AGNR-311.vasp\")\n\n\ntmp_a = p.Poscar(os.path.join(data_dir, \"POSCAR-9AGNR-smaller.vasp\"))\ntmp_a.parse()\n\n# Convert positions to Cartesian coordinates for visualization\natoms_before = pv.PolyData(np.dot(tmp_a.dpos, tmp_a.lat))\natoms_before[\"atoms\"] = tmp_a.elm\nlabels_before = [elm for elm, point in zip(tmp_a.elm, tmp_a.dpos)]\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] * tmp_a.lat[0] + comb[1] * tmp_a.lat[1] + comb[2] * tmp_a.lat[2]\n        for comb in unit_cell_comb\n    ]\n)\nunit_cell_before = pv.PolyData(unit_cell)\n\ntmp_a = p.Poscar(os.path.join(data_dir, \"POSCAR-9AGNR-311.vasp\"))\ntmp_a.parse()\n\n# Convert positions to Cartesian coordinates for visualization\natoms_after = pv.PolyData(np.dot(tmp_a.dpos, tmp_a.lat))\natoms_after[\"atoms\"] = tmp_a.elm\nlabels_after = [elm for elm, point in zip(tmp_a.elm, tmp_a.dpos)]\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] * tmp_a.lat[0] + comb[1] * tmp_a.lat[1] + comb[2] * tmp_a.lat[2]\n        for comb in unit_cell_comb\n    ]\n)\nunit_cell_after = pv.PolyData(unit_cell)\n\ncreate_gif(\n    atoms=atoms_before,\n    labels=labels_before,\n    unit_cell=unit_cell_before,\n    save_file=\"atoms_before_supercell.gif\",\n)\ncreate_gif(\n    atoms=atoms_after,\n    labels=labels_after,\n    unit_cell=unit_cell_after,\n    save_file=\"atoms_after_supercell.gif\",\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Introducing Defects\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print(\n    \"\\n\\nFinally I want to create a defect by changing atoms #28, #29 to B and N, respectively\"\n)\nc = p.poscarUtils.poscar_modify(\n    os.path.join(data_dir, \"POSCAR-9AGNR-311.vasp\"), verbose=False\n)\nc.change_elements(\n    indexes=[28, 29], newElements=[\"B\", \"N\"], human=True\n)  # Mind, without `human`, first is 0, second is 1, ...\nc.write(os.path.join(data_dir, \"POSCAR-AGNR-defect.vasp\"))\nprint(\"It was saves as POSCAR-AGNR-defect.vasp\")\n\n\ntmp_a = p.Poscar(os.path.join(data_dir, \"POSCAR-9AGNR-311.vasp\"))\ntmp_a.parse()\n\n# Convert positions to Cartesian coordinates for visualization\natoms_before = pv.PolyData(np.dot(tmp_a.dpos, tmp_a.lat))\natoms_before[\"atoms\"] = tmp_a.elm\nlabels_before = [elm for elm, point in zip(tmp_a.elm, tmp_a.dpos)]\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] * tmp_a.lat[0] + comb[1] * tmp_a.lat[1] + comb[2] * tmp_a.lat[2]\n        for comb in unit_cell_comb\n    ]\n)\nunit_cell_before = pv.PolyData(unit_cell)\n\ntmp_a = p.Poscar(os.path.join(data_dir, \"POSCAR-AGNR-defect.vasp\"))\ntmp_a.parse()\n\n# Convert positions to Cartesian coordinates for visualization\natoms_after = pv.PolyData(np.dot(tmp_a.dpos, tmp_a.lat))\natoms_after[\"atoms\"] = tmp_a.elm\nlabels_after = [elm for elm, point in zip(tmp_a.elm, tmp_a.dpos)]\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] * tmp_a.lat[0] + comb[1] * tmp_a.lat[1] + comb[2] * tmp_a.lat[2]\n        for comb in unit_cell_comb\n    ]\n)\nunit_cell_after = pv.PolyData(unit_cell)\n\ncreate_gif(\n    atoms=atoms_before,\n    labels=labels_before,\n    unit_cell=unit_cell_before,\n    save_file=\"atoms_before_defect.gif\",\n)\ncreate_gif(\n    atoms=atoms_after,\n    labels=labels_after,\n    unit_cell=unit_cell_after,\n    save_file=\"atoms_after_defect.gif\",\n)\n\n\nprint(\"\")\n\nprint(\"Loading an AGNR with a defects in the last two entries\")\na = p.poscar.Poscar(os.path.join(data_dir, \"POSCAR-AGNR-defect.vasp\"), verbose=False)\na.parse()\n\nprint(\"The nearest neighbors of the defects are:\")\nnn = p.latticeUtils.Neighbors(a)\nprint(a.elm[-2], \":\", a.Ntotal - 2, \"-->\", nn.nn_list[-2])\nprint(a.elm[-1], \":\", a.Ntotal - 1, \"-->\", nn.nn_list[-1])"
      ]
    }
  ],
  "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
}