Implementation of diffuse flux in the background

flarestack/cluster/submitter.py

@@ -95,8 +95,12 @@ def submit(self, mh_dict):
         if self.remove_old_results:
             self._clean_injection_values_and_pickled_results(self.mh_dict["name"])
         if self.use_cluster:
-            if mh_dict["inj_dict"]["injector_name"] == "low_memory_injector":
-                make_band_mask(mh_dict=copy.deepcopy(mh_dict))
+            try:
+                inj_name = mh_dict["inj_dict"]["injector_name"]
+                if inj_name == "low_memory_injector":
+                    make_band_mask(mh_dict=copy.deepcopy(mh_dict))
+            except KeyError:
+                pass
 
             self.submit_cluster(mh_dict)
         else:

flarestack/core/injector.py

@@ -15,6 +15,7 @@
 from flarestack.core.time_pdf import TimePDF, read_t_pdf_dict
 from flarestack.shared import band_mask_cache_name, k_to_flux
 from flarestack.utils.catalogue_loader import calculate_source_weight
+from flarestack.utils.make_SoB_splines import SoB_splines
 
 if TYPE_CHECKING:
     from flarestack.data import Season, SeasonWithMC
@@ -87,14 +88,35 @@ def __init__(self, season: "Season", sources: Table, **kwargs) -> None:
         if len(sources) > 0:
             self.weight_scale = calculate_source_weight(self.sources)
 
+        # get SoB splines dict
+        try:
+            sob_dict = kwargs["sob_dict"]
+        except KeyError:
+            sob_dict = {}
+
+        sob_dict.setdefault("smoothing_order", "flarestack")
+        sob_dict.setdefault("gamma_precision", "flarestack")
+        self.sob_name = sob_dict.setdefault("bkg_model_name", "no_difffuse")
+        self.sob = SoB_splines.create(season, sob_dict)
+
         try:
             self.sig_time_pdf = TimePDF.create(
                 kwargs["injection_sig_time_pdf"], season.get_time_pdf()
             )
             # self.bkg_time_pdf = TimePDF.create(kwargs["injection_bkg_time_pdf"],
             #                                    season.get_time_pdf())
             self.energy_pdf = EnergyPDF.create(kwargs["injection_energy_pdf"])
