From a12beeb2b227f6c101ba3e35bc1698307dfb0bb0 Mon Sep 17 00:00:00 2001 From: Robert James Date: Thu, 11 Jan 2024 00:05:40 +1100 Subject: [PATCH] Add discovery stuff (bit rough for now). --- flamedisx/non_asymptotic_inference.py | 45 ++++++++++++++++++++++++--- 1 file changed, 40 insertions(+), 5 deletions(-) diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py index bcec10578..388750172 100644 --- a/flamedisx/non_asymptotic_inference.py +++ b/flamedisx/non_asymptotic_inference.py @@ -211,7 +211,8 @@ def run_routine(self, mus_test=None, save_fits=False, observed_test_stats=None, generate_B_toys=False, simulate_dict_B=None, toy_data_B=None, constraint_extra_args_B=None, - toy_batch=0): + toy_batch=0, + discovery=False): """If observed_data is passed, evaluate observed test statistics. Otherwise, obtain test statistic distributions (for both S+B and B-only). @@ -315,7 +316,8 @@ def run_routine(self, mus_test=None, save_fits=False, # Case where we want test statistic distributions else: self.toy_test_statistic_dist(test_stat_dists_SB, test_stat_dists_B, - mu_test, signal_source, likelihood, save_fits=save_fits) + mu_test, signal_source, likelihood, + save_fits=save_fits, discovery=discovery) if observed_data is not None: observed_test_stats_collection[signal_source] = observed_test_stats @@ -366,7 +368,8 @@ def sample_data_constraints(self, mu_test, signal_source_name, likelihood): return simulate_dict, toy_data, constraint_extra_args def toy_test_statistic_dist(self, test_stat_dists_SB, test_stat_dists_B, - mu_test, signal_source_name, likelihood, save_fits=False): + mu_test, signal_source_name, likelihood, + save_fits=False, discovery=False): """Internal function to get test statistic distribution. """ ts_values_SB = [] @@ -396,7 +399,10 @@ def toy_test_statistic_dist(self, test_stat_dists_SB, test_stat_dists_B, if value < 0.1: guess_dict_SB[key] = 0.1 # Evaluate test statistic - ts_result_SB = test_statistic_SB(mu_test, signal_source_name, guess_dict_SB) + if discovery: + ts_result_SB = test_statistic_SB(0., signal_source_name, guess_dict_SB) + else: + ts_result_SB = test_statistic_SB(mu_test, signal_source_name, guess_dict_SB) # Save test statistic, and possibly fits ts_values_SB.append(ts_result_SB[0]) if save_fits: @@ -424,7 +430,10 @@ def toy_test_statistic_dist(self, test_stat_dists_SB, test_stat_dists_B, # Create test statistic test_statistic_B = self.test_statistic(likelihood) # Evaluate test statistic - ts_result_B = test_statistic_B(mu_test, signal_source_name, guess_dict_B) + if discovery: + ts_result_B = test_statistic_B(0., signal_source_name, guess_dict_B) + else: + ts_result_B = test_statistic_B(mu_test, signal_source_name, guess_dict_B) # Save test statistic, and possibly fits ts_values_B.append(ts_result_B[0]) if save_fits: @@ -667,3 +676,29 @@ def get_bands(self, conf_level=0.1, quantiles=[0, 1, -1, 2, -2], bands[signal_source] = these_bands return bands + + def get_bands_discovery(self, quantiles=[0, 1, -1]): + """ + """ + bands = dict() + + # Loop over signal sources + for signal_source in self.signal_source_names: + # Get test statistic distribitions + test_stat_dists_SB = self.test_stat_dists_SB[signal_source] + test_stat_dists_B = self.test_stat_dists_B[signal_source] + + assert len(test_stat_dists_SB.ts_dists.keys()) == 1, 'Currently only support a single signal strength' + + these_p_vals = (100. - stats.percentileofscore(list(test_stat_dists_B.ts_dists.values())[0], + list(test_stat_dists_SB.ts_dists.values())[0], + kind='weak')) / 100. + these_p_vals = these_p_vals[these_p_vals > 0.] + these_disco_sigs = stats.norm.ppf(1. - these_p_vals) + + these_bands = dict() + for quantile in quantiles: + these_bands[quantile] = np.quantile(np.sort(these_disco_sigs), stats.norm.cdf(quantile)) + bands[signal_source] = these_bands + + return bands