Z_2D.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#most python scripts start with 'import' statements like this that make\n",
    "#useful packages available. Examples are 'uproot' which reads in our data\n",
    "#files and allows them to be read, 'numpy' is a powefule and popular package\n",
    "#for fast manipultation of arrays and 'scip.stats' is useful for statitical\n",
    "#analysis, we'll use it to generate and fit fucntions to our data.\n",
    "import lorentz\n",
    "import uproot\n",
    "import uproot_methods.classes.TLorentzVector as LVepm\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy import stats\n",
    "from scipy.stats import cauchy\n",
    "from scipy.stats import expon\n",
    "from scipy.stats import norm\n",
    "\n",
    "from scipy.stats import crystalball\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2 µs, sys: 0 ns, total: 2 µs\n",
      "Wall time: 4.77 µs\n",
      "File has been successfully opened!\n",
      "muon_pt                    (no streamer)              asdtype(\"('>f4', (2,))\")\n",
      "muon_eta                   (no streamer)              asdtype(\"('>f4', (2,))\")\n",
      "muon_phi                   (no streamer)              asdtype(\"('>f4', (2,))\")\n",
      "muon_E                     (no streamer)              asdtype(\"('>f4', (2,))\")\n",
      "dimuon_mass                (no streamer)              asdtype('>f4')\n",
      "muon_etcone20              (no streamer)              asdtype(\"('>f8', (2,))\")\n",
      "muon_ptcone30              (no streamer)              asdtype(\"('>f8', (2,))\")\n"
     ]
    }
   ],
   "source": [
    "#the first thing we want to do is read our data file. The file is hosted on a CERN web-server so you\n",
    "#will see a full web address. As this file is large and we are reading over the web,\n",
    "#this step can take a few minutes\n",
    "#%timeit eventsData = uproot.open(\"https://atlas-opendata.web.cern.ch/atlas-opendata/samples/2020/2lep/Data/data_D.2lep.root\")[\"mini\"]\n",
    "\n",
    "%time\n",
    "eventsData = uproot.open(\"data_Skim_mumu.root\")[\"mini\"]\n",
    "\n",
    "print(\"File has been successfully opened!\")\n",
    "eventsData.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "muon_pt                    (no streamer)              asdtype(\"('>f4', (2,))\")\n",
      "muon_eta                   (no streamer)              asdtype(\"('>f4', (2,))\")\n",
      "muon_phi                   (no streamer)              asdtype(\"('>f4', (2,))\")\n",
      "muon_E                     (no streamer)              asdtype(\"('>f4', (2,))\")\n",
      "dimuon_mass                (no streamer)              asdtype('>f4')\n",
      "muon_etcone20              (no streamer)              asdtype(\"('>f8', (2,))\")\n",
      "muon_ptcone30              (no streamer)              asdtype(\"('>f8', (2,))\")\n"
     ]
    }
   ],
   "source": [
    "#the file is now represented by the variable \"eventsData\".\n",
    "#we can have a peek inside the file to see what kind of information is available\n",
    "#by running the following command\n",
    "\n",
    "eventsData.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "nEvents = 200000 \n",
    "#this is the total number of collisions/events we want to access from the file, you can reduce this\n",
    "#for quick test, for final results 100000 events should be enough\n",
    "\n",
    "muon_E, muon_pt, muon_phi, muon_eta, muon_etcone20, dimuon_mass, muon_ptcone30 = eventsData.arrays([\"muon_E\", \"muon_pt\", \"muon_phi\", \"muon_eta\", \"muon_etcone20\", \"dimuon_mass\", \"muon_ptcone30\"], outputtype=tuple, entrystop=nEvents)\n",
    "flatData = LVepm.TLorentzVectorArray.from_ptetaphi(muon_pt, muon_eta, muon_phi, muon_E)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "only size-1 arrays can be converted to Python scalars",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-8-32af2400c7c2>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     14\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mev\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mflatData\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     15\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0;34m(\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mmuon_ptcone30\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mev\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m/\u001b[0m\u001b[0mmuon_pt\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mev\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0;36m0.1\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m&\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mmuon_ptcone30\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mev\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m/\u001b[0m\u001b[0mmuon_pt\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mev\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0;36m0.1\u001b[0m \u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 16\u001b[0;31m             \u001b[0mmumu\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mflatData\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mev\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mflatData\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mev\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     17\u001b[0m             \u001b[0mmass\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmumu\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmass\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     18\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m \u001b[0mmass\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0mminMass\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m&\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mmass\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0mmaxMass\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.virtualenvs/vpyenv/lib/python3.7/site-packages/awkward/array/objects.py\u001b[0m in \u001b[0;36m__getitem__\u001b[0;34m(self, where)\u001b[0m\n\u001b[1;32m    191\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_util_isinteger\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhead\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    192\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtail\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtuple\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mtail\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 193\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgenerator\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcontent\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    194\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    195\u001b[0m                 \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgenerator\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcontent\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mtail\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.virtualenvs/vpyenv/lib/python3.7/site-packages/uproot_methods/classes/TLorentzVector.py\u001b[0m in \u001b[0;36m<lambda>\u001b[0;34m(row)\u001b[0m\n\u001b[1;32m    124\u001b[0m \u001b[0;32mclass\u001b[0m \u001b[0mArrayMethods\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mCommon\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0muproot_methods\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbase\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mROOTMethods\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    125\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_initObjectArray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtable\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 126\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mawkward\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mObjectArray\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__init__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtable\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mlambda\u001b[0m \u001b[0mrow\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mTLorentzVector\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrow\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"fX\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrow\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"fY\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrow\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"fZ\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrow\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"fE\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    127\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    128\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__awkward_serialize__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mserializer\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.virtualenvs/vpyenv/lib/python3.7/site-packages/uproot_methods/classes/TLorentzVector.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, x, y, z, t)\u001b[0m\n\u001b[1;32m    938\u001b[0m \u001b[0;32mclass\u001b[0m \u001b[0mTLorentzVector\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mMethods\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    939\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__init__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mz\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 940\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_fP\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0muproot_methods\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclasses\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTVector3\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTVector3\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfloat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfloat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfloat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mz\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    941\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_fE\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfloat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mt\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    942\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: only size-1 arrays can be converted to Python scalars"
     ]
    }
   ],
   "source": [
    "#create a few arrays to hold some specific information about the collisions\n",
    "massesData = np.array([])\n",
    "muPtData = np.array([])\n",
    "muEtaData = np.array([])\n",
    "muPhiData = np.array([])\n",
    "\n",
    "nData = 0\n",
    "nTarget = 200000\n",
    "\n",
    "minMass = 71000\n",
    "maxMass = 110000\n",
    "    \n",
    "#loop over events and extract info from the array\n",
    "for ev in range(0,len(flatData)):\n",
    "        if ( (muon_ptcone30[ev][0]/muon_pt[ev][0] < 0.1) & (muon_ptcone30[ev][1]/muon_pt[ev][1] < 0.1 )):\n",
    "            mumu = flatData[ev][0] + flatData[ev][1] \n",
    "            mass = mumu.mass\n",
    "            if (( mass > minMass) & (mass < maxMass)):\n",
    "                massesData = np.append(massesData,mass)\n",
    "                muPtData = np.append(muPtData,flatData[ev][0].pt)\n",
    "                muPtData = np.append(muPtData,flatData[ev][1].pt)\n",
    "                muEtaData = np.append(muEtaData,flatData[ev][0].eta)\n",
    "                muEtaData = np.append(muEtaData,flatData[ev][1].eta)\n",
    "                muPhiData = np.append(muPhiData,flatData[ev][0].phi)\n",
    "                muPhiData = np.append(muPhiData,flatData[ev][1].phi)\n",
    "                nData = nData + 1\n",
    "            if (nData == nTarget):\n",
    "                break\n",
    "        \n",
    "print(\"successfully analysed \" +str(nData) + \" events\")\n",
    "\n",
    "# plot the mu pt histogram.\n",
    "plt.figure()\n",
    "plt.xlabel(\"mu pt\")\n",
    "plt.ylabel(\"events per bin\")\n",
    "plt.hist(muPtData, bins=35, range=[0,120000], alpha=0.6, color='g')\n",
    "\n",
    "# plot the mu pseudorapdity histogram.\n",
    "plt.figure()\n",
    "plt.xlabel(\"mu eta\")\n",
    "plt.ylabel(\"events per bin\")\n",
    "plt.hist(muEtaData, bins=25, range=[-3.0,3.0], alpha=0.6, color='g')\n",
    "\n",
    "# plot the mu phi histogram.\n",
    "plt.figure()\n",
    "plt.xlabel(\"mu phi\")\n",
    "plt.ylabel(\"events per bin\")\n",
    "plt.hist(muPhiData, bins=15, range=[-3.14,3.14], alpha=0.6, color='g')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nBins =50 \n",
    "\n",
    "mZ  = 91000.0\n",
    "sigmaZ = 2100\n",
    "\n",
    "#redefine m(mumu) as a numpy histogram\n",
    "countsData, edges = np.histogram(massesData, bins=nBins, range=(minMass, maxMass))\n",
    "centres = (edges[1:] + edges[:-1]) / 2\n",
    "\n",
    "# get the width of each bin\n",
    "bin_width = edges[1] - edges[0]\n",
    "# sum over number in each bin and mult by bin width, which can be factored out\n",
    "integral = bin_width * sum(countsData[0:nBins])\n",
    "\n",
    "print(integral)\n",
    "\n",
    "#fit a Cauchy distributions to the dimuon mass histogram\n",
    "#massesData = massesData[ (massesData >= 60000) & (massesData <= 120000) ]\n",
    "mu, std = cauchy.fit(massesData)\n",
    "#beta, m, loc, scale = crystalball.fit(massesData)\n",
    "#mu, std = norm.fit(massesData)\n",
    "\n",
    "print(\"mu\" + str(mu))\n",
    "print(\"std\" + str(std))\n",
    "\n",
    "plt.figure()\n",
    "p = (cauchy.pdf(centres, mu, std) * integral)\n",
    "#p = (crystalball.pdf(centres, beta=beta, m=m, loc=loc, scale=scale) * integral)\n",
    "#p = (crystalball.pdf(x, beta, m, scale=scale, loc=loc) * integral)\n",
    "\n",
    "#p = (norm.pdf(centres, mu, std) * integral)\n",
    "\n",
    "plt.plot(centres, p, 'k', linewidth=2)\n",
    "plt.hist(massesData, bins=nBins, range=[minMass, maxMass], alpha=0.6, color='g')\n",
    "plt.xlabel(\"mumu mass\")\n",
    "plt.ylabel(\"events per bin\")\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#chi-squared function\n",
    "def calcChiSq(obs, preds):\n",
    "    chiSq = 0.0\n",
    "    ndf = len(obs)\n",
    "    for bin in range(0, len(obs)):\n",
    "        diff = preds[bin] - obs[bin]\n",
    "        #print(\"diff = \" + str(diff))\n",
    "        var = ( np.abs(preds[bin])) \n",
    "        if (var != 0):\n",
    "            chiSq += (diff**2)/(var)\n",
    "            #print(\"obs, pred, diff, var  chi contrib = \"+  str( obs[bin]) +\" \" +  str(preds[bin]) +\" \"+  str(diff) +\" \"+  str(var) +\" \" + str((diff**2)/(var)))\n",
    "    return chiSq, ndf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mZAr = np.array([])\n",
    "chi2Ar = np.array([])\n",
    "\n",
    "minMz = 90150\n",
    "maxMz = 90600\n",
    "step = 1\n",
    "\n",
    "bestFitMz1D = 0.0\n",
    "minChi2 = 10000000\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "for mZ in range(minMz, maxMz, step):\n",
    "    countsPDF = (cauchy.pdf(centres, mZ, sigmaZ) * integral)\n",
    "    chi2, ndf = calcChiSq(countsData, countsPDF)\n",
    "    mZAr = np.append(mZAr, mZ)\n",
    "    chi2Ar = np.append(chi2Ar, chi2)\n",
    "    if(chi2 < minChi2):\n",
    "        minChi2 = chi2\n",
    "        bestFitMz1D = mZ\n",
    "\n",
    "chi2Ar = chi2Ar - np.min(chi2Ar)\n",
    "    \n",
    "ax.plot(mZAr, chi2Ar, 'r+', lw=5, alpha=0.6, label=\"chi2 evaluations\")\n",
    "\n",
    "z = np.polyfit(mZAr, chi2Ar, 2)\n",
    "p = np.poly1d(z)\n",
    "\n",
    "ax.plot(mZAr, p(mZAr), 'b-', lw=5, alpha=0.6, label=\"parabolic fit\")\n",
    "\n",
    "# 'statitical uncertianty' on our mZ measurement via critical values of chi-squared curve\n",
    "y0 = 1.0\n",
    "crit = (p - y0).roots\n",
    "\n",
    "px=np.arange(crit[1],crit[0],0.001)\n",
    "\n",
    "ax.fill_between(px,p(px),alpha=0.5, color='g', label=\"uncertainty\")\n",
    "ax.legend()\n",
    "\n",
    "ax.set_xlim(minMz, maxMz)\n",
    "ax.set_ylim(0.0, 5.0)\n",
    "\n",
    "plt.xlabel(\"mZ\")\n",
    "plt.ylabel(\"delta chi-squared\")\n",
    "plt.show()\n",
    "\n",
    "#What is causing the fluctuations around the expected parabolic shape?\n",
    "#How could we smooth out this parabola?\n",
    "#Why are the chi-quared values so small? Could our fit be that good? No.\n",
    "\n",
    "#extract 1D result and uncertainty\n",
    "my1DResult = bestFitMz1D\n",
    "my1DUncertainty = bestFitMz1D-crit[1]\n",
    "print(\"best-fit = \" + str(my1DResult) + \" +/- \" + str(my1DUncertainty))\n",
    "\n",
    "\n",
    "plt.figure()\n",
    "p = (cauchy.pdf(centres, my1DResult, sigmaZ))\n",
    "\n",
    "pInt = bin_width * sum(p[0:nBins])\n",
    "scale = integral/pInt\n",
    "pScaled = (cauchy.pdf(centres, my1DResult, sigmaZ) *scale)\n",
    "\n",
    "plt.plot(centres, pScaled, 'k', linewidth=2)\n",
    "plt.hist(massesData, bins=nBins, range=[minMass, maxMass], alpha=0.6, color='g')\n",
    "plt.xlabel(\"mumu mass\")\n",
    "plt.ylabel(\"events per bin\")\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#How about fitting the mass and the width at the same time?\n",
    "#minMz = 90620\n",
    "#maxMz = 90730\n",
    "\n",
    "minSigmaZ = 2140\n",
    "maxSigmaZ = 2560\n",
    "step = 5\n",
    "\n",
    "mZAr = np.arange(minMz, maxMz, step)\n",
    "sigmaZAr = np.arange(minSigmaZ, maxSigmaZ, step)\n",
    "chi2Ar = np.empty( (len(sigmaZAr), len(mZAr) ))\n",
    "\n",
    "X, Y = np.meshgrid(mZAr, sigmaZAr)\n",
    "\n",
    "\n",
    "bestFitMz2D = 0.0\n",
    "minChi2 = 1000000\n",
    "\n",
    "for yi in range(0, len(sigmaZAr)):\n",
    "    for xi in range(0, len(mZAr)):\n",
    "        #print( str(X[xi,yi]) + \",\" + str(Y[xi,yi]) )\n",
    "        #print(chi2AtPoint(X[xi,yi],Y[xi,yi]))\n",
    "        countsPDF = (cauchy.pdf(centres, X[yi, xi],Y[yi, xi]) * integral)\n",
    "        chi2, ndf = calcChiSq(countsData, countsPDF)\n",
    "        chi2Ar[yi, xi] = chi2\n",
    "        if(chi2 < minChi2):\n",
    "            #print(\"chi2 = \" + str(chi2) + \" mz = \" + str(X[yi, xi]))\n",
    "            minChi2 = chi2\n",
    "            bestFitMz2D = X[yi, xi]\n",
    "\n",
    "print(bestFitMz2D)\n",
    "minChi2 = np.min(chi2Ar)\n",
    "deltaChi2Ar = chi2Ar - minChi2\n",
    "\n",
    "#plot \n",
    "levels = [2.3]\n",
    "fig = plt.figure(figsize=(6,6))\n",
    "ax = fig.gca()\n",
    "\n",
    "cset = plt.contour(mZAr,sigmaZAr,deltaChi2Ar, levels=levels, colors=['white'])\n",
    "cfset = plt.contourf(mZAr,sigmaZAr,deltaChi2Ar,  cmap='coolwarm')\n",
    "\n",
    "h = plt.contourf(mZAr,sigmaZAr,deltaChi2Ar)\n",
    "\n",
    "cbar = plt.colorbar(h)\n",
    "cbar.set_label('Delta Chi-squared', rotation=270)\n",
    "ax.clabel(cset, inline=1, fontsize=10)\n",
    "ax.set_xlabel('mZ')\n",
    "ax.set_ylabel('sigmaZ')\n",
    "plt.show()\n",
    "\n",
    "#extract 2D result and uncertainty\n",
    "my2DResult = bestFitMz2D\n",
    "my2DUncertainty = bestFitMz2D-crit[1]\n",
    "print(\"best-fit = \" + str(my2DResult) + \" +/- \" + str(my2DUncertainty))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#compare to PDG world average\n",
    "mZPDG =  91187.6\n",
    "sigmaZPDG = 2.1\n",
    "\n",
    "mZPDGAr = np.array([mZPDG]) \n",
    "sigZPDGAr = np.array([sigmaZPDG]) \n",
    "yPDGAr = np.array([1.0]) \n",
    "\n",
    "mZ1D = np.array([my1DResult]) \n",
    "sigZ1D = np.array([my1DUncertainty]) \n",
    "y1D = np.array([2.0]) \n",
    "\n",
    "mZ2D = np.array([my2DResult]) \n",
    "sigZ2D = np.array([my2DUncertainty]) \n",
    "y2D = np.array([3.0]) \n",
    "\n",
    "plt.figure()\n",
    "\n",
    "ax = plt.gca()\n",
    "ax.set_xlim(minMz, maxMz+700)#need to extend x axis range to include PDG value\n",
    "ax.set_ylim(0.0, 4.0)\n",
    "\n",
    "plt.errorbar(mZPDGAr, yPDGAr, xerr=sigZPDGAr, label=\"Particle Data Group\", elinewidth=5,  fmt=\"o\", mfc=\"blue\", ms=8)\n",
    "plt.errorbar(mZ1D, y1D, xerr=sigZ1D, fmt=\"o\", label=\"ATLAS Open Data 1D\", elinewidth=5, ms=8, mfc='blue')\n",
    "plt.errorbar(mZ2D, y2D, xerr=sigZ2D, fmt=\"o\", label=\"ATLAS Open Data 2D\", elinewidth=5, ms=8, mfc='blue')\n",
    "\n",
    "plt.legend()\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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}