GeneralDistributions.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook that takes dihadron ntuples and plots kinematics distributions also with MC"
   ]
  },
  {
   "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",
    "\n",
    "import seaborn as sns\n",
    "\n",
    "import root_pandas as rpd\n",
    "from root_pandas import read_root\n",
    "\n",
    "#mpl.rcParams['figure.figsize'] = 6.2,3.5\n",
    "#mpl.rcParams['axes.labelsize'] = 17\n",
    "#mpl.rcParams['lines.linewidth'] = 2\n",
    "#mpl.rcParams['lines.markersize'] = 6\n",
    "#mpl.rcParams['legend.fontsize']= 13"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#this needs to be read in chunks otherwise it uses all memory. (from https://github.com/scikit-hep/root_pandas)\n",
    "def getdatainChunks(filename,treename):\n",
    "    dataframe =pd.DataFrame()\n",
    "    for df in read_root(filename, treename, chunksize=100000) :\n",
    "        #print df.shape[0]\n",
    "        dataframe = pd.concat([dataframe,df])\n",
    "    \n",
    "    print dataframe.shape[0]\n",
    "    return dataframe\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = {}\n",
    "df_trigger = {}\n",
    "\n",
    "for target in ['C','Fe','Pb']:\n",
    "    print 'Getting dataframes from target %s'%target\n",
    "    df[target]=getdatainChunks('Pairs_%s.root'%target, target)\n",
    "    df['%s_mix'%target] = getdatainChunks('Pairs_%s.root'%target, '%s_mix'%target)\n",
    "    df['D_%s'%target]= getdatainChunks('Pairs_%s.root'%target, 'D_%s'%target)\n",
    "    df['D_%s_mix'%target] = getdatainChunks('Pairs_%s.root'%target, 'D_%s_mix'%target)\n",
    "    \n",
    "    print 'Get trigger dataframes from target%s'%target\n",
    "    df_trigger['%s_trigger'%target] = getdatainChunks('Pairs_%s.root'%target, '%s_trigger'%target)\n",
    "    df_trigger['D_%s_trigger'%target] = getdatainChunks('Pairs_%s.root'%target, 'D_%s_trigger'%target)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GiBUU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for target in ['D','C','Fe','Pb']:#,'Fe','C']:\n",
    "    #GiBUU\n",
    "    print target\n",
    "    print '----pairs'\n",
    "    df['GiBUU_%s'%target]= getdatainChunks('GiBUU_Pairs_%s.root'%target, target)\n",
    "    print '----trigger'\n",
    "    df_trigger['GiBUU_%s_trigger'%target]= getdatainChunks('GiBUU_Pairs_%s.root'%target, '%s_trigger'%target)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Combine deuterium datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df['D'] = pd.concat([df['D_Pb'],df['D_Fe'],df['D_C']]) #simply combine the deuterium results for all targets\n",
    "df['D_mix'] = pd.concat([df['D_Pb_mix'],df['D_Fe_mix'],df['D_C_mix']]) #simply combine the deuterium results for all targets\n",
    "df_trigger['D_trigger'] = pd.concat([df_trigger['D_Pb_trigger'],df_trigger['D_Fe_trigger'],df_trigger['D_C_trigger']]) #simply combine the deuterium results for all targets\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def applyCut(inputDataframe, cut, text=None):\n",
    "    nbeforecut = inputDataframe.shape[0]\n",
    "    cutDataframe = inputDataframe.query(cut)\n",
    "    if text:\n",
    "        print text, cutDataframe.shape[0], ' (%2.2f '%(100.0*cutDataframe.shape[0]/nbeforecut), '%)'\n",
    "    return cutDataframe\n",
    "def applyCuts(fullDataframe,name='default',isMC=False): \n",
    "    dataframe = fullDataframe\n",
    "    print 'Entries before cut ', dataframe.shape[0]\n",
    "    dataframe.eval('inelasticity = Nu/5.014', inplace=True)\n",
    "    dataframe.eval('E = z*Nu', inplace=True)\n",
    "    dataframe.eval('P = sqrt(E*E-0.140*0.140)', inplace=True)\n",
    "    \n",
    "    dataframe = applyCut(dataframe, 'Q2>1.0', 'Q2>1.0 :')\n",
    "    #dataframe = applyCut(dataframe, 'W>2.0', 'W>2.0 :')\n",
    "    #dataframe = applyCut(dataframe, 'Nu>3.5', 'Nu>3.5 :')\n",
    "\n",
    "    dataframe = applyCut(dataframe, 'z>0.5', 'z>0.5 :')\n",
    "    dataframe = applyCut(dataframe, 'P <2.5 ', 'P<2.5 ')\n",
    "    dataframe = applyCut(dataframe, 'inelasticity<0.85','inelasticity < 0.85')\n",
    "    return dataframe\n",
    "\n",
    "def applyCutsPair(fullDataframe,name='default'):\n",
    "    print 'Starting election on dipion variables'\n",
    "    dataframe = fullDataframe\n",
    "    dataframe.eval('z_tot = z+z_asso', inplace=True)\n",
    "    dataframe.eval('E_asso = z_asso*Nu', inplace=True)\n",
    "    dataframe.eval('P_asso = sqrt(E_asso*E_asso-0.140*0.140)', inplace=True)\n",
    "\n",
    "    dataframe = applyCut(dataframe, 'pid*pid_asso<0', 'Opposite sign pairs')\n",
    "    dataframe = applyCut(dataframe, 'P_asso>0.200', 'P_asso > 200 MeV')\n",
    "    dataframe = applyCut(dataframe, 'P_asso<2.5', 'P_asso <2.5 GeV')\n",
    "    \n",
    "    return dataframe\n",
    "\n",
    "def printPairBreakdown(dataframe):\n",
    "    \n",
    "    print 'Pairs with Leading pi+', dataframe.query('pid==211').shape[0]\n",
    "    print 'Pairs with Leading pi-', dataframe.query('pid==-211').shape[0]\n",
    "    print 'Pairs with Sub-Leading pi+', dataframe.query('pid_asso==211').shape[0]\n",
    "    print 'Pairs with Sub-Leading pi-', dataframe.query('pid_asso==-211').shape[0]\n",
    "    print 'pi+ pi+ pairs',dataframe.query('pid==211 and pid_asso==211').shape[0]\n",
    "    print 'pi- pi- pairs',dataframe.query('pid==-211 and pid_asso==-211').shape[0]\n",
    "    print 'pi+ pi- pairs',dataframe.query('pid==211 and pid_asso==-211').shape[0]\n",
    "    print 'pi- pi+ pairs',dataframe.query('pid==-211 and pid_asso==211').shape[0]\n",
    "    print '//////////////////////////////////////////////////////'\n",
    "    return \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## apply cuts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key in df.keys():\n",
    "    print key\n",
    "    df[key] = applyCuts(df[key])\n",
    "    #printPairBreakdown(df[key])\n",
    "    df[key] = applyCutsPair(df[key])\n",
    "    print '--------------'\n",
    "    \n",
    "for key in df_trigger.keys():\n",
    "    print key\n",
    "    df_trigger[key] = applyCuts(df_trigger[key])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Event Variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plotEventVariables(df):\n",
    "    fig = plt.figure( figsize=(20, 10))\n",
    "    plt.subplot(2,3,1)\n",
    "    plt.hist(df['Nu'], bins=25,range =(2,5), histtype='step', normed=True,label='Data')\n",
    "    plt.legend()\n",
    "    plt.xlabel(\"Nu [GeV]\", fontsize=20)\n",
    "\n",
    "    plt.subplot(2,3,2)\n",
    "    plt.hist(df['Q2'], bins=25,range =(1,5), histtype='step', normed=True,label='Data')\n",
    "    #plt.hist(MC_reco['Q2'], bins=25,range = (1,5),  normed = True,alpha=0.2,label='MC')\n",
    "    plt.xlabel(\"$Q^{2} [GeV^{2}]$\", fontsize=20)\n",
    "\n",
    "    #plt.subplot(2,3,3)\n",
    "    #plt.hist(df['W'], bins=25,range =(1,3), histtype='step', normed=True,label='Data')\n",
    "    #plt.hist(MC_reco['W'], bins=25,range = (1,3), normed = True,alpha=0.2,label='MC')\n",
    "    #plt.xlabel(\"$W$ [GeV]\", fontsize=20)\n",
    "\n",
    "    plt.subplot(2,3,4)\n",
    "    plt.hist(df['Xb'], bins=25,range =(0,1), histtype='step', normed=True,label='Data')\n",
    "    #plt.hist(MC_reco['Xb'], bins=25,range = (0,1),  normed = True,alpha=0.2,label='MC')\n",
    "    plt.xlabel(\"$x$\", fontsize=20)\n",
    "    \n",
    "    plt.subplot(2,3,5)\n",
    "    plt.hist(df['inelasticity'], bins=25,range =(0,1), histtype='step', normed=True,label='Data')\n",
    "    plt.legend()\n",
    "    plt.xlabel(\"inelasticity\", fontsize=20)\n",
    "    plt.axvline(x=.850,color='red',linestyle='--')\n",
    "\n",
    "\n",
    "    #plt.subplot(2,3,5)\n",
    "    #plt.hist(df['SectorEl'], bins = 12, range = (-0.5,5.5),histtype='step', normed=True,label='Data')\n",
    "    #plt.hist(MC_reco['SectorEl'], bins = 12, range=(-0.5,5.5), normed = True,alpha=0.2,label='MC')\n",
    "    #plt.xlabel('Sector El')\n",
    "\n",
    "    #plt.subplot(2,3,6)\n",
    "    #plt.hist(df['deltaZ'], bins = 100, range=(-50,50), histtype='step', normed=True,label='Data')\n",
    "    #plt.hist(MC_reco['deltaZ'], bins = 100, range=(-50,50), normed = True,alpha=0.2,label='MC')\n",
    "    #plt.xlabel('Delta z')\n",
    "    #plt.show()\n",
    "    \n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hexbin(df_trigger['D_trigger']['Q2'],df_trigger['D_trigger']['Nu'], cmap='inferno')\n",
    "plt.hexbin(df_trigger['GiBUU_D_trigger']['Q2'],df_trigger['GiBUU_D_trigger']['Nu'],cmap='inferno')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hexbin(df_trigger['GiBUU_D_trigger']['Q2'],df_trigger['GiBUU_D_trigger']['Nu'],cmap='inferno')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### This is accomplished with sns.regplot (seaborn package needed). It plots the average of a quantity in bins o the other quantities. This is more useful than a 2D distribution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Profile distributions to see correlation with kinematic variables with z2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure( figsize=(14, 5))\n",
    "\n",
    "for ivar,variable in enumerate(['Q2','Xb','Nu','W']):\n",
    "    plt.subplot(1,4,ivar+1)\n",
    "    query = '3.0 < Nu<3.5 and P_asso>0.250'\n",
    "    for i, target in enumerate(['D','C','Fe','Pb']):\n",
    "        sns.regplot(x=df[target].query(query)['z_asso'],y=df[target].query(query)[variable],x_bins=np.logspace(np.log10(0.06), np.log10(0.5),12),fit_reg=None,label=target)\n",
    "    \n",
    "    for i, target in enumerate(['GiBUU_D','GiBUU_C','GiBUU_Fe','GiBUU_Pb']):\n",
    "        sns.regplot(x=df[target].query(query)['z_asso'],y=df[target].query(query)[variable],x_bins=np.logspace(np.log10(0.06), np.log10(0.5),12),fit_reg=None,label=target)\n",
    "   \n",
    "    \n",
    "    plt.legend(loc='best',frameon=False)\n",
    "\n",
    "plt.xlim([0.0,0.5])\n",
    "plt.tight_layout()\n",
    "plt.savefig('ProfilePlots_1.png')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure( figsize=(14, 5))\n",
    "\n",
    "for ivar,variable in enumerate(['z','P','P_asso']):\n",
    "    plt.subplot(1,4,ivar+1)                                       \n",
    "    for i, target in enumerate(['D','C','Fe','Pb']):\n",
    "        sns.regplot(x=df[target]['z_asso'],y=df[target][variable],x_bins=np.logspace(np.log10(0.05), np.log10(0.5),12),fit_reg=None,label=target)\n",
    "    \n",
    "    for i, target in enumerate(['GiBUU_D','GiBUU_C','GiBUU_Fe','GiBUU_Pb']):\n",
    "        sns.regplot(x=df[target]['z_asso'],y=df[target][variable],x_bins=np.logspace(np.log10(0.05), np.log10(0.5),12),fit_reg=None,label=target)\n",
    "   \n",
    "    \n",
    "    plt.legend(loc='best',frameon=False)\n",
    "\n",
    "plt.xlim([0.0,0.5])\n",
    "plt.tight_layout()\n",
    "plt.savefig('ProfilePlots_2.png')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "                               \n",
    "for ivar,variable in enumerate(['xf','xf_asso','y','y_asso','dy']):  \n",
    "    plt.subplot(1,5,ivar+1)                   \n",
    "\n",
    "    for i, target in enumerate(['D','C','Fe','Pb']):\n",
    "        sns.regplot(x=df[target]['z_asso'],y=df[target][variable],x_bins=np.logspace(np.log10(0.05), np.log10(0.5),12),fit_reg=None,label=target)\n",
    "    plt.legend(loc='best',frameon=False)\n",
    "\n",
    "plt.xlim([0.0,0.5])\n",
    "plt.tight_layout()\n",
    "plt.savefig('ProfilePlots_3.png')     \n",
    "plt.clf() \n",
    "                               \n",
    "for ivar,variable in enumerate(['mass','missing_mass','dipion_pt','dphi']): \n",
    "    plt.subplot(1,4,ivar+1)                   \n",
    "    for i, target in enumerate(['D','C','Fe','Pb']):\n",
    "        sns.regplot(x=df[target]['z_asso'],y=df[target][variable],x_bins=np.logspace(np.log10(0.05), np.log10(0.5),12),fit_reg=None,label=target)\n",
    "    plt.legend(loc='best',frameon=False)\n",
    "\n",
    "plt.xlim([0.0,0.5])\n",
    "plt.tight_layout()\n",
    "plt.savefig('ProfilePlots_4.png')    \n",
    "plt.clf() \n",
    "                               \n",
    "           \n",
    "for ivar,variable in enumerate(['phi_lab_asso','theta_lab_asso','phi_lab','theta_lab']): \n",
    "    plt.subplot(1,4,ivar+1)                   \n",
    "    for i, target in enumerate(['D','C','Fe','Pb']):\n",
    "        sns.regplot(x=df[target]['z_asso'],y=df[target][variable],x_bins=np.logspace(np.log10(0.05), np.log10(0.5),12),fit_reg=None,label=target)\n",
    "    plt.legend(loc='best',frameon=False)\n",
    "\n",
    "plt.xlim([0.0,0.5])\n",
    "plt.tight_layout()\n",
    "plt.savefig('ProfilePlots_5.png')    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## plot for mixed events. In this case, key correlation should go away\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Profile plots but only for z1+z2<0.8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure( figsize=(21, 21))\n",
    "\n",
    "for ivar,variable in enumerate(['Q2','Nu','Xb','W','z','P','P_asso','mass','missing_mass','dy','dipion_pt','xf','xf_asso','y','y_asso','phi_pq']):\n",
    "    plt.subplot(4,4,ivar+1)\n",
    "    for i, target in enumerate(['Pb','D']):\n",
    "        sns.regplot(x=df[target].query('z_tot<0.8')['z_asso'],y=df[target].query('z_tot<0.8')[variable],x_bins=np.logspace(np.log10(0.05), np.log10(0.5),9),fit_reg=None,label=target)\n",
    "    plt.legend(loc='best',frameon=False)\n",
    "\n",
    "plt.xlim([0.0,0.5])\n",
    "plt.savefig('ProfilePlots_withztotcut.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## invariant mass distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(3, 4,sharex=True, sharey=True,figsize=(14,9), gridspec_kw={'wspace':0})\n",
    "\n",
    "for i, target in enumerate(['D','C','Fe','Pb']):\n",
    "    axs[0,i].hist(df[target]['mass'], bins=50, histtype='step', range=(0,2.0),normed=True,label='%s'%target,linewidth=2.0)\n",
    "    axs[0,i].hist(df['%s_mix'%target]['mass'], bins=50, range=(0,2.0),normed=True,alpha=0.4,label='%s mixed'%target)\n",
    "    axs[0,i].axvline(x=0.770,linestyle='--',color='black',alpha=0.6)\n",
    "    axs[0,i].xaxis.set_label_text(\"Dipion Mass  [GeV]\", fontsize=13)\n",
    "    axs[0,i].legend(loc='best',frameon=False)\n",
    "\n",
    "axs[0,0].yaxis.set_label_text(\"entries (normalized)\", fontsize=13)\n",
    "\n",
    "\n",
    "for i, target in enumerate(['D','C','Fe','Pb']):\n",
    "    axs[1,i].hist(df[target].query('z_tot<0.8')['mass'], bins=50, histtype='step', range=(0,2.0),normed=True,label='%s'%target,linewidth=2.0)\n",
    "    axs[1,i].hist(df['%s_mix'%target].query('z_tot<0.8')['mass'], bins=50, range=(0,2.0),normed=True,alpha=0.4,label='%s mixed'%target)\n",
    "    axs[1,i].axvline(x=0.770,linestyle='--',color='black',alpha=0.6)\n",
    "    axs[1,i].xaxis.set_label_text(\"Dipion Mass  [GeV]\", fontsize=13)\n",
    "    axs[1,i].legend(loc='best',frameon=False)\n",
    "\n",
    "for i, target in enumerate(['D','C','Fe','Pb']):\n",
    "    axs[2,i].hist(df[target].query('z_tot>0.8')['mass'], bins=50, histtype='step', range=(0,2.0),normed=True,label='%s'%target,linewidth=2.0)\n",
    "    axs[2,i].hist(df['%s_mix'%target].query('z_tot>0.8')['mass'], bins=50, range=(0,2.0),normed=True,alpha=0.4,label='%s mixed'%target)\n",
    "    axs[2,i].axvline(x=0.770,linestyle='--',color='black',alpha=0.6)\n",
    "    axs[2,i].xaxis.set_label_text(\"Dipion Mass  [GeV]\", fontsize=13)\n",
    "    axs[2,i].legend(loc='best',frameon=False)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig('Massforztotbins.png')    \n",
    "                               "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Missing mas distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 4,sharex=True, sharey=True,figsize=(14,4), gridspec_kw={'wspace':0})\n",
    "\n",
    "for i, target in enumerate(['D','C','Fe','Pb']):\n",
    "    axs[i].hist(df[target]['missing_mass'], bins=50, histtype='step', range=(0,2.0),normed=True,label='all',linewidth=2.0)\n",
    "    axs[i].hist(df[target].query('z_tot<0.8')['missing_mass'], bins=50, histtype='step', range=(0,2.0),normed=True,label='z1+z2<0.8',linewidth=2.0)\n",
    "    axs[i].hist(df[target].query('z_tot>0.8')['missing_mass'], bins=50, histtype='step', range=(0,2.0),normed=True,label='z1+z2>0.8',linewidth=2.0)\n",
    "\n",
    "    #axs[i].hist(df['%s_mix'%target]['missing_mass'], bins=50, range=(0,2.0),normed=True,alpha=0.4,label='%s mixed'%target)\n",
    "    axs[i].axvline(x=0.938,linestyle='--',color='black',alpha=0.6)\n",
    "    axs[i].xaxis.set_label_text(\"Missing mass  [GeV]\", fontsize=13)\n",
    "    axs[i].set_title(target)\n",
    "    axs[i].legend(loc='best',frameon=False)\n",
    "\n",
    "axs[0].yaxis.set_label_text(\"pdf\", fontsize=13)\n",
    "\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig('MissingMassforztotbins.png')    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Z and z_asso distributions without z1+z2 selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 2,sharey=True,figsize=(14,5), gridspec_kw={'wspace':0})\n",
    "\n",
    "axs[0].hist(df['D']['z_asso'], bins=30,range =(0,0.5), normed=True,label='D',alpha=0.5)\n",
    "\n",
    "for i, target in enumerate(['C','Fe','Pb']):\n",
    "    axs[0].hist(df[target]['z_asso'], bins=30,range =(0,0.5), histtype='step', normed=True,label=[target])\n",
    "    \n",
    "axs[1].hist(df['D']['z'], bins=50,range =(0.5,1.0), normed=True,label='D',alpha=0.5)\n",
    "for i, target in enumerate(['C','Fe','Pb']):\n",
    "    axs[1].hist(df[target]['z'], bins=50,range =(0.5,1.0), histtype='step', normed=True,label=[target]) \n",
    " \n",
    "\n",
    "axs[0].xaxis.set_label_text(\"$z_{2}$\", fontsize=13)\n",
    "axs[1].xaxis.set_label_text(\"$z_{1}$\", fontsize=13)\n",
    "axs[0].yaxis.set_label_text(\"entries (normalized)\", fontsize=13)\n",
    "\n",
    "plt.legend(loc='best',frameon=False)\n",
    "\n",
    "plt.savefig('Zdistributions_all.png')    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Z and z_asso distributions for z1+z2<0.8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 2,sharey=True,figsize=(14,5), gridspec_kw={'wspace':0})\n",
    "\n",
    "axs[0].hist(df['D'].query('z_tot<0.8')['z_asso'], bins=30,range =(0,0.5), normed=True,label='D',alpha=0.5)\n",
    "\n",
    "for i, target in enumerate(['C','Fe','Pb']):\n",
    "    axs[0].hist(df[target].query('z_tot<0.8')['z_asso'], bins=30,range =(0,0.5), histtype='step', normed=True,label=[target])\n",
    "    \n",
    "axs[1].hist(df['D'].query('z_tot<0.8')['z'], bins=50,range =(0.5,1.0), normed=True,label='D',alpha=0.5)\n",
    "for i, target in enumerate(['C','Fe','Pb']):\n",
    "    axs[1].hist(df[target].query('z_tot<0.8')['z'], bins=50,range =(0.5,1.0), histtype='step', normed=True,label=[target]) \n",
    " \n",
    "\n",
    "axs[0].xaxis.set_label_text(\"$z_{2}$\", fontsize=13)\n",
    "axs[1].xaxis.set_label_text(\"$z_{1}$\", fontsize=13)\n",
    "axs[0].yaxis.set_label_text(\"entries (normalized)\", fontsize=13)\n",
    "\n",
    "plt.legend(loc='best',frameon=False)\n",
    "plt.savefig('Zdistributions_cut.png')    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Angular distributions phi_lab eta_lab for all z2 ranges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure( figsize=(12, 4))\n",
    "\n",
    "plt.subplot(1,2,1)\n",
    "plt.hexbin(df['Pb']['phi_lab'],df['Pb']['theta_lab'],cmap='inferno')\n",
    "plt.xlabel(\"$\\phi_{lab}$ (rad)\", fontsize=14)\n",
    "plt.ylabel(\"$\\Theta_{lab}$ (rad) \", fontsize=14)\n",
    "plt.title('Leading pion, $z_{1}$> 0.5')\n",
    "\n",
    "plt.subplot(1,2,2)\n",
    "plt.hexbin(df['Pb']['phi_lab_asso'],df['Pb']['theta_lab_asso'],cmap='viridis')\n",
    "plt.xlabel(\"$\\phi_{lab}$ (rad)\", fontsize=14)\n",
    "plt.ylabel(\"$\\Theta_{lab}$ (rad) \", fontsize=14)\n",
    "plt.title('Subleading pion, $z_{1}$> 0.5')\n",
    "\n",
    "plt.savefig('Angles_leadingandsubleading.png')    \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## angles for z1 for various z2 bins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure( figsize=(12, 12))\n",
    "\n",
    "fig, axs = plt.subplots(1, 4,figsize=(14,4),sharey=True,sharex=True, gridspec_kw={'wspace':0, 'hspace':0})\n",
    "\n",
    "\n",
    "asso_Edges = [0.05,0.08, 0.15,0.3,0.5]\n",
    "asso_bins = [(x, y) for (x, y) in zip(asso_Edges[:-1], asso_Edges[1:])]\n",
    "\n",
    "for j, asso_bin in enumerate(asso_bins):\n",
    "    query_trigger = 'z> %2.2f and z<= %2.2f and theta_lab<1.0'%(0.5,1.0)\n",
    "    query_asso = 'z_asso>%2.2f and z_asso<=%2.2f'%(asso_bin[0],asso_bin[1])\n",
    "    query = query_trigger + ' and ' + query_asso\n",
    "    #plt.subplot(3,3,j+1)\n",
    "    axs[j].hexbin(df['D'].query(query)['phi_lab'],df['D'].query(query)['theta_lab'],cmap='inferno')\n",
    "    axs[j].xaxis.set_label_text(\"$\\phi_{lab}$ (rad)\", fontsize=14)\n",
    "    axs[0].yaxis.set_label_text(\"$\\Theta_{lab}$  (rad)\", fontsize=14)\n",
    "    axs[j].set_title('$z_{1}$> 0.5 \\n %2.2f < $z_{2}$ < %2.2f'%(asso_bin[0],asso_bin[1]), fontsize=13)\n",
    "    \n",
    "\n",
    "plt.savefig('z1LabAngles_z2bins_allpions.png')    \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Angle for z2 in different bins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure( figsize=(12, 12))\n",
    "\n",
    "fig, axs = plt.subplots(1, 4,figsize=(14,4),sharey=True,sharex=True, gridspec_kw={'wspace':0, 'hspace':0})\n",
    "\n",
    "\n",
    "asso_Edges = [0.05,0.08, 0.15,0.3,0.5]\n",
    "asso_bins = [(x, y) for (x, y) in zip(asso_Edges[:-1], asso_Edges[1:])]\n",
    "\n",
    "for j, asso_bin in enumerate(asso_bins):\n",
    "    query_trigger = 'z> %2.2f and z<= %2.2f '%(0.5,1.0)\n",
    "    query_asso = 'z_asso>%2.2f and z_asso<=%2.2f'%(asso_bin[0],asso_bin[1])\n",
    "    #plt.subplot(3,3,j+1)\n",
    "    axs[j].hexbin(df['D'].query(query_asso)['phi_lab_asso'],df['D'].query(query_asso)['theta_lab_asso'],cmap='viridis')\n",
    "    axs[j].xaxis.set_label_text(\"$\\phi_{lab}$ (rad)\", fontsize=14)\n",
    "    axs[0].yaxis.set_label_text(\"$\\Theta_{lab}$  (rad)\", fontsize=14)\n",
    "    axs[j].set_title('$z_{1}$> 0.5 \\n %2.2f < $z_{2}$ < %2.2f'%(asso_bin[0],asso_bin[1]), fontsize=13)\n",
    "    \n",
    "\n",
    "plt.savefig('LabAngles_z2bins_allpions.png')    \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Same as above but separately for pi+ and pi- of associated hadron"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure( figsize=(12, 12))\n",
    "\n",
    "fig, axs = plt.subplots(2, 4,figsize=(14,8),sharey=True,sharex=True, gridspec_kw={'wspace':0, 'hspace':0})\n",
    "\n",
    "\n",
    "asso_Edges = [0.05,0.08, 0.15,0.3,0.5]\n",
    "asso_bins = [(x, y) for (x, y) in zip(asso_Edges[:-1], asso_Edges[1:])]\n",
    "\n",
    "for j, asso_bin in enumerate(asso_bins):\n",
    "    query_trigger = 'z> %2.2f and z<= %2.2f '%(0.5,1.0)\n",
    "    query_asso = 'z_asso>%2.2f and z_asso<=%2.2f'%(asso_bin[0],asso_bin[1])\n",
    "    axs[0,j].hexbin(df['D'].query(query_asso + ' and pid_asso>0')['phi_lab_asso'],df['D'].query(query_asso + ' and pid_asso>0')['theta_lab_asso'],cmap='viridis')\n",
    "    axs[1,j].hexbin(df['D'].query(query_asso + ' and pid_asso<0')['phi_lab_asso'],df['D'].query(query_asso + ' and pid_asso<0')['theta_lab_asso'],cmap='viridis')\n",
    "    axs[1,j].xaxis.set_label_text(\"$\\phi_{lab}$ (rad)\", fontsize=14)\n",
    "    axs[0,0].yaxis.set_label_text(\"$\\Theta_{lab}$  Positive Pions (rad)\", fontsize=14)\n",
    "    axs[1,0].yaxis.set_label_text(\"$\\Theta_{lab}$  Negative Pions (rad)\", fontsize=14)\n",
    "    axs[0,j].set_title('$z_{1}$> 0.5 \\n %2.2f < $z_{2}$ < %2.2f'%(asso_bin[0],asso_bin[1]), fontsize=13)\n",
    "\n",
    "plt.savefig('LabAngles_z2bins_piplusminus.png')    \n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## dphi_lab vs phi_lab for same-event and mixed-events"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plotacceptance(df,df_mix,title='',lab=True):\n",
    "    if(lab==True):\n",
    "        var1 = 'phi_lab'\n",
    "        var1_asso = 'phi_lab_asso'\n",
    "        var2 = 'dphi_lab'\n",
    "        label = 'lab'\n",
    "    else:\n",
    "        var1 = 'phi_pq'\n",
    "        var1_asso = 'phi_pq_asso'\n",
    "        var2 = 'dphi'\n",
    "        label ='PQ'\n",
    "    fig = plt.figure( figsize=(12, 9))\n",
    "    plt.subplot(2,3,1)\n",
    "    plt.title(title)\n",
    "    plt.hexbin(df[var1],df[var2],cmap='inferno',gridsize=25)\n",
    "    plt.xlabel(\"$\\phi_{%s}$ (rad)\"%label, fontsize=14)\n",
    "    plt.ylabel(\"$|d\\phi_{%s}|$ (rad) \"%label, fontsize=14)\n",
    "\n",
    "    plt.subplot(2,3,2)\n",
    "    plt.hist(df[var1],bins=100,label='z1',alpha=0.5)\n",
    "    plt.hist(df[var1_asso],bins=100,label='z2',alpha=0.5)\n",
    "    plt.legend(loc='best',frameon=False)\n",
    "\n",
    "    plt.xlabel(\"$\\phi_{%s}$ (rad)\"%label, fontsize=14)\n",
    "    plt.title(title)\n",
    "\n",
    "    plt.subplot(2,3,3)\n",
    "    plt.hist(df[var2],bins=50)\n",
    "    plt.xlabel(\"$|d\\phi_{%s}|$ (rad)\"%label, fontsize=14)\n",
    "    plt.title(title)\n",
    "    #mixed events\n",
    "    plt.subplot(2,3,4)\n",
    "    plt.hexbin(df_mix[var1],df_mix[var2],cmap='inferno',gridsize=25)\n",
    "    plt.xlabel(\"$\\phi_{%s}$ (rad)\"%label, fontsize=14)\n",
    "    plt.ylabel(\"$|d\\phi_{%s}|$ (rad)\"%label, fontsize=14) \n",
    "    #plt.title(title)\n",
    "\n",
    "    plt.subplot(2,3,5)\n",
    "    plt.hist(df_mix[var1],bins=100,label='z1',alpha=0.5)\n",
    "    plt.hist(df_mix[var1_asso],bins=100,label='z2',alpha=0.5)\n",
    "    plt.legend(loc='best',frameon=False)\n",
    "    plt.xlabel(\"$\\phi_{%s}$ (rad)\"%label, fontsize=14)\n",
    "    #plt.title(title)\n",
    "\n",
    "    plt.subplot(2,3,6)\n",
    "    plt.hist(df_mix[var2],bins=50)\n",
    "    plt.xlabel(\"$|d\\phi_{%s}|$ (rad)\"%label, fontsize=14)\n",
    "    #plt.title(title)\n",
    "\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example for 0.1<z2 < 0.2 for both lab and pq frames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plotacceptance(df['D'].query('z_asso>0.1 and z_asso<0.2'),df['D_mix'].query('z_asso>0.1 and z_asso<0.2'),title='$0.1<z_2 <0.2$')\n",
    "plotacceptance(df['D'].query('z_asso>0.1 and z_asso<0.2'),df['D_mix'].query('z_asso>0.1 and z_asso<0.2'),title='$0.1<z_2 <0.2$',lab=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## dphi_pq vs phi_pq for same-event and mixed-events"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "asso_Edges = [0.05,0.08, 0.15,0.3,0.5]\n",
    "asso_bins = [(x, y) for (x, y) in zip(asso_Edges[:-1], asso_Edges[1:])]\n",
    "\n",
    "for j, asso_bin in enumerate(asso_bins):\n",
    "    query_asso = 'z_asso>%2.2f and z_asso<=%2.2f'%(asso_bin[0],asso_bin[1]) \n",
    "    label =  '$%2.2f<z_2<%2.2f$'%(asso_bin[0],asso_bin[1]) \n",
    "    plotacceptance(df['D'].query(query_asso),df['D_mix'].query(query_asso),title=label,lab=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## z1 vs angles in lab / z2 vs angles in lab\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure( figsize=(12, 8))\n",
    "plt.subplot(2,2,1)\n",
    "plt.hexbin(df['D']['z'],df['D']['theta_lab'], bins='log',cmap='viridis')\n",
    "plt.xlabel(\"$z_{1}$\", fontsize=14)\n",
    "plt.ylabel(\"$\\Theta_{lab}$ leading (rad) \", fontsize=14)\n",
    "\n",
    "plt.subplot(2,2,2)\n",
    "plt.hexbin(df['D']['z'],df['D']['phi_lab'], bins='log',cmap='viridis')\n",
    "plt.xlabel(\"$z_{1}$\", fontsize=14)\n",
    "plt.ylabel(\"$\\phi_{lab}$ (rad) \", fontsize=14)\n",
    "\n",
    "plt.subplot(2,2,3)\n",
    "plt.hexbin(df['D']['z_asso'],df['D']['theta_lab_asso'], bins='log',cmap='inferno')\n",
    "plt.xlabel(\"$z_2$\", fontsize=14)\n",
    "plt.ylabel(\"$\\Theta_{lab}$ subleading (rad) \", fontsize=14)\n",
    "\n",
    "plt.subplot(2,2,4)\n",
    "plt.hexbin(df['D']['z_asso'],df['D']['phi_lab_asso'], bins='log',cmap='inferno')\n",
    "plt.xlabel(\"$z_2$\", fontsize=14)\n",
    "plt.ylabel(\"$\\phi_{lab}$ (rad) \", fontsize=14)\n",
    "\n",
    "\n",
    "plt.savefig('CorrelationAnglesandz1z2.png')    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Same as above but sorted for pions of positive or negative charge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure( figsize=(12, 8))\n",
    "plt.subplot(2,2,1)\n",
    "plt.hexbin(df['D'].query('pid_asso>0')['z_asso'],df['D'].query('pid_asso>0')['theta_lab_asso'], bins='log',cmap='inferno')\n",
    "plt.xlabel(\"$z_2$\", fontsize=14)\n",
    "plt.ylabel(\"$\\Theta_{lab}$ subleading (rad) \", fontsize=14)\n",
    "\n",
    "plt.subplot(2,2,2)\n",
    "plt.hexbin(df['D'].query('pid_asso>0')['z_asso'],df['D'].query('pid_asso>0')['phi_lab_asso'], bins='log',cmap='inferno')\n",
    "plt.xlabel(\"$z_2$\", fontsize=14)\n",
    "plt.ylabel(\"$\\phi_{lab}$ (rad) \", fontsize=14)\n",
    "\n",
    "plt.subplot(2,2,3)\n",
    "plt.hexbin(df['D'].query('pid_asso<0')['z_asso'],df['D'].query('pid_asso<0')['theta_lab_asso'], bins='log',cmap='inferno')\n",
    "plt.xlabel(\"$z_2$\", fontsize=14)\n",
    "plt.ylabel(\"$\\Theta_{lab}$ subleading (rad) \", fontsize=14)\n",
    "\n",
    "plt.subplot(2,2,4)\n",
    "plt.hexbin(df['D'].query('pid_asso<0')['z_asso'],df['D'].query('pid_asso<0')['phi_lab_asso'], bins='log',cmap='inferno')\n",
    "plt.xlabel(\"$z_2$\", fontsize=14)\n",
    "plt.ylabel(\"$\\phi_{lab}$ (rad) \", fontsize=14)\n",
    "\n",
    "\n",
    "plt.savefig('CorrelationAnglesandz1z2.png')    \n"
   ]
  }
 ],
 "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
}