From 6a914e5f85985b2a88c96bdc422c798738b679f8 Mon Sep 17 00:00:00 2001
From: "Marcel R." <github.riga@icloud.com>
Date: Wed, 12 Feb 2025 09:43:39 +0100
Subject: [PATCH] Add back top pt weight.

---
 hbt/config/configs_hbt.py | 25 +++++++++++++++++++++----
 hbt/production/default.py |  9 +++++++++
 hbt/selection/default.py  |  5 +++++
 modules/columnflow        |  2 +-
 4 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/hbt/config/configs_hbt.py b/hbt/config/configs_hbt.py
index 0ec8bfac..6a2b5682 100644
--- a/hbt/config/configs_hbt.py
+++ b/hbt/config/configs_hbt.py
@@ -958,6 +958,25 @@ def if_not_era(*, values: list[str | None] | None = None, **kwargs) -> list[str]
             discriminator="btagPNetB",
         )
 
+    ################################################################################################
+    # dataset / process specific methods
+    ################################################################################################
+
+    # top pt reweighting
+    # https://twiki.cern.ch/twiki/bin/view/CMS/TopPtReweighting?rev=31
+    from columnflow.production.cms.top_pt_weight import TopPtWeightConfig
+    cfg.x.top_pt_weight = TopPtWeightConfig(
+        params={
+            "a": 0.0615,
+            "a_up": 0.0615 * 1.5,
+            "a_down": 0.0615 * 0.5,
+            "b": -0.0005,
+            "b_up": -0.0005 * 1.5,
+            "b_down": -0.0005 * 0.5,
+        },
+        pt_max=500.0,
+    )
+
     ################################################################################################
     # shifts
     ################################################################################################
@@ -1322,10 +1341,8 @@ def add_external(name, value):
 
     # define per-dataset event weights
     for dataset in cfg.datasets:
-        # skipped for now
-        # if dataset.has_tag("ttbar"):
-        #     dataset.x.event_weights = {"top_pt_weight": get_shifts("top_pt")}
-        pass
+        if dataset.has_tag("has_top"):
+            dataset.x.event_weights = {"top_pt_weight": get_shifts("top_pt")}
 
     cfg.x.shift_groups = {
         "jec": [
diff --git a/hbt/production/default.py b/hbt/production/default.py
index b39ba3ed..a889f336 100644
--- a/hbt/production/default.py
+++ b/hbt/production/default.py
@@ -9,6 +9,7 @@
 from columnflow.production.categories import category_ids
 from columnflow.production.cms.electron import electron_weights
 from columnflow.production.cms.muon import muon_weights
+from columnflow.production.cms.top_pt_weight import top_pt_weight
 from columnflow.util import maybe_import
 
 from hbt.production.weights import (
@@ -78,6 +79,10 @@ def default(self: Producer, events: ak.Array, **kwargs) -> ak.Array:
         if self.has_dep(trigger_weights):
             events = self[trigger_weights](events, **kwargs)
 
+        # top pt weight
+        if self.has_dep(top_pt_weight):
+            events = self[top_pt_weight](events, **kwargs)
+
     return events
 
 
@@ -85,6 +90,10 @@ def default(self: Producer, events: ak.Array, **kwargs) -> ak.Array:
 def default_init(self: Producer) -> None:
     if self.produce_weights:
         weight_producers = {tau_weights, electron_weights, muon_weights, trigger_weights}
+
+        if (dataset_inst := getattr(self, "dataset_inst", None)) and dataset_inst.has_tag("ttbar"):
+            weight_producers.add(top_pt_weight)
+
         self.uses |= weight_producers
         self.produces |= weight_producers
 
diff --git a/hbt/selection/default.py b/hbt/selection/default.py
index 99f650fa..c0f3c141 100644
--- a/hbt/selection/default.py
+++ b/hbt/selection/default.py
@@ -22,6 +22,7 @@
 from columnflow.production.cms.pileup import pu_weight
 from columnflow.production.cms.pdf import pdf_weights
 from columnflow.production.cms.scale import murmuf_weights
+from columnflow.production.cms.top_pt_weight import gen_parton_top
 from columnflow.production.util import attach_coffea_behavior
 from columnflow.columnar_util import full_like
 from columnflow.util import maybe_import
@@ -261,6 +262,10 @@ def default_init(self: Selector) -> None:
             # save it as an attribute
             setattr(self, prod_name, prod)
 
+    if self.dataset_inst.has_tag("ttbar"):
+        self.uses.add(gen_parton_top)
+        self.produces.add(gen_parton_top)
+
 
 empty = default.derive("empty", cls_dict={})
 
diff --git a/modules/columnflow b/modules/columnflow
index e9906f69..ce7fddd1 160000
--- a/modules/columnflow
+++ b/modules/columnflow
@@ -1 +1 @@
-Subproject commit e9906f690f320c58b8a5c04245fd95d96fced443
+Subproject commit ce7fddd1d3a42342b7f30f0badcb563b538db964