From 2ef46627bd91560c4ee0f2bc1df2e905e5291c17 Mon Sep 17 00:00:00 2001 From: Mathis Frahm Date: Tue, 9 Jan 2024 11:17:03 +0100 Subject: [PATCH] extend to both solutions and add Neutrino variables --- hbw/config/ml_variables.py | 6 +-- hbw/config/variables.py | 19 ++++++++++ hbw/production/neutrino.py | 76 ++++++++++++++++++++++++++------------ 3 files changed, 75 insertions(+), 26 deletions(-) diff --git a/hbw/config/ml_variables.py b/hbw/config/ml_variables.py index 6a6acd3f..b61c8083 100644 --- a/hbw/config/ml_variables.py +++ b/hbw/config/ml_variables.py @@ -233,7 +233,7 @@ def add_ml_variables(config: od.Config) -> None: name=f"mli_{obj}_{var}", expression=f"mli_{obj}_{var}", binning=default_var_binning[var], - unit=default_var_unit.get(var, var), + unit=default_var_unit.get(var, "1"), x_title="{obj} {var}".format(obj=obj, var=var), ) @@ -243,7 +243,7 @@ def add_ml_variables(config: od.Config) -> None: name=f"mli_{obj}_{var}", expression=f"mli_{obj}_{var}", binning=default_var_binning[var], - unit=default_var_unit.get(var, var), + unit=default_var_unit.get(var, "1"), x_title="{obj} {var}".format(obj=obj, var=var), ) @@ -253,6 +253,6 @@ def add_ml_variables(config: od.Config) -> None: name=f"mli_{obj}_{var}", expression=f"mli_{obj}_{var}", binning=default_var_binning[var], - unit=default_var_unit.get(var, var), + unit=default_var_unit.get(var, "1"), x_title="{obj} {var}".format(obj=obj, var=var), ) diff --git a/hbw/config/variables.py b/hbw/config/variables.py index f3c7cda7..149b9ad2 100644 --- a/hbw/config/variables.py +++ b/hbw/config/variables.py @@ -13,6 +13,7 @@ from columnflow.columnar_util import EMPTY_FLOAT from hbw.util import call_once_on_config +from hbw.config.styling import default_var_binning, default_var_unit @call_once_on_config() @@ -138,6 +139,24 @@ def add_feature_variables(config: od.Config) -> None: ) +@call_once_on_config() +def add_neutrino_variables(config: od.Config) -> None: + """ + Adds variables to a *config* that are produced as part of the `neutrino_reconstruction` producer. + """ + + for obj in ["Neutrino1", "Neutrino2"]: + # pt and phi should be the same as MET, mass should always be 0 + for var in ["pt", "eta", "phi", "mass"]: + config.add_variable( + name=f"{obj}_{var}", + expression=f"{obj}.{var}", + binning=default_var_binning[var], + unit=default_var_unit.get(var, "1"), + x_title="{obj} {var}".format(obj=obj, var=var), + ) + + @call_once_on_config() def add_variables(config: od.Config) -> None: """ diff --git a/hbw/production/neutrino.py b/hbw/production/neutrino.py index a559b68b..51c3ca26 100644 --- a/hbw/production/neutrino.py +++ b/hbw/production/neutrino.py @@ -15,6 +15,7 @@ from hbw.util import four_vec from hbw.production.prepare_objects import prepare_objects +from hbw.config.variables import add_neutrino_variables np = maybe_import("numpy") @@ -29,13 +30,15 @@ @producer( uses=prepare_objects, - produces=four_vec("Neutrino"), + produces=four_vec(["Neutrino1", "Neutrino2"]), ) def neutrino_reconstruction(self: Producer, events: ak.Array, **kwargs) -> ak.Array: """ Producer to reconstruct a neutrino orignating from a leptonically decaying W boson. Assumes that Neutrino pt can be reconstructed via MET and that the W boson has been produced on-shell. + + TODO: reference """ # add behavior and define new collections (e.g. Lepton) events = self[prepare_objects](events, **kwargs) @@ -49,13 +52,11 @@ def neutrino_reconstruction(self: Producer, events: ak.Array, **kwargs) -> ak.Ar pz_l = events.Lepton.pz[:, 0] pt_nu = events.MET.pt - delta_phi = events.Lepton[:, 0].delta_phi(events.MET) - + delta_phi = abs(events.Lepton[:, 0].delta_phi(events.MET)) mu = w_mass**2 / 2 + pt_nu * pt_l * np.cos(delta_phi) - # pz_nu = A +- sqrt(B-C) + # Neutrino pz will be calculated as: pz_nu = A +- sqrt(B-C) A = mu * pz_l / pt_l**2 - B = mu**2 * pz_l**2 / pt_l**4 C = (E_l**2 * pt_nu**2 - mu**2) / pt_l**2 @@ -66,26 +67,55 @@ def neutrino_reconstruction(self: Producer, events: ak.Array, **kwargs) -> ak.Ar # complex solution -> take only the real part A, ) - # pz_nu_2 = ak.where( - # B - C >= 0, - # # solution is real - # A - np.sqrt(B - C), - # # complex solution -> take only the real part - # A, - # ) - p_nu_1 = np.sqrt(pt_nu**2 + pz_nu_1**2) - eta_nu_1 = np.log((p_nu_1 + pz_nu_1) / (p_nu_1 - pz_nu_1)) / 2 + pz_nu_2 = ak.where( + B - C >= 0, + # solution is real + A - np.sqrt(B - C), + # complex solution -> take only the real part + A, + ) - # store Neutrino 4 vector components - events["Neutrino"] = events.MET - events = set_ak_column_f32(events, "Neutrino.eta", eta_nu_1) + pz_nu_solutions = [pz_nu_1, pz_nu_2] + + for i, pz_nu in enumerate(pz_nu_solutions, start=1): + # calculate Neutrino eta to define the Neutrino 4-vector + p_nu_1 = np.sqrt(pt_nu**2 + pz_nu**2) + eta_nu_1 = np.log((p_nu_1 + pz_nu) / (p_nu_1 - pz_nu)) / 2 + + # store Neutrino 4 vector components + events[f"Neutrino{i}"] = events.MET + events = set_ak_column_f32(events, f"Neutrino{i}.eta", eta_nu_1) + + # sanity check: Neutrino pz should be the same as pz_nu within rounding errors + sanity_check_1 = ak.sum(abs(events[f"Neutrino{i}"].pz - pz_nu) > abs(events[f"Neutrino{i}"].pz) / 100) + if sanity_check_1: + logger.warning( + "Number of events with Neutrino.pz that differs from pz_nu by more than 1 percent: " + f"{sanity_check_1} (solution {i})", + ) + + # sanity check: reconstructing W mass should always (if B-C>0) give the input W mass (80.4 GeV) + W_on_shell = events[f"Neutrino{i}"] + events.Lepton[:, 0] + sanity_check_2 = ak.sum(abs(ak.where(B - C >= 0, W_on_shell.mass, w_mass) - w_mass) > 1) + if sanity_check_2: + logger.warning( + "Number of events with W mass from reconstructed Neutrino (real solutions only) that " + f"differs by more than 1 GeV from the input W mass: {sanity_check_2} (solution {i})", + ) + + # sanity check: for complex solutions, only the real part is considered -> both solutions should be identical + sanity_check_3 = ak.sum(ak.where(B - C <= 0, events.Neutrino1.eta - events.Neutrino2.eta, 0)) + if sanity_check_3: + raise Exception( + "When finding complex neutrino solutions, both reconstructed Neutrinos should be identical", + ) - # testing: Neutrino pz should be the same as pz_nu_1 within rounding errors + return events - # testing: reconstructing W mass should always (if B-C>0) give the input W mass (80.4 GeV) - # but currently, this is at 114 most of the times, so something is off - W_on_shell = events.Neutrino + events.Lepton[:, 0] - print(W_on_shell.mass) - return events +@neutrino_reconstruction.init +def neutrino_reconstruction_init(self: Producer) -> None: + + # add variable instances to config + add_neutrino_variables(self.config_inst)