Update BTA workflow (#56)

* Update ptrel, ratio, ratioRel, and IP calculation * Update formula * Fix ptrel formula * Add track collection to BTA workflow * Incorporate Ming-Yan’s comment - correct CI bug - correct fill_none values - add README section for the additional two BTA workflows * Remove _algo HLTs as they are special for 2018
cms-btv-pog · Jul 19, 2023 · f61bb9a · f61bb9a
1 parent 895f84e
commit f61bb9a
Show file tree

Hide file tree

Showing 6 changed files with 298 additions and 52 deletions.
diff --git a/.github/workflows/BTA_workflow.yml b/.github/workflows/BTA_workflow.yml
@@ -87,6 +87,6 @@ jobs:
   
     - name: BTA workflow test
       run: |
-         python runner.py --wf BTA --json metadata/test_bta_run3.json --campaign Summer22EERun3 --executor iterative --isJERC
+         python runner.py --wf BTA --json metadata/test_bta_run3.json --campaign Summer22Run3 --executor iterative --isJERC
     
 
diff --git a/README.md b/README.md
@@ -63,7 +63,7 @@ More options for `runner.py`
 <p>
 
 ```
---wf {validation,ttcom,ttdilep_sf,ttsemilep_sf,emctag_ttdilep_sf,ctag_ttdilep_sf,ectag_ttdilep_sf,ctag_ttsemilep_sf,ectag_ttsemilep_sf,ctag_Wc_sf,ectag_Wc_sf,ctag_DY_sf,ectag_DY_sf,BTA}, --workflow {validation,ttcom,ttdilep_sf,ttsemilep_sf,emctag_ttdilep_sf,ctag_ttdilep_sf,ectag_ttdilep_sf,ctag_ttsemilep_sf,ectag_ttsemilep_sf,ctag_Wc_sf,ectag_Wc_sf,ctag_DY_sf,ectag_DY_sf,BTA}
+--wf {validation,ttcom,ttdilep_sf,ttsemilep_sf,emctag_ttdilep_sf,ctag_ttdilep_sf,ectag_ttdilep_sf,ctag_ttsemilep_sf,ectag_ttsemilep_sf,ctag_Wc_sf,ectag_Wc_sf,ctag_DY_sf,ectag_DY_sf,BTA,BTA_addPFMuons,BTA_addAllTracks}, --workflow {validation,ttcom,ttdilep_sf,ttsemilep_sf,emctag_ttdilep_sf,ctag_ttdilep_sf,ectag_ttdilep_sf,ctag_ttsemilep_sf,ectag_ttsemilep_sf,ctag_Wc_sf,ectag_Wc_sf,ctag_DY_sf,ectag_DY_sf,BTA,BTA_addPFMuons,BTA_addAllTracks}
                         Which processor to run
   -o OUTPUT, --output OUTPUT
                         Output histogram filename (default: hists.coffea)
@@ -208,10 +208,21 @@ Based on Congqiao's [development](notebooks/BTA_array_producer.ipynb) to produce
 <details><summary>details</summary>
 <p>
 
+Run with the nominal `BTA` workflow to include the basic event variables, jet observables, and GEN-level quarks, hadrons, leptons, and V0 variables. 
 ```
 python runner.py --wf BTA --json metadata/test_bta_run3.json --campaign Summer22EERun3 --isJERC
 ```
 
+Run with the `BTA_addPFMuons` workflow to additionally include the `PFMuon` and `TrkInc` collection, used by the b-tag SF derivation with the QCD(μ) methods.
+```
+python runner.py --wf BTA_addPFMuons --json metadata/test_bta_run3.json --campaign Summer22EERun3 --isJERC
+```
+
+Run with the `BTA_addAllTracks` workflow to additionally include the `Tracks` collection, used by the JP variable calibration.
+```
+python runner.py --wf BTA_addAllTracks --json metadata/test_bta_run3.json --campaign Summer22EERun3 --isJERC
+```
+
 </p>
 </details>
 

diff --git a/src/BTVNanoCommissioning/helpers/BTA_helper.py b/src/BTVNanoCommissioning/helpers/BTA_helper.py
@@ -2,6 +2,7 @@
 import numba as nb
 import numpy as np
 import pandas as pd
+import coffea.nanoevents.methods.vector as vector
 import os, psutil
 
 BTA_HLT = [
@@ -45,15 +46,6 @@
     "BTagMu_AK4Jet300_Mu5",
     "BTagMu_AK8DiJet170_Mu5",
     "BTagMu_AK8Jet300_Mu5",
-    "BTagMu_AK4DiJet20_Mu5_noalgo",
-    "BTagMu_AK4DiJet40_Mu5_noalgo",
-    "BTagMu_AK4DiJet70_Mu5_noalgo",
-    "BTagMu_AK4DiJet110_Mu5_noalgo",
-    "BTagMu_AK4DiJet170_Mu5_noalgo",
-    "BTagMu_AK4Jet300_Mu5_noalgo",
-    "BTagMu_AK8DiJet170_Mu5_noalgo",
-    "BTagMu_AK8Jet300_Mu5_noalgo",
-    "BTagMu_AK8Jet170_DoubleMu5_noalgo",
 ]
 # mass table from https://github.com/scikit-hep/particle/blob/master/src/particle/data/particle2022.csv and https://gitlab.cern.ch/lhcb-conddb/DDDB/-/blob/master/param/ParticleTable.txt
 df_main = pd.read_csv(
@@ -145,3 +137,48 @@ def cumsum(array):
     )
     cumsum_array = ak.fill_none(scan - ak.firsts(scan) + ak.firsts(array), [], axis=0)
     return cumsum_array
+
+
+def calc_ip_vector(obj, dxy, dz, is_3d=False):
+    '''Calculate the 2D or 3D impact parameter vector, given the track obj (with 4-mom), 
+       and its dxy and dz, taking the standard definition from NanoAOD'''
+
+    # 2D impact parameter
+    pvec = ak.zip(
+        {
+            'x': obj.px,
+            'y': obj.py,
+            'z': obj.pz,
+        },
+        behavior=vector.behavior,
+        with_name='ThreeVector'
+    )
+    zvec = ak.zip(
+        {
+            'x': ak.zeros_like(dxy),
+            'y': ak.zeros_like(dxy),
+            'z': ak.zeros_like(dxy) + 1,
+        },
+        behavior=vector.behavior,
+        with_name='ThreeVector'
+    )
+    # 2D impact parameter vector: (-py, px) / pt * dxy
+    ipvec_2d = zvec.cross(pvec) * dxy / obj.pt
+
+    if is_3d == False:
+        return ipvec_2d
+
+    # Then calculate the 3D impact parameter vector
+    # first, extend ipvec_2d to 3D space
+    ipvec_2d_ext = ak.zip(
+        {
+            'x': ipvec_2d.x,
+            'y': ipvec_2d.y,
+            'z': dz,
+        },
+        behavior=vector.behavior,
+        with_name='ThreeVector'
+    )
+    # then, get the closest distance to the track on 3D geometry
+    ipvec_3d = ipvec_2d_ext - ipvec_2d_ext.dot(pvec) / pvec.p2 * pvec
+    return ipvec_3d
diff --git a/src/BTVNanoCommissioning/helpers/update_branch.py b/src/BTVNanoCommissioning/helpers/update_branch.py
@@ -4,7 +4,7 @@
 
 
 def missing_branch(events):
-    events.fixedGridRhoFastjetAll = (
+    events['fixedGridRhoFastjetAll'] = (
         events.fixedGridRhoFastjetAll
         if hasattr(events, "fixedGridRhoFastjetAll")
         else events.Rho.fixedGridRhoFastjetAll