{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uproot\n",
    "import numpy as np\n",
    "import sys\n",
    "import os\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits import mplot3d\n",
    "import itertools\n",
    "import awkward as ak\n",
    "from scipy.optimize import curve_fit\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "410"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "file = uproot.open(\"tracking_losses_ntuple_Bd2KstEE.root:PrDebugTrackingLosses.PrDebugTrackingTool/Tuple;1\")\n",
    "\n",
    "\n",
    "\n",
    "#selektiere Kaonen und Pionen aus K*->Kpi - B->K*ee - und nur solche mit einem momentum von ueber 5 GeV \n",
    "allcolumns = file.arrays()\n",
    "found = allcolumns[((allcolumns.isPion) | (allcolumns.isKaon)) & (~allcolumns.lost) & (allcolumns.fromSignal) & (allcolumns.p>5e3)] #B: 8126\n",
    "lost = allcolumns[((allcolumns.isPion) | (allcolumns.isKaon)) & (allcolumns.lost) & (allcolumns.fromSignal) & (allcolumns.p>5e3)] #B: 410\n",
    "\n",
    "ak.num(lost, axis=0)\n",
    "#ak.count(found, axis=None)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9519681349578257"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def t_eff(found, lost):\n",
    "    sel = found[\"energy\"]\n",
    "    des = lost[\"energy\"]\n",
    "    return ak.count(sel,axis=None)/(ak.count(sel,axis=None)+ak.count(des,axis=None))\n",
    "\n",
    "t_eff(found, lost)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#ist die shape der teilspur im scifi anders? (koenntest du zum beispiel durch vergleich der verteilungen der fit parameter studieren,\n",
    "#in meiner thesis findest du das fitmodell -- ist einfach ein polynom dritten grades)\n",
    "z_ref=8520 #mm\n",
    "\n",
    "def scifi_track(z, a, b, c, d):\n",
    "    return a + b*(z-z_ref) + c*(z-z_ref)**2 + d*(z-z_ref)**3\n",
    "\n",
    "def z_mag(xv, zv, tx, a, b):\n",
    "    \"\"\" optical centre of the magnet is defined as the intersection between the trajectory tangents before and after the magnet\n",
    "\n",
    "    Args:\n",
    "        xv (double): velo x track\n",
    "        zv (double): velo z track\n",
    "        tx (double): velo x slope\n",
    "        a (double): ax parameter of track fit\n",
    "        b (double): bx parameter of track fit\n",
    "\n",
    "    Returns:\n",
    "        double: z_mag\n",
    "    \"\"\"\n",
    "    return (xv-tx*zv-a+b*z_ref)/(b-tx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "scifi_found = found[found[\"scifi_hit_pos_x_length\"]>3]\n",
    "scifi_lost = lost[lost[\"scifi_hit_pos_x_length\"]>3]\n",
    "\n",
    "scifi_x_found = scifi_found[\"scifi_hit_pos_x\"]\n",
    "scifi_z_found = scifi_found[\"scifi_hit_pos_z\"]\n",
    "\n",
    "tx_found = scifi_found[\"velo_track_tx\"]\n",
    "\n",
    "scifi_x_lost = scifi_lost[\"scifi_hit_pos_x\"]\n",
    "scifi_z_lost = scifi_lost[\"scifi_hit_pos_z\"]\n",
    "\n",
    "tx_lost = scifi_lost[\"velo_track_tx\"]\n",
    "\n",
    "xv_found = scifi_found[\"velo_track_x\"]\n",
    "zv_found = scifi_found[\"velo_track_z\"]\n",
    "\n",
    "xv_lost = scifi_lost[\"velo_track_x\"]\n",
    "zv_lost = scifi_lost[\"velo_track_z\"]\n",
    "\n",
    "\n",
    "\n",
    "#ak.num(scifi_found[\"energy\"], axis=0)\n",
    "#scifi_found.snapshot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "scifi_fitpars_found = ak.ArrayBuilder()\n",
    "\n",
    "for i in range(0,ak.num(scifi_found[\"energy\"], axis=0)):\n",
    "    popt, pcov = curve_fit(scifi_track,ak.to_numpy(scifi_z_found[i,:]),ak.to_numpy(scifi_x_found[i,:]))\n",
    "    scifi_fitpars_found.begin_list()\n",
    "    scifi_fitpars_found.real(popt[0])\n",
    "    scifi_fitpars_found.real(popt[1])\n",
    "    scifi_fitpars_found.real(popt[2])\n",
    "    scifi_fitpars_found.real(popt[3])\n",
    "    scifi_fitpars_found.end_list()\n",
    "\n",
    "scifi_fitpars_lost = ak.ArrayBuilder()\n",
    "\n",
    "for i in range(0,ak.num(scifi_lost[\"energy\"], axis=0)):\n",
    "    popt, pcov = curve_fit(scifi_track,ak.to_numpy(scifi_z_lost[i,:]),ak.to_numpy(scifi_x_lost[i,:]))\n",
    "    scifi_fitpars_lost.begin_list()\n",
    "    scifi_fitpars_lost.real(popt[0])\n",
    "    scifi_fitpars_lost.real(popt[1])\n",
    "    scifi_fitpars_lost.real(popt[2])\n",
    "    scifi_fitpars_lost.real(popt[3])\n",
    "    scifi_fitpars_lost.end_list()\n",
    "\n",
    "\n",
    "scifi_fitpars_lost = scifi_fitpars_lost.to_numpy()\n",
    "scifi_fitpars_found = scifi_fitpars_found.to_numpy()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 81376 characters]",
      "text/plain": [
       "<Figure size 1500x1000 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ((ax0, ax1), (ax2, ax3)) = plt.subplots(nrows=2, ncols=2, figsize=(15,10))\n",
    "\n",
    "ax0.hist(scifi_fitpars_found[:,0], bins=100, density=True, alpha=0.5, histtype='bar', color=\"blue\", label=r\"$a_x$ found\")\n",
    "ax0.hist(scifi_fitpars_lost[:,0], bins=100, density=True, alpha=0.5, histtype='bar', color=\"darkorange\", label=r\"$a_x$ lost\")\n",
    "ax0.set_xlabel(\"a\")\n",
    "ax0.set_ylabel(\"normed\")\n",
    "ax0.set_title(\"fitparameter a der scifi track\")\n",
    "ax0.legend()\n",
    "\n",
    "ax1.hist(scifi_fitpars_found[:,1], bins=100, density=True, alpha=0.5, histtype='bar', color=\"blue\", label=r\"$b_x$ found\")\n",
    "ax1.hist(scifi_fitpars_lost[:,1], bins=100, density=True, alpha=0.5, histtype='bar', color=\"darkorange\", label=r\"$b_x$ lost\")\n",
    "ax1.set_xlabel(\"b\")\n",
    "ax1.set_ylabel(\"normed\")\n",
    "ax1.set_title(\"fitparameter b der scifi track\")\n",
    "ax1.legend()\n",
    "\n",
    "ax2.hist(scifi_fitpars_found[:,2], bins=200, density=True, alpha=0.5, histtype='bar', color=\"blue\", label=r\"$c_x$ found\")\n",
    "ax2.hist(scifi_fitpars_lost[:,2], bins=200, density=True, alpha=0.5, histtype='bar', color=\"darkorange\", label=r\"$c_x$ lost\")\n",
    "ax2.set_xlim([-1e-5,1e-5])\n",
    "ax2.set_xticks(np.arange(-1e-5,1.5e-5,5e-6),minor=False)\n",
    "ax2.set_xlabel(\"c\")\n",
    "ax2.set_ylabel(\"normed\")\n",
    "ax2.set_title(\"fitparameter c der scifi track\")\n",
    "ax2.legend()\n",
    "\n",
    "ax3.hist(scifi_fitpars_found[:,3], bins=200, density=True, alpha=0.5, histtype='bar', color=\"blue\", label=r\"$d_x$ found\")\n",
    "ax3.hist(scifi_fitpars_lost[:,3], bins=200, density=True, alpha=0.5, histtype='bar', color=\"darkorange\", label=r\"$d_x$ lost\")\n",
    "ax3.set(xlim=(-2e-8,2e-8))\n",
    "ax3.set_xlabel(\"d\")\n",
    "ax3.set_ylabel(\"normed\")\n",
    "ax3.set_title(\"fitparameter d der scifi track\")\n",
    "ax3.legend()\n",
    "\n",
    "\"\"\"\n",
    "a_x: virtual hit on the reference plane\n",
    "\"\"\"\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "found\n",
      "a =  1.3758753561567658\n",
      "b =  1.4867317350363123e-05\n",
      "c =  1.0611928711675984e-09\n",
      "d =  2.524323901023556e-12\n",
      "lost\n",
      "a =  19.655460959039587\n",
      "b =  -0.0007972008149349264\n",
      "c =  -1.6113881305735422e-07\n",
      "d =  8.074893917737833e-11\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "Electrons\n",
    "found\n",
    "a =  -0.6718207391527037    b =  0.0013778237292529144      c =  3.3126998287416195e-08     d =  -1.0330674442255529e-10\n",
    "lost\n",
    "a =  -36.98764338200992     b =  -0.015685137956233643      c =  -8.265859479503501e-07     d =  -1.541510766903436e-11\n",
    "\n",
    "Kaon and Pions\n",
    "found\n",
    "a =  1.3758753561567658     b =  1.4867317350363123e-05     c =  1.0611928711675984e-09     d =  2.524323901023556e-12\n",
    "lost\n",
    "a =  19.655460959039587     b =  -0.0007972008149349264     c =  -1.6113881305735422e-07    d =  8.074893917737833e-11\n",
    "\n",
    "Kaons\n",
    "found\n",
    "a =  1.553371699292705\t\tb =  8.625619976952656e-05\t\tc =  3.1024420569192613e-09\t\td =  -2.028183118694356e-13\n",
    "lost\n",
    "a =  50.918713410777336\t\tb =  0.005102900113795297\t\tc =  -7.416239379215659e-09\t\td =  5.64633728533602e-11\n",
    "\n",
    "Pions\n",
    "found\n",
    "a =  1.1512709788236055\t\tb =  -7.54683774644083e-05\t\tc =  -1.5218096529826036e-09\t\td =  5.975257412125118e-12\n",
    "lost\n",
    "a =  -3.625684483743866\t\tb =  -0.005190892995904242\t\tc =  -2.7561307005171267e-07\t\td =  9.883393537611524e-11\n",
    "\"\"\"\n",
    "\n",
    "print(\"found\")\n",
    "print(\"a = \", str(np.mean(scifi_fitpars_found[:,0])))\n",
    "print(\"b = \", str(np.mean(scifi_fitpars_found[:,1])))\n",
    "print(\"c = \", str(np.mean(scifi_fitpars_found[:,2])))\n",
    "print(\"d = \", str(np.mean(scifi_fitpars_found[:,3])))\n",
    "\n",
    "print(\"lost\")\n",
    "print(\"a = \", str(np.mean(scifi_fitpars_lost[:,0])))\n",
    "print(\"b = \", str(np.mean(scifi_fitpars_lost[:,1])))\n",
    "print(\"c = \", str(np.mean(scifi_fitpars_lost[:,2])))\n",
    "print(\"d = \", str(np.mean(scifi_fitpars_lost[:,3])))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 939836 characters]",
      "text/plain": [
       "<Figure size 1500x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ((ax0, ax1)) = plt.subplots(nrows=1, ncols=2, figsize=(15,6))\n",
    "\n",
    "for i in range(0,ak.num(scifi_found[\"energy\"], axis=0)):\n",
    "    z_coord = np.linspace(scifi_z_found[i,0],12000,300)\n",
    "    fit = scifi_track(z_coord, *scifi_fitpars_found[i])\n",
    "    ax0.plot(z_coord, fit, \"-\", lw=0.5)\n",
    "    ax0.errorbar(ak.to_numpy(scifi_z_found[i,:]),ak.to_numpy(scifi_x_found[i,:]),fmt=\".\",ms=2)\n",
    "\n",
    "#ax0.legend()\n",
    "ax0.set_xlabel(\"z [mm]\")\n",
    "ax0.set_ylabel(\"x [mm]\")\n",
    "ax0.set_title(\"found tracks of scifi hits\")\n",
    "ax0.set(xlim=(7e3,12000), ylim=(-4000,4000))\n",
    "ax0.grid()\n",
    "\n",
    "for i in range(0,ak.num(scifi_lost[\"energy\"], axis=0)):\n",
    "    z_coord = np.linspace(scifi_z_lost[i,0],12000,300)\n",
    "    fit = scifi_track(z_coord, *scifi_fitpars_lost[i])\n",
    "    ax1.plot(z_coord, fit, \"-\", lw=0.5)\n",
    "    ax1.errorbar(ak.to_numpy(scifi_z_lost[i,:]),ak.to_numpy(scifi_x_lost[i,:]),fmt=\".\",ms=2)\n",
    "\n",
    "#ax1.legend()\n",
    "ax1.set_xlabel(\"z [mm]\")\n",
    "ax1.set_ylabel(\"x [mm]\")\n",
    "ax1.set_title(\"lost tracks of scifi hits\")\n",
    "ax1.set(xlim=(7e3,12000), ylim=(-4000,4000))\n",
    "ax1.grid()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "found \n",
      "zmag =  5196.312017664934\n",
      "lost \n",
      "zmag =  5200.71031871899\n"
     ]
    }
   ],
   "source": [
    "#vergleich der zmag werte\n",
    "zmag_found = z_mag(xv_found, zv_found, tx_found, scifi_fitpars_found[:,0], scifi_fitpars_found[:,1])\n",
    "zmag_lost = z_mag(xv_lost, zv_lost, tx_lost, scifi_fitpars_lost[:,0], scifi_fitpars_lost[:,1])\n",
    "zmag_lost = zmag_lost[~np.isnan(zmag_lost)]\n",
    "zmag_found = zmag_found[~np.isnan(zmag_found)]\n",
    "\n",
    "print(\"found \\nzmag = \", str(np.mean(zmag_found)))\n",
    "print(\"lost \\nzmag = \", str(np.mean(zmag_lost)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 31012 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.hist(zmag_found, bins=2000, density=True, alpha=0.5, histtype='bar', color=\"blue\", label=\"found\")\n",
    "plt.hist(zmag_lost, bins=500, density=True, alpha=0.5, histtype=\"bar\",color=\"darkorange\", label=\"lost\")\n",
    "plt.xlabel(\"$z_{mag}$ [mm]\")\n",
    "plt.ylabel(\"normed\")\n",
    "plt.title(\"magnet kick position $z_{mag}$ for $K$ and $\\pi$\")\n",
    "plt.legend()\n",
    "plt.xticks(np.arange(4900,5405,5), minor=True)\n",
    "plt.yticks(np.arange(0,0.015,0.001), minor=True)\n",
    "plt.xlim(4900,5400)\n",
    "\n",
    "\"\"\"\n",
    "the distribution for lost K and pi has a peak at around the same zmag value as found but its less compact and seems chaotic outside the range of the peak.\n",
    "the distribution for found particles appears compact around the peak.\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env1",
   "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.11.5"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
