diff --git a/flarestack/cluster/submitter.py b/flarestack/cluster/submitter.py index a9056ab1..8df5097e 100644 --- a/flarestack/cluster/submitter.py +++ b/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: diff --git a/flarestack/core/injector.py b/flarestack/core/injector.py index b01ac3c5..bce40923 100644 --- a/flarestack/core/injector.py +++ b/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,6 +88,17 @@ 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() @@ -94,7 +106,17 @@ def __init__(self, season: "Season", sources: Table, **kwargs) -> None: # 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: diff --git a/flarestack/core/llh.py b/flarestack/core/llh.py index 63aa37fe..eaf28b82 100644 --- a/flarestack/core/llh.py +++ b/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: diff --git a/flarestack/core/minimisation.py b/flarestack/core/minimisation.py index e4750def..b000541f 100644 --- a/flarestack/core/minimisation.py +++ b/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) diff --git a/flarestack/core/spatial_pdf.py b/flarestack/core/spatial_pdf.py index f61325ea..b9797ca3 100644 --- a/flarestack/core/spatial_pdf.py +++ b/flarestack/core/spatial_pdf.py @@ -13,7 +13,7 @@ from flarestack.core.astro import angular_distance, fast_angular_distance from flarestack.core.stats import king from flarestack.shared import bkg_spline_path -from flarestack.utils.make_SoB_splines import load_bkg_spatial_spline +from flarestack.utils.make_SoB_splines import SoB_splines try: from photospline import SplineTable @@ -28,9 +28,9 @@ class SpatialPDF: spatial PDF objects. """ - def __init__(self, spatial_pdf_dict, season): + def __init__(self, spatial_pdf_dict, season, SoB: Optional[SoB_splines]): self.signal = SignalSpatialPDF.create(spatial_pdf_dict) - self.background = BackgroundSpatialPDF.create(spatial_pdf_dict, season) + self.background = BackgroundSpatialPDF.create(spatial_pdf_dict, season, SoB) self.simulate_distribution = self.signal.simulate_distribution self.signal_spatial = self.signal.signal_spatial @@ -560,7 +560,7 @@ def _get_params( class BackgroundSpatialPDF: subclasses: "dict[str, type[BackgroundSpatialPDF]]" = {} - def __init__(self, spatial_pdf_dict, season): + def __init__(self, spatial_pdf_dict, season, SoB: Optional[SoB_splines]): pass @classmethod @@ -576,7 +576,9 @@ def decorator(subclass): return decorator @classmethod - def create(cls, s_pdf_dict, season) -> "BackgroundSpatialPDF": + def create( + cls, s_pdf_dict, season, SoB: Optional[SoB_splines] + ) -> "BackgroundSpatialPDF": try: s_pdf_name = s_pdf_dict["bkg_spatial_pdf"] except KeyError: @@ -588,7 +590,7 @@ def create(cls, s_pdf_dict, season) -> "BackgroundSpatialPDF": "Available names are {1}".format(s_pdf_name, cls.subclasses) ) - return cls.subclasses[s_pdf_name](s_pdf_dict, season) + return cls.subclasses[s_pdf_name](s_pdf_dict, season, SoB) def background_spatial(self, events): return np.ones(len(events)) @@ -614,7 +616,7 @@ class UniformSolidAngle(BackgroundSpatialPDF): """ def __init__(self, spatial_pdf_dict, season): - BackgroundSpatialPDF.__init__(self, spatial_pdf_dict, season) + BackgroundSpatialPDF.__init__(self, spatial_pdf_dict, season, None) try: self.solid_angle = spatial_pdf_dict["background_solid_angle"] @@ -644,18 +646,22 @@ class ZenithSpline(BackgroundSpatialPDF): spline is used to parameterise this distribution. """ - def __init__(self, spatial_pdf_dict, season): - BackgroundSpatialPDF.__init__(self, spatial_pdf_dict, season) - self.bkg_f = self.create_background_function(season) + def __init__(self, spatial_pdf_dict, season, SoB): + BackgroundSpatialPDF.__init__(self, spatial_pdf_dict, season, SoB) + if SoB is None: + raise TypeError( + "For 'zenith_spline' background spatial pdf you need the SoB" + ) + + self.bkg_f = self.create_background_function(season, SoB) @staticmethod - def create_background_function(season): + def create_background_function(season, SoB): # Checks if background spatial spline has been created + if not os.path.isfile(bkg_spline_path(season, SoB.spline_name)): + season.make_background_spatial(SoB) - if not os.path.isfile(bkg_spline_path(season)): - season.make_background_spatial() - - return load_bkg_spatial_spline(season) + return SoB.load_bkg_spatial_spline(season) def background_spatial(self, events): space_term = (1.0 / (2.0 * np.pi)) * np.exp(self.bkg_f(events["sinDec"])) diff --git a/flarestack/data/__init__.py b/flarestack/data/__init__.py index 42acfdd4..fe574f9d 100644 --- a/flarestack/data/__init__.py +++ b/flarestack/data/__init__.py @@ -3,6 +3,7 @@ import copy import logging import os +from typing import Callable import numpy as np from astropy.table import Table @@ -16,7 +17,7 @@ TimePDF, ) from flarestack.utils.create_acceptance_functions import make_acceptance_season -from flarestack.utils.make_SoB_splines import make_background_spline +from flarestack.utils.make_SoB_splines import SoB_splines logger = logging.getLogger(__name__) @@ -232,6 +233,11 @@ def simulate_background( data = np.random.choice(data, int(len(data) * self._subselection_fraction)) return data, 0 + def simulate_bkg_with_diffuse( + self, sources: Table, spatial_box_width: None | float, diffuse_weights: Callable + ) -> tuple[Table, int]: + raise NotImplementedError("Diffuse background model is implemented only for NT") + def get_exp_data(self, **kwargs) -> Table: return self.load_data(self.exp_path, **kwargs) @@ -290,8 +296,8 @@ def make_injector(self, sources, **inj_kwargs): def return_name(self): return self.sample_name + "/" + self.season_name - def make_background_spatial(self): - make_background_spline(self) + def make_background_spatial(self, SoB: SoB_splines): + SoB.make_background_spline(self) def get_pseudo_mc(self, **kwargs): return self.load_data(self.pseudo_mc_path, **kwargs) diff --git a/flarestack/data/icecube/northern_tracks/__init__.py b/flarestack/data/icecube/northern_tracks/__init__.py index 90fa20fd..a16cd806 100644 --- a/flarestack/data/icecube/northern_tracks/__init__.py +++ b/flarestack/data/icecube/northern_tracks/__init__.py @@ -1,3 +1,5 @@ +from typing import Callable + import numpy as np from astropy.table import Table @@ -194,6 +196,90 @@ def simulate_background( analysis_keys = list(self.get_background_dtype().names or []) return sim_bkg[analysis_keys], n_excluded + def simulate_bkg_with_diffuse( + self, sources: Table, spatial_box_width: None | float, bkg_weights: Callable + ) -> tuple[Table, int]: + """Same method as simulating atmospheric-only background, + only that diffuse flux is added to the background model. + The diffuse flux spectrum is designated in the injector dict, + the events' diffuse + atm weights are used to estimate the + expected number of and subsequently select the background events. + + Args: + sources (Table): sources when you want to choose + MC events in a spatial box around each + bkg_weights (Callable): method to get the events' + diffuse + atm weights, for the chosen diffuse flux. + Takes MC Table as arg + + Returns: + tuple[Table, int]: bkg MC, no. excluded events if spatial box is on + """ + rng = np.random + + if self.loaded_background_model is None: + raise RuntimeError( + "Monte Carlo background is not loaded. Call `load_background_model` before `simulate_background`." + ) + + if spatial_box_width is None: + # Draw from full MC sample + mc = self.loaded_background_model + n_excluded = 0 + else: + # Draw from MC sample within a spatial box around the sources + if ( + sources is not self._prev_sources + or spatial_box_width != self._prev_spatial_box_width + ): + self._prev_mc = self.loaded_background_model[ + self._get_spatially_coincident_indices( + np.asarray(self.loaded_background_model["ra"]), + np.asarray(self.loaded_background_model["dec"]), + np.asarray(sources["ra_rad"]), + np.asarray(sources["dec_rad"]), + np.deg2rad(spatial_box_width), + ) + ] + self._prev_sources = sources + self._prev_spatial_box_width = spatial_box_width + mc = self._prev_mc + # Draw a number of events that would have been outside the box + n_excluded = rng.poisson( + np.sum(bkg_weights(self.loaded_background_model)) + - np.sum(bkg_weights(mc)) + ) + + weights = bkg_weights(mc) # diffuse + atmospheric bkg weights + n_exp = np.sum(weights) + + n_bkg = rng.poisson(n_exp) + print( + f"Simulating {n_bkg} bkg diffuse + atmospheric events, excluding {n_excluded}." + ) + + p_select = weights.cumsum() / n_exp + ind = np.searchsorted(p_select, np.sort(rng.uniform(size=n_bkg)), side="right") + + sim_bkg = mc[ind] + + # Simulates random times + time_pdf = self.get_time_pdf() + sim_bkg["time"] = time_pdf.simulate_times(source=None, n_s=n_bkg) + + # Check that the time pdf evaluates to 1 for all the simulated times. + pdf_sum = np.sum(time_pdf.season_f(sim_bkg["time"])) + if pdf_sum < n_bkg: + raise RuntimeError( + f"The time PDF does not evaluate to 1 for all generated event times.\n \ + The sum of the PDF values over {n_bkg} events is {pdf_sum}.\n \ + This means the sampling of background times is not reliable and must be fixed." + ) + + # Reduce the data to the relevant fields for analysis. + analysis_keys = list(self.get_background_dtype().names or []) + return sim_bkg[analysis_keys], n_excluded + class NTSeasonNewStyle(NTSeason): def get_background_model(self): diff --git a/flarestack/data/public/icecube/all_sky_point_source/all_sky_3_year.py b/flarestack/data/public/icecube/all_sky_point_source/all_sky_3_year.py index 3329dbc1..18e1817b 100644 --- a/flarestack/data/public/icecube/all_sky_point_source/all_sky_3_year.py +++ b/flarestack/data/public/icecube/all_sky_point_source/all_sky_3_year.py @@ -15,7 +15,7 @@ med_ang_res_path, public_dataset_dir, ) -from flarestack.utils.make_SoB_splines import make_individual_spline_set +from flarestack.utils.make_SoB_splines import SoB_splines logger = logging.getLogger(__name__) @@ -228,10 +228,17 @@ def run_all(): parse_numpy_dataset() parse_angular_resolution() + sob_dict = { + "bkg_model_name": "no_difffuse", + "smoothing_order": "flarestack", + "gamma_precision": "flarestack", + } + for season in icecube_ps_3_year.get_seasons().values(): season.map_energy_proxy() season.plot_effective_area() - make_individual_spline_set(season, SoB_spline_path(season)) + sob = SoB_splines.create(season, sob_dict) + sob.make_individual_spline_set(season, SoB_spline_path(season, **sob.sob_dict)) # If data has not been extracted, then extract from zip file diff --git a/flarestack/shared.py b/flarestack/shared.py index 15ecf13d..707735a8 100644 --- a/flarestack/shared.py +++ b/flarestack/shared.py @@ -173,7 +173,7 @@ def pull_pickle(pull_dict): return pull_dir + str(unique_key) + ".pkl" -def llh_energy_hash_pickles(llh_dict, season): +def llh_energy_hash_pickles(llh_dict, season, bkg_model_name="no_difffuse"): hash_dict = dict(llh_dict["llh_energy_pdf"]) hash_dict["llh_name"] = llh_dict["llh_name"] @@ -187,7 +187,7 @@ def llh_energy_hash_pickles(llh_dict, season): + season.sample_name + "/" + season.season_name - + smoothing_precision_string(smoothing_order, precision) + + smoothing_precision_string(smoothing_order, precision, bkg_model_name) + ".pkl" ) SoB_path = SoB_spline_dir + season_path @@ -275,7 +275,9 @@ def get_base_sob_plot_dir(season): ) -def smoothing_precision_string(smoothing_order="flarestack", gamma_precision="skylab"): +def smoothing_precision_string( + smoothing_order="flarestack", gamma_precision="skylab", bkg_model_name="no_difffuse" +): if isinstance(smoothing_order, str): if smoothing_order in default_smoothing_order.keys(): smoothing_order = default_smoothing_order[smoothing_order] @@ -292,6 +294,8 @@ def smoothing_precision_string(smoothing_order="flarestack", gamma_precision="sk f"smoothing order is {smoothing_order}, gamma precision is {gamma_precision}" ) s = "" + if bkg_model_name != "no_difffuse": + s += f"_{bkg_model_name}" if smoothing_order != default_smoothing_order["flarestack"]: s += f"_smoothing{smoothing_order}" if gamma_precision != default_gamma_precision["flarestack"]: @@ -314,8 +318,13 @@ def SoB_spline_path(season, *args, **kwargs): ) -def bkg_spline_path(season): - return bkg_spline_dir + season.sample_name + "/" + season.season_name + ".pkl" +def bkg_spline_path(season, bkg_model_name="no_difffuse"): + path = bkg_spline_dir + season.sample_name + "/" + season.season_name + if bkg_model_name != "no_difffuse": + path += f"_{bkg_model_name}.pkl" + else: + path += ".pkl" + return path def energy_proxy_path(season): diff --git a/flarestack/utils/create_acceptance_functions.py b/flarestack/utils/create_acceptance_functions.py index d2b06632..6e18b5e8 100644 --- a/flarestack/utils/create_acceptance_functions.py +++ b/flarestack/utils/create_acceptance_functions.py @@ -6,7 +6,7 @@ from flarestack.core.energy_pdf import PowerLaw from flarestack.shared import acceptance_path, get_base_sob_plot_dir -from flarestack.utils.make_SoB_splines import make_plot +from flarestack.utils.make_SoB_splines import SoB_splines logger = logging.getLogger(__name__) @@ -73,7 +73,7 @@ def make_acceptance_season(season, acc_path): except OSError: pass - make_plot( + SoB_splines.make_plot( acc, savepath, gamma_vals, diff --git a/flarestack/utils/make_SoB_splines.py b/flarestack/utils/make_SoB_splines.py index a74a5b65..c54c27e7 100644 --- a/flarestack/utils/make_SoB_splines.py +++ b/flarestack/utils/make_SoB_splines.py @@ -1,11 +1,11 @@ import logging import os import pickle as Pickle -import shutil import matplotlib.pyplot as plt import numpy as np import scipy.interpolate +from astropy.table import Table from flarestack.core.energy_pdf import PowerLaw from flarestack.shared import ( @@ -17,17 +17,14 @@ get_base_sob_plot_dir, ) -environment_smoothing_key = "FLARESATCK_SMOOTHING_ORDER" -environment_precision_key = "FLARESTACK_PRECISION" logger = logging.getLogger(__name__) energy_pdf = PowerLaw() def get_gamma_precision(precision=flarestack_gamma_precision): """Returns the precision in gamma that is used. - Returns default value if the environment_precision_key is not present in th environ dictionary - :param precision: Specify which precision to use. Default to the standard precision. + :param precision: Specify which precision to use. Default is 0.025 for 'flarestack'. Can also provide default of llh codes by name, either 'skylab' or 'flarestack'. :return: Precision as a float """ @@ -66,567 +63,730 @@ def _around(value, precision=flarestack_gamma_precision): def get_gamma_support_points(precision=flarestack_gamma_precision): """Return the gamma support points based on the gamma precision - :param precision: Specify which precision to use. Default to the standard precision. + :param precision: Specify which precision to use. Can also provide default of llh codes by name, either 'skylab' or 'flarestack'. + Default is 0.025 for 'flarestack'. :return: Gamma support points """ gamma_points = np.arange(0.7, 4.3, get_gamma_precision(precision=precision)) return set([_around(i, precision=precision) for i in gamma_points]) -def create_2d_hist(sin_dec, log_e, sin_dec_bins, log_e_bins, weights): - """Creates a 2D histogram for a set of data (Experimental or Monte - Carlo), in which the dataset is binned by sin(Declination) and - Log(Energy). Weights the histogram by the values in the weights array. - Normalises the histogram, such that the sum of each sin(Declination) - column is equal to 1. - - :param sin_dec: Sin(Declination) array - :param log_e: Log(Energy/GeV) array - :param sin_dec_bins: Bins of Sin(Declination to be used) - :param weights: Array of weights for event - :return: Normalised histogram +class SoB_splines: + """Base class for Signal over Background splines + which are used as the energy pdf in the llh. + The splines are created per season from the + 2D histograms of the S/B ratio in sin(DEC) and Log(Energy/GeV). + The background and signal histograms are created separately, + weighted according to the chosen background flux model and + the power-law injected signal respectively. + The ratio of these 2D histograms is used to make the spline. """ - # Produces the histogram - # rangee = tuple([(wB_i[0], wB_i[-1]) for wB_i in [log_e_bins, sin_dec_bins]]) - # if np.all(weights == 1): - # weights = None - # print('no weights') - hist_2d, binedges = np.histogramdd( - (log_e, sin_dec), bins=(log_e_bins, sin_dec_bins), weights=weights - ) # , range=rangee, normed=True) - - # n_dimensions = hist_2d.ndim - - # # Normalises histogram - # norms = np.sum(hist_2d, axis=n_dimensions - 2) - # norms[norms == 0.] = 1. - # hist_2d /= norms - - return hist_2d - - -# def create_2d_gamma_energy(sin_dec, log_e, sin_dec_bins, weights): -# """Creates a 2D histogram for a set of data (Experimental or Monte -# Carlo), in which the dataset is binned by sin(Declination) and -# Log(Energy). Weights the histogram by the values in the weights array. -# Normalises the histogram, such that the sum of each sin(Declination) -# column is equal to 1. -# -# :param sin_dec: Sin(Declination) array -# :param log_e: Log(Energy/GeV) array -# :param sin_dec_bins: Bins of Sin(Declination to be used) -# :param weights: Array of weights for event -# :return: Normalised histogram -# """ -# # Produces the histogram -# hist_2d, binedges = np.histogramdd( -# (log_e, sin_dec), bins=(energy_bins, sin_dec_bins), weights=weights) -# n_dimensions = hist_2d.ndim -# -# # Normalises histogram -# norms = np.sum(hist_2d, axis=n_dimensions - 2) -# norms[norms == 0.] = 1. -# hist_2d /= norms -# -# return hist_2d + subclasses: dict[str, type["SoB_splines"]] = {} -def create_bkg_2d_hist(exp, sin_dec_bins, log_e_bins): - """Creates a background 2D logE/sinDec histogram. + def __init__(self, season, SoB_dict) -> None: + self.season = season + self.sob_dict = SoB_dict + self.spline_name = SoB_dict["bkg_model_name"] - :param exp: Experimental data - :param sin_dec_bins: Bins of Sin(Declination to be used) - :return: 2D histogram - """ - return create_2d_hist( - exp["sinDec"], exp["logE"], sin_dec_bins, log_e_bins, weights=exp["weight"] - ) + self.smoothing_order = SoB_dict["smoothing_order"] + if isinstance(self.smoothing_order, str): + if self.smoothing_order in default_smoothing_order.keys(): + self.smoothing_order = default_smoothing_order[self.smoothing_order] + else: + raise ValueError( + f"Smoothing order for {self.smoothing_order} not known!" + ) + self.gamma_precision = SoB_dict["gamma_precision"] -def create_sig_2d_hist(mc, sin_dec_bins, log_e_bins, weight_function): - """Creates a signal 2D logE/sinDec histogram. + @classmethod + def register_subclass(cls, spline_name): + """Adds a new subclass of SoB_splines, with class name equal to + "bkg_model_name". + """ - :param mc: MC Simulations - :param sin_dec_bins: Bins of Sin(Declination) to be used - :param weight_function: Weight Function - :return: 2D histogram - """ + def decorator(subclass): + cls.subclasses[spline_name] = subclass + return subclass - return create_2d_hist( - mc["sinDec"], mc["logE"], sin_dec_bins, log_e_bins, weights=weight_function(mc) - ) - - -def create_2d_ratio_hist(exp, mc, sin_dec_bins, log_e_bins, weight_function): - """Creates a 2D histogram for both data and MC, in which the seasons - are binned by Sin(Declination) and Log(Energy/GeV). Each histogram is - normalised in Sin(Declination) bands. Then creates a histogram of the - ratio of the Signal/Background histograms. In bins where there is - simulation but no data, a count of 1 is assigned to the background - histogram. This is broadly unimportant for unblinding archival searches, - because there will never be a situation in which a bin without any data - will be queried. In all other cases, the ratio is set to 1. - - :param exp: Experimental data - :param mc: MC Simulations - :param sin_dec_bins: Bins of Sin(Declination) to be used - :param weight_function: Weight Function - :return: ratio histogram - """ + return decorator - bkg_hist = create_bkg_2d_hist(exp, sin_dec_bins, log_e_bins) - sig_hist = create_sig_2d_hist(mc, sin_dec_bins, log_e_bins, weight_function) - n_dimensions = sig_hist.ndim - norms = np.sum(sig_hist, axis=n_dimensions - 2) - norms[norms == 0.0] = 1.0 - sig_hist /= norms - # bkg_norms = np.sum(bkg_hist, axis=tuple(range(n_dimensions - 1))) - # bkg_norms[bkg_norms == 0.] = 1 - # bkg_hist /= bkg_norms - - ratio = np.ones_like(bkg_hist, dtype=float) - - # wSd = sig_hist > 0 - # wB_domain = bkg_hist > 0 - # ratio[wSd & wB_domain] = (sig_hist[wSd & wB_domain] / bkg_hist[wSd & wB_domain]) - # # values outside of the exp domain, but inside the MC one are mapped to - # # the most signal-like value - # if np.any(ratio > 1): - # min_ratio = np.percentile(ratio[ratio > 1.], 99.) - # np.copyto(ratio, min_ratio, where=wSd & ~wB_domain) - - for i, bkg_row in enumerate(bkg_hist.T): - sig_row = sig_hist.T[i] - - fill_mask = (bkg_row == 0.0) & (sig_row > 0.0) - bkg_row[fill_mask] = 1.0 - # bkg_row /= np.sum(bkg_row) - - mask = (bkg_row > 0.0) & (sig_row > 0.0) - r = np.ones_like(bkg_row) - r[mask] = sig_row[mask] / (bkg_row[mask] / np.sum(bkg_row)) - - ratio.T[i] = r - # print(max(sig_row), np.median(bkg_row), np.sum(bkg_row)) - # print(r) - # input("prompt") - # - # input("prompt") - - # print('background signal ratio') - # for j, (b, s, r) in enumerate(zip(bkg_hist[:, 0], sig_hist[:, 0], ratio[:, 0])): - # print(f'{j}: {b:.4f} {s:.4f} {r:.4f}') - # input('continue? ') - - return ratio - - -def create_2d_ratio_spline( - exp, mc, sin_dec_bins, log_e_bins, weight_function, smoothing_order -): - """Creates 2D histograms for both data and MC, in which the seasons - are binned by Sin(Declination) and Log(Energy/GeV). Each histogram is - normalised in Sin(Declination) bands. Then creates a histogram of the - ratio of the Signal/Background histograms. In bins where there is - simulation but no data, the ratio is set to the highest ratio - value found for cases with both data and MC. This is broadly - unimportant for unblinded archival searches, because there will never - be a situation in which a bin without any data will be queried. In all - other cases, the ratio is set to 1. - - A 2D spline, of 2nd order in x and y, is then fit to the Log(Ratio), - and returned. - - :param exp: Experimental data - :param mc: MC Simulations - :param sin_dec_bins: Bins of Sin(Declination) to be used - :param weight_function: Weight Function - :return: 2D spline function - """ + @classmethod + def create(cls, season, SoB_dict) -> "SoB_splines": + try: + spline_name = SoB_dict["bkg_model_name"] + except KeyError: + spline_name = "no_difffuse" + logger.info("No 'bkg_model_name' in sob dict, default to atmospheric-only") + + if spline_name not in cls.subclasses: + raise ValueError("Bad SoB spline name {}".format(spline_name)) + + return cls.subclasses[spline_name](season, SoB_dict) + + def bkg_weights(self, exp: Table) -> np.ndarray: + raise NotImplementedError + + # ============================================================================== + # BACKGROUND SPATIAL PDF + # ============================================================================== + + def create_bkg_spatial_spline(self, exp, sin_dec_bins): + """Creates the spatial PDF for background. + Generates a histogram for the exp. distribution in sin declination. + Fits a spline function to the distribution, giving a spatial PDF. + Returns this spatial PDF. + + :param exp: Experimental data (background) + :param sin_dec_bins: Bins of Sin(Declination) to be used + :return: Background spline function + """ + sin_dec_range = (np.min(sin_dec_bins), np.max(sin_dec_bins)) + hist, bins = np.histogram( + exp["sinDec"], + density=True, + bins=sin_dec_bins, + range=sin_dec_range, + weights=self.bkg_weights(exp), + ) - ratio = create_2d_ratio_hist(exp, mc, sin_dec_bins, log_e_bins, weight_function) + bins = np.concatenate([bins[:1], bins, bins[-1:]]) + hist = np.concatenate([hist[:1], hist, hist[-1:]]) - spline = make_2d_spline_from_hist(ratio, sin_dec_bins, log_e_bins, smoothing_order) + bkg_spline = scipy.interpolate.InterpolatedUnivariateSpline( + (bins[1:] + bins[:-1]) / 2.0, np.log(hist), k=2 + ) + return bkg_spline - return spline + def make_background_spline(self, season): + bkg_path = bkg_spline_path(season, self.spline_name) + bkg = season.get_background_model() + sin_dec_bins = season.sin_dec_bins + bkg_spline = self.create_bkg_spatial_spline(bkg, sin_dec_bins) -def make_2d_spline_from_hist(ratio, sin_dec_bins, log_e_bins, smoothing_order): - # Sets bin centers, and order of spline (for x and y) - sin_bin_center = (sin_dec_bins[:-1] + sin_dec_bins[1:]) / 2.0 - log_e_bin_center = (log_e_bins[:-1] + log_e_bins[1:]) / 2.0 + logger.info(f"Saving bakcground spatial spline to {bkg_path}") + try: + os.makedirs(os.path.dirname(bkg_path)) + except OSError: + pass - # the order of the splines defaults to 2 - # default_order = 2 - if not isinstance(smoothing_order, int): - smoothing_order = default_smoothing_order[smoothing_order] + with open(bkg_path, "wb") as f: + Pickle.dump(bkg_spline, f) - # if the environment_smoothing_key is present in the environ dictionary use the corresponding value instead - # _order = os.environ.get(environment_smoothing_key, default_order) - # order = None if _order == 'None' else int(_order) + x_range = np.linspace(sin_dec_bins[0], sin_dec_bins[-1], 101) + plt.figure() + plt.plot(x_range, np.exp(bkg_spline(x_range))) + plt.ylabel(r"$P_{bkg}$ (spatial)") + plt.xlabel(r"$\sin(\delta)$") + savepath = get_base_sob_plot_dir(season) - # when setting the emviron value to 'None', order will be None and no splines will be produced - if isinstance(smoothing_order, type(None)): - logger.warning(f"{environment_smoothing_key} is None! Not making splines!") - return + try: + os.makedirs(os.path.dirname(savepath)) + except OSError: + pass - # Fits a 2D spline function to the log of ratio array - if smoothing_order == 0: - # use nearest-neighbor interpolation to ensure that ratio retains the normalization of the underlying histograms - spline = scipy.interpolate.RegularGridInterpolator( - (log_e_bin_center, sin_bin_center), - np.log(ratio), - method="nearest", - bounds_error=False, - fill_value=None, - ) + sfx = f"{self.spline_name}_" if self.spline_name != "no_difffuse" else "" + plt.savefig(savepath + sfx + "bkg_spatial.pdf") + plt.close() - # If the splines are of order one the RegularGridInterpolator is used to match the SkyLab behavior - elif smoothing_order == 1: - sin_bin_center[0], sin_bin_center[-1] = sin_dec_bins[0], sin_dec_bins[-1] - # log_e_bins[0], log_e_bins[-1] = log_e_bins[0], log_e_bins[-1] + def load_bkg_spatial_spline(self, season): + path = bkg_spline_path(season, self.spline_name) - binmids = (log_e_bin_center, sin_bin_center) + logger.debug(f"Loading background spatial spline from {path}") - spline = scipy.interpolate.RegularGridInterpolator( - binmids, np.log(ratio), method="linear", bounds_error=False, fill_value=0.0 + try: + with open(path, "rb") as f: + res = Pickle.load(f) + except FileNotFoundError as err: + logger.info( + f"No cached spline found at {path}. Creating this file instead." + ) + logger.info(f"Cause: {err}") + self.make_background_spline(season) + with open(path, "rb") as f: + res = Pickle.load(f) + + except ModuleNotFoundError as err: + logger.error( + "A spline was found but it seems incompatible with this installation." + ) + logger.error(f"Cause: {err}") + raise + + return res + + # ============================================================================== + # ENERGY PDF SPLINES + # ============================================================================== + + @staticmethod + def create_2d_hist(sin_dec, log_e, sin_dec_bins, log_e_bins, weights): + """Creates a 2D histogram for a set of data (Experimental or Monte + Carlo), in which the dataset is binned by sin(Declination) and + Log(Energy). Weights the histogram by the values in the weights array. + Normalises the histogram, such that the sum of each sin(Declination) + column is equal to 1. + + :param sin_dec: Sin(Declination) array + :param log_e: Log(Energy/GeV) array + :param sin_dec_bins: Bins of Sin(Declination) to be used + :param log_e_bins: bins of log(E/GeV) to be used + :param weights: Array of weights for events + :return: Normalised histogram + """ + hist_2d, _ = np.histogramdd( + (log_e, sin_dec), bins=(log_e_bins, sin_dec_bins), weights=weights ) - # If the interpolating splines are of order greater than 1, use RectBivariateSpline - else: - # This is order-th order in both dimensions - spline = scipy.interpolate.RectBivariateSpline( - log_e_bin_center, - sin_bin_center, - np.log(ratio), - kx=smoothing_order, - ky=smoothing_order, - s=0, - ) + return hist_2d - return spline + def create_sig_2d_hist(self, mc, sin_dec_bins, log_e_bins, weight_function): + """Creates a signal 2D logE/sinDec weighted histogram. + :param mc: MC Simulations + :param weight_function: Weight Function + :return: 2D histogram + """ -def create_gamma_2d_ratio_spline( - exp, mc, sin_dec_bins, log_e_bins, gamma, smoothing_order -): - """Creates a 2D gamma ratio spline by creating a function that weights MC - assuming a power law of spectral index gamma. + return self.create_2d_hist( + mc["sinDec"], + mc["logE"], + sin_dec_bins, + log_e_bins, + weights=weight_function(mc), + ) - :param exp: Experimental data - :param mc: MC Simulations - :param sin_dec_bins: Bins of Sin(Declination) to be used - :param gamma: Spectral Index - :return: 2D spline function - """ + def create_bkg_2d_hist(self, exp, sin_dec_bins, log_e_bins): + """Creates a background 2D logE/sinDec weighted histogram. + The weights depend on the chosen background model. + + :param exp: Experimental data (or MC depending on season) + :return: 2D histogram + """ + + w = self.bkg_weights(exp) + return self.create_2d_hist( + exp["sinDec"], + exp["logE"], + sin_dec_bins, + log_e_bins, + weights=w, + ) - def weight_function(sig_mc): - return energy_pdf.weight_mc(sig_mc, gamma) + def create_2d_ratio_hist( + self, + exp, + mc, + sin_dec_bins, + log_e_bins, + weight_f, + ): + """Creates a 2D histogram for both data and MC, in which the seasons + are binned by Sin(Declination) and Log(Energy/GeV). Each histogram is + normalised in Sin(Declination) bands. Then creates a histogram of the + ratio of the Signal/Background histograms. In bins where there is + simulation but no data, a count of 1 is assigned to the background + histogram. This is broadly unimportant for unblinding archival searches, + because there will never be a situation in which a bin without any data + will be queried. In all other cases, the ratio is set to 1. + + :param exp: Experimental data (or MC) used for bkg hist + :param mc: MC Simulations used for sig hist + :param weight_f: Weight Function used for sig hist + :return: ratio histogram + """ + + bkg_hist = self.create_bkg_2d_hist(exp, sin_dec_bins, log_e_bins) + sig_hist = self.create_sig_2d_hist(mc, sin_dec_bins, log_e_bins, weight_f) + n_dimensions = sig_hist.ndim + norms = np.sum(sig_hist, axis=n_dimensions - 2) + norms[norms == 0.0] = 1.0 + sig_hist /= norms - return create_2d_ratio_spline( - exp, mc, sin_dec_bins, log_e_bins, weight_function, smoothing_order - ) + ratio = np.ones_like(bkg_hist, dtype=float) + for i, bkg_row in enumerate(bkg_hist.T): + sig_row = sig_hist.T[i] -def create_2d_splines(exp, mc, sin_dec_bins, log_e_bins, **kwargs): - """If gamma will not be fit, then calculates the Log(Signal/Background) - 2D PDF for the fixed value self.default_gamma. Fits a spline to each - histogram, and saves the spline in a dictionary. + fill_mask = (bkg_row == 0.0) & (sig_row > 0.0) + bkg_row[fill_mask] = 1.0 + # bkg_row /= np.sum(bkg_row) - If gamma should be fit, instead loops over each value of gamma in - self.gamma_support_points. For each gamma value, the spline creation - is repeated, and saved as a dictionary entry. + mask = (bkg_row > 0.0) & (sig_row > 0.0) + r = np.ones_like(bkg_row) + r[mask] = sig_row[mask] / (bkg_row[mask] / np.sum(bkg_row)) - In either case, returns the dictionary of spline/splines. + ratio.T[i] = r + return ratio - :param exp: Experimental data - :param mc: MC Simulations - :param sin_dec_bins: Bins of Sin(Declination) to be used - :return: Dictionary of 2D Log(Signal/Background) splines - """ - splines = dict() - gamma_precision = kwargs.get("gamma_precision", "flarestack") - smoothing_order = kwargs.get("smoothing_order", "flarestack") - gamma_support_points = get_gamma_support_points(precision=gamma_precision) - - for gamma in gamma_support_points: - splines[gamma] = create_gamma_2d_ratio_spline( - exp, mc, sin_dec_bins, log_e_bins, gamma, smoothing_order - ) + def make_2d_spline_from_hist( + self, ratio, sin_dec_bins, log_e_bins, smoothing_order: int + ): - if not np.any(list(splines.values())): - logger.warning("No splines!") - return + # Sets bin centers, and order of spline (for x and y) + sin_bin_center = (sin_dec_bins[:-1] + sin_dec_bins[1:]) / 2.0 + log_e_bin_center = (log_e_bins[:-1] + log_e_bins[1:]) / 2.0 - return splines + # Fits a 2D spline function to the log of ratio array + if smoothing_order == 0: + # use nearest-neighbor interpolation to ensure that ratio retains the normalization of the underlying histograms + spline = scipy.interpolate.RegularGridInterpolator( + (log_e_bin_center, sin_bin_center), + np.log(ratio), + method="nearest", + bounds_error=False, + fill_value=None, + ) + # If the splines are of order one the RegularGridInterpolator is used to match the SkyLab behavior + elif smoothing_order == 1: + sin_bin_center[0], sin_bin_center[-1] = sin_dec_bins[0], sin_dec_bins[-1] + # log_e_bins[0], log_e_bins[-1] = log_e_bins[0], log_e_bins[-1] -def create_bkg_spatial_spline(exp, sin_dec_bins): - """Creates the spatial PDF for background. - Generates a histogram for the exp. distribution in sin declination. - Fits a spline function to the distribution, giving a spatial PDF. - Returns this spatial PDF. + binmids = (log_e_bin_center, sin_bin_center) - :param exp: Experimental data (background) - :param sin_dec_bins: Bins of Sin(Declination) to be used - :return: Background spline function - """ - sin_dec_range = (np.min(sin_dec_bins), np.max(sin_dec_bins)) - hist, bins = np.histogram( - exp["sinDec"], - density=True, - bins=sin_dec_bins, - range=sin_dec_range, - weights=exp["weight"], - ) + spline = scipy.interpolate.RegularGridInterpolator( + binmids, + np.log(ratio), + method="linear", + bounds_error=False, + fill_value=0.0, + ) - bins = np.concatenate([bins[:1], bins, bins[-1:]]) - hist = np.concatenate([hist[:1], hist, hist[-1:]]) + # If the interpolating splines are of order greater than 1, use RectBivariateSpline + else: + # This is order-th order in both dimensions + spline = scipy.interpolate.RectBivariateSpline( + log_e_bin_center, + sin_bin_center, + np.log(ratio), + kx=smoothing_order, + ky=smoothing_order, + s=0, + ) - bkg_spline = scipy.interpolate.InterpolatedUnivariateSpline( - (bins[1:] + bins[:-1]) / 2.0, np.log(hist), k=2 - ) - return bkg_spline + return spline + + def create_2d_ratio_spline( + self, exp, mc, sin_dec_bins, log_e_bins, weight_f, smoothing_order: int + ): + """Creates 2D histograms for both data and MC, in which the seasons + are binned by Sin(Declination) and Log(Energy/GeV). Each histogram is + normalised in Sin(Declination) bands. Then creates a histogram of the + ratio of the Signal/Background histograms. In bins where there is + simulation but no data, the ratio is set to the highest ratio + value found for cases with both data and MC. This is broadly + unimportant for unblinded archival searches, because there will never + be a situation in which a bin without any data will be queried. In all + other cases, the ratio is set to 1. + + A 2D spline, of chosen order in x and y, is then fit to the Log(Ratio), + and returned. + + :param exp: Experimental data (or MC) used for bkg hist + :param mc: MC Simulations used for sig hist + :param weight_f: weight function used for sig hist + :param bkg_weights: array of weights for bkg events + :param smoothing_order: order of the spline + :return: 2D spline function + """ + + ratio = self.create_2d_ratio_hist( + exp, + mc, + sin_dec_bins, + log_e_bins, + weight_f, + ) + spline = self.make_2d_spline_from_hist( + ratio, sin_dec_bins, log_e_bins, smoothing_order + ) -def make_spline(seasons, **kwargs): - """Make the S/B splines for each season, as well as the background spline. + return spline - :param seasons: Seasons to iterate over - """ + def create_gamma_2d_ratio_spline( + self, exp, mc, sin_dec_bins, log_e_bins, gamma, smoothing_order: int + ): + """Creates a 2D gamma ratio spline by creating a function that weights MC + assuming a power law of spectral index gamma. - logger.info( - f"Splines will be made to calculate the Signal/Background ratio of " - "the MC to data. The MC will be weighted with a power law, for each" - " gamma in: {list(get_gamma_support_points(**kwargs))}" - ) - - for season in seasons.values(): - SoB_path = SoB_spline_path(season, **kwargs) - make_individual_spline_set(season, SoB_path, **kwargs) - make_background_spline(season) - - -def make_plot( - hist, - savepath, - x_bins, - y_bins, - normed=True, - log_min=5, - label_x=r"$\sin(\delta)$", - label_y="log(Energy)", -): - if normed: - norms = np.sum(hist, axis=hist.ndim - 2) - norms[norms == 0.0] = 1.0 - hist /= norms - else: - hist = np.log(np.array(hist)) - plt.figure() - ax = plt.subplot(111) - X, Y = np.meshgrid(x_bins, y_bins) - if not normed: - max_col = min( - abs(min([min(row) for row in hist.T])), max([max(row) for row in hist.T]) + :param exp: Experimental data + :param mc: MC Simulations + :param sin_dec_bins: Bins of Sin(Declination) to be used + :param gamma: Spectral Index + :return: 2D spline function + """ + + def weight_function(sig_mc): + return energy_pdf.weight_mc(sig_mc, gamma) + + return self.create_2d_ratio_spline( + exp, mc, sin_dec_bins, log_e_bins, weight_function, smoothing_order ) - cbar = ax.pcolormesh(X, Y, hist, cmap="seismic", vmin=-5, vmax=5) - plt.colorbar(cbar, label="Log(Signal/Background)") - else: - hist[hist == 0.0] = np.nan - cbar = ax.pcolormesh(X, Y, hist) - plt.colorbar(cbar, label="Column-normalised density") - plt.xlabel(label_x) - plt.ylabel(label_y) - plt.savefig(savepath) - plt.close() + def create_2d_splines(self, exp, mc, sin_dec_bins, log_e_bins): + """If gamma will not be fit, then calculates the Log(Signal/Background) + 2D PDF for the fixed value self.default_gamma. Fits a spline to each + histogram, and saves the spline in a dictionary. -def make_individual_spline_set(season, SoB_path, **kwargs): - try: - logger.info(f"Making splines for {season.season_name}") - # path = SoB_spline_path(season) + If gamma should be fit, instead loops over each value of gamma in + self.gamma_support_points. For each gamma value, the spline creation + is repeated, and saved as a dictionary entry. - exp = season.get_background_model() - mc = season.get_pseudo_mc() + In either case, returns the dictionary of spline/splines. - sin_dec_bins = season.sin_dec_bins - log_e_bins = season.log_e_bins + :param exp: Experimental data + :param mc: MC Simulations + :param sin_dec_bins: Bins of Sin(Declination) to be used + :return: Dictionary of 2D Log(Signal/Background) splines + """ + splines = dict() - splines = create_2d_splines(exp, mc, sin_dec_bins, log_e_bins, **kwargs) + gamma_support_points = get_gamma_support_points(self.gamma_precision) + for gamma in gamma_support_points: + splines[gamma] = self.create_gamma_2d_ratio_spline( + exp, mc, sin_dec_bins, log_e_bins, gamma, self.smoothing_order + ) - logger.info(f"Saving to {SoB_path}.") + if not np.any(list(splines.values())): + logger.warning("No splines!") + return + return splines + + @staticmethod + def make_plot( + hist, + savepath, + x_bins, + y_bins, + normed=True, + log_min=5, + label_x=r"$\sin(\delta)$", + label_y="log(Energy)", + ): + """Plot a 2D histogram""" + if normed: + norms = np.sum(hist, axis=hist.ndim - 2) + norms[norms == 0.0] = 1.0 + hist /= norms + else: + hist = np.log(np.array(hist)) + plt.figure() + ax = plt.subplot(111) + X, Y = np.meshgrid(x_bins, y_bins) + if not normed: + max_col = min( + abs(min([min(row) for row in hist.T])), + max([max(row) for row in hist.T]), + ) + cbar = ax.pcolormesh(X, Y, hist, cmap="seismic", vmin=-5, vmax=5) + plt.colorbar(cbar, label="Log(Signal/Background)") + else: + hist[hist == 0.0] = np.nan + cbar = ax.pcolormesh(X, Y, hist) + plt.colorbar(cbar, label="Column-normalised density") + plt.xlabel(label_x) + plt.ylabel(label_y) + plt.savefig(savepath) + plt.close() + + def make_individual_spline_set(self, season, SoB_path): + """Create the SoB splines dictionary for given season, + and plot the normalized 2D histograms for background + and signal, as well as the log(S/B) 2D hist. + The latter two are plotted for gamma values in + np.linspace(1.0, 4.0, 7). + """ try: - os.makedirs(os.path.dirname(SoB_path)) - except OSError: - pass + logger.info(f"Making SoB splines for {season.season_name}") - with open(SoB_path, "wb") as f: - Pickle.dump(splines, f) - - if isinstance(splines, type(None)): - return + exp = season.get_background_model() + mc = season.get_pseudo_mc() - base_plot_path = get_base_sob_plot_dir(season) + sin_dec_bins = season.sin_dec_bins + log_e_bins = season.log_e_bins - exp_hist = create_bkg_2d_hist(exp, sin_dec_bins, log_e_bins) + ##### MAKE SPLINES ##### + splines = self.create_2d_splines(exp, mc, sin_dec_bins, log_e_bins) - # Generate plots - for gamma in np.linspace(1.0, 4.0, 7): - plot_path = ( - base_plot_path - + "gamma=" - + str(gamma) - + "/" - + f'precision{kwargs.get("gamma_precision", "flarestack")}_' - f'smoothing{kwargs.get("smoothing_order", "flarestack")}' - ) + logger.info(f"Saving SoB splines to {SoB_path}.") try: - os.makedirs(plot_path) + os.makedirs(os.path.dirname(SoB_path)) except OSError: pass - def weight_function(sig_mc): - return energy_pdf.weight_mc(sig_mc, gamma) - - mc_hist = create_sig_2d_hist(mc, sin_dec_bins, log_e_bins, weight_function) + with open(SoB_path, "wb") as f: + Pickle.dump(splines, f) - make_plot(mc_hist, plot_path + "sig.pdf", sin_dec_bins, log_e_bins) - make_plot( - create_2d_ratio_hist( - exp, mc, sin_dec_bins, log_e_bins, weight_function - ), - plot_path + "SoB.pdf", - sin_dec_bins, - log_e_bins, - normed=False, - ) + if isinstance(splines, type(None)): + return - Z = [] - for s in sin_dec_bins: - z_line = [] - for e in log_e_bins: - # logging.debug(f'{e}, {s}') - try: - z_line.append(splines[gamma](e, s)[0][0]) - except: - z_line.append(splines[gamma]((e, s))) - Z.append(z_line) + ##### GENERATE PLOTS ##### + base_plot_path = get_base_sob_plot_dir(season) + sfx = f"{self.spline_name}_" if self.spline_name != "no_difffuse" else "" - Z = np.array(Z).T - - max_col = min( - abs(min([min(row) for row in Z])), max([max(row) for row in Z]) + # plot 2D background hist + bkg_hist = self.create_bkg_2d_hist(exp, sin_dec_bins, log_e_bins) + self.make_plot( + bkg_hist, + savepath=base_plot_path + sfx + "bkg.pdf", + x_bins=sin_dec_bins, + y_bins=log_e_bins, ) - plt.figure() - ax = plt.subplot(111) - X, Y = np.meshgrid(sin_dec_bins, log_e_bins) - cbar = ax.pcolormesh( - X, Y, Z, cmap="seismic", vmin=-max_col, vmax=max_col, shading="auto" - ) - plt.colorbar(cbar, label="Log(Signal/Background)") - plt.xlabel(r"$\sin(\delta)$") - plt.ylabel("log(Energy)") - plt.savefig(plot_path + "spline.pdf") - plt.close() - - make_plot( - exp_hist, - savepath=base_plot_path + "bkg.pdf", - x_bins=sin_dec_bins, - y_bins=log_e_bins, - ) - - del mc - - except IOError: - pass + # plot 2D signal & S/B hist, and spline + for gamma in np.linspace(1.0, 4.0, 7): + plot_path = ( + base_plot_path + + "gamma=" + + str(gamma) + + "/" + + f"precision{get_gamma_precision(self.gamma_precision)}_" + f"smoothing{self.smoothing_order}" + "/" + ) + + try: + os.makedirs(plot_path) + except OSError: + pass + + def weight_function(sig_mc): + return energy_pdf.weight_mc(sig_mc, gamma) + + mc_hist = self.create_sig_2d_hist( + mc, sin_dec_bins, log_e_bins, weight_function + ) + self.make_plot( + mc_hist, + plot_path + sfx + "sig.pdf", + sin_dec_bins, + log_e_bins, + ) + + self.make_plot( + self.create_2d_ratio_hist( + exp, mc, sin_dec_bins, log_e_bins, weight_function + ), + plot_path + sfx + "SoB.pdf", + sin_dec_bins, + log_e_bins, + normed=False, + ) + + Z = [] + for s in sin_dec_bins: + z_line = [] + for e in log_e_bins: + # logging.debug(f'{e}, {s}') + try: + z_line.append(splines[gamma](e, s)[0][0]) + except: + z_line.append(splines[gamma]((e, s))) + Z.append(z_line) + + Z = np.array(Z).T + + max_col = min( + abs(min([min(row) for row in Z])), max([max(row) for row in Z]) + ) + + plt.figure() + ax = plt.subplot(111) + X, Y = np.meshgrid(sin_dec_bins, log_e_bins) + cbar = ax.pcolormesh( + X, Y, Z, cmap="seismic", vmin=-max_col, vmax=max_col, shading="auto" + ) + plt.colorbar(cbar, label="Log(Signal/Background)") + plt.xlabel(r"$\sin(\delta)$") + plt.ylabel("log(Energy)") + plt.savefig(plot_path + sfx + "spline.pdf") + plt.close() + + del mc + + except IOError: + pass + def make_splines(self, seasons): + """Make the S/B splines for all dataset seasons, + as well as the background spatial spline. + """ -def make_background_spline(season): - bkg_path = bkg_spline_path(season) - bkg = season.get_background_model() - sin_dec_bins = season.sin_dec_bins + logger.info( + "Splines will be made to calculate the Signal/Background ratio of " + + "the MC to data. The MC will be weighted with a power law for the signal, " + + f"for each gamma in: {list(get_gamma_support_points(self.gamma_precision))}. " + + "For the background, the MC is weighted according to the model assumption" + ) - bkg_spline = create_bkg_spatial_spline(bkg, sin_dec_bins) + for season in seasons.values(): + SoB_path = SoB_spline_path(season, **self.sob_dict) + self.make_individual_spline_set(season, SoB_path) + self.make_background_spline(season) - logger.info(f"Saving to {bkg_path}") - try: - os.makedirs(os.path.dirname(bkg_path)) - except OSError: - pass + def load_spline(self): + path = SoB_spline_path(self.season, **self.sob_dict) - with open(bkg_path, "wb") as f: - Pickle.dump(bkg_spline, f) + logger.debug(f"Loading SoB spline from {path}") - x_range = np.linspace(sin_dec_bins[0], sin_dec_bins[-1], 101) - plt.figure() - plt.plot(x_range, np.exp(bkg_spline(x_range))) - plt.ylabel(r"$P_{bkg}$ (spatial)") - plt.xlabel(r"$\sin(\delta)$") - savepath = get_base_sob_plot_dir(season) + try: + with open(path, "rb") as f: + res = Pickle.load(f) + except FileNotFoundError as err: + logger.info( + f"No cached spline found at {path}. Creating this file instead." + ) + logger.info(f"Cause: {err}") + self.make_individual_spline_set(self.season, path) + with open(path, "rb") as f: + res = Pickle.load(f) + + except ModuleNotFoundError as err: + logger.error( + "A spline was found but it seems incompatible with this installation." + ) + logger.error(f"Cause: {err}") + raise - try: - os.makedirs(os.path.dirname(savepath)) - except OSError: - pass + return res - plt.savefig(savepath + "bkg_spatial.pdf") - plt.close() +@SoB_splines.register_subclass("no_difffuse") +class NoDiffuseSpline(SoB_splines): + """Atmospheric-only background model. + The weights for creating the background + 2D histogram that goes into creating + the SoB splines are simply the + conventional atmospheric weights in the MC + """ -def load_spline(season, **kwargs): - path = SoB_spline_path(season, **kwargs) + def bkg_weights(self, exp: Table) -> np.ndarray: + logger.debug("Atmospheric-only background weights") + return np.asarray(exp["weight"]) + + +@SoB_splines.register_subclass("spl_difffuse") +class SPLDiffuseSpline(SoB_splines): + """Single powerlaw diffuse + atmospheric background. + Provided the (per-flavour) best-fit spectral parameters + phi & gamma, the per-flavour diffuse flux following SPL is + Phi_nu_antinu = phi * (E/100 TeV)**-gamma * 1e-18 /GeV cm2 s sr + The flux is multiplied with the livetime + to get the fluence and then with "ow" + in order to get the weights that go into + the background histogram. + The histogram is weighted w/ atmospheric + diffuse weights, + where atmospheric weights are the + conventional atmospheric weights in the MC. + If not provided, the default SPL parameters + are taken from the 10y ESTES SPL + (https://journals.aps.org/prd/abstract/10.1103/PhysRevD.110.022001) + w/ best-fit phi_0 = 1.68 & gamma = 2.58 + """ - logger.debug(f"Loading from {path}") + def __init__(self, season, SoB_dict) -> None: + super().__init__(season, SoB_dict) - try: - with open(path, "rb") as f: - res = Pickle.load(f) - except FileNotFoundError as err: - logger.info(f"No cached spline found at {path}. Creating this file instead.") - logger.info(f"Cause: {err}") - make_individual_spline_set(season, path, **kwargs) - with open(path, "rb") as f: - res = Pickle.load(f) + time_pdf = self.season.get_time_pdf() + self.livetime = time_pdf.get_livetime() * 3600 * 24 - except ModuleNotFoundError as err: - logger.error( - f"A spline was found but it seems incompatible with this installation." + def get_diffuse_flux(self, exp: Table) -> np.ndarray: + phi = self.sob_dict.get("phi0", 1.68) + gamma = self.sob_dict.get("spl_gamma", 2.58) + logger.info( + f"SPL diffuse flux used with best-fit phi = {phi} & gamma = {gamma}" ) - logger.error(f"Cause: {err}") - raise - - return res - - -def load_bkg_spatial_spline(season): - path = bkg_spline_path(season) - - logger.debug(f"Loading from {path}") - - try: - with open(path, "rb") as f: - res = Pickle.load(f) - except FileNotFoundError as err: - logger.info(f"No cached spline found at {path}. Creating this file instead.") - logger.info(f"Cause: {err}") - make_background_spline(season) - with open(path, "rb") as f: - res = Pickle.load(f) + return 1e-18 * phi * (np.asarray(exp["trueE"]) / 1e5) ** -gamma # /GeV cm2 s sr + + def bkg_weights(self, exp: Table) -> np.ndarray: + logger.debug("Background weights w/ SPL diffuse flux") + flux = self.get_diffuse_flux(exp) + fluence = flux * self.livetime + diff_weights = fluence * np.asarray(exp["ow"]) + return np.array(diff_weights + exp["weight"]) + + +@SoB_splines.register_subclass("bpl_difffuse") +class BPLDiffuseSpline(SoB_splines): + """Broken powerlaw diffuse + atmospheric background. + Provided the (per-flavour) best-fit spectral parameters + phi, gamma1, gamma2, and E_break the + per-flavour diffuse flux following a BPL is + Phi_nu_antinu = C * phi * (E_break/100 TeV)**-gamma1 * (E/E_break)**-gamma1 + for E < E_break & E_break > 100 TeV, and + Phi_nu_antinu = C * phi * (E_break/100 TeV)**-gamma2 * (E/E_break)**-gamma2 + for E >= E_break & E_break <= 100 TeV, + where C = 10**-18 /GeV cm2 s sr. + The flux is multiplied with the livetime + to get the fluence and then with "ow" + in order to get the weights that go into + the background histogram. + The histogram is weighted w/ atmospheric + diffuse weights, + where atmospheric weights are the + conventional atmospheric weights in the MC. + If not provided, the default BPL parameters + are taken from the 11y MESE BPL + (https://arxiv.org/abs/2507.22233), chosen against + the Combined Fit (CF) since it has a higher TS (see table 3), + w/ best-fit phi_0 = 2.28, gamma1 = 1.72, gamma2 = 2.84, E_break = 33.11 TeV + """ - except ModuleNotFoundError as err: - logger.error( - f"A spline was found but it seems incompatible with this installation." + def __init__(self, season, SoB_dict) -> None: + super().__init__(season, SoB_dict) + + time_pdf = self.season.get_time_pdf() + self.livetime = time_pdf.get_livetime() * 3600 * 24 + + def get_diffuse_flux( + self, exp: Table + ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + phi0 = self.sob_dict.get("phi0", 2.28) + gamma1 = self.sob_dict.get("gamma1", 1.72) + gamma2 = self.sob_dict.get("gamma2", 2.84) + log_Ebreak = self.sob_dict.get("logEbreak", 4.52) # log(E_break/GeV) + logger.info( + f"BPL diffuse flux used with best-fit phi = {phi0}, gamma1 = {gamma1}, " + + f"gamma2 = {gamma2}, and E_break = {10**log_Ebreak/1e3} TeV" ) - logger.error(f"Cause: {err}") - raise - - return res + E_break = 10**log_Ebreak # in GeV + if E_break <= 1e5: + phi = 1e-18 * phi0 * (E_break / 1e5) ** -gamma2 + else: + phi = 1e-18 * phi0 * (E_break / 1e5) ** -gamma1 + mask1 = np.nonzero(exp["trueE"] < E_break) + mask2 = np.nonzero(exp["trueE"] >= E_break) + phi1 = ( + phi * (np.asarray(exp[mask1]["trueE"]) / E_break) ** -gamma1 + ) # /GeV cm2 s sr + phi2 = ( + phi * (np.asarray(exp[mask2]["trueE"]) / E_break) ** -gamma2 + ) # /GeV cm2 s sr + return phi1, mask1[0], phi2, mask2[0] + + def bkg_weights(self, exp: Table) -> np.ndarray: + logger.debug("Background weights w/ BPL diffuse flux") + w = np.empty(len(exp), dtype=float) + flux1, mask1, flux2, mask2 = self.get_diffuse_flux(exp) + + # get diffuse + atm weights when E < E_break + fluence1 = flux1 * self.livetime + bpl1_w = fluence1 * np.asarray(exp[mask1]["ow"]) + w1 = np.array(bpl1_w + exp[mask1]["weight"]) + + # get diffuse + atm weights when E >= E_break + fluence2 = flux2 * self.livetime + bpl2_w = fluence2 * np.asarray(exp[mask2]["ow"]) + w2 = np.array(bpl2_w + exp[mask2]["weight"]) + + w[mask1] = w1 + w[mask2] = w2 + return w # def delete_old_splines():