EG2_CreateSingleHadronNtuple.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ROOT\n",
    "from ROOT import TFile\n",
    "import numpy as np\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import copy\n",
    "import pandas as pd \n",
    "import time\n",
    "from os import listdir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class particle:\n",
    "    def __init__(self, pid, fourvector, virtual_photon,ThetaPQ):\n",
    "        #work in frame that is aligned with photon\n",
    "        #the photon is aligned with z-direction and has 4-momentum q= (0,0,sqrt(Nu2+Q2),Nu) (note q2 = -Q2)        \n",
    "        self.virtual_photon = virtual_photon\n",
    "        Nu = virtual_photon.E()\n",
    "        Q2 = -virtual_photon.M2()\n",
    "        self.proton = ROOT.TLorentzVector()\n",
    "        self.proton.SetPxPyPzE(0,0,0, 0.938)\n",
    "        self.W = (virtual_photon + self.proton).M()\n",
    "        incoming_e = ROOT.TLorentzVector()\n",
    "        incoming_e.SetPxPyPzE(0,0,5.014,5.014)\n",
    "        part1 = virtual_photon.Vect().Cross(incoming_e.Vect()).Unit()\n",
    "        part2 = virtual_photon.Vect().Cross(fourvector.Vect()).Unit()\n",
    "        sign  = np.sign(part1.Dot(fourvector.Vect()))\n",
    "        self.PhiPQ = sign*np.arccos(part1.Dot(part2))\n",
    "        #print ' Angle my way ', self.PhiPQ\n",
    "        #print 'Hayk (radians) ' , PhiPQ*np.pi/180.0\n",
    "        \n",
    "        photon_pz = np.sqrt(Nu*Nu+Q2) #direction is positive by definition\n",
    "        self.bcm = photon_pz/(Nu + 0.938)#photon-nucleon center-of-mass velocity \n",
    "        self.ycm = 0.5*np.log(( 1+self.bcm)/(1-self.bcm)) #photon-nucleon center-of-mass rapidity\n",
    "        \n",
    "        self.LorentzVector = fourvector #hadron four-vector\n",
    "        self.PhiLab = self.LorentzVector.Phi()\n",
    "        self.E = self.LorentzVector.E() #energy in lab frame\n",
    "        self.vector = self.LorentzVector.Vect()\n",
    "        self.Pt = self.vector.Perp(virtual_photon.Vect().Unit()) #pT with respect to photon direction\n",
    "        self.Pl  = self.vector.Dot(virtual_photon.Vect().Unit()) #pL with respect to photon direction (in lab frame)\n",
    "        self.y =  0.5*np.log( (self.E+self.Pl)/(self.E-self.Pl)) #rapidity in lab frame\n",
    "        self.mT = np.sqrt(self.LorentzVector.M2() + self.Pt*self.Pt)\n",
    "        self.y_star = self.y - self.ycm\n",
    "        self.Pl_star = self.mT*np.sinh(self.y_star)\n",
    "        self.Xf = 2.0*self.Pl_star/self.W \n",
    "        self.pid = pid\n",
    "        self.Zh = self.E/Nu\n",
    "        self.ThetaPQ = np.arctan(self.Pt/self.Pl)\n",
    "        #print 'My COSPQ' , np.arctan(self.Pt/self.Pl)\n",
    "        #print 'Hayk COSPQ', self.ThetaPQ\n",
    "        \n",
    "    def redefine(self, new_virtual_photon):\n",
    "        #this methods recalculates variables that depend on virtual photon direction\n",
    "        #this helps for event mixing. \n",
    "        incoming_e = ROOT.TLorentzVector()\n",
    "        incoming_e.SetPxPyPzE(0,0,5.014,5.014)\n",
    "        part1 = new_virtual_photon.Vect().Cross(incoming_e.Vect()).Unit()\n",
    "        part2 = new_virtual_photon.Vect().Cross(self.LorentzVector.Vect()).Unit()\n",
    "       \n",
    "        sign  = np.sign(part1.Dot(self.LorentzVector.Vect()))\n",
    "        self.PhiPQ = sign*np.arccos(part1.Dot(part2)) \n",
    "        self.Pt = self.LorentzVector.Vect().Perp(new_virtual_photon.Vect().Unit()) #pT with respect to photon direction\n",
    "        self.Pl  = self.LorentzVector.Vect().Dot(new_virtual_photon.Vect().Unit()) #pL with respect to photon direction (in lab frame)\n",
    "        self.y =  0.5*np.log( (self.E+self.Pl)/(self.E-self.Pl)) #rapidity in lab frame\n",
    "        self.ThetaPQ = np.arctan(self.Pt/self.Pl)\n",
    "        self.virtual_photon = new_virtual_photon\n",
    "        \n",
    "        return\n",
    "        \n",
    "        \n",
    "        \n",
    "    def print_properties(self):\n",
    "        print ('Hello, let me introduce myself, i am particle pid = ' , self.pid, ' with index ', self.index, ', from event  #', self.ievt, ' Nu and W', self.Nu, ' ' , self.W)\n",
    "        print ('zh = ', self.Zh, ' phi_pq= ', self.PhiPQ, ' theta_pq=' , self.ThetaPQ, 'E = ', self.E, ' xf', self.Xf,'Pt ', self.Pt, ' Pl= ', self.Pl, ' rapidity=' ,  self.y)\n",
    "        print ('pid = ' , self.pid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def getDataframes(filename, Target=1,maxevents=1e9,tree_name='tree_data'):\n",
    "    dphi = np.array([])\n",
    "    ParticlesFromPrevious = []\n",
    "    try:\n",
    "        myfile = TFile.Open('%s'%filename,'READ')\n",
    "    except:\n",
    "        print(\"could not open file\")\n",
    "        \n",
    "    myfile.Print()\n",
    "    mytree = myfile.Get('ntuple_data')\n",
    "    e_tree = myfile.Get('ntuple_data_electrons')\n",
    "        \n",
    "    print (filename, ' has ', mytree.GetEntries(), ' entries')\n",
    "    \n",
    "    tupla_electron = {}\n",
    "    tupla_electron = ['Q2']\n",
    "    tupla_electron = ['Xb']\n",
    "    tupla_electron = ['Nu']\n",
    "    tupla_electron = ['W']\n",
    "\n",
    "    \n",
    "    tupla_hadron = {}\n",
    "    tupla_hadron['pid'] = []\n",
    "    tupla_hadron['xf'] = []\n",
    "    tupla_hadron['xf_default'] = []\n",
    "    tupla_hadron['z']  = []\n",
    "    tupla_hadron['y']  = []\n",
    "    tupla_hadron['ycm'] = []\n",
    "    tupla_hadron['Q2'] = []\n",
    "    tupla_hadron['Xb'] = []\n",
    "    tupla_hadron['Nu'] = []\n",
    "    tupla_hadron['W'] = []\n",
    "    tupla_hadron['phi_pq'] = []\n",
    "    tupla_hadron['theta_pq'] = []\n",
    "    tupla_hadron['TargType'] = []\n",
    "    tupla_hadron['mass'] = []\n",
    "    tupla_hadron['phi_lab'] = []\n",
    "    tupla_hadron['theta_lab'] = []\n",
    "    tupla_hadron['pt'] = []\n",
    "    \n",
    "    start = time.time()\n",
    "\n",
    "    for ievt  in range(mytree.GetEntries()):\n",
    "        mytree.GetEntry(ievt)   \n",
    "        e_tree.GetEntry(ievt)\n",
    "        if mytree.Q2<1.0: continue\n",
    "        if ievt>maxevents: break\n",
    "        if mytree.TargType!=Target: continue\n",
    "        W = 0\n",
    "        Nu = mytree.Nu\n",
    "        #get electron momentum:\n",
    "        Pe = np.sqrt(e_tree.Px*e_tree.Px + e_tree.Py*e_tree.Py+ e_tree.Pz*e_tree.Pz)\n",
    "        scattered_e = ROOT.TLorentzVector()\n",
    "        scattered_e.SetPxPyPzE(e_tree.Px, e_tree.Py, e_tree.Pz, Pe)\n",
    "        incoming_e = ROOT.TLorentzVector()\n",
    "        incoming_e.SetPxPyPzE(0,0,5.014,5.014)\n",
    "        virtual_photon  = incoming_e - scattered_e \n",
    "        virtual_photon_unitvector = virtual_photon.Vect().Unit()\n",
    "        proton = ROOT.TLorentzVector()\n",
    "        proton.SetPxPyPzE(0,0,0, 0.938)\n",
    "        \n",
    "        ##alternative calculation of virtual photon\n",
    "        #scattered electron\n",
    "        #E    = 5.014\n",
    "        #Eprime = E-Nu\n",
    "        #phi_e = -scattered_e.Vect().Phi()\n",
    "        #alt_scattered_e = ROOT.TLorentzVector()\n",
    "        #cos_thetae = 1-mytree.Q2/(2*E*Eprime)\n",
    "        #sin_thetae = np.sqrt(1-cos_thetae*cos_thetae)\n",
    "        #alt_scattered_e.SetPxPyPzE(Eprime*sin_thetae*np.cos(phi_e),\n",
    "        #                       Eprime*sin_thetae*np.sin(phi_e),\n",
    "        #                       Eprime*cos_thetae,\n",
    "        #                       Eprime)\n",
    "        #alt_virtual_photon  = incoming_e - alt_scattered_e\n",
    "        \n",
    "        #print 'Comparing normal and alternative virtual photons momentum'\n",
    "        #print 'virtual ',virtual_photon.Px(), ' ', virtual_photon.Py(), ' ', virtual_photon.Pz(), ' ',virtual_photon.E()\n",
    "        #print 'alt ' ,alt_virtual_photon.Px(), ' ', alt_virtual_photon.Py(), ' ',alt_virtual_photon.Pz(), ' ',alt_virtual_photon.E()\n",
    "        #################################################\n",
    "        \n",
    "        \n",
    "        \n",
    "        #photon-nucleon center-of-mass energy and rapidity\n",
    "        #the photon is aligned with z-direction and has 4-momentum q= (0,0,sqrt(Nu2+Q2),Nu) (note q2 = -Q2)\n",
    "        \n",
    "     \n",
    "        #print 'W ',  W , 'electron W ', e_tree.W\n",
    "        #print (virtual_photon + proton).M()\n",
    "        #print ' Entering main loop over particles'\n",
    "        for i in range(len(mytree.PID)):\n",
    "            if abs(mytree.PID[i]) !=211: continue\n",
    "            i_lv = ROOT.TLorentzVector()    \n",
    "            i_lv.SetPxPyPzE(mytree.Px[i],mytree.Py[i],mytree.Pz[i],mytree.Zh[i]*Nu)\n",
    "#            print i_lv.Phi(), '   ', mytree.Phi[i]*np.pi/180.0\n",
    "            \n",
    "            \n",
    "            i_part = particle(mytree.PID[i], i_lv, virtual_photon, mytree.ThetaPQ[i] )\n",
    "            #print i_part.LorentzVector.Phi()\n",
    "            #print mytree.Phi[i]*np.pi/180.0\n",
    "            #print i_part.LorentzVector.Theta(), '   ', mytree.Theta[i]*np.pi/180.0\n",
    "\n",
    "\n",
    "            if i_part.Zh > 0.0: \n",
    "                tupla_hadron['TargType'].append(mytree.TargType)\n",
    "                tupla_hadron['pid'].append(i_part.pid)\n",
    "                tupla_hadron['xf'].append(i_part.Xf)\n",
    "                tupla_hadron['xf_default'].append(mytree.Xf[i])\n",
    "                tupla_hadron['z'].append(i_part.Zh)\n",
    "                tupla_hadron['y'].append(i_part.y_star)\n",
    "                tupla_hadron['ycm'].append(i_part.ycm)\n",
    "                tupla_hadron['phi_pq'].append(i_part.PhiPQ)\n",
    "                tupla_hadron['theta_pq'].append(i_part.ThetaPQ)\n",
    "                tupla_hadron['mass'].append(i_part.LorentzVector.M())\n",
    "                tupla_hadron['Q2'].append(mytree.Q2)\n",
    "                tupla_hadron['Xb'].append(mytree.Xb)\n",
    "                tupla_hadron['Nu'].append(mytree.Nu)\n",
    "                tupla_hadron['W'].append(mytree.W)\n",
    "                tupla_hadron['phi_lab'].append(i_part.LorentzVector.Phi())\n",
    "                tupla_hadron['theta_lab'].append(mytree.Theta[i])\n",
    "                tupla_hadron['pt'].append(i_part.Pt)\n",
    "                \n",
    "                #i_part.LorentzVector.Theta())\n",
    "                #print 'mytree.Pt[i] ' , mytree.Pt[i], ' check: ' ,i_part.Vector.Perp(virtual_photon_unitvector)\n",
    "                \n",
    "        #print ' going for next event'    \n",
    "        #print ' particles in event', len(particles\n",
    "        ##end loop over events correlations    \n",
    "    end = time.time()\n",
    "    print ('Processed in',  end-start, 'seconds')\n",
    "    \n",
    "    df_electron = pd.DataFrame(tupla_electron)\n",
    "\n",
    "    df_hadron = pd.DataFrame(tupla_hadron)\n",
    "    myfile.Close()\n",
    "    return df_electron,df_hadron"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#for a particle on-shell:\n",
    "# E = mT cosh y \n",
    "# pz = mT sinh y\n",
    "#mT2 = m2 + pt2\n",
    "# rapidity in the ion-at-rest frame, with the photon direction is \n",
    "# y = 0.5 * ln( (E+pz)/(E-pz))\n",
    "#rapidity of the target in the lab frame = 0 (at rest)\n",
    "# y = y* + ycm\n",
    "#bcm = az + bz/ a0 + b0 \n",
    "#bcm = photon_z / Nu + m_{p}\n",
    "#ycm = 0.5 * ln( 1+bcm/1-bcm)\n",
    "#energy of the photon = Nu\n",
    "#momentum of the photon (magnitude) = sqrt(Nu2 + Q2) \n",
    "#q2 = -Q2. (four-momentum transfer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create dataframe df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pb data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = '/home/miguel/EG2_DATA/sebastian/'\n",
    "\n",
    "Files = listdir(path) \n",
    "df['Pb_electron'], df['Pb_hadron'] = [None,None]\n",
    "df['D_Pb_electron'], df['D_Pb_hadron'] = [None,None]\n",
    "\n",
    "for name in Files:\n",
    "    filename = path+name\n",
    "    print (filename)\n",
    "    e, h = getDataframes(filename,Target=2,maxevents=1e3)\n",
    "    df['Pb_electron'] = pd.concat([ df['Pb_electron'], e])\n",
    "    df['Pb_hadron'] = pd.concat([ df['Pb_hadron'], h])\n",
    "\n",
    "    e,h = getDataframes(filename,Target=1,maxevents=1e3)\n",
    "    df['D_Pb_electron'] = pd.concat([ df['D_Pb_electron'], e])\n",
    "    df['D_Pb_hadron'] = pd.concat([ df['D_Pb_hadron'], h])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df['Pb_hadron'].head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fe Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = '/mnt/c/Users/marratia/Linux/Fe_data/'\n",
    "\n",
    "Files = listdir(path) \n",
    "df['Fe'],df['Fe_mix'], df['Fe_trigger'] = [None,None,None]\n",
    "df['D_Fe'],df['D_Fe_mix'], df['D_Fe_trigger'] = [None,None,None]\n",
    "\n",
    "for name in Files:\n",
    "    filename = path+name\n",
    "    pairs, pairs_mix, trigger = getDataframes(filename,Target=2)\n",
    "    df['Fe'] = pd.concat([ df['Fe'], pairs])\n",
    "    df['Fe_mix'] = pd.concat([ df['Fe_mix'], pairs_mix])\n",
    "    df['Fe_trigger'] = pd.concat([ df['Fe_trigger'], trigger])\n",
    "    pairs, pairs_mix, trigger = getDataframes(filename,Target=1)\n",
    "    df['D_Fe'] = pd.concat([ df['D_Fe'], pairs])\n",
    "    df['D_Fe_mix'] = pd.concat([ df['D_Fe_mix'], pairs_mix])\n",
    "    df['D_Fe_trigger'] = pd.concat([ df['D_Fe_trigger'], trigger])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# C data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = '/home/miguel/EG2_DATA/sebastian/C/'\n",
    "Files = listdir(path) \n",
    "df['C'],df['C_mix'], df['C_trigger'] = [None,None,None]\n",
    "df['D_C'],df['D_C_mix'], df['D_C_trigger'] = [None,None,None]\n",
    "\n",
    "for name in Files:\n",
    "    filename = path+name\n",
    "    pairs, pairs_mix, trigger = getDataframes(filename,Target=2,maxevents=1e3)\n",
    "    df['C'] = pd.concat([ df['C'], pairs])\n",
    "    df['C_mix'] = pd.concat([ df['C_mix'], pairs_mix])\n",
    "    df['C_trigger'] = pd.concat([ df['C_trigger'], trigger])\n",
    "    pairs, pairs_mix, trigger = getDataframes(filename,Target=1,maxevents=1e3)\n",
    "    df['D_C'] = pd.concat([ df['D_C'], pairs])\n",
    "    df['D_C_mix'] = pd.concat([ df['D_C_mix'], pairs_mix])\n",
    "    df['D_C_trigger'] = pd.concat([ df['D_C_trigger'], trigger])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = '/mnt/c/Users/marratia/Linux/MC_C/'\n",
    "\n",
    "Files = listdir(path) \n",
    "df['MC'],df['MC_mix'], df['MC_trigger'] = [None,None,None]\n",
    "df['MC_true'],df['MC_true_mix'], df['MC_true_trigger'] = [None,None,None]\n",
    "\n",
    "for name in Files:\n",
    "    filename = path+name\n",
    "    pairs, pairs_mix, trigger = getDataframes(filename,Target=2,tree_name='tree_accept')\n",
    "    df['MC'] = pd.concat([ df['MC'], pairs])\n",
    "    df['MC_mix'] = pd.concat([ df['MC_mix'], pairs_mix])\n",
    "    df['MC_trigger'] = pd.concat([ df['MC_trigger'], trigger])\n",
    "\n",
    "    pairs, pairs_mix, trigger = getDataframes(filename,Target=2,tree_name='tree_thrown')\n",
    "    df['MC_true'] = pd.concat([ df['MC_true'], pairs])\n",
    "    df['MC_true_mix'] = pd.concat([ df['MC_true_mix'], pairs_mix])\n",
    "    df['MC_true_trigger'] = pd.concat([ df['MC_true_trigger'], trigger])\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import root_pandas\n",
    "from root_pandas import read_root\n",
    "from root_pandas import to_root "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save dataframes to ROOT files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "to_root(df['Pb'],'Single_Pb.root', key='Pb')\n",
    "to_root(df['D_Pb'],'Single_Pb.root', key='D_Pb',mode='a')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "to_root(df['Fe'],'Single_Fe.root', key='Fe')\n",
    "to_root(df['D_Fe'],'Single_Fe.root', key='D_Fe',mode='a')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "to_root(df['C'],'Single_C.root', key='C')\n",
    "to_root(df['D_C'],'Single_C.root', key='D_C',mode='a')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df['Pb_trigger'].hist(figsize=(12,12),bins=100)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df['MC_true'].hist(figsize=(12,12),bins=100)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df['Pb_mix'].hist(figsize=(12,12),bins=100)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## See distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df['D_Pb_hadron'].hist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.colors as mcolors\n",
    "\n",
    "plt.hist2d(df['D_Pb_hadron']['z'],df['D_Pb_hadron']['theta_lab'],range=([0,1],[0,140]), norm=mcolors.PowerNorm(0.3), bins=100)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist2d(df['D_Pb_hadron'].query('pid==-211')['z'],df['D_Pb_hadron'].query('pid==-211')['theta_lab'],range=([0,1],[0,140]), norm=mcolors.PowerNorm(0.3), bins=100)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist2d(df['D_Pb_hadron'].query('pid==211')['z'],df['D_Pb_hadron'].query('pid==211')['theta_lab'],range=([0,1],[0,140]), norm=mcolors.PowerNorm(0.3), bins=100)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(df['D_Pb_hadron'].query('pid==211')['z'],range=(0.0,1.0),bins=100)\n",
    "plt.hist(df['Pb_hadron'].query('pid==211')['z'],range=(0.0,1.0),bins=100)\n",
    "\n",
    "plt.yscale('log')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist2d(df['D_Pb_hadron'].query('pid==211')['W'],df['D_Pb_hadron'].query('pid==211')['Q2'],range=([0,5],[0,5]), norm=mcolors.PowerNorm(0.3), bins=100)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(4, 5, figsize=(14,8),sharex=True, sharey=True,gridspec_kw={'hspace': 0, 'wspace':0})\n",
    "\n",
    "\n",
    "WEdges = np.linspace(2.2,3.0,6)\n",
    "Q2Edges = [1.0,1.4,1.6,2.0,4.1]\n",
    "\n",
    "print(WEdges)\n",
    "Wbins = [(x, y) for (x, y) in zip(WEdges[:-1], WEdges[1:])]\n",
    "Q2bins = [(x, y) for (x, y) in zip(Q2Edges[:-1], Q2Edges[1:])]\n",
    "\n",
    "df['D_Pb_hadron'].eval('pt2 = pt*pt',inplace=True)\n",
    "df['Pb_hadron'].eval('pt2 = pt*pt',inplace=True)\n",
    "\n",
    "for iq2, q2bin in enumerate(Q2bins):\n",
    "    print(iq2)\n",
    "    for i, ibin in enumerate(Wbins):\n",
    "        row = 0 \n",
    "        if(i>4): row = 1\n",
    "        col = i\n",
    "        if col>4: col = col-4\n",
    "        #print(row,col)\n",
    "        query = 'pid==211 and W>%2.2f and W<%2.2f and Q2>%2.2f and Q2<%2.2f'%(ibin[0],ibin[1],q2bin[0],q2bin[1])\n",
    "        #print(query)\n",
    "        datacut= df['D_Pb_hadron'].query(query)\n",
    "        axs[0,col].set_title(' W>%2.2f and W<%2.2f'%(ibin[0],ibin[1]))\n",
    "        axs[iq2,0].yaxis.set_label_text('%2.2f<Q2<%2.2f'%(q2bin[0],q2bin[1]), fontsize=13)\n",
    "\n",
    "        axs[iq2,col].hist2d(datacut['z'],datacut['pt2'],range=([0,1],[0,2.0]), norm=mcolors.LogNorm(), bins=50)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 2, figsize=(12,6),sharex=True, sharey=True,gridspec_kw={'hspace': 0, 'wspace':0})\n",
    "axs[0].hist2d(df['D_Pb_hadron'].query('pid==211')['z'],df['D_Pb_hadron'].query('pid==211')['pt2'],range=([0,1.2],[0,2.0]), norm=mcolors.LogNorm(), bins=50)\n",
    "axs[1].hist2d(df['Pb_hadron'].query('pid==211')['z'],df['Pb_hadron'].query('pid==211')['pt2'],range=([0,1.2],[0,2.0]), norm=mcolors.LogNorm(), bins=50)\n",
    "axs[0].set_title('Deuterium',fontsize=20)\n",
    "axs[1].set_title('Lead',fontsize=20)\n",
    "axs[0].yaxis.set_label_text('pt2',fontsize=20)\n",
    "axs[1].xaxis.set_label_text('z',fontsize=20)\n",
    "axs[0].xaxis.set_label_text('z',fontsize=20)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist2d(df['D_Pb_hadron'].query('pid==-211')['z'],df['D_Pb_hadron'].query('pid==-211')['pt2'],range=([0,1],[0,2.0]), norm=mcolors.LogNorm(), bins=50)\n",
    "plt.show()"
   ]
  }
 ],
 "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.7.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}