{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uproot\t\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits import mplot3d\n",
    "import awkward as ak\n",
    "from scipy.optimize import curve_fit\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "# found\n",
    "\n",
    "brem_e = electrons[\"brem_photons_pe\"]\n",
    "brem_z = electrons[\"brem_vtx_z\"]\n",
    "brem_x = electrons[\"brem_vtx_x\"]\n",
    "e = electrons[\"energy\"]\n",
    "length = electrons[\"brem_vtx_z_length\"]\n",
    "\n",
    "\n",
    "\n",
    "brem_f = ak.ArrayBuilder()\n",
    "\n",
    "for itr in range(ak.num(found, axis=0)):\n",
    "    brem_f.begin_record()\n",
    "    brem_f.field(\"lost\").boolean()\n",
    "    # [:,\"energy\"] energy\n",
    "    brem_f.field(\"energy\").append(e_f[itr])\n",
    "    # [:,\"photon_length\"] number of vertices\n",
    "    brem_f.field(\"photon_length\").integer(length_f[itr])\n",
    "    # [:,\"brem_photons_pe\",:] photon energy\n",
    "    brem_f.field(\"brem_photons_pe\").append(brem_e_f[itr])\n",
    "    # [:,\"brem_vtx_z\",:] brem vtx z\n",
    "    brem_f.field(\"brem_vtx_x\").append(brem_x_f[itr])\n",
    "    brem_f.field(\"brem_vtx_z\").append(brem_z_f[itr])\n",
    "    brem_f.end_record()\n",
    "\n",
    "brem_f = ak.Array(brem_f)\n",
    "\n",
    "# lost\n",
    "\n",
    "brem_e_l = lost[\"brem_photons_pe\"]\n",
    "brem_z_l = lost[\"brem_vtx_z\"]\n",
    "brem_x_l = lost[\"brem_vtx_x\"]\n",
    "e_l = lost[\"energy\"]\n",
    "length_l = lost[\"brem_vtx_z_length\"]\n",
    "\n",
    "brem_l = ak.ArrayBuilder()\n",
    "\n",
    "for itr in range(ak.num(lost, axis=0)):\n",
    "    brem_l.begin_record()\n",
    "    # [:,\"energy\"] energy\n",
    "    brem_l.field(\"energy\").append(e_l[itr])\n",
    "    # [:,\"photon_length\"] number of vertices\n",
    "    brem_l.field(\"photon_length\").integer(length_l[itr])\n",
    "    # [:,\"brem_photons_pe\",:] photon energy\n",
    "    brem_l.field(\"brem_photons_pe\").append(brem_e_l[itr])\n",
    "    # [:,\"brem_vtx_z\",:] brem vtx z\n",
    "    brem_l.field(\"brem_vtx_x\").append(brem_x_l[itr])\n",
    "    brem_l.field(\"brem_vtx_z\").append(brem_z_l[itr])\n",
    "    brem_l.end_record()\n",
    "\n",
    "brem_l = ak.Array(brem_l)\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "photon_cut = 0\n",
    "photon_cut_ratio = 0.2\n",
    "\n",
    "cut_brem_found = ak.ArrayBuilder()\n",
    "\n",
    "for itr in range(ak.num(brem_f, axis=0)):\n",
    "    cut_brem_found.begin_record()\n",
    "    cut_brem_found.field(\"energy\").real(brem_f[itr, \"energy\"])\n",
    "\n",
    "    tmp_energy = brem_f[itr, \"energy\"]\n",
    "\n",
    "    cut_brem_found.field(\"brem_photons_pe\")\n",
    "    cut_brem_found.begin_list()\n",
    "    for jentry in range(brem_f[itr, \"photon_length\"]):\n",
    "        if (\n",
    "            brem_f[itr, \"brem_vtx_z\", jentry] > 5000\n",
    "            or brem_f[itr, \"brem_photons_pe\", jentry] < photon_cut\n",
    "            or brem_f[itr, \"brem_photons_pe\", jentry] < photon_cut_ratio * tmp_energy\n",
    "        ):\n",
    "            continue\n",
    "        else:\n",
    "            cut_brem_found.real(brem_f[itr, \"brem_photons_pe\", jentry])\n",
    "            tmp_energy -= brem_f[itr, \"brem_photons_pe\", jentry]\n",
    "    cut_brem_found.end_list()\n",
    "\n",
    "    tmp_energy = brem_f[itr, \"energy\"]\n",
    "\n",
    "    cut_brem_found.field(\"brem_vtx_x\")\n",
    "    cut_brem_found.begin_list()\n",
    "    for jentry in range(brem_f[itr, \"photon_length\"]):\n",
    "        if (\n",
    "            brem_f[itr, \"brem_vtx_z\", jentry] > 5000\n",
    "            or brem_f[itr, \"brem_photons_pe\", jentry] < photon_cut\n",
    "            or brem_f[itr, \"brem_photons_pe\", jentry] < photon_cut_ratio * tmp_energy\n",
    "        ):\n",
    "            continue\n",
    "        else:\n",
    "            cut_brem_found.real(brem_f[itr, \"brem_vtx_x\", jentry])\n",
    "            tmp_energy -= brem_f[itr, \"brem_photons_pe\", jentry]\n",
    "    cut_brem_found.end_list()\n",
    "\n",
    "    tmp_energy = brem_f[itr, \"energy\"]\n",
    "\n",
    "    cut_brem_found.field(\"brem_vtx_z\")\n",
    "    cut_brem_found.begin_list()\n",
    "    for jentry in range(brem_f[itr, \"photon_length\"]):\n",
    "        if (\n",
    "            brem_f[itr, \"brem_vtx_z\", jentry] > 5000\n",
    "            or brem_f[itr, \"brem_photons_pe\", jentry] < photon_cut\n",
    "            or brem_f[itr, \"brem_photons_pe\", jentry] < photon_cut_ratio * tmp_energy\n",
    "        ):\n",
    "            continue\n",
    "        else:\n",
    "            cut_brem_found.real(brem_f[itr, \"brem_vtx_z\", jentry])\n",
    "            tmp_energy -= brem_f[itr, \"brem_photons_pe\", jentry]\n",
    "    cut_brem_found.end_list()\n",
    "\n",
    "    cut_brem_found.end_record()\n",
    "\n",
    "tuple_found = ak.Array(cut_brem_found)\n",
    "\n",
    "cut_brem_lost = ak.ArrayBuilder()\n",
    "\n",
    "for itr in range(ak.num(brem_l, axis=0)):\n",
    "    cut_brem_lost.begin_record()\n",
    "    cut_brem_lost.field(\"energy\").real(brem_l[itr, \"energy\"])\n",
    "\n",
    "    tmp_energy = brem_l[itr, \"energy\"]\n",
    "\n",
    "    cut_brem_lost.field(\"brem_photons_pe\")\n",
    "    cut_brem_lost.begin_list()\n",
    "    for jentry in range(brem_l[itr, \"photon_length\"]):\n",
    "        if (\n",
    "            brem_l[itr, \"brem_vtx_z\", jentry] > 5000\n",
    "            or brem_l[itr, \"brem_photons_pe\", jentry] < photon_cut\n",
    "            or brem_l[itr, \"brem_photons_pe\", jentry] < photon_cut_ratio * tmp_energy\n",
    "        ):\n",
    "            continue\n",
    "        else:\n",
    "            cut_brem_lost.real(brem_l[itr, \"brem_photons_pe\", jentry])\n",
    "            tmp_energy -= brem_l[itr, \"brem_photons_pe\", jentry]\n",
    "    cut_brem_lost.end_list()\n",
    "\n",
    "    tmp_energy = brem_l[itr, \"energy\"]\n",
    "\n",
    "    cut_brem_lost.field(\"brem_vtx_x\")\n",
    "    cut_brem_lost.begin_list()\n",
    "    for jentry in range(brem_l[itr, \"photon_length\"]):\n",
    "        if (\n",
    "            brem_l[itr, \"brem_vtx_z\", jentry] > 5000\n",
    "            or brem_l[itr, \"brem_photons_pe\", jentry] < photon_cut\n",
    "            or brem_l[itr, \"brem_photons_pe\", jentry] < photon_cut_ratio * tmp_energy\n",
    "        ):\n",
    "            continue\n",
    "        else:\n",
    "            cut_brem_lost.real(brem_l[itr, \"brem_vtx_x\", jentry])\n",
    "            tmp_energy -= brem_l[itr, \"brem_photons_pe\", jentry]\n",
    "    cut_brem_lost.end_list()\n",
    "\n",
    "    tmp_energy = brem_l[itr, \"energy\"]\n",
    "\n",
    "    cut_brem_lost.field(\"brem_vtx_z\")\n",
    "    cut_brem_lost.begin_list()\n",
    "    for jentry in range(brem_l[itr, \"photon_length\"]):\n",
    "        if (\n",
    "            brem_l[itr, \"brem_vtx_z\", jentry] > 5000\n",
    "            or brem_l[itr, \"brem_photons_pe\", jentry] < photon_cut\n",
    "            or brem_l[itr, \"brem_photons_pe\", jentry] < photon_cut_ratio * tmp_energy\n",
    "        ):\n",
    "            continue\n",
    "        else:\n",
    "            cut_brem_lost.real(brem_l[itr, \"brem_vtx_z\", jentry])\n",
    "            tmp_energy -= brem_l[itr, \"brem_photons_pe\", jentry]\n",
    "    cut_brem_lost.end_list()\n",
    "\n",
    "    cut_brem_lost.end_record()\n",
    "\n",
    "tuple_lost = ak.Array(cut_brem_lost)\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "electrons_lost = ak.ArrayBuilder()\n",
    "\n",
    "for jelec in range(ak.num(tuple_lost, axis=0)):\n",
    "    electrons_lost.begin_record()\n",
    "    electrons_lost.field(\"energy\").real(tuple_lost[jelec, \"energy\"])\n",
    "\n",
    "    tmp_velo = 0\n",
    "    tmp_richut = 0\n",
    "    for jphoton in range(ak.num(tuple_lost[jelec][\"brem_photons_pe\"], axis=0)):\n",
    "        if tuple_lost[jelec, \"brem_vtx_z\", jphoton] <= 800:\n",
    "            tmp_velo += tuple_lost[jelec, \"brem_photons_pe\", jphoton]\n",
    "        elif (tuple_lost[jelec, \"brem_vtx_z\", jphoton] > 800) and (\n",
    "            tuple_lost[jelec, \"brem_vtx_z\", jphoton] <= 3000\n",
    "        ):\n",
    "            tmp_richut += tuple_lost[jelec, \"brem_photons_pe\", jphoton]\n",
    "\n",
    "    electrons_lost.field(\"velo\").real(tmp_velo)\n",
    "\n",
    "    electrons_lost.field(\"rich\").real(tmp_richut)\n",
    "    \n",
    "    if (tmp_velo==0) and (tmp_richut==0):\n",
    "        electrons_lost.field(\"quality\").integer(0)\n",
    "    else:\n",
    "        electrons_lost.field(\"quality\").integer(1)\n",
    "\n",
    "    electrons_lost.end_record()\n",
    "\n",
    "electrons_lost = ak.Array(electrons_lost)\n",
    "\n",
    "electrons_found = ak.ArrayBuilder()\n",
    "\n",
    "for jelec in range(ak.num(tuple_found, axis=0)):\n",
    "    electrons_found.begin_record()\n",
    "    electrons_found.field(\"energy\").real(tuple_found[jelec, \"energy\"])\n",
    "\n",
    "    tmp_velo = 0\n",
    "    tmp_richut = 0\n",
    "    for jphoton in range(ak.num(tuple_found[jelec][\"brem_photons_pe\"], axis=0)):\n",
    "        if tuple_found[jelec, \"brem_vtx_z\", jphoton] <= 850:\n",
    "            tmp_velo += tuple_found[jelec, \"brem_photons_pe\", jphoton]\n",
    "        elif (tuple_found[jelec, \"brem_vtx_z\", jphoton] > 850) and (\n",
    "            tuple_found[jelec, \"brem_vtx_z\", jphoton] <= 3000\n",
    "        ):\n",
    "            tmp_richut += tuple_found[jelec, \"brem_photons_pe\", jphoton]\n",
    "\n",
    "    electrons_found.field(\"velo\").real(tmp_velo)\n",
    "\n",
    "    electrons_found.field(\"rich\").real(tmp_richut)\n",
    "    \n",
    "    if (tmp_velo==0) and (tmp_richut==0):\n",
    "        electrons_found.field(\"quality\").integer(0)\n",
    "    else:\n",
    "        electrons_found.field(\"quality\").integer(1)\n",
    "\n",
    "    electrons_found.end_record()\n",
    "\n",
    "electrons_found = ak.Array(electrons_found)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def ratio(nominator, denom):\n",
    "    denominator = ak.num(denom[\"energy\"], axis=0)\n",
    "    return nominator / denominator\n",
    "\n",
    "\n",
    "def eff(found, lost):\n",
    "    return found / (found + lost)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_elec_lost = ak.num(tuple_lost, axis=0)\n",
    "\n",
    "velo_lost = 0\n",
    "\n",
    "richut_lost = 0\n",
    "\n",
    "for jelec in range(ak.num(tuple_lost, axis=0)):\n",
    "    veloemitted = False\n",
    "    richemitted = False\n",
    "    for jphoton in range(ak.num(tuple_lost[jelec][\"brem_photons_pe\"], axis=0)):\n",
    "        if (tuple_lost[jelec, \"brem_vtx_z\", jphoton] <= 770) and (veloemitted == False):\n",
    "            velo_lost += 1\n",
    "            veloemitted = True\n",
    "        elif (\n",
    "            (tuple_lost[jelec, \"brem_vtx_z\", jphoton] > 770)\n",
    "            and (tuple_lost[jelec, \"brem_vtx_z\", jphoton] <= 3000)\n",
    "            and (richemitted == False)\n",
    "        ):\n",
    "            richut_lost += 1\n",
    "            richemitted = True\n",
    "\n",
    "print(\"ratio of lost electrons emitting in Velo: \", ratio(velo_lost, brem[brem.lost]))\n",
    "print(\"ratio of lost electrons emitting in Rich1+UT: \", ratio(richut_lost, brem[brem.lost]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# efficiency berechnen als found in velo oder rich über alle elektronen\n",
    "# dann kann man zusammenrechnen mit velo, rich, und allen anderen elektronen\n",
    "\n",
    "num_velo_found = 0\n",
    "num_rich_found = 0\n",
    "for itr in range(ak.num(electrons_found, axis=0)):\n",
    "    if (electrons_found[itr, \"quality\"]==1):\n",
    "        if (electrons_found[itr, \"velo\"] >= electrons_found[itr, \"rich\"]):\n",
    "            num_velo_found += 1\n",
    "        else:\n",
    "            num_rich_found += 1\n",
    "\n",
    "num_velo_lost = 0\n",
    "num_rich_lost = 0\n",
    "for itr in range(ak.num(electrons_lost, axis=0)):\n",
    "    if (electrons_lost[itr, \"quality\"]==1):\n",
    "        if electrons_lost[itr, \"velo\"] >= electrons_lost[itr, \"rich\"]:\n",
    "            num_velo_lost += 1\n",
    "        else:\n",
    "            num_rich_lost += 1\n",
    "\n",
    "lost_elec = ak.num(electrons[electrons.lost],axis=0)\n",
    "found_elec = ak.num(electrons[~(electrons.lost)],axis=0)\n",
    "denom = lost_elec+found_elec\n",
    "\n",
    "eff_velo = num_velo_found/denom\n",
    "\n",
    "eff_rich = num_rich_found/denom\n",
    "\n",
    "eff_other = ak.num(electrons_found[electrons_found[\"quality\"]==0],axis=0)/denom\n",
    "\n",
    "print(\"VELO energy emission, eff: \", eff_velo)\n",
    "\n",
    "print(\"RICH1+UT energy emission, eff: \", eff_rich)\n",
    "\n",
    "print(\"Neither, eff: \", eff_other)\n",
    "\n",
    "print(\"total efficiency: \", eff_velo + eff_rich + eff_other)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}