{
 "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": [
    "file = uproot.open(\n",
    "    \"tracking_losses_ntuple_B_default.root:PrDebugTrackingLosses.PrDebugTrackingTool/Tuple;1\"\n",
    ")\n",
    "\n",
    "# selektiere nur elektronen von B->K*ee\n",
    "allcolumns = file.arrays()\n",
    "found = allcolumns[\n",
    "    (allcolumns.isElectron) & (~allcolumns.lost) & (allcolumns.fromB)\n",
    "]  # B: 9056\n",
    "lost = allcolumns[\n",
    "    (allcolumns.isElectron) & (allcolumns.lost) & (allcolumns.fromB)\n",
    "]  # B: 1466\n",
    "\n",
    "electrons = allcolumns[(allcolumns.isElectron) & (allcolumns.fromB)]\n",
    "\n",
    "print(ak.num(found, axis=0), ak.num(lost, axis=0))\n",
    "print(ak.num(electrons,axis=0))\n",
    "# ak.count(found, axis=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lost_e = electrons[\"lost\"]\n",
    "e = electrons[\"energy\"]\n",
    "brem_pe = electrons[\"brem_photons_pe\"]\n",
    "brem_z = electrons[\"brem_vtx_z\"]\n",
    "brem_x = electrons[\"brem_vtx_x\"]\n",
    "length = electrons[\"brem_vtx_z_length\"]\n",
    "\n",
    "\n",
    "\n",
    "brem = ak.ArrayBuilder()\n",
    "\n",
    "for itr in range(ak.num(electrons, axis=0)):\n",
    "    brem.begin_record()\n",
    "    brem.field(\"lost\").boolean(lost_e[itr])\n",
    "    # [:,\"energy\"] energy\n",
    "    brem.field(\"energy\").append(e[itr])\n",
    "    # [:,\"photon_length\"] number of vertices\n",
    "    brem.field(\"photon_length\").integer(length[itr])\n",
    "    # [:,\"brem_photons_pe\",:] photon energy\n",
    "    brem.field(\"brem_photons_pe\").append(brem_pe[itr])\n",
    "    # [:,\"brem_vtx_z\",:] brem vtx z\n",
    "    brem.field(\"brem_vtx_x\").append(brem_x[itr])\n",
    "    brem.field(\"brem_vtx_z\").append(brem_z[itr])\n",
    "    brem.end_record()\n",
    "\n",
    "brem = ak.Array(brem)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tuner",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}