{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Examples for config.yml file\n\nConfiguration file `config.yml` use YAML (https://yaml.org) format to describe decay process.\n\nThe main parts of config.yml is `decay` and `particle`.\n\nThe `decay` part describe the particle (or an id of a list of particle)\ndecay into which particles, it can be a list of a list of list.\nA list means that there is ony one decay mode, A list of list is the list of possible decay mode.\nThe list item can be the particle name (or a dict to describe the decay parameters).\nAll name should appear in `particle` part.\n\nThe `particle` part describe the parameters of particles.\nThere are two special parts `$top` and `$finals` describe the top and finals particles.\nThe other parts are lists of particle name or dicts of particle parameters.\nThe list is the same type particle in decay part.\nThe dict is the parameters of the particle name.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "config_str = \"\"\"\n\ndecay:\n    A:\n       - [R1, B]\n       - [R2, C]\n       - [R3, D]\n    R1: [C, D]\n    R2: [B, D]\n    R3: [B, C]\n\nparticle:\n    $top:\n       A: { mass: 1.86, J: 0, P: -1}\n    $finals:\n       B: { mass: 0.494, J: 0, P: -1}\n       C: { mass: 0.139, J: 0, P: -1}\n       D: { mass: 0.139, J: 0, P: -1}\n    R1: [ R1_a, R1_b ]\n    R1_a: { mass: 0.7, width: 0.05, J: 1, P: -1}\n    R1_b: { mass: 0.5, width: 0.05, J: 0, P: +1}\n    R2: { mass: 0.824, width: 0.05, J: 0, P: +1}\n    R3: { mass: 0.824, width: 0.05, J: 0, P: +1}\n\n\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The config file can be loaded by `yaml` library.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\nimport yaml\n\nfrom tf_pwa.config_loader import ConfigLoader\nfrom tf_pwa.histogram import Hist1D\n\nconfig = ConfigLoader(yaml.full_load(config_str))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We set parameters to a blance value. And we can generate some toy data and calclute the weights\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "input_params = {\n    \"A->R1_a.BR1_a->C.D_total_0r\": 6.0,\n    \"A->R1_b.BR1_b->C.D_total_0r\": 1.0,\n    \"A->R2.CR2->B.D_total_0r\": 2.0,\n    \"A->R3.DR3->B.C_total_0r\": 1.0,\n}\nconfig.set_params(input_params)\n\ndata = config.generate_toy(1000)\nphsp = config.generate_phsp(10000)\n\n# You can also fit the data fit to the data\nfit_result = config.fit([data], [phsp])\nerr = config.get_params_error(fit_result, [data], [phsp])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "we can see that thre fit results consistant with inputs, the first one is fixed.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "for var in input_params:\n    print(\n        f\"in: {input_params[var]} => out: {fit_result.params[var]} +/- {err.get(var, 0.)}\"\n    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can use the amplitude to plot the fit results\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "amp = config.get_amplitude()\nweight = amp(phsp)\npartial_weight = amp.partial_weight(phsp)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can plot the data, Hist1D include some plot method base on matplotlib.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "data_hist = Hist1D.histogram(\n    data.get_mass(\"(C, D)\"), bins=60, range=(0.25, 1.45)\n)\n\nmass_phsp = phsp.get_mass(\"(C, D)\")\nphsp_hist = Hist1D.histogram(\n    mass_phsp, weights=weight, bins=60, range=(0.25, 1.45)\n)\nscale = phsp_hist.scale_to(data_hist)\n\npw_hist = []\nfor w in partial_weight:\n    # here we used more bins for a smooth plot\n    hist = Hist1D.histogram(\n        mass_phsp, weights=w, bins=60 * 2, range=(0.25, 1.45)\n    )\n    pw_hist.append(scale * hist * 2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Then we can plot the histogram into matplotlib\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "for hist, dec in zip(pw_hist, config.get_decay()):\n    hist.draw_kde(label=str(dec))\nphsp_hist.draw(label=\"total amplitude\")\ndata_hist.draw_error(label=\"toy data\", color=\"black\")\n\nplt.legend()\nplt.ylim((0, None))\nplt.xlabel(\"M(CD)/ GeV\")\nplt.ylabel(\"Events/ 20 MeV\")\nplt.show()"
      ]
    }
  ],
  "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
}