{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2bce8c80",
   "metadata": {},
   "source": [
    "# Selecting atoms with categories\n",
    "\n",
    "Many sisl methods take a set of atoms: `Geometry.sub`, `Geometry.remove`, `Geometry.translate`,\n",
    "plotting styles and more. Atomic indices work, but they change as soon as the geometry changes.\n",
    "Atom *categories* select atoms by what they are, where they are, and what surrounds them.\n",
    "The same category works on a small test system and on a large production geometry.\n",
    "\n",
    "This tutorial introduces categories one step at a time:\n",
    "\n",
    "1. one category and its result,\n",
    "2. categories where sisl expects atoms,\n",
    "3. combining categories,\n",
    "4. the atomic environment (neighbors),\n",
    "5. the position of atoms,\n",
    "6. the short keyword form,\n",
    "7. changing a geometry with categories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbe98dc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import sisl as si\n",
    "from sisl.geom import AtomCategory, AtomNeighbors, AtomXYZ, AtomZ"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0eb8384",
   "metadata": {},
   "source": [
    "Our test system is a zigzag graphene nanoribbon. It is periodic along the ribbon and has two edges.\n",
    "We replace two carbon atoms with nitrogen, one at an edge and one inside the ribbon,\n",
    "and we remove one atom to make a vacancy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddb69721",
   "metadata": {},
   "outputs": [],
   "source": [
    "ribbon = si.geom.zgnr(6).tile(4, 0)\n",
    "ribbon.atoms[[17, 20]] = si.Atom(\"N\", R=1.4342)\n",
    "ribbon = ribbon.remove(38)\n",
    "print(ribbon)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f202c99d",
   "metadata": {},
   "outputs": [],
   "source": [
    "ribbon.plot(axes=\"xy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96db08fe",
   "metadata": {},
   "source": [
    "## 1. One category\n",
    "\n",
    "`AtomZ` selects atoms by their atomic number. It accepts the number or the symbol.\n",
    "The method `mask` returns one boolean value per atom: `True` for the atoms in the category."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50dff862",
   "metadata": {},
   "outputs": [],
   "source": [
    "nitrogen = AtomZ(\"N\")\n",
    "mask = nitrogen.mask(ribbon)\n",
    "print(mask)\n",
    "print(\"indices:\", np.flatnonzero(mask))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ca269aa",
   "metadata": {},
   "source": [
    "A category does not store the geometry. You can use the same `nitrogen` category on any geometry."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d182fd8",
   "metadata": {},
   "source": [
    "## 2. Use a category where sisl expects atoms\n",
    "\n",
    "Methods that take atomic indices also take a category.\n",
    "Here we extract the nitrogen atoms, and we highlight them in a plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7c7bacb",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(ribbon.sub(nitrogen))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c72bc3be",
   "metadata": {},
   "outputs": [],
   "source": [
    "ribbon.plot(axes=\"xy\", atoms_style=[{\"atoms\": nitrogen, \"color\": \"blue\"}])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cea2374",
   "metadata": {},
   "source": [
    "## 3. Combine categories\n",
    "\n",
    "Categories combine with the operators of Python sets:\n",
    "\n",
    "| Operator | Meaning |\n",
    "|---|---|\n",
    "| `A & B` | in `A` and in `B` |\n",
    "| `A \\| B` | in `A` or in `B` |\n",
    "| `A ^ B` | in exactly one of `A` and `B` |\n",
    "| `~A` | not in `A` |\n",
    "\n",
    "For example, all atoms that are not carbon, or the carbon atoms only:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1278160c",
   "metadata": {},
   "outputs": [],
   "source": [
    "carbon = AtomZ(\"C\")\n",
    "print(\"not carbon:\", np.flatnonzero((~carbon).mask(ribbon)))\n",
    "print(\"carbon count:\", carbon.mask(ribbon).sum(), \"of\", ribbon.na)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c14a71e",
   "metadata": {},
   "source": [
    "## 4. The atomic environment\n",
    "\n",
    "`AtomNeighbors` counts the neighbors of each atom within a radius `R` (in Ångström).\n",
    "Atoms in the bulk of graphene have 3 neighbors.\n",
    "The atoms at the edges, and the atoms around the vacancy, have only 2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83425204",
   "metadata": {},
   "outputs": [],
   "source": [
    "two_bonds = AtomNeighbors(max=2, R=1.5)\n",
    "print(np.flatnonzero(two_bonds.mask(ribbon)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3465aa29",
   "metadata": {},
   "outputs": [],
   "source": [
    "ribbon.plot(axes=\"xy\", atoms_style=[{\"atoms\": two_bonds, \"color\": \"red\"}])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9f05ebd",
   "metadata": {},
   "source": [
    "`AtomNeighbors` takes one or two numbers: `AtomNeighbors(3)` means exactly 3 neighbors,\n",
    "`AtomNeighbors(1, 2)` means 1 or 2 neighbors. The keywords `min` and `max` give only one bound.\n",
    "\n",
    "The argument `neighbor` counts only the neighbors in another category.\n",
    "These are the atoms with at least one nitrogen neighbor:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d542673",
   "metadata": {},
   "outputs": [],
   "source": [
    "next_to_N = AtomNeighbors(min=1, neighbor=nitrogen, R=1.5)\n",
    "print(np.flatnonzero(next_to_N.mask(ribbon)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbac5b1a",
   "metadata": {},
   "source": [
    "## 5. The position of atoms\n",
    "\n",
    "`AtomXYZ` selects atoms by their coordinates. A tuple `(min, max)` gives an interval,\n",
    "and `None` means no bound. The atoms near the vacancy are not at the edges,\n",
    "so a position interval separates them from the edge atoms:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e1acbf9",
   "metadata": {},
   "outputs": [],
   "source": [
    "y = ribbon.xyz[:, 1]\n",
    "inside = AtomXYZ(y=(y.min() + 1, y.max() - 1))\n",
    "vacancy = two_bonds & inside\n",
    "print(np.flatnonzero(vacancy.mask(ribbon)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14baf440",
   "metadata": {},
   "source": [
    "`AtomXYZ` also has a short form with `<` and `>`, also for fractional coordinates (`fx`, `fy`, `fz`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3439736",
   "metadata": {},
   "outputs": [],
   "source": [
    "lower_half = AtomXYZ.fy < 0.5\n",
    "print(lower_half.mask(ribbon).sum(), \"atoms in the lower half of the cell\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17e93e9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "ribbon.plot(\n",
    "    axes=\"xy\",\n",
    "    atoms_style=[\n",
    "        {\"atoms\": two_bonds & ~inside, \"color\": \"red\"},\n",
    "        {\"atoms\": vacancy, \"color\": \"green\"},\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "657fa8f5",
   "metadata": {},
   "source": [
    "## 6. The keyword form\n",
    "\n",
    "`AtomCategory` builds categories from keywords. Each keyword is the name of a category\n",
    "without `Atom`, and many keywords combine with `&`.\n",
    "A `dict` with the same keywords works directly as an index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0c75f7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "edge_carbon = AtomCategory(Z=\"C\", neighbors={\"max\": 2, \"R\": 1.5})\n",
    "print(edge_carbon)\n",
    "print(np.flatnonzero(edge_carbon.mask(ribbon)))\n",
    "print(ribbon.sub({\"Z\": \"N\"}).na, \"nitrogen atoms\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00ec6df7",
   "metadata": {},
   "source": [
    "An unknown keyword shows the valid keywords:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d2a65bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    AtomCategory(neigbors=2)\n",
    "except ValueError as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf60d6a4",
   "metadata": {},
   "source": [
    "## 7. Change a geometry with categories\n",
    "\n",
    "A category describes *which* atoms to change. The method describes *how*.\n",
    "Here we move the atoms next to nitrogen out of the plane:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b495d043",
   "metadata": {},
   "outputs": [],
   "source": [
    "buckled = ribbon.translate([0, 0, 0.3], atoms=next_to_N)\n",
    "print(np.unique(buckled.xyz[:, 2] - ribbon.xyz[:, 2]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f2ebc4e",
   "metadata": {},
   "source": [
    "`sisl.geom.saturate` adds atoms to the missing bonds.\n",
    "Each rule is a pair of a category and a `Saturate` object, and the first rule that\n",
    "selects an atom wins. Here the edge carbon atoms get hydrogen, and the edge nitrogen atom\n",
    "gets nothing because it already has a lone pair:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c22b9a63",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sisl.geom import Saturate, saturate\n",
    "\n",
    "edge = two_bonds & ~inside\n",
    "passivated = saturate(\n",
    "    ribbon,\n",
    "    [\n",
    "        (edge & nitrogen, Saturate(\"H\", coordination=2, R=1.5)),\n",
    "        (edge, Saturate(\"H\", bond=1.09, R=1.5)),\n",
    "    ],\n",
    ")\n",
    "print(passivated.na - ribbon.na, \"hydrogen atoms added\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "852a9545",
   "metadata": {},
   "outputs": [],
   "source": [
    "passivated.plot(axes=\"xy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cc63450",
   "metadata": {},
   "source": [
    "`Saturate` finds the number of missing bonds from the most frequent number of bonds of the\n",
    "same element, here 3 for carbon. The argument `coordination` sets it directly:\n",
    "`coordination=2` for the nitrogen means that it is already saturated.\n",
    "Without `bond`, the bond length is the sum of the atomic radii, which is approximate.\n",
    "\n",
    "## Summary\n",
    "\n",
    "- A category selects atoms by element, position or environment, and it does not depend on the indices.\n",
    "- `mask` returns one boolean per atom. Methods that take atoms also take a category or a `dict`.\n",
    "- Combine categories with `&`, `|`, `^` and `~`.\n",
    "- `AtomCategory(...)` and plain `dict`s build categories from keywords.\n",
    "\n",
    "The reference of all categories is in the API documentation: [Indexing atoms](../api/geom/indexing.rst)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}