import json
from pathlib import Path
import correctionlib
import awkward as ak
import numpy as np
from attrs import define, field
from analyzer.core.analysis_modules import (
AnalyzerModule,
IsSampleType,
MetadataExpr,
ModuleParameterSpec,
ParameterSpec,
)
from analyzer.core.columns import Column
from analyzer.utils.structure_tools import SimpleCache
from analyzer.modules.common.categories import addCategory
from analyzer.modules.common.axis import IntegerAxis
from analyzer.core.columns import addSelection
[docs]
class BNNEnsemble:
def __init__(self, num_inputs, mean, sigma, hidden_width, data):
hidden_length = hidden_width * (num_inputs + 1)
hidden_data = data[:, 0:hidden_length].reshape(len(data), -1, num_inputs + 1)
[docs]
self.output_bias = data[:, hidden_length]
[docs]
self.output_weights = data[:, hidden_length + 1 :]
[docs]
self.hidden_biases = hidden_data[:, :, 0]
[docs]
self.hidden_weights = hidden_data[:, :, 1:]
[docs]
def __call__(self, data):
data = data.to_numpy()
data = (data - self.mean) / self.sigma
l1 = self.hidden_biases + np.matvec(self.hidden_weights, data[:, None])
l2 = self.output_bias + np.vecdot(self.output_weights, np.tanh(l1))
vals = 1 / (1 + np.exp(-l2))
return (
np.median(vals, axis=1),
np.percentile(vals, 15.865, axis=1),
np.percentile(vals, 84.135, axis=1),
)
@staticmethod
[docs]
def fromFile(path):
with open(path, "r") as f:
data = json.load(f)
return BNNEnsemble(
data["num_inputs"],
np.array(data["mean"]),
np.array(data["sigma"]),
data["hidden_width"],
np.array(data["parameters"]),
)
@define
[docs]
class TriggerBNN(AnalyzerModule):
"""
Compute trigger efficiency weights using a BNN ensemble.
This analyzer evaluates a Bayesian Neural Network on HT and leading
fat jet pt to produce a trigger weight for MC samples.
Parameters
----------
base_path : str
Base directory where BNN JSON files are stored.
net_pattern : str
Pattern for the network filename, formatted with the era name.
weight_name : str, optional
Name of the output weight column, by default "trigger_eff".
should_run : MetadataExpr, optional
Condition to determine if the module should run. By default runs on MC samples.
"""
[docs]
weight_name: str = "trigger_eff"
[docs]
should_run: MetadataExpr = field(factory=lambda: IsSampleType("MC"))
__bnns: dict = field(factory=dict)
__bnn_res_cache: dict = field(factory=SimpleCache)
[docs]
def outputs(self, metadata):
return [Column(fields=("Weights", self.weight_name))]
[docs]
def neededResources(self, metadata):
return [self.base_path]
[docs]
def getParameterSpec(self, metadata):
return ModuleParameterSpec(
{
"triggereff-variation": ParameterSpec(
default_value="central",
possible_values=["central", "up", "down"],
tags={
"weight_variation",
},
),
}
)
[docs]
def getBNN(self, metadata):
name = metadata["era"]["name"]
path = str(Path(self.base_path) / self.net_pattern.format(era=name))
if path in self.__bnns:
return self.__bnns[path]
bnn = BNNEnsemble.fromFile(path)
self.__bnns[path] = bnn
return bnn
[docs]
def run(self, columns, params):
k = self.getKeyNoParams(columns)
systematic = params["triggereff-variation"]
if k in self.__bnn_res_cache:
central, down, up = self.__bnn_res_cache[k]
else:
ht = columns["HT"]
fj = columns["GoodFatJet"]
fjpt = fj[:, 0].pt
bnn = self.getBNN(columns.metadata)
out = bnn(ak.concatenate([x[:, np.newaxis] for x in [ht, fjpt]], axis=1))
central, down, up = out[0], out[1], out[2]
self.__bnn_res_cache[k] = (central, down, up)
if systematic == "central":
w = central
elif systematic == "up":
w = up
elif systematic == "down":
w = down
columns[Column(("Weights", self.weight_name))] = w
return columns, []
@define
[docs]
class TriggerBNNCorrection(AnalyzerModule):
"""
Compute trigger efficiency weights using correctionlib for BNN.
Parameters
----------
base_path : str
Base directory where correction files are stored.
correction_pattern : str
Pattern for the correction filename, formatted with the era name (e.g. bnn_correction_{era}.json.gz).
correction_name : str, optional
Name of the correction in the file, by default "BNN_Trigger_Efficiency".
weight_name : str, optional
Name of the output weight column, by default "trigger_eff".
should_run : MetadataExpr, optional
Condition to determine if the module should run. By default runs on MC samples.
"""
[docs]
correction_pattern: str
[docs]
correction_name: str = "trigger_eff"
[docs]
weight_name: str = "trigger_eff"
[docs]
should_run: MetadataExpr = field(factory=lambda: IsSampleType("MC"))
__corrections: dict = field(factory=dict)
[docs]
def outputs(self, metadata):
return [Column(fields=("Weights", self.weight_name))]
[docs]
def neededResources(self, metadata):
return [self.base_path]
[docs]
def getParameterSpec(self, metadata):
return ModuleParameterSpec(
{
"triggereff-variation": ParameterSpec(
default_value="central",
possible_values=["central", "up", "down", "disabled"],
tags={
"weight_variation",
},
),
}
)
[docs]
def getCorrection(self, metadata):
name = metadata["era"]["name"]
path = str(Path(self.base_path) / self.correction_pattern.format(era=name))
if path in self.__corrections:
return self.__corrections[path]
cset = correctionlib.CorrectionSet.from_file(path)
self.__corrections[path] = cset
return cset
[docs]
def run(self, columns, params):
systematic = params["triggereff-variation"]
if systematic == "disabled":
columns["Weights", self.weight_name] = ak.ones_like(columns["HT"])
return columns, []
syst_map = {"central": "nominal", "up": "up", "down": "down"}
corr_syst = syst_map.get(systematic, "nominal")
ht = columns["HT"]
fj = ak.pad_none(columns["GoodFatJet"], 1)[:, 0]
fjpt = ak.fill_none(fj.pt, 0)
fjmsd = ak.fill_none(fj.msoftdrop, 0)
cset = self.getCorrection(columns.metadata)
corr = cset[f"{self.correction_name}_{corr_syst}"]
w = corr.evaluate(ht, fjpt)
columns[Column(("Weights", self.weight_name))] = w
return columns, []
@define
[docs]
class MSDCleanerCategory(AnalyzerModule):
[docs]
def outputs(self, metadata):
return [Column(fields=("Categories", "PassMSDCleaner"))]
[docs]
def run(self, columns, params):
metadata = columns.metadata
trigger_names = metadata["era"]["trigger_names"]
ht = columns["HT"]
fj = ak.pad_none(columns["FatJet"], 1)[:, 0]
fjpt = ak.fill_none(fj.pt, 0)
fjmsd = ak.fill_none(fj.msoftdrop, 0)
hlt = columns["HLT"]
# bad = (
# ~hlt[trigger_names["HT"]]
# & ~hlt[trigger_names["AK8SingleJetPtNoTrim"]]
# & hlt[trigger_names["AK8SingleJetPt"]]
# & (fjpt > 550)
# & (fjmsd < 70)
# )
bad = (
~(ht > 1200)
& ~(fjpt > 700)
& (fjmsd < 70)
)
addCategory(columns, "PassMSDCleaner", ~bad, IntegerAxis("PassMSDCleaner", 0, 2))
return columns, []
@define
[docs]
class MSDCleanerSelection(AnalyzerModule):
[docs]
selection_name: str = "PassMSDCleaner"
[docs]
def outputs(self, metadata):
return [Column(fields=("Selection", self.selection_name))]
[docs]
def run(self, columns, params):
metadata = columns.metadata
ht = columns["HT"]
fj = ak.pad_none(columns["FatJet"], 1)[:, 0]
fjpt = ak.fill_none(fj.pt, 0)
fjmsd = ak.fill_none(fj.msoftdrop, 0)
hlt = columns["HLT"]
# Require events be on the softdrop plateau
bad = (
~((ht > 1050) | (fjpt > 550)) # probably didn't pass the other triggers
& ((fjmsd < 60) & (fjpt > 450))
)
addSelection(columns, self.selection_name, ~bad)
return columns, []