notebooks/BTA_array_producer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from coffea import processor, hist\n",
    "import awkward as ak\n",
    "import numpy as np\n",
    "import uproot\n",
    "import numba as nb\n",
    "\n",
    "from functools import partial\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "from cycler import cycler\n",
    "\n",
    "# mpl.use('Agg')\n",
    "# mpl.rcParams['axes.prop_cycle'] = cycler(color=['blue', 'red', 'green', 'violet', 'darkorange', 'black', 'cyan', 'yellow'])\n",
    "\n",
    "import boost_histogram as bh\n",
    "import mplhep as hep\n",
    "\n",
    "plt.style.use(hep.style.CMS)\n",
    "\n",
    "from itertools import product\n",
    "import shutil\n",
    "import yaml\n",
    "from types import SimpleNamespace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from coffea.nanoevents import NanoEventsFactory, PFNanoAODSchema\n",
    "from coffea.nanoevents.methods import vector\n",
    "\n",
    "events = NanoEventsFactory.from_root(\n",
    "    \"../testfiles/nano_mc2018_ULv2.root\",\n",
    "    schemaclass=PFNanoAODSchema,\n",
    "    metadata={\"dataset\": \"test\"},\n",
    ").events()\n",
    "\n",
    "dataset = events.metadata[\"dataset\"]\n",
    "isRealData = not hasattr(events, \"genWeight\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Basic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "basic_vars = {\n",
    "    \"Run\": events.run,\n",
    "    \"Evt\": events.event,\n",
    "    \"LumiBlock\": events.luminosityBlock,\n",
    "    \"nPU\": events.Pileup.nPU,\n",
    "    \"nPUtrue\": events.Pileup.nTrueInt,\n",
    "    \"rho\": events.fixedGridRhoFastjetAll,\n",
    "    \"pthat\": events.Generator.binvar,\n",
    "}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# HLT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "## HLT\n",
    "\n",
    "\n",
    "@nb.njit\n",
    "def to_bitwise_trigger(pass_trig, builder):\n",
    "    for it in pass_trig:\n",
    "        # group by every 32 bits\n",
    "        builder.begin_list()\n",
    "        for bitidx in range(len(it) // 32 + 1):\n",
    "            trig = 0\n",
    "            start = bitidx * 32\n",
    "            end = min((bitidx + 1) * 32, len(it))\n",
    "            for i, b in enumerate(it[start:end]):\n",
    "                trig += b << i\n",
    "            builder.integer(trig)\n",
    "        builder.end_list()\n",
    "\n",
    "    return builder\n",
    "\n",
    "\n",
    "triggers = [\n",
    "    \"PFJet40\",\n",
    "    \"PFJet60\",\n",
    "    \"PFJet80\",\n",
    "    \"PFJet140\",\n",
    "    \"PFJet200\",\n",
    "    \"PFJet260\",\n",
    "    \"PFJet320\",\n",
    "    \"PFJet400\",\n",
    "    \"PFJet450\",\n",
    "    \"PFJet500\",\n",
    "    \"PFJet550\",\n",
    "    \"AK8PFJet40\",\n",
    "    \"AK8PFJet60\",\n",
    "    \"AK8PFJet80\",\n",
    "    \"AK8PFJet140\",\n",
    "    \"AK8PFJet200\",\n",
    "    \"AK8PFJet260\",\n",
    "    \"AK8PFJet320\",\n",
    "    \"AK8PFJet400\",\n",
    "    \"AK8PFJet450\",\n",
    "    \"AK8PFJet500\",\n",
    "    \"AK8PFJet550\",\n",
    "    \"PFHT180\",\n",
    "    \"PFHT250\",\n",
    "    \"PFHT370\",\n",
    "    \"PFHT430\",\n",
    "    \"PFHT510\",\n",
    "    \"PFHT590\",\n",
    "    \"PFHT680\",\n",
    "    \"PFHT780\",\n",
    "    \"PFHT890\",\n",
    "    \"PFHT1050\",\n",
    "    \"BTagMu_AK4DiJet20_Mu5\",\n",
    "    \"BTagMu_AK4DiJet40_Mu5\",\n",
    "    \"BTagMu_AK4DiJet70_Mu5\",\n",
    "    \"BTagMu_AK4DiJet110_Mu5\",\n",
    "    \"BTagMu_AK4DiJet170_Mu5\",\n",
    "    \"BTagMu_AK4Jet300_Mu5\",\n",
    "    \"BTagMu_AK8DiJet170_Mu5\",\n",
    "    \"BTagMu_AK8Jet300_Mu5\",\n",
    "    \"BTagMu_AK4DiJet20_Mu5_noalgo\",\n",
    "    \"BTagMu_AK4DiJet40_Mu5_noalgo\",\n",
    "    \"BTagMu_AK4DiJet70_Mu5_noalgo\",\n",
    "    \"BTagMu_AK4DiJet110_Mu5_noalgo\",\n",
    "    \"BTagMu_AK4DiJet170_Mu5_noalgo\",\n",
    "    \"BTagMu_AK4Jet300_Mu5_noalgo\",\n",
    "    \"BTagMu_AK8DiJet170_Mu5_noalgo\",\n",
    "    \"BTagMu_AK8Jet300_Mu5_noalgo\",\n",
    "    \"BTagMu_AK8Jet170_DoubleMu5_noalgo\",\n",
    "]\n",
    "checkHLT = ak.Array([hasattr(events.HLT, _trig) for _trig in triggers])\n",
    "if ak.all(checkHLT == False):\n",
    "    raise ValueError(\"HLT paths:\", triggers, \" are all invalid in\", dataset)\n",
    "elif ak.any(checkHLT == False):\n",
    "    print(\n",
    "        np.array(triggers)[~checkHLT], \" does not exist in\", dataset, \", please check!\"\n",
    "    )\n",
    "\n",
    "# pass_trigger dim: (n_evt, n_trig), append False for non-exist triggers\n",
    "pass_trig = np.array(\n",
    "    [\n",
    "        getattr(events.HLT, _trig, np.zeros(len(events), dtype=\"bool\"))\n",
    "        for _trig in triggers\n",
    "    ]\n",
    ").T\n",
    "\n",
    "# convert to bitwise trigger\n",
    "# basic_vars['nBitTrigger'] = len(triggers) // 32 + 1\n",
    "basic_vars[\"BitTrigger\"] = to_bitwise_trigger(pass_trig, ak.ArrayBuilder()).snapshot()\n",
    "\n",
    "# PV\n",
    "PV = ak.zip(\n",
    "    {\n",
    "        \"x\": events.PV.x,\n",
    "        \"y\": events.PV.y,\n",
    "        \"z\": events.PV.z,\n",
    "        \"ndof\": events.PV.ndof,\n",
    "        \"chi2\": events.PV.chi2,\n",
    "    }\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quarks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Quarks\n",
    "\n",
    "\n",
    "# @nb.njit # sometime has bug... strange!\n",
    "def is_from_GSP(\n",
    "    target_idx_list, pdgid_list, mom_idx_list, mom2_idx_list, status_list, builder\n",
    "):\n",
    "    # check if the quark is from gluon-splitting\n",
    "    for target_idx, pdgid, mom_idx, mom2_idx, status in zip(\n",
    "        target_idx_list, pdgid_list, mom_idx_list, mom2_idx_list, status_list\n",
    "    ):\n",
    "        # loop over events\n",
    "\n",
    "        builder.begin_list()\n",
    "        for idx in target_idx:\n",
    "            # loop over target (b/c quark) index\n",
    "\n",
    "            i = idx\n",
    "            quark_pdgid = abs(pdgid[idx])\n",
    "            from_GSP = False\n",
    "            while (\n",
    "                mom2_idx[i] == -1\n",
    "                and mom_idx[i] >= 0\n",
    "                and (\n",
    "                    abs(pdgid[mom_idx[i]]) == 21\n",
    "                    or abs(pdgid[mom_idx[i]]) == quark_pdgid\n",
    "                )\n",
    "                and not (\n",
    "                    abs(pdgid[mom_idx[i]]) == quark_pdgid\n",
    "                    and status[mom_idx[i]] >= 21\n",
    "                    and status[mom_idx[i]] <= 29\n",
    "                )\n",
    "            ):\n",
    "                # tracing mother upwards, but the mother cannot be a hard process quark\n",
    "                if pdgid[mom_idx[i]] == 21:  # is gluon\n",
    "                    from_GSP = True\n",
    "                    break\n",
    "                i = mom_idx[i]\n",
    "            builder.boolean(from_GSP)\n",
    "\n",
    "        builder.end_list()\n",
    "    return builder\n",
    "\n",
    "\n",
    "for quark_pdgid in [4, 5]:\n",
    "    # finding b or c quarks\n",
    "\n",
    "    quark_sel = (abs(events.GenPart.pdgId) == quark_pdgid) & (\n",
    "        (events.GenPart.statusFlags & (1 << 13)) > 0\n",
    "    )  # is b/c-quark & is last copy\n",
    "    quark_index = ak.local_index(events.GenPart.pdgId)[quark_sel]\n",
    "    quarks = events.GenPart[quark_sel]\n",
    "    out = ak.zip(\n",
    "        {\n",
    "            \"pT\": quarks.pt,\n",
    "            \"eta\": quarks.eta,\n",
    "            \"phi\": quarks.phi,\n",
    "            \"pdgId\": quarks.pdgId,\n",
    "            \"status\": quarks.status,\n",
    "            \"fromGSP\": ak.values_astype(\n",
    "                is_from_GSP(\n",
    "                    quark_index,\n",
    "                    events.GenPart.pdgId,\n",
    "                    events.GenPart.genPartIdxMother,\n",
    "                    events.GenPart.genPartIdxMother2,\n",
    "                    events.GenPart.status,\n",
    "                    ak.ArrayBuilder(),\n",
    "                ).snapshot(),\n",
    "                int,\n",
    "            ),\n",
    "        }\n",
    "    )\n",
    "    if quark_pdgid == 4:\n",
    "        cQuark = out\n",
    "    elif quark_pdgid == 5:\n",
    "        bQuark = out"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hadrons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Hadrons\n",
    "\n",
    "# mass table from https://github.com/scikit-hep/particle/blob/master/src/particle/data/particle2022.csv\n",
    "hadron_mass_table = {\n",
    "    411: 1.86966,\n",
    "    413: 2.01026,\n",
    "    415: 2.4611,\n",
    "    421: 1.86484,\n",
    "    423: 2.00685,\n",
    "    425: 2.4611,\n",
    "    431: 1.96835,\n",
    "    433: 2.1122,\n",
    "    435: 2.5691,\n",
    "    441: 2.9839,\n",
    "    443: 3.0969,\n",
    "    445: 3.55617,\n",
    "    511: 5.27966,\n",
    "    513: 5.32471,\n",
    "    515: 5.7395,\n",
    "    521: 5.27934,\n",
    "    523: 5.32471,\n",
    "    525: 5.7372,\n",
    "    531: 5.36692,\n",
    "    533: 5.4154,\n",
    "    535: 5.83986,\n",
    "    541: 6.27447,\n",
    "    553: 9.4603,\n",
    "    555: 9.9122,\n",
    "    4101: -0.001,\n",
    "    4103: -0.001,\n",
    "    4112: 2.45375,\n",
    "    4114: 2.51848,\n",
    "    4122: 2.28646,\n",
    "    4132: 2.47044,\n",
    "    4201: -0.001,\n",
    "    4203: -0.001,\n",
    "    4212: 2.45265,\n",
    "    4214: 2.5174,\n",
    "    4222: 2.45397,\n",
    "    4224: 2.51841,\n",
    "    4232: 2.46771,\n",
    "    4301: -0.001,\n",
    "    4303: -0.001,\n",
    "    4312: 2.5787,\n",
    "    4314: 2.64616,\n",
    "    4322: 2.5782,\n",
    "    4324: 2.6451,\n",
    "    4332: 2.6952,\n",
    "    4334: 2.7659,\n",
    "    4403: -0.001,\n",
    "    5101: -0.001,\n",
    "    5103: -0.001,\n",
    "    5112: 5.81564,\n",
    "    5114: 5.83474,\n",
    "    5122: 5.6196,\n",
    "    5132: 5.797,\n",
    "    5201: -0.001,\n",
    "    5203: -0.001,\n",
    "    5222: 5.81056,\n",
    "    5224: 5.83032,\n",
    "    5232: 5.7919,\n",
    "    5301: -0.001,\n",
    "    5303: -0.001,\n",
    "    5332: 6.0452,\n",
    "    5401: -0.001,\n",
    "    5403: -0.001,\n",
    "    5503: -0.001,\n",
    "}\n",
    "hadron_pids = ak.Array(list(hadron_mass_table.keys()))\n",
    "hadron_masses = ak.Array(list(hadron_mass_table.values()))\n",
    "\n",
    "\n",
    "@nb.njit\n",
    "def get_hadron_mass(pids_list, hadron_pids, hadron_masses, builder):\n",
    "    for pids in pids_list:\n",
    "        # loop over events\n",
    "        builder.begin_list()\n",
    "        for p in pids:\n",
    "            for i in range(len(hadron_pids)):\n",
    "                if hadron_pids[i] == abs(p):\n",
    "                    builder.real(hadron_masses[i])\n",
    "                    break\n",
    "            else:\n",
    "                builder.real(-1)\n",
    "        builder.end_list()\n",
    "    return builder\n",
    "\n",
    "\n",
    "is_heavy_hadron = lambda p, pid: (abs(p.pdgId) // 100 == pid) | (\n",
    "    abs(p.pdgId) // 1000 == pid\n",
    ")\n",
    "\n",
    "# finding D hadrons\n",
    "sel = is_heavy_hadron(events.GenPart, 4) & (\n",
    "    (events.GenPart.statusFlags & (1 << 13)) > 0\n",
    ")  # PID match, is last copy\n",
    "chadrons = events.GenPart[sel]\n",
    "chadrons = chadrons[\n",
    "    ~ak.any(is_heavy_hadron(chadrons.children, 4), axis=-1)\n",
    "]  # should not select D hadron whose daughters also includes D hadrons\n",
    "DHadron = ak.zip(\n",
    "    {\n",
    "        \"pT\": chadrons.pt,\n",
    "        \"eta\": chadrons.eta,\n",
    "        \"phi\": chadrons.phi,\n",
    "        \"pdgID\": chadrons.pdgId,\n",
    "        \"mass\": get_hadron_mass(\n",
    "            chadrons.pdgId, hadron_pids, hadron_masses, ak.ArrayBuilder()\n",
    "        ).snapshot(),\n",
    "        \"muInDaughter\": ak.any(abs(chadrons.children.pdgId) == 13, axis=-1),\n",
    "    }\n",
    ")\n",
    "\n",
    "# finding B hadrons\n",
    "sel = is_heavy_hadron(events.GenPart, 5) & (\n",
    "    (events.GenPart.statusFlags & (1 << 13)) > 0\n",
    ")  # PID match, is last copy\n",
    "bhadrons = events.GenPart[sel]\n",
    "BHadron = ak.zip(\n",
    "    {\n",
    "        \"pT\": bhadrons.pt,\n",
    "        \"eta\": bhadrons.eta,\n",
    "        \"phi\": bhadrons.phi,\n",
    "        \"pdgID\": bhadrons.pdgId,\n",
    "        \"mass\": get_hadron_mass(\n",
    "            bhadrons.pdgId, hadron_pids, hadron_masses, ak.ArrayBuilder()\n",
    "        ).snapshot(),\n",
    "        \"hasBdaughter\": ak.values_astype(\n",
    "            ak.any(is_heavy_hadron(bhadrons.children, 5), axis=-1), int\n",
    "        ),  # B hadrons with B-daughters not removed\n",
    "    }\n",
    ")\n",
    "\n",
    "## find all daughter stable D hadrons and matched with D hadron list\n",
    "bhad_ch1G = bhadrons.children\n",
    "bhad_ch2G = bhadrons.children.children\n",
    "bhad_ch3G = bhadrons.children.children.children\n",
    "bhad_ch4G = bhadrons.children.children.children.children\n",
    "# for each B hadron: form a list of it daugthers which are:\n",
    "# D hadron & D does not include D hadron daugthers\n",
    "# i.e., find B -> D\n",
    "chad_asdau_1G = bhad_ch1G[\n",
    "    is_heavy_hadron(bhad_ch1G, 4) & ~ak.any(is_heavy_hadron(bhad_ch2G, 4), axis=-1)\n",
    "]\n",
    "# then find B -> D* -> D (store the final D)\n",
    "chad_asdau_2G = ak.flatten(\n",
    "    bhad_ch2G[\n",
    "        is_heavy_hadron(bhad_ch2G, 4) & ~ak.any(is_heavy_hadron(bhad_ch3G, 4), axis=-1)\n",
    "    ],\n",
    "    axis=3,\n",
    ")\n",
    "# then find B -> D** -> D* -> D (store the final D)\n",
    "chad_asdau_3G = ak.flatten(\n",
    "    ak.flatten(\n",
    "        bhad_ch3G[\n",
    "            is_heavy_hadron(bhad_ch3G, 4)\n",
    "            & ~ak.any(is_heavy_hadron(bhad_ch4G, 4), axis=-1)\n",
    "        ],\n",
    "        axis=4,\n",
    "    ),\n",
    "    axis=3,\n",
    ")\n",
    "# combine all final D hadrons together\n",
    "# now in dim: (evt, bhad, chad_asdau)\n",
    "chad_asdau_final = ak.concatenate([chad_asdau_1G, chad_asdau_2G, chad_asdau_3G], axis=2)\n",
    "\n",
    "chadrons_broadcast = ak.cartesian([bhadrons, chadrons], axis=1, nested=True)[\n",
    "    \"1\"\n",
    "]  # broadcast to include bhad dimension: (evt, bhad, chad)\n",
    "\n",
    "# pair the D hadrons (as daughters) with the previous D hadron list\n",
    "pair = ak.cartesian(\n",
    "    [chad_asdau_final, chadrons_broadcast], axis=2, nested=True\n",
    ")  # dim: (evt, bhad, chad_asdau, chad)\n",
    "chad_index = ak.argmax(\n",
    "    pair[\"0\"].pt == pair[\"1\"].pt, axis=-1\n",
    ")  # dim: (evt, bhad, chad_asdau)\n",
    "\n",
    "# fill D hadron index (exclude those with hasBdaughter==False)\n",
    "BHadron[\"DHadron1\"] = ak.fill_none(\n",
    "    ak.where(\n",
    "        (ak.num(chad_index, axis=-1) > 0) & ~BHadron.hasBdaughter,\n",
    "        ak.pad_none(chad_index, 2, axis=2)[:, :, 0],\n",
    "        ak.zeros_like(BHadron.pT, dtype=int) - 1,\n",
    "    ),\n",
    "    -1,\n",
    ")\n",
    "BHadron[\"DHadron2\"] = ak.fill_none(\n",
    "    ak.where(\n",
    "        (ak.num(chad_index, axis=-1) > 1) & ~BHadron.hasBdaughter,\n",
    "        ak.pad_none(chad_index, 2, axis=2)[:, :, 1],\n",
    "        ak.zeros_like(BHadron.pT, dtype=int) - 1,\n",
    "    ),\n",
    "    -1,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Genlep & V0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Genlep\n",
    "\n",
    "_fix = lambda x: ak.fill_none(x, 0)\n",
    "is_lep = (\n",
    "    lambda p: (abs(_fix(p.pdgId)) == 11)\n",
    "    | (abs(_fix(p.pdgId)) == 13)\n",
    "    | (abs(_fix(p.pdgId)) == 15)\n",
    ")\n",
    "is_WZ = lambda p: (abs(_fix(p.pdgId)) == 23) | (abs(_fix(p.pdgId)) == 24)\n",
    "is_heavy_hadron = lambda p, pid: (abs(_fix(p.pdgId)) // 100 == pid) | (\n",
    "    abs(_fix(p.pdgId)) // 1000 == pid\n",
    ")\n",
    "\n",
    "sel = (\n",
    "    is_lep(events.GenPart)\n",
    "    & ((events.GenPart.statusFlags & (1 << 13)) > 0)\n",
    "    & (events.GenPart.pt > 3.0)\n",
    ")  # requires pT > 3 GeV\n",
    "genlep = events.GenPart[sel]\n",
    "\n",
    "# trace parents up to 4 generations (from BTA code)\n",
    "genlep_pa1G = genlep.parent\n",
    "genlep_pa2G = genlep.parent.parent\n",
    "genlep_pa3G = genlep.parent.parent.parent\n",
    "genlep_pa4G = genlep.parent.parent.parent.parent\n",
    "istau = abs(genlep_pa1G.pdgId) == 15\n",
    "isWZ = is_WZ(genlep_pa1G) | is_WZ(genlep_pa2G)\n",
    "isD = is_heavy_hadron(genlep_pa1G, 4) | is_heavy_hadron(genlep_pa2G, 4)\n",
    "isB = (\n",
    "    is_heavy_hadron(genlep_pa1G, 5)\n",
    "    | is_heavy_hadron(genlep_pa2G, 5)\n",
    "    | is_heavy_hadron(genlep_pa3G, 5)\n",
    "    | is_heavy_hadron(genlep_pa4G, 5)\n",
    ")\n",
    "\n",
    "Genlep = ak.zip(\n",
    "    {\n",
    "        \"pT\": genlep.pt,\n",
    "        \"eta\": genlep.eta,\n",
    "        \"phi\": genlep.phi,\n",
    "        \"pdgID\": genlep.pdgId,\n",
    "        \"status\": genlep.status,\n",
    "        \"mother\": ak.fill_none(\n",
    "            ak.where(isB | isD, 5 * isB + 4 * isD, 10 * istau + 100 * isWZ), 0\n",
    "        ),\n",
    "    }\n",
    ")\n",
    "\n",
    "# V0\n",
    "\n",
    "is_V0 = lambda p: (abs(p.pdgId) == 310) | (abs(p.pdgId) == 3122)\n",
    "sel = is_V0(events.GenPart) & ((events.GenPart.statusFlags & (1 << 13)) > 0)\n",
    "genV0 = events.GenPart[sel]\n",
    "\n",
    "# finding charged daughters with pT > 1\n",
    "is_lep = (\n",
    "    lambda p: (abs(_fix(p.pdgId)) == 11)\n",
    "    | (abs(_fix(p.pdgId)) == 13)\n",
    "    | (abs(_fix(p.pdgId)) == 15)\n",
    ")\n",
    "is_charged_light_hadron = (\n",
    "    lambda p: (abs(p.pdgId) == 211)\n",
    "    | (abs(p.pdgId) == 213)\n",
    "    | (abs(p.pdgId) == 321)\n",
    "    | (abs(p.pdgId) == 323)\n",
    ")\n",
    "genV0_ch_charged = genV0.children[\n",
    "    (is_charged_light_hadron(genV0.children) | is_lep(genV0.children))\n",
    "    & (genV0.children.pt > 1.0)\n",
    "]\n",
    "ncharged = ak.num(genV0_ch_charged, axis=2)\n",
    "\n",
    "# check acceptance (from the BTA code)\n",
    "svx = ak.fill_none(ak.firsts(genV0_ch_charged.vx, axis=-1), 0)\n",
    "svy = ak.fill_none(ak.firsts(genV0_ch_charged.vy, axis=-1), 0)\n",
    "svz = ak.fill_none(ak.firsts(genV0_ch_charged.vz, axis=-1), 0)\n",
    "radius = np.hypot(svx, svy)\n",
    "abseta = abs(genV0.eta)\n",
    "\n",
    "in_acceptance = (ncharged > 0) & (\n",
    "    ((abseta < 2.0) & (radius < 7.3))\n",
    "    | ((abseta < 2.5) & (radius < 4.4))\n",
    "    | ((abseta > 1.8) & (abseta < 2.5) & (abs(svz) < 34.5))\n",
    ")\n",
    "\n",
    "genV0_inaccept = genV0[in_acceptance]\n",
    "GenV0 = ak.zip(\n",
    "    {\n",
    "        \"pT\": genV0_inaccept.pt,\n",
    "        \"eta\": genV0_inaccept.eta,\n",
    "        \"phi\": genV0_inaccept.phi,\n",
    "        \"pdgId\": genV0_inaccept.pdgId,\n",
    "        \"nCharged\": ncharged[in_acceptance],\n",
    "        \"SVx\": svx[in_acceptance],\n",
    "        \"SVy\": svy[in_acceptance],\n",
    "        \"SVz\": svz[in_acceptance],\n",
    "    }\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Jets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Jets\n",
    "\n",
    "jet = events.Jet[\n",
    "    (events.Jet.pt > 20.0) & (abs(events.Jet.eta) < 2.5)\n",
    "]  # basic selection\n",
    "zeros = ak.zeros_like(jet.pt, dtype=int)\n",
    "Jet = ak.zip(\n",
    "    {\n",
    "        # basic kinematics\n",
    "        \"pt\": jet.pt,\n",
    "        \"eta\": jet.eta,\n",
    "        \"phi\": jet.phi,\n",
    "        \"mass\": jet.mass,\n",
    "        \"uncorrpt\": (1 - jet.rawFactor) * jet.pt,\n",
    "        # jet ID/pileup ID // !!!\n",
    "        \"looseID\": ak.values_astype(\n",
    "            zeros + 1, int\n",
    "        ),  # jet.jetId & (1 << 0) > 0, # has some problem..? always 0\n",
    "        \"tightID\": ak.values_astype(jet.jetId & (1 << 1) > 0, int),\n",
    "        # 'tightlepvetoID': jet.jetId & (1 << 2) > 0,\n",
    "        # pileup ID (essentially userInt('puId106XUL18Id') for UL18)\n",
    "        \"pileup_tightID\": ak.values_astype(\n",
    "            (jet.puId & (1 << 0) > 0) | (jet.pt > 50.0), int\n",
    "        ),\n",
    "        \"pileup_mediumID\": ak.values_astype(\n",
    "            (jet.puId & (1 << 1) > 0) | (jet.pt > 50.0), int\n",
    "        ),\n",
    "        \"pileup_looseID\": ak.values_astype(\n",
    "            (jet.puId & (1 << 2) > 0) | (jet.pt > 50.0), int\n",
    "        ),\n",
    "        # taggers/vars // JP/JBP to be calibrated.. !!!\n",
    "        \"area\": jet.area,\n",
    "        \"Proba\": jet.Proba,\n",
    "        \"ProbaN\": jet.ProbaN,\n",
    "        \"Bprob\": jet.Bprob,\n",
    "        \"BprobN\": jet.BprobN,\n",
    "        # taggers\n",
    "        \"DeepFlavourBDisc\": jet.btagDeepFlavB,\n",
    "        \"DeepFlavourCvsLDisc\": jet.btagDeepFlavCvL,\n",
    "        \"DeepFlavourCvsBDisc\": jet.btagDeepFlavCvB,\n",
    "        \"DeepFlavourBDiscN\": jet.btagNegativeDeepFlavB,\n",
    "        \"DeepFlavourCvsLDiscN\": jet.btagNegativeDeepFlavCvL,\n",
    "        \"DeepFlavourCvsBDiscN\": jet.btagNegativeDeepFlavCvB,\n",
    "    }\n",
    ")\n",
    "\n",
    "if not isRealData:\n",
    "    Jet[\"nbHadrons\"] = jet.nBHadrons\n",
    "    Jet[\"ncHadrons\"] = jet.nCHadrons\n",
    "    Jet[\"partonFlavour\"] = jet.partonFlavour\n",
    "    Jet[\"hadronFlavour\"] = jet.hadronFlavour\n",
    "    # corrected flavours defined in BTA\n",
    "    Jet[\"flavour\"] = ak.where(\n",
    "        jet.hadronFlavour != 0,\n",
    "        jet.hadronFlavour,\n",
    "        ak.where(\n",
    "            (abs(jet.partonFlavour) == 4) | (abs(jet.partonFlavour) == 5),\n",
    "            zeros,\n",
    "            jet.partonFlavour,\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    # genJet pT\n",
    "    genJetIdx = ak.where(\n",
    "        jet.genJetIdx < ak.num(events.GenJet), jet.genJetIdx, zeros - 1\n",
    "    )  # in case the genJet index out of range\n",
    "    Jet[\"genpt\"] = ak.where(genJetIdx >= 0, events.GenJet[genJetIdx].pt, zeros - 1)\n",
    "\n",
    "    # gen-level jet cleaning aginst prompt leptons\n",
    "    genlep_prompt = genlep[(Genlep.mother != 0) & (Genlep.mother % 10 == 0)]\n",
    "    overlapped = ak.zeros_like(jet.pt, dtype=bool)\n",
    "\n",
    "    for lep_pid, dr_thres in [(11, 0.2), (13, 0.2), (15, 0.3)]:\n",
    "        lep = genlep_prompt[abs(genlep_prompt.pdgId) == lep_pid]\n",
    "        pair = ak.cartesian([jet, lep], axis=1, nested=True)  # dim: (event, jet, lep)\n",
    "        dr_min = ak.min(pair[\"0\"].delta_r(pair[\"1\"]), axis=-1)  # dim: (event, jet)\n",
    "        dr_min = ak.fill_none(dr_min, 99.0)  # for 0 lepton case, assign dr=99\n",
    "        overlapped = overlapped | (dr_min < dr_thres)\n",
    "    Jet[\"flavourCleaned\"] = ak.where(\n",
    "        overlapped, zeros - 999, Jet[\"flavour\"]\n",
    "    )  # fill -999 for jets with lepton overlapping\n",
    "\n",
    "    # jet cleaning aginst pileup\n",
    "    Jet[\"flavourCleaned\"] = ak.where(Jet[\"genpt\"] < 8.0, zeros, Jet[\"flavourCleaned\"])\n",
    "\n",
    "\n",
    "# CSV inputs per jets: retreived from the DeepCSV input given they are the same\n",
    "TagVarCSV = ak.zip(\n",
    "    {\n",
    "        \"trackJetPt\": jet.DeepCSV_trackJetPt,\n",
    "        \"vertexCategory\": jet.DeepCSV_vertexCategory,\n",
    "        \"jetNSecondaryVertices\": jet.DeepCSV_jetNSecondaryVertices,\n",
    "        \"trackSumJetEtRatio\": jet.DeepCSV_trackSumJetEtRatio,\n",
    "        \"trackSumJetDeltaR\": jet.DeepCSV_trackSumJetDeltaR,\n",
    "        \"vertexMass\": jet.DeepCSV_vertexMass,\n",
    "        \"vertexNTracks\": jet.DeepCSV_vertexNTracks,\n",
    "        \"vertexEnergyRatio\": jet.DeepCSV_vertexEnergyRatio,\n",
    "        \"vertexJetDeltaR\": jet.DeepCSV_vertexJetDeltaR,\n",
    "    }\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TrkInc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "@nb.njit\n",
    "def cumsum(array, builder):\n",
    "    for arr in array:\n",
    "        builder.begin_list()\n",
    "        sum = 0\n",
    "        for i in arr:\n",
    "            sum += i\n",
    "            builder.integer(sum)\n",
    "        builder.end_list()\n",
    "    return builder\n",
    "\n",
    "\n",
    "# # match jets with PFCands\n",
    "# trk = events.PFCands[(events.PFCands.trkQuality != 0) & (events.PFCands.pt > 5.)]\n",
    "# pair = ak.cartesian([jet, trk], axis=1, nested=True) # dim: (event, jet, pfcand)\n",
    "# matched = pair['0'].delta_r(pair['1']) < 0.4\n",
    "# trk_jetmatch = pair['1'][matched]\n",
    "\n",
    "## match with JetPFCands\n",
    "# form the track selection from JetPFCands\n",
    "trkj = events.JetPFCands[\n",
    "    (events.JetPFCands.pf.trkQuality != 0) & (events.JetPFCands.pt > 1.0)\n",
    "]\n",
    "\n",
    "# selection for \"TrkInc\" in BTA\n",
    "trkj = trkj[\n",
    "    (trkj.pf.trkHighPurity == 1)\n",
    "    & (trkj.pf.trkAlgo != 9)\n",
    "    & (trkj.pf.trkAlgo != 10)\n",
    "    & (trkj.pt > 5.0)\n",
    "    & (trkj.pf.numberOfHits >= 11)\n",
    "    & (trkj.pf.numberOfPixelHits >= 2)\n",
    "    & (trkj.pf.trkChi2 < 10)\n",
    "    & (trkj.pf.lostOuterHits <= 2)\n",
    "    & (trkj.pf.dz < 1.0)\n",
    "]\n",
    "\n",
    "pair = ak.cartesian([jet, trkj], axis=1, nested=True)  # dim: (event, jet, pfcand)\n",
    "# matched = pair['0'].pt == pair['1'].jet.pt\n",
    "matched = (pair[\"0\"].pt == pair[\"1\"].jet.pt) & (pair[\"0\"].delta_r(pair[\"1\"].pf) < 0.4)\n",
    "trkj_jetbased = pair[\"1\"][matched]  # dim: (event, jet, pfcand)\n",
    "\n",
    "# calculate pTrel\n",
    "vec = ak.zip(\n",
    "    {\n",
    "        \"pt\": trkj_jetbased.pf.trkPt,\n",
    "        \"eta\": trkj_jetbased.pf.trkEta,\n",
    "        \"phi\": trkj_jetbased.pf.trkPhi,\n",
    "        \"mass\": trkj_jetbased.pf.mass,\n",
    "    },\n",
    "    behavior=vector.behavior,\n",
    "    with_name=\"PtEtaPhiMLorentzVector\",\n",
    ")\n",
    "ptperp = vec.dot(jet) / jet.p\n",
    "trkj_jetbased[\"ptrel\"] = np.sqrt(np.maximum(vec.p**2 - ptperp**2, 0))\n",
    "\n",
    "# flatten jet-based track arrays\n",
    "trkj_jetbased_flat = ak.flatten(trkj_jetbased, axis=2)  # dim: (event, pfcand)\n",
    "trkj_jetbased_num = ak.num(trkj_jetbased, axis=2)\n",
    "\n",
    "TrkInc = ak.zip(\n",
    "    {\n",
    "        \"pt\": ak.fill_none(trkj_jetbased_flat.pf.trkPt, 0),\n",
    "        \"eta\": ak.fill_none(trkj_jetbased_flat.pf.trkEta, 0),\n",
    "        \"phi\": ak.fill_none(trkj_jetbased_flat.pf.trkPhi, 0),\n",
    "        \"ptrel\": ak.fill_none(trkj_jetbased_flat.ptrel, 0),\n",
    "    }\n",
    ")\n",
    "\n",
    "# assign the first and last index of matched muons to Jet\n",
    "lastidx = cumsum(trkj_jetbased_num, ak.ArrayBuilder()).snapshot()\n",
    "firstidx = ak.fill_none(\n",
    "    ak.concatenate(\n",
    "        [\n",
    "            ak.singletons(ak.zeros_like(events.run, dtype=int)),\n",
    "            lastidx.mask[ak.num(lastidx) > 0][:, :-1],\n",
    "        ],\n",
    "        axis=1,\n",
    "    ),\n",
    "    0,\n",
    ")\n",
    "Jet[\"nFirstTrkInc\"] = firstidx\n",
    "Jet[\"nLastTrkInc\"] = lastidx"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PFMuon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PFMuon\n",
    "\n",
    "# find muon PFCands inside a jet\n",
    "# based on SoftPFMuonTagInfoProducer:\n",
    "# https://github.com/cms-sw/cmssw/blob/10_6_X/RecoBTag/SoftLepton/plugins/SoftPFMuonTagInfoProducer.cc\n",
    "\n",
    "mutrkj = events.JetPFCands[\n",
    "    (events.JetPFCands.pf.trkQuality != 0)\n",
    "    & (events.JetPFCands.pt > 2.0)\n",
    "    & (abs(events.JetPFCands.pf.pdgId) == 13)\n",
    "]\n",
    "\n",
    "# match muon PFCands to muons\n",
    "pair = ak.cartesian(\n",
    "    [mutrkj, events.Muon], axis=1, nested=True\n",
    ")  # dim: (event, pfcand, mu)\n",
    "matched = (\n",
    "    (pair[\"0\"].pf.pdgId == pair[\"1\"].pdgId)\n",
    "    & (pair[\"0\"].pf.delta_r(pair[\"1\"]) < 0.01)\n",
    "    & ((pair[\"0\"].pt - pair[\"1\"].pt) / pair[\"0\"].pt < 0.1)\n",
    ")\n",
    "mutrkj[\"mu\"] = ak.firsts(\n",
    "    pair[\"1\"][matched], axis=2\n",
    ")  # dim same with mutrkj: (event, matched-pfcand), None if not matching with a muon\n",
    "\n",
    "# match muon PFCands also to jets\n",
    "pair = ak.cartesian([jet, mutrkj], axis=1, nested=True)  # dim: (event, jet, pfcand)\n",
    "matched = pair[\"0\"].pt == pair[\"1\"].jet.pt\n",
    "mutrkj_jetbased = pair[\"1\"][matched]  # dim: (event, jet, matched-pfcand)\n",
    "\n",
    "# further requirement: must match with a muon and require the loose ID\n",
    "mutrkj_jetbased = mutrkj_jetbased[\n",
    "    ak.fill_none(mutrkj_jetbased.mu.looseId, False)\n",
    "]  # dim: (event, jet, matched-muons)\n",
    "mu_jetbased = mutrkj_jetbased.mu\n",
    "\n",
    "# calculate pTrel and other kinematics\n",
    "ptperp = mu_jetbased.dot(jet) / jet.p\n",
    "mu_jetbased[\"ptrel\"] = np.sqrt(mu_jetbased.p**2 - ptperp**2)\n",
    "# ratio and ratioRel seems different from the BTA value..\n",
    "mu_jetbased[\"ratio\"] = mu_jetbased.pt / jet.pt\n",
    "mu_jetbased[\"ratioRel\"] = (mu_jetbased.t * jet.t - mu_jetbased.dot(jet)) / jet.mass2\n",
    "\n",
    "mu_jetbased[\"deltaR\"] = mu_jetbased.delta_r(jet)\n",
    "\n",
    "# quality\n",
    "zeros = ak.zeros_like(mu_jetbased.pt, dtype=int)\n",
    "mu_jetbased[\"GoodQuality\"] = zeros\n",
    "mu_jetbased[\"GoodQuality\"] = ak.where(\n",
    "    mu_jetbased.isGlobal, zeros + 1, mu_jetbased[\"GoodQuality\"]\n",
    ")\n",
    "mu_jetbased[\"GoodQuality\"] = ak.where(\n",
    "    (mu_jetbased[\"GoodQuality\"] == 1)\n",
    "    & (mu_jetbased.nStations >= 2)\n",
    "    & (mutrkj_jetbased.pf.numberOfHits >= 11)\n",
    "    & (mutrkj_jetbased.pf.numberOfPixelHits >= 2)\n",
    "    & (mutrkj_jetbased.pf.trkChi2 < 10),\n",
    "    zeros + 2,\n",
    "    mu_jetbased[\"GoodQuality\"],\n",
    ")\n",
    "\n",
    "# flatten jet-based track arrays\n",
    "mu_jetbased_flat = ak.flatten(mu_jetbased, axis=2)  # dim: (event, pfcand)\n",
    "mu_jetbased_jetidx_flat = ak.flatten(\n",
    "    ak.broadcast_arrays(ak.local_index(jet, axis=1), mu_jetbased.pt)[0], axis=2\n",
    ")\n",
    "mu_jetbased_num = ak.num(mu_jetbased, axis=2)\n",
    "\n",
    "PFMuon = ak.zip(\n",
    "    {\n",
    "        \"IdxJet\": ak.fill_none(mu_jetbased_jetidx_flat, 0),\n",
    "        \"pt\": ak.fill_none(mu_jetbased_flat.pt, 0),\n",
    "        \"eta\": ak.fill_none(mu_jetbased_flat.eta, 0),\n",
    "        \"phi\": ak.fill_none(mu_jetbased_flat.phi, 0),\n",
    "        \"ptrel\": ak.fill_none(mu_jetbased_flat.ptrel, 0),\n",
    "        \"ratio\": ak.fill_none(mu_jetbased_flat.ratio, 0),\n",
    "        \"ratioRel\": ak.fill_none(mu_jetbased_flat.ratioRel, 0),\n",
    "        \"deltaR\": ak.fill_none(mu_jetbased_flat.deltaR, 0),\n",
    "        \"IP\": ak.fill_none(\n",
    "            mu_jetbased_flat.ip3d * np.sign(mu_jetbased_flat.dxy + mu_jetbased_flat.dz),\n",
    "            0,\n",
    "        ),\n",
    "        \"IPsig\": ak.fill_none(\n",
    "            mu_jetbased_flat.sip3d\n",
    "            * np.sign(mu_jetbased_flat.dxy + mu_jetbased_flat.dz),\n",
    "            0,\n",
    "        ),\n",
    "        \"IP2D\": ak.fill_none(mu_jetbased_flat.dxy, 0),\n",
    "        \"IP2Dsig\": ak.fill_none(mu_jetbased_flat.dxy / mu_jetbased_flat.dxyErr, 0),\n",
    "        \"GoodQuality\": ak.fill_none(mu_jetbased_flat.GoodQuality, 0),\n",
    "    }\n",
    ")\n",
    "\n",
    "# assign the first and last index of matched muons to Jet\n",
    "lastidx = cumsum(mu_jetbased_num, ak.ArrayBuilder()).snapshot()\n",
    "firstidx = ak.fill_none(\n",
    "    ak.concatenate(\n",
    "        [\n",
    "            ak.singletons(ak.zeros_like(events.run, dtype=int)),\n",
    "            lastidx.mask[ak.num(lastidx) > 0][:, :-1],\n",
    "        ],\n",
    "        axis=1,\n",
    "    ),\n",
    "    0,\n",
    ")\n",
    "Jet[\"nFirstSM\"] = firstidx\n",
    "Jet[\"nLastSM\"] = lastidx"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Write to ROOT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Run [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ... 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]\n",
      "Evt [112068, 112135, 112128, 112124, 112073, ... 114279, 114255, 114237, 114267, 114268]\n",
      "LumiBlock [13, 13, 13, 13, 13, 13, 13, 13, 13, 13, ... 13, 13, 13, 13, 13, 13, 13, 13, 13, 13]\n",
      "nPU [19, 34, 27, 28, 22, 30, 24, 56, 29, 13, ... 38, 25, 34, 24, 46, 12, 37, 42, 35, 37]\n",
      "nPUtrue [19, 39, 34, 27, 13, 30, 23, 48, 27, 19, ... 38, 20, 37, 32, 41, 14, 33, 28, 32, 42]\n",
      "rho [9.18, 24.3, 21.8, 18.4, 20.8, 22.2, 12.9, ... 30.3, 14.3, 29.1, 20.8, 19.5, 25.9]\n",
      "pthat [550, 595, 569, 542, 485, 493, 511, 585, ... 494, 476, 485, 516, 500, 592, 559, 496]\n",
      "BitTrigger [[2144335999, 57311], [4294967295, 65535, ... 65280], [2144336383, 65535]]\n",
      "PV x [0.0107, 0.0105, 0.0112, 0.011, 0.0107, ... 0.00977, 0.0102, 0.0112, 0.0107, 0.0101]\n",
      "PV y [0.0411, 0.0411, 0.0418, 0.0407, 0.0406, ... 0.0408, 0.0418, 0.0422, 0.0424, 0.0415]\n",
      "PV z [-6.12, 2.12, -3.42, 1.9, -3.55, 2.33, ... -0.554, -1.02, 1.76, 0.541, 4.99, 2.54]\n",
      "PV ndof [106, 178, 168, 145, 164, 122, 59.6, 170, ... 132, 202, 186, 202, 162, 80.8, 186]\n",
      "PV chi2 [0.689, 0.768, 0.768, 0.801, 0.697, 0.576, ... 0.816, 0.803, 0.777, 0.705, 0.705]\n",
      "bQuark pT [[], [], [62.2, 374], [], [], [14, 26.4], ... [], [19.9, 89.8], [], [], [], [], []]\n",
      "bQuark eta [[], [], [0.111, -0.0452], [], [], [1.84, 1.7, ... [0.729, 1.14], [], [], [], [], []]\n",
      "bQuark phi [[], [], [0.0801, 0.162], [], [], [-0.584, ... [-2.6, -1.81], [], [], [], [], []]\n",
      "bQuark pdgId [[], [], [-5, 5], [], [], [5, -5], [], ... -5], [], [5, -5], [], [], [], [], []]\n",
      "bQuark status [[], [], [71, 71], [], [], [71, 71], [], ... 71], [], [71, 71], [], [], [], [], []]\n",
      "bQuark fromGSP [[], [], [0, 1], [], [], [0, 0], [], [1, ... [0, 0], [], [0, 0], [], [], [], [], []]\n",
      "cQuark pT [[3.95, 1.55], [0.215, 3.09], [], ... [29, 5.36, 1.96, 1.03, 160, 6.91, 4.73, 0.723]]\n",
      "cQuark eta [[-1.23, -0.523], [2.58, 0.497], ... -5.55, -3.17, -0.723, -4.02, -3.99, -6.45]]\n",
      "cQuark phi [[0.347, -2.09], [0.877, 2.43], ... 1.04, -2.95, 2.22, 0.723, 0.13, -2.05, -0.0862]]\n",
      "cQuark pdgId [[-4, 4], [4, -4], [], [4, 4, -4, -4], ... -4], [], [], [-4, 4, 4, -4, -4, 4, -4, 4]]\n",
      "cQuark status [[71, 71], [71, 71], [], [71, 71, 71, ... [], [], [71, 71, 71, 71, 71, 71, 71, 71]]\n",
      "cQuark fromGSP [[1, 1], [0, 0], [], [0, 0, 0, 0], [0, ... 0, 1], [], [], [0, 0, 0, 0, 0, 0, 0, 0]]\n",
      "BHadron pT [[], [], [52.4, 316], [], [], [13.6, 18.3, ... [16.8, 88.2, 88], [], [], [], [], []]\n",
      "BHadron eta [[], [], [0.102, -0.0432], [], [], [1.88, ... 0.768, 1.14, 1.14], [], [], [], [], []]\n",
      "BHadron phi [[], [], [0.0798, 0.16], [], [], [-0.559, ... -1.81, -1.81], [], [], [], [], []]\n",
      "BHadron pdgID [[], [], [521, -521], [], [], [-513, 523, ... [-521, 523, 521], [], [], [], [], []]\n",
      "BHadron mass [[], [], [5.28, 5.28], [], [], [5.32, ... 5.28, 5.32, 5.28], [], [], [], [], []]\n",
      "BHadron hasBdaughter [[], [], [0, 0], [], [], [1, 1, 0, 0], [], ... 0], [], [0, 1, 0], [], [], [], [], []]\n",
      "BHadron DHadron1 [[], [], [0, 1], [], [], [-1, -1, 0, 1], ... [], [0, -1, 1], [], [], [], [], []]\n",
      "BHadron DHadron2 [[], [], [-1, -1], [], [], [-1, -1, -1, 2], ... [], [-1, -1, 2], [], [], [], [], []]\n",
      "DHadron pT [[4.64, 0.85], [4.59, 0.688], ... 3.28, 87, 5.89, 3.88, 0.317, 60.1, 2.6, 1.12]]\n",
      "DHadron eta [[-1.29, -0.344], [0.405, 2.3], ... -0.727, -4.19, -4, -7.64, 0.783, -5.33, -3.47]]\n",
      "DHadron phi [[0.246, -2.16], [2.62, -0.992], ... 0.73, 0.172, -2.13, -1.76, -2.38, 3.13, 1.71]]\n",
      "DHadron pdgID [[-431, 421], [-431, 421], [-421, ... [421, -421, 411, -421, 421, -411, 421, -421]]\n",
      "DHadron mass [[1.97, 1.86], [1.97, 1.86], ... 1.86, 1.86, 1.87, 1.86, 1.86, 1.87, 1.86, 1.86]]\n",
      "DHadron muInDaughter [[False, False], [False, False], ... True, True, False, False, True, False, False]]\n",
      "Genlep pT [[6.23], [8.16], [27.6, 6.77, 5.17], [210], ... [63.6], [5.88], [11.4], [24.6, 12.2]]\n",
      "Genlep eta [[-1.65], [-0.742], [-0.0524, 0.163, -0.0483, ... -0.456], [-0.251], [-0.695, 0.805]]\n",
      "Genlep phi [[1.93], [-1.25], [0.192, 0.143, 0.194], ... [-0.936], [1.84], [0.729, -2.39]]\n",
      "Genlep pdgID [[-13], [-13], [11, 13, -13], [13], [11, -13, ... [-13], [13], [-13], [13], [13, 13]]\n",
      "Genlep status [[1], [1], [1, 1, 1], [1], [1, 1], [1], ... [2, 1], [1], [1], [1], [1], [1, 1]]\n",
      "Genlep mother [[0], [0], [5, 9, 0], [4], [0, 0], [0], ... [9, 9], [0], [4], [0], [0], [4, 4]]\n",
      "GenV0 pT [[], [], [], [], [], [], [], [], [], ... [2.59], [], [7.08], [1.65], [], [], []]\n",
      "GenV0 eta [[], [], [], [], [], [], [], [], ... [0.0404], [], [-0.0608], [-1.35], [], [], []]\n",
      "GenV0 phi [[], [], [], [], [], [], [], [], [], ... [0.105], [], [0.396], [-1.16], [], [], []]\n",
      "GenV0 pdgId [[], [], [], [], [], [], [], [], [], [], ... [], [310], [], [310], [310], [], [], []]\n",
      "GenV0 nCharged [[], [], [], [], [], [], [], [], [], [], ... [], [1], [], [1], [1], [], [], []]\n",
      "GenV0 SVx [[], [], [], [], [], [], [], [], [], ... [5.45], [], [2.99], [2.35], [], [], []]\n",
      "GenV0 SVy [[], [], [], [], [], [], [], [], [], ... [0.613], [], [1.29], [-5.29], [], [], []]\n",
      "GenV0 SVz [[], [], [], [], [], [], [], [], [], ... [-1.54], [], [-1.22], [-8.77], [], [], []]\n",
      "Jet pt [[383, 306, 131, 63.9, 40.8, 25.4, 24.3], ... [546, 493, 87.4, 36.5], [482, 446]]\n",
      "Jet eta [[0.188, -1.64, -0.109, -1.9, -1.39, -1.89, 1.26, ... 0.296, 1.59], [-0.689, 0.782]]\n",
      "Jet phi [[-1.68, 1.94, -1.17, 1.39, 0.266, 0.272, -0.699, ... 1.67, -1.97], [0.748, -2.38]]\n",
      "Jet mass [[23.8, 21.4, 10.7, 8.95, 5.04, 2.93, 4.81], ... 44.5, 22.1, 8.3], [60.1, 26.7]]\n",
      "Jet uncorrpt [[359, 284, 124, 55.6, 36.4, 21.6, 21.7], ... [507, 463, 85.2, 34.1], [453, 419]]\n",
      "Jet looseID [[1, 1, 1, 1, 1, 1, 1], [1, 1], [1, 1, 1, ... 1, 1, 1, 1, 1], [1, 1, 1, 1], [1, 1]]\n",
      "Jet tightID [[1, 1, 1, 1, 1, 1, 1], [1, 1], [1, 1, 1, ... 1, 1, 1, 1, 1], [1, 1, 1, 1], [1, 1]]\n",
      "Jet pileup_tightID [[1, 1, 1, 1, 1, 1, 1], [1, 1], [1, 1, 1, ... 1, 1, 1, 1, 1], [1, 1, 1, 0], [1, 1]]\n",
      "Jet pileup_mediumID [[1, 1, 1, 1, 1, 1, 1], [1, 1], [1, 1, 1, ... 1, 1, 1, 1, 1], [1, 1, 1, 0], [1, 1]]\n",
      "Jet pileup_looseID [[1, 1, 1, 1, 1, 1, 1], [1, 1], [1, 1, 1, ... 1, 1, 1, 1, 1], [1, 1, 1, 0], [1, 1]]\n",
      "Jet area [[0.479, 0.479, 0.449, 0.419, 0.529, 0.359, 0.529, ... 0.469, 0.459], [0.459, 0.519]]\n",
      "Jet Proba [[0.638, 0.0831, 0.379, 0.041, 0.00845, 0.329, ... 0.00138, -0], [0.121, 0.266]]\n",
      "Jet ProbaN [[0.199, 0.226, 0.0165, 0.0495, -0, 0.0164, ... 0.0471, 0.0372], [0.073, 0.00751]]\n",
      "Jet Bprob [[3.31, 0.656, 1.21, 0.141, 0.0836, 1.23, 0.291, ... 0.0845, 1.1], [1.01, 0.735]]\n",
      "Jet BprobN [[1.14, 1.27, 0.0758, 0.137, -0, 0.0754, 0.421, ... 0.217, 0.171], [1.04, 0.0936]]\n",
      "Jet DeepFlavourBDisc [[0.0436, 0.364, 0.00597, 0.00434, 0.00649, 0.0346, ... 0.355], [0.0639, 0.0349]]\n",
      "Jet DeepFlavourCvsLDisc [[0.111, 0.265, 0.0697, 0.038, 0.0829, 0.135, ... 0.0856, 0.339], [0.136, 0.533]]\n",
      "Jet DeepFlavourCvsBDisc [[0.709, 0.317, 0.92, 0.897, 0.927, 0.791, 0.889, ... 0.791, 0.382], [0.666, 0.937]]\n",
      "Jet DeepFlavourBDiscN [[0.0448, 0.133, 0.00454, 0.0039, 0.00419, 0.0188, ... 0.0428], [0.0124, 0.0048]]\n",
      "Jet DeepFlavourCvsLDiscN [[0.202, 0.239, 0.0541, 0.0337, 0.0527, 0.084, ... 0.0844, 0.0848], [0.0525, 0.0567]]\n",
      "Jet DeepFlavourCvsBDiscN [[0.812, 0.609, 0.922, 0.896, 0.926, 0.814, 0.838, ... 0.799, 0.655], [0.807, 0.922]]\n",
      "Jet nbHadrons [[0, 0, 0, 0, 0, 0, 0], [0, 0], [0, 2, 0, ... 0, 0, 0, 0, 0], [0, 0, 0, 0], [0, 0]]\n",
      "Jet ncHadrons [[0, 0, 0, 0, 1, 0, 0], [0, 0], [0, 2, 0, ... 0, 0, 0, 0, 0], [0, 0, 0, 0], [1, 1]]\n",
      "Jet partonFlavour [[21, 1, 21, 21, -4, 21, 21], [21, 2], ... 21, 21, 21], [1, 1, 21, 21], [-4, -4]]\n",
      "Jet hadronFlavour [[0, 0, 0, 0, 4, 0, 0], [0, 0], [0, 5, 0, ... 0, 0, 0, 0, 0], [0, 0, 0, 0], [4, 4]]\n",
      "Jet flavour [[21, 1, 21, 21, 4, 21, 21], [21, 2], [21, ... 21, 21, 21], [1, 1, 21, 21], [4, 4]]\n",
      "Jet genpt [[324, 306, 121, 62, 58.3, 29.7, 20.1], [602, ... [544, 463, 102, -1], [482, 486]]\n",
      "Jet flavourCleaned [[21, 1, 21, 21, 4, 21, 21], [21, 2], [21, ... 21, 21, 21], [1, 1, 21, 0], [4, 4]]\n",
      "Jet nFirstTrkInc [[0, 10, 13, 16, 19, 21, 22], [0, 9], [0, ... 15, 25, 30, 34], [0, 1, 6, 8], [0, 10]]\n",
      "Jet nLastTrkInc [[10, 13, 16, 19, 21, 22, 23], [9, 16], [6, ... 30, 34, 37], [1, 6, 8, 8], [10, 17]]\n",
      "Jet nFirstSM [[0, 0, 1, 1, 1, 1, 1], [0, 1], [0, 0, 2, ... 0, 0, 1, 1, 1], [0, 0, 1, 1], [0, 1]]\n",
      "Jet nLastSM [[0, 1, 1, 1, 1, 1, 1], [1, 1], [0, 2, 2, ... 0, 1, 1, 1, 1], [0, 1, 1, 1], [1, 2]]\n",
      "TagVarCSV trackJetPt [[155, 38.3, 66.2, 45.5, 18.8, 22.4, 15.6], ... 41.2, 160, 28.2, 1.4], [119, 305]]\n",
      "TagVarCSV vertexCategory [[0, 2, 2, 2, 2, 0, 2], [0, 2], [2, 0, 2, ... 1, 1, 0, 0, 2], [2, 2, 2, 2], [0, 0]]\n",
      "TagVarCSV jetNSecondaryVertices [[1, 0, 0, 0, 0, 1, 0], [2, 0], [0, 1, 0, ... 0, 0, 2, 1, 0], [0, 0, 0, 0], [1, 1]]\n",
      "TagVarCSV trackSumJetEtRatio [[0.572, 0.283, 0.535, 0.818, 0.517, 1.04, 0.717, ... 0.329, 0.0409], [0.522, 0.728]]\n",
      "TagVarCSV trackSumJetDeltaR [[0.00573, 0.0204, 0.0185, 0.0276, 0.024, 0.0133, ... 0.15, 0.279], [0.0109, 0.0065]]\n",
      "TagVarCSV vertexMass [[1.24, -999, -999, -999, -999, 0.948, -999], ... -999, -999, -999], [2.26, 2.6]]\n",
      "TagVarCSV vertexNTracks [[3, 0, 0, 0, 0, 3, 0], [5, 0], [0, 3, 0, ... 3, 3, 6, 2, 0], [0, 0, 0, 0], [3, 2]]\n",
      "TagVarCSV vertexEnergyRatio [[0.293, -999, -999, -999, -999, 0.603, -999], ... -999, -999, -999], [0.301, 0.385]]\n",
      "TagVarCSV vertexJetDeltaR [[0.023, -999, -999, -999, -999, 0.0363, -999, ... -999, -999], [0.0465, 0.0122]]\n",
      "TrkInc pt [[16.5, 12.3, 56.7, 15.1, 13.1, 35.6, 27.5, ... 14.6, 52.7, 31.6, 50.2, 106, 12.2]]\n",
      "TrkInc eta [[0.221, 0.214, 0.192, 0.18, 0.179, 0.178, ... 0.786, 0.785, 0.777, 0.776, 0.804]]\n",
      "TrkInc phi [[-1.69, -1.63, -1.68, -1.7, -1.71, -1.7, ... -2.39, -2.37, -2.37, -2.38, -2.39]]\n",
      "TrkInc ptrel [[0.608, 0.652, 0.497, 0.319, 0.425, 0.78, ... 0.269, 0.471, 0.586, 0.714, 0.295]]\n",
      "PFMuon IdxJet [[1], [0], [1, 1], [0], [0, 0, 0], [], [], ... [0], [3], [], [1], [2], [1], [0, 1]]\n",
      "PFMuon pt [[6.45], [8.27], [6.76, 5.25], [206], ... [63.9], [6.16], [11.4], [24.9, 12.2]]\n",
      "PFMuon eta [[-1.65], [-0.744], [0.162, -0.0539], ... [-0.456], [-0.263], [-0.696, 0.804]]\n",
      "PFMuon phi [[1.96], [-1.25], [0.136, 0.2], [-0.829, ... [-0.925], [1.83], [0.728, -2.39]]\n",
      "PFMuon ptrel [[0.118], [1.33], [1.3, 0.257], [0.177], ... [1.98], [0.451], [0.293], [0.54, 0.293]]\n",
      "PFMuon ratio [[0.021], [0.0133], [0.0157, 0.0122], ... [0.0346], [0.0231], [0.0515, 0.0275]]\n",
      "PFMuon ratioRel [[0.0119], [0.0153], [0.0472, 0.00838], ... [0.0253], [0.0126], [0.0266, 0.0164]]\n",
      "PFMuon deltaR [[0.0182], [0.16], [0.192, 0.0489], ... [0.0743], [0.0257], [0.0217, 0.0237]]\n",
      "PFMuon IP [[-0.134], [0.00216], [-0.289, -0.107], ... [-0.0468], [-0.891], [-0.0433, -0.0221]]\n",
      "PFMuon IPsig [[-21.6], [1.56], [-43.4, -25.2], [2.78], ... 0.458], [-11.1], [-20], [-16.8, -8.49]]\n",
      "PFMuon IP2D [[-0.134], [0.00216], [-0.205, -0.104], ... [-0.0468], [0.216], [0.00221, 0.0076]]\n",
      "PFMuon IP2Dsig [[-21.7], [1.55], [-53.2, -24.8], [2.95], ... [0.011], [-11.1], [18.8], [1.59, 4.4]]\n",
      "PFMuon GoodQuality [[2], [2], [2, 2], [2], [2, 1, 0], [], [], ... [2], [2], [], [2], [2], [1], [2, 2]]\n"
     ]
    }
   ],
   "source": [
    "output = {\n",
    "    **basic_vars,\n",
    "    \"PV\": PV,\n",
    "    \"bQuark\": bQuark,\n",
    "    \"cQuark\": cQuark,\n",
    "    \"BHadron\": BHadron,\n",
    "    \"DHadron\": DHadron,\n",
    "    \"Genlep\": Genlep,\n",
    "    \"GenV0\": GenV0,\n",
    "    \"Jet\": Jet,\n",
    "    \"TagVarCSV\": TagVarCSV,\n",
    "    \"TrkInc\": TrkInc,\n",
    "    \"PFMuon\": PFMuon,\n",
    "}\n",
    "for key in output:\n",
    "    if key in basic_vars:\n",
    "        print(key, output[key])\n",
    "    else:\n",
    "        for b in ak.fields(output[key]):\n",
    "            print(key, b, output[key][b])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "with uproot.recreate(\"btatuple_qcdmu_v3.5.root\") as f:\n",
    "    f[\"btagana/ttree\"] = output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "######################### END ##########################"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pt 2.484375\n",
      "eta -1.56640625\n",
      "phi 2.87451171875\n",
      "jetIdx 0\n",
      "pt 5.5234375\n",
      "eta -1.44140625\n",
      "phi 3.00048828125\n",
      "jetIdx 0\n",
      "pt 2.255859375\n",
      "eta -1.40380859375\n",
      "phi 3.10791015625\n",
      "jetIdx 0\n",
      "pt 1.552734375\n",
      "eta -1.545166015625\n",
      "phi -2.861328125\n",
      "jetIdx 0\n",
      "pt 2.498046875\n",
      "eta -1.631591796875\n",
      "phi -3.08154296875\n",
      "jetIdx 0\n",
      "pt 8.7265625\n",
      "eta -1.634521484375\n",
      "phi 3.13720703125\n",
      "jetIdx 0\n",
      "pt 215.375\n",
      "eta -1.60693359375\n",
      "phi 3.05224609375\n",
      "jetIdx 0\n",
      "pt 0.30224609375\n",
      "eta -1.8349609375\n",
      "phi 3.03515625\n",
      "jetIdx 0\n",
      "pt 0.495849609375\n",
      "eta -1.5908203125\n",
      "phi 3.06982421875\n",
      "jetIdx 0\n",
      "pt 0.235107421875\n",
      "eta -1.56005859375\n",
      "phi 2.95556640625\n",
      "jetIdx 0\n",
      "pt 0.286376953125\n",
      "eta -1.29541015625\n",
      "phi -2.99658203125\n",
      "jetIdx 0\n",
      "pt 4.375\n",
      "eta -1.3486328125\n",
      "phi 3.05859375\n",
      "jetIdx 0\n",
      "pt 1.7099609375\n",
      "eta -1.354736328125\n",
      "phi 2.7451171875\n",
      "jetIdx 0\n",
      "pt 0.2247314453125\n",
      "eta -1.430908203125\n",
      "phi -3.013671875\n",
      "jetIdx 0\n",
      "pt 1.134765625\n",
      "eta -1.43115234375\n",
      "phi 3.07958984375\n",
      "jetIdx 0\n",
      "pt 0.28466796875\n",
      "eta -1.431640625\n",
      "phi 3.0107421875\n",
      "jetIdx 0\n",
      "pt 5.41015625\n",
      "eta -1.5166015625\n",
      "phi 3.05224609375\n",
      "jetIdx 0\n",
      "pt 0.328125\n",
      "eta -1.53515625\n",
      "phi -3.13037109375\n",
      "jetIdx 0\n",
      "pt 22.078125\n",
      "eta -1.572265625\n",
      "phi 3.06005859375\n",
      "jetIdx 0\n",
      "pt 1.7177734375\n",
      "eta -1.58349609375\n",
      "phi -3.09228515625\n",
      "jetIdx 0\n",
      "pt 4.8125\n",
      "eta -1.59716796875\n",
      "phi 3.12939453125\n",
      "jetIdx 0\n",
      "pt 135.375\n",
      "eta -1.601318359375\n",
      "phi 3.0419921875\n",
      "jetIdx 0\n",
      "pt 5.046875\n",
      "eta -1.626708984375\n",
      "phi 3.06494140625\n",
      "jetIdx 0\n",
      "pt 9.34375\n",
      "eta -1.657958984375\n",
      "phi 2.84765625\n",
      "jetIdx 0\n",
      "pt 0.302734375\n",
      "eta -1.694580078125\n",
      "phi 2.9853515625\n",
      "jetIdx 0\n",
      "pt 0.76123046875\n",
      "eta -1.966796875\n",
      "phi 3.126953125\n",
      "jetIdx 0\n",
      "pt 0.578125\n",
      "eta 1.57421875\n",
      "phi 0.3714599609375\n",
      "jetIdx 1\n",
      "pt 0.81005859375\n",
      "eta 1.9853515625\n",
      "phi 0.39788818359375\n",
      "jetIdx 1\n",
      "pt 0.90673828125\n",
      "eta 1.785400390625\n",
      "phi -0.172088623046875\n",
      "jetIdx 1\n",
      "pt 0.88671875\n",
      "eta 1.78857421875\n",
      "phi -0.233795166015625\n",
      "jetIdx 1\n",
      "pt 4.56640625\n",
      "eta 1.698974609375\n",
      "phi 0.094818115234375\n",
      "jetIdx 1\n",
      "pt 1.5126953125\n",
      "eta 1.574951171875\n",
      "phi 0.022846221923828125\n",
      "jetIdx 1\n",
      "pt 1.130859375\n",
      "eta 1.8798828125\n",
      "phi 0.21728515625\n",
      "jetIdx 1\n",
      "pt 1.2255859375\n",
      "eta 1.862060546875\n",
      "phi -0.025577545166015625\n",
      "jetIdx 1\n",
      "pt 3.828125\n",
      "eta 1.839599609375\n",
      "phi 0.03476715087890625\n",
      "jetIdx 1\n",
      "pt 24.15625\n",
      "eta 1.81982421875\n",
      "phi 0.070220947265625\n",
      "jetIdx 1\n",
      "pt 29.0\n",
      "eta 1.8115234375\n",
      "phi 0.0595703125\n",
      "jetIdx 1\n",
      "pt 50.40625"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "eta 1.798095703125\n",
      "phi 0.0646514892578125\n",
      "jetIdx 1\n",
      "pt 3.822265625\n",
      "eta 1.79296875\n",
      "phi 0.1066436767578125\n",
      "jetIdx 1\n",
      "pt 5.85546875\n",
      "eta 1.7919921875\n",
      "phi 0.115234375\n",
      "jetIdx 1\n",
      "pt 15.46875\n",
      "eta 1.76806640625\n",
      "phi 0.1193389892578125\n",
      "jetIdx 1\n",
      "pt 2.732421875\n",
      "eta 1.767822265625\n",
      "phi -0.165435791015625\n",
      "jetIdx 1\n",
      "pt 26.40625\n",
      "eta 1.737060546875\n",
      "phi 0.0631866455078125\n",
      "jetIdx 1\n",
      "pt 13.2421875\n",
      "eta 1.736083984375\n",
      "phi 0.0935516357421875\n",
      "jetIdx 1\n",
      "pt 11.7890625\n",
      "eta 1.73486328125\n",
      "phi 0.0758819580078125\n",
      "jetIdx 1\n",
      "pt 14.8984375\n",
      "eta 1.71435546875\n",
      "phi 0.05381011962890625\n",
      "jetIdx 1\n",
      "pt 2.2734375\n",
      "eta 1.710205078125\n",
      "phi -0.0666961669921875\n",
      "jetIdx 1\n",
      "pt 22.546875\n",
      "eta 1.701904296875\n",
      "phi 0.0777435302734375\n",
      "jetIdx 1\n",
      "pt 1.5146484375\n",
      "eta 1.662353515625\n",
      "phi -0.008991241455078125\n",
      "jetIdx 1\n",
      "pt 2.53515625\n",
      "eta 2.0185546875\n",
      "phi -0.063873291015625\n",
      "jetIdx 1\n",
      "pt 1.736328125\n",
      "eta 1.797119140625\n",
      "phi 0.25860595703125\n",
      "jetIdx 1\n",
      "pt 8.09375\n",
      "eta 1.794677734375\n",
      "phi 0.1025390625\n",
      "jetIdx 1\n",
      "pt 1.03515625\n",
      "eta 1.996337890625\n",
      "phi 0.09454345703125\n",
      "jetIdx 1\n",
      "pt 4.27734375\n",
      "eta 1.8623046875\n",
      "phi -0.0695343017578125\n",
      "jetIdx 1\n",
      "pt 3.81640625\n",
      "eta 1.7412109375\n",
      "phi 0.11700439453125\n",
      "jetIdx 1\n",
      "pt 2.072265625\n",
      "eta 1.786376953125\n",
      "phi -0.008686065673828125\n",
      "jetIdx 1\n",
      "pt 1.892578125\n",
      "eta 1.669677734375\n",
      "phi 0.24267578125\n",
      "jetIdx 1\n",
      "pt 1.251953125\n",
      "eta 1.539794921875\n",
      "phi 0.03858184814453125\n",
      "jetIdx 1\n",
      "pt 7.81640625\n",
      "eta 1.892578125\n",
      "phi 0.04111480712890625\n",
      "jetIdx 1\n",
      "pt 0.39501953125\n",
      "eta 1.98095703125\n",
      "phi -0.015336990356445312\n",
      "jetIdx 1\n",
      "pt 0.4619140625\n",
      "eta 1.634033203125\n",
      "phi -0.16815185546875\n",
      "jetIdx 1\n",
      "pt 0.161865234375\n",
      "eta 1.95458984375\n",
      "phi -0.1081085205078125\n",
      "jetIdx 1\n",
      "pt 0.296875\n",
      "eta 1.88525390625\n",
      "phi -0.1149444580078125\n",
      "jetIdx 1\n",
      "pt 0.35546875\n",
      "eta 2.0390625\n",
      "phi -0.05400848388671875\n",
      "jetIdx 1\n",
      "pt 0.236572265625\n",
      "eta 2.0576171875\n",
      "phi -0.126068115234375\n",
      "jetIdx 1\n",
      "pt 0.85546875\n",
      "eta 1.98681640625\n",
      "phi 0.149627685546875\n",
      "jetIdx 1\n",
      "pt 0.31982421875\n",
      "eta 1.980712890625\n",
      "phi 0.3189697265625\n",
      "jetIdx 1\n",
      "pt 0.70458984375\n",
      "eta 1.899169921875\n",
      "phi -0.228240966796875\n",
      "jetIdx 1\n",
      "pt 5.79296875\n",
      "eta 1.859375\n",
      "phi 0.03369140625\n",
      "jetIdx 1\n",
      "pt 1.72265625\n",
      "eta 1.782470703125\n",
      "phi -0.13360595703125\n",
      "jetIdx 1\n",
      "pt 41.28125\n",
      "eta 1.76708984375\n",
      "phi 0.0690460205078125\n",
      "jetIdx 1\n",
      "pt 4.171875\n",
      "eta 1.748046875\n",
      "phi 0.26190185546875\n",
      "jetIdx 1\n",
      "pt 0.7451171875\n",
      "eta 1.73583984375\n",
      "phi -0.193267822265625\n",
      "jetIdx 1\n",
      "pt 1.732421875\n",
      "eta 1.631591796875\n",
      "phi 0.19549560546875\n",
      "jetIdx 1\n",
      "pt 0.5166015625\n",
      "eta 1.625732421875\n",
      "phi 0.21600341796875\n",
      "jetIdx 1\n",
      "pt 0.53125\n",
      "eta 1.3974609375\n",
      "phi -0.058990478515625\n",
      "jetIdx 1\n",
      "pt 0.751953125\n",
      "eta 1.388427734375\n",
      "phi -0.36328125\n",
      "jetIdx 2\n",
      "pt 0.67431640625\n",
      "eta 1.44775390625\n",
      "phi -0.32745361328125\n",
      "jetIdx 2\n",
      "pt 0.91748046875\n",
      "eta 1.513916015625\n",
      "phi -0.94970703125\n",
      "jetIdx 2\n",
      "pt 3.224609375\n",
      "eta 1.748779296875\n",
      "phi -0.42767333984375\n",
      "jetIdx 2\n",
      "pt 0.434326171875\n",
      "eta 1.818359375\n",
      "phi -0.63037109375\n",
      "jetIdx 2\n",
      "pt 3.712890625\n",
      "eta 1.765869140625\n",
      "phi -0.866943359375\n",
      "jetIdx 2\n",
      "pt 1.451171875\n",
      "eta 1.6552734375\n",
      "phi -0.6014404296875\n",
      "jetIdx 2\n",
      "pt 0.496826171875\n",
      "eta 1.574462890625\n",
      "phi -0.39813232421875\n",
      "jetIdx 2\n",
      "pt 1.2998046875\n",
      "eta 1.5693359375\n",
      "phi -0.5711669921875\n",
      "jetIdx 2\n",
      "pt 3.974609375\n",
      "eta 1.5126953125\n",
      "phi -0.6392822265625\n",
      "jetIdx 2\n",
      "pt 3.23046875\n",
      "eta 1.798095703125\n",
      "phi -0.525634765625\n",
      "jetIdx 2\n",
      "pt 1.4169921875\n",
      "eta 1.72607421875\n",
      "phi -0.7518310546875\n",
      "jetIdx 2\n",
      "pt 1.83984375\n",
      "eta 1.373291015625\n",
      "phi -0.47900390625\n",
      "jetIdx 2\n",
      "pt 0.319091796875\n",
      "eta 1.59912109375\n",
      "phi -0.6373291015625\n",
      "jetIdx 2\n",
      "pt 0.441162109375\n",
      "eta 1.853759765625\n",
      "phi -0.6942138671875\n",
      "jetIdx 2\n",
      "pt 0.43310546875\n",
      "eta 1.324462890625\n",
      "phi -0.630859375\n",
      "jetIdx 2\n",
      "pt 0.67626953125\n",
      "eta 2.0078125\n",
      "phi -0.6761474609375\n",
      "jetIdx 2\n",
      "pt 1.2109375\n",
      "eta 1.92919921875\n",
      "phi -0.7841796875\n",
      "jetIdx 2\n",
      "pt 0.93603515625\n",
      "eta 1.6015625\n",
      "phi -0.7425537109375\n",
      "jetIdx 2\n",
      "pt 2.609375\n",
      "eta 1.546142578125\n",
      "phi -0.7764892578125\n",
      "jetIdx 2\n",
      "pt 3.28125\n",
      "eta 1.5029296875\n",
      "phi -0.6546630859375\n",
      "jetIdx 2\n",
      "pt 1.3984375\n",
      "eta -3.78759765625\n",
      "phi -3.0537109375\n",
      "jetIdx 3\n",
      "pt 1.4951171875\n",
      "eta -3.7880859375\n",
      "phi 3.05517578125\n",
      "jetIdx 3\n",
      "pt 1.63671875\n",
      "eta -3.7880859375\n",
      "phi 2.880859375\n",
      "jetIdx 3\n",
      "pt 0.95263671875\n",
      "eta -3.96240234375\n",
      "phi -2.87890625\n",
      "jetIdx 3\n",
      "pt 4.02734375\n",
      "eta -3.96240234375\n",
      "phi -3.05322265625\n",
      "jetIdx 3\n",
      "pt 8.53125\n",
      "eta -3.962890625\n",
      "phi 3.05517578125\n",
      "jetIdx 3\n",
      "pt 1.8466796875\n",
      "eta -3.962890625\n",
      "phi 2.880859375\n",
      "jetIdx 3\n",
      "pt 1.3662109375\n",
      "eta -3.962890625\n",
      "phi 2.7060546875\n",
      "jetIdx 3\n",
      "pt 0.12457275390625\n",
      "eta -4.1162109375\n",
      "phi -2.87890625\n",
      "jetIdx 3\n",
      "pt 3.32421875\n",
      "eta -4.1376953125\n",
      "phi -3.05322265625\n",
      "jetIdx 3\n",
      "pt 3.357421875\n",
      "eta -4.1376953125\n",
      "phi 3.0556640625\n",
      "jetIdx 3\n",
      "pt 3.189453125\n",
      "eta -4.138671875\n",
      "phi 2.880859375\n",
      "jetIdx 3\n",
      "pt 0.3681640625\n",
      "eta -4.138671875\n",
      "phi 2.70654296875\n",
      "jetIdx 3\n",
      "pt 0.533203125\n",
      "eta -4.3125\n",
      "phi -3.05322265625\n",
      "jetIdx 3\n",
      "pt 0.5615234375\n",
      "eta -4.3134765625\n",
      "phi 3.0556640625\n",
      "jetIdx 3\n",
      "pt 0.8271484375\n",
      "eta -4.3134765625\n",
      "phi 2.88134765625\n",
      "jetIdx 3\n",
      "pt 0.85595703125\n",
      "eta 2.6318359375\n",
      "phi -1.468505859375\n",
      "jetIdx 4\n",
      "pt 0.59765625\n",
      "eta 2.85986328125\n",
      "phi -0.9765625\n",
      "jetIdx 4\n",
      "pt 0.560546875\n",
      "eta 2.77783203125\n",
      "phi -1.182373046875\n",
      "jetIdx 4\n",
      "pt 0.43994140625\n",
      "eta 2.82080078125\n",
      "phi -1.2666015625\n",
      "jetIdx 4\n",
      "pt 4.76171875\n",
      "eta 3.0830078125\n",
      "phi -1.309326171875\n",
      "jetIdx 4\n",
      "pt 2.076171875\n",
      "eta 3.0830078125\n",
      "phi -1.483642578125\n",
      "jetIdx 4\n",
      "pt 0.248779296875\n",
      "eta 3.0615234375\n",
      "phi -1.134765625\n",
      "jetIdx 4\n",
      "pt 0.18701171875\n",
      "eta 3.0615234375\n",
      "phi -0.9603271484375\n",
      "jetIdx 4\n",
      "pt 0.81689453125\n",
      "eta 2.9404296875\n",
      "phi -0.960205078125\n",
      "jetIdx 4\n",
      "pt 2.23046875\n",
      "eta 2.8447265625\n",
      "phi -1.0693359375\n",
      "jetIdx 4\n",
      "pt 8.5859375\n",
      "eta 2.81201171875\n",
      "phi -1.323486328125\n",
      "jetIdx 4\n",
      "pt 0.6591796875\n",
      "eta 2.2216796875\n",
      "phi -0.5850830078125\n",
      "jetIdx 5\n",
      "pt 0.5615234375\n",
      "eta 2.41064453125\n",
      "phi -0.40545654296875\n",
      "jetIdx 5\n",
      "pt 0.95751953125\n",
      "eta 2.36962890625\n",
      "phi -0.17657470703125\n",
      "jetIdx 5\n",
      "pt 2.96484375\n",
      "eta 2.3486328125\n",
      "phi -0.3134765625\n",
      "jetIdx 5\n",
      "pt 1.0751953125\n",
      "eta 2.4208984375\n",
      "phi -0.131744384765625\n",
      "jetIdx 5\n",
      "pt 0.99658203125\n",
      "eta 2.3466796875\n",
      "phi -0.100311279296875\n",
      "jetIdx 5\n",
      "pt 1.4296875\n",
      "eta 2.52587890625\n",
      "phi -0.40087890625\n",
      "jetIdx 5\n",
      "pt 1.2900390625\n",
      "eta 2.5556640625\n",
      "phi -0.46923828125\n",
      "jetIdx 5\n",
      "pt 4.0703125\n",
      "eta 2.50439453125\n",
      "phi -0.1095733642578125\n",
      "jetIdx 5\n",
      "pt 1.1005859375\n",
      "eta 2.40185546875\n",
      "phi -0.6541748046875\n",
      "jetIdx 5\n",
      "pt 1.45703125\n",
      "eta 2.29833984375\n",
      "phi -0.240325927734375\n",
      "jetIdx 5\n",
      "pt 2.8359375\n",
      "eta 2.1552734375\n",
      "phi -0.40447998046875\n",
      "jetIdx 5\n",
      "pt 1.0654296875\n",
      "eta 2.02734375\n",
      "phi -0.3912353515625\n",
      "jetIdx 5\n",
      "pt 0.93212890625\n",
      "eta -3.240234375\n",
      "phi -2.181640625\n",
      "jetIdx 6\n",
      "pt 2.474609375\n",
      "eta -3.26171875\n",
      "phi -2.0068359375\n",
      "jetIdx 6\n",
      "pt 0.58447265625\n",
      "eta -3.41552734375\n",
      "phi -2.181640625\n",
      "jetIdx 6\n",
      "pt 0.428466796875\n",
      "eta -3.41552734375\n",
      "phi -2.35595703125\n",
      "jetIdx 6\n",
      "pt 0.416748046875\n",
      "eta -3.41552734375\n",
      "phi -1.658203125\n",
      "jetIdx 6\n",
      "pt 1.828125\n",
      "eta -3.43701171875\n",
      "phi -1.83251953125\n",
      "jetIdx 6\n",
      "pt 3.5703125\n",
      "eta -3.43701171875\n",
      "phi -2.0068359375\n",
      "jetIdx 6\n",
      "pt 1.9482421875\n",
      "eta -3.6123046875\n",
      "phi -2.35595703125\n",
      "jetIdx 6\n",
      "pt 1.4794921875\n",
      "eta -3.6123046875\n",
      "phi -2.181640625\n",
      "jetIdx 6\n",
      "pt 1.470703125\n",
      "eta -3.6123046875\n",
      "phi -2.0068359375\n",
      "jetIdx 6\n",
      "pt 4.1171875\n",
      "eta -3.6123046875\n",
      "phi -1.83251953125\n",
      "jetIdx 6\n",
      "pt 0.69482421875\n",
      "eta -3.787109375\n",
      "phi -1.83251953125\n",
      "jetIdx 6\n",
      "pt 2.529296875\n",
      "eta -3.787109375\n",
      "phi -2.181640625\n",
      "jetIdx 6\n",
      "pt 1.46875\n",
      "eta -3.787109375\n",
      "phi -2.0068359375\n",
      "jetIdx 6\n"
     ]
    }
   ],
   "source": [
    "lt = events.JetPFCands[5]\n",
    "for i in range(len(lt.pt)):\n",
    "    for b in [\"pt\", \"eta\", \"phi\"]:\n",
    "        print(b, lt.pf[b][i])\n",
    "    print(\"jetIdx\", lt.jetIdx[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 ------\n",
      "pt [8.73, 215]\n",
      "eta [-1.63, -1.61]\n",
      "phi [3.14, 3.05]\n",
      "ptrel [0.846, 2.06]\n",
      "1 ------\n",
      "pt [24.2, 29, 50.4, 5.86, 15.5, 26.4, 13.2, 11.8, 14.9, 22.5, 8.09, 7.82]\n",
      "eta [1.82, 1.81, 1.8, 1.79, 1.77, 1.74, 1.74, 1.74, 1.71, 1.7, 1.79, 1.89]\n",
      "phi [0.0702, 0.0596, 0.0646, 0.115, 0.119, ... 0.0759, 0.0538, 0.0778, 0.103, 0.0411]\n",
      "ptrel [1.01, 0.977, 0.953, 0.291, 0.813, 1.09, 0.652, 0.517, 0.966, 1.69, 0.312, 0.96]\n",
      "2 ------\n",
      "pt []\n",
      "eta []\n",
      "phi []\n",
      "ptrel []\n",
      "3 ------\n",
      "pt []\n",
      "eta []\n",
      "phi []\n",
      "ptrel []\n"
     ]
    }
   ],
   "source": [
    "ie = 5\n",
    "for i, (start, end) in enumerate(zip(Jet.nFirstTrkInc[ie], Jet.nLastTrkInc[ie])):\n",
    "    print(i, \"------\")\n",
    "    for b in ak.fields(TrkInc):\n",
    "        print(b, TrkInc[b][ie][start:end])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for ie in range(20):\n",
    "ie = 12\n",
    "print(ie, \"----------------\")\n",
    "for i, (start, end) in enumerate(zip(Jet.nFirstSM[ie], Jet.nLastSM[ie])):\n",
    "    print(\"jet\", i, \"------\")\n",
    "    for b in ak.fields(PFMuon):\n",
    "        print(b, PFMuon[b][ie][start:end])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dir(mu_jetbased)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = uproot.lazy(\"btatuple_qcdmu_v3.root\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Array [6.76, 5.25] type='2 * float64'>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df[2].PFMuon_pt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for ie in range(len(events)):\n",
    "    print(ie, \"--------------------\")\n",
    "    for b in ak.fields(GenV0):\n",
    "        print(b, GenV0[b][ie])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test with BTA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = uproot.lazy(\"../testfiles/JetTree_mc.root:btagana/ttree\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for ie in range(len(events)):\n",
    "#     if len(bQuark[ie]) != len(df[ie].bQuark_pT) or (len(bQuark[ie]) > 0 and (sum(abs(bQuark[ie].pT - df[ie].bQuark_pT)) > 2 or sum(bQuark[ie].fromGSP != ak.values_astype(df[ie].bQuark_fromGSP, 'bool')) > 0)):\n",
    "#         print(ie)\n",
    "\n",
    "# for ie in range(len(events)):\n",
    "#     if len(cQuark[ie]) != len(df[ie].cQuark_pT) or (len(cQuark[ie]) > 0 and (sum(abs(cQuark[ie].pT - df[ie].cQuark_pT)) > 2 or sum(cQuark[ie].fromGSP != ak.values_astype(df[ie].cQuark_fromGSP, 'bool')) > 0)):\n",
    "#         print(ie)\n",
    "\n",
    "\n",
    "# for ie in range(100):\n",
    "#     if len(DHadron[ie]) != len(df[ie].DHadron_pT) or (len(DHadron[ie]) > 0 and (sum(abs(DHadron[ie].pT - df[ie].DHadron_pT)) > 2)):\n",
    "#         print(ie)\n",
    "\n",
    "# for ie in range(200):\n",
    "#     if len(Genlep[ie]) != len(df[ie].Genlep_pT) or (len(Genlep[ie]) > 0 and (sum(abs(Genlep[ie].pT - df[ie].Genlep_pT)) > 2) and (sum(abs(Genlep[ie].mother - df[ie].Genlep_mother)) > 0)):\n",
    "#         print(ie, len(Genlep[ie]), len(df[ie].Genlep_pT), sum(abs(Genlep[ie].pT - df[ie].Genlep_pT)) > 2, sum(abs(Genlep[ie].mother - df[ie].Genlep_mother)) > 0)\n",
    "\n",
    "# for ie in range(100):\n",
    "#     if len(GenV0[ie]) != len(df[ie].GenV0_pT) or (len(GenV0[ie]) > 0 and (sum(abs(GenV0[ie].pT - df[ie].GenV0_pT)) > 2) and (sum(abs(GenV0[ie].SVx - df[ie].GenV0_SVx)) > 0)):\n",
    "#         print(ie)\n",
    "\n",
    "\n",
    "# # GenV0\n",
    "# for ie in range(200):\n",
    "#     if len(Genlep[ie]) != len(df[ie].Genlep_pT):\n",
    "#         print(ie, 'length doesnt match', len(Genlep[ie]), len(df[ie].Genlep_pT))\n",
    "#     else:\n",
    "#         for b in ak.fields(Genlep):\n",
    "#             if len(Genlep[ie]) > 0 and any(abs(Genlep[ie][b] - df[ie]['Genlep_'+b]) / (df[ie]['Genlep_'+b]) > 1e-2):\n",
    "#                 print(ie, b, 'doesnt match', Genlep[ie][b], df[ie]['Genlep_'+b])\n",
    "\n",
    "# Jets\n",
    "for ie in range(100):\n",
    "    if len(Jet[ie]) != len(df[ie].Jet_pt):\n",
    "        print(ie, \"length doesnt match\", len(Jet[ie]), len(df[ie].Jet_pt))\n",
    "    else:\n",
    "        for b in ak.fields(Jet):\n",
    "            if len(Jet[ie]) > 0 and any(\n",
    "                abs(Jet[ie][b] - df[ie][\"Jet_\" + b]) / (df[ie][\"Jet_\" + b] + 1e-3)\n",
    "                > 1e-3\n",
    "            ):\n",
    "                if \"Prob\" not in b and \"Bprob\" not in b:\n",
    "                    print(ie, b, \"doesnt match\", Jet[ie][b], df[ie][\"Jet_\" + b])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for ie in range(len(events)):\n",
    "    for var in basic_vars:\n",
    "        if abs(basic_vars[var][ie] - df[ie][var]) / (df[ie][var] + 1e-3) > 1e-3:\n",
    "            print(ie, var, \"doesnt match\", basic_vars[var][ie], df[ie][var])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 ptrel doesnt match [0.608, 0.652, 0.497, 0.319, 0.425, 0.78, ... 0.203, 0.226, 0.419, 0.372, 0.385] [0.609, 0.655, 0.511, 0.324, 0.421, 0.784, ... 0.201, 0.226, 0.419, 0.371, 0.388]\n",
      "1 ptrel doesnt match [1.67, 10.6, 0.642, 0.61, 1.23, 1.96, 0.289, ... 3.9, 0.83, 0.827, 0.591, 0.56, 1.73] [1.66, 10.6, 0.642, 0.606, 1.24, 1.96, ... 3.95, 0.819, 0.823, 0.589, 0.559, 1.73]\n",
      "2 length doesnt match 18 17\n",
      "3 ptrel doesnt match [0.944, 0.095, 0.041, 0.188, 0.288, 0.611, ... 0.497, 0.136, 1.51, 0.359, 0.544] [0.952, 0.0975, 0.0396, 0.211, 0.275, 0.606, ... 0.498, 0.129, 1.51, 0.358, 0.544]\n",
      "4 ptrel doesnt match [0.864, 1.64, 0.322, 0.259, 0.296, 0.678, ... 0.554, 0.682, 3.38, 0.0898, 0.182] [0.863, 1.64, 0.323, 0.255, 0.295, 0.676, ... 0.551, 0.678, 3.32, 0.0863, 0.183]\n",
      "5 length doesnt match 14 17\n",
      "6 length doesnt match 14 18\n",
      "7 ptrel doesnt match [0.618, 1.49, 0.467, 1.56, 0.43, 3.48, 0.843, ... 3.39, 1.68, 1.56, 1.4, 1.25, 1.51] [0.615, 1.47, 0.465, 1.55, 0.427, 3.46, 0.84, ... 3.38, 1.68, 1.56, 1.4, 1.25, 1.51]\n",
      "8 ptrel doesnt match [1.42, 0.82, 0.662, 0.903, 4.68, 3.15, 2.98, 0.888, 0.974, 2.15, 6.84, 3.55] [1.48, 0.828, 0.663, 0.9, 4.68, 3.14, 2.99, 0.88, 0.962, 2.15, 6.85, 3.55]\n",
      "9 ptrel doesnt match [0.887, 9.81, 1.56, 4.23, 0.186, 0.751, ... 0.493, 0.843, 0.862, 0.889, 1.25, 0.339] [0.888, 9.78, 1.56, 4.23, 0.187, 0.746, ... 0.507, 0.841, 0.868, 0.885, 1.25, 0.338]\n",
      "10 length doesnt match 26 27\n",
      "11 ptrel doesnt match [0.393, 0.0241, 0.689, 0.558, 1.01, 1.44, ... 1.57, 0.382, 0.957, 1.4, 0.933, 1.24] [0.389, 0.0227, 0.687, 0.557, 1.01, 1.43, ... 1.57, 0.39, 0.951, 1.4, 0.933, 1.24]\n",
      "12 length doesnt match 29 30\n",
      "13 length doesnt match 18 17\n",
      "14 ptrel doesnt match [2.52, 0.644, 1.08, 1.62, 3.44, 3.28, ... 0.49, 0.479, 0.0668, 0.541, 0.21, 0.531] [2.5, 0.643, 1.07, 1.62, 3.42, 3.28, ... 0.488, 0.478, 0.0644, 0.542, 0.211, 0.529]\n",
      "15 pt doesnt match [5.82, 61.9, 15.6, 29.3, 22.2, 10.5, 6.39, ... 70.4, 17.7, 5.3, 8.77, 8.02, 6.61] [5.82, 61.9, 15.6, 29.3, 22.2, 10.5, 6.39, ... 8.04, 70.4, 17.7, 5.3, 8.77, 8.02]\n",
      "15 eta doesnt match [1.47, 1.46, 1.44, 1.38, 1.34, 1.45, ... -1.07, -1.1, 0.935, 0.885, 0.869, 0.914] [1.47, 1.46, 1.44, 1.38, 1.34, 1.45, 2.2, ... 2.17, -1.07, -1.1, 0.935, 0.885, 0.869]\n",
      "15 phi doesnt match [1.61, 1.45, 1.46, 1.6, 1.54, 1.45, ... -2.01, -2.04, 0.182, 0.253, 0.173, 0.229] [1.61, 1.45, 1.46, 1.6, 1.54, 1.45, ... -2.07, -2.01, -2.04, 0.182, 0.253, 0.173]\n",
      "15 ptrel doesnt match [0.747, 2.05, 0.458, 3.66, 2.64, 0.358, ... 0.588, 0.868, 0.27, 0.411, 0.327, 0.209] [0.748, 2.03, 0.454, 3.65, 2.65, 0.353, ... 0.385, 0.628, 0.868, 0.27, 0.41, 0.326]\n",
      "16 ptrel doesnt match [0.562, 0.443, 0.515, 0.887, 0.616, 1.96, ... 0.35, 0.399, 0.0808, 1.62, 0.345] [0.562, 0.444, 0.518, 0.886, 0.616, 1.94, ... 0.347, 0.399, 0.0835, 1.62, 0.356]\n",
      "17 length doesnt match 28 29\n",
      "18 ptrel doesnt match [0.318, 1.45, 1.98, 0.378, 0.311, 0.449, ... 0.161, 0.463, 0.72, 0.31, 0.636, 0.155] [0.317, 1.46, 1.96, 0.382, 0.318, 0.447, ... 0.459, 0.724, 0.307, 0.639, 0.154]\n",
      "19 ptrel doesnt match [1.58, 2.18, 1.08, 3.92, 0.115, 0.499, ... 0.203, 1.29, 0.217, 0.724, 1.03, 1.23] [1.58, 2.18, 1.08, 3.91, 0.116, 0.5, ... 0.208, 1.29, 0.207, 0.726, 1.04, 1.23]\n",
      "20 ptrel doesnt match [0.493, 1.01, 2.34, 0.606, 0.61, 0.17, 0.146, ... 1.17, 1.52, 2.79, 4.73, 2.33, 1.46] [0.497, 0.997, 2.31, 0.597, 0.611, 0.175, ... 1.18, 1.52, 2.79, 4.74, 2.32, 1.47]\n",
      "21 ptrel doesnt match [0.808, 0.549, 0.806, 0.204, 0.343, 1.03, 0.908, 0.593, 0.656] [0.808, 0.547, 0.809, 0.205, 0.342, 1.04, 0.934, 0.595, 0.653]\n",
      "22 length doesnt match 28 27\n",
      "23 length doesnt match 28 29\n",
      "24 ptrel doesnt match [0.381, 0.834, 0.48, 0.455, 0.276, 0.562, ... 0.834, 2.33, 1.28, 1.7, 1.05, 0.418] [0.377, 0.831, 0.482, 0.452, 0.284, 0.565, ... 0.833, 2.32, 1.28, 1.67, 1.05, 0.423]\n",
      "25 ptrel doesnt match [0.521, 0.73, 0.337, 0.187, 0.291, 0.273, ... 0.355, 0.408, 0.348, 0.0764, 0.379] [0.52, 0.726, 0.333, 0.186, 0.296, 0.284, ... 0.353, 0.408, 0.356, 0.0765, 0.366]\n",
      "26 ptrel doesnt match [0.219, 0.762, 0.686, 1.75, 0.876, 1.19, ... 0.376, 1.65, 0.61, 2.32, 2.39, 1.04] [0.219, 0.756, 0.682, 1.75, 0.876, 1.19, ... 0.381, 1.64, 0.604, 2.32, 2.4, 1.04]\n",
      "27 ptrel doesnt match [0.699, 0.423, 0.355, 0.565, 0.412, 1.32, ... 0.398, 0.657, 0.945, 0.529, 0.842] [0.695, 0.438, 0.366, 0.571, 0.454, 1.31, ... 0.394, 0.654, 0.958, 0.529, 0.828]\n",
      "28 ptrel doesnt match [5.41, 2.71, 0.912, 3.82, 2.04, 0.617, ... 0.673, 1.06, 0.752, 0.906, 0.1, 1.29] [5.41, 2.71, 0.916, 3.82, 2.04, 0.618, ... 0.676, 1.06, 0.754, 0.909, 0.0977, 1.29]\n",
      "29 ptrel doesnt match [0.168, 0.316, 0.299, 2.04, 5.57, 2.11, ... 0.0777, 0.606, 2.31, 3.53, 2.98, 1.53] [0.176, 0.314, 0.299, 2.04, 5.56, 2.11, ... 0.0778, 0.606, 2.31, 3.53, 2.98, 1.53]\n",
      "30 length doesnt match 26 25\n",
      "31 length doesnt match 17 16\n",
      "32 ptrel doesnt match [0.179, 0.155, 0.718, 0.699, 0.198, 0.464, 0.545, 0.622] [0.183, 0.129, 0.772, 0.687, 0.224, 0.46, 0.541, 0.63]\n",
      "33 pt doesnt match [26.3, 19.7, 36.5, 39.1, 19.8, 6.54, 37.4, ... 29.7, 6.82, 6.13, 7.2, 5.2, 10.6, 8.4] [26.3, 19.7, 36.5, 39.1, 19.8, 6.54, 37.4, ... 29.7, 6.82, 6.13, 7.2, 5.2, 10.6, 8.4]\n",
      "33 eta doesnt match [0.177, 0.167, 0.167, 0.146, 0.146, 0.142, ... 0.891, 0.859, 0.373, 0.373, 0.303] [0.177, 0.167, 0.167, 0.146, 0.146, 0.142, ... 0.891, 0.859, 0.373, 0.372, 0.303]\n",
      "33 phi doesnt match [-2.63, -2.64, -2.62, -2.62, -2.63, -2.67, ... 0.0765, 0.0953, -1.54, -1.57, -1.95] [-2.64, -2.65, -2.62, -2.62, -2.63, -2.67, ... 0.0766, 0.0953, -1.54, -1.57, -1.95]\n",
      "33 ptrel doesnt match [1.63, 1.02, 2.06, 1.66, 0.665, 0.26, ... 0.29, 0.397, 0.672, 0.369, 0.425, 2.86] [1.63, 1.02, 2.05, 1.63, 0.659, 0.259, ... 0.29, 0.397, 0.674, 0.367, 0.424, 2.86]\n",
      "34 ptrel doesnt match [0.616, 0.745, 0.279, 0.397, 1.13, 1.54, ... 1.44, 1.15, 0.453, 0.308, 0.724, 0.515] [0.616, 0.714, 0.28, 0.397, 1.15, 1.54, ... 1.45, 1.16, 0.446, 0.306, 0.726, 0.516]\n",
      "35 length doesnt match 6 27\n",
      "36 ptrel doesnt match [1.38, 2.93, 1.42, 0.584, 15.3, 1.31, ... 0.84, 0.961, 0.302, 0.433, 0.521, 0.881] [1.38, 2.93, 1.41, 0.582, 15.3, 1.31, ... 0.841, 0.962, 0.304, 0.44, 0.523, 0.875]\n",
      "37 length doesnt match 20 21\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "38 ptrel doesnt match [0.698, 0.227, 0.416, 0.838, 0.751, 0.505, 0.874, 0.642, 0.591] [0.671, 0.23, 0.416, 0.837, 0.755, 0.522, 0.88, 0.642, 0.592]\n",
      "39 ptrel doesnt match [1, 1.33, 1.62, 0.768, 3.01, 0.934, 5.16, ... 0.732, 1.07, 6.31, 4.05, 2.32, 0.524] [1, 1.33, 1.62, 0.768, 3.01, 0.934, 5.14, ... 0.732, 1.07, 6.31, 4.05, 2.32, 0.523]\n",
      "40 length doesnt match 19 20\n",
      "41 ptrel doesnt match [1.21, 0.555, 0.73, 0.985, 0.913, 0.591, ... 0.265, 0.855, 0.427, 0.6, 0.581, 0.198] [1.21, 0.555, 0.731, 0.983, 0.912, 0.59, ... 0.864, 0.425, 0.599, 0.582, 0.203]\n",
      "42 length doesnt match 15 16\n",
      "43 ptrel doesnt match [0.476, 0.329, 0.565, 0.482, 0.689, 0.46, ... 0.384, 0.663, 0.189, 0.936, 0.618] [0.479, 0.323, 0.556, 0.466, 0.7, 0.447, ... 0.383, 0.666, 0.187, 0.935, 0.625]\n",
      "44 length doesnt match 16 19\n",
      "45 pt doesnt match [44, 26.9, 22.5, 6.08, 10.7, 8.88, 10.3, ... 16.8, 6.17, 5.36, 94.4, 7.29, 10, 5.05] [44, 26.9, 22.5, 6.08, 10.7, 8.88, 10.3, ... 70.5, 16.8, 6.17, 94.4, 7.29, 10, 5.05]\n",
      "45 eta doesnt match [0.839, 0.831, 0.829, 0.827, 0.821, 0.812, ... 1.93, 1.85, 1.91, 1.9, 1.96, 0.722] [0.839, 0.831, 0.829, 0.827, 0.821, 0.813, ... 1.95, 1.93, 1.91, 1.9, 1.96, 0.722]\n",
      "45 phi doesnt match [1.3, 1.35, 1.32, 1.35, 1.37, 1.32, ... -1.65, -2.06, -1.65, -1.63, -1.65, 2.13] [1.3, 1.35, 1.32, 1.35, 1.37, 1.32, ... -2.04, -1.65, -1.65, -1.63, -1.65, 2.13]\n",
      "45 ptrel doesnt match [1.97, 0.778, 0.594, 0.168, 0.395, 0.134, ... 1.21, 1.23, 17.8, 1.56, 1.96, 0.302] [1.97, 0.777, 0.597, 0.17, 0.393, 0.131, ... 3.32, 1.21, 17.8, 1.56, 1.96, 0.302]\n",
      "46 length doesnt match 24 25\n",
      "47 ptrel doesnt match [1.24, 2.33, 4.21, 5.04, 3.74, 1.35, ... 0.651, 1.38, 1.21, 1.24, 0.0735, 0.233] [1.24, 2.34, 4.22, 5.07, 3.73, 1.35, ... 0.651, 1.39, 1.21, 1.24, 0.0751, 0.231]\n",
      "48 ptrel doesnt match [1.42, 1.35, 4.7, 0.0156, 1.63, 2.63, ... 0.689, 3.4, 0.275, 0.312, 1.41, 1.89] [1.41, 1.35, 4.7, 0.0213, 1.64, 2.6, 0.256, ... 0.692, 3.39, 0.273, 0.312, 1.4, 1.89]\n",
      "49 ptrel doesnt match [1.75, 1.66, 0.963, 2.41, 0.744, 7.88, ... 0.274, 2.1, 0.64, 0.834, 1.34, 0.892] [1.75, 1.65, 0.964, 2.4, 0.748, 7.89, ... 0.273, 2.11, 0.64, 0.832, 1.34, 0.888]\n",
      "50 length doesnt match 24 25\n",
      "51 length doesnt match 14 16\n",
      "52 ptrel doesnt match [0.405, 0.571, 0.837, 4.47, 0.576, 0.328, 0.1, 0.386, 0.233, 0.687, 0.267, 0.599] [0.408, 0.567, 0.836, 4.46, 0.574, 0.319, ... 0.37, 0.232, 0.687, 0.268, 0.598]\n",
      "53 ptrel doesnt match [0.432, 1.59, 0.651, 0.841, 0.428, 1.72, 1.19, 0.54, 1.17] [0.431, 1.65, 0.643, 0.841, 0.425, 1.72, 1.19, 0.541, 1.17]\n",
      "54 length doesnt match 29 30\n",
      "55 length doesnt match 18 17\n",
      "56 ptrel doesnt match [0.512, 0.421, 0.384, 0.633, 0.0522, 0.302, ... 0.253, 0.688, 0.561, 0.225, 0.303] [0.514, 0.425, 0.389, 0.639, 0.0518, 0.307, ... 0.254, 0.688, 0.558, 0.227, 0.303]\n",
      "57 ptrel doesnt match [7.19, 2.4, 2.59, 0.439, 1.21, 0.183, ... 0.162, 0.688, 1.34, 0.392, 0.89, 0.84] [7.18, 2.4, 2.59, 0.437, 1.28, 0.184, ... 0.165, 0.69, 1.34, 0.394, 0.893, 0.84]\n",
      "58 length doesnt match 17 20\n",
      "59 length doesnt match 25 28\n",
      "60 ptrel doesnt match [0.555, 0.179, 0.85, 0.0726, 0.29, 0.715, ... 0.319, 0.646, 0.612, 1.27, 0.115] [0.555, 0.179, 0.85, 0.0728, 0.291, 0.72, ... 0.312, 0.651, 0.609, 1.22, 0.112]\n",
      "61 ptrel doesnt match [0.932, 1.31, 3.69, 4.72, 0.781, 2.29, ... 0.271, 1.18, 0.149, 0.874, 0.935, 0.53] [0.933, 1.31, 3.71, 4.74, 0.778, 2.28, ... 0.269, 1.2, 0.15, 0.873, 0.936, 0.527]\n",
      "62 length doesnt match 30 35\n",
      "63 ptrel doesnt match [0.89, 0.597, 0.296, 0.84, 0.107, 0.513, ... 3.32, 0.227, 1.95, 0.25, 0.27, 0.87] [0.893, 0.597, 0.307, 0.815, 0.1, 0.459, ... 3.32, 0.227, 1.97, 0.249, 0.269, 0.873]\n",
      "64 ptrel doesnt match [3.07, 1.31, 2.51, 2.61, 1.88, 0.501, ... 0.14, 0.86, 0.526, 1.1, 0.901, 0.418] [3.07, 1.31, 2.51, 2.61, 1.88, 0.501, ... 0.133, 0.852, 0.516, 1.06, 0.901, 0.418]\n",
      "65 ptrel doesnt match [0.25, 0.359, 0.38, 0.393, 0.364, 1.78, ... 1.01, 0.982, 0.713, 0.279, 0.692, 0.374] [0.251, 0.358, 0.373, 0.42, 0.363, 1.78, ... 1.01, 0.984, 0.718, 0.292, 0.679, 0.37]\n",
      "66 length doesnt match 25 28\n",
      "67 length doesnt match 29 31\n",
      "68 ptrel doesnt match [0.691, 2.17, 0.88, 0.229, 0.264, 1.78, ... 0.768, 0.204, 1.98, 1.02, 2.22, 1.42] [0.692, 2.17, 0.878, 0.23, 0.264, 1.79, ... 0.767, 0.205, 1.97, 1.02, 2.22, 1.42]\n",
      "69 ptrel doesnt match [2.3, 11.7, 0.302, 0.516, 5.93, 1.27, ... 0.861, 5.19, 0.439, 0.409, 0.811, 2.39] [2.29, 11.7, 0.303, 0.519, 5.92, 1.27, ... 0.86, 5.2, 0.439, 0.412, 0.811, 2.39]\n",
      "70 ptrel doesnt match [0.178, 0.0625, 0.575, 0.386, 0.36, 0.988, ... 0.599, 0.965, 0.47, 0.779, 0.431] [0.169, 0.0602, 0.583, 0.391, 0.361, 0.985, ... 0.603, 0.954, 0.473, 0.78, 0.429]\n",
      "71 ptrel doesnt match [0.194, 1.58, 0.521, 0.22, 0.488, 0.705, ... 0.13, 0.252, 0.428, 0.282, 0.382, 0.737] [0.193, 1.56, 0.522, 0.222, 0.49, 0.702, ... 0.252, 0.434, 0.284, 0.382, 0.736]\n",
      "72 ptrel doesnt match [0.401, 0.177, 0, 2.4, 0.257, 0.0651, ... 0.167, 0.394, 0.402, 0.231, 0.469, 1.15] [0.398, 0.198, 0.0694, 2.77, 0.26, 0.0652, ... 0.394, 0.413, 0.231, 0.469, 1.15]\n",
      "73 ptrel doesnt match [0.391, 0.312, 0.349, 0.394, 1.24, 0.911, ... 1.91, 1.01, 0.937, 0.827, 0.963, 3.83] [0.392, 0.313, 0.352, 0.396, 1.24, 0.911, ... 1.91, 1.01, 0.936, 0.826, 0.962, 3.84]\n",
      "74 ptrel doesnt match [0.676, 2.84, 0.748, 1.34, 0.175, 0.285, ... 1.43, 1.09, 1.02, 1.7, 1.68, 0.523] [0.676, 2.84, 0.743, 1.35, 0.18, 0.286, ... 1.43, 1.08, 1.02, 1.7, 1.68, 0.523]\n",
      "75 ptrel doesnt match [0.459, 0.237, 1.15, 1.83, 0.416, 0.551, ... 2.42, 1.41, 2.59, 6.36, 1.63, 0.523] [0.459, 0.238, 1.13, 1.88, 0.414, 0.555, ... 2.42, 1.41, 2.59, 6.37, 1.63, 0.522]\n",
      "76 ptrel doesnt match [1.74, 0.624, 0.143, 1.5, 0.0884, 0.157, ... 1.63, 0.208, 1.13, 2.68, 1.47, 0.8] [1.74, 0.627, 0.143, 1.51, 0.105, 0.153, ... 1.62, 0.203, 1.14, 2.68, 1.47, 0.8]\n",
      "77 ptrel doesnt match [0.39, 0.808, 0.306, 0.427, 0.372, 0.904, ... 0.887, 1.22, 1.59, 0.906, 1.03, 0.64] [0.39, 0.809, 0.306, 0.427, 0.37, 0.897, ... 0.884, 1.22, 1.6, 0.902, 1.03, 0.642]\n",
      "78 ptrel doesnt match [0.804, 0.389, 0.887, 0.562, 2.73, 0.622, 0.514, 1.01] [0.81, 0.387, 0.875, 0.562, 2.75, 0.625, 0.514, 1.01]\n",
      "79 length doesnt match 31 32\n",
      "80 ptrel doesnt match [1.29, 0.254, 0.11, 1.24, 1.72, 1.46, 1.51, 0.542, 2.1, 1.45, 1.84, 0.914, 0.248] [1.31, 0.251, 0.11, 1.23, 1.73, 1.47, 1.51, 0.543, 2.1, 1.46, 1.84, 0.905, 0.25]\n",
      "81 ptrel doesnt match [2.41, 1.17, 3.24, 2.11, 1.23, 4.35, ... 0.406, 0.245, 0.69, 0.854, 1.36, 0.609] [2.41, 1.17, 3.24, 2.1, 1.23, 4.33, ... 0.403, 0.246, 0.688, 0.853, 1.36, 0.616]\n",
      "82 ptrel doesnt match [0.695, 5.16, 1.13, 0.714, 1.1, 0.669, ... 0.506, 0.373, 0.926, 0.181, 0.58, 0.387] [0.691, 5.18, 1.13, 0.728, 1.11, 0.667, ... 0.506, 0.374, 0.926, 0.183, 0.586, 0.388]\n",
      "83 ptrel doesnt match [0.977, 0.5, 0.274, 0.626, 0.99, 0.32, ... 0.708, 0.514, 0.377, 0.633, 0.365, 1.02] [0.976, 0.502, 0.274, 0.627, 0.99, 0.325, ... 0.513, 0.376, 0.641, 0.367, 1.02]\n",
      "84 length doesnt match 38 40\n",
      "85 ptrel doesnt match [8.83, 0.638, 0.982, 2.08, 1.42, 0.863, ... 1.63, 4.46, 1.59, 1.4, 0.879, 0.148] [8.87, 0.638, 0.984, 2.07, 1.43, 0.863, 1.2, ... 1.64, 4.46, 1.59, 1.4, 0.879, 0.148]\n",
      "86 ptrel doesnt match [0.385, 0.867, 0.434, 0.515, 0.177, 0.23, ... 0.364, 0.239, 0.272, 0.813, 0.243] [0.388, 0.863, 0.433, 0.535, 0.17, 0.227, ... 1.29, 0.362, 0.24, 0.274, 0.812, 0.244]\n",
      "87 ptrel doesnt match [3.02, 1.1, 1.9, 0.723, 0.704, 0.303, 0.514, ... 0.37, 0.512, 2.49, 1.21, 0.476, 1.5] [3.02, 1.1, 1.9, 0.722, 0.709, 0.305, ... 0.385, 0.522, 2.52, 1.19, 0.472, 1.5]\n",
      "88 length doesnt match 20 24\n",
      "89 length doesnt match 27 31\n",
      "90 ptrel doesnt match [0.34, 0.291, 0.186, 1.03, 0.63, 1.76, ... 1.02, 1.18, 0.827, 1.42, 0.59, 0.659] [0.341, 0.294, 0.185, 1.03, 0.629, 1.74, ... 1.02, 1.18, 0.829, 1.42, 0.584, 0.657]\n",
      "91 ptrel doesnt match [0.645, 0.443, 0.566, 0.478, 0.337, 0.31, ... 0.581, 0.51, 0.394, 0.351, 0.635] [0.646, 0.444, 0.567, 0.477, 0.337, 0.309, ... 0.57, 0.515, 0.401, 0.355, 0.641]\n",
      "92 ptrel doesnt match [0.245, 0.445, 2.04, 1.09, 5.6, 0.988, ... 0.64, 1.79, 1.46, 0.414, 0.679, 0.655] [0.248, 0.447, 2.04, 1.08, 5.59, 0.976, ... 0.642, 1.79, 1.48, 0.414, 0.68, 0.655]\n",
      "93 ptrel doesnt match [0.707, 0.883, 0.905, 0.73, 0.168, 0.334, ... 0.348, 0.572, 0.623, 0.485, 0.285] [0.702, 0.879, 0.912, 0.758, 0.166, 0.333, ... 0.345, 0.571, 0.624, 0.491, 0.284]\n",
      "94 pt doesnt match [12.3, 6.97, 7.41, 9.96, 7.75, 6.74, 6.88, ... 20.5, 5.5, 10.5, 11.2, 7.95, 5.09] [12.3, 6.97, 7.41, 9.96, 7.75, 6.74, 6.88, ... 20.5, 5.5, 10.5, 7.95, 12.3, 5.09]\n",
      "94 eta doesnt match [0.203, 0.00842, -0.0374, -0.0558, -0.0738, ... -1.62, -1.62, -1.55, 0.418, 0.452] [0.203, 0.00842, -0.0374, -0.0558, -0.0738, ... -1.62, -1.62, 0.418, 0.203, 0.452]\n",
      "94 ptrel doesnt match [3.92, 0.796, 0.516, 0.711, 1.05, 0.301, ... 1.71, 0.665, 1.21, 0.784, 0.489, 0.37] [3.92, 0.794, 0.515, 0.71, 1.05, 0.305, ... 1.72, 0.667, 1.21, 0.486, 3.22, 0.371]\n",
      "95 ptrel doesnt match [1.46, 0.713, 1.1, 0.231, 0.343, 0.49, 0.575, ... 2.03, 1.01, 1.85, 11.3, 2.9, 1.15] [1.47, 0.716, 1.1, 0.234, 0.345, 0.484, 0.575, ... 2.04, 1.01, 1.85, 11.3, 2.9, 1.15]\n",
      "96 ptrel doesnt match [0.715, 0.894, 0.785, 1.93, 0.969, 0.442, ... 0.364, 0.524, 0.335, 0.862, 0.936] [0.714, 0.915, 0.785, 1.87, 1.05, 0.437, ... 2.83, 0.362, 0.522, 0.335, 0.862, 0.935]\n",
      "97 ptrel doesnt match [0.739, 1.46, 0.809, 0.591, 1.47, 5.47, ... 1.17, 0.186, 0.993, 0.678, 1.24, 1.26] [0.737, 1.46, 0.813, 0.585, 1.48, 5.47, ... 1.17, 0.182, 0.993, 0.681, 1.24, 1.26]\n",
      "98 ptrel doesnt match [4.22, 4.66, 0.746, 1.01, 2.64, 4.02, 1.33, ... 1, 0.42, 0.585, 1.78, 0.928, 2.35] [4.23, 4.65, 0.749, 1.01, 2.65, 4.02, 1.34, ... 1, 0.425, 0.587, 1.78, 0.924, 2.35]\n",
      "99 length doesnt match 39 41\n"
     ]
    }
   ],
   "source": [
    "for ie in range(100):\n",
    "    if len(TrkInc[ie]) != len(df[ie].TrkInc_pt):\n",
    "        print(ie, \"length doesnt match\", len(TrkInc[ie]), len(df[ie].TrkInc_pt))\n",
    "    else:\n",
    "        for b in ak.fields(TrkInc):\n",
    "            if len(TrkInc[ie]) > 0 and any(\n",
    "                abs(TrkInc[ie][b] - df[ie][\"TrkInc_\" + b])\n",
    "                / (df[ie][\"TrkInc_\" + b] + 1e-3)\n",
    "                > 1e-3\n",
    "            ):\n",
    "                print(ie, b, \"doesnt match\", TrkInc[ie][b], df[ie][\"TrkInc_\" + b])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Array [0.0429, 0.0429, 0.0429, ... 132, 150] type='76 * float32'>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 \t 2 \t 0.0 \t 21 \t False \t -1 -1 0.04285188611507102\n",
      "1 \t 21 \t 0.0 \t 21 \t False \t -1 -1 0.04285188611507102\n",
      "2 \t 2 \t 540.0 \t 23 \t False \t 0 1 0.04285188611507102\n",
      "3 \t 21 \t 540.0 \t 23 \t False \t 0 1 0.04285188611507102\n",
      "4 \t 2 \t 594.0 \t 44 \t False \t 2 -1 0.04285188611507102\n",
      "5 \t 21 \t 24.1875 \t 51 \t True \t 3 -1 0.04285188611507102\n",
      "6 \t 21 \t 1.27734375 \t 62 \t True \t 3 -1 0.04285188611507102\n",
      "7 \t 211 \t 0.102294921875 \t 2 \t True \t -1 -1 0.04285188611507102\n",
      "8 \t 2 \t 23.0 \t 71 \t True \t 5 -1 0.04285188611507102\n",
      "9 \t 21 \t 3.46875 \t 71 \t True \t 3 -1 0.04285188611507102\n",
      "10 \t 2 \t 243.5 \t 71 \t True \t 4 -1 0.04285188611507102\n",
      "11 \t 21 \t 24.75 \t 71 \t True \t 4 -1 0.04285188611507102\n",
      "12 \t 21 \t 202.0 \t 71 \t True \t 4 -1 0.04285188611507102\n",
      "13 \t 21 \t 19.9375 \t 71 \t True \t 4 -1 0.04285188611507102\n",
      "14 \t -3 \t 31.875 \t 71 \t True \t 4 -1 0.04285188611507102\n",
      "15 \t -211 \t 0.541015625 \t 2 \t True \t -1 -1 0.04285188611507102\n",
      "16 \t 3 \t 9.9375 \t 71 \t True \t 4 -1 0.04285188611507102\n",
      "17 \t 21 \t 8.6875 \t 71 \t True \t 4 -1 0.04285188611507102\n",
      "18 \t 21 \t 22.6875 \t 71 \t True \t 4 -1 0.04285188611507102\n",
      "19 \t 21 \t 1.1640625 \t 71 \t True \t 3 -1 0.04285188611507102\n",
      "20 \t 21 \t 2.5546875 \t 71 \t True \t 3 -1 0.04285188611507102\n",
      "21 \t 21 \t 0.955078125 \t 71 \t True \t 3 -1 0.04285188611507102\n",
      "22 \t 21 \t 1.08203125 \t 71 \t True"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " \t 3 -1 0.04285188611507102\n",
      "23 \t 21 \t 5.765625 \t 71 \t True \t 3 -1 0.04285188611507102\n",
      "24 \t 21 \t 0.564453125 \t 71 \t True \t 3 -1 0.04285188611507102\n",
      "25 \t 21 \t 10.1875 \t 71 \t True \t 3 -1 0.04285188611507102\n",
      "26 \t 21 \t 15.8125 \t 71 \t True \t 3 -1 0.04285188611507102\n",
      "27 \t 21 \t 75.25 \t 71 \t True \t 3 -1 0.04285188611507102\n",
      "28 \t 21 \t 3.359375 \t 71 \t True \t 3 -1 0.04285188611507102\n",
      "29 \t 21 \t 8.75 \t 71 \t True \t 3 -1 0.04285188611507102\n",
      "30 \t 21 \t 3.2421875 \t 71 \t True \t 3 -1 0.04285188611507102\n",
      "31 \t 21 \t 12.09375 \t 71 \t True \t 3 -1 0.04285188611507102\n",
      "32 \t 21 \t 8.0625 \t 71 \t True \t 3 -1 0.04285188611507102\n",
      "33 \t 21 \t 5.296875 \t 71 \t True \t 3 -1 0.04285188611507102\n",
      "34 \t 21 \t 5.703125 \t 71 \t True \t 3 -1 0.04285188611507102\n",
      "35 \t 21 \t 2.1875 \t 71 \t True \t 3 -1 0.04285188611507102\n",
      "36 \t 21 \t 68.75 \t 71 \t True \t 3 -1 0.04285188611507102\n",
      "37 \t 21 \t 246.5 \t 71 \t True \t 3 -1 0.04285188611507102\n",
      "38 \t 21 \t 16.625 \t 71 \t True \t 3 -1 0.04285188611507102\n",
      "39 \t 21 \t 9.96875 \t 71 \t True \t 3 -1 0.04285188611507102\n",
      "40 \t -2 \t 3.34375 \t 71 \t True \t 5 6 0.04285188611507102\n",
      "41 \t 211 \t 0.87109375 \t 2 \t True \t 4 16 0.04285188611507102\n",
      "42 \t 321 \t 53.25 \t 2 \t True \t 4 16 0.04285188611507102\n",
      "43 \t 111 \t 32.875 \t 2 \t True \t 4 16 0.04285188611507102\n",
      "44 \t 21 \t 19.6875 \t 71 \t True \t -1 -1 0.04285188611507102\n",
      "45 \t 21 \t 21.6875 \t 71 \t True \t -1 -1 0.04285188611507102\n",
      "46 \t 310 \t 30.625 \t 2 \t True \t 4 10 0.04285188611507102\n",
      "47 \t 310 \t 7.1875 \t 2 \t True \t 4 16 0.04285188611507102\n",
      "48 \t 310 \t 20.4375 \t 2 \t True \t 4 16 0.04285188611507102\n",
      "49 \t 310 \t 2.71875 \t 2 \t True \t 44 45 0.04285188611507102\n",
      "50 \t -13 \t 0.07861328125 \t 1 \t True \t 7 -1 83.68255966012464\n",
      "51 \t 14 \t 0.02386474609375 \t 1 \t True \t 7 -1 83.68255966012464\n",
      "52 \t 211 \t 3.046875 \t 2 \t True \t 4 10 0.04285188611507102\n",
      "53 \t 13 \t 0.537109375 \t 1 \t True \t 15 -1 78.91785770481191\n",
      "54 \t -14 \t 0.011993408203125 \t 1 \t True \t 15 -1 78.91785770481191\n",
      "55 \t 13 \t 6.703125 \t 91 \t False \t 4 16 78.87669767528214\n",
      "56 \t -14 \t 2.390625 \t 91 \t False \t 4 16 78.87669767528214\n",
      "57 \t 13 \t 2.6875 \t 1 \t True \t 55 -1 78.87669767528214\n",
      "58 \t -14 \t 2.0859375 \t 1 \t True \t 56 -1 78.87669767528214\n",
      "59 \t -13 \t 0.66796875 \t 1 \t True \t 41 -1 18.45510363404564\n",
      "60 \t 14 \t 0.2041015625 \t 1 \t True \t 41 -1 18.45510363404564\n",
      "61 \t -13 \t 14.46875 \t 1 \t True \t 42 -1 71.60326147348555\n",
      "62 \t 14 \t 38.75 \t 1 \t True \t 42 -1 71.60326147348555\n",
      "63 \t 22 \t 27.0625 \t 1 \t True \t 43 -1 0.0428922950616115\n",
      "64 \t 211 \t 0.4091796875 \t 2 \t True \t -1 -1 0.04285188611507102\n",
      "65 \t -211 \t 0.1083984375 \t 2 \t True \t 44 45 0.04285188611507102\n",
      "66 \t -12 \t 0.22412109375 \t 1 \t True \t 4 16 10.615321459872037\n",
      "67 \t 11 \t 0.15234375 \t 1 \t True \t 4 16 10.615321459872037\n",
      "68 \t 111 \t 13.53125 \t 2 \t True \t 48 -1 150.45401344003713\n",
      "69 \t -13 \t 1.76953125 \t 1 \t True \t 52 -1 176.2615826179248\n",
      "70 \t 14 \t 1.28125 \t 1 \t True \t 52 -1 176.2615826179248\n",
      "71 \t -13 \t 0.2509765625 \t 1 \t True \t 64 -1 123.05194005239545\n",
      "72 \t 14 \t 0.1591796875 \t 1 \t True \t 64 -1 123.05194005239545\n",
      "73 \t 13 \t 0.0560302734375 \t 1 \t True \t 65 -1 132.15468145465334\n",
      "74 \t -14 \t 0.054443359375 \t 1 \t True \t 65 -1 132.15468145465334\n",
      "75 \t 22 \t 11.34375 \t 1 \t True \t 68 -1 150.45410639210937\n"
     ]
    }
   ],
   "source": [
    "ievt = 54\n",
    "for i in range(len(events.GenPart[ievt].pt)):\n",
    "    print(\n",
    "        i,\n",
    "        \"\\t\",\n",
    "        events.GenPart[ievt].pdgId[i],\n",
    "        \"\\t\",\n",
    "        events.GenPart[ievt].pt[i],\n",
    "        \"\\t\",\n",
    "        events.GenPart[ievt].status[i],\n",
    "        \"\\t\",\n",
    "        bool((events.GenPart[ievt].statusFlags & (1 << 13))[i]),\n",
    "        \"\\t\",\n",
    "        events.GenPart[ievt].genPartIdxMother[i],\n",
    "        events.GenPart[ievt].genPartIdxMother2[i],\n",
    "        np.hypot(events.GenPart[ievt].vx[i], events.GenPart[ievt].vy[i]),\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "1 << 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Array [3, 22, 21, 12, 38, ... 35, 27, 29, 21] type='200 * int64'>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ak.num(events.LostTrk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ... -0.00107, -0.0018, 0.0714, 0, 0, 0, 0, 0, 0]]\n"
     ]
    }
   ],
   "source": [
    "print(events.JetPFCands.pf.d0 - events.JetPFCands.dxyFromPV)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Array [[-0.00424, -0.00171, ... -1, -1]] type='200 * var * float32[parameters={...'>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}