{
 "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. The particle database that the `expertsystem` loads by default is therefore often not sufficient for your analysis. In this notebook, we go through a few ways to add or overwrite 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 function `.load_default_particle_list` is called, which fills the :code:`expertsystem.state.particle.DATABASE` instance with particle definitions as defined in `particle_list.yml <https://github.com/ComPWA/expertsystem/blob/master/expertsystem/particle_list.yml>`_. Here, we call this method directly to illustrate what happens:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from expertsystem.state import particle\n",
    "from expertsystem.ui import load_default_particle_list\n",
    "\n",
    "print(\"Before loading:\", len(particle.DATABASE))\n",
    "load_default_particle_list()\n",
    "print(\"After loading:\", len(particle.DATABASE))"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "The functions in the :mod:`.particle` module can be used to add or modify particle definitions in the particle database. There are a few ways to add or overwrite particle definitions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finding particles"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "There are also functions to search for particles, either by name or by PID (see `.ParticleCollection.find`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "particle.DATABASE.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.DATABASE.find_subset(\"Delta\")\n",
    "list(subset)"
   ]
  },
  {
   "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": [
    "omega = particle.DATABASE[\"omega(782)\"]\n",
    "omega"
   ]
  },
  {
   "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, width, you have to create a new `.Particle` instance from the particle you want to modify and `~.ParticleCollection.add` it back to the database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from expertsystem.data import Particle\n",
    "\n",
    "new_omega = Particle(\n",
    "    name=omega.name,\n",
    "    pid=omega.pid,\n",
    "    mass=omega.mass,\n",
    "    width=0.01,\n",
    "    state=omega.state,\n",
    ")\n",
    "particle.DATABASE.add(new_omega)\n",
    "particle.DATABASE[\"omega(782)\"].width"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "When adding additional particles you may need for your research, it is also useful 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.DATABASE[\"EpEm (4230 MeV)\"]\n",
    "for counter, energy_mev in enumerate(energies_mev, 1):\n",
    "    energy_gev = energy_mev / 1e3\n",
    "    new_particle = Particle(\n",
    "        name=f\"EpEm ({energy_mev} MeV)\",\n",
    "        pid=template_particle.pid,\n",
    "        mass=energy_gev,\n",
    "        width=template_particle.width,\n",
    "        state=template_particle.state,\n",
    "    )\n",
    "    particle.DATABASE.add(new_particle)\n",
    "len(particle.DATABASE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "list(particle.DATABASE.find_subset(\"EpEm\"))"
   ]
  },
  {
   "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 `.load_particles`. Existing entries remain and if the loaded particle list file contains a particle with the same name, it is overwritten in the database. 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": [
    "print(\"Sigmas before loading:\", list(particle.DATABASE.find_subset(\"Sigma\")))\n",
    "particle.load_particles(\"additional_particles.yml\")\n",
    "print(\"Sigmas after loading: \", list(particle.DATABASE.find_subset(\"Sigma\")))"
   ]
  },
  {
   "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.DATABASE, filename=\"dumped_particle_list.xml\")\n",
    "io.write(instance=particle.DATABASE, 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.DATABASE[\"J/psi\"]\n",
    "output += particle.DATABASE.find(22)  # gamma\n",
    "output += particle.DATABASE.find_subset(\"f\")\n",
    "output += particle.DATABASE.find_subset(\"pi\")\n",
    "io.write(output, \"particle_list_selection.yml\")\n",
    "list(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",
    "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": 2
}