from __future__ import annotations
import functools as ft
from analyzer.core.columns import addSelection
import operator as op
from analyzer.core.analysis_modules import (
AnalyzerModule,
ModuleParameterSpec,
ParameterSpec,
MetadataExpr,
MetadataAnd,
IsRun,
IsSampleType,
)
from analyzer.core.columns import Column
from attrs import define, field
import correctionlib
import correctionlib
import awkward as ak
import correctionlib.convert
from analyzer.core.adl import ADLBlock, ADLStatement
@define
[docs]
class PileupSF(AnalyzerModule):
"""
Apply pileup scale factors to Monte Carlo events.
This analyzer evaluates the pileup weight based on the number of true
interactions in the event and applies systematic variations if requested.
Parameters
----------
weight_name : str, optional
Name of the column where the pileup weights are stored, by default "pileup_sf".
should_run : MetadataExpr, optional
Condition to determine if the module should run. By default runs only
on MC samples.
"""
[docs]
weight_name: str = "pileup_sf"
[docs]
should_run: MetadataExpr = field(factory=lambda: IsSampleType("MC"))
__corrections: dict = field(factory=dict)
[docs]
def getParameterSpec(self, metadata):
return ModuleParameterSpec(
{
"pusf-variation": ParameterSpec(
default_value="nominal",
possible_values=["nominal", "up", "down", "disabled"],
tags={
"weight_variation",
},
),
},
)
[docs]
def run(self, columns, params):
if params["pusf-variation"] == "disabled":
columns[Column(("Weights", self.weight_name))] = ak.ones_like(
columns["Pileup.nTrueInt"], dtype=float
)
return columns, []
corr = self.getCorrection(columns.metadata)
n_pu = columns["Pileup.nTrueInt"]
correction = corr.evaluate(n_pu, params["pusf-variation"])
columns[Column(("Weights", self.weight_name))] = correction
return columns, []
[docs]
def getCorrection(self, metadata):
file_path = metadata["era"]["pileup_scale_factors"]["file"]
name = metadata["era"]["pileup_scale_factors"]["name"]
key = (name, file_path)
if key in self.__corrections:
return self.__corrections[key]
cset = correctionlib.CorrectionSet.from_file(file_path)
ret = cset[name]
self.__corrections[key] = ret
return ret
[docs]
def outputs(self, metadata):
return [Column(fields=("Weights", self.weight_name))]
@define
[docs]
class L1PrefiringSF(AnalyzerModule):
"""
Apply L1 prefiring scale factors to Monte Carlo events.
This analyzer retrieves the L1 prefiring weight and adds it to a named weight column, optionally applying systematic variations.
Parameters
----------
weight_name : str, optional
Name of the output weight column, by default "l1_prefiring".
should_run : MetadataExpr, optional
Condition to determine if the module should run. By default, it runs
on MC samples for Run 2.
"""
[docs]
weight_name: str = "l1_prefiring"
[docs]
should_run: MetadataExpr = field(
factory=lambda: MetadataAnd([IsSampleType("MC"), IsRun(2)])
)
__corrections: dict = field(factory=dict)
[docs]
def getParameterSpec(self, metadata):
return ModuleParameterSpec(
{
"l1prefiring-variation": ParameterSpec(
default_value="Nom",
possible_values=["Nom", "Up", "Dn"],
tags={"weight_variation"},
),
}
)
[docs]
def run(self, columns, params):
variation = params["l1prefiring-variation"]
if "L1PreFiringWeight" in columns.fields:
columns["Weights", self.weight_name] = columns["L1PreFiringWeight"][
variation
]
else:
columns["Weights", self.weight_name] = ak.ones_like(
columns["run"], dtype=float
)
return columns, []
[docs]
def outputs(self, metadata):
return [Column(fields=("Weights", self.weight_name))]
@define
[docs]
class PosNegGenWeight(AnalyzerModule):
[docs]
should_run: MetadataExpr = field(factory=lambda: IsSampleType("MC"))
[docs]
weight_name: str = "pos_neg_weight"
[docs]
def run(self, columns, params):
gw = ak.where(columns["genWeight"] > 0, 1.0, -1.0)
columns["Weights", self.weight_name] = gw
return columns, []
[docs]
def outputs(self, metadata):
return [Column(fields=("Weights", self.weight_name))]
@define
[docs]
class GoldenLumi(AnalyzerModule):
"""
Apply a golden JSON luminosity selection for data events.
This analyzer filters events to only include those that are within
certified good luminosity sections for the given data-taking era.
Parameters
----------
selection_name : str, optional
Name of the selection column to store the golden JSON filter,
by default "golden_lumi".
should_run : MetadataExpr, optional
Condition to determine if the module should run. By default, only
runs on real data samples.
Notes
-----
- The certified luminosity sections are read from the metadata for
the given era under "golden_json".
"""
[docs]
selection_name: str = "golden_lumi"
[docs]
should_run: MetadataExpr = field(factory=lambda: IsSampleType("Data"))
[docs]
def outputs(self, metadata):
return [Column(("Selection", self.selection_name))]
[docs]
def run(self, columns, params):
import coffea.lumi_tools as ltools
metadata = columns.metadata
lumi_json = metadata["era"]["golden_json"]
lmask = ltools.LumiMask(lumi_json)
addSelection(
columns,
self.selection_name,
lmask(columns["run"], columns["luminosityBlock"]),
)
return columns, []
@define
[docs]
class NoiseFilter(AnalyzerModule):
"""
Apply standard noise filters to data events.
This analyzer combines multiple event-level noise flags and produces a
single selection column indicating events that pass all required noise
filters.
Parameters
----------
selection_name : str, optional
Name of the selection column to store the combined noise filter,
by default "noise_filters".
"""
[docs]
selection_name: str = "noise_filters"
[docs]
def outputs(self, metadata):
return [Column(("Selection", self.selection_name))]
[docs]
def run(self, columns, params):
metadata = columns.metadata
noise_flags = metadata["era"]["noise_filters"]
sel = ft.reduce(op.and_, [columns["Flag"][x] for x in noise_flags])
addSelection(columns, self.selection_name, sel)
return columns, []
[docs]
def adlExport(self, metadata):
noise_flags = metadata["era"]["noise_filters"]
statements = [
ADLStatement("select", f"Flag_{flag} == True") for flag in noise_flags
]
return [ADLBlock(block_type="region_statement", name="", statements=statements)]