From e33015ace150288646673d9ba0e6f114b8e0d96a Mon Sep 17 00:00:00 2001
From: Mason Proffitt <mason.proffitt@cern.ch>
Date: Tue, 23 Jul 2024 14:44:17 +0200
Subject: [PATCH] add background uncertainty #45

---
 src/abcd_pyhf/abcd.py      | 11 ++++++--
 src/abcd_pyhf/pyhf_util.py | 51 ++++++++++++++++++++++----------------
 2 files changed, 39 insertions(+), 23 deletions(-)

diff --git a/src/abcd_pyhf/abcd.py b/src/abcd_pyhf/abcd.py
index f63fe19..d3acaf3 100644
--- a/src/abcd_pyhf/abcd.py
+++ b/src/abcd_pyhf/abcd.py
@@ -39,6 +39,8 @@ class ABCD:
         arbitrary.
     signal_uncertainty : float, optional
         Total fractional uncertainty of the signal yields
+    background_uncertainty : float, optional
+        Total fractional uncertainty of the background estimation
 
     Attributes
     ----------
@@ -65,11 +67,12 @@ class ABCD:
     """
 
     def __init__(
-        self, observed_yields, signal_yields=None, signal_uncertainty=None
+        self, observed_yields, signal_yields=None, signal_uncertainty=None, background_uncertainty=None,
     ):
         self._observed_yields = observed_yields
         self._signal_yields = signal_yields
         self._signal_uncertainty = signal_uncertainty
+        self._background_uncertainty = background_uncertainty
 
     @property
     def observed_yields(self):
@@ -83,6 +86,10 @@ def signal_yields(self):
     def signal_uncertainty(self):
         return self._signal_uncertainty
 
+    @property
+    def background_uncertainty(self):
+        return self._background_uncertainty
+
     @property
     def blinded(self):
         return signal_region not in self.observed_yields
@@ -90,7 +97,7 @@ def blinded(self):
     @property
     def model(self):
         return create_model(
-            self.signal_yields, self.signal_uncertainty, self.blinded
+            self.signal_yields, self.signal_uncertainty, self.background_uncertainty, self.blinded
         )
 
     @property
diff --git a/src/abcd_pyhf/pyhf_util.py b/src/abcd_pyhf/pyhf_util.py
index 981202c..2b8d727 100644
--- a/src/abcd_pyhf/pyhf_util.py
+++ b/src/abcd_pyhf/pyhf_util.py
@@ -13,7 +13,8 @@
 all_regions = [signal_region] + control_regions
 
 poi_name = 'mu'
-signal_uncertainty_name = 'systematic_uncertainty'
+signal_uncertainty_name = 'signal_uncertainty'
+bkg_uncertainty_name = 'background_uncertainty'
 
 bkg_normalization_name = 'mu_b'
 bkg_scale_factor_1_name = 'tau_B'
@@ -24,26 +25,7 @@ def normfactor(name):
     return {'name': name, 'type': 'normfactor', 'data': None}
 
 
-bkg_normalization = normfactor(bkg_normalization_name)
-bkg_modifiers = {
-    signal_region: [bkg_normalization],
-    control_regions[0]: [
-        bkg_normalization,
-        normfactor(bkg_scale_factor_1_name),
-    ],
-    control_regions[1]: [
-        bkg_normalization,
-        normfactor(bkg_scale_factor_2_name),
-    ],
-    control_regions[2]: [
-        bkg_normalization,
-        normfactor(bkg_scale_factor_1_name),
-        normfactor(bkg_scale_factor_2_name),
-    ],
-}
-
-
-def create_model(signal_yield, signal_uncertainty, blinded):
+def create_model(signal_yield, signal_uncertainty, bkg_uncertainty, blinded):
     signal_modifiers = [
         normfactor('mu'),
         {
@@ -55,6 +37,33 @@ def create_model(signal_yield, signal_uncertainty, blinded):
             },
         },
     ]
+    bkg_normalization = normfactor(bkg_normalization_name)
+    bkg_modifiers = {
+        signal_region: [
+            bkg_normalization,
+            {
+                'name': bkg_uncertainty_name,
+                'type': 'normsys',
+                'data': {
+                    'hi': 1 + bkg_uncertainty,
+                    'lo': 1 - bkg_uncertainty,
+                },
+            },
+        ],
+        control_regions[0]: [
+            bkg_normalization,
+            normfactor(bkg_scale_factor_1_name),
+        ],
+        control_regions[1]: [
+            bkg_normalization,
+            normfactor(bkg_scale_factor_2_name),
+        ],
+        control_regions[2]: [
+            bkg_normalization,
+            normfactor(bkg_scale_factor_1_name),
+            normfactor(bkg_scale_factor_2_name),
+        ],
+    }
     regions_to_include = control_regions if blinded is True else all_regions
     spec = {
         'channels': [