Make_GST_DataFrames_v2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "python tools loaded.\n"
     ]
    }
   ],
   "source": [
    "import os, sys\n",
    "from python_tools import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib\n",
    "import pickle\n",
    "\n",
    "\n",
    "class PickleProtocol:\n",
    "    def __init__(self, level):\n",
    "        self.previous = pickle.HIGHEST_PROTOCOL\n",
    "        self.level = level\n",
    "\n",
    "    def __enter__(self):\n",
    "        importlib.reload(pickle)\n",
    "        pickle.HIGHEST_PROTOCOL = self.level\n",
    "\n",
    "    def __exit__(self, *exc):\n",
    "        importlib.reload(pickle)\n",
    "        pickle.HIGHEST_PROTOCOL = self.previous\n",
    "\n",
    "\n",
    "def pickle_protocol(level):\n",
    "    return PickleProtocol(level)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#function for mergning in event weights\n",
    "def merge_weights(file,tree_names,var_suffix,df):\n",
    "    for i in range(len(tree_names)):\n",
    "        w_df = uproot.open(file+\":\"+tree_names[i]).arrays(library=\"pd\")\n",
    "        wght_name=\"wght_\"+var_suffix[i]\n",
    "        w_df = w_df.rename(columns={\"wght\": wght_name})\n",
    "        df = df.merge(w_df[[\"iev\",wght_name]],how=\"left\",on=[\"iev\"])\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#function for combining all the event weights from the root files I've made...\n",
    "def merge_weights_simple(filepath,knob_name,df,qual=\"\"):\n",
    "    filename=\"%s/%s_weights_%s.root\"%(filepath,qual,knob_name)\n",
    "    tree_names = [ \"%s_twk%d\"%(knob_name,i) for i in [0,1,3,4] ]\n",
    "    var_suffixes = [ \"%s_n2\"%knob_name,\"%s_n1\"%knob_name,\"%s_p1\"%knob_name,\"%s_p2\"%knob_name ]\n",
    "    return merge_weights(filename,tree_names,var_suffixes,df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "mass_dict = {}\n",
    "#mass_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "knob_names = [\"FrCEx_N\",\"FrAbs_N\",\"FrInel_N\",\"FrPiProd_N\",#\"MFP_N\",\n",
    "              \"FrCEx_pi\",\"FrAbs_pi\",\"FrInel_pi\",\"FrPiProd_pi\"]#\"MFP_pi\",\n",
    "              #\"FormZone\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running run 6 ... \n",
      "\tProcess knob FrCEx_N\n",
      "\tProcess knob FrAbs_N\n",
      "\tProcess knob FrInel_N\n",
      "\tProcess knob FrPiProd_N\n",
      "\tProcess knob FrCEx_pi\n",
      "\tProcess knob FrAbs_pi\n",
      "\tProcess knob FrInel_pi\n",
      "\tProcess knob FrPiProd_pi\n",
      "Running run 7 ... \n",
      "\tProcess knob FrCEx_N\n",
      "\tProcess knob FrAbs_N\n",
      "\tProcess knob FrInel_N\n",
      "\tProcess knob FrPiProd_N\n",
      "\tProcess knob FrCEx_pi\n",
      "\tProcess knob FrAbs_pi\n",
      "\tProcess knob FrInel_pi\n",
      "\tProcess knob FrPiProd_pi\n",
      "Running run 8 ... \n",
      "\tProcess knob FrCEx_N\n",
      "\tProcess knob FrAbs_N\n",
      "\tProcess knob FrInel_N\n",
      "\tProcess knob FrPiProd_N\n",
      "\tProcess knob FrCEx_pi\n",
      "\tProcess knob FrAbs_pi\n",
      "\tProcess knob FrInel_pi\n",
      "\tProcess knob FrPiProd_pi\n",
      "Running run 9 ... \n",
      "\tProcess knob FrCEx_N\n",
      "\tProcess knob FrAbs_N\n",
      "\tProcess knob FrInel_N\n",
      "\tProcess knob FrPiProd_N\n",
      "\tProcess knob FrCEx_pi\n",
      "\tProcess knob FrAbs_pi\n",
      "\tProcess knob FrInel_pi\n",
      "\tProcess knob FrPiProd_pi\n",
      "Running run 10 ... \n",
      "\tProcess knob FrCEx_N\n",
      "\tProcess knob FrAbs_N\n",
      "\tProcess knob FrInel_N\n",
      "\tProcess knob FrPiProd_N\n",
      "\tProcess knob FrCEx_pi\n",
      "\tProcess knob FrAbs_pi\n",
      "\tProcess knob FrInel_pi\n",
      "\tProcess knob FrPiProd_pi\n"
     ]
    }
   ],
   "source": [
    "for run in range(6,11):\n",
    "\n",
    "    print(\"Running run %d ... \"%run)\n",
    "    \n",
    "    rstring = \"run_%d_gntp\"%run\n",
    "    gst_file = uproot.open(\"/Users/wketchum/Data/LDMX/eN_Ti_Base/%s.0.gst.root:gst\"%rstring)\n",
    "    gst_keys = gst_file.keys()\n",
    "    \n",
    "    df_gst_keys = gst_keys[0:74]+gst_keys[80:81]+gst_keys[88:]\n",
    "    df_gst_ip_keys = gst_keys[0:1]+gst_keys[74:80]\n",
    "    df_gst_fp_keys = gst_keys[0:1]+gst_keys[81:88]\n",
    "    \n",
    "    df_gst = gst_file.arrays(df_gst_keys,library=\"pd\")\n",
    "    df_gst_ip = gst_file.arrays(df_gst_ip_keys,library=\"pd\")\n",
    "    df_gst_fp = gst_file.arrays(df_gst_fp_keys,library=\"pd\")\n",
    "    \n",
    "    for kn in knob_names:\n",
    "        print(\"\\tProcess knob %s\"%kn)\n",
    "        df_gst = merge_weights_simple(filepath=\"/Users/wketchum/Data/LDMX/eN_Ti_Rw_FSI/%s\"%rstring,\n",
    "                                      knob_name=kn,df=df_gst,qual=rstring)\n",
    "    \n",
    "    for pdg in set(df_gst_ip[\"pdgi\"].values):\n",
    "#        print(pdg)\n",
    "        m=particle.Particle.from_pdgid(pdg).mass\n",
    "        if not m: m=0.\n",
    "        mass_dict[pdg] = m*1e-3\n",
    "    for pdg in set(df_gst_fp[\"pdgf\"].values):\n",
    "        m=particle.Particle.from_pdgid(pdg).mass\n",
    "        if not m: m=0.\n",
    "        mass_dict[pdg] = m*1e-3\n",
    "    \n",
    "    df_gst_ip[\"mass\"] = df_gst_ip[\"pdgi\"].apply(lambda x: mass_dict[x])\n",
    "    df_gst_fp[\"mass\"] = df_gst_fp[\"pdgf\"].apply(lambda x: mass_dict[x])\n",
    "    \n",
    "    df_gst_ip[\"kei\"] = df_gst_ip[\"Ei\"] - df_gst_ip[\"mass\"]\n",
    "    df_gst_fp[\"kef\"] = df_gst_fp[\"Ef\"] - df_gst_fp[\"mass\"]\n",
    "    \n",
    "    df_gst[\"ptl\"] = np.sqrt(df_gst[\"pxl\"]*df_gst[\"pxl\"]+df_gst[\"pyl\"]*df_gst[\"pyl\"])\n",
    "  \n",
    "    outfile=\"/Users/wketchum/Data/LDMX/eN_DataFrames_11Aug2021/%s_FSI_df.hdf\"%rstring\n",
    "    with pickle_protocol(4):\n",
    "        df_gst.to_hdf(outfile,\"gst_df\",mode=\"w\")\n",
    "        df_gst_ip.to_hdf(outfile,\"gst_ip_df\",mode=\"a\")\n",
    "        df_gst_fp.to_hdf(outfile,\"gst_fp_df\",mode=\"a\")\n",
    "      "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gst_file = uproot.open(\"/Users/wketchum/Data/LDMX/gntp.0.gst.root:gst\")\n",
    "gst_keys = gst_file.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_gst_keys = gst_keys[0:74]+gst_keys[80:81]+gst_keys[88:]\n",
    "df_gst_ip_keys = gst_keys[0:1]+gst_keys[74:80]\n",
    "df_gst_fp_keys = gst_keys[0:1]+gst_keys[81:88]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_gst = gst_file.arrays(df_gst_keys,library=\"pd\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_gst_ip = gst_file.arrays(df_gst_ip_keys,library=\"pd\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_gst_fp = gst_file.arrays(df_gst_fp_keys,library=\"pd\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "knob_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for kn in knob_names:\n",
    "    df_gst = merge_weights_simple(filepath=\"/Users/wketchum/Data/LDMX\",knob_name=kn,df=df_gst,qual=\"FSIFix\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_gst"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_gst_ip[\"mass\"] = df_gst_ip[\"pdgi\"].apply(lambda x: mass_dict[x])\n",
    "df_gst_fp[\"mass\"] = df_gst_fp[\"pdgf\"].apply(lambda x: mass_dict[x])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_gst_ip[\"kei\"] = df_gst_ip[\"Ei\"] - df_gst_ip[\"mass\"]\n",
    "df_gst_fp[\"kef\"] = df_gst_fp[\"Ef\"] - df_gst_fp[\"mass\"]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_gst_fp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with pickle_protocol(4):\n",
    "    df_gst.to_hdf(\"/Users/wketchum/Data/LDMX/base_generation_with_weights_13Jun_FSIFix.hdf\",\"gst_df\",mode=\"w\")\n",
    "    df_gst_ip.to_hdf(\"/Users/wketchum/Data/LDMX/base_generation_with_weights_13Jun_FSIFix.hdf\",\"gst_ip_df\",mode=\"a\")\n",
    "    df_gst_fp.to_hdf(\"/Users/wketchum/Data/LDMX/base_generation_with_weights_13Jun_FSIFix.hdf\",\"gst_fp_df\",mode=\"a\")\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_gst[\"ptl\"] = np.sqrt(df_gst[\"pxl\"]*df_gst[\"pxl\"]+df_gst[\"pyl\"]*df_gst[\"pyl\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_gst[[\"pxl\",\"pyl\",\"pzl\",\"ptl\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_gst_ptlcut_index = df_gst[df_gst[\"ptl\"]>0.4].index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_gst_ptlcut_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_gst_fp.join(df_gst[\"ptl\"],on='entry').query(\"ptl>0.4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_gst_fp.join(df_gst[\"ptl\"],on='entry')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_gst_ip.join(df_gst[\"wght_FrCEx_pi_p1\"],on='entry').query(\"wght_FrCEx_pi_p1!=1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_gst[\"cthl\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}