Make_GST_DataFrames.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": [
    "#take in hadron list, and calculate basic properties\n",
    "#return theta_xz, theta_yz, theta_z, ke\n",
    "def calc_hadrons(pdg=None,px=None,py=None,pz=None,energy=None,s=None):\n",
    "    \n",
    "    if s is not None:\n",
    "        pdg = s[0]\n",
    "        px = s[1]\n",
    "        py = s[2]\n",
    "        pz = s[3]\n",
    "        energy = s[4]\n",
    "\n",
    "    thetaxz = np.arctan2(px,pz)\n",
    "    thetayz = np.arctan2(py,pz)\n",
    "    thetaz = np.arctan2(np.sqrt(np.square(px)+np.square(py)),pz)\n",
    "    ke = np.array( [ (energy[i] - particle.Particle.from_pdgid(pdg[i]).mass*1e-3) for i in range(len(pdg)) ],\n",
    "                   dtype=np.dtype(np.float) )\n",
    "    \n",
    "    return thetaxz,thetayz,thetaz,ke"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#utility function to get the max element of a sub-array\n",
    "def get_max_val_by_mask(val_arr,mask):\n",
    "    sub_idx = np.argmax(val_arr[mask])\n",
    "    max_idx = np.arange(val_arr.shape[0])[mask][sub_idx]\n",
    "    return max_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#apply an acceptance based on particle type, angle, and KE thresholds\n",
    "#return acceptance for hadron list, number of passing hadrons, and properties of leading (by KE) hadron\n",
    "def acceptance(pdg=None,px=None,py=None,pz=None,energy=None,thetaz=None,ke=None,\n",
    "               pdgcode=2212,thetaz_thresh_deg=40,ke_thresh_mev=100.,\n",
    "               s=None):\n",
    "    \n",
    "    if s is not None:\n",
    "        pdg = s[0]\n",
    "        px = s[1]\n",
    "        py = s[2]\n",
    "        pz = s[3]\n",
    "        energy = s[4]\n",
    "        thetaz = s[5]\n",
    "        ke = s[6]\n",
    "    \n",
    "    #put in rad and GeV\n",
    "    THETAZ_THRESH_RAD=np.radians(thetaz_thresh_deg)\n",
    "    KE_THRESH_GEV=ke_thresh_mev*1e-3\n",
    "\n",
    "    mask = (pdg==pdgcode)\n",
    "    accept = np.logical_and(mask,np.abs(thetaz)<THETAZ_THRESH_RAD)\n",
    "    \n",
    "    px_1=np.nan\n",
    "    py_1=np.nan\n",
    "    pz_1=np.nan\n",
    "    E_1=np.nan\n",
    "    ke_1=np.nan\n",
    "    thetaz_1=np.nan\n",
    "    if np.any(accept):\n",
    "        max_idx = get_max_val_by_mask(ke,accept)\n",
    "        px_1=px[max_idx]\n",
    "        py_1=py[max_idx]\n",
    "        pz_1=pz[max_idx]\n",
    "        E_1=energy[max_idx]\n",
    "        ke_1=ke[max_idx]\n",
    "        thetaz_1=thetaz[max_idx]\n",
    "            \n",
    "    return accept,np.sum(np.logical_and(accept,ke>KE_THRESH_GEV)),px_1,py_1,pz_1,E_1,ke_1,thetaz_1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "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]].pandas.df()\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": 6,
   "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/weights_%s_%s.root\"%(filepath,knob_name,qual)\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": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#functions for some lepton quantities\n",
    "def add_lepton_calcs(df):\n",
    "    df['ptl'] = np.sqrt(np.square(df['pxl'])+np.square(df['pyl']))\n",
    "    df['thetal'] = np.arctan2(df['ptl'],df['pzl'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "#functions for performing hadron calcs and acceptances\n",
    "def add_hadron_calcs(df):\n",
    "    df['thetaxzf'], df['thetayzf'], df['thetazf'], df['kef'] \\\n",
    "        = zip(*df[[\"pdgf\",\"pxf\",\"pyf\",\"pzf\",\"Ef\"]].apply((lambda x: calc_hadrons(s=x)),axis=1))\n",
    "    \n",
    "def add_neutron_acceptance(df):\n",
    "    df['accept_n'], df['n_n'], df['px_n1'], df['py_n1'], df['pz_n1'], df['E_n1'], df['ke_n1'], df['thetaz_n1'] \\\n",
    "        = zip(*df[[\"pdgf\",\"pxf\",\"pyf\",\"pzf\",\"Ef\",\"thetazf\",\"kef\"]].apply((lambda x: acceptance(s=x,\n",
    "                                                                                               pdgcode=2112,\n",
    "                                                                                               thetaz_thresh_deg=360.,\n",
    "                                                                                               ke_thresh_mev=0.)),\n",
    "                                                                         axis=1))\n",
    "\n",
    "def add_proton_acceptance(df):\n",
    "    df['accept_p'], df['n_p'], df['px_p1'], df['py_p1'], df['pz_p1'], df['E_p1'], df['ke_p1'], df['thetaz_p1'] \\\n",
    "        = zip(*df[[\"pdgf\",\"pxf\",\"pyf\",\"pzf\",\"Ef\",\"thetazf\",\"kef\"]].apply((lambda x: acceptance(s=x,\n",
    "                                                                                               pdgcode=2212,\n",
    "                                                                                               thetaz_thresh_deg=360.,\n",
    "                                                                                               ke_thresh_mev=0.)),\n",
    "                                                                         axis=1))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_base_df(root_filename,out_name=None,max_entries=-1,verbose=True):\n",
    "    \n",
    "    print(\"Reading %d entries from root file %s\"%(max_entries,root_filename))\n",
    "    gst_df = uproot.open(root_filename)['gst'].pandas.df(flatten=False)\n",
    "    \n",
    "    if(max_entries>0):\n",
    "        gst_df = gst_df.loc[:max_entries]\n",
    "    \n",
    "    print(\"Adding lepton quantities.\")\n",
    "    #lepton quantities\n",
    "    add_lepton_calcs(gst_df)\n",
    "    \n",
    "    print(\"Adding hadron quantities.\")\n",
    "    #hadron quantities\n",
    "    add_hadron_calcs(gst_df)\n",
    "    \n",
    "    print(\"Adding hadron acceptances.\")\n",
    "    #acceptances\n",
    "    add_neutron_acceptance(gst_df)\n",
    "    add_proton_acceptance(gst_df)\n",
    "    \n",
    "    if(out_name):\n",
    "        print(\"Writing output file %s\"%out_name)\n",
    "        gst_df.to_hdf(out_name,key=\"gst_df\")\n",
    "    \n",
    "    return gst_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "gst_df = uproot.open(\"/Users/wketchum/Data/LDMX/gntp.0.gst.root\")['gst'].pandas.df(flatten=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Index(['iev', 'neu', 'fspl', 'tgt', 'Z', 'A', 'hitnuc', 'hitqrk', 'resid',\n",
       "       'sea', 'qel', 'mec', 'res', 'dis', 'coh', 'dfr', 'imd', 'imdanh',\n",
       "       'singlek', 'nuel', 'em', 'cc', 'nc', 'charm', 'amnugamma', 'neut_code',\n",
       "       'nuance_code', 'wght', 'xs', 'ys', 'ts', 'Q2s', 'Ws', 'x', 'y', 't',\n",
       "       'Q2', 'W', 'EvRF', 'Ev', 'pxv', 'pyv', 'pzv', 'En', 'pxn', 'pyn', 'pzn',\n",
       "       'El', 'pxl', 'pyl', 'pzl', 'pl', 'cthl', 'nfp', 'nfn', 'nfpip', 'nfpim',\n",
       "       'nfpi0', 'nfkp', 'nfkm', 'nfk0', 'nfem', 'nfother', 'nip', 'nin',\n",
       "       'nipip', 'nipim', 'nipi0', 'nikp', 'nikm', 'nik0', 'niem', 'niother',\n",
       "       'ni', 'pdgi', 'resc', 'Ei', 'pxi', 'pyi', 'pzi', 'nf', 'pdgf', 'Ef',\n",
       "       'pxf', 'pyf', 'pzf', 'pf', 'cthf', 'vtxx', 'vtxy', 'vtxz', 'vtxt',\n",
       "       'sumKEf', 'calresp0'],\n",
       "      dtype='object')"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gst_df.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>iev</th>\n",
       "      <th>ni</th>\n",
       "      <th>pdgi</th>\n",
       "      <th>resc</th>\n",
       "      <th>Ei</th>\n",
       "      <th>pxi</th>\n",
       "      <th>pyi</th>\n",
       "      <th>pzi</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>entry</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>[2112, 111]</td>\n",
       "      <td>[1, 7]</td>\n",
       "      <td>[0.9617546426360942, 0.6555073152025024]</td>\n",
       "      <td>[0.0457321474779242, -0.14466688283247225]</td>\n",
       "      <td>[-0.03015256614920707, -0.1162190789148202]</td>\n",
       "      <td>[0.19796014944401197, 0.6140323903868259]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>[2212]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[0.9725774517657677]</td>\n",
       "      <td>[0.1263301507117572]</td>\n",
       "      <td>[0.19938262288017303]</td>\n",
       "      <td>[0.09919508260842358]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "      <td>[2112, 111]</td>\n",
       "      <td>[1, 2]</td>\n",
       "      <td>[1.1118738559110037, 0.26367230653318857]</td>\n",
       "      <td>[-0.37268949030114623, -0.0217469271935761]</td>\n",
       "      <td>[0.12064437792810033, 0.07995593610690205]</td>\n",
       "      <td>[0.44724465792891954, -0.21080446419342747]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>2</td>\n",
       "      <td>[2212, 211]</td>\n",
       "      <td>[2, 2]</td>\n",
       "      <td>[0.9947425097045423, 0.33418229063236526]</td>\n",
       "      <td>[0.1359989998127492, 0.10373864951607942]</td>\n",
       "      <td>[-0.23426565972115085, 0.15065100688841299]</td>\n",
       "      <td>[0.18916140027650094, 0.24236446717113797]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>4</td>\n",
       "      <td>1</td>\n",
       "      <td>[2212]</td>\n",
       "      <td>[1]</td>\n",
       "      <td>[0.9788680916071129]</td>\n",
       "      <td>[-0.1595306987006702]</td>\n",
       "      <td>[-0.22535854804818672]</td>\n",
       "      <td>[0.039896413074052006]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>5</td>\n",
       "      <td>3</td>\n",
       "      <td>[2112, 111, 111]</td>\n",
       "      <td>[1, 2, 2]</td>\n",
       "      <td>[1.1786630134557343, 0.32716008581848793, 0.23...</td>\n",
       "      <td>[-0.24773881057001107, -0.16275019858550982, 0...</td>\n",
       "      <td>[0.27303580841489067, 0.1356730302253644, -0.0...</td>\n",
       "      <td>[0.6087201552303894, 0.20957155611116918, -0.0...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>6</td>\n",
       "      <td>1</td>\n",
       "      <td>[2212]</td>\n",
       "      <td>[2]</td>\n",
       "      <td>[0.9750088500743606]</td>\n",
       "      <td>[-0.13935469153763638]</td>\n",
       "      <td>[-0.160346926829927]</td>\n",
       "      <td>[0.15860924544901556]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>7</td>\n",
       "      <td>2</td>\n",
       "      <td>[2212, -211]</td>\n",
       "      <td>[1, 1]</td>\n",
       "      <td>[1.6538507098452568, 0.31708058020849444]</td>\n",
       "      <td>[-0.08750094945471691, 0.22525150096807683]</td>\n",
       "      <td>[-0.015516544631886162, -0.028576637929007824]</td>\n",
       "      <td>[1.35903292570813, 0.1717713488228637]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>8</td>\n",
       "      <td>4</td>\n",
       "      <td>[111, 111, 2212, 111]</td>\n",
       "      <td>[2, 1, 1, 4]</td>\n",
       "      <td>[0.37047206161945667, 1.164044835061437, 2.291...</td>\n",
       "      <td>[0.019229258834476813, -0.7193155149876883, 0....</td>\n",
       "      <td>[-0.3434288307375938, 0.12046501043563534, 0.3...</td>\n",
       "      <td>[-0.026773803758052588, 0.8971366294573256, 1....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>9</td>\n",
       "      <td>3</td>\n",
       "      <td>[2212, 211, -211]</td>\n",
       "      <td>[1, 1, 2]</td>\n",
       "      <td>[1.6360782894461952, 1.6039415748735817, 0.344...</td>\n",
       "      <td>[0.5628625365147639, 0.1737637643346316, -0.31...</td>\n",
       "      <td>[-0.10707090375982621, -0.2024995374855611, -0...</td>\n",
       "      <td>[1.2116597338683646, 1.575420200423499, -0.052...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       iev  ni                   pdgi          resc  \\\n",
       "entry                                                 \n",
       "0        0   2            [2112, 111]        [1, 7]   \n",
       "1        1   1                 [2212]           [1]   \n",
       "2        2   2            [2112, 111]        [1, 2]   \n",
       "3        3   2            [2212, 211]        [2, 2]   \n",
       "4        4   1                 [2212]           [1]   \n",
       "5        5   3       [2112, 111, 111]     [1, 2, 2]   \n",
       "6        6   1                 [2212]           [2]   \n",
       "7        7   2           [2212, -211]        [1, 1]   \n",
       "8        8   4  [111, 111, 2212, 111]  [2, 1, 1, 4]   \n",
       "9        9   3      [2212, 211, -211]     [1, 1, 2]   \n",
       "\n",
       "                                                      Ei  \\\n",
       "entry                                                      \n",
       "0               [0.9617546426360942, 0.6555073152025024]   \n",
       "1                                   [0.9725774517657677]   \n",
       "2              [1.1118738559110037, 0.26367230653318857]   \n",
       "3              [0.9947425097045423, 0.33418229063236526]   \n",
       "4                                   [0.9788680916071129]   \n",
       "5      [1.1786630134557343, 0.32716008581848793, 0.23...   \n",
       "6                                   [0.9750088500743606]   \n",
       "7              [1.6538507098452568, 0.31708058020849444]   \n",
       "8      [0.37047206161945667, 1.164044835061437, 2.291...   \n",
       "9      [1.6360782894461952, 1.6039415748735817, 0.344...   \n",
       "\n",
       "                                                     pxi  \\\n",
       "entry                                                      \n",
       "0             [0.0457321474779242, -0.14466688283247225]   \n",
       "1                                   [0.1263301507117572]   \n",
       "2            [-0.37268949030114623, -0.0217469271935761]   \n",
       "3              [0.1359989998127492, 0.10373864951607942]   \n",
       "4                                  [-0.1595306987006702]   \n",
       "5      [-0.24773881057001107, -0.16275019858550982, 0...   \n",
       "6                                 [-0.13935469153763638]   \n",
       "7            [-0.08750094945471691, 0.22525150096807683]   \n",
       "8      [0.019229258834476813, -0.7193155149876883, 0....   \n",
       "9      [0.5628625365147639, 0.1737637643346316, -0.31...   \n",
       "\n",
       "                                                     pyi  \\\n",
       "entry                                                      \n",
       "0            [-0.03015256614920707, -0.1162190789148202]   \n",
       "1                                  [0.19938262288017303]   \n",
       "2             [0.12064437792810033, 0.07995593610690205]   \n",
       "3            [-0.23426565972115085, 0.15065100688841299]   \n",
       "4                                 [-0.22535854804818672]   \n",
       "5      [0.27303580841489067, 0.1356730302253644, -0.0...   \n",
       "6                                   [-0.160346926829927]   \n",
       "7         [-0.015516544631886162, -0.028576637929007824]   \n",
       "8      [-0.3434288307375938, 0.12046501043563534, 0.3...   \n",
       "9      [-0.10707090375982621, -0.2024995374855611, -0...   \n",
       "\n",
       "                                                     pzi  \n",
       "entry                                                     \n",
       "0              [0.19796014944401197, 0.6140323903868259]  \n",
       "1                                  [0.09919508260842358]  \n",
       "2            [0.44724465792891954, -0.21080446419342747]  \n",
       "3             [0.18916140027650094, 0.24236446717113797]  \n",
       "4                                 [0.039896413074052006]  \n",
       "5      [0.6087201552303894, 0.20957155611116918, -0.0...  \n",
       "6                                  [0.15860924544901556]  \n",
       "7                 [1.35903292570813, 0.1717713488228637]  \n",
       "8      [-0.026773803758052588, 0.8971366294573256, 1....  \n",
       "9      [1.2116597338683646, 1.575420200423499, -0.052...  "
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gst_df[0:10][[\"iev\",\"ni\",\"pdgi\",\"resc\",\"Ei\",\"pxi\",\"pyi\",\"pzi\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "gst_ipart_df = pd.DataFrame(columns=[\"iev\",\"ipart\",\"pdg\",\"resc\",\"E\",\"px\",\"py\",\"pz\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  iev ipart   pdg resc         E        px        py        pz\n",
      "0   0     0  2112    1  0.961755  0.045732 -0.030153  0.197960\n",
      "1   0     1   111    7  0.655507 -0.144667 -0.116219  0.614032\n",
      "0   1     0  2212    1  0.972577  0.126330  0.199383  0.099195\n",
      "0   2     0  2112    1  1.111874 -0.372689  0.120644  0.447245\n",
      "1   2     1   111    2  0.263672 -0.021747  0.079956 -0.210804\n",
      "0   3     0  2212    2  0.994743  0.135999 -0.234266  0.189161\n",
      "1   3     1   211    2  0.334182  0.103739  0.150651  0.242364\n",
      "0   4     0  2212    1  0.978868 -0.159531 -0.225359  0.039896\n",
      "0   5     0  2112    1  1.178663 -0.247739  0.273036  0.608720\n",
      "1   5     1   111    2  0.327160 -0.162750  0.135673  0.209572\n",
      "2   5     2   111    2  0.238910  0.158513 -0.088161 -0.077204\n",
      "0   6     0  2212    2  0.975009 -0.139355 -0.160347  0.158609\n",
      "0   7     0  2212    1  1.653851 -0.087501 -0.015517  1.359033\n",
      "1   7     1  -211    1  0.317081  0.225252 -0.028577  0.171771\n",
      "0   8     0   111    2  0.370472  0.019229 -0.343429 -0.026774\n",
      "1   8     1   111    1  1.164045 -0.719316  0.120465  0.897137\n",
      "2   8     2  2212    1  2.291845  0.587502  0.329423  1.979527\n",
      "3   8     3   111    4  0.430553 -0.148354  0.004830  0.380952\n",
      "0   9     0  2212    1  1.636078  0.562863 -0.107071  1.211660\n",
      "1   9     1   211    1  1.603942  0.173764 -0.202500  1.575420\n",
      "2   9     2  -211    2  0.344908 -0.310595 -0.017394 -0.052055\n"
     ]
    }
   ],
   "source": [
    "for row in gst_df[[\"iev\",\"ni\",\"pdgi\",\"Ei\",\"pxi\",\"pyi\",\"pzi\",\"resc\"]][0:10].itertuples(index=False):\n",
    "#    print(row[0],row[1],row[2])\n",
    "    iev = np.full(row[1],row[0])\n",
    "    ipart = np.arange(row[1])\n",
    "#    print(len(iev),len(ipart))\n",
    "    gst_ipart_df = pd.concat([gst_ipart_df,\n",
    "                              pd.DataFrame({\"iev\":np.full(row[1],row[0]),\n",
    "                                            \"ipart\":np.arange(row[1]),\n",
    "                                            \"pdg\":row[2],\n",
    "                                            \"E\":row[3],\n",
    "                                            \"px\":row[4],\n",
    "                                            \"py\":row[5],\n",
    "                                            \"pz\":row[6],\n",
    "                                            \"resc\":row[7]})])\n",
    "print(gst_ipart_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "gst_ipart_df = \\\n",
    "pd.concat([pd.DataFrame({\"iev\":np.full(row[1],row[0]),\n",
    "                                            \"ipart\":np.arange(row[1]),\n",
    "                                            \"pdg\":row[2],\n",
    "                                            \"E\":row[3],\n",
    "                                            \"px\":row[4],\n",
    "                                            \"py\":row[5],\n",
    "                                            \"pz\":row[6],\n",
    "                                            \"resc\":row[7]}) \n",
    "           for row in gst_df[[\"iev\",\"ni\",\"pdgi\",\"Ei\",\"pxi\",\"pyi\",\"pzi\",\"resc\"]].itertuples(index=False)],\n",
    "         ignore_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>iev</th>\n",
       "      <th>ipart</th>\n",
       "      <th>pdg</th>\n",
       "      <th>E</th>\n",
       "      <th>px</th>\n",
       "      <th>py</th>\n",
       "      <th>pz</th>\n",
       "      <th>resc</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>2112</td>\n",
       "      <td>0.961755</td>\n",
       "      <td>0.045732</td>\n",
       "      <td>-0.030153</td>\n",
       "      <td>0.197960</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>111</td>\n",
       "      <td>0.655507</td>\n",
       "      <td>-0.144667</td>\n",
       "      <td>-0.116219</td>\n",
       "      <td>0.614032</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>2212</td>\n",
       "      <td>0.972577</td>\n",
       "      <td>0.126330</td>\n",
       "      <td>0.199383</td>\n",
       "      <td>0.099195</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "      <td>2112</td>\n",
       "      <td>1.111874</td>\n",
       "      <td>-0.372689</td>\n",
       "      <td>0.120644</td>\n",
       "      <td>0.447245</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>111</td>\n",
       "      <td>0.263672</td>\n",
       "      <td>-0.021747</td>\n",
       "      <td>0.079956</td>\n",
       "      <td>-0.210804</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1740088</th>\n",
       "      <td>999995</td>\n",
       "      <td>0</td>\n",
       "      <td>2212</td>\n",
       "      <td>0.985852</td>\n",
       "      <td>-0.079529</td>\n",
       "      <td>-0.068857</td>\n",
       "      <td>0.283695</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1740089</th>\n",
       "      <td>999996</td>\n",
       "      <td>0</td>\n",
       "      <td>2212</td>\n",
       "      <td>0.980588</td>\n",
       "      <td>0.227694</td>\n",
       "      <td>-0.093408</td>\n",
       "      <td>0.143630</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1740090</th>\n",
       "      <td>999997</td>\n",
       "      <td>0</td>\n",
       "      <td>2212</td>\n",
       "      <td>0.991671</td>\n",
       "      <td>0.060129</td>\n",
       "      <td>-0.306679</td>\n",
       "      <td>-0.073415</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1740091</th>\n",
       "      <td>999998</td>\n",
       "      <td>0</td>\n",
       "      <td>2212</td>\n",
       "      <td>0.996182</td>\n",
       "      <td>0.035054</td>\n",
       "      <td>0.150963</td>\n",
       "      <td>0.296659</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1740092</th>\n",
       "      <td>999999</td>\n",
       "      <td>0</td>\n",
       "      <td>2212</td>\n",
       "      <td>1.186637</td>\n",
       "      <td>0.622530</td>\n",
       "      <td>0.323088</td>\n",
       "      <td>0.189271</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>1740093 rows × 8 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "            iev  ipart   pdg         E        px        py        pz  resc\n",
       "0             0      0  2112  0.961755  0.045732 -0.030153  0.197960     1\n",
       "1             0      1   111  0.655507 -0.144667 -0.116219  0.614032     7\n",
       "2             1      0  2212  0.972577  0.126330  0.199383  0.099195     1\n",
       "3             2      0  2112  1.111874 -0.372689  0.120644  0.447245     1\n",
       "4             2      1   111  0.263672 -0.021747  0.079956 -0.210804     2\n",
       "...         ...    ...   ...       ...       ...       ...       ...   ...\n",
       "1740088  999995      0  2212  0.985852 -0.079529 -0.068857  0.283695     3\n",
       "1740089  999996      0  2212  0.980588  0.227694 -0.093408  0.143630     2\n",
       "1740090  999997      0  2212  0.991671  0.060129 -0.306679 -0.073415     2\n",
       "1740091  999998      0  2212  0.996182  0.035054  0.150963  0.296659     1\n",
       "1740092  999999      0  2212  1.186637  0.622530  0.323088  0.189271     1\n",
       "\n",
       "[1740093 rows x 8 columns]"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gst_ipart_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading -1 entries from root file /Users/wketchum/Data/LDMX/gntp.0.gst.root\n",
      "Adding lepton quantities.\n",
      "Adding hadron quantities.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Library/Python/3.8/site-packages/numpy/core/_asarray.py:83: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
      "  return array(a, dtype, copy=False, order=order)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adding hadron acceptances.\n"
     ]
    }
   ],
   "source": [
    "gst_df = create_base_df(\"/Users/wketchum/Data/LDMX/gntp.0.gst.root\",\n",
    "                        max_entries=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "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": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "for kn in knob_names:\n",
    "    gst_df = merge_weights_simple(filepath=\"/Users/wketchum/Data/LDMX\",knob_name=kn,df=gst_df,qual=\"FSIFix\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>iev</th>\n",
       "      <th>neu</th>\n",
       "      <th>fspl</th>\n",
       "      <th>tgt</th>\n",
       "      <th>Z</th>\n",
       "      <th>A</th>\n",
       "      <th>hitnuc</th>\n",
       "      <th>hitqrk</th>\n",
       "      <th>resid</th>\n",
       "      <th>sea</th>\n",
       "      <th>...</th>\n",
       "      <th>wght_FrInel_pi_p1</th>\n",
       "      <th>wght_FrInel_pi_p2</th>\n",
       "      <th>wght_FrPiProd_pi_n2</th>\n",
       "      <th>wght_FrPiProd_pi_n1</th>\n",
       "      <th>wght_FrPiProd_pi_p1</th>\n",
       "      <th>wght_FrPiProd_pi_p2</th>\n",
       "      <th>wght_FormZone_n2</th>\n",
       "      <th>wght_FormZone_n1</th>\n",
       "      <th>wght_FormZone_p1</th>\n",
       "      <th>wght_FormZone_p2</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>11</td>\n",
       "      <td>11</td>\n",
       "      <td>1000220480</td>\n",
       "      <td>22</td>\n",
       "      <td>48</td>\n",
       "      <td>2112</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>False</td>\n",
       "      <td>...</td>\n",
       "      <td>0.926723</td>\n",
       "      <td>0.853446</td>\n",
       "      <td>0.6</td>\n",
       "      <td>0.8</td>\n",
       "      <td>1.2</td>\n",
       "      <td>1.4</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>11</td>\n",
       "      <td>11</td>\n",
       "      <td>1000220480</td>\n",
       "      <td>22</td>\n",
       "      <td>48</td>\n",
       "      <td>2212</td>\n",
       "      <td>0</td>\n",
       "      <td>-99</td>\n",
       "      <td>False</td>\n",
       "      <td>...</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2</td>\n",
       "      <td>11</td>\n",
       "      <td>11</td>\n",
       "      <td>1000220480</td>\n",
       "      <td>22</td>\n",
       "      <td>48</td>\n",
       "      <td>2112</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>False</td>\n",
       "      <td>...</td>\n",
       "      <td>0.865019</td>\n",
       "      <td>0.730039</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>11</td>\n",
       "      <td>11</td>\n",
       "      <td>1000220480</td>\n",
       "      <td>22</td>\n",
       "      <td>48</td>\n",
       "      <td>2212</td>\n",
       "      <td>0</td>\n",
       "      <td>6</td>\n",
       "      <td>False</td>\n",
       "      <td>...</td>\n",
       "      <td>0.779041</td>\n",
       "      <td>0.585628</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>4</td>\n",
       "      <td>11</td>\n",
       "      <td>11</td>\n",
       "      <td>1000220480</td>\n",
       "      <td>22</td>\n",
       "      <td>48</td>\n",
       "      <td>2212</td>\n",
       "      <td>0</td>\n",
       "      <td>-99</td>\n",
       "      <td>False</td>\n",
       "      <td>...</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>999995</th>\n",
       "      <td>999995</td>\n",
       "      <td>11</td>\n",
       "      <td>11</td>\n",
       "      <td>1000220480</td>\n",
       "      <td>22</td>\n",
       "      <td>48</td>\n",
       "      <td>2212</td>\n",
       "      <td>0</td>\n",
       "      <td>-99</td>\n",
       "      <td>False</td>\n",
       "      <td>...</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>999996</th>\n",
       "      <td>999996</td>\n",
       "      <td>11</td>\n",
       "      <td>11</td>\n",
       "      <td>1000220480</td>\n",
       "      <td>22</td>\n",
       "      <td>48</td>\n",
       "      <td>2212</td>\n",
       "      <td>0</td>\n",
       "      <td>-99</td>\n",
       "      <td>False</td>\n",
       "      <td>...</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>999997</th>\n",
       "      <td>999997</td>\n",
       "      <td>11</td>\n",
       "      <td>11</td>\n",
       "      <td>1000220480</td>\n",
       "      <td>22</td>\n",
       "      <td>48</td>\n",
       "      <td>2212</td>\n",
       "      <td>0</td>\n",
       "      <td>-99</td>\n",
       "      <td>False</td>\n",
       "      <td>...</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>999998</th>\n",
       "      <td>999998</td>\n",
       "      <td>11</td>\n",
       "      <td>11</td>\n",
       "      <td>1000220480</td>\n",
       "      <td>22</td>\n",
       "      <td>48</td>\n",
       "      <td>2212</td>\n",
       "      <td>0</td>\n",
       "      <td>-99</td>\n",
       "      <td>False</td>\n",
       "      <td>...</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>999999</th>\n",
       "      <td>999999</td>\n",
       "      <td>11</td>\n",
       "      <td>11</td>\n",
       "      <td>1000220480</td>\n",
       "      <td>22</td>\n",
       "      <td>48</td>\n",
       "      <td>2212</td>\n",
       "      <td>0</td>\n",
       "      <td>-99</td>\n",
       "      <td>False</td>\n",
       "      <td>...</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>1000000 rows × 152 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "           iev  neu  fspl         tgt   Z   A  hitnuc  hitqrk  resid    sea  \\\n",
       "0            0   11    11  1000220480  22  48    2112       0      0  False   \n",
       "1            1   11    11  1000220480  22  48    2212       0    -99  False   \n",
       "2            2   11    11  1000220480  22  48    2112       0      0  False   \n",
       "3            3   11    11  1000220480  22  48    2212       0      6  False   \n",
       "4            4   11    11  1000220480  22  48    2212       0    -99  False   \n",
       "...        ...  ...   ...         ...  ..  ..     ...     ...    ...    ...   \n",
       "999995  999995   11    11  1000220480  22  48    2212       0    -99  False   \n",
       "999996  999996   11    11  1000220480  22  48    2212       0    -99  False   \n",
       "999997  999997   11    11  1000220480  22  48    2212       0    -99  False   \n",
       "999998  999998   11    11  1000220480  22  48    2212       0    -99  False   \n",
       "999999  999999   11    11  1000220480  22  48    2212       0    -99  False   \n",
       "\n",
       "        ...  wght_FrInel_pi_p1  wght_FrInel_pi_p2  wght_FrPiProd_pi_n2  \\\n",
       "0       ...           0.926723           0.853446                  0.6   \n",
       "1       ...           1.000000           1.000000                  1.0   \n",
       "2       ...           0.865019           0.730039                  1.0   \n",
       "3       ...           0.779041           0.585628                  1.0   \n",
       "4       ...           1.000000           1.000000                  1.0   \n",
       "...     ...                ...                ...                  ...   \n",
       "999995  ...           1.000000           1.000000                  1.0   \n",
       "999996  ...           1.000000           1.000000                  1.0   \n",
       "999997  ...           1.000000           1.000000                  1.0   \n",
       "999998  ...           1.000000           1.000000                  1.0   \n",
       "999999  ...           1.000000           1.000000                  1.0   \n",
       "\n",
       "        wght_FrPiProd_pi_n1  wght_FrPiProd_pi_p1  wght_FrPiProd_pi_p2  \\\n",
       "0                       0.8                  1.2                  1.4   \n",
       "1                       1.0                  1.0                  1.0   \n",
       "2                       1.0                  1.0                  1.0   \n",
       "3                       1.0                  1.0                  1.0   \n",
       "4                       1.0                  1.0                  1.0   \n",
       "...                     ...                  ...                  ...   \n",
       "999995                  1.0                  1.0                  1.0   \n",
       "999996                  1.0                  1.0                  1.0   \n",
       "999997                  1.0                  1.0                  1.0   \n",
       "999998                  1.0                  1.0                  1.0   \n",
       "999999                  1.0                  1.0                  1.0   \n",
       "\n",
       "        wght_FormZone_n2  wght_FormZone_n1  wght_FormZone_p1  wght_FormZone_p2  \n",
       "0                    1.0               1.0               1.0               1.0  \n",
       "1                    1.0               1.0               1.0               1.0  \n",
       "2                    1.0               1.0               1.0               1.0  \n",
       "3                    1.0               1.0               1.0               1.0  \n",
       "4                    1.0               1.0               1.0               1.0  \n",
       "...                  ...               ...               ...               ...  \n",
       "999995               1.0               1.0               1.0               1.0  \n",
       "999996               1.0               1.0               1.0               1.0  \n",
       "999997               1.0               1.0               1.0               1.0  \n",
       "999998               1.0               1.0               1.0               1.0  \n",
       "999999               1.0               1.0               1.0               1.0  \n",
       "\n",
       "[1000000 rows x 152 columns]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gst_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "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": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Library/Python/3.8/site-packages/pandas/core/generic.py:2606: PerformanceWarning: \n",
      "your performance may suffer as PyTables will pickle object types that it cannot\n",
      "map directly to c-types [inferred_type->mixed,key->block4_values] [items->Index(['pdgi', 'resc', 'Ei', 'pxi', 'pyi', 'pzi', 'pdgf', 'Ef', 'pxf', 'pyf',\n",
      "       'pzf', 'pf', 'cthf', 'thetaxzf', 'thetayzf', 'thetazf', 'kef',\n",
      "       'accept_n', 'accept_p'],\n",
      "      dtype='object')]\n",
      "\n",
      "  pytables.to_hdf(\n"
     ]
    }
   ],
   "source": [
    "with pickle_protocol(4):\n",
    "    gst_df.to_hdf(\"/Users/wketchum/Data/LDMX/base_generation_with_weights_9Jun_FSIFix.hdf\",\"gst_df\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: '/Users/wketchum/Data/LDMX/gntp.0.gst.root:gst'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-62-80adb10a5303>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mgst\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0muproot\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"/Users/wketchum/Data/LDMX/gntp.0.gst.root:gst\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m/Library/Python/3.8/site-packages/uproot/rootio.py\u001b[0m in \u001b[0;36mopen\u001b[0;34m(path, localsource, xrootdsource, httpsource, **options)\u001b[0m\n\u001b[1;32m     52\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     53\u001b[0m             \u001b[0mopenfcn\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlocalsource\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 54\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mROOTDirectory\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mread\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mopenfcn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0moptions\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     55\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     56\u001b[0m     \u001b[0;32melif\u001b[0m \u001b[0m_bytesid\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mparsed\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mscheme\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34mb\"root\"\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/Library/Python/3.8/site-packages/uproot/rootio.py\u001b[0m in \u001b[0;36m<lambda>\u001b[0;34m(path)\u001b[0m\n\u001b[1;32m     49\u001b[0m                 \u001b[0;32mif\u001b[0m \u001b[0mn\u001b[0m \u001b[0;32min\u001b[0m \u001b[0moptions\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     50\u001b[0m                     \u001b[0mkwargs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mn\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0moptions\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpop\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 51\u001b[0;31m             \u001b[0mopenfcn\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mlambda\u001b[0m \u001b[0mpath\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mMemmapSource\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     52\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     53\u001b[0m             \u001b[0mopenfcn\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlocalsource\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/Library/Python/3.8/site-packages/uproot/source/memmap.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, path)\u001b[0m\n\u001b[1;32m     19\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__init__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpath\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     20\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexpanduser\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 21\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_source\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnumpy\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmemmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnumpy\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muint8\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmode\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"r\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     22\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclosed\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     23\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/Library/Python/3.8/site-packages/numpy/core/memmap.py\u001b[0m in \u001b[0;36m__new__\u001b[0;34m(subtype, filename, dtype, mode, offset, shape, order)\u001b[0m\n\u001b[1;32m    223\u001b[0m             \u001b[0mf_ctx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcontextlib_nullcontext\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfilename\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    224\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 225\u001b[0;31m             \u001b[0mf_ctx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mos_fspath\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfilename\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;34m'r'\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mmode\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'c'\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0mmode\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0;34m'b'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    226\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    227\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mf_ctx\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mfid\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: '/Users/wketchum/Data/LDMX/gntp.0.gst.root:gst'"
     ]
    }
   ],
   "source": [
    "gst = uproot.open(\"/Users/wketchum/Data/LDMX/gntp.0.gst.root:gst\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[b'gst;16', b'gst;15']"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "file.keys()"
   ]
  },
  {
   "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
}