MakePrefitFigures.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from cycler import cycler\n",
    "from coffea import util\n",
    "from coffea.processor import accumulate\n",
    "import hist\n",
    "import hist.intervals\n",
    "\n",
    "from ttgamma.utils.plotting import (\n",
    "    RebinHist,\n",
    "    SetRangeHist,\n",
    "    GroupBy,\n",
    "    DictToHist,\n",
    "    plotWithRatio,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load output files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set up groupings of samples.\n",
    "Order of dictionary is used for determining plotting order"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from save_to_root import outputMC, outputData\n",
    "\n",
    "groupingCategory = {\n",
    "    \"NonPrompt\": [3j, 4j],\n",
    "    \"MisID\": [2j],\n",
    "    \"Prompt\": [1j],\n",
    "}\n",
    "\n",
    "groupingFineCategory = {\n",
    "    \"Genuine Photon\": [1j],\n",
    "    \"Misidentified Electron\": [2j],\n",
    "    \"Hadronic Photon\": [3j],\n",
    "    \"Hadronic Fake\": [4j],\n",
    "}\n",
    "\n",
    "groupingMCDatasets = {\n",
    "    \"ZG\": [\n",
    "        \"ZGamma_01J_5f_lowMass\",\n",
    "    ],\n",
    "    \"WG\": [\n",
    "        \"WGamma\",\n",
    "    ],\n",
    "    \"other\": [\n",
    "        \"TTbarPowheg_Dilepton\",\n",
    "        \"TTbarPowheg_Semilept\",\n",
    "        \"TTbarPowheg_Hadronic\",\n",
    "        \"W2jets\",\n",
    "        \"W3jets\",\n",
    "        \"W4jets\",\n",
    "        \"DYjetsM50\",\n",
    "        \"ST_s_channel\",\n",
    "        \"ST_tW_channel\",\n",
    "        \"ST_tbarW_channel\",\n",
    "        \"ST_tbar_channel\",\n",
    "        \"ST_t_channel\",\n",
    "        \"TTWtoLNu\",\n",
    "        \"TTWtoQQ\",\n",
    "        \"TTZtoLL\",\n",
    "        \"GJets_HT200To400\",\n",
    "        \"GJets_HT400To600\",\n",
    "        \"GJets_HT600ToInf\",\n",
    "        \"ZZ\",\n",
    "        \"WZ\",\n",
    "        \"WW\",\n",
    "        \"TGJets\",\n",
    "    ],\n",
    "    \"ttgamma\": [\n",
    "        \"TTGamma_Dilepton\",\n",
    "        \"TTGamma_SingleLept\",\n",
    "        \"TTGamma_Hadronic\",\n",
    "    ],\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Group MC histograms\n",
    "outputMC = {\n",
    "    process: accumulate(outputMC[ds] for ds in datasets)\n",
    "    for process, datasets in groupingMCDatasets.items()\n",
    "}\n",
    "\n",
    "\n",
    "def getHisto(name):\n",
    "    return DictToHist(\n",
    "        {process: histos[name] for process, histos in outputMC.items()}, name=\"dataset\"\n",
    "    )\n",
    "\n",
    "\n",
    "# Group data histograms\n",
    "outputDataHist = accumulate([histo for key, histo in outputData.items()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get photon pt distribution from coffea output\n",
    "# sum over lepton flavors (get both electron and muon)\n",
    "h = getHisto(\"photon_pt\")[{\"lepFlavor\": sum}]\n",
    "\n",
    "# integrate over systematics, selecting only \"nominal\"\n",
    "h = h[{\"systematic\": \"nominal\"}]\n",
    "# group the photon category axis into the category types listed above\n",
    "h = GroupBy(h, \"category\", \"category\", groupingCategory)\n",
    "# rebin the pt axis\n",
    "new_bins = np.array(\n",
    "    [\n",
    "        20,\n",
    "        25,\n",
    "        30,\n",
    "        35,\n",
    "        40,\n",
    "        45,\n",
    "        50,\n",
    "        60,\n",
    "        70,\n",
    "        80,\n",
    "        90,\n",
    "        100,\n",
    "        120,\n",
    "        140,\n",
    "        160,\n",
    "        180,\n",
    "        200,\n",
    "        250,\n",
    "        300,\n",
    "        400,\n",
    "        500,\n",
    "    ]\n",
    ")\n",
    "pt_axis = hist.axis.Variable(new_bins, name=\"pt\", label=r\"$p_{T}$ [GeV]\")\n",
    "h = RebinHist(h, pt=pt_axis)\n",
    "\n",
    "hData = outputDataHist[\"photon_pt\"][\n",
    "    {\"lepFlavor\": sum, \"systematic\": sum, \"category\": sum}\n",
    "]\n",
    "hData = RebinHist(hData, pt=pt_axis)\n",
    "\n",
    "plotWithRatio(\n",
    "    h[{\"dataset\": sum}],\n",
    "    hData,\n",
    "    overlay=\"category\",\n",
    "    xRange=[20, 500],\n",
    "    yRange=[0.01, None],\n",
    "    logY=True,\n",
    "    leg=\"right\",\n",
    "    binwnorm=1,\n",
    ")\n",
    "plotWithRatio(\n",
    "    h[{\"category\": sum}],\n",
    "    hData,\n",
    "    overlay=\"dataset\",\n",
    "    xRange=[20, 500],\n",
    "    yRange=[0.01, None],\n",
    "    logY=True,\n",
    "    leg=\"right\",\n",
    "    binwnorm=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Can you draw similar plots for the photon eta?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h = getHisto(\"photon_chIso\")[{\"lepFlavor\": sum, \"systematic\": \"nominal\"}]\n",
    "h = GroupBy(h, \"category\", \"category\", groupingCategory)\n",
    "\n",
    "s = hist.tag.Slicer()\n",
    "h = h[{\"chIso\": s[:: hist.rebin(5)]}]\n",
    "\n",
    "hData = outputDataHist[\"photon_chIso\"]\n",
    "hData = hData[{\"lepFlavor\": sum, \"systematic\": sum, \"category\": sum}]\n",
    "hData = hData[{\"chIso\": s[:: hist.rebin(5)]}]\n",
    "\n",
    "plotWithRatio(\n",
    "    h[{\"dataset\": sum}].sort(\"category\"),\n",
    "    hData,\n",
    "    overlay=\"category\",\n",
    "    xRange=[0, 10],\n",
    "    logY=True,\n",
    "    leg=\"right\",\n",
    ")\n",
    "plotWithRatio(\n",
    "    h[{\"category\": sum}],\n",
    "    hData,\n",
    "    overlay=\"dataset\",\n",
    "    xRange=[0, 10],\n",
    "    logY=True,\n",
    "    leg=\"right\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h = getHisto(\"M3\")[{\"lepFlavor\": \"muon\", \"systematic\": \"nominal\"}]\n",
    "h = GroupBy(h, \"category\", \"category\", groupingCategory)\n",
    "\n",
    "s = hist.tag.Slicer()\n",
    "h = h[{\"M3\": s[:: hist.rebin(4)]}]\n",
    "h = SetRangeHist(h, \"M3\", 50, 500)\n",
    "\n",
    "hData = outputDataHist[\"M3\"][{\"lepFlavor\": \"muon\", \"systematic\": sum, \"category\": sum}]\n",
    "hData = hData[{\"M3\": s[:: hist.rebin(4)]}]\n",
    "hData = SetRangeHist(hData, \"M3\", 50, 500)\n",
    "\n",
    "plotWithRatio(h[{\"dataset\": sum}], hData, overlay=\"category\", leg=\"right\")\n",
    "plotWithRatio(h[{\"category\": sum}], hData, overlay=\"dataset\", leg=\"right\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h = getHisto(\"M3\")[{\"lepFlavor\": sum, \"systematic\": \"nominal\"}]\n",
    "h = GroupBy(h, \"category\", \"category\", groupingCategory)\n",
    "\n",
    "s = hist.tag.Slicer()\n",
    "h = h[{\"M3\": s[:: hist.rebin(4)]}]\n",
    "h = SetRangeHist(h, \"M3\", 50, 500)\n",
    "\n",
    "hData = outputDataHist[\"M3\"][{\"lepFlavor\": sum, \"systematic\": sum, \"category\": sum}]\n",
    "hData = hData[{\"M3\": s[:: hist.rebin(4)]}]\n",
    "hData = SetRangeHist(hData, \"M3\", 50, 500)\n",
    "\n",
    "plotWithRatio(h[{\"dataset\": sum}], hData, overlay=\"category\", leg=\"right\")\n",
    "plotWithRatio(h[{\"category\": sum}], hData, overlay=\"dataset\", leg=\"right\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h = getHisto(\"M3\")[{\"lepFlavor\": sum, \"systematic\": \"nominal\"}]\n",
    "h = GroupBy(h, \"category\", \"category\", groupingFineCategory)\n",
    "\n",
    "s = hist.tag.Slicer()\n",
    "h = h[{\"M3\": s[:: hist.rebin(4)]}]\n",
    "h = SetRangeHist(h, \"M3\", 50, 500)\n",
    "\n",
    "hData = outputDataHist[\"M3\"][{\"lepFlavor\": sum, \"systematic\": sum, \"category\": sum}]\n",
    "hData = hData[{\"M3\": s[:: hist.rebin(4)]}]\n",
    "hData = SetRangeHist(hData, \"M3\", 50, 500)\n",
    "\n",
    "plotWithRatio(h[{\"dataset\": sum}], hData, overlay=\"category\", leg=\"right\")\n",
    "plotWithRatio(h[{\"category\": sum}], hData, overlay=\"dataset\", leg=\"right\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h = getHisto(\"photon_lepton_mass_3j0t\")[\n",
    "    {\"lepFlavor\": \"electron\", \"systematic\": \"nominal\"}\n",
    "]\n",
    "h = GroupBy(h, \"category\", \"category\", groupingCategory)\n",
    "\n",
    "s = hist.tag.Slicer()\n",
    "h = h[{\"mass\": s[:: hist.rebin(10)]}]\n",
    "h = SetRangeHist(h, \"mass\", 10, 200)\n",
    "\n",
    "hData = outputDataHist[\"photon_lepton_mass_3j0t\"][\n",
    "    {\"lepFlavor\": \"electron\", \"systematic\": sum, \"category\": sum}\n",
    "]\n",
    "hData = hData[{\"mass\": s[:: hist.rebin(10)]}]\n",
    "hData = SetRangeHist(hData, \"mass\", 10, 200)\n",
    "\n",
    "plotWithRatio(h[{\"dataset\": sum}], hData, overlay=\"category\", leg=\"right\")\n",
    "plotWithRatio(h[{\"category\": sum}], hData, overlay=\"dataset\", leg=\"right\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}