Visualize pymatgen's structure with NGLView

nglview_pymatgen.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "778b450bb7f64ce58ece0d839432d641",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import nglview\n",
    "from pymatgen.core import Structure, Lattice\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Refs.\n",
    "- https://github.com/pyiron/pyiron/blob/9ba4f1efa57b286d69ee2cdb6b865782eda755f5/pyiron/atomistics/structure/_visualize.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot3d(structure, spacefill=True, show_axes=True):\n",
    "    from itertools import product\n",
    "    from pymatgen.core import Structure\n",
    "    from pymatgen.core.sites import PeriodicSite\n",
    "    \n",
    "    eps = 1e-8\n",
    "    sites = []\n",
    "    for site in structure:\n",
    "        species = site.species\n",
    "        frac_coords = np.remainder(site.frac_coords, 1)\n",
    "        for jimage in product([0, 1 - eps], repeat=3):\n",
    "            new_frac_coords = frac_coords + np.array(jimage)\n",
    "            if np.all(new_frac_coords < 1 + eps):\n",
    "                new_site = PeriodicSite(species=species, coords=new_frac_coords, lattice=structure.lattice)\n",
    "                sites.append(new_site)\n",
    "    structure_display = Structure.from_sites(sites)\n",
    "    \n",
    "    view = nglview.show_pymatgen(structure_display)\n",
    "    view.add_unitcell()\n",
    "    \n",
    "    if spacefill:\n",
    "        view.add_spacefill(radius_type='vdw', radius=0.5, color_scheme='element')\n",
    "        view.remove_ball_and_stick()\n",
    "    else:\n",
    "        view.add_ball_and_stick()\n",
    "        \n",
    "    if show_axes:\n",
    "        view.shape.add_arrow([-4, -4, -4], [0, -4, -4], [1, 0, 0], 0.5, \"x-axis\")\n",
    "        view.shape.add_arrow([-4, -4, -4], [-4, 0, -4], [0, 1, 0], 0.5, \"y-axis\")\n",
    "        view.shape.add_arrow([-4, -4, -4], [-4, -4, 0], [0, 0, 1], 0.5, \"z-axis\")\n",
    "        \n",
    "    view.camera = \"perspective\"\n",
    "    return view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# rutile structure taken from mp-2657\n",
    "a = 4.653\n",
    "c = 2.969\n",
    "x_4f = 0.3046\n",
    "\n",
    "lattice = Lattice.from_parameters(a, a, c, 90, 90, 90)\n",
    "species = [\"Ti\", \"Ti\", \"O\", \"O\", \"O\", \"O\"]\n",
    "frac_coords = np.array([\n",
    "    [0, 0, 0],                      # Ti(2a)\n",
    "    [0.5, 0.5, 0.5],                # Ti(2a)\n",
    "    [x_4f, x_4f, 0],                # O(4f)\n",
    "    [1 - x_4f, 1 - x_4f, 0],        # O(4f)\n",
    "    [0.5 - x_4f, 0.5 + x_4f, 0.5],  # O(4f)\n",
    "    [0.5 + x_4f, 0.5 - x_4f, 0.5],  # O(4f)\n",
    "])\n",
    "structure = Structure(lattice, species, frac_coords)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Structure Summary\n",
       "Lattice\n",
       "    abc : 4.653 4.653 2.969\n",
       " angles : 90.0 90.0 90.0\n",
       " volume : 64.28006432099998\n",
       "      A : 4.653 0.0 2.849140778216317e-16\n",
       "      B : 7.48259468894539e-16 4.653 2.849140778216317e-16\n",
       "      C : 0.0 0.0 2.969\n",
       "PeriodicSite: Ti (0.0000, 0.0000, 0.0000) [0.0000, 0.0000, 0.0000]\n",
       "PeriodicSite: Ti (2.3265, 2.3265, 1.4845) [0.5000, 0.5000, 0.5000]\n",
       "PeriodicSite: O (1.4173, 1.4173, 0.0000) [0.3046, 0.3046, 0.0000]\n",
       "PeriodicSite: O (3.2357, 3.2357, 0.0000) [0.6954, 0.6954, 0.0000]\n",
       "PeriodicSite: O (0.9092, 3.7438, 1.4845) [0.1954, 0.8046, 0.5000]\n",
       "PeriodicSite: O (3.7438, 0.9092, 1.4845) [0.8046, 0.1954, 0.5000]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4f6db327678e4a2eb349b6ed2b73d046",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "NGLWidget()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot3d(structure, spacefill=True)"
   ]
  }
 ],
 "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.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}