diff --git a/Corrections.py b/Corrections.py
index 0465869..3bfa724 100644
--- a/Corrections.py
+++ b/Corrections.py
@@ -408,26 +408,40 @@ def getNormalisationCorrections(
     ):
         isCentral = unc_source == central
         all_weights = []
-        if "MC_Lumi_pu" in self.to_apply:
+        lumi_weight_name = "weight_lumi"
+        if "lumi" in self.to_apply:
             lumi = self.global_params["luminosity"]
+            df = df.Define(lumi_weight_name, f"float({lumi})")
+            all_weights.append(lumi_weight_name)
 
+        crossSectionBranch = "weight_xs"
+        if "xs" in self.to_apply:
+            df = self.defineCrossSection(df, crossSectionBranch)
+            all_weights.append(crossSectionBranch)
+
+        gen_weight_name = "weight_gen"
+        if "gen" in self.to_apply:
             genWeight_def = (
-                "std::copysign<double>(1., genWeight)"
+                "std::copysign<float>(1.f, genWeight)"
                 if use_genWeight_sign_only
-                else "double(genWeight)"
+                else "genWeight"
             )
-            df = df.Define("genWeightD", genWeight_def)
-            crossSectionBranch = "crossSection"
-            df = self.defineCrossSection(df, crossSectionBranch)
-
-            shape_weights_dict = {(central, central): []}
-            if "pu" in self.to_apply:
-                df = self.pu.getWeight(
-                    df,
-                    shape_weights_dict=shape_weights_dict,
-                    return_variations=return_variations and isCentral,
-                )
+            df = df.Define(gen_weight_name, genWeight_def)
+            all_weights.append(gen_weight_name)
+
+        shape_weights_dict = {(central, central): []}
+        if "pu" in self.to_apply:
+            pu_enabled = self.to_apply["pu"].get("enabled", {}).get(self.stage, True)
+            df, weight_pu_branches = self.pu.getWeight(
+                df,
+                shape_weights_dict=shape_weights_dict,
+                return_variations=return_variations and isCentral,
+                return_list_of_branches=True,
+                enabled=pu_enabled,
+            )
+            all_weights.extend(weight_pu_branches)
 
+        if "base" in self.to_apply:
             for (
                 shape_unc_source,
                 shape_unc_scale,
@@ -440,22 +454,19 @@ def getNormalisationCorrections(
                 shape_weights_product = (
                     " * ".join(shape_weights) if len(shape_weights) > 0 else "1.0"
                 )
-                weight_name = (
-                    f"weight_{shape_unc_name}"
-                    if shape_unc_name != central
-                    else "weight_MC_Lumi_pu"
-                )
-                weight_rel_name = f"weight_MC_Lumi_{shape_unc_name}_rel"
-                weight_out_name = (
-                    weight_name if shape_unc_name == central else weight_rel_name
-                )
-                weight_formula = f"genWeightD * {lumi} * {crossSectionBranch} * {shape_weights_product} / {denomBranch}"
+                weight_name_central = "weight_base"
+                if shape_unc_name == central:
+                    weight_name = weight_name_central
+                    weight_out_name = weight_name
+                else:
+                    weight_name = f"{weight_name_central}_{shape_unc_name}"
+                    weight_out_name = f"{weight_name}_rel"
+                weight_formula = f"{gen_weight_name} * {lumi_weight_name} * {crossSectionBranch} * {shape_weights_product} / {denomBranch}"
                 df = df.Define(weight_name, f"static_cast<float>({weight_formula})")
-
                 if shape_unc_name != central:
                     df = df.Define(
                         weight_out_name,
-                        f"static_cast<float>(weight_{shape_unc_name}/weight_MC_Lumi_pu)",
+                        f"static_cast<float>({weight_name}/{weight_name_central})",
                     )
                 all_weights.append(weight_out_name)
 
diff --git a/pu.py b/pu.py
index b63c1d4..7dc8754 100644
--- a/pu.py
+++ b/pu.py
@@ -43,19 +43,33 @@ def __init__(self, period):
             )
             puWeightProducer.initialized = True
 
-    def getWeight(self, df, shape_weights_dict=None, return_variations=True):
+    def getWeight(
+        self,
+        df,
+        shape_weights_dict=None,
+        return_variations=True,
+        return_list_of_branches=False,
+        enabled=True,
+    ):
         sf_sources = puWeightProducer.uncSource if return_variations else []
+        branches = []
         for source in [central] + sf_sources:
             for scale in getScales(source):
-                branch_name = f"puWeight_{scale}"
-                df = df.Define(
-                    branch_name,
-                    f"""::correction::puCorrProvider::getGlobal().getWeight(
-                                ::correction::UncScale::{scale}, Pileup_nTrueInt)""",
-                )
+                branch_name = f"weight_pu_{scale}"
+                if enabled:
+                    df = df.Define(
+                        branch_name,
+                        f"""::correction::puCorrProvider::getGlobal().getWeight(
+                                    ::correction::UncScale::{scale}, Pileup_nTrueInt)""",
+                    )
+                    branches.append(branch_name)
+
                 key = (source, scale)
                 if shape_weights_dict:
                     if key not in shape_weights_dict:
                         shape_weights_dict[key] = []
                     shape_weights_dict[key].append(branch_name)
+
+        if return_list_of_branches:
+            return df, branches
         return df