#!/usr/bin/env python3
##########################################################################
# basf2 (Belle II Analysis Software Framework) #
# Author: The Belle II Collaboration #
# #
# See git log for contributors and copyright holders. #
# This file is licensed under LGPL-3.0, see LICENSE.md. #
##########################################################################
import re
import basf2 as b2
import modularAnalysis as ma
import pdg
from ROOT import Belle2
Const = Belle2.Const
# define arrays to interpret cut matrix
_chargednames = ['pi', 'K', 'p', 'e', 'mu']
_pidnames = ['pionID', 'kaonID', 'protonID', 'electronID', 'muonID']
_stdnames = ['all', 'loose', 'loosepid', 'good', 'higheff']
_effnames = ['95eff', '90eff', '85eff']
# default particle list for stdPi() and similar functions
_defaultlist = 'good'
_mostLikelyList = 'mostlikely'
def _stdChargedEffCuts(particletype, listtype):
"""
Provides the PID cut corresponding to a given efficiency percentile
@param particletype type of charged particle (pi, K, p, e, mu)
@param listtype efficiency percentile for the list (95eff, 90eff, 85eff)
"""
particleindex = _chargednames.index(particletype)
effindex = _effnames.index(listtype)
# efficiency cuts = [.95,.90,.85] efficiency; values outside (0,1) mean the cut does not exist and an error will be thrown
effcuts = [[0.001, 0.019, 0.098],
[5e-6, 0.027, 0.167],
[0.000, 0.043, 0.251],
[0.093, 0.301, 0.709],
[0.187, 0.418, 0.909]]
#
return effcuts[particleindex][effindex]
[docs]def stdCharged(particletype, listtype, path):
"""
Function to prepare one of several standardized types of charged particle lists:
- 'all' with no cuts on track
- 'good' high purity lists for data studies
- 'loosepid' loose selections for skimming, PID cut only
- 'loose' loose selections for skimming
- 'higheff' high efficiency list with loose global ID cut for data studies
- 'mostlikely' list with the highest PID likelihood
Also the following lists, which may or may not be available depending on the release
- '99eff' with 99% selection efficiency (calculated for 1<p<4 GeV) and good track (MC only)
- '95eff' with 95% selection efficiency (calculated for 1<p<4 GeV) and good track (MC only)
- '90eff' with 90% selection efficiency (calculated for 1<p<4 GeV) and good track (MC only)
- '85eff' with 85% selection efficiency (calculated for 1<p<4 GeV) and good track (MC only)
@param particletype type of charged particle to make a list of
@param listtype name of standard list
@param path modules are added to this path
"""
# basic quality cut strings
trackQuality = 'thetaInCDCAcceptance and nCDCHits>20'
ipCut = 'dr < 0.5 and abs(dz) < 2'
goodTrack = trackQuality + ' and ' + ipCut
if particletype not in _chargednames:
b2.B2ERROR("The requested list is not a standard charged particle. Use one of pi, K, e, mu, p.")
if listtype == 'all':
ma.fillParticleList(particletype + '+:all', '', True, path=path)
elif listtype == 'good':
ma.fillParticleList(
particletype + '+:good',
_pidnames[_chargednames.index(particletype)] + ' > 0.5 and ' + goodTrack,
True,
path=path)
elif listtype == 'loose':
ma.fillParticleList(
particletype + '+:loose',
_pidnames[_chargednames.index(particletype)] + ' > 0.1 and ' + goodTrack,
True,
path=path)
elif listtype == 'loosepid':
ma.fillParticleList(
particletype + '+:loosepid',
_pidnames[_chargednames.index(particletype)] + ' > 0.1',
True,
path=path)
elif listtype == 'higheff':
ma.fillParticleList(
particletype + '+:higheff',
_pidnames[_chargednames.index(particletype)] + ' > 0.002 and ' + goodTrack,
True,
path=path)
elif listtype not in _effnames:
b2.B2ERROR("The requested list is not defined. Please refer to the stdCharged documentation.")
else:
pidcut = _stdChargedEffCuts(particletype, listtype)
if 0.0 < pidcut < 1.0:
ma.fillParticleList(
particletype +
'+:' +
listtype,
_pidnames[_chargednames.index(particletype)] +
' > ' +
str(pidcut) +
' and ' +
goodTrack,
True,
path=path)
else:
b2.B2ERROR('The requested standard particle list ' + particletype +
'+:' + listtype + ' is not available in this release.')
[docs]def stdPi(listtype=_defaultlist, path=None):
"""
Function to prepare standard pion lists, refer to `stdCharged` for details
@param listtype name of standard list
@param path modules are added to this path
"""
stdCharged('pi', listtype, path)
[docs]def stdK(listtype=_defaultlist, path=None):
"""
Function to prepare standard kaon lists, refer to `stdCharged` for details
@param listtype name of standard list
@param path modules are added to this path
"""
stdCharged('K', listtype, path)
[docs]def stdPr(listtype=_defaultlist, path=None):
"""
Function to prepare standard proton lists, refer to `stdCharged` for details
@param listtype name of standard list
@param path modules are added to this path
"""
stdCharged('p', listtype, path)
[docs]def stdLep(pdgId, listtype, method, classification, path=None):
"""
Function to prepare one of several standardized types of lepton (:math:`e,\\mu`) lists:
* 'UniformEff60' 60% lepton efficiency list, uniform in a given multi-dimensional parametrisation.
* 'UniformEff70' 70% lepton efficiency list, uniform in a given multi-dimensional parametrisation.
* 'UniformEff80' 80% lepton efficiency list, uniform in a given multi-dimensional parametrisation.
* 'UniformEff90' 90% lepton efficiency list, uniform in a given multi-dimensional parametrisation.
* 'UniformEff95' 95% lepton efficiency list, uniform in a given multi-dimensional parametrisation.
The function will select particles according to the chosen ``listtype``, and decorate each candidate
with the nominal Data/MC :math:`\\ell` ID efficiency and :math:`\\pi,K` fake rate
correction factors and their stat, syst uncertainty, reading the info from the Conditions Database.
Parameters:
pdgId (int): the lepton pdg code.
listtype (str): name of standard list. Choose among the above values.
method (str): the PID method: 'likelihood' or 'bdt'.
classification (str): the type of classifier: 'binary' (one-vs-pion) or 'global' (one-vs-all).
path (basf2.Path): modules are added to this path.
"""
std_lepton_list_names = (
"UniformEff60",
"UniformEff70",
"UniformEff80",
"UniformEff90",
"UniformEff95",
)
available_methods = ("likelihood", "bdt")
available_classificators = ("global", "binary")
# We stick to positive pdgId by convention.
# Anyway, the particle list will be filled for anti-particles too.
pdgId = abs(pdgId)
if listtype not in std_lepton_list_names:
b2.B2ERROR("The requested lepton list is not defined. Please refer to the stdLep and stdCharged documentation.")
return
if pdgId not in (Const.electron.getPDGCode(), Const.muon.getPDGCode()):
b2.B2ERROR(f"{pdgId} is not that of a light charged lepton.")
return
if method not in available_methods:
b2.B2ERROR(f"method: {method}. Must be any of: {available_methods}.")
return
if classification not in available_classificators:
b2.B2ERROR(f"classification: {classification}. Must be any of: {available_classificators}.")
return
pid_variables = {
"likelihood": {
# TEMP: use 'electronID_noTOP' for electrons to circumvent bug in TOP electron PDFs in release 5.
"global": "electronID_noTOP" if pdgId == Const.electron.getPDGCode() else "muonID",
# TEMP: use 'binaryPID_noTOP' for electrons to circumvent bug in TOP electron PDFs in release 5.
"binary": f"binaryPID_noTOP({pdgId}, {Const.pion.getPDGCode()})" if pdgId == Const.electron.getPDGCode() \
else f"binaryPID({pdgId}, {Const.pion.getPDGCode()})"
},
"bdt": {
"global": f"pidChargedBDTScore({pdgId}, ALL)",
"binary": f"pidPairChargedBDTScore({pdgId}, 211, ALL)"
}
}
# Start creating the particle list, w/o any selection.
plistname = f"{pdg.to_name(pdgId)}:{listtype}"
ma.fillParticleList(plistname, "", path=path)
# The PID variable name, as it appears in the VariableManager.
pid_var = pid_variables[method][classification]
# Remove non-alphanumeric chars from the variable name, and strip last "_" if present.
# This is needed to match the name of the payload in the CDB.
pid_var_stripped = re.sub(r"[\W]+", "_", pid_var).rstrip("_")
# The names of the payloads w/ efficiency and mis-id corrections.
payload_eff = f"ParticleReweighting:{pid_var_stripped}_eff_combination_{listtype}"
payload_misid_pi = f"ParticleReweighting:{pid_var_stripped}_misid_pi_combination_{listtype}"
payload_misid_K = f"ParticleReweighting:{pid_var_stripped}_misid_K_combination_{listtype}"
# Configure weighting module(s).
reweighter_eff = path.add_module("ParticleWeighting",
particleList=plistname,
tableName=payload_eff).set_name(f"ParticleWeighting_eff_{plistname}")
reweighter_misid_pi = path.add_module("ParticleWeighting",
particleList=plistname,
tableName=payload_misid_pi).set_name(f"ParticleWeighting_misid_pi_{plistname}")
reweighter_misid_K = path.add_module("ParticleWeighting",
particleList=plistname,
tableName=payload_misid_K).set_name(f"ParticleWeighting_misid_K_{plistname}")
# Apply the PID selection cut, which is read from the efficiency payload.
cut = f"{pid_var} > extraInfo({payload_eff}_threshold)"
ma.applyCuts(plistname, cut, path=path)
[docs]def stdE(listtype=_defaultlist, method=None, classification=None, path=None):
""" Function to prepare one of several standardized types of electron lists.
See the documentation of `stdLep` for details.
It also accepts any of the legacy definitions
for the ``listtype`` parameter to fall back to the `stdCharged` behaviour:
* 'all'
* 'good'
* 'loosepid'
* 'loose'
* 'higheff'
* '95eff'
* '90eff'
* '85eff'
"""
if listtype in _stdnames + _effnames:
stdCharged("e", listtype, path)
return
stdLep(Const.electron.getPDGCode(), listtype, method, classification, path=path)
[docs]def stdMu(listtype=_defaultlist, method=None, classification=None, path=None):
""" Function to prepare one of several standardized types of muon lists.
See the documentation of `stdLep` for details.
It also accepts any of the legacy definitions
for the ``listtype`` parameter to fall back to the `stdCharged` behaviour:
* 'all'
* 'good'
* 'loosepid'
* 'loose'
* 'higheff'
* '95eff'
* '90eff'
* '85eff'
"""
if listtype in _stdnames + _effnames:
stdCharged("mu", listtype, path)
return
stdLep(Const.muon.getPDGCode(), listtype, method, classification, path=path)
[docs]def stdMostLikely(pidPriors=None, suffix='', custom_cuts='', path=None):
"""
Function to prepare most likely particle lists according to PID likelihood, refer to stdCharged for details
@param pidPriors list of 6 float numbers used to reweight PID likelihoods
@param suffix string added to the end of particle list names
@param custom_cuts custom selection cut string, if empty, standard track quality cuts will be applied
@param path modules are added to this path
"""
# Here we need basic track quality cuts to be applied,
# otherwise, we get a lot of badly reconstructed particles,
# which will end up filled as a random type
args = ''
if pidPriors is not None:
args = str(pidPriors)[1:-1] # remove brackets
trackQuality = 'thetaInCDCAcceptance and nCDCHits>20'
if custom_cuts != '':
trackQuality = custom_cuts
for name in _chargednames:
ma.fillParticleList(f'{name}+:{_mostLikelyList}{suffix}',
f'pidIsMostLikely({args}) > 0 and {trackQuality}', True, path=path)