combination_resolved_FD.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "382dc870",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Welcome to JupyROOT 6.28/04\n"
     ]
    }
   ],
   "source": [
    "import ROOT\n",
    "ROOT.EnableImplicitMT()\n",
    "RDF = ROOT.ROOT.RDataFrame\n",
    "from matplotlib import pyplot as plt\n",
    "import mplhep as hep\n",
    "plt.style.use(hep.style.ROOT)\n",
    "import glob, os\n",
    "import numpy as np\n",
    "from sklearn.gaussian_process import GaussianProcessRegressor\n",
    "from sklearn.gaussian_process   .kernels import RBF, ConstantKernel as C,  DotProduct, Matern\n",
    "import matplotlib.gridspec as gridspec\n",
    "from gp_signal import GPnew, Signal as ESC\n",
    "cols = ['#4285f4','#ea4335','#fbbc05','#34a853', '#a00498', '#536267']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e7e3cb3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "tag=''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7f8cf1dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# with open('root_files.txt') as f:\n",
    "#     lines = f.readlines()\n",
    "\n",
    "lines = glob.glob('/uscms/home/abhijith/nobackup/trijet_boost/Analysis/slim_out_trip3/*root')\n",
    "# lines = glob.glob('test/*root')\n",
    "    \n",
    "data2018 = ROOT.TChain('events')\n",
    "data = ROOT.TChain('events')\n",
    "data2016 = ROOT.TChain('events')\n",
    "ttbar = ROOT.TChain('events')\n",
    "data2017 = ROOT.TChain('events')\n",
    "\n",
    "for file in lines:\n",
    "    # print(file)\n",
    "    if ('trip4' in file and '2018' in file):\n",
    "        data2018.Add(file)\n",
    "\n",
    "    if ('trip4' in file and '2017' in file):\n",
    "        data.Add(file)\n",
    "\n",
    "    if ('trip4' in file and '2016' in file):\n",
    "        data2016.Add(file)\n",
    "\n",
    "    if ('trip4' in file and 'TT' in file):\n",
    "        ttbar.Add(file)\n",
    "    \n",
    "    data2017.Add(file)\n",
    "\n",
    "df = {}\n",
    "# df['2018'] = RDF(data2018)\n",
    "df['2017'] = RDF(data)\n",
    "# df['2017'] = RDF(data2017)\n",
    "# df['2016'] = RDF(data2016)\n",
    "# df['ttbar'] = RDF(ttbar)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "01889962",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['2017'])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df['2017'].GetColumnNames()\n",
    "df.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7c7295ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "regions=[]\n",
    "regions.append({'name':'region_1','cuts':'ak4_HT>550 && sixpt>40 && jmds6332<1.25 && jtrip_masym <0.25 && jtrip_delta>250 && jtrip_mds<0.05 && jtrip_qgl>1 && jtrip_qgt>0','bins':[102*2,100,610]}) #\n",
    "regions.append({'name':'region_2','cuts':'ak4_HT>550 && sixpt>50 && jmds6332<1.25 && jtrip_masym <0.175 && jtrip_delta>180 && jtrip_mds<0.175 && jtrip_qgl>1 && jtrip_qgt>0','bins':[90*2,300,1200]}) #\n",
    "regions.append({'name':'region_3','cuts':'ak4_HT>550 && sixpt>100 && jmds6332<0.9 && jtrip_masym <0.15 && jtrip_delta>150 && jtrip_mds<0.2 && jtrip_qgl>1 && jtrip_qgt>0','bins':[99*2,520,2500]}) #\n",
    "# regions.append({'name':'run2_invqgl','cuts':'ak4_HT>550 && sixpt>100 && jmds6332<0.9 && jtrip_masym <0.15 && jtrip_delta>150 && jtrip_mds<0.2 && jtrip_qgl>1 && jtrip_qgt>0','bins':[99,520,2500]})\n",
    "\n",
    "mass_hists={}\n",
    "\n",
    "jtrip_shift = {}\n",
    "jtrip_shift['2016'] = '(18/16.5)*jtrip_mass-18.55'\n",
    "jtrip_shift['2017'] = 'jtrip_mass'\n",
    "jtrip_shift['ttbar'] = 'jtrip_mass'\n",
    "jtrip_shift['2018'] = '1.276*jtrip_mass-31.42'\n",
    "\n",
    "for reg in regions:\n",
    "    mass_hists[reg['name']]={}\n",
    "    for years in df.keys():\n",
    "        mass_hists[reg['name']][years] = df[years].Define('corr_mass_'+years,jtrip_shift[years]).Filter(reg['cuts']).Histo1D(('mass_hists_'+reg['name']+'_'+years,'mass_hists_'+reg['name']+'_'+years,reg['bins'][0],reg['bins'][1],reg['bins'][2]),'corr_mass_'+years)\n",
    "\n",
    "for reg in regions:\n",
    "    mass_hists[reg['name']]['data'] = mass_hists[reg['name']]['2017'].GetValue().Clone()\n",
    "    #mass_hists[reg['name']]['data'].Add(mass_hists[reg['name']]['2017'].GetValue())\n",
    "    #mass_hists[reg['name']]['data'].Add(mass_hists[reg['name']]['2018'].GetValue())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "420ca35f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'str'>\n",
      "<class 'str'>\n",
      "<class 'str'>\n"
     ]
    }
   ],
   "source": [
    "regions=[]\n",
    "regions.append({'name':'region_1','cuts':'ak4_HT>550 && sixpt>40 && jmds6332<1.25 && jtrip_masym <0.25 && jtrip_delta>250 && jtrip_mds<0.05 && jtrip_qgl>1 && jtrip_qgt>0','bins':[102*2,100,610]}) #\n",
    "regions.append({'name':'region_2','cuts':'ak4_HT>550 && sixpt>50 && jmds6332<1.25 && jtrip_masym <0.175 && jtrip_delta>180 && jtrip_mds<0.175 && jtrip_qgl>1 && jtrip_qgt>0','bins':[90*2,300,1200]}) #\n",
    "regions.append({'name':'region_3','cuts':'ak4_HT>550 && sixpt>100 && jmds6332<0.9 && jtrip_masym <0.15 && jtrip_delta>150 && jtrip_mds<0.2 && jtrip_qgl>1 && jtrip_qgt>0','bins':[99*2,520,2500]}) #\n",
    "# regions.append({'name':'run2_invqgl','cuts':'ak4_HT>550 && sixpt>100 && jmds6332<0.9 && jtrip_masym <0.15 && jtrip_delta>150 && jtrip_mds<0.2 && jtrip_qgl>1 && jtrip_qgt>0','bins':[99,520,2500]})\n",
    "\n",
    "mass_hists={}\n",
    "\n",
    "jtrip_shift = {}\n",
    "jtrip_shift['2016'] = '(18/16.5)*jtrip_mass-18.55'\n",
    "jtrip_shift['2017'] = 'jtrip_mass'\n",
    "jtrip_shift['ttbar'] = 'jtrip_mass'\n",
    "jtrip_shift['2018'] = '1.276*jtrip_mass-31.42'\n",
    "\n",
    "for reg in regions:\n",
    "    mass_hists[reg['name']]={}\n",
    "    for years in df.keys():\n",
    "        print(type(years))\n",
    "        mass_hists[reg['name']][years] = df[years].Define('corr_mass_'+years,jtrip_shift[years]).Filter(reg['cuts']).Histo1D(('mass_hists_'+reg['name']+'_'+years,'mass_hists_'+reg['name']+'_'+years,reg['bins'][0],reg['bins'][1],reg['bins'][2]),'corr_mass_'+years)\n",
    "\n",
    "#mass_hists['region_1']['data'].Add(mass_hists['region_1']['2017'].GetValue())\n",
    "for reg in regions:\n",
    "    mass_hists[reg['name']]['data'] = mass_hists[reg['name']]['2017'].GetValue().Clone()\n",
    "    #mass_hists[reg['name']]['data'].Add(mass_hists[reg['name']]['2017'].GetValue())\n",
    "    #mass_hists[reg['name']]['data'].Add(mass_hists[reg['name']]['2018'].GetValue())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "79b0b37e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<cppyy.gbl.ROOT.RDF.RResultPtr<TH1D> object at 0x55ebc98e7500>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mass_hists['region_1']['2017']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f82af080",
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyError",
     "evalue": "'2018'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[8], line 10\u001b[0m\n\u001b[1;32m      8\u001b[0m mass_hists[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mregion_1\u001b[39m\u001b[38;5;124m'\u001b[39m][\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m2017\u001b[39m\u001b[38;5;124m'\u001b[39m]\u001b[38;5;241m.\u001b[39mRebin()\n\u001b[1;32m      9\u001b[0m mass_hists[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mregion_1\u001b[39m\u001b[38;5;124m'\u001b[39m][\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m2017\u001b[39m\u001b[38;5;124m'\u001b[39m]\u001b[38;5;241m.\u001b[39mDraw(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124msame\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[0;32m---> 10\u001b[0m \u001b[43mmass_hists\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mregion_1\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43m2018\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241m.\u001b[39mSetLineColor(\u001b[38;5;241m46\u001b[39m)\n\u001b[1;32m     11\u001b[0m mass_hists[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mregion_1\u001b[39m\u001b[38;5;124m'\u001b[39m][\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m2018\u001b[39m\u001b[38;5;124m'\u001b[39m]\u001b[38;5;241m.\u001b[39mSetLineWidth(\u001b[38;5;241m2\u001b[39m)\n\u001b[1;32m     12\u001b[0m mass_hists[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mregion_1\u001b[39m\u001b[38;5;124m'\u001b[39m][\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m2018\u001b[39m\u001b[38;5;124m'\u001b[39m]\u001b[38;5;241m.\u001b[39mRebin()\n",
      "\u001b[0;31mKeyError\u001b[0m: '2018'"
     ]
    }
   ],
   "source": [
    "c1=ROOT.TCanvas()\n",
    "c1.cd()\n",
    "# regions[0]['h_a_nq'].SetLineColor(38)\n",
    "# regions[0]['h_a_nq'].Draw()\n",
    "# regions[0]['h_a_nq_cor'].Draw('same')\n",
    "mass_hists['region_1']['2017'].SetLineColor(30)\n",
    "mass_hists['region_1']['2017'].SetLineWidth(2)\n",
    "mass_hists['region_1']['2017'].Rebin()\n",
    "mass_hists['region_1']['2017'].Draw('same')\n",
    "mass_hists['region_1']['2018'].SetLineColor(46)\n",
    "mass_hists['region_1']['2018'].SetLineWidth(2)\n",
    "mass_hists['region_1']['2018'].Rebin()\n",
    "mass_hists['region_1']['2018'].Draw('same')\n",
    "mass_hists['region_1']['2016'].SetLineColor(38)\n",
    "mass_hists['region_1']['2016'].SetLineWidth(2)\n",
    "mass_hists['region_1']['2016'].Rebin()\n",
    "mass_hists['region_1']['2016'].Draw('same')\n",
    "# regions[0]['h_a_nq_ttbar'].Draw('same')\n",
    "c1.Draw()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b545d78",
   "metadata": {},
   "outputs": [],
   "source": [
    "file = ROOT.TFile('resolved_FDhist.root',\"RECREATE\")                                \n",
    "file.cd()\n",
    "for i in range(len(regions)):\n",
    "    regions[i]['h_a_nq_2017'].Write()\n",
    "    regions[i]['h_a_nq_2018'].Write()\n",
    "    regions[i]['h_a_nq_2016'].Write()\n",
    "file.Close()   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6ed6381",
   "metadata": {},
   "outputs": [],
   "source": [
    "for plot in ['h_a_nq_2018','h_a_nq_2016','h_a_nq_2017']:\n",
    "    x=[]\n",
    "    y=[]\n",
    "    dy=[]\n",
    "    xc=[]\n",
    "    yc=[]\n",
    "    dyc=[]\n",
    "    h = regions[0][plot]\n",
    "    sig_m=170\n",
    "    sig_s=15\n",
    "    sig_m_bounds=(150,190)\n",
    "    sig_s_bounds=(7, 20)\n",
    "    for i in range(1,h.GetNbinsX()+1):\n",
    "        if(h.GetBinContent(i)):\n",
    "            x.append(h.GetBinCenter(i))\n",
    "            y.append(h.GetBinContent(i))\n",
    "            dy.append(h.GetBinError(i))\n",
    "            if((h.GetBinCenter(i)<130 or h.GetBinCenter(i)>200)):\n",
    "                xc.append(h.GetBinCenter(i))\n",
    "                yc.append(h.GetBinContent(i))\n",
    "                dyc.append(h.GetBinError(i))\n",
    "    print(len(x),len(y),len(dy),len(xc),len(yc),len(dyc))\n",
    "\n",
    "\n",
    "    X = np.atleast_2d(x).T\n",
    "    Y = np.atleast_2d(y).T\n",
    "    dy = np.array(dy)\n",
    "\n",
    "    Xc = np.atleast_2d(xc).T\n",
    "    Yc = np.atleast_2d(yc).T\n",
    "    dyc = np.array(dyc)\n",
    "    # kernel = C(0.1, (0.01, 1000.0)) * Matern(length_scale=40.0, length_scale_bounds=(40, 100.0), nu=11.5)\n",
    "    # kernel = kernel + \n",
    "    # kernel = C(1.0, (1e-3, 1e3)) * RBF(30, (30, 1e3))\n",
    "\n",
    "    kernel = C(1.0, (1e-3, 1e16)) * RBF(200, (10, 1e16))\n",
    "    gp = GaussianProcessRegressor(kernel=kernel, alpha=dyc**2, n_restarts_optimizer=20)\n",
    "    gp.fit(Xc, Yc)\n",
    "    print(gp.kernel_)\n",
    "    print('Using Lenght scale', float(str(gp.kernel_).split('=')[-1].split(')')[0]))\n",
    "\n",
    "    ls  = float(str(gp.kernel_).split('=')[-1].split(')')[0])\n",
    "\n",
    "    kernel = C(1.0, (1e-3, 1e16)) * RBF(ls, (ls, 500)) + C(1.0, (1e-3, 1e16)) * ESC(mass = sig_m, mass_bounds = sig_m_bounds , width = sig_s, width_bounds = sig_s_bounds)\n",
    "    # Fit for unmasked Data\n",
    "    print(X.shape, Y.shape, dy.shape)\n",
    "    gp = GPnew(kernel=kernel, alpha=dy**2, n_restarts_optimizer=30)\n",
    "    gp.fit(X, Y)\n",
    "    XC = np.atleast_2d(np.linspace(x[0], x[-1], num=1000)).T\n",
    "    Y_pred, sigma = gp.predict(X, return_std=True)\n",
    "    YC_pred, sigmaC = gp.predict(XC, return_std=True)\n",
    "    print(gp.kernel_)\n",
    "    YC_sig, sigmaC_sig = gp.predict_partial(XC,return_std=True,kx = 'k2')\n",
    "    YC_bkg, sigmaC_bkg = gp.predict_partial(XC,return_std=True,kx = 'k1')\n",
    "    Y_bkg, sigma_bkg = gp.predict_partial(X,return_std=True,kx = 'k1')\n",
    "    # kernel = C(0.1, (0.01, 1000.0)) * Matern(length_scale=100.0, length_scale_bounds=(100, 2000.0), nu=10.5)\n",
    "\n",
    "    kernel = C(1.0, (1e-3, 1e16)) * RBF(200, (135, 1e3))\n",
    "    #Fit for masked Data\n",
    "    gp = GaussianProcessRegressor(kernel=kernel, alpha=dyc**2, n_restarts_optimizer=20)\n",
    "    gp.fit(Xc, Yc)\n",
    "    XCc = np.atleast_2d(np.linspace(xc[0], xc[-1], num=1000)).T\n",
    "    Yc_pred, sigmac = gp.predict(Xc, return_std=True)\n",
    "    YCc_pred, sigmaCc = gp.predict(XCc, return_std=True)\n",
    "    Y_p, sig = gp.predict(X, return_std=True)\n",
    "\n",
    "    # with plt.xkcd():\n",
    "    fig = plt.figure(figsize=(10, 13))\n",
    "    spec = fig.add_gridspec(ncols=1, nrows=3,height_ratios=[0.8,0.2,0.3])\n",
    "    main = fig.add_subplot(spec[0,0])\n",
    "    main.plot(XCc.flatten(), YCc_pred.flatten(), '--', color='gainsboro', alpha=0.5)\n",
    "    # main.fill_between(XCc.flatten(),YCc_pred.flatten()-1*sigmaCc, YCc_pred.flatten()+1*sigmaCc, alpha=.5, fc='navy', ec='None', label='95% confidence interval')\n",
    "    main.text(0.0, 1.05,'CMS Prelimininary ', ha='left', va='top', transform=main.transAxes, fontsize=15)\n",
    "    main.text(0.99, 1.05,r'2018', ha='right', va='top', transform=main.transAxes, fontsize=15)\n",
    "    main.plot(XC.flatten(), YC_pred.flatten(), 'b-',color=cols[1], label='Total Prediction')\n",
    "    main.fill_between(XC.flatten(),YC_pred.flatten()-1*sigmaC, YC_pred.flatten()+1*sigmaC, alpha=.5, fc='k', ec='None',color=cols[1], label='Fit error')\n",
    "    main.errorbar(X.flatten(), Y.flatten(), dy, fmt='.', color='k',markersize=8,alpha=1.0, label='Data')\n",
    "    # main.errorbar(Xc.flatten(), Yc.flatten(), dyc, fmt='.',color='k', markersize=8,alpha=.8, label='Masked Data')\n",
    "    main.set_ylabel('Events', fontsize=15)\n",
    "    main.axvline(x=130, color='lightgrey', linestyle='--')\n",
    "    main.axvline(x=200, color='lightgrey', linestyle='--')\n",
    "    main.set_xlabel('Avg. Mass', fontsize=15)\n",
    "    main.tick_params(axis='y', labelsize=15)\n",
    "    main.tick_params(axis='x', labelsize=15)\n",
    "    main.plot(XC.flatten(),YC_sig.flatten(),color=cols[3],label='Gaussian')\n",
    "    main.plot(XC.flatten(),YC_bkg.flatten(),color=cols[0],label='Background Prediction')\n",
    "    main.fill_between(XC.flatten(),YC_sig.flatten()-sigmaC_sig.flatten(),YC_sig.flatten()+sigmaC_sig.flatten(),color=cols[3],alpha=.5,label='Fit error')\n",
    "    main.fill_between(XC.flatten(),YC_bkg.flatten()-sigmaC_bkg.flatten(),YC_bkg.flatten()+sigmaC_bkg.flatten(),color=cols[0],alpha=.5,label='Fit error')\n",
    "    main.legend(loc='best', fontsize=15)\n",
    "    # main.set_yscale('log')\n",
    "    # main.set_ylim(0.9,1000)\n",
    "    # main.plot(XCc.flatten(), YCc_pred.flatten()-YC_pred.flatten(), 'm-', label='Difference B/n predictions')\n",
    "    # main.fill_between(XCc.flatten(), (YCc_pred.flatten()-YC_pred.flatten())-1*sigmaCc, (YCc_pred.flatten()-YC_pred.flatten())+1*sigmaCc,alpha=.5, fc='m' )\n",
    "    pull = fig.add_subplot(spec[1,0])\n",
    "    pull.bar(X.flatten(),((Y.flatten()-Y_p.flatten())/dy),width=h.GetBinWidth(1),color='r', alpha=1)\n",
    "    pull.bar(X.flatten(),((Y.flatten()-Y_pred.flatten())/dy),width=h.GetBinWidth(1),color='k', alpha=0.7)\n",
    "    pull.set_ylabel('Pull',fontsize=15)\n",
    "    pull.axvline(x=130, color='lightgrey', linestyle='--')\n",
    "    pull.axvline(x=200, color='lightgrey', linestyle='--')\n",
    "    print(np.sum(((Y.flatten()-Y_pred.flatten())/dy)**2)/(len(x)-1))\n",
    "    print(np.sum(((Y.flatten()-Y_p.flatten())/dy)**2)/(len(x)-1))\n",
    "    print(np.sum(((Yc.flatten()-Yc_pred.flatten())/dyc)**2)/(len(x)-1))\n",
    "    diff = fig.add_subplot(spec[2,0])\n",
    "    diff.plot(XCc.flatten(), YC_pred.flatten()-YC_bkg.flatten(), 'm-', label='Gaussian')\n",
    "    # diff.fill_between(XCc.flatten(), (YC_pred.flatten()-YC_bkg.flatten())-1*sigmaCc, (YC_pred.flatten()-YC_bkg.flatten())+1*sigmaCc,alpha=.5, fc='m',label='Error on difference' )\n",
    "    diff.fill_between(XC.flatten(), np.sqrt(YC_pred.flatten()), -np.sqrt(YC_pred.flatten()),alpha=.4, fc='b' , label='Statistical Error band')\n",
    "    diff.set_ylabel('Difference',fontsize=15)\n",
    "    # diff.errorbar(xt, yt, dyt, fmt='.',color='b', markersize=8,alpha=1, label='ttbar MC prediction')\n",
    "    diff.errorbar(Xc.flatten(), Yc.flatten()-Yc_pred.flatten(), dyc, fmt='.',color='k', markersize=8,alpha=.5)\n",
    "    diff.errorbar(X.flatten(), Y.flatten()-Y_bkg.flatten(), dy, fmt='.', color='#3c0008',markersize=8,alpha=0.5, label='Data - Prediction')\n",
    "    diff.legend(loc='right',fontsize=13)\n",
    "    diff.axvline(x=130, color='lightgrey', linestyle='--')\n",
    "    diff.axvline(x=200, color='lightgrey', linestyle='--')\n",
    "    np.max(YC_pred.flatten()-YC_bkg.flatten())\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98f3149b",
   "metadata": {},
   "outputs": [],
   "source": [
    "x=[]\n",
    "y=[]\n",
    "dy=[]\n",
    "xc=[]\n",
    "yc=[]\n",
    "dyc=[]\n",
    "h = mass_hists['region_1']['data'].Clone('h_mass')\n",
    "h.Rebin(2)\n",
    "# regions[0]['h_a_nq_2018'].Clone('h_mass')\n",
    "# h.Add(regions[0]['h_a_nq_2017'])\n",
    "# h.Add(regions[0]['h_a_nq_2016'])\n",
    "\n",
    "# myFile = ROOT.TFile(\"resolved_FDhist.root\", \"RECREATE\")\n",
    "# myFile.cd()\n",
    "# regions[0]['h_a_nq_2017'].\n",
    "# myFile.Close()  \n",
    "\n",
    "sig_m=170\n",
    "sig_s=15\n",
    "sig_m_bounds=(150,190)\n",
    "sig_s_bounds=(7, 20)\n",
    "for i in range(1,h.GetNbinsX()+1):\n",
    "    if(h.GetBinContent(i)):\n",
    "        x.append(h.GetBinCenter(i))\n",
    "        y.append(h.GetBinContent(i))\n",
    "        dy.append(h.GetBinError(i))\n",
    "        if((h.GetBinCenter(i)<130 or h.GetBinCenter(i)>210)):\n",
    "            xc.append(h.GetBinCenter(i))\n",
    "            yc.append(h.GetBinContent(i))\n",
    "            dyc.append(h.GetBinError(i))\n",
    "print(len(x),len(y),len(dy),len(xc),len(yc),len(dyc))\n",
    "\n",
    "\n",
    "X = np.atleast_2d(x).T\n",
    "Y = np.atleast_2d(y).T\n",
    "dy = np.array(dy)\n",
    "\n",
    "Xc = np.atleast_2d(xc).T\n",
    "Yc = np.atleast_2d(yc).T\n",
    "dyc = np.array(dyc)\n",
    "# kernel = C(0.1, (0.01, 1000.0)) * Matern(length_scale=40.0, length_scale_bounds=(40, 100.0), nu=11.5)\n",
    "# kernel = kernel + \n",
    "# kernel = C(1.0, (1e-3, 1e3)) * RBF(30, (30, 1e3))\n",
    "\n",
    "kernel = C(1.0, (1e-3, 1e10)) * RBF(200, (10, 1e16))\n",
    "gp = GaussianProcessRegressor(kernel=kernel, alpha=dyc**2, n_restarts_optimizer=20)\n",
    "gp.fit(Xc, Yc)\n",
    "print(gp.kernel_)\n",
    "print('Using Lenght scale', float(str(gp.kernel_).split('=')[-1].split(')')[0]))\n",
    "\n",
    "ls  = float(str(gp.kernel_).split('=')[-1].split(')')[0])\n",
    "\n",
    "kernel = C(523**2, \"fixed\") * RBF(ls, \"fixed\") + C(1.0, (1e-3, 1e10)) * ESC(mass = sig_m, mass_bounds = sig_m_bounds , width = sig_s, width_bounds = sig_s_bounds)\n",
    "# Fit for unmasked Data\n",
    "print(X.shape, Y.shape, dy.shape)\n",
    "gp = GPnew(kernel=kernel, alpha=dy**2, n_restarts_optimizer=30)\n",
    "gp.fit(X, Y)\n",
    "XC = np.atleast_2d(np.linspace(x[0], x[-1], num=1000)).T\n",
    "Y_pred, sigma = gp.predict(X, return_std=True)\n",
    "YC_pred, sigmaC = gp.predict(XC, return_std=True)\n",
    "print(gp.kernel_)\n",
    "YC_sig, sigmaC_sig = gp.predict_partial(XC,return_std=True,kx = 'k2')\n",
    "YC_bkg, sigmaC_bkg = gp.predict_partial(XC,return_std=True,kx = 'k1')\n",
    "Y_sig, sigma_sig = gp.predict_partial(X,return_std=True,kx = 'k2')\n",
    "Y_bkg, sigma_bkg = gp.predict_partial(X,return_std=True,kx = 'k1')\n",
    "# kernel = C(0.1, (0.01, 1000.0)) * Matern(length_scale=100.0, length_scale_bounds=(100, 2000.0), nu=10.5)\n",
    "\n",
    "kernel = C(1.0, (1e-3, 1e10)) * RBF(200, (135, 1e3))\n",
    "#Fit for masked Data\n",
    "gp = GaussianProcessRegressor(kernel=kernel, alpha=dyc**2, n_restarts_optimizer=20)\n",
    "gp.fit(Xc, Yc)\n",
    "XCc = np.atleast_2d(np.linspace(xc[0], xc[-1], num=1000)).T\n",
    "Yc_pred, sigmac = gp.predict(Xc, return_std=True)\n",
    "YCc_pred, sigmaCc = gp.predict(XCc, return_std=True)\n",
    "Y_p, sig = gp.predict(X, return_std=True)\n",
    "\n",
    "# with plt.xkcd():\n",
    "fig = plt.figure(figsize=(10, 13))\n",
    "spec = fig.add_gridspec(ncols=1, nrows=3,height_ratios=[0.8,0.2,0.3])\n",
    "main = fig.add_subplot(spec[0,0])\n",
    "main.plot(XCc.flatten(), YCc_pred.flatten(), '--', color='gainsboro', alpha=0.5)\n",
    "# main.fill_between(XCc.flatten(),YCc_pred.flatten()-1*sigmaCc, YCc_pred.flatten()+1*sigmaCc, alpha=.5, fc='navy', ec='None', label='95% confidence interval')\n",
    "main.text(0.0, 1.05,'CMS Prelimininary ', ha='left', va='top', transform=main.transAxes, fontsize=15)\n",
    "main.text(0.99, 1.05,r'Run-2 ($9.837~fb^{-1}$)', ha='right', va='top', transform=main.transAxes, fontsize=15)\n",
    "main.plot(XC.flatten(), YC_pred.flatten(), 'b-',color=cols[1], label='Total Prediction')\n",
    "main.fill_between(XC.flatten(),YC_pred.flatten()-1*sigmaC, YC_pred.flatten()+1*sigmaC, alpha=.5, fc='k', ec='None',color=cols[1], label='Fit error')\n",
    "main.errorbar(X.flatten(), Y.flatten(), dy, fmt='.', color='k',markersize=8,alpha=1.0, label='Data')\n",
    "# main.errorbar(Xc.flatten(), Yc.flatten(), dyc, fmt='.',color='k', markersize=8,alpha=.8, label='Masked Data')\n",
    "main.set_ylabel('Events', fontsize=15)\n",
    "main.axvline(x=130, color='lightgrey', linestyle='--')\n",
    "main.axvline(x=210, color='lightgrey', linestyle='--')\n",
    "main.set_xlabel('Avg. Mass', fontsize=15)\n",
    "main.tick_params(axis='y', labelsize=15)\n",
    "main.tick_params(axis='x', labelsize=15)\n",
    "main.plot(XC.flatten(),YC_sig.flatten(),color=cols[3],label='Gaussian')\n",
    "main.plot(XC.flatten(),YC_bkg.flatten(),color=cols[0],label='Background Prediction')\n",
    "main.fill_between(XC.flatten(),YC_sig.flatten()-sigmaC_sig.flatten(),YC_sig.flatten()+sigmaC_sig.flatten(),color=cols[3],alpha=.5,label='Fit error')\n",
    "main.fill_between(XC.flatten(),YC_bkg.flatten()-sigmaC_bkg.flatten(),YC_bkg.flatten()+sigmaC_bkg.flatten(),color=cols[0],alpha=.5,label='Fit error')\n",
    "main.legend(loc='best', fontsize=15)\n",
    "# main.set_yscale('log')\n",
    "# main.set_ylim(0.9,1000)\n",
    "# main.plot(XCc.flatten(), YCc_pred.flatten()-YC_pred.flatten(), 'm-', label='Difference B/n predictions')\n",
    "# main.fill_between(XCc.flatten(), (YCc_pred.flatten()-YC_pred.flatten())-1*sigmaCc, (YCc_pred.flatten()-YC_pred.flatten())+1*sigmaCc,alpha=.5, fc='m' )\n",
    "pull = fig.add_subplot(spec[1,0])\n",
    "pull.bar(X.flatten(),((Y.flatten()-Y_p.flatten())/dy),width=h.GetBinWidth(1),color='r', alpha=1)\n",
    "pull.bar(X.flatten(),((Y.flatten()-Y_pred.flatten())/dy),width=h.GetBinWidth(1),color='k', alpha=0.7)\n",
    "pull.set_ylabel('Pull',fontsize=15)\n",
    "pull.axvline(x=130, color='lightgrey', linestyle='--')\n",
    "pull.axvline(x=210, color='lightgrey', linestyle='--')\n",
    "print(np.sum(((Y.flatten()-Y_pred.flatten())/dy)**2)/(len(x)-1))\n",
    "print(np.sum(((Y.flatten()-Y_p.flatten())/dy)**2)/(len(x)-1))\n",
    "print(np.sum(((Yc.flatten()-Yc_pred.flatten())/dyc)**2)/(len(x)-1))\n",
    "diff = fig.add_subplot(spec[2,0])\n",
    "diff.plot(XCc.flatten(), YC_pred.flatten()-YC_bkg.flatten(), 'm-', label='Gaussian')\n",
    "# diff.fill_between(XCc.flatten(), (YC_pred.flatten()-YC_bkg.flatten())-1*sigmaCc, (YC_pred.flatten()-YC_bkg.flatten())+1*sigmaCc,alpha=.5, fc='m',label='Error on difference' )\n",
    "diff.fill_between(XC.flatten(), np.sqrt(YC_pred.flatten()), -np.sqrt(YC_pred.flatten()),alpha=.4, fc='b' , label='Statistical Error band')\n",
    "diff.set_ylabel('Difference',fontsize=15)\n",
    "# diff.errorbar(xt, yt, dyt, fmt='.',color='b', markersize=8,alpha=1, label='ttbar MC prediction')\n",
    "diff.errorbar(Xc.flatten(), Yc.flatten()-Yc_pred.flatten(), dyc, fmt='.',color='k', markersize=8,alpha=.5)\n",
    "diff.errorbar(X.flatten(), Y.flatten()-Y_bkg.flatten(), dy, fmt='.', color='#3c0008',markersize=8,alpha=0.5, label='Data - Prediction')\n",
    "diff.legend(loc='right',fontsize=13)\n",
    "diff.axvline(x=130, color='lightgrey', linestyle='--')\n",
    "diff.axvline(x=210, color='lightgrey', linestyle='--')\n",
    "np.max(YC_pred.flatten()-YC_bkg.flatten())\n",
    "# fig.savefig('Mass_ranges_1.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e5b2c51",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.style.use(hep.style.ROOT)\n",
    "x=[]\n",
    "y=[]\n",
    "dy=[]\n",
    "xc=[]\n",
    "yc=[]\n",
    "dyc=[]\n",
    "h = mass_hists['region_1']['data'].Clone('h_mass')\n",
    "h.Rebin(2)\n",
    "\n",
    "# myFile = ROOT.TFile(\"resolved_FDhist.root\", \"RECREATE\")\n",
    "# myFile.cd()\n",
    "# regions[0]['h_a_nq_2017'].\n",
    "# myFile.Close()  \n",
    "\n",
    "sig_m=170\n",
    "sig_s=15\n",
    "sig_m_bounds=(150,190)\n",
    "sig_s_bounds=(7, 20)\n",
    "for i in range(1,h.GetNbinsX()+1):\n",
    "    if(h.GetBinContent(i)):\n",
    "        x.append(h.GetBinCenter(i))\n",
    "        y.append(h.GetBinContent(i))\n",
    "        dy.append(h.GetBinError(i))\n",
    "        if((h.GetBinCenter(i)<130 or h.GetBinCenter(i)>210)):\n",
    "            xc.append(h.GetBinCenter(i))\n",
    "            yc.append(h.GetBinContent(i))\n",
    "            dyc.append(h.GetBinError(i))\n",
    "print(len(x),len(y),len(dy),len(xc),len(yc),len(dyc))\n",
    "\n",
    "\n",
    "X = np.atleast_2d(x).T\n",
    "Y = np.atleast_2d(y).T\n",
    "dy = np.array(dy)\n",
    "\n",
    "Xc = np.atleast_2d(xc).T\n",
    "Yc = np.atleast_2d(yc).T\n",
    "dyc = np.array(dyc)\n",
    "# kernel = C(0.1, (0.01, 1000.0)) * Matern(length_scale=40.0, length_scale_bounds=(40, 100.0), nu=11.5)\n",
    "# kernel = kernel + \n",
    "# kernel = C(1.0, (1e-3, 1e3)) * RBF(30, (30, 1e3))\n",
    "\n",
    "kernel = C(1.0, (1e-3, 1e10)) * RBF(200, (10, 1e16))\n",
    "gp = GaussianProcessRegressor(kernel=kernel, alpha=dyc**2, n_restarts_optimizer=20)\n",
    "gp.fit(Xc, Yc)\n",
    "print(gp.kernel_)\n",
    "print('Using Lenght scale', float(str(gp.kernel_).split('=')[-1].split(')')[0]))\n",
    "\n",
    "ls  = float(str(gp.kernel_).split('=')[-1].split(')')[0])\n",
    "\n",
    "kernel = C(523**2, \"fixed\") * RBF(ls, \"fixed\") + C(1.0, (1e-3, 1e10)) * ESC(mass = sig_m, mass_bounds = sig_m_bounds , width = sig_s, width_bounds = sig_s_bounds)\n",
    "# Fit for unmasked Data\n",
    "print(X.shape, Y.shape, dy.shape)\n",
    "gp = GPnew(kernel=kernel, alpha=dy**2, n_restarts_optimizer=30)\n",
    "gp.fit(X, Y)\n",
    "XC = np.atleast_2d(np.linspace(x[0], x[-1], num=1000)).T\n",
    "Y_pred, sigma = gp.predict(X, return_std=True)\n",
    "YC_pred, sigmaC = gp.predict(XC, return_std=True)\n",
    "print(gp.kernel_)\n",
    "YC_sig, sigmaC_sig = gp.predict_partial(XC,return_std=True,kx = 'k2')\n",
    "YC_bkg, sigmaC_bkg = gp.predict_partial(XC,return_std=True,kx = 'k1')\n",
    "Y_sig, sigma_sig = gp.predict_partial(X,return_std=True,kx = 'k2')\n",
    "Y_bkg, sigma_bkg = gp.predict_partial(X,return_std=True,kx = 'k1')\n",
    "# kernel = C(0.1, (0.01, 1000.0)) * Matern(length_scale=100.0, length_scale_bounds=(100, 2000.0), nu=10.5)\n",
    "\n",
    "kernel = C(1.0, (1e-3, 1e10)) * RBF(200, (135, 1e3))\n",
    "#Fit for masked Data\n",
    "gp = GaussianProcessRegressor(kernel=kernel, alpha=dyc**2, n_restarts_optimizer=20)\n",
    "gp.fit(Xc, Yc)\n",
    "XCc = np.atleast_2d(np.linspace(xc[0], xc[-1], num=1000)).T\n",
    "Yc_pred, sigmac = gp.predict(Xc, return_std=True)\n",
    "YCc_pred, sigmaCc = gp.predict(XCc, return_std=True)\n",
    "Y_p, sig = gp.predict(X, return_std=True)\n",
    "\n",
    "# with plt.xkcd():\n",
    "fig = plt.figure(figsize=(10, 13))\n",
    "spec = fig.add_gridspec(ncols=1, nrows=2,height_ratios=[0.8,0.2])\n",
    "main = fig.add_subplot(spec[0,0])\n",
    "hep.cms.label(loc=0,data=True,llabel='Preliminary', rlabel=\"128 fb$^{-1}$ (13 TeV)\")\n",
    "main.axvline(x=200, color='grey', linestyle='--')\n",
    "main.axvline(x=500, color='grey', linestyle='--')\n",
    "main.plot(X.flatten(), Y_pred.flatten(), 'b-',color=cols[1],drawstyle='steps-mid', label=r'Background$+t\\bar{t}$ fit')\n",
    "main.errorbar(X.flatten(), Y.flatten(), dy, fmt='.', color='k',markersize=8,alpha=1.0, label='Data')\n",
    "main.set_ylabel('Triplets / {} GeV'.format(int(x[1]-x[0])),  fontsize='medium')\n",
    "\n",
    "main.tick_params(axis='y', )\n",
    "main.tick_params(axis='x', )\n",
    "main.plot(XC.flatten(),YC_bkg.flatten(),color=cols[0],label='Background Fit')\n",
    "xc_sig = np.linspace(195-1.5*16.5, 195+1.5*16.5, 1000)\n",
    "main.plot(xc_sig,2500*gauss_array(xc_sig,195,16.5,x[1]-x[0]),color=cols[3],label=r'Signal $M_{\\tilde{g}} = 200$ GeV') \n",
    "xc_sig = np.linspace(290-1.5*23.1, 290+1.5*23.1, 1000)\n",
    "main.plot(xc_sig,2500*gauss_array(xc_sig,290,23.1,x[1]-x[0]),color=cols[2],label=r'Signal $M_{\\tilde{g}} = 300$ GeV') \n",
    "xc_sig = np.linspace(387-1.5*29.2, 387+1.5*29.2, 1000)\n",
    "main.plot(xc_sig,2500*gauss_array(xc_sig,387,29.2,x[1]-x[0]),color=cols[4],label=r'Signal $M_{\\tilde{g}} = 400$ GeV') \n",
    "main.plot([],[],' ',label=r'$\\chi^2$/ndof = {:.1f}/{}'.format(np.sum(((Y.flatten()-Y_pred.flatten())/dy)**2),int(len(x)-1)))\n",
    "main.fill_between(XC.flatten(),YC_bkg.flatten()-sigmaC_bkg.flatten(),YC_bkg.flatten()+sigmaC_bkg.flatten(),color=cols[0],alpha=.5)\n",
    "handles,labels = main.get_legend_handles_labels()\n",
    "handles = [handles[6],handles[0],handles[5],handles[1],handles[2],handles[3],handles[4]]\n",
    "labels = [labels[6],labels[0],labels[5],labels[1],labels[2],labels[3],labels[4]]\n",
    "main.legend(handles,labels,loc='upper right', )\n",
    "# main.legend(loc='best', )\n",
    "# main.set_yscale('log')\n",
    "main.set_ylim(0,900)\n",
    "# main.plot(XCc.flatten(), YCc_pred.flatten()-YC_pred.flatten(), 'm-', label='Difference B/n predictions')\n",
    "# main.fill_between(XCc.flatten(), (YCc_pred.flatten()-YC_pred.flatten())-1*sigmaCc, (YCc_pred.flatten()-YC_pred.flatten())+1*sigmaCc,alpha=.5, fc='m' )\n",
    "pull = fig.add_subplot(spec[1,0],sharex=main)\n",
    "# pull.bar(X.flatten(),((Y.flatten()-Y_bkg.flatten())/dy),width=h_a.GetBinWidth(1),color='r', alpha=1)\n",
    "pull.bar(X.flatten(),((Y.flatten()-Y_pred.flatten())/dy),width=h.GetBinWidth(1),color='k', alpha=0.7)\n",
    "pull.set_ylabel('Pull', fontsize='medium', loc='center')\n",
    "pull.set_xlabel(r'$m_{jjj}$ [GeV]',  fontsize='medium')\n",
    "pull.set_ylim(-3.5,3.5)\n",
    "# pull.axvline(x=(95-(2*11)), color='lightgrey', linestyle='--')\n",
    "# pull.axvline(x=(95+(2*11)), color='lightgrey', linestyle='--')\n",
    "print(np.sum(((Y.flatten()-Y_pred.flatten())/dy)**2)/(len(x)-1))\n",
    "print(np.sum(((Y.flatten()-Y_bkg.flatten())/dy)**2)/(len(x)-1))\n",
    "print(np.sum(((Yc.flatten()-Yc_pred.flatten())/dyc)**2)/(len(x)-1))\n",
    "# diff = fig.add_subplot(spec[2,0])\n",
    "# diff.plot(XCc.flatten(), YC_pred.flatten()-YC_bkg.flatten(), 'm-', label='Gaussian')\n",
    "# # diff.fill_between(XCc.flatten(), (YC_pred.flatten()-YC_bkg.flatten())-1*sigmaCc, (YC_pred.flatten()-YC_bkg.flatten())+1*sigmaCc,alpha=.5, fc='m',label='Error on difference' )\n",
    "# diff.fill_between(XC.flatten(), np.sqrt(YC_pred.flatten()), -np.sqrt(YC_pred.flatten()),alpha=.4, fc='b' , label='Statistical Error band')\n",
    "# diff.set_ylabel('Difference',)\n",
    "# diff.errorbar(xt, yt, dyt, fmt='.',color='b', markersize=8,alpha=1, label='ttbar MC prediction')\n",
    "# # diff.errorbar(Xc.flatten(), Yc.flatten()-Yc_pred.flatten(), dyc, fmt='.',color='k', markersize=8,alpha=.5)\n",
    "# diff.errorbar(X.flatten(), Y.flatten()-Y_bkg.flatten(), dy, fmt='.', color='#3c0008',markersize=8,alpha=0.5, label='Data - Prediction')\n",
    "# diff.errorbar(X.flatten(), Y.flatten()-Y_p.flatten(), dy, fmt='.', color='k',markersize=8,alpha=0.5, label='Data - Prediction')\n",
    "# diff.legend(loc='right',)\n",
    "# diff.axvline(x=(95-(2*11)), color='lightgrey', linestyle='--')\n",
    "# diff.axvline(x=(95+(2*11)), color='lightgrey', linestyle='--')\n",
    "np.max(YC_pred.flatten()-YC_bkg.flatten())\n",
    "plt.subplots_adjust(wspace=0, hspace=0)\n",
    "fig.savefig('Mass_ranges_1_sig_'+tag+'.pdf',bbox_inches='tight')\n",
    "plt.show()\n",
    "fig = plt.figure(figsize=(10,10))\n",
    "plt.hist(((Y.flatten()-Y_pred.flatten())/dy),bins=20,density=True)\n",
    "plt.plot([],[],' ',label='Mean: {:.2f}'.format(np.mean(((Y.flatten()-Y_pred.flatten())/dy))))\n",
    "plt.plot([],[],' ',label='Std: {:.2f}'.format(np.std(((Y.flatten()-Y_pred.flatten())/dy))))\n",
    "plt.xlabel('Pull',  fontsize='medium')\n",
    "plt.ylabel('Fraction',)\n",
    "plt.tick_params(axis='y', )\n",
    "plt.tick_params(axis='x', )\n",
    "plt.legend()\n",
    "print(np.mean(((Y.flatten()-Y_pred.flatten())/dy)),np.std(((Y.flatten()-Y_pred.flatten())/dy)))\n",
    "# fig.savefig('Mass_ranges_1_pull_'+tag+'.pdf',bbox_inches='tight')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "multijets_env",
   "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}