-            self.spatial_pdf = SpatialPDF(kwargs["injection_spatial_pdf"], season)
+
+            # bkg_spatial_pdf = zenith_spline if not specified
+            # for this case SoB dict is needed
+            if "bkg_spatial_pdf" not in kwargs["injection_spatial_pdf"].keys():
+                self.spatial_pdf = SpatialPDF(
+                    kwargs["injection_spatial_pdf"], season, self.sob
+                )
+            else:
+                self.spatial_pdf = SpatialPDF(
+                    kwargs["injection_spatial_pdf"], season, None
+                )
         except KeyError:
             raise Exception(
                 "Injection Arguments missing. \n "
@@ -166,9 +188,19 @@ def create_dataset(
         :param angular_error_modifier: AngularErrorModifier to change angular errors
         :return: Simulated dataset
         """
-        bkg_events, n_excluded = self.season.simulate_background(
-            self.sources, self.spatial_box_width
-        )
+        if self.sob_name != "no_difffuse":
+            try:
+                bkg_events, n_excluded = self.season.simulate_bkg_with_diffuse(
+                    self.sources, self.spatial_box_width, self.sob.bkg_weights
+                )
+            except NotImplementedError:
+                bkg_events, n_excluded = self.season.simulate_background(
+                    self.sources, self.spatial_box_width
+                )
+        else:
+            bkg_events, n_excluded = self.season.simulate_background(
+                self.sources, self.spatial_box_width
+            )
 
         if scale > 0.0:
             sig_events = self.inject_signal(scale)
@@ -729,8 +761,7 @@ class MockUnblindedInjector:
 
     def __init__(self, season: "Season", sources=np.nan, **kwargs):
         self.season = season
-        self._raw_data = season.get_exp_data()
-        season.load_background_model()
+        self.season.use_data_for_trials()
 
     def create_dataset(
         self, scale: float, angular_error_modifier=None
@@ -742,11 +773,18 @@ def create_dataset(
         seed = int(123456)
         np.random.seed(seed)
 
-        simulated_data, n_excluded = self.season.simulate_background(Table(), None)
+        # copy Season.simulate_background()
+        scrambled_data = self.season.pseudo_background()
+        if self.season._subselection_fraction is not None:
+            scrambled_data = np.random.choice(
+                scrambled_data,
+                int(len(scrambled_data) * self.season._subselection_fraction),
+            )
+
         if angular_error_modifier is not None:
-            simulated_data = angular_error_modifier.pull_correct_static(simulated_data)
+            scrambled_data = angular_error_modifier.pull_correct_static(scrambled_data)
 
-        return simulated_data, n_excluded
+        return scrambled_data, 0
 
 
 class TrueUnblindedInjector:

flarestack/core/llh.py

@@ -20,20 +20,13 @@
 from flarestack.shared import (
     SoB_spline_path,
     acceptance_path,
-    default_gamma_precision,
-    default_smoothing_order,
     llh_energy_hash_pickles,
 )
 from flarestack.utils.create_acceptance_functions import (
     dec_range,
     make_acceptance_season,
 )
-from flarestack.utils.make_SoB_splines import (
-    create_2d_ratio_hist,
-    load_spline,
-    make_2d_spline_from_hist,
-    make_individual_spline_set,
-)
+from flarestack.utils.make_SoB_splines import SoB_splines, get_gamma_precision
 
 logger = logging.getLogger(__name__)
 
@@ -100,7 +93,26 @@ def __init__(self, season, sources, llh_dict):
         self.season = season
         self.sources = sources
         self.llh_dict = llh_dict
-        self.spatial_pdf = SpatialPDF(llh_dict["llh_spatial_pdf"], season)
+
+        try:
+            sob_dict = llh_dict["sob_dict"]
+            sob_dict.setdefault("smoothing_order", "flarestack")
+            sob_dict.setdefault("gamma_precision", "flarestack")
+        except KeyError:
+            sob_dict = dict()
+            smoothing = llh_dict.get("smoothing_order", "flarestack")
+            precision = llh_dict.get("gamma_precision", "flarestack")
+            sob_dict = {"smoothing_order": smoothing, "gamma_precision": precision}
+
+        self.sob_name = sob_dict.setdefault("bkg_model_name", "no_difffuse")
+        self.sob = SoB_splines.create(season, sob_dict)
+        self.smoothing_order = self.sob.smoothing_order
+        self.precision = get_gamma_precision(self.sob.gamma_precision)
+
+        if "bkg_spatial_pdf" not in llh_dict["llh_spatial_pdf"].keys():
+            self.spatial_pdf = SpatialPDF(llh_dict["llh_spatial_pdf"], season, self.sob)
+        else:
+            self.spatial_pdf = SpatialPDF(llh_dict["llh_spatial_pdf"], season, None)
         self.spatial_box_width = llh_dict.get(
             "spatial_box_width", default_spacial_box_width
         )
@@ -423,7 +435,7 @@ def create_energy_function(self):
         del exp
 
         # return lambda x: data_rate
-        return lambda x: np.exp(self.bkg_spatial(np.sin(x))) * data_rate
+        return lambda x: self.spatial_pdf.background_spatial(x) * data_rate
 
     def create_llh_function(
         self, data: Table, n_excluded: int, pull_corrector, weight_f=None
@@ -522,22 +534,6 @@ def __init__(self, season, sources, llh_dict):
                 "again."
             )
 
-        # defines the order of the splines used in the building of the energy PDF
-        self.smoothing_order = None
-        smoothing_order = llh_dict.get("smoothing_order", "flarestack")
-        if isinstance(smoothing_order, str):
-            self.smoothing_order = default_smoothing_order[smoothing_order]
-        else:
-            self.smoothing_order = smoothing_order
-
-        # used to construct the support points in gamma when the energy PDF is built
-        self.precision = None
-        precision = llh_dict.get("gamma_precision", "flarestack")
-        if isinstance(precision, str):
-            self.precision = default_gamma_precision[precision]
-        else:
-            self.precision = precision
-
         LLH.__init__(self, season, sources, llh_dict)
 
     def create_energy_functions(self):
@@ -548,7 +544,9 @@ def create_energy_functions(self):
         :return: Acceptance function, energy_weighting_function
         """
 
-        SoB_path, acc_path = llh_energy_hash_pickles(self.llh_dict, self.season)
+        SoB_path, acc_path = llh_energy_hash_pickles(
+            self.llh_dict, self.season, self.sob_name
+        )
 
         # Set up acceptance function, creating values if they have not been
         # created before
@@ -577,7 +575,7 @@ def acc_f(source, params=None):
         with open(SoB_path, "rb") as f:
             [dec_vals, ratio_hist] = pickle.load(f)
 
-        spline = make_2d_spline_from_hist(
+        spline = self.sob.make_2d_spline_from_hist(
             np.array(ratio_hist), dec_vals, self.season.log_e_bins, self.smoothing_order
         )
 
@@ -636,12 +634,12 @@ def create_energy_weighting_function(self, SoB_path):
 
         # dec_range = self.season["sinDec bins"]
 
-        ratio_hist = create_2d_ratio_hist(
+        ratio_hist = self.sob.create_2d_ratio_hist(
             exp=self.season.get_background_model(),
             mc=self.season.get_pseudo_mc(),
-            sin_dec_bins=dec_range,
+            sin_dec_bins=self.season.sin_dec_bins,
             log_e_bins=self.season.log_e_bins,
-            weight_function=self.energy_pdf.weight_mc,
+            weight_f=self.energy_pdf.weight_mc,
         )
 
         try:
@@ -652,7 +650,7 @@ def create_energy_weighting_function(self, SoB_path):
         logger.info("Saving to {0}".format(SoB_path))
 
         with open(SoB_path, "wb") as f:
-            pickle.dump([dec_range, ratio_hist], f)
+            pickle.dump([self.season.sin_dec_bins, ratio_hist], f)
 
     def create_kwargs(
         self,
@@ -742,11 +740,7 @@ def __init__(self, season, sources, llh_dict):
         # Sets precision for energy SoB
         # self.precision = .1
 
-        self.SoB_spline_2Ds = load_spline(
-            self.season,
-            smoothing_order=self.smoothing_order,
-            gamma_precision=self.precision,
-        )
+        self.SoB_spline_2Ds = self.sob.load_spline()
 
         if self.SoB_spline_2Ds:
             logger.debug("Loaded {0} splines.".format(len(self.SoB_spline_2Ds)))
@@ -778,11 +772,7 @@ def create_energy_functions(self):
         :return: Acceptance function, energy_weighting_function
         """
 
-        SoB_path = SoB_spline_path(
-            self.season,
-            smoothing_order=self.smoothing_order,
-            gamma_precision=self.precision,
-        )
+        SoB_path = SoB_spline_path(self.season, **self.sob.sob_dict)
         acc_path = acceptance_path(self.season)
 
         # Set up acceptance function, creating values if they have not been
@@ -807,12 +797,7 @@ def create_energy_functions(self):
         # Checks if energy weighting functions have been created
 
         if not os.path.isfile(SoB_path):
-            make_individual_spline_set(
-                self.season,
-                SoB_path,
-                smoothing_order=self.smoothing_order,
-                gamma_precision=self.precision,
-            )
+            self.sob.make_individual_spline_set(self.season, SoB_path)
 
         return acc_f, None
 
@@ -1708,8 +1693,12 @@ def generate_dynamic_flare_class(season, sources, llh_dict):
     except KeyError:
         raise KeyError("No LLH specified.")
 
-    # Set up dynamic inheritance
+    if mh_name in ["spatial", "fixed_energy"]:
+        raise NotImplementedError(
+            "Need an LLh that creates SoB energy cache, choose another"
+        )
 
+    # Set up dynamic inheritance
     try:
         ParentLLH = LLH.subclasses[mh_name]
     except KeyError:

flarestack/core/minimisation.py

@@ -222,10 +222,18 @@ def __init__(self, mh_dict):
                 "gamma" in self.llh_dict["llh_energy_pdf"].keys()
             ), "Running trials with fixed gamma but no gamma passed in the llh energy pdf"
             self.llh_gamma = self.llh_dict["llh_energy_pdf"]["gamma"]
-            if self.llh_gamma != self.inj_dict["injection_energy_pdf"]["gamma"]:
-                logger.warning(
-                    f"Fixing gamma to {self.llh_gamma} for llh but injection is with {self.inj_dict['injection_energy_pdf']['gamma']}"
-                )
+
+        # check if SoB spline dict is same in the llh & inj dicts
+        if (
+            "sob_dict" in self.inj_dict
+            and "sob_dict" in self.llh_dict
+            and self.llh_dict["sob_dict"] != self.inj_dict["sob_dict"]
+        ):
+            logger.warning(
+                "Passed 'sob_dict' aren't the same for injector and llh: "
+                + f"inj = {self.inj_dict['sob_dict']}, llh = {self.llh_dict['sob_dict']}"
+                + " Proceed with caution!"
+            )
 
         p0, bounds, names = self.return_parameter_info(mh_dict)