{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Particle and amplitude\n\n    Amplitude = DecayGroup + Variable\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We will use following parameters for a toy model\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from tf_pwa.amp import DecayChain, DecayGroup, get_decay, get_particle\n\nresonances = {\n    \"R0\": {\"J\": 0, \"P\": 1, \"mass\": 1.0, \"width\": 0.07},\n    \"R1\": {\"J\": 0, \"P\": 1, \"mass\": 1.0, \"width\": 0.07},\n    \"R2\": {\"J\": 1, \"P\": -1, \"mass\": 1.225, \"width\": 0.08},\n}\n\na, b, c, d = [get_particle(i, J=0, P=-1) for i in \"ABCD\"]\nr1, r2, r3 = [get_particle(i, **resonances[i]) for i in resonances.keys()]\n\n\ndecay_group = DecayGroup(\n    [\n        DecayChain([get_decay(a, [r1, c]), get_decay(r1, [b, d])]),\n        DecayChain([get_decay(a, [r2, b]), get_decay(r2, [c, d])]),\n        DecayChain([get_decay(a, [r3, b]), get_decay(r3, [c, d])]),\n    ]\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The above parts can be represented as config.yml used by ConfigLoader.\n\nWe can get AmplitudeModel form decay_group and a optional Variables Managerager.\nIt has parameters, so we can get and set parameters for the amplitude model\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from tf_pwa.amp import AmplitudeModel\nfrom tf_pwa.variable import VarsManager\n\nvm = VarsManager()\namp = AmplitudeModel(decay_group, vm=vm)\n\nprint(amp.get_params())\namp.set_params(\n    {\n        \"A->R0.CR0->B.D_total_0r\": 1.0,\n        \"A->R1.BR1->C.D_total_0r\": 1.0,\n        \"A->R2.BR2->C.D_total_0r\": 7.0,\n    }\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "For the calculation, we generate some phase space data.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from tf_pwa.phasespace import PhaseSpaceGenerator\n\nm_A, m_B, m_C, m_D = 1.8, 0.18, 0.18, 0.18\np1, p2, p3 = PhaseSpaceGenerator(m_A, [m_B, m_C, m_D]).generate(100000)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "and the calculate helicity angle from the data\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from tf_pwa.cal_angle import cal_angle_from_momentum\n\ndata = cal_angle_from_momentum({b: p1, c: p2, d: p3}, decay_group)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "we can index mass from data as\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from tf_pwa.data import data_index\n\nm_bd = data_index(data, (\"particle\", \"(B, D)\", \"m\"))\n# m_bc = data_index(data, (\"particle\", \"(B, C)\", \"m\"))\nm_cd = data_index(data, (\"particle\", \"(C, D)\", \"m\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class=\"alert alert-info\"><h4>Note</h4><p>If DecayGroup do not include resonant of (B, C), the data will not include its mass too.\n    We can use different DecayGroup for cal_angle and AmplitudeModel\n    when they have the same initial and final particle.</p></div>\n\nThe amplitde square is calculated by amplitude model simply as\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "amp_s2 = amp(data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now by using matplotlib we can get the Dalitz plot as\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n\nplt.clf()\nplt.hist2d(\n    m_bd.numpy() ** 2,\n    m_cd.numpy() ** 2,\n    weights=amp_s2.numpy(),\n    bins=60,\n    cmin=1,\n    cmap=\"jet\",\n)\nplt.colorbar()\nplt.xlabel(\"$m^2(BD)$\")\nplt.ylabel(\"$m^2(CD)$\")\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
}