{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "%config Completer.use_jedi = False\n",
    "%config InlineBackend.figure_formats = ['svg']\n",
    "\n",
    "# Install on Google Colab\n",
    "import subprocess\n",
    "import sys\n",
    "\n",
    "from IPython import get_ipython\n",
    "\n",
    "install_packages = \"google.colab\" in str(get_ipython())\n",
    "if install_packages:\n",
    "    for package in [\"expertsystem\", \"graphviz\"]:\n",
    "        subprocess.check_call(\n",
    "            [sys.executable, \"-m\", \"pip\", \"install\", package]\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Particle database"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{warning}\n",
    "The {doc}`PWA Expert System <index>` has been split up into {doc}`QRules <qrules:index>` and {doc}`AmpForm <ampform:index>`. Please use these packages instead!\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{note}\n",
    "The formulas and figures on this page have been generated with the lineshapes in the {mod}`.lineshape` module, so as to [glue](https://myst-nb.readthedocs.io/en/latest/use/glue.html) them back in to the API of that module.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In PWA, you usually want to search for special resonances, possibly even some not listed in the PDG. In this notebook, we go through a few ways to add or overwrite {class}`.Particle` instances in the database with your own particle definitions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading the default database"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In {doc}`reaction`, we made use of the {class}`.StateTransitionManager`. By default, if you do not specify the `particles` argument, the {class}`.StateTransitionManager` calls the function {func}`.load_default_particles`. This functions returns a {class}`.ParticleCollection` instance with {class}`.Particle` definitions from the [PDG](https://pdg.lbl.gov), along with additional definitions that are provided in the file {download}`additional_definitions.yml <../../src/expertsystem/reaction/additional_definitions.yml>`.\n",
    "\n",
    "Here, we call this method directly to illustrate what happens (we use {func}`.load_pdg`, which loads a subset):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from expertsystem.reaction.particle import load_pdg\n",
    "\n",
    "particle_db = load_pdg()\n",
    "print(\"Number of loaded particles:\", len(particle_db))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following, we illustrate how to use the methods of the {class}`.ParticleCollection` class to find and 'modify' {class}`.Particle`s and {meth}`~.ParticleCollection.add` them back to the {class}`.ParticleCollection`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finding particles"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The {class}`.ParticleCollection` class offers some methods to search for particles by name or by PID (see {meth}`~.ParticleCollection.find`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "particle_db.find(333)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With {meth}`~.ParticleCollection.filter`, you can perform more sophisticated searches. This is done by either passing a function or [lambda](https://docs.python.org/3/tutorial/controlflow.html#lambda-expressions)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subset = particle_db.filter(lambda p: p.name.startswith(\"f(2)\"))\n",
    "subset.names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subset = particle_db.filter(\n",
    "    lambda p: p.strangeness == 1\n",
    "    and p.spin >= 1\n",
    "    and p.mass > 1.8\n",
    "    and p.mass < 1.9\n",
    ")\n",
    "subset.names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subset = particle_db.filter(lambda p: p.is_lepton())\n",
    "subset.names"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that in each of these examples, we call the {attr}`~.ParticleCollection.names` property. This is just to only display the names, sorted alphabetically, otherwise the output becomes a bit of a mess:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "particle_db.filter(lambda p: p.name.startswith(\"pi\") and len(p.name) == 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LaTeX representation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "{class}`.Particle`s also contain a {attr}`~.Particle.latex` tag. Here, we use [ipython](https://ipython.readthedocs.io/en/stable/api/generated/IPython.display.html#IPython.display.Math) to render them nicely as mathematical symbols:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Math\n",
    "\n",
    "sigmas = particle_db.filter(\n",
    "    lambda p: p.name.startswith(\"Sigma\") and p.charmness == 1\n",
    ")\n",
    "Math(\", \".join([p.latex for p in sigmas]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding custom particle definitions through Python"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A quick way to modify or overwrite particles, is through your Python script or notebook. Notice that the instances in the database are {class}`.Particle` instances:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N1650_plus = particle_db[\"N(1650)+\"]\n",
    "N1650_plus"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The instances in the database are [immutable](https://en.wikipedia.org/wiki/Immutable_object). Therefore, if you want to modify, say, the width, you have to create a new {class}`.Particle` instance from the particle you want to modify and {meth}`~.ParticleCollection.add` it back to the database. You can do this with {func}`.create_particle`:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{margin} Duplicate names or PIDs\n",
    "The warning that you see here comes from the fact that names and PIDs are considered mere labels of a {obj}`.Particle` instance ― it is defined uniquely only by its quantum numbers, such as {attr}`~.Particle.spin` and {attr}`~.Particle.charge`. \n",
    "    \n",
    "The warning is suppressed in the rest of this page.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-warn"
    ]
   },
   "outputs": [],
   "source": [
    "from expertsystem.reaction.particle import create_particle\n",
    "\n",
    "new_N1650_plus = create_particle(\n",
    "    template_particle=N1650_plus, name=\"Modified N(1650)+\", width=0.2\n",
    ")\n",
    "\n",
    "particle_db.add(new_N1650_plus)\n",
    "particle_db[\"Modified N(1650)+\"].width"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You often also want to add the antiparticle of the particle you modified to the database. Using {func}`.create_antiparticle`, it is easy to create the corresponding antiparticle object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from expertsystem.reaction.particle import create_antiparticle\n",
    "\n",
    "new_N1650_minus = create_antiparticle(\n",
    "    new_N1650_plus, new_name=\"Modified N(1650)-\"\n",
    ")\n",
    "\n",
    "particle_db.add(new_N1650_minus)\n",
    "particle_db[\"Modified N(1650)-\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When adding additional particles you may need for your research, it is easiest to work with an existing particle as template. Let's say we want to study $e^+e^-$ collisions of several energies:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````{margin}\n",
    "```{note}\n",
    "By convention, the {mod}`expertsystem` uses $\\mathrm{GeV}/c^2$ as energy unit.\n",
    "```\n",
    "````"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "energies_mev = {4180, 4220, 4420, 4600}\n",
    "template_particle = particle_db[\"J/psi(1S)\"]\n",
    "for energy_mev in energies_mev:\n",
    "    energy_gev = energy_mev / 1e3\n",
    "    new_particle = create_particle(\n",
    "        template_particle, name=f\"EpEm ({energy_mev} MeV)\", mass=energy_gev\n",
    "    )\n",
    "    particle_db.add(new_particle)\n",
    "len(particle_db)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "particle_db.filter(lambda p: \"EpEm\" in p.name).names"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Of course, it's also possible to add any kind of custom {class}`.Particle`, as long as its quantum numbers comply with the {func}`.gellmann_nishijima` rule:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from expertsystem.reaction.particle import Particle\n",
    "\n",
    "custom = Particle(\n",
    "    name=\"custom\",\n",
    "    pid=99999,\n",
    "    latex=R\"p_\\mathrm{custom}\",\n",
    "    spin=1.0,\n",
    "    mass=1,\n",
    "    charge=1,\n",
    "    isospin=(1.5, 0.5),\n",
    "    charmness=1,\n",
    ")\n",
    "custom"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "particle_db += custom\n",
    "len(particle_db)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading custom definitions from a YAML file"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's also possible to add particles from a config file, with {func}`.io.load`. Existing entries remain and if the imported file of particle definitions contains a particle with the same name, it is overwritten in the database.\n",
    "\n",
    "It's easiest to work with YAML. Here, we use the provided {download}`additional_particles.yml` example file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from expertsystem import io\n",
    "\n",
    "particle_db += io.load(\"additional_particles.yml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Writing to YAML"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also dump the existing particle lists to YAML. You do this with the {func}`.io.write` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "io.write(instance=particle_db, filename=\"dumped_particle_list.yaml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the function {func}`write <.io.write>` can dump any {class}`.ParticleCollection` to an output file, also a specific subset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from expertsystem.reaction.particle import ParticleCollection\n",
    "\n",
    "output = ParticleCollection()\n",
    "output += particle_db[\"J/psi(1S)\"]\n",
    "output += particle_db.find(22)  # gamma\n",
    "output += particle_db.filter(lambda p: p.name.startswith(\"f(0)\"))\n",
    "output += particle_db[\"pi0\"]\n",
    "output += particle_db[\"pi+\"]\n",
    "output += particle_db[\"pi-\"]\n",
    "output += particle_db[\"custom\"]\n",
    "io.write(output, \"particle_list_selection.yml\")\n",
    "output.names"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a side note, the {mod}`expertsystem` provides [JSON schemas](https://json-schema.org) ({download}`reaction/particle-validation.json <../../src/expertsystem/reaction/particle-validation.json>`) to validate your particle list files (see also {func}`jsonschema.validate`). If you have installed the {mod}`expertsystem` as an {ref}`pwa:develop:Editable installation` and {ref}`use VSCode <develop:Visual Studio code>`, your YAML particle list are checked automatically in the GUI."
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Raw Cell Format",
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}