Skip to content

Instantly share code, notes, and snippets.

@nsmith-

nsmith-/btagSFexample.ipynb Secret

Last active November 22, 2023 00:18
Show Gist options
  • Save nsmith-/8d3d41aaffda92148ebc7bfcc5c827f5 to your computer and use it in GitHub Desktop.
Save nsmith-/8d3d41aaffda92148ebc7bfcc5c827f5 to your computer and use it in GitHub Desktop.
Download ZIP
Display the source blob
Display the rendered blob
Raw
btagSFexample.ipynb
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"id": "c225b0b3",
"metadata": {},
"outputs": [],
"source": [
"import correctionlib\n",
"import hist\n",
"import awkward as ak\n",
"import numpy as np\n",
"\n",
"from coffea.nanoevents import NanoEventsFactory, NanoAODSchema\n",
"from coffea.lookup_tools.correctionlib_wrapper import correctionlib_wrapper\n",
"from coffea.lookup_tools.dense_lookup import dense_lookup"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "9216f46d",
"metadata": {},
"outputs": [],
"source": [
"events = NanoEventsFactory.from_root(\n",
" \"root://cmsxrootd-site.fnal.gov//store/mc/RunIISummer20UL18NanoAODv9/TTToSemiLeptonic_TuneCP5_13TeV-powheg-pythia8/NANOAODSIM/106X_upgrade2018_realistic_v16_L1v1-v1/120000/0520A050-AF68-EF43-AA5B-5AA77C74ED73.root\",\n",
" entry_stop=100_000,\n",
" schemaclass=NanoAODSchema,\n",
").events()"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "2a8cad60",
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"<html>\n",
"<div style=\"display:flex; align-items:center;\">\n",
"<div style=\"width:290px;\">\n",
"<svg xmlns=\"http://www.w3.org/2000/svg\" viewBox=\"-10 -10 220 220\">\n",
"<rect x=\"0\" y=\"0\" width=\"160\" height=\"160\" style=\"fill:white;opacity:.5;stroke-width:2;stroke:currentColor;\"/>\n",
"<rect x=\"20\" y=\"20\" width=\"160\" height=\"160\" style=\"fill:white;opacity:.5;stroke-width:2;stroke:currentColor;\"/>\n",
"<rect x=\"40\" y=\"40\" width=\"160\" height=\"160\" style=\"fill:white;opacity:.5;stroke-width:2;stroke:currentColor;\"/>\n",
"<text x=\"120.0\" y=\"120.0\" style=\"font-size: 26pt; font-family: verdana; font-style: bold; fill: black;\" text-anchor=\"middle\" alignment-baseline=\"middle\">\n",
"4D\n",
"</text>\n",
"</svg>\n",
"</div>\n",
"<div style=\"flex=grow:1;\">\n",
"Regular(20, 40, 300, name='pt', label='pt')<br/>\n",
"Regular(4, 0, 2.5, name='abseta', label='abseta')<br/>\n",
"IntCategory([0, 4, 5], name='flavor', label='flavor')<br/>\n",
"Boolean(name='passWP', label='passWP')<br/>\n",
"<hr style=\"margin-top:.2em; margin-bottom:.2em;\"/>\n",
"Double() Σ=331859.0 <em>(334836.0 with flow)</em>\n",
"\n",
"</div>\n",
"</div>\n",
"</html>"
],
"text/plain": [
"Hist(\n",
" Regular(20, 40, 300, name='pt', label='pt'),\n",
" Regular(4, 0, 2.5, name='abseta', label='abseta'),\n",
" IntCategory([0, 4, 5], name='flavor', label='flavor'),\n",
" Boolean(name='passWP', label='passWP'),\n",
" storage=Double()) # Sum: 331859.0 (334836.0 with flow)"
]
},
"execution_count": 3,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"phasespace_cuts = (\n",
" (abs(events.Jet.eta) < 2.5)\n",
" & (events.Jet.pt > 40.)\n",
")\n",
"jets = ak.flatten(events.Jet[phasespace_cuts])\n",
"\n",
"efficiencyinfo = (\n",
" hist.Hist.new\n",
" .Reg(20, 40, 300, name=\"pt\")\n",
" .Reg(4, 0, 2.5, name=\"abseta\")\n",
" .IntCat([0, 4, 5], name=\"flavor\")\n",
" .Bool(name=\"passWP\")\n",
" .Double()\n",
" .fill(\n",
" pt=jets.pt,\n",
" abseta=abs(jets.eta),\n",
" flavor=jets.hadronFlavour,\n",
" passWP=jets.btagDeepFlavB > 0.2783, # UL 2018 medium WP\n",
" )\n",
")\n",
"efficiencyinfo"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "d1906161",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"3 dimensional histogram with axes:\n",
"\t1: [ 40. 53. 66. 79. 92. 105. 118. 131. 144. 157. 170. 183. 196. 209.\n",
" 222. 235. 248. 261. 274. 287. 300.]\n",
"\t2: [0. 0.625 1.25 1.875 2.5 ]\n",
"\t3: [0. 1. 2. 3.]"
]
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"eff = efficiencyinfo[{\"passWP\": True}] / efficiencyinfo[{\"passWP\": sum}]\n",
"# note this seems to turn 0,4,5 into 0,1,2\n",
"efflookup = dense_lookup(eff.values(), [ax.edges for ax in eff.axes])\n",
"efflookup"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "e2f81919",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"array([0.01634657, 0.17330046, 0.83132001])"
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# Efficiency at 42 GeV, |eta|=0.2, for light, c, and b quark respectively\n",
"efflookup(42, 0.2, np.array([0, 1, 2]))"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "3a458bce",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"['deepCSV_comb', 'deepCSV_incl', 'deepCSV_mujets', 'deepCSV_shape', 'deepJet_comb', 'deepJet_incl', 'deepJet_mujets', 'deepJet_shape']\n"
]
}
],
"source": [
"cset = correctionlib.CorrectionSet.from_file(\"/cvmfs/cms.cern.ch/rsync/cms-nanoAOD/jsonpog-integration/POG/BTV/2018_UL/btagging.json.gz\")\n",
"print([c for c in cset])\n",
"# more docs at https://cms-nanoaod-integration.web.cern.ch/commonJSONSFs/BTV_btagging_Run2_UL/BTV_btagging_2018_UL.html"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "9eda57cc",
"metadata": {},
"outputs": [],
"source": [
"def lighttagSF(j, syst=\"central\"):\n",
" # until correctionlib handles jagged data natively we have to flatten and unflatten\n",
" j, nj = ak.flatten(j), ak.num(j)\n",
" sf = cset[\"deepJet_incl\"].evaluate(syst, \"M\", np.array(j.hadronFlavour), np.array(abs(j.eta)), np.array(j.pt))\n",
" return ak.unflatten(sf, nj)\n",
"\n",
"\n",
"def btagSF(j, syst=\"central\"):\n",
" # until correctionlib handles jagged data natively we have to flatten and unflatten\n",
" j, nj = ak.flatten(j), ak.num(j)\n",
" sf = cset[\"deepJet_comb\"].evaluate(syst, \"M\", np.array(j.hadronFlavour), np.array(abs(j.eta)), np.array(j.pt))\n",
" return ak.unflatten(sf, nj)"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "3218073d",
"metadata": {},
"outputs": [],
"source": [
"lightJets = events.Jet[phasespace_cuts & (events.Jet.hadronFlavour == 0)]\n",
"bcJets = events.Jet[phasespace_cuts & (events.Jet.hadronFlavour > 0)]\n",
"\n",
"lightEff = efflookup(lightJets.pt, abs(lightJets.eta), lightJets.hadronFlavour)\n",
"bcEff = efflookup(bcJets.pt, abs(bcJets.eta), bcJets.hadronFlavour-3)"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "7bf389ca",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"<Array [0.986, 1.67, 1.15, ... 1.13, 0.982] type='100000 * float64'>"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# BTagSFMethod \"1a\"\n",
"\n",
"def combine(eff, sf, passbtag):\n",
" # tagged SF = SF*eff / eff = SF\n",
" tagged_sf = ak.prod(sf[passbtag], axis=-1)\n",
" # untagged SF = (1 - SF*eff) / (1 - eff)\n",
" untagged_sf = ak.prod(((1 - sf*eff) / (1 - eff))[~passbtag], axis=-1)\n",
" return tagged_sf * untagged_sf\n",
"\n",
"lightweight = combine(\n",
" lightEff,\n",
" lighttagSF(lightJets),\n",
" lightJets.btagDeepB > 0.2783,\n",
")\n",
"bcweight = combine(\n",
" bcEff,\n",
" btagSF(bcJets),\n",
" bcJets.btagDeepB > 0.2783,\n",
")\n",
"eventweight = lightweight * bcweight\n",
"eventweight"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "89d4597b",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"<Array [0.992, 1.15, 0.997, ... 0.99, 1, 0.99] type='100000 * float64'>"
]
},
"execution_count": 10,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"eventweight_lightflavUp = combine(\n",
" lightEff,\n",
" lighttagSF(lightJets, \"up\"),\n",
" lightJets.btagDeepB > 0.2783,\n",
") * bcweight\n",
"\n",
"eventweight_lightflavUp / eventweight"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3d2902e1",
"metadata": {},
"outputs": [],
"source": []
}
],
"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.8.12"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment