{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Examples for particle model\n\ndecay system is model as\n\nDecayGroup\n    DecayChain\n\n        Decay\n\n        Particle\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n\nfrom tf_pwa.amp import Decay, DecayChain, DecayGroup, Particle\nfrom tf_pwa.vis import plot_decay_struct"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can easy create some instance of Particle\nand then combine them as Decay\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "a = Particle(\"A\")\nb = Particle(\"B\")\nc = Particle(\"C\")\nd = Particle(\"D\")\n\nr = Particle(\"R\")\n\ndec1 = Decay(a, [r, b])\ndec2 = Decay(r, [c, d])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "DecayChain is a list of Decays.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "decay_chain = DecayChain([dec1, dec2])\ndecay_chain"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can plot it using matplotlib.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plot_decay_struct(decay_chain)\nplt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "DecayGroup is a list of DecayChain with the same initial and final states\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "decay_group = DecayGroup([decay_chain])\ndecay_group"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can build a simple function to infer the charge from final states.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def charge_infer(dec, charge_map):\n    # use particle equal condition\n    cached_charge = {Particle(k): v for k, v in charge_map.items()}\n    # loop for all decays in decay chain\n    for i, dec_i in dec.depth_first(False):\n        # all out particles has charge\n        assert all(i in cached_charge for i in dec_i.outs)\n        # the charge or core particle is the sum of\n        cached_charge[dec_i.core] = sum(cached_charge[i] for i in dec_i.outs)\n    return cached_charge\n\n\ncharges = {\n    \"B\": -1,\n    \"C\": 0,\n    \"D\": 1,\n}\n\ncharge_infer(decay_chain, charges)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "See more in `~tf_pwa.cal_angle.cal_chain_boost`.\n\n"
      ]
    }
  ],
  "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.7.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}