{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "Download this notebook :download:`here <particles.ipynb>`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Particle database"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "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 `.Particle` instances in the database with your own particle definitions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading the default database"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "In the usual :doc:`workflow </usage/quickstart>`, you start by calling the `.StateTransitionManager`. Upon construction, the the `.StateTransitionManager` calls the function `.load_default_particles`, which fills the :code:`expertsystem.state.particle.DATABASE` (simply a `.ParticleCollection` instance) with `.Particle` definitions from the `PDG <https://pdg.lbl.gov/>`_. Here, we call this method directly to illustrate what happens:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from expertsystem.ui import load_default_particles\n",
    "\n",
    "particle_db = load_default_particles()\n",
    "print(\"Number of loaded particles:\", len(particle_db))"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "In the following, we illustrate how to use the methods of the `.ParticleCollection` class and the `~expertsystem.state.particle` module to modify or add particle definitions in the particle database."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finding particles"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "The `.ParticleCollection` class offers some methods to search for particles, either by name or by PID (see `.ParticleCollection.find`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "particle_db.find(333)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "`.ParticleCollection.find_subset` returns a `.ParticleCollection` subset instead of one `.Particle` instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subset = particle_db.find_subset(\"f(2)\")\n",
    "list(subset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "More complicated searches can be done with list comprehensions. This is especially useful when searching for specific quantum numbers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "{\n",
    "    item.name\n",
    "    for item in particle_db.values()\n",
    "    if item.state.strangeness == 1\n",
    "    and item.state.spin >= 1\n",
    "    and item.mass > 1.8\n",
    "    and item.mass < 1.9\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding custom particle definitions through Python"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "A quick way to modify or overwrite particles, is through your Python script of notebook. Notice that the instances in the database are  `.Particle` instances:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N1650_plus = particle_db[\"N(1650)+\"]\n",
    "N1650_plus"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "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 `.Particle` instance from the particle you want to modify and `~.ParticleCollection.add` it back to the database. You can do this with `.create_particle`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from expertsystem.data 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": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "You often also want to add the antiparticle of the particle you modified to the database. Using `.create_antiparticle`, it is easy to create the corresponding antiparticle object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from expertsystem.data import create_antiparticle\n",
    "\n",
    "new_N1650_minus = create_antiparticle(new_N1650_plus, new_name=\"Modified N(1650)-\")\n",
    "\n",
    "particle_db.add(new_N1650_minus)\n",
    "particle_db[\"Modified N(1650)-\"]"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "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 :math:`e^+e^-` collisions of several energies:"
   ]
  },
  {
   "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": [
    "set(particle_db.find_subset(\"EpEm\"))"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    ".. note::\n",
    "  By convention, the `expertsystem` uses GeV as energy unit."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading custom definitions from a YAML file"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "It's also possible to add particles from a config file, with `~.io.load_particle_collection`. 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.io import load_particle_collection\n",
    "\n",
    "particle_db += load_particle_collection(\"additional_particles.yml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Writing to XML or YAML"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "You can also dump the existing particle lists to either YAML or XML. You do this with the `expertsystem.io.write` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from expertsystem import io\n",
    "\n",
    "io.write(instance=particle_db, filename=\"dumped_particle_list.xml\")\n",
    "io.write(instance=particle_db, filename=\"dumped_particle_list.yaml\")"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "Note that the function `~expertsystem.io.write` can dump any `.ParticleCollection` to an output file, also a specific subset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from expertsystem.data import ParticleCollection\n",
    "\n",
    "output = ParticleCollection()\n",
    "output += particle_db[\"J/psi(1S)\"]\n",
    "output += particle_db.find(22)  # gamma\n",
    "output += particle_db.find_subset(\"f(0)\")\n",
    "output += particle_db[\"pi0\"]\n",
    "output += particle_db[\"pi+\"]\n",
    "output += particle_db[\"pi-\"]\n",
    "io.write(output, \"particle_list_selection.yml\")\n",
    "set(output)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "As a side note, the `expertsystem` provides `JSON schemas <https://json-schema.org/>`_ (e.g. :download:`yaml/particle-list.json <../../expertsystem/schemas/yaml/particle-list.json>`) to validate your particle list files (see also `jsonschema.validate`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml\n",
    "\n",
    "from expertsystem import io\n",
    "\n",
    "with open(\"dumped_particle_list.yaml\") as stream:\n",
    "    definition = yaml.load(stream, Loader=yaml.SafeLoader)\n",
    "io.yaml.validation.particle_list(definition)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "In addition, if you have installed the `expertsystem` in :ref:`install:Development mode` and :ref:`use VSCode <contribute: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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}