01-Introduction-MC-v1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Import the packages\n",
    "import ROOT as r\n",
    "import numpy as np\n",
    "import math"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 0. The datasets\n",
    " \n",
    " 1. **Z to $\\mu\\mu$ simultion**\n",
    "\n",
    "We will begin by looking at simulation data corresponding to a Z boson decaying into two muons, which we will be comparing with real CMS data later on after getting familiar with the software.\n",
    "\n",
    "These simulated datasets correspond to the collision data collected by the CMS experiment in 2012.\n",
    "\n",
    "First, let's understand what we are looking at.\n",
    "\n",
    "The dataset name is: \n",
    "\n",
    "`/DYToMuMu_M-20_CT10_8TeV-powheg-pythia6/Summer12_DR53X-PU_S10_START53_V19-v1/AODSIM`\n",
    "\n",
    "The dataset name follows a specific convention to provide a self-explanatory description of the dataset.\n",
    "\n",
    "It follows a pattern like this:\n",
    "\n",
    "`/Primary dataset/Condition/Format/`\n",
    "\n",
    " - The `Primary dataset` tells you about the process for a simulation dataset, and the data taking period and trigger stream that the data was collected if it is data.\n",
    "\n",
    " - The `Condition` part tells you what reconstruction condition has been applied to the dataset.\n",
    "\n",
    " - Lastly, the `Format` part tells you the dataformat. We will work with the simplified version of `MINIAOD` format, called `NanoAOD`, which I have prepared for you in advance.**\n",
    "\n",
    "The dataset is downloaded from here. The files `dy_mc.root`/`dy_mc_small.root` is a tiny fraction of the whole dataset, as we do not need the whole 14TB of events.  \n",
    "http://opendata.cern.ch/record/7741\n",
    "\n",
    " 2. **Data - Double Muon data**\n",
    "\n",
    "We will be using proton-proton collision data at $\\sqrt{s}=8 TeV$ taken during 2012 by the CMS.\n",
    "\n",
    "Of course, we will not be using all of them, the dataset name \n",
    "\n",
    "`/DoubleMuParked/Run2012B-22Jan2013-v1/AOD`\n",
    "\n",
    "tells you with subset of data we will be using.\n",
    "\n",
    "Here `DoubleMuParked` means the data collected was triggered by a stream of High Level Trigger(HLT) trigger paths that targets muon, such as\n",
    "```\n",
    "HLT_DoubleMu5_IsoMu5\n",
    "```\n",
    "which means it requires two muons with at least 5 GeV and one of them is isolated from other particles.\n",
    "\n",
    "Like the simulation data, a skimmed version have been made for you: `Run2012B_DoubleMuParked_nanoaod.root`\n",
    "\n",
    "More details can be found in this record:\n",
    "http://opendata.cern.ch/record/6004\n",
    "\n",
    "\n",
    "<br>\n",
    "<br>\n",
    "<br>\n",
    "\n",
    "\n",
    "** In case you are curious, the conversion is done also with cms open software here:\n",
    "\n",
    "https://github.com/cms-opendata-analyses/AOD2NanoAODOutreachTool"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 1 Exploring ROOT files with pyROOT\n",
    "\n",
    "In this section, we will go over some of the basic operations to explore a root file, working with the small simulation file first.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Opening a ROOT file\n",
    "f = r.TFile(\"dy_mc_small.root\")\n",
    "# list the content of the file\n",
    "f.ls()\n",
    "dir = f.GetDirectory(\"aod2nanoaod\")\n",
    "dir.ls()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TTree\n",
    "TTree is the main object that holds the data. \n",
    "\n",
    "In general, all the objects in a TFile can be obtained with the `Get` Method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = f.Get(\"aod2nanoaod/Events\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## See how many events are there in this tree\n",
    "nEntries = t.GetEntries()\n",
    "print(nEntries)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TTree.Show()\n",
    " There a several methods to see what's inside a tree\n",
    "\n",
    "`Show(n)` prints out the n-th event in the tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t.Show(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TTree.Print()\n",
    "\n",
    "`Print()` gives you more information about the branches, esp dataformat. \n",
    "\n",
    "Here, `run/I` tells you that the branch `run` is of type integer\n",
    "\n",
    "whereas, `PV_x/F` means PV_x is a float\n",
    "\n",
    "Finally, `Muon_pt[nMuon]/F ` is a array of floats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t.Print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Question: Why do we need a different data format for the Muon_pt branch?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TTree::Draw\n",
    "\n",
    "Another very handy way to see the data is the `Draw` method.\n",
    "\n",
    "This is an internal method to loop over all the events on a branch (very efficiently!),\n",
    "\n",
    "and fill an histogram from the branch content\n",
    "\n",
    "To visualize the histogram, we first need to create a Canvas.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c = r.TCanvas('c1','c1',800,600)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example, we can look at the branch called `nMuon`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t.Draw('nMuon')\n",
    "c.Draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Question: Can you guess what does this histogram represents?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The branch names usually contains good description of what does it contain.\n",
    "\n",
    "Similarly, we can explore another branch called `PV_npvs` as well"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t.Draw(\"PV_npvs\")\n",
    "c.Draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Question: Can you guess what `PV_npvs` represents then?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Looping a TTree with a histogram\n",
    "\n",
    "`TTree.Draw()` can only offer a simple plotting, for more custom analysis, we can go through the `TTree` event-by-event by building a for-loop.\n",
    "\n",
    "Let's use a simple for-loop to re-create the `nMuon` histogram.\n",
    "\n",
    " - Learn how to create a histogram\n",
    " - Learn how to write a loop structure for a TTree"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining your own 1D histogram\n",
    "\n",
    "`ROOT` has a very powerful library of histogram class called `TH1` with extensive functionalities.\n",
    "\n",
    "You can construct a very simple 1D histogram like this:\n",
    "\n",
    "`h = r.TH1F(\"name\",\"title\", nBin, x_low, x_hi)` \n",
    "\n",
    " - `TH1F` is a sub-class of `TH1`, using a `float` to store the data of each bin of the histogram\n",
    " - `name` is the name of the histogram object, it is used to identify the histogram object (h, in this case) in a file\n",
    " - `title` is the title to be displayed in plotting\n",
    " - `nBin` is the number of bins \n",
    " - `x_low,x_hi` are the lower and upper range of the 1D histogram respectively.\n",
    " \n",
    " \n",
    "Further documentation about `TH1` can be found here:\n",
    "https://root.cern.ch/doc/master/classTH1.html"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Let's define our first histogram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_nMuon = r.TH1F('nMuon','nMuon',10,0,10)\n",
    "h_nMuon.Draw()\n",
    "c.Update()\n",
    "c.Draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Filling a histogram\n",
    "\n",
    "Say, if you have an event with 4 muons in it. You will want to put 1 event in the 4-th bin."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_nMuon.Fill(4)\n",
    "c.Update()\n",
    "c.Draw()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "## now let's clear the bin content, by setting the 5th bin to 0\n",
    "h_nMuon.SetBinContent(5,0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Looping a TTree\n",
    "\n",
    "By default, all the branches are loaded as the methods of the TTree object.\n",
    "\n",
    "For example, you can simply access the `nMuon` branch like this:\n",
    "\n",
    "`nMuon = t.nMuon`\n",
    "\n",
    "**It will return the value of the tree is currently located\n",
    "\n",
    "To access different events, we can do:\n",
    "\n",
    "`t.GetEntry(i)`\n",
    "\n",
    "This will set the memeory address to load the i-th entry in the Tree.\n",
    "\n",
    "Combining both, a TTree can be looped over with the simple loop structure like this:\n",
    "\n",
    "```\n",
    "for i in range(0,t.GetEntries()):\n",
    "    t.GetEntry(i)\n",
    "    nMuon = t.nMuon\n",
    "    h_nMuon.Fill(nMuon)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(0,t.GetEntries()):\n",
    "    t.GetEntry(i)\n",
    "    nMuon = t.nMuon\n",
    "    h_nMuon.Fill(nMuon)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_nMuon.Draw()\n",
    "c.Update()\n",
    "c.Draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 3. Plotting the di-muon mass spectrum\n",
    "\n",
    "Now we have all the ingredients to plot the di-muon mass spectrum.\n",
    "\n",
    "First, we will need a histogram."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_dimuon_mass = r.TH1F('DiMuon_mass','DiMuon_mass',200,0,200)\n",
    "h_dimuon_mass.Sumw2()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You may find the following formulas useful:\n",
    "\n",
    "Coordinate transformation between pesudeo-rapidity and cartesian coordinates:\n",
    "\n",
    "$\n",
    "p_x = p_T cos(\\phi)\\\\\n",
    "p_y = p_T sin(\\phi)\\\\\n",
    "p_z = p_T sinh(\\eta)\n",
    "$\n",
    "\n",
    "Vector additions and Energy-momentum relations:\n",
    "\n",
    "$\n",
    "E^2 = |\\vec{p}|^2 + m^2\\\\\n",
    "\\vec{p_1}+\\vec{p_2} =  (p1_x+p2_x) \\hat{x}+ (p1_y+p2_y) \\hat{y}+ (p1_z+p2_z) \\hat{z}\\\\\n",
    "$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "## write a function that computes the invariant mass given the event\n",
    "## pt, phi, eta, mass are the arrays of the muons in each event\n",
    "##\n",
    "def invmass(pt,phi,eta,mass):\n",
    "    m = 0\n",
    "    ####################\n",
    "    ## YOUR CODE HERE ##\n",
    "    ####################\n",
    "    return m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "## write a function that select \n",
    "def selectMuons(muons):\n",
    "    select = False\n",
    "\n",
    "    # Minimum selection for di-muon\n",
    "    if muons['n']==2: select = True\n",
    "    return select"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(0,nEntries):\n",
    "    t.GetEntry(i)\n",
    "    nMuon   = t.nMuon\n",
    "    pt      = np.array(t.Muon_pt)          # note: this is an array\n",
    "    phi     = np.array(t.Muon_phi)         # note: this is an array\n",
    "    eta     = np.array(t.Muon_eta)         # note: this is an array \n",
    "    mass    = np.array(t.Muon_mass)        # note: this is an array\n",
    "    isTight = np.array(t.Muon_tightId)     # note: this is an array\n",
    "    charge  = np.array(t.Muon_charge)      # note: this is an array\n",
    "    muons   = {\n",
    "        'n':nMuon,\n",
    "        'pt':pt,\n",
    "        'phi':phi,\n",
    "        'eta':eta,\n",
    "        'isTight':isTight,\n",
    "        'charge':charge,\n",
    "    }\n",
    "    if selectMuons(muons):\n",
    "        h_dimuon_mass.Fill(invmass(pt,phi,eta,mass))\n",
    "        \n",
    "h_dimuon_mass.Draw('hist')\n",
    "h_dimuon_mass.GetXaxis().SetTitle('m_{#mu#mu}[GeV]')\n",
    "c.Update()\n",
    "c.SetLogy(1)\n",
    "c.Draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RESET the histogram if you want to try again!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(0,201): h_dimuon_mass.SetBinContent(i,0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 3 Extracting the Z-boson mass\n",
    "\n",
    "Let's try to extract the Z-boson mass by fitting the peak with a Gaussian function.\n",
    "\n",
    "Mathematical functions in ROOT can be constructed with the class `TF1`*\n",
    "Gaussian function is one of the built-in function in ROOT with the following form:\n",
    "\n",
    "$f(x) = [0] exp\\left(-\\frac{1}{2}((x-[1])/[2]\\right)^2$\n",
    "\n",
    "Here `[0],[1],[2]` are the parameters of the function, corresponding to the `amplitude`, `mean` and `width` of the Gaussian function.\n",
    "\n",
    "Let's create a gaussian function, initialize it with some values, and draw it.\n",
    "\n",
    "* Full documentation of the `TF1` class\n",
    "https://root.cern.ch/doc/master/classTF1.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Here, 'f1' is the name of the object, \n",
    "## 'gaus' is a special string that corresponds to the gaussian formula above\n",
    "## 85 and 95 is the range of the function.\n",
    "## i.e. we are insterested only in that mass window.\n",
    "gaus = r.TF1('f1','gaus',75,105)\n",
    "gaus.SetParameters(1,90,10)\n",
    "gaus.Draw()\n",
    "c.SetLogy(0)\n",
    "c.Draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`TH1` has basic fit function implemented.\n",
    "\n",
    "The default option is the least square fit, \n",
    "\n",
    "$\n",
    "\\chi^2 = \\sum \\left((y_i - f(x_i) )/e_i \\right)^2\n",
    "$\n",
    "\n",
    "The option `R` tells the fit to restrict to the range defined by the function\n",
    "\n",
    "The options `S` simply return the fit result as an object instead of just print-outs.\n",
    "\n",
    "The full options can be found here:\n",
    "\n",
    "https://root.cern.ch/doc/master/classTH1.html#a63eb028df86bc86c8e20c989eb23fb2a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fitResult = h_dimuon_mass.Fit(gaus,'RS')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_dimuon_mass.Draw()\n",
    "c.Draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 3a Quality of the fit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chi2 = gaus.GetChisquare()\n",
    "NDF  = gaus.GetNDF()\n",
    "print('Checking fit quality:')\n",
    "print(\"chi2 = %.3f\\nNDF = %.3f\\nchi2/NDF = %.3f\"% (chi2,NDF, chi2/NDF))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Question:\n",
    "- How does the width of the Gaussian (calculate the full width at half maximum from sigma) compares with the known intrinsic width of the Z boson?\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 4 Take-home exercise"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 4a Knowing more about the muons\n",
    "\n",
    "Now we have the basic tools to explore a ROOT file and analysis. Let's learn deeper about the muons!\n",
    "\n",
    "Here are some questions to think about:\n",
    "\n",
    "- plot muon pT for the leading, subleading, and third-leading muon (when it's there) and explain the features seen in the plot;\n",
    "- plot muon eta for the leading and subleading muons and explain the feature seen in the plot;\n",
    "- plot isTight flag for the three muons;\n",
    "- plot the muon multiplicity for isTight muons;\n",
    "\n",
    "These questions can be solved in a similar way as the `nMuon` example, with the addition of some control statements during the histogram filling.\n",
    "\n",
    "### Hint: all the `Muon_*` branch has already been sorted in decending order of `pt` \n",
    "i.e `Muon_pt[0]` corresponds to the pt of the leading muon, `Muon_eta[0]` would be the $\\eta$ of the same muon\n",
    "\n",
    "\n",
    "To help you get started, let's define some histograms for you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Define the histograms\n",
    "h_mu0_pt      = r.TH1F('h_mu0_pt'     ,'Leading muon pt',200,0,200)\n",
    "h_mu0_eta     = r.TH1F('h_mu0_eta'    ,'Leading muon eta',100,-5,5)\n",
    "h_mu0_isTight = r.TH1F('h_mu0_isTight','Leading muon isTight',2,0,2)\n",
    "\n",
    "h_mu1_pt      = r.TH1F('h_mu1_pt'  ,'Sub-leading muon pt',200,0,200)\n",
    "h_mu1_eta     = r.TH1F('h_mu1_eta' ,'Sub-leading muon eta',100,-5,5)\n",
    "h_mu1_isTight = r.TH1F('h_mu1_isTight','Sub-leading muon isTight',2,0,2)\n",
    "\n",
    "h_mu2_pt      = r.TH1F('h_mu2_pt','Third-Leading muon pt',200,0,200)\n",
    "h_mu2_isTight = r.TH1F('h_mu2_isTight','Third-leading muon isTight',2,0,2)\n",
    "\n",
    "h_nTightMuon  = r.TH1F('h_nTightMuon','Number of tight muons',10,0,10)\n",
    "\n",
    "#collect the histograms into a dic\n",
    "histograms = {\n",
    "        'h_mu0_pt'     :h_mu0_pt,\n",
    "        'h_mu0_eta'    :h_mu0_eta,\n",
    "        'h_mu0_isTight':h_mu0_isTight,\n",
    "        'h_mu1_pt'     :h_mu1_pt,\n",
    "        'h_mu1_eta'    :h_mu1_eta,\n",
    "        'h_mu1_isTight':h_mu1_isTight,\n",
    "        'h_mu2_pt'     :h_mu2_pt,\n",
    "        'h_mu2_isTight':h_mu2_isTight,\n",
    "        'h_nTightMuon' :h_nTightMuon,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(0,nEntries):\n",
    "    t.GetEntry(i)\n",
    "    nMuon   = t.nMuon\n",
    "    pt      = np.array(t.Muon_pt)          # note: this is an array\n",
    "    phi     = np.array(t.Muon_phi)         # note: this is an array\n",
    "    eta     = np.array(t.Muon_eta)         # note: this is an array \n",
    "    mass    = np.array(t.Muon_mass)        # note: this is an array\n",
    "    isTight = np.array(t.Muon_tightId)     # note: this is an array\n",
    "    charge  = np.array(t.Muon_charge)      # note: this is an array\n",
    "    muons   = {\n",
    "        'n':nMuon,\n",
    "        'pt':pt,\n",
    "        'phi':phi,\n",
    "        'eta':eta,\n",
    "        'isTight':isTight,\n",
    "        'charge':charge,\n",
    "    }\n",
    "    #### YOUR CODE HERE ###"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Drawing the muon pT histograms\n",
    "\n",
    "You can execute the following cells to check the histograms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_mu0_pt.SetLineColor(r.kGreen)   ## r.kRed is a global constant stored in ROOT, corresponding to the color code\n",
    "h_mu1_pt.SetLineColor(r.kRed)     ## r.kRed is a global constant stored in ROOT, corresponding to the color code\n",
    "h_mu2_pt.SetLineColor(r.kBlue)\n",
    "h_mu0_pt.Draw()\n",
    "h_mu1_pt.Draw('same')             ## the \"same\" option overlays the histograms on the same canvas\n",
    "h_mu2_pt.Draw('same')\n",
    "h_mu0_pt.GetXaxis().SetTitle('pT[GeV]')\n",
    "c.Update()\n",
    "c.Draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Question: What do you observe? Can you understand the distribution?\n",
    "\n",
    "YOUR ANS HERE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Drawing the muon eta histograms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_mu0_eta.Draw()\n",
    "h_mu0_eta.SetLineColor(r.kRed)     ## r.kRed is a global constant stored in ROOT, corresponding to the color code\n",
    "h_mu1_eta.SetLineColor(r.kBlue)\n",
    "h_mu1_eta.Draw('same')             ## the \"same\" option overlays the histograms on the same canvas\n",
    "h_mu0_eta.GetXaxis().SetTitle('#eta')\n",
    "c.Update()\n",
    "c.Draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Drawing isTight histograms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_mu0_isTight.Draw()\n",
    "h_mu0_isTight.SetLineColor(r.kRed)     ## r.kRed is a global constant stored in ROOT, corresponding to the color code\n",
    "h_mu1_isTight.SetLineColor(r.kBlue)\n",
    "h_mu2_isTight.SetLineColor(r.kGreen)\n",
    "h_mu1_isTight.Draw('same')             ## the \"same\" option overlays the histograms on the same canvas\n",
    "h_mu2_isTight.Draw('same')             ## the \"same\" option overlays the histograms on the same canvas\n",
    "h_mu0_isTight.GetXaxis().SetTitle('isTight')\n",
    "c.Update()\n",
    "c.SetLogy(1)\n",
    "c.Draw()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "h_nTightMuon.Draw()\n",
    "c.SetLogy(0)\n",
    "c.Update()\n",
    "c.Draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Remember to reset the histograms if you want to re-fill them!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "for hname,h in histograms.items():\n",
    "    for b in range(0,h.GetNbinsX()):\n",
    "        h.SetBinContent(b,0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 4b Di-muon selections:\n",
    " - Can we increase the number of events by using events with exactly two isTight muons, not exactly two muons?\n",
    "   - Try modifying the `selectMuon` function to test your ideas!\n",
    " - histogram the product of electric charges for the events with two tight muons; How often do we have same-sign muons?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 4c Z-boson decays\n",
    "\n",
    "Now that you have observe a Z boson, you do more in-depth analysis of those events coming from a Z boson decay.\n",
    "\n",
    "- For events within an invariant mass window near the Z peak, plot pT spectrum of the Z boson.\n",
    "- For these events, plot a correlation (2D histogram) of the opening azimuthal angle (phi) between the two muons and the pT of the Z boson.\n",
    "\n",
    "You can use the following histograms:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_pT_z    = r.TH1F('pT_z','pT of Z boson',100,0,100)\n",
    "h_pTz_phi = r.TH2F('pTz_phi','pT of Z boson v.s. #Delta#phi(#mu#mu)',100,0,100,50,0,math.pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We need to be careful when we calculate azimuthal angle difference, since computers do not know 2pi+phi = phi\n",
    "# so we need to have a function that returns deltaPhi to within 0 to pi\n",
    "def deltaPhi(phi1, phi2):\n",
    "    x = np.abs(phi1 - phi2)\n",
    "    if x >=  math.pi: \n",
    "        x = 2.*math.pi - x\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(0,nEntries):\n",
    "    t.GetEntry(i)\n",
    "    nMuon   = t.nMuon\n",
    "    pt      = np.array(t.Muon_pt)          # note: this is an array\n",
    "    phi     = np.array(t.Muon_phi)         # note: this is an array\n",
    "    eta     = np.array(t.Muon_eta)         # note: this is an array \n",
    "    mass    = np.array(t.Muon_mass)        # note: this is an array\n",
    "    isTight = np.array(t.Muon_tightId)     # note: this is an array\n",
    "    charge  = np.array(t.Muon_charge)      # note: this is an array\n",
    "    muons   = {\n",
    "        'n':nMuon,\n",
    "        'pt':pt,\n",
    "        'phi':phi,\n",
    "        'eta':eta,\n",
    "        'isTight':isTight,\n",
    "        'charge':charge,\n",
    "    }\n",
    "    ### YOUR CODE HERE###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_pT_z.Draw()\n",
    "c.Update()\n",
    "c.Draw()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_pTz_phi.Draw('COLZ')\n",
    "c.Update()\n",
    "c.Draw()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 4d Improving fit quality\n",
    "\n",
    "- Modify the fit to have a sum of a Gaussian and an exponent to describe the background under the peak \n",
    "- Does this improve the fit chi2?\n",
    "- Can you think of other ways to improve the fit quality?\n",
    "\n",
    "You may want to read the `TF1` documentation to think of which functions to use!\n",
    "\n",
    "Here are some skeleton code for you to use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[p0]*exp(-0.5*((x-[p1])/[p2])*((x-[p1])/[p2]))+exp([p3]+[p4]*x)\n"
     ]
    }
   ],
   "source": [
    "## Define the sum of gaussian and exponential:\n",
    "f2 = r.TF1('f2','gaus(0)+expo(3)',75,105)\n",
    "print(f2.GetFormula().GetExpFormula()) \n",
    "f2.SetParameters(1,90,10,6,-0.02)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## ADD your fitting code here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_dimuon_mass.Draw()\n",
    "c.Update()\n",
    "c.Draw()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Now check your fit quality. Does it improve?\n",
    "chi2 = f2.GetChisquare()\n",
    "NDF  = f2.GetNDF()\n",
    "print('Checking fit quality:')\n",
    "print(\"chi2 = %.3f\\nNDF = %.3f\\nchi2/NDF = %.3f\"% (chi2,NDF, chi2/NDF))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}