From f527ca5a59e50449a12d841ae6091f5bccbf658a Mon Sep 17 00:00:00 2001 From: Joe Biagio McLaughlin Date: Tue, 19 Mar 2024 04:57:36 -0700 Subject: [PATCH] Passing shape-varying parameters to likelihood in bkg-only toy generation. --- flamedisx/non_asymptotic_inference.py | 30 ++++++++++++++++++++++++--- 1 file changed, 27 insertions(+), 3 deletions(-) diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py index 3bcdceac4..d16fa2477 100644 --- a/flamedisx/non_asymptotic_inference.py +++ b/flamedisx/non_asymptotic_inference.py @@ -339,7 +339,7 @@ def run_routine(self, mus_test=None, save_fits=False, if observed_data is not None: return observed_test_stats_collection else: - return test_stat_dists_SB_collection, test_stat_dists_B_collection + return test_stat_dists_SB_collection, testx_stat_dists_B_collection def sample_data_constraints(self, mu_test, signal_source_name, likelihood): """Internal function to sample the toy data and constraint central values @@ -372,8 +372,32 @@ def sample_data_constraints(self, mu_test, signal_source_name, likelihood): constraint_extra_args[f'{background_source}_expected_counts'] = tf.cast(draw, fd.float_type()) simulate_dict[f'{background_source}_rate_multiplier'] = expected_background_counts - simulate_dict[f'{signal_source_name}_rate_multiplier'] = mu_test - + + for svp in self.shape_varying_params.keys(): + if self.observed_test_stats is not None: + try: + conditional_bfs_observed = self.observed_test_stats[signal_source_name].conditional_best_fits + expected_svp = \ + conditional_bfs_observed[mu_test][f'{svp}'] + except Exception: + raise RuntimeError("Could not find observed conditional best fits") + # Case where we use the prior expected counts as constraint centers and simualted values + else: + expected_svp = self.common_defaults[svp] + + if svp in self.gaussian_constraint_widths: + draw = stats.norm.rvs(loc=expected_svp, + scale=self.gaussian_constraint_widths[svp]) + constraint_extra_args[f'{svp}'] = tf.cast(draw, fd.float_type()) + + elif svp in self.sample_other_constraints: + draw = self.sample_other_constraints[svp](expected_svp) + constraint_extra_args[f'{svp}'] = tf.cast(draw, fd.float_type()) + + simulate_dict[f'{svp}'] = expected_svp + simulate_dict[f'{signal_source_name}_rate_multiplier'] = mu_test + for k, v in simulate_dict.items(): + print(k, v) toy_data = likelihood.simulate(**simulate_dict) return simulate_dict, toy_data, constraint_extra_args