{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Numpy backup for scaling in the beginning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\users\\sa_li\\.conda\\envs\\rmd\\lib\\site-packages\\zfit\\util\\execution.py:53: UserWarning: Not running on Linux. Determining available cpus for thread can failand be overestimated. Workaround (only if too many cpus are used):`zfit.run.set_n_cpu(your_cpu_number)`\n",
      "  warnings.warn(\"Not running on Linux. Determining available cpus for thread can fail\"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "WARNING: The TensorFlow contrib module will not be included in TensorFlow 2.0.\n",
      "For more information, please see:\n",
      "  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n",
      "  * https://github.com/tensorflow/addons\n",
      "If you depend on functionality not listed there, please file an issue.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from pdg_const import pdg\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import pickle as pkl\n",
    "import sys\n",
    "import time\n",
    "from helperfunctions import display_time, prepare_plot\n",
    "import cmath as c\n",
    "import scipy.integrate as integrate\n",
    "from scipy.optimize import fminbound\n",
    "from array import array as arr\n",
    "import collections\n",
    "from itertools import compress\n",
    "import tensorflow as tf\n",
    "import zfit\n",
    "from zfit import ztf\n",
    "from IPython.display import clear_output\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def formfactor( q2, subscript): #returns real value\n",
    "    #check if subscript is viable\n",
    "\n",
    "    if subscript != \"0\" and subscript != \"+\" and subscript != \"T\":\n",
    "        raise ValueError('Wrong subscript entered, choose either 0, + or T')\n",
    "\n",
    "    #get constants\n",
    "\n",
    "    mK = pdg['Ks_M']\n",
    "    mbstar0 = pdg[\"mbstar0\"]\n",
    "    mbstar = pdg[\"mbstar\"]\n",
    "    b0 = pdg[\"b0\"]\n",
    "    bplus = pdg[\"bplus\"]\n",
    "    bT = pdg[\"bT\"]\n",
    "\n",
    "    mmu = pdg['muon_M']\n",
    "    mb = pdg['bquark_M']\n",
    "    ms = pdg['squark_M']\n",
    "    mB = pdg['Bplus_M']\n",
    "\n",
    "    #N comes from derivation in paper\n",
    "\n",
    "    N = 3\n",
    "\n",
    "    #some helperfunctions\n",
    "\n",
    "    tpos = (mB - mK)**2\n",
    "    tzero = (mB + mK)*(np.sqrt(mB)-np.sqrt(mK))**2\n",
    "\n",
    "    z_oben = np.sqrt(tpos - q2) - np.sqrt(tpos - tzero)\n",
    "    z_unten = np.sqrt(tpos - q2) + np.sqrt(tpos - tzero)\n",
    "    z = z_oben/ z_unten\n",
    "\n",
    "    #calculate f0\n",
    "\n",
    "    if subscript == \"0\":\n",
    "        prefactor = 1/(1 - q2/(mbstar0**2))\n",
    "        _sum = 0\n",
    "\n",
    "        for i in range(N):\n",
    "            _sum += b0[i]*(np.power(z,i))\n",
    "\n",
    "        return prefactor * _sum\n",
    "\n",
    "    #calculate f+ or fT\n",
    "\n",
    "    else:\n",
    "        prefactor = 1/(1 - q2/(mbstar**2))\n",
    "        _sum = 0\n",
    "\n",
    "        if subscript == \"T\":\n",
    "            b = bT\n",
    "        else:\n",
    "            b = bplus\n",
    "\n",
    "        for i in range(N):\n",
    "            _sum += b[i] * (np.power(z, i) - ((-1)**(i-N)) * (i/N) * np.power(z, N))\n",
    "\n",
    "        return prefactor * _sum\n",
    "\n",
    "def resonance(q, _mass, width, phase, scale):\n",
    "\n",
    "    q2 = np.power(q, 2)\n",
    "\n",
    "    mmu = pdg['muon_M']\n",
    "\n",
    "    p = 0.5 * np.sqrt(q2 - 4*(mmu**2))\n",
    "\n",
    "    p0 =  0.5 * np.sqrt(_mass**2 - 4*mmu**2)\n",
    "\n",
    "    gamma_j = p/ q2 * _mass * width / p0\n",
    "\n",
    "    #Calculate the resonance\n",
    "\n",
    "    _top = _mass * width\n",
    "\n",
    "    _bottom = np.vectorize(complex)(_mass**2 - q2, -_mass*gamma_j)\n",
    "\n",
    "    com = _top/_bottom\n",
    "\n",
    "    #Rotate by the phase\n",
    "\n",
    "    r = scale*np.abs(com)\n",
    "\n",
    "    _phase = np.angle(com)\n",
    "\n",
    "    _phase += phase\n",
    "\n",
    "    com = r * np.exp(np.vectorize(complex)(0.0 , _phase))\n",
    "\n",
    "    return com\n",
    "\n",
    "\n",
    "def axiv_nonres(q):\n",
    "\n",
    "    GF = pdg['GF']\n",
    "    alpha_ew = pdg['alpha_ew']\n",
    "    Vtb = pdg['Vtb']\n",
    "    Vts = pdg['Vts']\n",
    "    C10eff = pdg['C10eff']\n",
    "\n",
    "    mmu = pdg['muon_M']\n",
    "    mb = pdg['bquark_M']\n",
    "    ms = pdg['squark_M']\n",
    "    mK = pdg['Ks_M']\n",
    "    mB = pdg['Bplus_M']\n",
    "\n",
    "    q2 = np.power(q, 2)\n",
    "\n",
    "    #Some helperfunctions\n",
    "\n",
    "    beta = np.sqrt(np.abs(1. - 4. * mmu**2. / q2))\n",
    "\n",
    "    kabs = np.sqrt(mB**2. +np.power(q2, 2)/mB**2. + mK**4./mB**2. - 2. * (mB**2. * mK**2. + mK**2. * q2 + mB**2. * q2) / mB**2.)\n",
    "\n",
    "    #prefactor in front of whole bracket\n",
    "\n",
    "    prefactor1 = GF**2. *alpha_ew**2. * (np.abs(Vtb*Vts))**2. * kabs * beta / (128. * np.pi**5.)\n",
    "\n",
    "    #left term in bracket\n",
    "\n",
    "    bracket_left = 2./3. * kabs**2. * beta**2. *np.abs(np.vectorize(complex)(C10eff, 0.0)*formfactor(q2, \"+\"))**2.\n",
    "\n",
    "    #middle term in bracket\n",
    "\n",
    "    _top = 4. * mmu**2. * (mB**2. - mK**2.) * (mB**2. - mK**2.)\n",
    "\n",
    "    _under = q2 * mB**2.\n",
    "\n",
    "    bracket_middle = _top/_under *np.power(np.abs(np.vectorize(complex)(C10eff, 0.0) * formfactor(q2, \"0\")), 2)\n",
    "\n",
    "    #Note sqrt(q2) comes from derivation as we use q2 and plot q\n",
    "\n",
    "    return prefactor1 * (bracket_left + bracket_middle) * 2 *np.sqrt(q2)\n",
    "\n",
    "def vec(q, funcs):\n",
    "    \n",
    "    q2 = np.power(q, 2)\n",
    "\n",
    "    GF = pdg['GF']\n",
    "    alpha_ew = pdg['alpha_ew']\n",
    "    Vtb = pdg['Vtb']\n",
    "    Vts = pdg['Vts']\n",
    "    C7eff = pdg['C7eff']\n",
    "\n",
    "    mmu = pdg['muon_M']\n",
    "    mb = pdg['bquark_M']\n",
    "    ms = pdg['squark_M']\n",
    "    mK = pdg['Ks_M']\n",
    "    mB = pdg['Bplus_M']\n",
    "\n",
    "    #Some helperfunctions\n",
    "\n",
    "    beta = np.sqrt(np.abs(1. - 4. * mmu**2. / q2))\n",
    "\n",
    "    kabs = np.sqrt(mB**2. + np.power(q2, 2)/mB**2. + mK**4./mB**2. - 2 * (mB**2 * mK**2 + mK**2 * q2 + mB**2 * q2) / mB**2)\n",
    "\n",
    "    #prefactor in front of whole bracket\n",
    "\n",
    "    prefactor1 = GF**2. *alpha_ew**2. * (np.abs(Vtb*Vts))**2 * kabs * beta / (128. * np.pi**5.)\n",
    "\n",
    "    #right term in bracket\n",
    "\n",
    "    prefactor2 = kabs**2 * (1. - 1./3. * beta**2)\n",
    "\n",
    "    abs_bracket = np.abs(c9eff(q, funcs) * formfactor(q2, \"+\") + np.vectorize(complex)(2.0 * C7eff * (mb + ms)/(mB + mK), 0.0) * formfactor(q2, \"T\"))**2\n",
    "\n",
    "    bracket_right = prefactor2 * abs_bracket\n",
    "\n",
    "    #Note sqrt(q2) comes from derivation as we use q2 and plot q\n",
    "\n",
    "    return prefactor1 * bracket_right * 2 * np.sqrt(q2)\n",
    "\n",
    "def c9eff(q, funcs):\n",
    "\n",
    "    C9eff_nr = np.vectorize(complex)(pdg['C9eff'], 0.0)\n",
    "\n",
    "    c9 = C9eff_nr\n",
    "\n",
    "    c9 = c9 + funcs\n",
    "\n",
    "    return c9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "jpsi_mass, jpsi_width, jpsi_phase, jpsi_scale = pdg[\"jpsi\"]\n",
    "psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale = pdg[\"psi2s\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def t_f(x):\n",
    "        \n",
    "    def jpsi_res(q):\n",
    "        return resonance(q, _mass = jpsi_mass, scale = jpsi_scale, phase = jpsi_phase, width = jpsi_width)\n",
    "\n",
    "    def psi2s_res(q):\n",
    "        return resonance(q, _mass = psi2s_mass, scale = psi2s_scale, phase = psi2s_phase, width = psi2s_width)\n",
    "\n",
    "    funcs = jpsi_res(x) + psi2s_res(x)\n",
    "\n",
    "    vec_f = vec(x, funcs)\n",
    "\n",
    "    axiv_nr = axiv_nonres(x)\n",
    "\n",
    "    tot = vec_f + axiv_nr\n",
    "    \n",
    "    return tot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Area: 7.126305680132978e-06, rel. err.: 0.1894614654746863%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\users\\sa_li\\.conda\\envs\\rmd\\lib\\site-packages\\ipykernel_launcher.py:4: IntegrationWarning: The integral is probably divergent, or slowly convergent.\n",
      "  after removing the cwd from sys.path.\n"
     ]
    }
   ],
   "source": [
    "x_min = 2*pdg['muon_M']\n",
    "x_max = (pdg[\"Bplus_M\"]-pdg[\"Ks_M\"]-0.1)\n",
    "\n",
    "result, err = integrate.quad(lambda x: t_f(x), 3600, 3800, limit = 100000000)\n",
    "print(\"Area: {0}, rel. err.: {1}%\".format(result, err/result))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tot:  Area: 0.0013270139056058355, rel. err.: 0.001434977783938156%\n",
      "jpsi: Area: 6.527491388105805e-06, rel. err.: 0.0012105326520570792% (before scaling)\n"
     ]
    }
   ],
   "source": [
    "print(\"tot:  Area: 0.0013270139056058355, rel. err.: 0.001434977783938156%\")\n",
    "print(\"jpsi: Area: 6.527491388105805e-06, rel. err.: 0.0012105326520570792% (before scaling)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 8048 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "test_q = np.linspace(x_min, x_max, 2000000)\n",
    "calcs_test = t_f(test_q)\n",
    "plt.clf()\n",
    "# plt.plot(x_part, calcs, '.')\n",
    "plt.plot(test_q, calcs_test, label = 'pdf')\n",
    "# plt.plot(test_q, res_y, label = 'res')\n",
    "plt.legend()\n",
    "plt.ylim(0.0, 1e-8)\n",
    "# plt.yscale('log')\n",
    "# plt.xlim(3080, 3110)\n",
    "plt.savefig('np-test.png')\n",
    "# print(jpsi_width)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
