From 769800ffed5613fb2f56ad16be11e400f27433ed Mon Sep 17 00:00:00 2001 From: Wei Zha Date: Thu, 19 Jun 2025 12:10:24 -0700 Subject: [PATCH 1/9] Initial push Add basic code to float parameters in toy --- flamedisx/non_asymptotic_inference.py | 32 +++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py index ea916968..ec24c53c 100644 --- a/flamedisx/non_asymptotic_inference.py +++ b/flamedisx/non_asymptotic_inference.py @@ -349,6 +349,38 @@ def sample_data_constraints(self, mu_test, signal_source_name, likelihood): simulate_dict[f'{background_source}_rate_multiplier'] = tf.cast(expected_background_counts, fd.float_type()) simulate_dict[f'{signal_source_name}_rate_multiplier'] = tf.cast(mu_test, fd.float_type()) + for param_name in likelihood.param_names: + # For all other parameters + if '_rate_multiplier' in param_name: + continue + + # Initialize default parameter and bounds. Needs improvement later on. + try: + param_expect = likelihood.param_defaults[param_name] + except Exception: + raise RuntimeError(f"Default value of parameter {param_name} not found") + try: + param_bounds = likelihood.default_bounds[param_name] + except Exception: + raise RuntimeError(f"Bounds of parameter {param_name} not found") + + # Case of the conditional best fits + if self.observed_test_stats is not None: + try: + param_expect = conditional_bfs_observed[mu_test][param_name] + except Exception: + raise RuntimeError(f"Could not find observed conditional best fits for parameter {param_name}") + + # Sample constraint centers + if param_name in self.sample_other_constraints.keys(): + draw = self.sample_other_constraints[param_name](param_expect) + constraint_extra_args[param_name] = tf.cast(draw, fd.float_type()) + else: + # Hard-coded for now. Needs fixing. + draw = stats.norm.rvs(loc=param_expect, scale=0.04) + constraint_extra_args[param_name] = tf.cast(draw, fd.float_type()) + + if self.observed_test_stats is not None: conditional_bfs_observed = self.observed_test_stats[signal_source_name].conditional_best_fits[mu_test] non_rate_params_added = [] From a2d93a6dbfb4ac6afd7af6451b3aff803ce62557 Mon Sep 17 00:00:00 2001 From: Wei Zha Date: Thu, 19 Jun 2025 12:27:43 -0700 Subject: [PATCH 2/9] Texting Minor detail of parameter bounds --- flamedisx/non_asymptotic_inference.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py index ec24c53c..c223e000 100644 --- a/flamedisx/non_asymptotic_inference.py +++ b/flamedisx/non_asymptotic_inference.py @@ -373,11 +373,13 @@ def sample_data_constraints(self, mu_test, signal_source_name, likelihood): # Sample constraint centers if param_name in self.sample_other_constraints.keys(): + # Given the parameter center, use {sample_other_constraint} as draw draw = self.sample_other_constraints[param_name](param_expect) constraint_extra_args[param_name] = tf.cast(draw, fd.float_type()) else: - # Hard-coded for now. Needs fixing. - draw = stats.norm.rvs(loc=param_expect, scale=0.04) + # Hard-coded for now. If {sample_other_constraint} not provided, use 10% of bounds + param_range = tf.math.reduce_max(param_bounds) - tf.math.reduce_min(param_bounds) + draw = stats.norm.rvs(loc=param_expect, scale=0.1*param_range) constraint_extra_args[param_name] = tf.cast(draw, fd.float_type()) From 1ae41139c911aab294b256051e28253a326948f1 Mon Sep 17 00:00:00 2001 From: Wei Zha Date: Thu, 19 Jun 2025 12:36:07 -0700 Subject: [PATCH 3/9] Fix Forgot to add simulate_dict --- flamedisx/non_asymptotic_inference.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py index c223e000..398e8d82 100644 --- a/flamedisx/non_asymptotic_inference.py +++ b/flamedisx/non_asymptotic_inference.py @@ -382,6 +382,8 @@ def sample_data_constraints(self, mu_test, signal_source_name, likelihood): draw = stats.norm.rvs(loc=param_expect, scale=0.1*param_range) constraint_extra_args[param_name] = tf.cast(draw, fd.float_type()) + simulate_dict[param_name] = tf.cast(param_expect, fd.float_type()) + if self.observed_test_stats is not None: conditional_bfs_observed = self.observed_test_stats[signal_source_name].conditional_best_fits[mu_test] From 8e6b19511131ec9d661c5192328bf47d00735ff1 Mon Sep 17 00:00:00 2001 From: Wei Zha Date: Wed, 9 Jul 2025 22:09:42 -0700 Subject: [PATCH 4/9] Add fix_dict for parameters --- flamedisx/non_asymptotic_inference.py | 31 ++++++++++++++++++--------- flamedisx/templates.py | 2 +- 2 files changed, 22 insertions(+), 11 deletions(-) diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py index 398e8d82..c5e324ed 100644 --- a/flamedisx/non_asymptotic_inference.py +++ b/flamedisx/non_asymptotic_inference.py @@ -23,9 +23,11 @@ def __init__(self, likelihood): self.likelihood = likelihood def __call__(self, mu_test, signal_source_name, guess_dict, - asymptotic=False): + asymptotic=False, fix_dict_param=None): # To fix the signal RM in the conditional fit fix_dict = {f'{signal_source_name}_rate_multiplier': tf.cast(mu_test, fd.float_type())} + if fix_dict_param is not None: + fix_dict = {**fix_dict, **fix_dict_param} guess_dict_nuisance = guess_dict.copy() guess_dict_nuisance.pop(f'{signal_source_name}_rate_multiplier') @@ -180,7 +182,8 @@ def __init__( gaussian_constraint_widths: ty.Dict[str, float] = None, sample_other_constraints: ty.Dict[str, ty.Callable] = None, likelihood=None, - ntoys=1000): + ntoys=1000, + fix_dict=None): if gaussian_constraint_widths is None: gaussian_constraint_widths = dict() @@ -199,6 +202,7 @@ def __init__( self.expected_background_counts = expected_background_counts self.gaussian_constraint_widths = gaussian_constraint_widths self.sample_other_constraints = sample_other_constraints + self.fix_dict = fix_dict def run_routine(self, mus_test=None, save_fits=False, observed_data=None, @@ -374,15 +378,22 @@ def sample_data_constraints(self, mu_test, signal_source_name, likelihood): # Sample constraint centers if param_name in self.sample_other_constraints.keys(): # Given the parameter center, use {sample_other_constraint} as draw - draw = self.sample_other_constraints[param_name](param_expect) + if param_name not in self.fix_dict.keys(): + draw = self.sample_other_constraints[param_name](param_expect) + else: + draw = self.fix_dict[param_name] constraint_extra_args[param_name] = tf.cast(draw, fd.float_type()) else: - # Hard-coded for now. If {sample_other_constraint} not provided, use 10% of bounds - param_range = tf.math.reduce_max(param_bounds) - tf.math.reduce_min(param_bounds) - draw = stats.norm.rvs(loc=param_expect, scale=0.1*param_range) + # Hard-coded for now. If {sample_other_constraint} not provided, use 10% of bounds as gaussian width + if param_name not in self.fix_dict.keys(): + param_range = tf.math.reduce_max(param_bounds) - tf.math.reduce_min(param_bounds) + draw = stats.norm.rvs(loc=param_expect, scale=0.1*param_range) + else: + draw = self.fix_dict[param_name] constraint_extra_args[param_name] = tf.cast(draw, fd.float_type()) - simulate_dict[param_name] = tf.cast(param_expect, fd.float_type()) + if param_name not in self.fix_dict.keys(): + simulate_dict[param_name] = tf.cast(param_expect, fd.float_type()) if self.observed_test_stats is not None: @@ -449,8 +460,8 @@ 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 statistics - ts_result_SB = test_statistic_SB(mu_test, signal_source_name, guess_dict_SB) - ts_result_SB_disco = test_statistic_SB(0., signal_source_name, guess_dict_SB) + ts_result_SB = test_statistic_SB(mu_test, signal_source_name, guess_dict_SB, fix_dict_param = self.fix_dict) + ts_result_SB_disco = test_statistic_SB(0., signal_source_name, guess_dict_SB, fix_dict_param = self.fix_dict) # Save test statistics, and possibly fits ts_values_SB.append(ts_result_SB[0]) ts_values_SB_disco.append(ts_result_SB_disco[0]) @@ -493,7 +504,7 @@ 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) + ts_result_B = test_statistic_B(mu_test, signal_source_name, guess_dict_B, fix_dict_param = self.fix_dict) # Save test statistic, and possibly fits ts_values_B.append(ts_result_B[0]) if save_fits: diff --git a/flamedisx/templates.py b/flamedisx/templates.py index 6aa15b32..a6a18239 100644 --- a/flamedisx/templates.py +++ b/flamedisx/templates.py @@ -402,7 +402,7 @@ def simulate(self, n_events, fix_truth=None, full_annotate=False, # TODO: all other arguments are ignored, they make no sense # for this source. Should we warn about this? Remove them from def? - assert len(self.defaults) == 1 + #assert len(self.defaults) == 1 template_weights = tfp.math.batch_interp_regular_1d_grid( x=params[next(iter(self.defaults))], From ee913175fd9cc00609c515db51681ccd6cc869b5 Mon Sep 17 00:00:00 2001 From: Wei Zha Date: Wed, 9 Jul 2025 22:23:51 -0700 Subject: [PATCH 5/9] Small fix; Naming change --- flamedisx/non_asymptotic_inference.py | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py index c5e324ed..548b981f 100644 --- a/flamedisx/non_asymptotic_inference.py +++ b/flamedisx/non_asymptotic_inference.py @@ -202,7 +202,7 @@ def __init__( self.expected_background_counts = expected_background_counts self.gaussian_constraint_widths = gaussian_constraint_widths self.sample_other_constraints = sample_other_constraints - self.fix_dict = fix_dict + self.fix_dict_param = fix_dict def run_routine(self, mus_test=None, save_fits=False, observed_data=None, @@ -327,6 +327,11 @@ def sample_data_constraints(self, mu_test, signal_source_name, likelihood): """ simulate_dict = dict() constraint_extra_args = dict() + if self.fix_dict_param = None: + fix_dict_param = dict() + else: + fix_dict_param = self.fix_dict_param + for background_source in self.background_source_names: # Case where we use the conditional best fits as constraint centers and simulated values if self.observed_test_stats is not None: @@ -378,21 +383,21 @@ def sample_data_constraints(self, mu_test, signal_source_name, likelihood): # Sample constraint centers if param_name in self.sample_other_constraints.keys(): # Given the parameter center, use {sample_other_constraint} as draw - if param_name not in self.fix_dict.keys(): + if param_name not in fix_dict_param.keys(): draw = self.sample_other_constraints[param_name](param_expect) else: - draw = self.fix_dict[param_name] + draw = fix_dict_param[param_name] constraint_extra_args[param_name] = tf.cast(draw, fd.float_type()) else: # Hard-coded for now. If {sample_other_constraint} not provided, use 10% of bounds as gaussian width - if param_name not in self.fix_dict.keys(): + if param_name not in fix_dict_param.keys(): param_range = tf.math.reduce_max(param_bounds) - tf.math.reduce_min(param_bounds) draw = stats.norm.rvs(loc=param_expect, scale=0.1*param_range) else: - draw = self.fix_dict[param_name] + draw = fix_dict_param[param_name] constraint_extra_args[param_name] = tf.cast(draw, fd.float_type()) - if param_name not in self.fix_dict.keys(): + if param_name not in fix_dict_param.keys(): simulate_dict[param_name] = tf.cast(param_expect, fd.float_type()) @@ -460,8 +465,8 @@ 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 statistics - ts_result_SB = test_statistic_SB(mu_test, signal_source_name, guess_dict_SB, fix_dict_param = self.fix_dict) - ts_result_SB_disco = test_statistic_SB(0., signal_source_name, guess_dict_SB, fix_dict_param = self.fix_dict) + ts_result_SB = test_statistic_SB(mu_test, signal_source_name, guess_dict_SB, fix_dict_param = self.fix_dict_param) + ts_result_SB_disco = test_statistic_SB(0., signal_source_name, guess_dict_SB, fix_dict_param = self.fix_dict_param) # Save test statistics, and possibly fits ts_values_SB.append(ts_result_SB[0]) ts_values_SB_disco.append(ts_result_SB_disco[0]) @@ -504,7 +509,7 @@ 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, fix_dict_param = self.fix_dict) + ts_result_B = test_statistic_B(mu_test, signal_source_name, guess_dict_B, fix_dict_param = self.fix_dict_param) # Save test statistic, and possibly fits ts_values_B.append(ts_result_B[0]) if save_fits: From 9dfe413dd5b58f84ae36e7592b81f4911178cb49 Mon Sep 17 00:00:00 2001 From: Wei Zha Date: Wed, 9 Jul 2025 22:28:44 -0700 Subject: [PATCH 6/9] Oops, dumb bug --- flamedisx/non_asymptotic_inference.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py index 548b981f..989854dc 100644 --- a/flamedisx/non_asymptotic_inference.py +++ b/flamedisx/non_asymptotic_inference.py @@ -327,7 +327,7 @@ def sample_data_constraints(self, mu_test, signal_source_name, likelihood): """ simulate_dict = dict() constraint_extra_args = dict() - if self.fix_dict_param = None: + if self.fix_dict_param is None: fix_dict_param = dict() else: fix_dict_param = self.fix_dict_param From ed6c03eac992a51cf7f13fc35b940fb583b3b379 Mon Sep 17 00:00:00 2001 From: Wei Zha Date: Thu, 10 Jul 2025 19:43:35 -0700 Subject: [PATCH 7/9] Add in B8 uncertainty in toys --- flamedisx/non_asymptotic_inference.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py index 989854dc..11ae6b0a 100644 --- a/flamedisx/non_asymptotic_inference.py +++ b/flamedisx/non_asymptotic_inference.py @@ -210,7 +210,8 @@ def run_routine(self, mus_test=None, save_fits=False, generate_B_toys=False, simulate_dict_B=None, toy_data_B=None, constraint_extra_args_B=None, toy_batch=0, - asymptotic=False): + asymptotic=False, + vary_signal_dict=None): """If observed_data is passed, evaluate observed test statistics. Otherwise, obtain test statistic distributions (for both S+B and B-only). @@ -305,7 +306,8 @@ def run_routine(self, mus_test=None, save_fits=False, self.toy_test_statistic_dist(test_stat_dists_SB, test_stat_dists_B, test_stat_dists_SB_disco, mu_test, signal_source, likelihood, - save_fits=save_fits) + save_fits=save_fits, + vary_signal_dict=vary_signal_dict) if observed_data is not None: observed_test_stats_collection[signal_source] = observed_test_stats @@ -420,7 +422,7 @@ def sample_data_constraints(self, mu_test, signal_source_name, likelihood): def toy_test_statistic_dist(self, test_stat_dists_SB, test_stat_dists_B, test_stat_dists_SB_disco, mu_test, signal_source_name, likelihood, - save_fits=False): + save_fits=False, vary_signal_dict=None): """Internal function to get test statistic distribution. """ ts_values_SB = [] @@ -434,8 +436,12 @@ def toy_test_statistic_dist(self, test_stat_dists_SB, test_stat_dists_B, # Loop over toys for toy in tqdm(range(self.ntoys), desc='Doing toys'): + if vary_signal_dict is not None: + mu_sim = sps.norm.rvs(loc=mu_test, scale=vary_signal_dict[signal_source_name]) + else: + mu_sim = mu_test simulate_dict_SB, toy_data_SB, constraint_extra_args_SB = \ - self.sample_data_constraints(mu_test, signal_source_name, likelihood) + self.sample_data_constraints(mu_sim, signal_source_name, likelihood) # S+B toys From 6cc8419a7beae39b92ca9e82062d989a907b8c60 Mon Sep 17 00:00:00 2001 From: Wei Zha Date: Thu, 10 Jul 2025 19:54:39 -0700 Subject: [PATCH 8/9] Take Rob's advice, use gaussian constraint width as parameter uncertainty --- flamedisx/non_asymptotic_inference.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py index 11ae6b0a..a5790abc 100644 --- a/flamedisx/non_asymptotic_inference.py +++ b/flamedisx/non_asymptotic_inference.py @@ -365,7 +365,7 @@ def sample_data_constraints(self, mu_test, signal_source_name, likelihood): if '_rate_multiplier' in param_name: continue - # Initialize default parameter and bounds. Needs improvement later on. + # Initialize default parameter and bounds. try: param_expect = likelihood.param_defaults[param_name] except Exception: @@ -383,15 +383,16 @@ def sample_data_constraints(self, mu_test, signal_source_name, likelihood): raise RuntimeError(f"Could not find observed conditional best fits for parameter {param_name}") # Sample constraint centers - if param_name in self.sample_other_constraints.keys(): - # Given the parameter center, use {sample_other_constraint} as draw + if param_name in self.gaussian_constraint_widths.keys(): + # Given the parameter center, use {gaussian_constraint_widths} to get draw if param_name not in fix_dict_param.keys(): - draw = self.sample_other_constraints[param_name](param_expect) + draw = stats.norm.rvs(loc=param_expect, + scale=self.gaussian_constraint_widths[param_name]) else: draw = fix_dict_param[param_name] constraint_extra_args[param_name] = tf.cast(draw, fd.float_type()) else: - # Hard-coded for now. If {sample_other_constraint} not provided, use 10% of bounds as gaussian width + # If {gaussian_constraint_widths} not provided, use 10% of parameter bound as gaussian width if param_name not in fix_dict_param.keys(): param_range = tf.math.reduce_max(param_bounds) - tf.math.reduce_min(param_bounds) draw = stats.norm.rvs(loc=param_expect, scale=0.1*param_range) From 212ce8c06ab6d97f2bf5851fb2e376c2cf253712 Mon Sep 17 00:00:00 2001 From: Wei Zha Date: Thu, 10 Jul 2025 22:18:18 -0700 Subject: [PATCH 9/9] Sorry, think I'm going to switch back to sampling func --- flamedisx/non_asymptotic_inference.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py index a5790abc..fc0787ea 100644 --- a/flamedisx/non_asymptotic_inference.py +++ b/flamedisx/non_asymptotic_inference.py @@ -383,16 +383,21 @@ def sample_data_constraints(self, mu_test, signal_source_name, likelihood): raise RuntimeError(f"Could not find observed conditional best fits for parameter {param_name}") # Sample constraint centers - if param_name in self.gaussian_constraint_widths.keys(): +""" if param_name in self.gaussian_constraint_widths.keys(): # Given the parameter center, use {gaussian_constraint_widths} to get draw if param_name not in fix_dict_param.keys(): draw = stats.norm.rvs(loc=param_expect, scale=self.gaussian_constraint_widths[param_name]) else: draw = fix_dict_param[param_name] + constraint_extra_args[param_name] = tf.cast(draw, fd.float_type()) """ + + if param_name in self.sample_other_constraints.keys(): + # Given the parameter center, use {sample_other_constraint} as draw + draw = self.sample_other_constraints[param_name](param_expect) constraint_extra_args[param_name] = tf.cast(draw, fd.float_type()) else: - # If {gaussian_constraint_widths} not provided, use 10% of parameter bound as gaussian width + # If not provided, use 10% of parameter bound as gaussian width if param_name not in fix_dict_param.keys(): param_range = tf.math.reduce_max(param_bounds) - tf.math.reduce_min(param_bounds) draw = stats.norm.rvs(loc=param_expect, scale=0.1*param_range)