From ad5915afe15fa4989c6db0737a9a708d0316e7d7 Mon Sep 17 00:00:00 2001
From: Jonathon Langford <jonathon.langford@cern.ch>
Date: Tue, 19 Dec 2023 16:33:18 +0100
Subject: [PATCH] Enabling STXS model to work with inclusive cards

---
 python/STXS.py       |  8 ++++----
 python/STXSModels.py | 20 ++++++++++++++++++--
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/python/STXS.py b/python/STXS.py
index 29848a361da..7e864b406c3 100644
--- a/python/STXS.py
+++ b/python/STXS.py
@@ -9,14 +9,14 @@
 
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 stage0_procs = {
-    "ggH": ["ggH"],
+    "ggH": ["ggH","bbH","ggZH_had"],
     "qqH": ["qqH"],
     "WH_lep": ["WH_lep"],
     "ZH_lep": ["ZH_lep"],
-    #'ggZH_lep':['ggZH_lep']
+    'ggZH_lep':['ggZH_lep'],
     "VH_had": ["WH_had", "ZH_had"],
-    "ttH": ["ttH"]  # ,
-    #'other':['tHq','tHW','bbH']
+    "ttH": ["ttH"],
+    'other':['tHq','tHW','bbH',"WH","ZH"]
 }
 
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/python/STXSModels.py b/python/STXSModels.py
index 4885b82196d..0c2548513a5 100644
--- a/python/STXSModels.py
+++ b/python/STXSModels.py
@@ -2,6 +2,7 @@
 
 import fnmatch
 import json
+import re
 
 from HiggsAnalysis.CombinedLimit.PhysicsModel import *
 from HiggsAnalysis.CombinedLimit.SMHiggsBuilder import SMHiggsBuilder
@@ -29,7 +30,7 @@
     "hcc": "cc",
     "htt": "tautau",
     "hmm": "mumu",
-    "hzg": "gamgam",
+    "hzg": "Zgam",
     "hgluglu": "gluglu",
     "hinv": "inv",
 }
@@ -82,6 +83,7 @@ def __init__(self, denominator="WW"):
         self.denominator = denominator
         self.mergeBins = False
         self.mergeJson = ""
+        self.addStage0 = False
 
     def preProcessNuisances(self, nuisances):
         # add here any pre-processed nuisances such as constraint terms for the mass profiling?
@@ -101,6 +103,9 @@ def setPhysicsOptions(self, physOptions):
             if po.startswith("mergejson="):
                 self.mergeBins = True
                 self.mergeJson = po.replace("mergejson=", "")
+            if po.startswith("addStage0="):
+                self.addStage0 = po.replace("addStage0=", "") in ["yes", "1", "Yes", "True", "true"]
+
 
     def doMH(self):
         if self.floatMass:
@@ -217,7 +222,6 @@ def __init__(self):
         STXSBaseModel.__init__(self)  # not using 'super(x,self).__init__' since I don't understand it
         self.POIs = "mu"
         from HiggsAnalysis.CombinedLimit.STXS import stage1_2_procs, stage1_2_fine_procs
-
         self.stage1_2_fine_procs = stage1_2_fine_procs
         self.PROCESSES = [x for v in six.itervalues(stage1_2_procs) for x in v]
         self.FINEPROCESSES = [x for v in six.itervalues(stage1_2_fine_procs) for x in v]
@@ -240,6 +244,14 @@ def doParametersOfInterest(self):
                 self.modelBuilder.stringout = json.dumps(self.mergeSchemes)
             f.close()
 
+        # Add stage 0 processes
+        if self.addStage0:
+            from HiggsAnalysis.CombinedLimit.STXS import stage0_procs
+            PROCESSES_STAGE0 = [x for v in six.itervalues(stage0_procs) for x in v]
+            for proc_stage0 in PROCESSES_STAGE0:
+                if proc_stage0 not in self.PROCESSES: self.PROCESSES.append(proc_stage0)
+
+
         allProds = []
         for registered_proc in self.PROCESSES:
             P = registered_proc
@@ -314,6 +326,10 @@ def setup(self):
                 self.modelBuilder.out.function("scaling_%s_%s_13TeV" % (P, D)).Print("")
 
     def getHiggsSignalYieldScale(self, production, decay, energy):
+        # Catch for H->Zgam: taken from STXStoSMEFT model
+        if (decay == "Zgam") | ("bkg" in production):
+            production = production.split("_")[0]
+
         for regproc in self.PROCESSES:
             if fnmatch.fnmatch(production, regproc):
                 return "scaling_%s_%s_%s" % (regproc, decay, energy)