From 71d0d3a3493915a71c6fcff30ac204c0395a468e Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Tue, 9 Apr 2024 17:05:11 +0100
Subject: [PATCH 01/48] Allow for mu rescaling when simulating from likelihood.

---
 flamedisx/likelihood.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/flamedisx/likelihood.py b/flamedisx/likelihood.py
index a2198f383..d3a05b1e2 100644
--- a/flamedisx/likelihood.py
+++ b/flamedisx/likelihood.py
@@ -340,7 +340,8 @@ def set_data(self,
                 np.concatenate([[0], stop_idx[:-1]]),
                 stop_idx])
 
-    def simulate(self, fix_truth=None, **params):
+    def simulate(self, fix_truth=None, alter_source_mus=False,
+                 **params):
         """Simulate events from sources.
         """
         params = self.prepare_params(params, free_all_rates=True)
@@ -352,6 +353,8 @@ def simulate(self, fix_truth=None, **params):
             rm = self._get_rate_mult(sname, params)
             mu = rm * s.mu_before_efficiencies(
                 **self._filter_source_kwargs(params, sname))
+            if alter_source_mus:
+                mu *= self.mu_estimators[sname](**self._filter_source_kwargs(params, sname))
             # Simulate this many events from source
             n_to_sim = np.random.poisson(mu)
             if n_to_sim == 0:

From f4c6726fe3148d5347e266309eacf8167cc0d24f Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Tue, 9 Apr 2024 22:13:18 +0100
Subject: [PATCH 02/48] Towards passing in likelihood directly in
 non-asymptotic inference routine.

---
 flamedisx/non_asymptotic_inference.py | 69 +++------------------------
 1 file changed, 6 insertions(+), 63 deletions(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 388750172..7b21636e2 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -135,9 +135,6 @@ class TSEvaluation():
         - signal_source_names: tuple of names for signal sources (e.g. WIMPs of different
             masses)
         - background_source_names: tuple of names for background sources
-        - sources: dictionary {sourcename: class} of all signal and background source classes
-        - arguments: dictionary {sourcename: {kwarg1: value, ...}, ...}, for
-            passing keyword arguments to source constructors
         - expected_background_counts: dictionary of expected counts for background sources
         - gaussian_constraint_widths: dictionary giving the constraint width for all sources
             using Gaussian constraints for their rate nuisance parameters
@@ -145,31 +142,19 @@ class TSEvaluation():
             in the toys for any sources using non-Gaussian constraints for their rate nuisance
             parameters. Argument to the function will be either the prior expected counts,
             or the number of counts at the conditional MLE, depending on the mode
-        - rm_bounds: dictionary {sourcename: (lower, upper)} to set fit bounds on the rate multipliers
-        - log_constraint_fn: logarithm of the constraint function used in the likelihood. Any arguments
-            which aren't fit parameters, such as those determining constraint means for toys, will need
-            passing via the set_constraint_extra_args() function
+        - likelihood: BLAH
         - ntoys: number of toys that will be run to get test statistic distributions
-        - batch_size: batch size that will be used for the RM fits
     """
     def __init__(
             self,
             test_statistic: TestStatistic.__class__,
             signal_source_names: ty.Tuple[str],
             background_source_names: ty.Tuple[str],
-            sources: ty.Dict[str, fd.Source.__class__],
-            arguments: ty.Dict[str, ty.Dict[str, ty.Union[int, float]]] = None,
             expected_background_counts: ty.Dict[str, float] = None,
             gaussian_constraint_widths: ty.Dict[str, float] = None,
             sample_other_constraints: ty.Dict[str, ty.Callable] = None,
-            rm_bounds: ty.Dict[str, ty.Tuple[float, float]] = None,
-            log_constraint_fn: ty.Callable = None,
-            ntoys=1000,
-            batch_size=10000):
-
-        for key in sources.keys():
-            if key not in arguments.keys():
-                arguments[key] = dict()
+            likelihood=None,
+            ntoys=1000):
 
         if gaussian_constraint_widths is None:
             gaussian_constraint_widths = dict()
@@ -177,34 +162,17 @@ def __init__(
         if sample_other_constraints is None:
             sample_other_constraints = dict()
 
-        if rm_bounds is None:
-            rm_bounds = dict()
-        else:
-            for bounds in rm_bounds.values():
-                assert bounds[0] >= 0., 'Currently do not support negative rate multipliers'
-
-        if log_constraint_fn is None:
-            def log_constraint_fn(**kwargs):
-                return 0.
-            self.log_constraint_fn = log_constraint_fn
-        else:
-            self.log_constraint_fn = log_constraint_fn
-
         self.ntoys = ntoys
-        self.batch_size = batch_size
 
+        self.likelihood = likelihood
         self.test_statistic = test_statistic
 
         self.signal_source_names = signal_source_names
         self.background_source_names = background_source_names
 
-        self.sources = sources
-        self.arguments = arguments
-
         self.expected_background_counts = expected_background_counts
         self.gaussian_constraint_widths = gaussian_constraint_widths
         self.sample_other_constraints = sample_other_constraints
-        self.rm_bounds = rm_bounds
 
     def run_routine(self, mus_test=None, save_fits=False,
                     observed_data=None,
@@ -266,33 +234,8 @@ def run_routine(self, mus_test=None, save_fits=False,
             test_stat_dists_SB = TestStatisticDistributions()
             test_stat_dists_B = TestStatisticDistributions()
 
-            sources = dict()
-            arguments = dict()
-            for background_source in self.background_source_names:
-                sources[background_source] = self.sources[background_source]
-                arguments[background_source] = self.arguments[background_source]
-            sources[signal_source] = self.sources[signal_source]
-            arguments[signal_source] = self.arguments[signal_source]
-
-            # Create likelihood of TemplateSources
-            likelihood = fd.LogLikelihood(sources=sources,
-                                          arguments=arguments,
-                                          progress=False,
-                                          batch_size=self.batch_size,
-                                          free_rates=tuple([sname for sname in sources.keys()]))
-
-            rm_bounds = dict()
-            if signal_source in self.rm_bounds.keys():
-                rm_bounds[signal_source] = self.rm_bounds[signal_source]
-            for background_source in self.background_source_names:
-                if background_source in self.rm_bounds.keys():
-                    rm_bounds[background_source] = self.rm_bounds[background_source]
-
-            # Pass rate multiplier bounds to likelihood
-            likelihood.set_rate_multiplier_bounds(**rm_bounds)
-
-            # Pass constraint function to likelihood
-            likelihood.set_log_constraint(self.log_constraint_fn)
+            # Get likelihood
+            likelihood = self.likelihood
 
             # Where we want to generate B-only toys
             if generate_B_toys:

From 41a8f06e2c4d9204a82cfc01a1d0d9669290d175 Mon Sep 17 00:00:00 2001
From: Makayla Trask <mtrask@ucsb.edu>
Date: Tue, 9 Apr 2024 22:50:57 -0700
Subject: [PATCH 03/48] change set_data to handle new method

---
 flamedisx/non_asymptotic_inference.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 7b21636e2..ce0e9750c 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -28,9 +28,9 @@ def __call__(self, mu_test, signal_source_name, guess_dict):
         guess_dict_nuisance.pop(f'{signal_source_name}_rate_multiplier')
 
         # Conditional fit
-        bf_conditional = self.likelihood.bestfit(fix=fix_dict, guess=guess_dict_nuisance, suppress_warnings=True)
+        bf_conditional = self.likelihood.bestfit(fix=fix_dict, guess=guess_dict_nuisance, suppress_warnings=True, use_hessian=False)
         # Uncnditional fit
-        bf_unconditional = self.likelihood.bestfit(guess=guess_dict, suppress_warnings=True)
+        bf_unconditional = self.likelihood.bestfit(guess=guess_dict, suppress_warnings=True, use_hessian=False)
 
         # Return the test statistic, unconditional fit and conditional fit
         return self.evaluate(bf_unconditional, bf_conditional), bf_unconditional, bf_conditional
@@ -407,7 +407,14 @@ def get_observed_test_stat(self, observed_test_stats, observed_data,
         likelihood.set_constraint_extra_args(**constraint_extra_args)
 
         # Set data
-        likelihood.set_data(observed_data)
+        try:
+            for key in likelihood.likelihoods.keys():
+                observed_data_ll = observed_data[observed_data['component'] == key]
+                likelihood.set_data(observed_data_ll, key)
+        
+        except:
+            likelihood.set_data(observed_data)
+
         # Create test statistic
         test_statistic = self.test_statistic(likelihood)
         # Guesses for fit

From 632845c5e37b800459b185aa4aeb8dbb3029861a Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Wed, 10 Apr 2024 17:25:26 +0100
Subject: [PATCH 04/48] Remove sources not used in a given fit from likelihood.
 NOTE: non_asymptotic_inference routines now currently require a combined
 likelihood...

---
 flamedisx/non_asymptotic_inference.py | 22 +++++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index ce0e9750c..7ebc00cad 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -4,6 +4,8 @@
 from tqdm.auto import tqdm
 import typing as ty
 
+from copy import deepcopy
+
 import tensorflow as tf
 
 export, __all__ = fd.exporter()
@@ -235,7 +237,18 @@ def run_routine(self, mus_test=None, save_fits=False,
             test_stat_dists_B = TestStatisticDistributions()
 
             # Get likelihood
-            likelihood = self.likelihood
+            likelihood = deepcopy(self.likelihood)
+
+            assert hasattr(likelihood, 'likelihoods'), 'Logic only currently works for combined likelihood'
+            for ll in likelihood.likelihoods.values():
+                sources_remove = []
+                params_remove = []
+                for sname in ll.sources:
+                    if (sname != signal_source) and (sname not in self.background_source_names):
+                        sources_remove.append(sname)
+                        params_remove.append(f'{sname}_rate_multiplier')
+            likelihood.rebuild(sources_remove=sources_remove,
+                               params_remove=params_remove)
 
             # Where we want to generate B-only toys
             if generate_B_toys:
@@ -407,18 +420,17 @@ def get_observed_test_stat(self, observed_test_stats, observed_data,
         likelihood.set_constraint_extra_args(**constraint_extra_args)
 
         # Set data
-        try:
+        if hasattr(likelihood, 'likelihoods'):
             for key in likelihood.likelihoods.keys():
                 observed_data_ll = observed_data[observed_data['component'] == key]
                 likelihood.set_data(observed_data_ll, key)
-        
-        except:
+        else:
             likelihood.set_data(observed_data)
 
         # Create test statistic
         test_statistic = self.test_statistic(likelihood)
         # Guesses for fit
-        guess_dict = {f'{signal_source_name}_rate_multiplier': mu_test}
+        guess_dict = {f'{signal_source_name}_rate_multiplier': 0.1}
         for background_source in self.background_source_names:
             guess_dict[f'{background_source}_rate_multiplier'] = self.expected_background_counts[background_source]
         for key, value in guess_dict.items():

From 46d2d99cde1cd369afb1debf5bb091d7f364998f Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Thu, 11 Apr 2024 11:44:58 +0100
Subject: [PATCH 05/48] Hessian back on.

---
 flamedisx/non_asymptotic_inference.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 7ebc00cad..bbf78d0ec 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -30,9 +30,9 @@ def __call__(self, mu_test, signal_source_name, guess_dict):
         guess_dict_nuisance.pop(f'{signal_source_name}_rate_multiplier')
 
         # Conditional fit
-        bf_conditional = self.likelihood.bestfit(fix=fix_dict, guess=guess_dict_nuisance, suppress_warnings=True, use_hessian=False)
+        bf_conditional = self.likelihood.bestfit(fix=fix_dict, guess=guess_dict_nuisance, suppress_warnings=True)
         # Uncnditional fit
-        bf_unconditional = self.likelihood.bestfit(guess=guess_dict, suppress_warnings=True, use_hessian=False)
+        bf_unconditional = self.likelihood.bestfit(guess=guess_dict, suppress_warnings=True)
 
         # Return the test statistic, unconditional fit and conditional fit
         return self.evaluate(bf_unconditional, bf_conditional), bf_unconditional, bf_conditional

From 179d0ee5645d14ef07d2f1616b44957aeba31e8d Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Thu, 11 Apr 2024 13:31:21 +0100
Subject: [PATCH 06/48] Handle setting data in batch routines for combined
 likelihood.

---
 flamedisx/non_asymptotic_inference.py | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index bbf78d0ec..94de1e7c8 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -346,7 +346,11 @@ def toy_test_statistic_dist(self, test_stat_dists_SB, test_stat_dists_B,
             # Shift the constraint in the likelihood based on the background RMs we drew
             likelihood.set_constraint_extra_args(**constraint_extra_args_SB)
             # Set data
-            likelihood.set_data(toy_data_SB)
+            if hasattr(likelihood, 'likelihoods'):
+                for component, data in toy_data_SB.items():
+                    likelihood.set_data(data, component)
+            else:
+                likelihood.set_data(toy_data_SB)
             # Create test statistic
             test_statistic_SB = self.test_statistic(likelihood)
             # Guesses for fit
@@ -382,7 +386,11 @@ def toy_test_statistic_dist(self, test_stat_dists_SB, test_stat_dists_B,
             # Shift the constraint in the likelihood based on the background RMs we drew
             likelihood.set_constraint_extra_args(**constraint_extra_args_B)
             # Set data
-            likelihood.set_data(toy_data_B)
+            if hasattr(likelihood, 'likelihoods'):
+                for component, data in toy_data_B.items():
+                    likelihood.set_data(data, component)
+            else:
+                likelihood.set_data(toy_data_B)
             # Create test statistic
             test_statistic_B = self.test_statistic(likelihood)
             # Evaluate test statistic
@@ -421,9 +429,8 @@ def get_observed_test_stat(self, observed_test_stats, observed_data,
 
         # Set data
         if hasattr(likelihood, 'likelihoods'):
-            for key in likelihood.likelihoods.keys():
-                observed_data_ll = observed_data[observed_data['component'] == key]
-                likelihood.set_data(observed_data_ll, key)
+            for component, data in observed_data.items():
+                likelihood.set_data(data, component)
         else:
             likelihood.set_data(observed_data)
 

From 9de483b9cffdd79695ab8476987bc71ad90c3c1b Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Wed, 17 Apr 2024 17:30:54 +0100
Subject: [PATCH 07/48] Add asymptotic p-value calculation.

---
 flamedisx/non_asymptotic_inference.py | 83 ++++++++++++++++++++-------
 1 file changed, 63 insertions(+), 20 deletions(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 94de1e7c8..bcfaff10d 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -22,7 +22,8 @@ class TestStatistic():
     def __init__(self, likelihood):
         self.likelihood = likelihood
 
-    def __call__(self, mu_test, signal_source_name, guess_dict):
+    def __call__(self, mu_test, signal_source_name, guess_dict,
+                 asymptotic=False):
         # To fix the signal RM in the conditional fit
         fix_dict = {f'{signal_source_name}_rate_multiplier': mu_test}
 
@@ -35,7 +36,11 @@ def __call__(self, mu_test, signal_source_name, guess_dict):
         bf_unconditional = self.likelihood.bestfit(guess=guess_dict, suppress_warnings=True)
 
         # Return the test statistic, unconditional fit and conditional fit
-        return self.evaluate(bf_unconditional, bf_conditional), bf_unconditional, bf_conditional
+        if not asymptotic:
+            return self.evaluate(bf_unconditional, bf_conditional), bf_unconditional, bf_conditional
+        else:
+            return self.evaluate_asymptotic_pval(bf_unconditional, bf_conditional,
+                                                 mu_test), bf_unconditional, bf_conditional
 
 
 @export
@@ -55,6 +60,23 @@ def evaluate(self, bf_unconditional, bf_conditional):
         else:
             return ts
 
+    def evaluate_asymptotic_pval(self, bf_unconditional, bf_conditional, mu_test):
+        ll_conditional = self.likelihood(**bf_conditional)
+        ll_unconditional = self.likelihood(**bf_unconditional)
+
+        ts = -2. * (ll_conditional - ll_unconditional)
+
+        cov = 2. * self.likelihood.inverse_hessian(bf_unconditional)
+        sigma_mu = np.sqrt(cov[0][0])
+
+        if ts < (mu_test**2 / sigma_mu**2):
+            F = 2. * stats.norm.cdf(np.sqrt(ts)) - 1.
+        else:
+            F = stats.norm.cdf(np.sqrt(ts)) + stats.norm.cdf((ts + (mu_test**2 / sigma_mu**2)) / (2. * mu_test / sigma_mu)) - 1.
+
+        pval = 1. - F
+        return pval
+
 
 @export
 class TestStatisticDistributions():
@@ -182,7 +204,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,
-                    discovery=False):
+                    discovery=False,
+                    asymptotic=False):
         """If observed_data is passed, evaluate observed test statistics. Otherwise,
         obtain test statistic distributions (for both S+B and B-only).
 
@@ -268,7 +291,8 @@ def run_routine(self, mus_test=None, save_fits=False,
                 # Case where we want observed test statistics
                 if observed_data is not None:
                     self.get_observed_test_stat(observed_test_stats, observed_data,
-                                                mu_test, signal_source, likelihood, save_fits=save_fits)
+                                                mu_test, signal_source, likelihood, save_fits=save_fits,
+                                                asymptotic=asymptotic)
                 # Case where we want test statistic distributions
                 else:
                     self.toy_test_statistic_dist(test_stat_dists_SB, test_stat_dists_B,
@@ -416,7 +440,8 @@ def toy_test_statistic_dist(self, test_stat_dists_SB, test_stat_dists_B,
             test_stat_dists_B.add_conditional_best_fit(mu_test, conditional_bfs_B)
 
     def get_observed_test_stat(self, observed_test_stats, observed_data,
-                               mu_test, signal_source_name, likelihood, save_fits=False):
+                               mu_test, signal_source_name, likelihood, save_fits=False,
+                               asymptotic=False):
         """Internal function to evaluate observed test statistic.
         """
         # The constraints are centered on the expected values
@@ -444,7 +469,8 @@ def get_observed_test_stat(self, observed_test_stats, observed_data,
             if value < 0.1:
                 guess_dict[key] = 0.1
         # Evaluate test statistic
-        ts_result = test_statistic(mu_test, signal_source_name, guess_dict)
+        ts_result = test_statistic(mu_test, signal_source_name, guess_dict,
+                                   asymptotic=asymptotic)
 
         # Add to the test statistic collection
         observed_test_stats.add_test_stat(mu_test, ts_result[0])
@@ -505,7 +531,7 @@ def interp_helper(x, y, crossing_points, crit_val,
         else:
             return (crit_val - x_left) * gradient + y_left
 
-    def get_p_vals(self, conf_level, use_CLs=False):
+    def get_p_vals(self, conf_level, use_CLs=False, asymptotic=False):
         """Internal function to get p-value curves.
         """
         p_sb_collection = dict()
@@ -513,6 +539,10 @@ def get_p_vals(self, conf_level, use_CLs=False):
         p_b_collection = dict()
         # Loop over signal sources
         for signal_source in self.signal_source_names:
+            if asymptotic:
+                p_sb_collection[signal_source] = self.observed_test_stats[signal_source]
+                continue
+
             # Get test statistic distribitions and observed test statistics
             test_stat_dists_SB = self.test_stat_dists_SB[signal_source]
             test_stat_dists_B = self.test_stat_dists_B[signal_source]
@@ -529,13 +559,17 @@ def get_p_vals(self, conf_level, use_CLs=False):
                 powers = test_stat_dists_B.get_p_vals(crit_vals)
                 powers_collection[signal_source] = powers
 
+        if asymptotic:
+            return p_sb_collection
+
         if use_CLs:
             return p_sb_collection, p_b_collection
         else:
             return p_sb_collection, powers_collection
 
     def get_interval(self, conf_level=0.1, pcl_level=0.16,
-                     use_CLs=False):
+                     use_CLs=False,
+                     asymptotic=False):
         """Get frequentist confidence interval.
 
         Arguments:
@@ -547,26 +581,33 @@ def get_interval(self, conf_level=0.1, pcl_level=0.16,
                 (https://inspirehep.net/literature/599622), and the final return value
                 will be the p-value curves under H1
         """
-        if use_CLs:
-            p_sb, p_b = self.get_p_vals(conf_level, use_CLs=True)
+        if not asymptotic:
+            if use_CLs:
+                p_sb, p_b = self.get_p_vals(conf_level, use_CLs=True)
+            else:
+                p_sb, powers = self.get_p_vals(conf_level, use_CLs=False)
         else:
-            p_sb, powers = self.get_p_vals(conf_level, use_CLs=False)
+            p_sb = self.get_p_vals(conf_level, use_CLs=True, asymptotic=True)
 
         lower_lim_all = dict()
         upper_lim_all = dict()
         # Loop over signal sources
         for signal_source in self.signal_source_names:
-            these_p_sb = p_sb[signal_source]
+            if not asymptotic:
+                these_p_sb = p_sb[signal_source]
+            else:
+                these_p_sb = p_sb[signal_source].test_stats
             mus = np.array(list(these_p_sb.keys()))
             p_vals = np.array(list(these_p_sb.values()))
 
-            if use_CLs:
-                these_p_b = p_b[signal_source]
-                p_vals_b = np.array(list(these_p_b.values()))
-                p_vals = p_vals / (1. - p_vals_b + 1e-10)
-            else:
-                these_powers = powers[signal_source]
-                pws = np.array(list(these_powers.values()))
+            if not asymptotic:
+                if use_CLs:
+                    these_p_b = p_b[signal_source]
+                    p_vals_b = np.array(list(these_p_b.values()))
+                    p_vals = p_vals / (1. - p_vals_b + 1e-10)
+                else:
+                    these_powers = powers[signal_source]
+                    pws = np.array(list(these_powers.values()))
 
             # Find points where the p-value curve cross the critical value, decreasing
             upper_lims = np.argwhere(np.diff(np.sign(p_vals - np.ones_like(p_vals) * conf_level)) < 0.).flatten()
@@ -586,7 +627,7 @@ def get_interval(self, conf_level=0.1, pcl_level=0.16,
             upper_lim = self.interp_helper(mus, p_vals, upper_lims, conf_level,
                                            rising_edge=False, inverse=True)
 
-            if use_CLs is False:
+            if use_CLs is False and not asymptotic:
                 M0 = self.interp_helper(mus, pws, upper_lims, upper_lim,
                                         rising_edge=False, inverse=False)
                 if M0 < pcl_level:
@@ -599,6 +640,8 @@ def get_interval(self, conf_level=0.1, pcl_level=0.16,
             lower_lim_all[signal_source] = lower_lim
             upper_lim_all[signal_source] = upper_lim
 
+        if asymptotic:
+            return lower_lim_all, upper_lim_all
         if use_CLs is False:
             return lower_lim_all, upper_lim_all, p_sb, powers
         else:

From 32247b62d669ddff88df6a5c0d407f1e89fb8242 Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Thu, 18 Apr 2024 12:44:10 +0100
Subject: [PATCH 08/48] Fix constraint batching issue.

---
 flamedisx/likelihood.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/flamedisx/likelihood.py b/flamedisx/likelihood.py
index d3a05b1e2..6e36399cc 100644
--- a/flamedisx/likelihood.py
+++ b/flamedisx/likelihood.py
@@ -524,10 +524,14 @@ def _log_likelihood(self,
             0.)
         if dsetname == self.dsetnames[0]:
             if constraint_extra_args is None:
-                ll += self.log_constraint(**params_unstacked)
+                ll += tf.where(tf.equal(i_batch, tf.constant(0, dtype=fd.int_type())),
+                               self.log_constraint(**params_unstacked),
+                               0.)
             else:
                 kwargs = {**params_unstacked, **constraint_extra_args}
-                ll += self.log_constraint(**kwargs)
+                ll += tf.where(tf.equal(i_batch, tf.constant(0, dtype=fd.int_type())),
+                               self.log_constraint(**kwargs),
+                               0.)
 
         # Autodifferentiation. This is why we use tensorflow:
         grad = tf.gradients(ll, grad_par_stack)[0]

From 0ebe58dc3ab94737da77fdd688575651670d68a6 Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Wed, 8 May 2024 20:47:48 -0700
Subject: [PATCH 09/48] Extract discovery significance.

---
 flamedisx/non_asymptotic_inference.py | 32 ++++++++++++---------------
 1 file changed, 14 insertions(+), 18 deletions(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index bcfaff10d..2989c322c 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -689,28 +689,24 @@ def get_bands(self, conf_level=0.1, quantiles=[0, 1, -1, 2, -2],
 
         return bands
 
-    def get_bands_discovery(self, quantiles=[0, 1, -1]):
+    def get_disco_sig(self):
         """
         """
-        bands = dict()
+        disco_sigs = 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)
+            # Get observed (mu = 0) test statistic and B (m = 0) test statistic distribition
+            try:
+                observed_test_stat = self.observed_test_stats[signal_source].test_stats[0.]
+                test_stat_dist_B = self.test_stat_dists_B[signal_source].ts_dists[0.]
+            except Exception:
+                raise RuntimeError("Error: did you scan over mu = 0?")
 
-            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
+            p_val = (100. - stats.percentileofscore(test_stat_dist_B,
+                                                    observed_test_stat,
+                                                    kind='weak')) / 100.
+            disco_sig = stats.norm.ppf(1. - p_val)
+            disco_sigs[signal_source] = disco_sig
 
-        return bands
+        return disco_sigs

From 73f83ab5f649cd0b66d143efd689926aa2ddda8a Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Tue, 14 May 2024 17:41:37 +1000
Subject: [PATCH 10/48] Extra TS evaluation for expexted discovery significance
 plot.

---
 flamedisx/non_asymptotic_inference.py | 31 +++++++++++++++------------
 1 file changed, 17 insertions(+), 14 deletions(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 2989c322c..02154d3fa 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -204,7 +204,6 @@ 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,
-                    discovery=False,
                     asymptotic=False):
         """If observed_data is passed, evaluate observed test statistics. Otherwise,
         obtain test statistic distributions (for both S+B and B-only).
@@ -251,12 +250,14 @@ def run_routine(self, mus_test=None, save_fits=False,
 
         observed_test_stats_collection = dict()
         test_stat_dists_SB_collection = dict()
+        test_stat_dists_SB_disco_collection = dict()
         test_stat_dists_B_collection = dict()
 
         # Loop over signal sources
         for signal_source in self.signal_source_names:
             observed_test_stats = ObservedTestStatistics()
             test_stat_dists_SB = TestStatisticDistributions()
+            test_stat_dists_SB_disco = TestStatisticDistributions()
             test_stat_dists_B = TestStatisticDistributions()
 
             # Get likelihood
@@ -296,19 +297,22 @@ 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,
+                                                 test_stat_dists_SB_disco,
                                                  mu_test, signal_source, likelihood,
-                                                 save_fits=save_fits, discovery=discovery)
+                                                 save_fits=save_fits)
 
             if observed_data is not None:
                 observed_test_stats_collection[signal_source] = observed_test_stats
             else:
                 test_stat_dists_SB_collection[signal_source] = test_stat_dists_SB
+                test_stat_dists_SB_disco_collection[signal_source] = test_stat_dists_SB_disco
                 test_stat_dists_B_collection[signal_source] = test_stat_dists_B
 
         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, test_stat_dists_SB_disco_collection, \
+                test_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
@@ -348,11 +352,13 @@ 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,
+                                test_stat_dists_SB_disco,
                                 mu_test, signal_source_name, likelihood,
-                                save_fits=False, discovery=False):
+                                save_fits=False):
         """Internal function to get test statistic distribution.
         """
         ts_values_SB = []
+        ts_values_SB_disco = []
         ts_values_B = []
         if save_fits:
             unconditional_bfs_SB = []
@@ -382,13 +388,12 @@ def toy_test_statistic_dist(self, test_stat_dists_SB, test_stat_dists_B,
             for key, value in guess_dict_SB.items():
                 if value < 0.1:
                     guess_dict_SB[key] = 0.1
-            # Evaluate test statistic
-            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
+            # 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)
+            # Save test statistics, and possibly fits
             ts_values_SB.append(ts_result_SB[0])
+            ts_values_SB_disco.append(ts_result_SB_disco[0])
             if save_fits:
                 unconditional_bfs_SB.append(ts_result_SB[1])
                 conditional_bfs_SB.append(ts_result_SB[2])
@@ -418,10 +423,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
-            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)
+            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:
@@ -430,6 +432,7 @@ def toy_test_statistic_dist(self, test_stat_dists_SB, test_stat_dists_B,
 
         # Add to the test statistic distributions
         test_stat_dists_SB.add_ts_dist(mu_test, ts_values_SB)
+        test_stat_dists_SB_disco.add_ts_dist(mu_test, ts_values_SB_disco)
         test_stat_dists_B.add_ts_dist(mu_test, ts_values_B)
 
         # Possibly save the fits

From b41c38258a900e164d6233981754262c64e934fa Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Wed, 15 May 2024 14:23:36 +1000
Subject: [PATCH 11/48] Add temporary interpolator fix...

---
 flamedisx/templates.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/flamedisx/templates.py b/flamedisx/templates.py
index 610b9be59..c62da7e1e 100644
--- a/flamedisx/templates.py
+++ b/flamedisx/templates.py
@@ -98,7 +98,9 @@ def differential_rates_numpy(self, data):
 
         if self._interpolator:
             # transpose since RegularGridInterpolator expects (n_points, n_dims)
-            return self._interpolator(data.T)
+            interp_diff_rates = self._interpolator(data.T)
+            lookup_diff_rates = self._mh_diff_rate.lookup(*data)
+            return np.where(interp_diff_rates <= 0., lookup_diff_rates, interp_diff_rates)
         else:
             return self._mh_diff_rate.lookup(*data)
 

From c0dca874e75de63159025a6928286b3f29587dd4 Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Mon, 20 May 2024 04:39:16 -0700
Subject: [PATCH 12/48] Asymptotic calculation of median discovery potential
 curve.

---
 flamedisx/non_asymptotic_inference.py | 38 +++++++++++++++++++++++++--
 1 file changed, 36 insertions(+), 2 deletions(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 02154d3fa..197908bd9 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -507,12 +507,14 @@ def __init__(
             signal_source_names: ty.Tuple[str],
             observed_test_stats: ObservedTestStatistics,
             test_stat_dists_SB: TestStatisticDistributions,
-            test_stat_dists_B: TestStatisticDistributions):
+            test_stat_dists_B: TestStatisticDistributions,
+            test_stat_dists_SB_disco: TestStatisticDistributions=None):
 
         self.signal_source_names = signal_source_names
         self.observed_test_stats = observed_test_stats
         self.test_stat_dists_SB = test_stat_dists_SB
         self.test_stat_dists_B = test_stat_dists_B
+        self.test_stat_dists_SB_disco = test_stat_dists_SB_disco
 
     @staticmethod
     def interp_helper(x, y, crossing_points, crit_val,
@@ -655,6 +657,11 @@ def upper_lims_bands(self, pval_curve, mus, conf_level):
         return self.interp_helper(mus, pval_curve, upper_lims, conf_level,
                                   rising_edge=False, inverse=True)
 
+    def critical_disco_value(self, disco_pot_curve, mus, discovery_sigma):
+        crossing_point = np.argwhere(np.diff(np.sign(disco_pot_curve - np.ones_like(disco_pot_curve) * discovery_sigma)) > 0.).flatten()
+        return self.interp_helper(mus, disco_pot_curve, crossing_point, discovery_sigma,
+                                  rising_edge=True, inverse=True)
+
     def get_bands(self, conf_level=0.1, quantiles=[0, 1, -1, 2, -2],
                   use_CLs=False):
         """
@@ -680,7 +687,7 @@ def get_bands(self, conf_level=0.1, quantiles=[0, 1, -1, 2, -2],
                                                              kind='weak') / 100.
                     these_p_vals = these_p_vals / (1. - these_p_vals_b + 1e-10)
                 mus.append(mu_test)
-                p_val_curves.append(these_p_vals)
+                p_val_curves.append(these_p_vals[:2130])
 
             p_val_curves = np.transpose(np.stack(p_val_curves, axis=0))
             upper_lims_bands = np.apply_along_axis(self.upper_lims_bands, 1, p_val_curves, mus, conf_level)
@@ -713,3 +720,30 @@ def get_disco_sig(self):
             disco_sigs[signal_source] = disco_sig
 
         return disco_sigs
+
+    def get_median_disco_asymptotic(self, sigma_level=3):
+        """
+        """
+        medians = dict()
+
+        # Loop over signal sources
+        for signal_source in self.signal_source_names:
+            # Get test statistic distribitions
+            test_stat_dists_SB_disco = self.test_stat_dists_SB_disco[signal_source]
+
+            mus = []
+            disco_sig_curves = []
+            # Loop over signal rate multipliers
+            for mu_test, ts_values in test_stat_dists_SB_disco.ts_dists.items():
+                these_disco_sigs = np.sqrt(ts_values)
+
+                mus.append(mu_test)
+                disco_sig_curves.append(these_disco_sigs[:2130])
+
+            disco_sig_curves = np.stack(disco_sig_curves, axis=0)
+            median_disco_sigs = [np.median(disco_sigs) for disco_sigs in disco_sig_curves]
+
+            median_crossing_point = self.critical_disco_value(median_disco_sigs, mus, 3)
+            medians[signal_source] = median_crossing_point
+
+        return medians

From ba5d9f36816cf9e9151f30cef91962ed243407cd Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Fri, 24 May 2024 17:06:49 +1000
Subject: [PATCH 13/48] Fix, I think, for memory leak...

---
 flamedisx/inference.py                |  7 +++++--
 flamedisx/likelihood.py               | 13 ++++---------
 flamedisx/non_asymptotic_inference.py |  6 +++---
 3 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/flamedisx/inference.py b/flamedisx/inference.py
index cfeb9ca82..67541e24d 100644
--- a/flamedisx/inference.py
+++ b/flamedisx/inference.py
@@ -176,10 +176,13 @@ def _dict_to_array(self, x: dict) -> np.array:
 
     def _array_to_dict(self, x: ty.Union[np.ndarray, tf.Tensor]) -> dict:
         """Convert from array/tensor to {parameter: value} dictionary"""
+        x = tf.cast(x, fd.float_type())
         assert isinstance(x, (np.ndarray, tf.Tensor))
         assert len(x) == len(self.arg_names)
-        return {k: x[i]
-                for i, k in enumerate(self.arg_names)}
+        param_dict = dict()
+        for i, k in enumerate(self.arg_names):
+            param_dict[k] = tf.gather(x, i)
+        return param_dict
 
     def normalize(self,
                   x: ty.Union[dict, np.ndarray],
diff --git a/flamedisx/likelihood.py b/flamedisx/likelihood.py
index 6e36399cc..75ed217d3 100644
--- a/flamedisx/likelihood.py
+++ b/flamedisx/likelihood.py
@@ -379,6 +379,7 @@ def __call__(self, **kwargs):
         assert 'second_order' not in kwargs, 'Roep gewoon log_likelihood aan'
         return self.log_likelihood(second_order=False, **kwargs)[0]
 
+    @tf.function
     def log_likelihood(self, second_order=False,
                        omit_grads=tuple(), **kwargs):
         params = self.prepare_params(kwargs)
@@ -414,14 +415,14 @@ def log_likelihood(self, second_order=False,
                     empty_batch=empty_batch,
                     constraint_extra_args=self.constraint_extra_args,
                     **params)
-                ll += results[0].numpy().astype(np.float64)
+                ll += results[0]
 
                 if self.param_names:
                     if results[1] is None:
                         raise ValueError("TensorFlow returned None as gradient!")
-                    llgrad += results[1].numpy().astype(np.float64)
+                    llgrad += results[1]
                     if second_order:
-                        llgrad2 += results[2].numpy().astype(np.float64)
+                        llgrad2 += results[2]
 
         if second_order:
             return ll, llgrad, llgrad2
@@ -434,11 +435,6 @@ def minus2_ll(self, *, omit_grads=tuple(), **kwargs):
         return -2 * ll, -2 * grad, hess
 
     def prepare_params(self, kwargs, free_all_rates=False):
-        for k in kwargs:
-            if k not in self.param_defaults:
-                if k.endswith('_rate_multiplier') and free_all_rates:
-                    continue
-                raise ValueError(f"Unknown parameter {k}")
         return {**self.param_defaults, **fd.values_to_constants(kwargs)}
 
     def _get_rate_mult(self, sname, kwargs):
@@ -488,7 +484,6 @@ def mu(self, *,
                    * self.mu_estimators[sname](**filtered_params))
         return mu
 
-    @tf.function
     def _log_likelihood(self,
                         i_batch, dsetname, data_tensor, batch_info,
                         omit_grads=tuple(), second_order=False,
diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 197908bd9..4c17f6b13 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -25,7 +25,7 @@ def __init__(self, likelihood):
     def __call__(self, mu_test, signal_source_name, guess_dict,
                  asymptotic=False):
         # To fix the signal RM in the conditional fit
-        fix_dict = {f'{signal_source_name}_rate_multiplier': mu_test}
+        fix_dict = {f'{signal_source_name}_rate_multiplier': tf.cast(mu_test, fd.float_type())}
 
         guess_dict_nuisance = guess_dict.copy()
         guess_dict_nuisance.pop(f'{signal_source_name}_rate_multiplier')
@@ -344,8 +344,8 @@ def sample_data_constraints(self, mu_test, signal_source_name, likelihood):
                 draw = self.sample_other_constraints[background_source](expected_background_counts)
                 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
+            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())
 
         toy_data = likelihood.simulate(**simulate_dict)
 

From e5039f8714599b36bb5d96346edae5b50c87bae4 Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Fri, 24 May 2024 17:25:49 +1000
Subject: [PATCH 14/48] Bug fix.

---
 flamedisx/non_asymptotic_inference.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 4c17f6b13..1889f0143 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -465,12 +465,12 @@ def get_observed_test_stat(self, observed_test_stats, observed_data,
         # Create test statistic
         test_statistic = self.test_statistic(likelihood)
         # Guesses for fit
-        guess_dict = {f'{signal_source_name}_rate_multiplier': 0.1}
+        guess_dict = {f'{signal_source_name}_rate_multiplier': tf.cast(0.1, fd.float_type())}
         for background_source in self.background_source_names:
-            guess_dict[f'{background_source}_rate_multiplier'] = self.expected_background_counts[background_source]
+            guess_dict[f'{background_source}_rate_multiplier'] = tf.cast(self.expected_background_counts[background_source], fd.float_type())
         for key, value in guess_dict.items():
             if value < 0.1:
-                guess_dict[key] = 0.1
+                guess_dict[key] = tf.cast(0.1, fd.float_type())
         # Evaluate test statistic
         ts_result = test_statistic(mu_test, signal_source_name, guess_dict,
                                    asymptotic=asymptotic)

From 196ac25b4fb0f4dfee7d38ea3b14f6058e9dda6b Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Mon, 27 May 2024 16:51:12 +1000
Subject: [PATCH 15/48] Old (linear) version of 1D temlate morphing.

---
 flamedisx/templates.py | 87 ++++++++++++++++++++++++------------------
 1 file changed, 50 insertions(+), 37 deletions(-)

diff --git a/flamedisx/templates.py b/flamedisx/templates.py
index c62da7e1e..b92e44d57 100644
--- a/flamedisx/templates.py
+++ b/flamedisx/templates.py
@@ -11,6 +11,8 @@
 
 import flamedisx as fd
 
+from copy import deepcopy
+
 export, __all__ = fd.exporter()
 
 
@@ -101,6 +103,7 @@ def differential_rates_numpy(self, data):
             interp_diff_rates = self._interpolator(data.T)
             lookup_diff_rates = self._mh_diff_rate.lookup(*data)
             return np.where(interp_diff_rates <= 0., lookup_diff_rates, interp_diff_rates)
+
         else:
             return self._mh_diff_rate.lookup(*data)
 
@@ -211,6 +214,9 @@ def __init__(
                 template, bin_edges, axis_names, events_per_bin, interpolate)
             for _, template in params_and_templates]
 
+        # We assume that mu does not change as we interpolate
+        self.mu = self._templates[0].mu
+
         # Grab parameter names. Promote first set of values to defaults.
         self.n_templates = n_templates = len(self._templates)
         assert n_templates > 0
@@ -226,9 +232,10 @@ def __init__(
         #
         # When evaluated at the exact location of a template, the result has 1
         # in the corresponding template's position, and zeros elsewhere.
-        _template_weights = scipy.interpolate.LinearNDInterpolator(
-            points=np.asarray([list(params.values()) for params, _ in params_and_templates]),
-            values=np.eye(n_templates))
+
+        _template_weights = scipy.interpolate.interp1d(
+            x=np.asarray([list(params.values())[0] for params, _ in params_and_templates]),
+            y=np.eye(n_templates))
 
         # Unfortunately TensorFlow has no equivalent of LinearNDInterpolator,
         # only interpolators that work on rectilinear grids. Thus, instead of
@@ -249,24 +256,6 @@ def __init__(
         # for use in tensorflow interpolation.
         _grid_weights = _template_weights(*_full_grid_coordinates)
 
-        # The expected number of events must also be interpolated.
-        # For consistency, it must be done in the same way (first interpolate
-        # to a regular grid, then linearly from there).
-        # (n_templates,) array
-        _template_mus = np.asarray([
-            template.mu for template in self._templates])
-        self._grid_mus = np.average(
-            # numpy won't let us get away with a size-1 axis here, we have to
-            # actually repeat the values. (If we had jax we could just vmap...)
-            np.repeat(_template_mus[:, None], n_grid_points, axis=1),
-            axis=0,
-            weights=_grid_weights.reshape(n_templates, n_grid_points))
-        assert self._grid_mus.shape == (n_templates,)
-        self._mu_interpolator = scipy.interpolate.RegularGridInterpolator(
-            points=_grid_coordinates,
-            values=self._grid_mus.reshape(_full_grid_coordinates[0].shape),
-            method='linear')
-
         # Generate a random column name to use to store the diff rates
         # of observed events under every template
         self.column = (
@@ -280,7 +269,7 @@ def __init__(
         # usual Source.scan_model_functions.
         self.f_dims = dict()
         self.f_params = dict()
-        self.defaults = defaults
+        self.defaults = {k: tf.cast(v, fd.float_type()) for k, v in defaults.items()}
 
         # This is needed in tensorflow, so convert it now
         self._grid_coordinates = tuple([fd.np_to_tf(np.asarray(g)) for g in _grid_coordinates])
@@ -305,28 +294,20 @@ def _annotate(self):
             for template in self._templates]).T)
 
     def mu_before_efficiencies(self, **params):
-        return self.estimate_mu(self, **params)
+        return self.mu
 
-    def estimate_mu(self, **params):
-        """Estimate the number of events expected from the template source.
-        """
-        # TODO: maybe need .item or something here?
-        return self._mu_interpolator([
-            params.get(param, default)
-            for param, default in self.defaults.items()])
+    def estimate_mu(self, n_trials=None, **params):
+        return self.mu
 
     def _differential_rate(self, data_tensor, ptensor):
         # Compute template weights at this parameter point
         # (n_templates,) tensor
         # (The axis order is weird here. It seems to work...)
-        permutation = (
-            [self._grid_weights.ndim - 1]
-            + list(range(0, self._grid_weights.ndim - 1)))
-        template_weights = tfp.math.batch_interp_rectilinear_nd_grid(
+        template_weights = tfp.math.batch_interp_regular_1d_grid(
             x=ptensor[None, :],
-            x_grid_points=self._grid_coordinates,
-            y_ref=tf.transpose(self._grid_weights, permutation),
-            axis=1,
+            x_ref_min=self._grid_coordinates[0][0],
+            x_ref_max=self._grid_coordinates[0][-1],
+            y_ref=self._grid_weights,
         )[:, 0]
         # Ensure template weights sum to one.
         template_weights /= tf.reduce_sum(template_weights)
@@ -340,3 +321,35 @@ def _differential_rate(self, data_tensor, ptensor):
         return tf.reduce_sum(
             template_diffrates * template_weights[None, :],
             axis=1)
+
+    def simulate(self, n_events, fix_truth=None, full_annotate=False,
+                 keep_padding=False, **params):
+        """Simulate n events.
+        """
+        if fix_truth:
+            raise NotImplementedError("TemplateSource does not yet support fix_truth")
+        assert isinstance(n_events, (int, float)), \
+            f"n_events must be an int or float, not {type(n_events)}"
+
+        # 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
+
+        template_weights = tfp.math.batch_interp_regular_1d_grid(
+            x=params[next(iter(self.defaults))],
+            x_ref_min=self._grid_coordinates[0][0],
+            x_ref_max=self._grid_coordinates[0][-1],
+            y_ref=self._grid_weights,
+        )
+
+        template_weights /= tf.reduce_sum(template_weights)
+
+        template_epb = [template._mh_events_per_bin for template in self._templates]
+        template_epb_combine = deepcopy(template_epb[0])
+        template_epb_combine.histogram = np.sum([template.histogram * weight for template, weight in
+                                                 zip(template_epb, template_weights)], axis=0)
+
+        return pd.DataFrame(dict(zip(
+            self._templates[0].axis_names,
+            template_epb_combine.get_random(n_events).T)))
\ No newline at end of file

From a7e18cbad77761f8842393725defaf772666a3c9 Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Tue, 28 May 2024 14:50:31 +1000
Subject: [PATCH 16/48] Fixes to morphing implementation (simple).

---
 flamedisx/templates.py | 12 +-----------
 1 file changed, 1 insertion(+), 11 deletions(-)

diff --git a/flamedisx/templates.py b/flamedisx/templates.py
index b92e44d57..fa51465e9 100644
--- a/flamedisx/templates.py
+++ b/flamedisx/templates.py
@@ -103,7 +103,6 @@ def differential_rates_numpy(self, data):
             interp_diff_rates = self._interpolator(data.T)
             lookup_diff_rates = self._mh_diff_rate.lookup(*data)
             return np.where(interp_diff_rates <= 0., lookup_diff_rates, interp_diff_rates)
-
         else:
             return self._mh_diff_rate.lookup(*data)
 
@@ -232,7 +231,6 @@ def __init__(
         #
         # When evaluated at the exact location of a template, the result has 1
         # in the corresponding template's position, and zeros elsewhere.
-
         _template_weights = scipy.interpolate.interp1d(
             x=np.asarray([list(params.values())[0] for params, _ in params_and_templates]),
             y=np.eye(n_templates))
@@ -265,21 +263,13 @@ def __init__(
         # ... this column will hold an array, with one entry per template
         self.array_columns = ((self.column, n_templates),)
 
-        # This source has parameters but no model functions, so we can't do the
-        # usual Source.scan_model_functions.
-        self.f_dims = dict()
-        self.f_params = dict()
-        self.defaults = {k: tf.cast(v, fd.float_type()) for k, v in defaults.items()}
-
         # This is needed in tensorflow, so convert it now
         self._grid_coordinates = tuple([fd.np_to_tf(np.asarray(g)) for g in _grid_coordinates])
         self._grid_weights = fd.np_to_tf(_grid_weights)
 
         super().__init__(*args, **kwargs)
 
-    def scan_model_functions(self):
-        # Don't do anything here, already set defaults etc. in __init__ above
-        pass
+        self.defaults = {**self.defaults,**{k: tf.cast(v, fd.float_type()) for k, v in defaults.items()}}
 
     def extra_needed_columns(self):
         return super().extra_needed_columns() + [self.column]

From 2c963e0a38b141f7d91af01d1b0ed071d802cadd Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Wed, 29 May 2024 20:52:26 +1000
Subject: [PATCH 17/48] Fix batched_differential_rate with morphing source.

---
 flamedisx/templates.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/flamedisx/templates.py b/flamedisx/templates.py
index fa51465e9..2b70814b0 100644
--- a/flamedisx/templates.py
+++ b/flamedisx/templates.py
@@ -205,6 +205,7 @@ def __init__(
             axis_names=None,
             events_per_bin=False,
             interpolate=False,
+            _skip_tf_init=False,
             *args,
             **kwargs):
 
@@ -270,6 +271,9 @@ def __init__(
         super().__init__(*args, **kwargs)
 
         self.defaults = {**self.defaults,**{k: tf.cast(v, fd.float_type()) for k, v in defaults.items()}}
+        self.parameter_index = fd.index_lookup_dict(self.defaults.keys())
+        if not _skip_tf_init:
+            self.trace_differential_rate()
 
     def extra_needed_columns(self):
         return super().extra_needed_columns() + [self.column]

From e01fe36a6e0c3563ccecced976a9d42c14f2259c Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Fri, 31 May 2024 17:13:22 +1000
Subject: [PATCH 18/48] Fix memory leak fix... Data tensor was not updating
 before.

---
 flamedisx/likelihood.py | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/flamedisx/likelihood.py b/flamedisx/likelihood.py
index 75ed217d3..97862217e 100644
--- a/flamedisx/likelihood.py
+++ b/flamedisx/likelihood.py
@@ -379,14 +379,10 @@ def __call__(self, **kwargs):
         assert 'second_order' not in kwargs, 'Roep gewoon log_likelihood aan'
         return self.log_likelihood(second_order=False, **kwargs)[0]
 
-    @tf.function
     def log_likelihood(self, second_order=False,
                        omit_grads=tuple(), **kwargs):
         params = self.prepare_params(kwargs)
         n_grads = len(self.param_defaults) - len(omit_grads)
-        ll = 0.
-        llgrad = np.zeros(n_grads, dtype=np.float64)
-        llgrad2 = np.zeros((n_grads, n_grads), dtype=np.float64)
 
         for dsetname in self.dsetnames:
             # Getting this from the batch_info tensor is much slower
@@ -399,14 +395,18 @@ def log_likelihood(self, second_order=False,
             else:
                 empty_batch = False
 
+            ll = {i_batch: 0. for i_batch in range(n_batches)}
+            llgrad = np.zeros(n_grads, dtype=np.float64)
+            llgrad2 = np.zeros((n_grads, n_grads), dtype=np.float64)
+
             for i_batch in range(n_batches):
                 # Iterating over tf.range seems much slower!
                 if empty_batch:
                     batch_data_tensor = None
                 else:
-                    batch_data_tensor = self.data_tensors[dsetname][i_batch]
+                    batch_data_tensor = tf.gather(self.data_tensors[dsetname], i_batch)
                 results = self._log_likelihood(
-                    tf.constant(i_batch, dtype=fd.int_type()),
+                    i_batch,
                     dsetname=dsetname,
                     data_tensor=batch_data_tensor,
                     batch_info=self.batch_info,
@@ -415,7 +415,7 @@ def log_likelihood(self, second_order=False,
                     empty_batch=empty_batch,
                     constraint_extra_args=self.constraint_extra_args,
                     **params)
-                ll += results[0]
+                ll[i_batch] = results[0]
 
                 if self.param_names:
                     if results[1] is None:
@@ -425,8 +425,8 @@ def log_likelihood(self, second_order=False,
                         llgrad2 += results[2]
 
         if second_order:
-            return ll, llgrad, llgrad2
-        return ll, llgrad, None
+            return np.sum(list(ll.values())), llgrad, llgrad2
+        return np.sum(list(ll.values())), llgrad, None
 
     def minus2_ll(self, *, omit_grads=tuple(), **kwargs):
         result = self.log_likelihood(omit_grads=omit_grads, **kwargs)
@@ -484,6 +484,7 @@ def mu(self, *,
                    * self.mu_estimators[sname](**filtered_params))
         return mu
 
+    @tf.function
     def _log_likelihood(self,
                         i_batch, dsetname, data_tensor, batch_info,
                         omit_grads=tuple(), second_order=False,

From 7a4b1086029e25994fa95c91e5723b26ab8f6c54 Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Fri, 31 May 2024 18:37:50 +1000
Subject: [PATCH 19/48] Add in normalisation variation with morphing (Josh).

---
 flamedisx/templates.py | 31 +++++++++++++++++++++++++++----
 1 file changed, 27 insertions(+), 4 deletions(-)

diff --git a/flamedisx/templates.py b/flamedisx/templates.py
index 2b70814b0..0a2fece46 100644
--- a/flamedisx/templates.py
+++ b/flamedisx/templates.py
@@ -201,6 +201,7 @@ class MultiTemplateSource(fd.Source):
     def __init__(
             self,
             params_and_templates: ty.Tuple[ty.Dict[str, float], ty.Any],
+            params_and_normalisations:ty.Tuple[ty.Dict[str, float], float],
             bin_edges=None,
             axis_names=None,
             events_per_bin=False,
@@ -214,7 +215,7 @@ def __init__(
                 template, bin_edges, axis_names, events_per_bin, interpolate)
             for _, template in params_and_templates]
 
-        # We assume that mu does not change as we interpolate
+        # We will include mu variation separately
         self.mu = self._templates[0].mu
 
         # Grab parameter names. Promote first set of values to defaults.
@@ -248,7 +249,6 @@ def __init__(
                                   for params, _ in params_and_templates)))
             for param in defaults])
         _full_grid_coordinates = np.meshgrid(*_grid_coordinates, indexing='ij')
-        n_grid_points = np.prod([len(x) for x in _grid_coordinates])
 
         # Evaluate our irregular-grid scipy-interpolator on the grid.
         # This gives an array of shape (n_templates, ngrid_dim0, ngrid_dim1, ...)
@@ -268,6 +268,15 @@ def __init__(
         self._grid_coordinates = tuple([fd.np_to_tf(np.asarray(g)) for g in _grid_coordinates])
         self._grid_weights = fd.np_to_tf(_grid_weights)
 
+        param_vals = np.asarray([list(params.values())[0] for params, _ in params_and_templates])
+        self.pmin = tf.constant(min(param_vals), fd.float_type())
+        self.pmax = tf.constant(max(param_vals), fd.float_type())
+
+        normalisations = np.array([norm for _, norm in params_and_normalisations])
+        self.normalisations = tf.convert_to_tensor(normalisations / normalisations[0],
+                                                   fd.float_type())
+
+
         super().__init__(*args, **kwargs)
 
         self.defaults = {**self.defaults,**{k: tf.cast(v, fd.float_type()) for k, v in defaults.items()}}
@@ -291,7 +300,14 @@ def mu_before_efficiencies(self, **params):
         return self.mu
 
     def estimate_mu(self, n_trials=None, **params):
-        return self.mu
+        norm = tfp.math.batch_interp_regular_1d_grid(
+                    x=params[self.param_name],
+                    x_ref_min=self.pmin,
+                    x_ref_max=self.pmax,
+                    y_ref=self.normalisations,
+                    )
+
+        return tf.reshape(norm, shape=[]) * self.mu
 
     def _differential_rate(self, data_tensor, ptensor):
         # Compute template weights at this parameter point
@@ -310,10 +326,17 @@ def _differential_rate(self, data_tensor, ptensor):
         # (n_events, n_templates) tensor
         template_diffrates = self._fetch(self.column, data_tensor)
 
+        norm = tfp.math.batch_interp_regular_1d_grid(
+                x=ptensor[None, :],
+                x_ref_min=self.pmin,
+                x_ref_max=self.pmax,
+                y_ref=self.normalisations,
+                )
+
         # Compute weighted average of diff rates
         # (n_events,) tensor
         return tf.reduce_sum(
-            template_diffrates * template_weights[None, :],
+            norm * template_diffrates * template_weights[None, :],
             axis=1)
 
     def simulate(self, n_events, fix_truth=None, full_annotate=False,

From bcb95567d6a8e1f56f4e026378871718713f4a1f Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Fri, 31 May 2024 21:32:11 +1000
Subject: [PATCH 20/48] Cubic spline template morphing: very hacky, for now
 (Josh).

---
 flamedisx/__init__.py           |   3 +
 flamedisx/templates.py          |  89 ++++++++----
 flamedisx/tfbspline/__init__.py |   1 +
 flamedisx/tfbspline/bspline.py  | 237 ++++++++++++++++++++++++++++++++
 4 files changed, 307 insertions(+), 23 deletions(-)
 create mode 100644 flamedisx/tfbspline/__init__.py
 create mode 100644 flamedisx/tfbspline/bspline.py

diff --git a/flamedisx/__init__.py b/flamedisx/__init__.py
index fa5169098..3b89e39a2 100644
--- a/flamedisx/__init__.py
+++ b/flamedisx/__init__.py
@@ -41,3 +41,6 @@
 # Custom TFP files
 # Access through fd.tfp_files.xxx
 from . import tfp_files
+
+# TEMPORARY,I HOPE
+from . import tfbspline
diff --git a/flamedisx/templates.py b/flamedisx/templates.py
index 0a2fece46..f0803eda6 100644
--- a/flamedisx/templates.py
+++ b/flamedisx/templates.py
@@ -9,6 +9,8 @@
 import tensorflow as tf
 import tensorflow_probability as tfp
 
+from flamedisx.tfbspline import bspline
+
 import flamedisx as fd
 
 from copy import deepcopy
@@ -271,11 +273,17 @@ def __init__(
         param_vals = np.asarray([list(params.values())[0] for params, _ in params_and_templates])
         self.pmin = tf.constant(min(param_vals), fd.float_type())
         self.pmax = tf.constant(max(param_vals), fd.float_type())
+        self.pvals = tf.convert_to_tensor(param_vals, fd.float_type())
 
         normalisations = np.array([norm for _, norm in params_and_normalisations])
         self.normalisations = tf.convert_to_tensor(normalisations / normalisations[0],
                                                    fd.float_type())
 
+        # Assume equi-spacing!
+        self.dstep=self.pvals[1]-self.pvals[0]
+        # Need to pad domain.. four might be excessive
+        self.pvals=list(np.arange(self.pvals[0]-4*(self.dstep),self.pvals[-1]+5*(self.dstep),self.dstep))
+        self.array_columns = ((self.column, n_templates+8),)
 
         super().__init__(*args, **kwargs)
 
@@ -290,12 +298,51 @@ def extra_needed_columns(self):
     def _annotate(self):
         """Add columns needed in inference to self.data
         """
-        # Get array of differential rates for each template.
-        # Outer list() is to placate pandas, which does not like array columns..
+        #construct tensor of knots
+        #requires a tensor of elements
+        #data is stored as [[d_evt1^h1,d_evt1^h2..],[d_evt2^h1,d_evt2^h2..]]
+        # so just need to construct and x-values object and let data column handle y-values
+        #with some padding for the domain!
+        Nk=len(self.pvals)
+        knot_range=self.pvals[-1]-self.pvals[0]
+        linear_shift=2*self.dstep/knot_range
+        start=min(self.pvals)
+        end=max(self.pvals)
+        self.original_range=tf.constant(end-start,dtype=fd.float_type())
+        self.max_pos=tf.constant(Nk- 2,dtype=fd.float_type())
+
+        self.start=tf.constant(start,dtype=fd.float_type())
+        self.linear_shift=tf.constant(linear_shift,dtype=fd.float_type())
+        self.linear_shift_shift=tf.constant(knot_range/2,dtype=fd.float_type())
+
         self.data[self.column] = list(np.asarray([
             template.differential_rates_numpy(self.data)
             for template in self._templates]).T)
 
+        linear_interp_padded_diff_rates=[]
+        for diff_rate_per_hist in self.data[self.column]:
+
+            if np.sum(diff_rate_per_hist[:2])>0:
+                left_edge=scipy.interpolate.interp1d(
+                    self.pvals[4:6],diff_rate_per_hist[:2],
+                    kind='linear',fill_value="extrapolate",
+                    bounds_error=False)(self.pvals[:4])
+            else:
+                left_edge=list(np.repeat(diff_rate_per_hist[0],4))
+
+            if np.sum(diff_rate_per_hist[-2:])>0:
+                right_edge=scipy.interpolate.interp1d(
+                    self.pvals[-6:-4],diff_rate_per_hist[-2:],
+                    kind='linear',fill_value="extrapolate",
+                    bounds_error=False)(self.pvals[-4:])
+            else:
+                right_edge=list(np.repeat(diff_rate_per_hist[-1],4))
+
+            linear_interp_padded_diff_rates.append(np.concatenate([left_edge,diff_rate_per_hist,right_edge]))
+
+        self.data[self.column]=linear_interp_padded_diff_rates
+        self.tensor_xvals=tf.convert_to_tensor([self.pvals for d in self.data[self.column]],dtype=fd.float_type())
+
     def mu_before_efficiencies(self, **params):
         return self.mu
 
@@ -309,23 +356,19 @@ def estimate_mu(self, n_trials=None, **params):
 
         return tf.reshape(norm, shape=[]) * self.mu
 
-    def _differential_rate(self, data_tensor, ptensor):
-        # Compute template weights at this parameter point
-        # (n_templates,) tensor
-        # (The axis order is weird here. It seems to work...)
-        template_weights = tfp.math.batch_interp_regular_1d_grid(
-            x=ptensor[None, :],
-            x_ref_min=self._grid_coordinates[0][0],
-            x_ref_max=self._grid_coordinates[0][-1],
-            y_ref=self._grid_weights,
-        )[:, 0]
-        # Ensure template weights sum to one.
-        template_weights /= tf.reduce_sum(template_weights)
-
-        # Fetch precomputed diff rates for each template.
-        # (n_events, n_templates) tensor
-        template_diffrates = self._fetch(self.column, data_tensor)
+    def bspline_interpolate_per_bin(self, param,knots):
+        def interp(knots_for_event):
+            #second order non-cyclical b-spline with varying knots
+            #returns [x,y] so ignore x
+            #hackiest shit ever
+            return tf.reduce_sum(bspline.interpolate(knots_for_event,
+                                                     self.max_pos*(param-self.start)/self.original_range+self.linear_shift*(param -self.linear_shift_shift), 2, False) \
+                                * tf.constant([0,1],dtype=fd.float_type()))
+        #vectorized map over all events
+        y=tf.vectorized_map(interp,elems=knots)
+        return y
 
+    def _differential_rate(self, data_tensor, ptensor):
         norm = tfp.math.batch_interp_regular_1d_grid(
                 x=ptensor[None, :],
                 x_ref_min=self.pmin,
@@ -333,11 +376,11 @@ def _differential_rate(self, data_tensor, ptensor):
                 y_ref=self.normalisations,
                 )
 
-        # Compute weighted average of diff rates
-        # (n_events,) tensor
-        return tf.reduce_sum(
-            norm * template_diffrates * template_weights[None, :],
-            axis=1)
+        knots_per_event=tf.convert_to_tensor([self.tensor_xvals,self._fetch(self.column, data_tensor)],dtype=fd.float_type())
+        bspline_diff_rates=self.bspline_interpolate_per_bin(ptensor[None, :], tf.transpose(knots_per_event,perm=[1,0,2]))
+        dr=tf.squeeze(norm)*bspline_diff_rates
+
+        return dr
 
     def simulate(self, n_events, fix_truth=None, full_annotate=False,
                  keep_padding=False, **params):
diff --git a/flamedisx/tfbspline/__init__.py b/flamedisx/tfbspline/__init__.py
new file mode 100644
index 000000000..e5c1a7763
--- /dev/null
+++ b/flamedisx/tfbspline/__init__.py
@@ -0,0 +1 @@
+from .bspline import *
diff --git a/flamedisx/tfbspline/bspline.py b/flamedisx/tfbspline/bspline.py
new file mode 100644
index 000000000..6ae91d0e7
--- /dev/null
+++ b/flamedisx/tfbspline/bspline.py
@@ -0,0 +1,237 @@
+import enum
+
+import tensorflow as tf
+
+class Degree(enum.IntEnum):
+  """Defines valid degrees for B-spline interpolation."""
+  CONSTANT = 0
+  LINEAR = 1
+  QUADRATIC = 2
+  CUBIC = 3
+  QUARTIC = 4
+
+
+def _constant(position: tf.Tensor) -> tf.Tensor:
+  """B-Spline basis function of degree 0 for positions in the range [0, 1]."""
+  # A piecewise constant spline is discontinuous at the knots.
+  return tf.expand_dims(tf.clip_by_value(1.0 + position, 1.0, 1.0), axis=-1)
+
+
+def _linear(position: tf.Tensor) -> tf.Tensor:
+  """B-Spline basis functions of degree 1 for positions in the range [0, 1]."""
+  # Piecewise linear splines are C0 smooth.
+  return tf.stack((1.0 - position, position), axis=-1)
+
+
+def _quadratic(position: tf.Tensor) -> tf.Tensor:
+  """B-Spline basis functions of degree 2 for positions in the range [0, 1]."""
+  # We pre-calculate the terms that are used multiple times.
+  pos_sq = tf.pow(position, 2.0)
+
+  # Piecewise quadratic splines are C1 smooth.
+  return tf.stack((tf.pow(1.0 - position, 2.0) / 2.0, -pos_sq + position + 0.5,
+                   pos_sq / 2.0),
+                  axis=-1)
+
+
+def _cubic(position: tf.Tensor) -> tf.Tensor:
+  """B-Spline basis functions of degree 3 for positions in the range [0, 1]."""
+  # We pre-calculate the terms that are used multiple times.
+  neg_pos = 1.0 - position
+  pos_sq = tf.pow(position, 2.0)
+  pos_cb = tf.pow(position, 3.0)
+
+  # Piecewise cubic splines are C2 smooth.
+  return tf.stack(
+      (tf.pow(neg_pos, 3.0) / 6.0, (3.0 * pos_cb - 6.0 * pos_sq + 4.0) / 6.0,
+       (-3.0 * pos_cb + 3.0 * pos_sq + 3.0 * position + 1.0) / 6.0,
+       pos_cb / 6.0),
+      axis=-1)
+
+
+def _quartic(position: tf.Tensor) -> tf.Tensor:
+  """B-Spline basis functions of degree 4 for positions in the range [0, 1]."""
+  # We pre-calculate the terms that are used multiple times.
+  neg_pos = 1.0 - position
+  pos_sq = tf.pow(position, 2.0)
+  pos_cb = tf.pow(position, 3.0)
+  pos_qt = tf.pow(position, 4.0)
+
+  # Piecewise quartic splines are C3 smooth.
+  return tf.stack(
+      (tf.pow(neg_pos, 4.0) / 24.0,
+       (-4.0 * tf.pow(neg_pos, 4.0) + 4.0 * tf.pow(neg_pos, 3.0) +
+        6.0 * tf.pow(neg_pos, 2.0) + 4.0 * neg_pos + 1.0) / 24.0,
+       (pos_qt - 2.0 * pos_cb - pos_sq + 2.0 * position) / 4.0 + 11.0 / 24.0,
+       (-4.0 * pos_qt + 4.0 * pos_cb + 6.0 * pos_sq + 4.0 * position + 1.0) /
+       24.0, pos_qt / 24.0),
+      axis=-1)
+
+
+def knot_weights(
+    positions,
+    num_knots,
+    degree: int,
+    cyclical: bool,
+    sparse_mode: bool = False,
+    name: str = "bspline_knot_weights"
+):
+  """Function that converts cardinal B-spline positions to knot weights.
+
+  Note:
+    In the following, A1 to An are optional batch dimensions.
+
+  Args:
+    positions: A tensor with shape `[A1, .. An]`. Positions must be between
+      `[0, C - D)` for non-cyclical and `[0, C)` for cyclical splines, where `C`
+      is the number of knots and `D` is the spline degree.
+    num_knots: A strictly positive `int` describing the number of knots in the
+      spline.
+    degree: An `int` describing the degree of the spline, which must be smaller
+      than `num_knots`.
+    cyclical: A `bool` describing whether the spline is cyclical.
+    sparse_mode: A `bool` describing whether to return a result only for the
+      knots with nonzero weights. If set to True, the function returns the
+      weights of only the `degree` + 1 knots that are non-zero, as well as the
+      indices of the knots.
+    name: A name for this op. Defaults to "bspline_knot_weights".
+
+  Returns:
+    A tensor with dense weights for each control point, with the shape
+    `[A1, ... An, C]` if `sparse_mode` is False.
+    Otherwise, returns a tensor of shape `[A1, ... An, D + 1]` that contains the
+    non-zero weights, and a tensor with the indices of the knots, with the type
+    tf.int32.
+
+  Raises:
+    ValueError: If degree is greater than 4 or num_knots - 1, or less than 0.
+    InvalidArgumentError: If positions are not in the right range.
+  """
+  with tf.name_scope(name):
+    positions = tf.convert_to_tensor(value=positions)
+
+    all_basis_functions = {
+        # Maps valid degrees to functions.
+        Degree.CONSTANT: _constant,
+        Degree.LINEAR: _linear,
+        Degree.QUADRATIC: _quadratic,
+        Degree.CUBIC: _cubic,
+        Degree.QUARTIC: _quartic
+    }
+    basis_functions = all_basis_functions[degree]
+
+    if not cyclical and num_knots - degree == 1:
+      # In this case all weights are non-zero and we can just return them.
+      if not sparse_mode:
+        return basis_functions(positions)
+      else:
+        shift = tf.zeros_like(positions, dtype=tf.int32)
+        return basis_functions(positions), shift
+
+    # shape_batch = positions.shape.as_list()
+    shape_batch = tf.shape(input=positions)
+    positions = tf.reshape(positions, shape=(-1,))
+
+    # Calculate the nonzero weights from the decimal parts of positions.
+    shift = tf.floor(positions)
+    sparse_weights = basis_functions(positions - shift)
+    shift = tf.cast(shift, tf.int32)
+
+    if sparse_mode:
+      # Returns just the weights and the shift amounts, so that tf.gather_nd on
+      # the knots can be used to sparsely activate knots if needed.
+      shape_weights = tf.concat(
+          (shape_batch, tf.constant((degree + 1,), dtype=tf.int32)), axis=0)
+      sparse_weights = tf.reshape(sparse_weights, shape=shape_weights)
+      shift = tf.reshape(shift, shape=shape_batch)
+      return sparse_weights, shift
+
+    num_positions = tf.size(input=positions)
+    ind_row, ind_col = tf.meshgrid(
+        tf.range(num_positions, dtype=tf.int32),
+        tf.range(degree + 1, dtype=tf.int32),
+        indexing="ij")
+
+    tiled_shifts = tf.reshape(
+        tf.tile(tf.expand_dims(shift, axis=-1), multiples=(1, degree + 1)),
+        shape=(-1,))
+    ind_col = tf.reshape(ind_col, shape=(-1,)) + tiled_shifts
+    if cyclical:
+      ind_col = tf.math.mod(ind_col, num_knots)
+    indices = tf.stack((tf.reshape(ind_row, shape=(-1,)), ind_col), axis=-1)
+    shape_indices = tf.concat((tf.reshape(
+        num_positions, shape=(1,)), tf.constant(
+            (degree + 1, 2), dtype=tf.int32)),
+                              axis=0)
+    indices = tf.reshape(indices, shape=shape_indices)
+    shape_scatter = tf.concat((tf.reshape(
+        num_positions, shape=(1,)), tf.constant((num_knots,), dtype=tf.int32)),
+                              axis=0)
+    weights = tf.scatter_nd(indices, sparse_weights, shape_scatter)
+    shape_weights = tf.concat(
+        (shape_batch, tf.constant((num_knots,), dtype=tf.int32)), axis=0)
+    return tf.reshape(weights, shape=shape_weights)
+
+
+def interpolate_with_weights(
+    knots,
+    weights,
+    name: str = "bspline_interpolate_with_weights") -> tf.Tensor:
+  """Interpolates knots using knot weights.
+
+  Note:
+    In the following, A1 to An, and B1 to Bk are optional batch dimensions.
+
+  Args:
+    knots: A tensor with shape `[B1, ..., Bk, C]` containing knot values, where
+      `C` is the number of knots.
+    weights: A tensor with shape `[A1, ..., An, C]` containing dense weights for
+      the knots, where `C` is the number of knots.
+    name: A name for this op. Defaults to "bspline_interpolate_with_weights".
+
+  Returns:
+    A tensor with shape `[A1, ..., An, B1, ..., Bk]`, which is the result of
+    spline interpolation.
+
+  Raises:
+    ValueError: If the last dimension of knots and weights is not equal.
+  """
+  with tf.name_scope(name):
+    knots = tf.convert_to_tensor(value=knots)
+    weights = tf.convert_to_tensor(value=weights)
+
+  return tf.tensordot(weights, knots, (-1, -1))
+
+
+def interpolate(knots,
+                positions,
+                degree: int,
+                cyclical: bool,
+                name: str = "bspline_interpolate") -> tf.Tensor:
+  """Applies B-spline interpolation to input control points (knots).
+
+  Note:
+    In the following, A1 to An, and B1 to Bk are optional batch dimensions.
+
+  Args:
+    knots: A tensor with shape `[B1, ..., Bk, C]` containing knot values, where
+      `C` is the number of knots.
+    positions: Tensor with shape `[A1, .. An]`. Positions must be between
+      `[0, C - D)` for non-cyclical and `[0, C)` for cyclical splines, where `C`
+      is the number of knots and `D` is the spline degree.
+    degree: An `int` between 0 and 4, or an enumerated constant from the Degree
+      class, which is the degree of the splines.
+    cyclical: A `bool`, whether the splines are cyclical.
+    name: A name for this op. Defaults to "bspline_interpolate".
+
+  Returns:
+    A tensor of shape `[A1, ... An, B1, ..., Bk]`, which is the result of spline
+    interpolation.
+  """
+  with tf.name_scope(name):
+    knots = tf.convert_to_tensor(value=knots)
+    positions = tf.convert_to_tensor(value=positions)
+
+    num_knots = knots.get_shape().as_list()[-1]
+    weights = knot_weights(positions, num_knots, degree, cyclical, False, name)
+    return interpolate_with_weights(knots, weights)

From d03376e7553a6e018c0e2f08c3a74000accf0cfb Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Mon, 3 Jun 2024 12:26:02 +1000
Subject: [PATCH 21/48] Fix for batch size > 1 in cubic spline morphing.

---
 flamedisx/templates.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/flamedisx/templates.py b/flamedisx/templates.py
index f0803eda6..2382d5e73 100644
--- a/flamedisx/templates.py
+++ b/flamedisx/templates.py
@@ -341,7 +341,7 @@ def _annotate(self):
             linear_interp_padded_diff_rates.append(np.concatenate([left_edge,diff_rate_per_hist,right_edge]))
 
         self.data[self.column]=linear_interp_padded_diff_rates
-        self.tensor_xvals=tf.convert_to_tensor([self.pvals for d in self.data[self.column]],dtype=fd.float_type())
+        self.tensor_xvals=tf.convert_to_tensor([self.pvals for _ in range(self.batch_size)],dtype=fd.float_type())
 
     def mu_before_efficiencies(self, **params):
         return self.mu

From 06761168a0c16acd85fc3a337d111e1a04d5517c Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Mon, 3 Jun 2024 22:03:39 +1000
Subject: [PATCH 22/48] Undo temporary change from before.

---
 flamedisx/non_asymptotic_inference.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 1889f0143..1be0e8844 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -687,7 +687,7 @@ def get_bands(self, conf_level=0.1, quantiles=[0, 1, -1, 2, -2],
                                                              kind='weak') / 100.
                     these_p_vals = these_p_vals / (1. - these_p_vals_b + 1e-10)
                 mus.append(mu_test)
-                p_val_curves.append(these_p_vals[:2130])
+                p_val_curves.append(these_p_vals)
 
             p_val_curves = np.transpose(np.stack(p_val_curves, axis=0))
             upper_lims_bands = np.apply_along_axis(self.upper_lims_bands, 1, p_val_curves, mus, conf_level)
@@ -738,7 +738,7 @@ def get_median_disco_asymptotic(self, sigma_level=3):
                 these_disco_sigs = np.sqrt(ts_values)
 
                 mus.append(mu_test)
-                disco_sig_curves.append(these_disco_sigs[:2130])
+                disco_sig_curves.append(these_disco_sigs)
 
             disco_sig_curves = np.stack(disco_sig_curves, axis=0)
             median_disco_sigs = [np.median(disco_sigs) for disco_sigs in disco_sig_curves]

From 793d26adeab027b1f60c9139b477bd1754a78b78 Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Thu, 6 Jun 2024 00:33:44 +1000
Subject: [PATCH 23/48] Fix to not setting non-rate NPs to conditional BF
 values in toys.

---
 flamedisx/non_asymptotic_inference.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 1be0e8844..78b6f3f52 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -347,8 +347,18 @@ 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())
 
+        conditional_bfs_observed = self.observed_test_stats[signal_source_name].conditional_best_fits[mu_test]
+        non_rate_params_added = []
+        for pname, fitval in conditional_bfs_observed.items():
+            if (pname not in simulate_dict) and (pname in likelihood.param_defaults):
+                simulate_dict[pname] = fitval
+                non_rate_params_added.append(pname)
+
         toy_data = likelihood.simulate(**simulate_dict)
 
+        for pname in non_rate_params_added:
+            simulate_dict.pop(pname)
+
         return simulate_dict, toy_data, constraint_extra_args
 
     def toy_test_statistic_dist(self, test_stat_dists_SB, test_stat_dists_B,

From 64896d7dcc7763399e0ed98b30f5b1db022d612f Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Sun, 9 Jun 2024 00:21:22 -0700
Subject: [PATCH 24/48] Fix for disco sig < 0.

---
 flamedisx/non_asymptotic_inference.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 78b6f3f52..8e618473d 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -727,6 +727,7 @@ def get_disco_sig(self):
                                                     observed_test_stat,
                                                     kind='weak')) / 100.
             disco_sig = stats.norm.ppf(1. - p_val)
+            disco_sig = np.where(disco_sig > 0., disco_sig, 0.)
             disco_sigs[signal_source] = disco_sig
 
         return disco_sigs

From eeb68cdfdda27509d649376ef8e4b22a5031bbe3 Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Mon, 17 Jun 2024 14:58:24 +1000
Subject: [PATCH 25/48] Fix to earlier commit (non-rate NPs) not working when
 in sensitivity mode.

---
 flamedisx/non_asymptotic_inference.py | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 8e618473d..59ea8dda1 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -347,17 +347,19 @@ 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())
 
-        conditional_bfs_observed = self.observed_test_stats[signal_source_name].conditional_best_fits[mu_test]
-        non_rate_params_added = []
-        for pname, fitval in conditional_bfs_observed.items():
-            if (pname not in simulate_dict) and (pname in likelihood.param_defaults):
-                simulate_dict[pname] = fitval
-                non_rate_params_added.append(pname)
+        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 = []
+            for pname, fitval in conditional_bfs_observed.items():
+                if (pname not in simulate_dict) and (pname in likelihood.param_defaults):
+                    simulate_dict[pname] = fitval
+                    non_rate_params_added.append(pname)
 
         toy_data = likelihood.simulate(**simulate_dict)
 
-        for pname in non_rate_params_added:
-            simulate_dict.pop(pname)
+        if self.observed_test_stats is not None:
+            for pname in non_rate_params_added:
+                simulate_dict.pop(pname)
 
         return simulate_dict, toy_data, constraint_extra_args
 

From 94d66ab7314006b8e963f2b79409c48e0caaeafe Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Mon, 17 Jun 2024 22:35:03 -0700
Subject: [PATCH 26/48] Update band calculation routine: handle failed toys.

---
 flamedisx/non_asymptotic_inference.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 59ea8dda1..c81bbccd8 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -665,9 +665,12 @@ def get_interval(self, conf_level=0.1, pcl_level=0.16,
             return lower_lim_all, upper_lim_all, p_sb, p_b
 
     def upper_lims_bands(self, pval_curve, mus, conf_level):
-        upper_lims = np.argwhere(np.diff(np.sign(pval_curve - np.ones_like(pval_curve) * conf_level)) < 0.).flatten()
-        return self.interp_helper(mus, pval_curve, upper_lims, conf_level,
-                                  rising_edge=False, inverse=True)
+        try:
+            upper_lims = np.argwhere(np.diff(np.sign(pval_curve - np.ones_like(pval_curve) * conf_level)) < 0.).flatten()
+            return self.interp_helper(mus, pval_curve, upper_lims, conf_level,
+                                    rising_edge=False, inverse=True)
+        except Exception:
+            return 0.
 
     def critical_disco_value(self, disco_pot_curve, mus, discovery_sigma):
         crossing_point = np.argwhere(np.diff(np.sign(disco_pot_curve - np.ones_like(disco_pot_curve) * discovery_sigma)) > 0.).flatten()
@@ -704,6 +707,10 @@ def get_bands(self, conf_level=0.1, quantiles=[0, 1, -1, 2, -2],
             p_val_curves = np.transpose(np.stack(p_val_curves, axis=0))
             upper_lims_bands = np.apply_along_axis(self.upper_lims_bands, 1, p_val_curves, mus, conf_level)
 
+            if len(upper_lims_bands[upper_lims_bands == 0.]) > 0.:
+                print(f'Found {len(upper_lims_bands[upper_lims_bands == 0.])} failed toy for {signal_source}; removing...')
+                upper_lims_bands = upper_lims_bands[upper_lims_bands > 0.]
+
             these_bands = dict()
             for quantile in quantiles:
                 these_bands[quantile] = np.quantile(np.sort(upper_lims_bands), stats.norm.cdf(quantile))

From 9d35ec61981798e58d4084e538171f327af7e884 Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Mon, 22 Jul 2024 06:40:13 -0700
Subject: [PATCH 27/48] Return UL without PCL, also.

---
 flamedisx/non_asymptotic_inference.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index c81bbccd8..6b2b9703f 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -608,6 +608,7 @@ def get_interval(self, conf_level=0.1, pcl_level=0.16,
 
         lower_lim_all = dict()
         upper_lim_all = dict()
+        upper_lim_all_raw = dict()
         # Loop over signal sources
         for signal_source in self.signal_source_names:
             if not asymptotic:
@@ -643,6 +644,7 @@ def get_interval(self, conf_level=0.1, pcl_level=0.16,
             # Take the highest decreasing crossing point, and interpolate to get an upper limit
             upper_lim = self.interp_helper(mus, p_vals, upper_lims, conf_level,
                                            rising_edge=False, inverse=True)
+            upper_lim_raw = upper_lim
 
             if use_CLs is False and not asymptotic:
                 M0 = self.interp_helper(mus, pws, upper_lims, upper_lim,
@@ -656,11 +658,12 @@ def get_interval(self, conf_level=0.1, pcl_level=0.16,
 
             lower_lim_all[signal_source] = lower_lim
             upper_lim_all[signal_source] = upper_lim
+            upper_lim_all_raw[signal_source] = upper_lim_raw
 
         if asymptotic:
             return lower_lim_all, upper_lim_all
         if use_CLs is False:
-            return lower_lim_all, upper_lim_all, p_sb, powers
+            return lower_lim_all, upper_lim_all, upper_lim_all_raw, p_sb, powers
         else:
             return lower_lim_all, upper_lim_all, p_sb, p_b
 

From 72df5b79f2612076bd7f5ed546bcacc4652055be Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Mon, 21 Oct 2024 00:14:59 -0700
Subject: [PATCH 28/48] Fix long-standing bug with indexing constraint centers
 in toys.

---
 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 6b2b9703f..5fcd8be95 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -420,7 +420,7 @@ def toy_test_statistic_dist(self, test_stat_dists_SB, test_stat_dists_B,
                     if value < 0.1:
                         guess_dict_B[key] = 0.1
                 toy_data_B = self.toy_data_B[toy+(self.toy_batch*self.ntoys)]
-                constraint_extra_args_B = self.constraint_extra_args_B[toy]
+                constraint_extra_args_B = self.constraint_extra_args_B[toy+(self.toy_batch*self.ntoys)]
             except Exception:
                 raise RuntimeError("Could not find background-only datasets")
 

From 0ed3bb16554059cce651d03881b0eeb8de715559 Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Mon, 17 Mar 2025 12:28:38 +1100
Subject: [PATCH 29/48] Bug fix with template morphing: tracking of mu
 variation.

---
 flamedisx/templates.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/flamedisx/templates.py b/flamedisx/templates.py
index 2382d5e73..3aab80ea2 100644
--- a/flamedisx/templates.py
+++ b/flamedisx/templates.py
@@ -216,6 +216,8 @@ def __init__(
             TemplateWrapper(
                 template, bin_edges, axis_names, events_per_bin, interpolate)
             for _, template in params_and_templates]
+        assert len(params_and_templates[0][0]) == 1, "This implementation currently only supports moprhing of 1 parameter"
+        self.param_name = list(params_and_templates[0][0].keys())[0]
 
         # We will include mu variation separately
         self.mu = self._templates[0].mu

From 4017cc918354fc0b71d0abd5ef3ca8f2724950fd Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Wed, 7 May 2025 04:25:18 -0700
Subject: [PATCH 30/48] allow_failure on by default for fits.

---
 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 5fcd8be95..83d18c065 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -31,9 +31,11 @@ def __call__(self, mu_test, signal_source_name, guess_dict,
         guess_dict_nuisance.pop(f'{signal_source_name}_rate_multiplier')
 
         # Conditional fit
-        bf_conditional = self.likelihood.bestfit(fix=fix_dict, guess=guess_dict_nuisance, suppress_warnings=True)
+        bf_conditional = self.likelihood.bestfit(fix=fix_dict, guess=guess_dict_nuisance, suppress_warnings=True,
+                                                 allow_failure=True)
         # Uncnditional fit
-        bf_unconditional = self.likelihood.bestfit(guess=guess_dict, suppress_warnings=True)
+        bf_unconditional = self.likelihood.bestfit(guess=guess_dict, suppress_warnings=True,
+                                                   allow_failure=True)
 
         # Return the test statistic, unconditional fit and conditional fit
         if not asymptotic:

From 27b5e9aecdd2026a9ad91e68c3439bbaffc201cb Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Wed, 7 May 2025 20:57:19 -0700
Subject: [PATCH 31/48] Handle case where toys end up with empty dataset(s):
 skip that toy.

---
 flamedisx/non_asymptotic_inference.py | 78 +++++++++++++++++----------
 1 file changed, 49 insertions(+), 29 deletions(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 83d18c065..ea9169687 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -390,27 +390,37 @@ def toy_test_statistic_dist(self, test_stat_dists_SB, test_stat_dists_B,
             # Shift the constraint in the likelihood based on the background RMs we drew
             likelihood.set_constraint_extra_args(**constraint_extra_args_SB)
             # Set data
+            empty_dataframe = False
             if hasattr(likelihood, 'likelihoods'):
                 for component, data in toy_data_SB.items():
-                    likelihood.set_data(data, component)
+                    if len(data) == 0:
+                        empty_dataframe = True
             else:
-                likelihood.set_data(toy_data_SB)
-            # Create test statistic
-            test_statistic_SB = self.test_statistic(likelihood)
-            # Guesses for fit
-            guess_dict_SB = simulate_dict_SB.copy()
-            for key, value in guess_dict_SB.items():
-                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)
-            # Save test statistics, and possibly fits
-            ts_values_SB.append(ts_result_SB[0])
-            ts_values_SB_disco.append(ts_result_SB_disco[0])
-            if save_fits:
-                unconditional_bfs_SB.append(ts_result_SB[1])
-                conditional_bfs_SB.append(ts_result_SB[2])
+                if len(toy_data_SB) == 0:
+                    empty_dataframe = True
+
+            if not empty_dataframe:
+                if hasattr(likelihood, 'likelihoods'):
+                    for component, data in toy_data_SB.items():
+                        likelihood.set_data(data, component)
+                else:
+                    likelihood.set_data(toy_data_SB)
+                # Create test statistic
+                test_statistic_SB = self.test_statistic(likelihood)
+                # Guesses for fit
+                guess_dict_SB = simulate_dict_SB.copy()
+                for key, value in guess_dict_SB.items():
+                    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)
+                # Save test statistics, and possibly fits
+                ts_values_SB.append(ts_result_SB[0])
+                ts_values_SB_disco.append(ts_result_SB_disco[0])
+                if save_fits:
+                    unconditional_bfs_SB.append(ts_result_SB[1])
+                    conditional_bfs_SB.append(ts_result_SB[2])
 
             # B-only toys
 
@@ -429,20 +439,30 @@ def toy_test_statistic_dist(self, test_stat_dists_SB, test_stat_dists_B,
             # Shift the constraint in the likelihood based on the background RMs we drew
             likelihood.set_constraint_extra_args(**constraint_extra_args_B)
             # Set data
+            empty_dataframe = False
             if hasattr(likelihood, 'likelihoods'):
                 for component, data in toy_data_B.items():
-                    likelihood.set_data(data, component)
+                    if len(data) == 0:
+                        empty_dataframe = True
             else:
-                likelihood.set_data(toy_data_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)
-            # Save test statistic, and possibly fits
-            ts_values_B.append(ts_result_B[0])
-            if save_fits:
-                unconditional_bfs_B.append(ts_result_SB[1])
-                conditional_bfs_B.append(ts_result_SB[2])
+                if len(toy_data_B) == 0:
+                    empty_dataframe = True
+
+            if not empty_dataframe:
+                if hasattr(likelihood, 'likelihoods'):
+                    for component, data in toy_data_B.items():
+                        likelihood.set_data(data, component)
+                else:
+                    likelihood.set_data(toy_data_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)
+                # Save test statistic, and possibly fits
+                ts_values_B.append(ts_result_B[0])
+                if save_fits:
+                    unconditional_bfs_B.append(ts_result_SB[1])
+                    conditional_bfs_B.append(ts_result_SB[2])
 
         # Add to the test statistic distributions
         test_stat_dists_SB.add_ts_dist(mu_test, ts_values_SB)

From 583e82170bb7d4b8d1a3879a073fd796f21c67a3 Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Thu, 29 May 2025 22:25:26 -0700
Subject: [PATCH 32/48] Fix to anchor padding for B-spline interpolation
 (template morphing): no longer asymmetric.

---
 flamedisx/templates.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/flamedisx/templates.py b/flamedisx/templates.py
index 3aab80ea2..cca0bdbaf 100644
--- a/flamedisx/templates.py
+++ b/flamedisx/templates.py
@@ -282,9 +282,11 @@ def __init__(
                                                    fd.float_type())
 
         # Assume equi-spacing!
-        self.dstep=self.pvals[1]-self.pvals[0]
+        self.dstep = self.pvals[1] - self.pvals[0]
         # Need to pad domain.. four might be excessive
-        self.pvals=list(np.arange(self.pvals[0]-4*(self.dstep),self.pvals[-1]+5*(self.dstep),self.dstep))
+        old_length = len(self.pvals)
+        self.pvals=list(np.arange(self.pvals[0] - 4. * self.dstep, self.pvals[-1] + 4. * self.dstep, self.dstep))
+        assert len(self.pvals) == (old_length + 8), "Something went wrong with the padding!"
         self.array_columns = ((self.column, n_templates+8),)
 
         super().__init__(*args, **kwargs)

From fd8d32610fa2263952eddd1105782d35f69bf076 Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Sun, 1 Jun 2025 20:01:15 -0700
Subject: [PATCH 33/48] Whoops, that last commit wasn't catching every case.

---
 flamedisx/templates.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/flamedisx/templates.py b/flamedisx/templates.py
index cca0bdbaf..fed44db28 100644
--- a/flamedisx/templates.py
+++ b/flamedisx/templates.py
@@ -275,18 +275,22 @@ def __init__(
         param_vals = np.asarray([list(params.values())[0] for params, _ in params_and_templates])
         self.pmin = tf.constant(min(param_vals), fd.float_type())
         self.pmax = tf.constant(max(param_vals), fd.float_type())
-        self.pvals = tf.convert_to_tensor(param_vals, fd.float_type())
+        pvals = tf.convert_to_tensor(param_vals, fd.float_type())
 
         normalisations = np.array([norm for _, norm in params_and_normalisations])
         self.normalisations = tf.convert_to_tensor(normalisations / normalisations[0],
                                                    fd.float_type())
 
         # Assume equi-spacing!
-        self.dstep = self.pvals[1] - self.pvals[0]
+        self.dstep = pvals[1] - pvals[0]
         # Need to pad domain.. four might be excessive
-        old_length = len(self.pvals)
-        self.pvals=list(np.arange(self.pvals[0] - 4. * self.dstep, self.pvals[-1] + 4. * self.dstep, self.dstep))
-        assert len(self.pvals) == (old_length + 8), "Something went wrong with the padding!"
+        try:
+            self.pvals = list(np.arange(pvals[0] - 4. * self.dstep, pvals[-1] + 4. * self.dstep, self.dstep))
+            assert len(self.pvals) == len(pvals) + 8, "Something went wrong with the padding!"
+        except:
+            self.pvals = list(np.arange(pvals[0] - 4. * self.dstep, pvals[-1] + 5. * self.dstep, self.dstep))
+            assert len(self.pvals) == len(pvals) + 8, "Something went wrong with the padding!"
+
         self.array_columns = ((self.column, n_templates+8),)
 
         super().__init__(*args, **kwargs)

From 044cb167385bae5016c5fef184342d7299b7ad68 Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Tue, 8 Jul 2025 09:53:35 -0700
Subject: [PATCH 34/48] Generalise morphing when the ptensor has other
 elements.

---
 flamedisx/templates.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/flamedisx/templates.py b/flamedisx/templates.py
index fed44db28..6aa15b324 100644
--- a/flamedisx/templates.py
+++ b/flamedisx/templates.py
@@ -378,14 +378,14 @@ def interp(knots_for_event):
 
     def _differential_rate(self, data_tensor, ptensor):
         norm = tfp.math.batch_interp_regular_1d_grid(
-                x=ptensor[None, :],
+                x=self._fetch_param(self.param_name, ptensor),
                 x_ref_min=self.pmin,
                 x_ref_max=self.pmax,
                 y_ref=self.normalisations,
                 )
-
-        knots_per_event=tf.convert_to_tensor([self.tensor_xvals,self._fetch(self.column, data_tensor)],dtype=fd.float_type())
-        bspline_diff_rates=self.bspline_interpolate_per_bin(ptensor[None, :], tf.transpose(knots_per_event,perm=[1,0,2]))
+        
+        knots_per_event=tf.convert_to_tensor([self.tensor_xvals, self._fetch(self.column, data_tensor)],dtype=fd.float_type())
+        bspline_diff_rates=self.bspline_interpolate_per_bin(self._fetch_param(self.param_name, ptensor), tf.transpose(knots_per_event,perm=[1,0,2]))
         dr=tf.squeeze(norm)*bspline_diff_rates
 
         return dr

From c92b18e7102b14e9c32f0429496760eb5574af67 Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Fri, 11 Jul 2025 15:18:18 +0100
Subject: [PATCH 35/48] Adding in sampling of non-rate parameter constraint
 centers (Wei).

---
 flamedisx/non_asymptotic_inference.py | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index ea9169687..0c9fa3f09 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -323,6 +323,8 @@ def sample_data_constraints(self, mu_test, signal_source_name, likelihood):
         """
         simulate_dict = dict()
         constraint_extra_args = dict()
+
+        # For rate multipliers
         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:
@@ -349,6 +351,29 @@ 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 all other parameters
+        for param_name in likelihood.param_names:
+            if '_rate_multiplier' in param_name:
+                continue
+
+            if param_name not in self.sample_other_constraints.keys():
+                continue
+
+            # Case where we use the conditional best fits as constraint centers and simulated values
+            if self.observed_test_stats is not None:
+                try:
+                    conditional_bfs_observed = self.observed_test_stats[signal_source_name].conditional_best_fits
+                    param_val_expected = conditional_bfs_observed[mu_test][param_name]
+                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:
+                param_val_expected = likelihood.param_defaults[param_name]
+
+            # Sample constraint centers
+            draw = self.sample_other_constraints[param_name](param_val_expected)
+            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 d5ba07485b9fcb6ad0f0ddf0e4435623bde1ee03 Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Sun, 13 Jul 2025 17:52:25 +0100
Subject: [PATCH 36/48] Non-asymptotic inference updates: better handling of
 empty dataset; remove asymptotic options (was messy, and made no sense being
 here); tidy up inference routines, with improved control over what is
 calcualted for different modes (see LZFlameFit).

---
 flamedisx/likelihood.py               |   2 +
 flamedisx/non_asymptotic_inference.py | 199 ++++++++++----------------
 2 files changed, 79 insertions(+), 122 deletions(-)

diff --git a/flamedisx/likelihood.py b/flamedisx/likelihood.py
index 97862217e..d8d8a0717 100644
--- a/flamedisx/likelihood.py
+++ b/flamedisx/likelihood.py
@@ -271,6 +271,8 @@ def set_data(self,
                               UserWarning)
             for s in self.sources.values():
                 s.set_data(None)
+                s.n_batches = 0
+                self.batch_info = None
                 return
 
         batch_info = np.zeros((len(self.dsetnames), 3), dtype=int)
diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 0c9fa3f09..2336788f7 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -22,8 +22,7 @@ class TestStatistic():
     def __init__(self, likelihood):
         self.likelihood = likelihood
 
-    def __call__(self, mu_test, signal_source_name, guess_dict,
-                 asymptotic=False):
+    def __call__(self, mu_test, signal_source_name, guess_dict):
         # To fix the signal RM in the conditional fit
         fix_dict = {f'{signal_source_name}_rate_multiplier': tf.cast(mu_test, fd.float_type())}
 
@@ -38,11 +37,7 @@ def __call__(self, mu_test, signal_source_name, guess_dict,
                                                    allow_failure=True)
 
         # Return the test statistic, unconditional fit and conditional fit
-        if not asymptotic:
-            return self.evaluate(bf_unconditional, bf_conditional), bf_unconditional, bf_conditional
-        else:
-            return self.evaluate_asymptotic_pval(bf_unconditional, bf_conditional,
-                                                 mu_test), bf_unconditional, bf_conditional
+        return self.evaluate(bf_unconditional, bf_conditional), bf_unconditional, bf_conditional
 
 
 @export
@@ -62,23 +57,6 @@ def evaluate(self, bf_unconditional, bf_conditional):
         else:
             return ts
 
-    def evaluate_asymptotic_pval(self, bf_unconditional, bf_conditional, mu_test):
-        ll_conditional = self.likelihood(**bf_conditional)
-        ll_unconditional = self.likelihood(**bf_unconditional)
-
-        ts = -2. * (ll_conditional - ll_unconditional)
-
-        cov = 2. * self.likelihood.inverse_hessian(bf_unconditional)
-        sigma_mu = np.sqrt(cov[0][0])
-
-        if ts < (mu_test**2 / sigma_mu**2):
-            F = 2. * stats.norm.cdf(np.sqrt(ts)) - 1.
-        else:
-            F = stats.norm.cdf(np.sqrt(ts)) + stats.norm.cdf((ts + (mu_test**2 / sigma_mu**2)) / (2. * mu_test / sigma_mu)) - 1.
-
-        pval = 1. - F
-        return pval
-
 
 @export
 class TestStatisticDistributions():
@@ -206,7 +184,7 @@ 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):
+                    SB_toys=False, B_toys=False, discovery_TS=False):
         """If observed_data is passed, evaluate observed test statistics. Otherwise,
         obtain test statistic distributions (for both S+B and B-only).
 
@@ -251,16 +229,20 @@ def run_routine(self, mus_test=None, save_fits=False,
             self.toy_batch = toy_batch
 
         observed_test_stats_collection = dict()
+
         test_stat_dists_SB_collection = dict()
         test_stat_dists_SB_disco_collection = dict()
         test_stat_dists_B_collection = dict()
+        test_stat_dists_B_disco_collection = dict()
 
         # Loop over signal sources
         for signal_source in self.signal_source_names:
             observed_test_stats = ObservedTestStatistics()
+
             test_stat_dists_SB = TestStatisticDistributions()
             test_stat_dists_SB_disco = TestStatisticDistributions()
             test_stat_dists_B = TestStatisticDistributions()
+            test_stat_dists_B_disco = TestStatisticDistributions()
 
             # Get likelihood
             likelihood = deepcopy(self.likelihood)
@@ -294,14 +276,14 @@ def run_routine(self, mus_test=None, save_fits=False,
                 # Case where we want observed test statistics
                 if observed_data is not None:
                     self.get_observed_test_stat(observed_test_stats, observed_data,
-                                                mu_test, signal_source, likelihood, save_fits=save_fits,
-                                                asymptotic=asymptotic)
+                                                mu_test, signal_source, likelihood, save_fits=save_fits)
                 # Case where we want test statistic distributions
                 else:
-                    self.toy_test_statistic_dist(test_stat_dists_SB, test_stat_dists_B,
-                                                 test_stat_dists_SB_disco,
+                    self.toy_test_statistic_dist(test_stat_dists_SB, test_stat_dists_SB_disco,
+                                                 test_stat_dists_B, test_stat_dists_B_disco,
                                                  mu_test, signal_source, likelihood,
-                                                 save_fits=save_fits)
+                                                 save_fits=save_fits,
+                                                 SB_toys=SB_toys, B_toys=B_toys, discovery_TS=discovery_TS)
 
             if observed_data is not None:
                 observed_test_stats_collection[signal_source] = observed_test_stats
@@ -309,12 +291,13 @@ def run_routine(self, mus_test=None, save_fits=False,
                 test_stat_dists_SB_collection[signal_source] = test_stat_dists_SB
                 test_stat_dists_SB_disco_collection[signal_source] = test_stat_dists_SB_disco
                 test_stat_dists_B_collection[signal_source] = test_stat_dists_B
+                test_stat_dists_B_disco_collection[signal_source] = test_stat_dists_B_disco
 
         if observed_data is not None:
             return observed_test_stats_collection
         else:
             return test_stat_dists_SB_collection, test_stat_dists_SB_disco_collection, \
-                test_stat_dists_B_collection
+                test_stat_dists_B_collection, test_stat_dists_B_disco_collection
 
     def sample_data_constraints(self, mu_test, signal_source_name, likelihood):
         """Internal function to sample the toy data and constraint central values
@@ -334,7 +317,7 @@ def sample_data_constraints(self, mu_test, signal_source_name, likelihood):
                         conditional_bfs_observed[mu_test][f'{background_source}_rate_multiplier']
                 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
+            # Case where we use the prior expected counts as constraint centers and simulated values
             else:
                 expected_background_counts = self.expected_background_counts[background_source]
 
@@ -366,7 +349,7 @@ def sample_data_constraints(self, mu_test, signal_source_name, likelihood):
                     param_val_expected = conditional_bfs_observed[mu_test][param_name]
                 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
+            # Case where we use the prior expected counts as constraint centers and simulated values
             else:
                 param_val_expected = likelihood.param_defaults[param_name]
 
@@ -390,15 +373,19 @@ 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,
-                                test_stat_dists_SB_disco,
+    def toy_test_statistic_dist(self,
+                                test_stat_dists_SB, test_stat_dists_SB_disco,
+                                test_stat_dists_B, test_stat_dists_B_disco,
                                 mu_test, signal_source_name, likelihood,
-                                save_fits=False):
+                                save_fits=False,
+                                SB_toys=False, B_toys=False, discovery_TS=False):
         """Internal function to get test statistic distribution.
         """
         ts_values_SB = []
         ts_values_SB_disco = []
         ts_values_B = []
+        ts_values_B_disco = []
+
         if save_fits:
             unconditional_bfs_SB = []
             conditional_bfs_SB = []
@@ -407,24 +394,14 @@ 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'):
-            simulate_dict_SB, toy_data_SB, constraint_extra_args_SB = \
-                self.sample_data_constraints(mu_test, signal_source_name, likelihood)
-
             # S+B toys
+            if SB_toys:
+                simulate_dict_SB, toy_data_SB, constraint_extra_args_SB = \
+                    self.sample_data_constraints(mu_test, signal_source_name, likelihood)
 
-            # Shift the constraint in the likelihood based on the background RMs we drew
-            likelihood.set_constraint_extra_args(**constraint_extra_args_SB)
-            # Set data
-            empty_dataframe = False
-            if hasattr(likelihood, 'likelihoods'):
-                for component, data in toy_data_SB.items():
-                    if len(data) == 0:
-                        empty_dataframe = True
-            else:
-                if len(toy_data_SB) == 0:
-                    empty_dataframe = True
-
-            if not empty_dataframe:
+                # Shift the constraint in the likelihood based on the background RMs we drew
+                likelihood.set_constraint_extra_args(**constraint_extra_args_SB)
+                # Set data
                 if hasattr(likelihood, 'likelihoods'):
                     for component, data in toy_data_SB.items():
                         likelihood.set_data(data, component)
@@ -437,43 +414,35 @@ def toy_test_statistic_dist(self, test_stat_dists_SB, test_stat_dists_B,
                 for key, value in guess_dict_SB.items():
                     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)
-                # Save test statistics, and possibly fits
-                ts_values_SB.append(ts_result_SB[0])
-                ts_values_SB_disco.append(ts_result_SB_disco[0])
+                # Evaluate and save test statistics
+                if discovery_TS:
+                    ts_result_SB_disco = test_statistic_SB(0., signal_source_name, guess_dict_SB)
+                    ts_values_SB_disco.append(ts_result_SB_disco[0])
+                else:
+                    ts_result_SB = test_statistic_SB(mu_test, signal_source_name, guess_dict_SB)
+                    ts_values_SB.append(ts_result_SB[0])
+                # Possibly save fits
                 if save_fits:
                     unconditional_bfs_SB.append(ts_result_SB[1])
                     conditional_bfs_SB.append(ts_result_SB[2])
 
             # B-only toys
+            if B_toys:
+                try:
+                    # Guesses for fit
+                    guess_dict_B = self.simulate_dict_B.copy()
+                    guess_dict_B[f'{signal_source_name}_rate_multiplier'] = 0.
+                    for key, value in guess_dict_B.items():
+                        if value < 0.1:
+                            guess_dict_B[key] = 0.1
+                    toy_data_B = self.toy_data_B[toy+(self.toy_batch*self.ntoys)]
+                    constraint_extra_args_B = self.constraint_extra_args_B[toy+(self.toy_batch*self.ntoys)]
+                except Exception:
+                    raise RuntimeError("Could not find background-only datasets")
 
-            try:
-                # Guesses for fit
-                guess_dict_B = self.simulate_dict_B.copy()
-                guess_dict_B[f'{signal_source_name}_rate_multiplier'] = 0.
-                for key, value in guess_dict_B.items():
-                    if value < 0.1:
-                        guess_dict_B[key] = 0.1
-                toy_data_B = self.toy_data_B[toy+(self.toy_batch*self.ntoys)]
-                constraint_extra_args_B = self.constraint_extra_args_B[toy+(self.toy_batch*self.ntoys)]
-            except Exception:
-                raise RuntimeError("Could not find background-only datasets")
-
-            # Shift the constraint in the likelihood based on the background RMs we drew
-            likelihood.set_constraint_extra_args(**constraint_extra_args_B)
-            # Set data
-            empty_dataframe = False
-            if hasattr(likelihood, 'likelihoods'):
-                for component, data in toy_data_B.items():
-                    if len(data) == 0:
-                        empty_dataframe = True
-            else:
-                if len(toy_data_B) == 0:
-                    empty_dataframe = True
-
-            if not empty_dataframe:
+                # Shift the constraint in the likelihood based on the background RMs we drew
+                likelihood.set_constraint_extra_args(**constraint_extra_args_B)
+                # Set data
                 if hasattr(likelihood, 'likelihoods'):
                     for component, data in toy_data_B.items():
                         likelihood.set_data(data, component)
@@ -481,10 +450,14 @@ def toy_test_statistic_dist(self, test_stat_dists_SB, test_stat_dists_B,
                     likelihood.set_data(toy_data_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)
-                # Save test statistic, and possibly fits
-                ts_values_B.append(ts_result_B[0])
+                # Evaluate and save test statistics
+                if discovery_TS:
+                    ts_result_B_disco = test_statistic_B(0., signal_source_name, guess_dict_B)
+                    ts_values_B_disco.append(ts_result_B_disco[0])
+                else:
+                    ts_result_B = test_statistic_B(mu_test, signal_source_name, guess_dict_B)
+                    ts_values_B.append(ts_result_B[0])
+                # Possibly save fits
                 if save_fits:
                     unconditional_bfs_B.append(ts_result_SB[1])
                     conditional_bfs_B.append(ts_result_SB[2])
@@ -493,6 +466,7 @@ def toy_test_statistic_dist(self, test_stat_dists_SB, test_stat_dists_B,
         test_stat_dists_SB.add_ts_dist(mu_test, ts_values_SB)
         test_stat_dists_SB_disco.add_ts_dist(mu_test, ts_values_SB_disco)
         test_stat_dists_B.add_ts_dist(mu_test, ts_values_B)
+        test_stat_dists_B_disco.add_ts_dist(mu_test, ts_values_B_disco)
 
         # Possibly save the fits
         if save_fits:
@@ -502,8 +476,7 @@ def toy_test_statistic_dist(self, test_stat_dists_SB, test_stat_dists_B,
             test_stat_dists_B.add_conditional_best_fit(mu_test, conditional_bfs_B)
 
     def get_observed_test_stat(self, observed_test_stats, observed_data,
-                               mu_test, signal_source_name, likelihood, save_fits=False,
-                               asymptotic=False):
+                               mu_test, signal_source_name, likelihood, save_fits=False):
         """Internal function to evaluate observed test statistic.
         """
         # The constraints are centered on the expected values
@@ -531,8 +504,7 @@ def get_observed_test_stat(self, observed_test_stats, observed_data,
             if value < 0.1:
                 guess_dict[key] = tf.cast(0.1, fd.float_type())
         # Evaluate test statistic
-        ts_result = test_statistic(mu_test, signal_source_name, guess_dict,
-                                   asymptotic=asymptotic)
+        ts_result = test_statistic(mu_test, signal_source_name, guess_dict)
 
         # Add to the test statistic collection
         observed_test_stats.add_test_stat(mu_test, ts_result[0])
@@ -595,7 +567,7 @@ def interp_helper(x, y, crossing_points, crit_val,
         else:
             return (crit_val - x_left) * gradient + y_left
 
-    def get_p_vals(self, conf_level, use_CLs=False, asymptotic=False):
+    def get_p_vals(self, conf_level, use_CLs=False):
         """Internal function to get p-value curves.
         """
         p_sb_collection = dict()
@@ -603,10 +575,6 @@ def get_p_vals(self, conf_level, use_CLs=False, asymptotic=False):
         p_b_collection = dict()
         # Loop over signal sources
         for signal_source in self.signal_source_names:
-            if asymptotic:
-                p_sb_collection[signal_source] = self.observed_test_stats[signal_source]
-                continue
-
             # Get test statistic distribitions and observed test statistics
             test_stat_dists_SB = self.test_stat_dists_SB[signal_source]
             test_stat_dists_B = self.test_stat_dists_B[signal_source]
@@ -623,17 +591,13 @@ def get_p_vals(self, conf_level, use_CLs=False, asymptotic=False):
                 powers = test_stat_dists_B.get_p_vals(crit_vals)
                 powers_collection[signal_source] = powers
 
-        if asymptotic:
-            return p_sb_collection
-
         if use_CLs:
             return p_sb_collection, p_b_collection
         else:
             return p_sb_collection, powers_collection
 
     def get_interval(self, conf_level=0.1, pcl_level=0.16,
-                     use_CLs=False,
-                     asymptotic=False):
+                     use_CLs=False):
         """Get frequentist confidence interval.
 
         Arguments:
@@ -645,34 +609,27 @@ def get_interval(self, conf_level=0.1, pcl_level=0.16,
                 (https://inspirehep.net/literature/599622), and the final return value
                 will be the p-value curves under H1
         """
-        if not asymptotic:
-            if use_CLs:
-                p_sb, p_b = self.get_p_vals(conf_level, use_CLs=True)
-            else:
-                p_sb, powers = self.get_p_vals(conf_level, use_CLs=False)
+        if use_CLs:
+            p_sb, p_b = self.get_p_vals(conf_level, use_CLs=True)
         else:
-            p_sb = self.get_p_vals(conf_level, use_CLs=True, asymptotic=True)
+            p_sb, powers = self.get_p_vals(conf_level, use_CLs=False)
 
         lower_lim_all = dict()
         upper_lim_all = dict()
         upper_lim_all_raw = dict()
         # Loop over signal sources
         for signal_source in self.signal_source_names:
-            if not asymptotic:
-                these_p_sb = p_sb[signal_source]
-            else:
-                these_p_sb = p_sb[signal_source].test_stats
+            these_p_sb = p_sb[signal_source]
             mus = np.array(list(these_p_sb.keys()))
             p_vals = np.array(list(these_p_sb.values()))
 
-            if not asymptotic:
-                if use_CLs:
-                    these_p_b = p_b[signal_source]
-                    p_vals_b = np.array(list(these_p_b.values()))
-                    p_vals = p_vals / (1. - p_vals_b + 1e-10)
-                else:
-                    these_powers = powers[signal_source]
-                    pws = np.array(list(these_powers.values()))
+            if use_CLs:
+                these_p_b = p_b[signal_source]
+                p_vals_b = np.array(list(these_p_b.values()))
+                p_vals = p_vals / (1. - p_vals_b + 1e-10)
+            else:
+                these_powers = powers[signal_source]
+                pws = np.array(list(these_powers.values()))
 
             # Find points where the p-value curve cross the critical value, decreasing
             upper_lims = np.argwhere(np.diff(np.sign(p_vals - np.ones_like(p_vals) * conf_level)) < 0.).flatten()
@@ -693,7 +650,7 @@ def get_interval(self, conf_level=0.1, pcl_level=0.16,
                                            rising_edge=False, inverse=True)
             upper_lim_raw = upper_lim
 
-            if use_CLs is False and not asymptotic:
+            if use_CLs is False:
                 M0 = self.interp_helper(mus, pws, upper_lims, upper_lim,
                                         rising_edge=False, inverse=False)
                 if M0 < pcl_level:
@@ -707,8 +664,6 @@ def get_interval(self, conf_level=0.1, pcl_level=0.16,
             upper_lim_all[signal_source] = upper_lim
             upper_lim_all_raw[signal_source] = upper_lim_raw
 
-        if asymptotic:
-            return lower_lim_all, upper_lim_all
         if use_CLs is False:
             return lower_lim_all, upper_lim_all, upper_lim_all_raw, p_sb, powers
         else:

From 4af5d3b746a1f79b80bfbcb01734ac6a36ef7acb Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Wed, 16 Jul 2025 16:14:42 +0100
Subject: [PATCH 37/48] Fixes: wasn't correctly sampling constraint centres for
 non-rate parameters in toys; wasn't allowing to save all fits; turn off
 minimum value for rate multiplier guess.

---
 flamedisx/non_asymptotic_inference.py | 38 +++++++++++++++++++++------
 1 file changed, 30 insertions(+), 8 deletions(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 2336788f7..4880369ed 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -355,7 +355,7 @@ def sample_data_constraints(self, mu_test, signal_source_name, likelihood):
 
             # Sample constraint centers
             draw = self.sample_other_constraints[param_name](param_val_expected)
-            constraint_extra_args[param_name] = tf.cast(draw, fd.float_type())
+            constraint_extra_args[f'{param_name}_expected'] = 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]
@@ -389,9 +389,16 @@ def toy_test_statistic_dist(self,
         if save_fits:
             unconditional_bfs_SB = []
             conditional_bfs_SB = []
+
+            unconditional_bfs_SB_disco = []
+            conditional_bfs_SB_disco = []
+
             unconditional_bfs_B = []
             conditional_bfs_B = []
 
+            unconditional_bfs_B_disco = []
+            conditional_bfs_B_disco = []
+
         # Loop over toys
         for toy in tqdm(range(self.ntoys), desc='Doing toys'):
             # S+B toys
@@ -411,9 +418,9 @@ def toy_test_statistic_dist(self,
                 test_statistic_SB = self.test_statistic(likelihood)
                 # Guesses for fit
                 guess_dict_SB = simulate_dict_SB.copy()
-                for key, value in guess_dict_SB.items():
-                    if value < 0.1:
-                        guess_dict_SB[key] = 0.1
+                # for key, value in guess_dict_SB.items():
+                #     if value < 0.1:
+                #         guess_dict_SB[key] = 0.1
                 # Evaluate and save test statistics
                 if discovery_TS:
                     ts_result_SB_disco = test_statistic_SB(0., signal_source_name, guess_dict_SB)
@@ -423,8 +430,12 @@ def toy_test_statistic_dist(self,
                     ts_values_SB.append(ts_result_SB[0])
                 # Possibly save fits
                 if save_fits:
-                    unconditional_bfs_SB.append(ts_result_SB[1])
-                    conditional_bfs_SB.append(ts_result_SB[2])
+                    if discovery_TS:
+                        unconditional_bfs_SB_disco.append(ts_result_SB_disco[1])
+                        conditional_bfs_SB_disco.append(ts_result_SB_disco[2])
+                    else:
+                        unconditional_bfs_SB.append(ts_result_SB[1])
+                        conditional_bfs_SB.append(ts_result_SB[2])
 
             # B-only toys
             if B_toys:
@@ -459,8 +470,12 @@ def toy_test_statistic_dist(self,
                     ts_values_B.append(ts_result_B[0])
                 # Possibly save fits
                 if save_fits:
-                    unconditional_bfs_B.append(ts_result_SB[1])
-                    conditional_bfs_B.append(ts_result_SB[2])
+                    if discovery_TS:
+                        unconditional_bfs_B_disco.append(ts_result_B_disco[1])
+                        conditional_bfs_B_disco.append(ts_result_B_disco[2])
+                    else:
+                        unconditional_bfs_B.append(ts_result_SB[1])
+                        conditional_bfs_B.append(ts_result_SB[2])
 
         # Add to the test statistic distributions
         test_stat_dists_SB.add_ts_dist(mu_test, ts_values_SB)
@@ -472,9 +487,16 @@ def toy_test_statistic_dist(self,
         if save_fits:
             test_stat_dists_SB.add_unconditional_best_fit(mu_test, unconditional_bfs_SB)
             test_stat_dists_SB.add_conditional_best_fit(mu_test, conditional_bfs_SB)
+
+            test_stat_dists_SB_disco.add_unconditional_best_fit(mu_test, unconditional_bfs_SB_disco)
+            test_stat_dists_SB_disco.add_conditional_best_fit(mu_test, conditional_bfs_SB_disco)
+
             test_stat_dists_B.add_unconditional_best_fit(mu_test, unconditional_bfs_B)
             test_stat_dists_B.add_conditional_best_fit(mu_test, conditional_bfs_B)
 
+            test_stat_dists_B_disco.add_unconditional_best_fit(mu_test, unconditional_bfs_B_disco)
+            test_stat_dists_B_disco.add_conditional_best_fit(mu_test, conditional_bfs_B_disco)
+
     def get_observed_test_stat(self, observed_test_stats, observed_data,
                                mu_test, signal_source_name, likelihood, save_fits=False):
         """Internal function to evaluate observed test statistic.

From 3881977440152c78002c23dba12a3454df8a760e Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Wed, 16 Jul 2025 16:16:26 +0100
Subject: [PATCH 38/48] For now, handle combined_rate_scaling differently in
 toys (throw when simulating, keep constraint centre fixed).

---
 flamedisx/non_asymptotic_inference.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 4880369ed..89e806e55 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -365,6 +365,10 @@ def sample_data_constraints(self, mu_test, signal_source_name, likelihood):
                     simulate_dict[pname] = fitval
                     non_rate_params_added.append(pname)
 
+        if 'combined_rate_scaling_expected' in constraint_extra_args:
+            simulate_dict['combined_rate_scaling'] = constraint_extra_args['combined_rate_scaling_expected']
+            constraint_extra_args['combined_rate_scaling_expected'] = 0.
+
         toy_data = likelihood.simulate(**simulate_dict)
 
         if self.observed_test_stats is not None:

From 28951088d95cba5dda55f9b9bad8f8e5308680d9 Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Wed, 16 Jul 2025 22:22:25 +0100
Subject: [PATCH 39/48] Okay, undo the earlier change to the guesses.

---
 flamedisx/non_asymptotic_inference.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 89e806e55..2b8aed063 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -422,9 +422,9 @@ def toy_test_statistic_dist(self,
                 test_statistic_SB = self.test_statistic(likelihood)
                 # Guesses for fit
                 guess_dict_SB = simulate_dict_SB.copy()
-                # for key, value in guess_dict_SB.items():
-                #     if value < 0.1:
-                #         guess_dict_SB[key] = 0.1
+                for key, value in guess_dict_SB.items():
+                    if value < 0.1:
+                        guess_dict_SB[key] = 0.1
                 # Evaluate and save test statistics
                 if discovery_TS:
                     ts_result_SB_disco = test_statistic_SB(0., signal_source_name, guess_dict_SB)

From 3ece905af41f20927d9787748def019998fbff15 Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Thu, 31 Jul 2025 03:04:16 -0700
Subject: [PATCH 40/48] Flag to control when to handle combined_rate_scaling
 differently.

---
 flamedisx/non_asymptotic_inference.py | 25 ++++++++++++++++---------
 1 file changed, 16 insertions(+), 9 deletions(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 2b8aed063..74bb4746f 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -184,7 +184,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,
-                    SB_toys=False, B_toys=False, discovery_TS=False):
+                    SB_toys=False, B_toys=False, discovery_TS=False,
+                    sample_certain_nuisance=False):
         """If observed_data is passed, evaluate observed test statistics. Otherwise,
         obtain test statistic distributions (for both S+B and B-only).
 
@@ -264,7 +265,8 @@ def run_routine(self, mus_test=None, save_fits=False,
                 constraint_extra_args_B_all = []
                 for i in tqdm(range(self.ntoys), desc='Background-only toys'):
                     simulate_dict_B, toy_data_B, constraint_extra_args_B = \
-                        self.sample_data_constraints(0., signal_source, likelihood)
+                        self.sample_data_constraints(0., signal_source, likelihood,
+                                                     sample_certain_nuisance=sample_certain_nuisance)
                     toy_data_B_all.append(toy_data_B)
                     constraint_extra_args_B_all.append(constraint_extra_args_B)
                 simulate_dict_B.pop(f'{signal_source}_rate_multiplier')
@@ -283,7 +285,8 @@ def run_routine(self, mus_test=None, save_fits=False,
                                                  test_stat_dists_B, test_stat_dists_B_disco,
                                                  mu_test, signal_source, likelihood,
                                                  save_fits=save_fits,
-                                                 SB_toys=SB_toys, B_toys=B_toys, discovery_TS=discovery_TS)
+                                                 SB_toys=SB_toys, B_toys=B_toys, discovery_TS=discovery_TS,
+                                                 sample_certain_nuisance=sample_certain_nuisance)
 
             if observed_data is not None:
                 observed_test_stats_collection[signal_source] = observed_test_stats
@@ -299,7 +302,8 @@ def run_routine(self, mus_test=None, save_fits=False,
             return test_stat_dists_SB_collection, test_stat_dists_SB_disco_collection, \
                 test_stat_dists_B_collection, test_stat_dists_B_disco_collection
 
-    def sample_data_constraints(self, mu_test, signal_source_name, likelihood):
+    def sample_data_constraints(self, mu_test, signal_source_name, likelihood,
+                                sample_certain_nuisance=False):
         """Internal function to sample the toy data and constraint central values
         following a frequentist procedure. Method taken depends on whether conditional
         best fits were passed.
@@ -365,9 +369,10 @@ def sample_data_constraints(self, mu_test, signal_source_name, likelihood):
                     simulate_dict[pname] = fitval
                     non_rate_params_added.append(pname)
 
-        if 'combined_rate_scaling_expected' in constraint_extra_args:
-            simulate_dict['combined_rate_scaling'] = constraint_extra_args['combined_rate_scaling_expected']
-            constraint_extra_args['combined_rate_scaling_expected'] = 0.
+        if sample_certain_nuisance:
+            if 'combined_rate_scaling_expected' in constraint_extra_args:
+                simulate_dict['combined_rate_scaling'] = constraint_extra_args['combined_rate_scaling_expected']
+                constraint_extra_args['combined_rate_scaling_expected'] = 0.
 
         toy_data = likelihood.simulate(**simulate_dict)
 
@@ -382,7 +387,8 @@ def toy_test_statistic_dist(self,
                                 test_stat_dists_B, test_stat_dists_B_disco,
                                 mu_test, signal_source_name, likelihood,
                                 save_fits=False,
-                                SB_toys=False, B_toys=False, discovery_TS=False):
+                                SB_toys=False, B_toys=False, discovery_TS=False,
+                                sample_certain_nuisance=False):
         """Internal function to get test statistic distribution.
         """
         ts_values_SB = []
@@ -408,7 +414,8 @@ def toy_test_statistic_dist(self,
             # S+B toys
             if SB_toys:
                 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_test, signal_source_name, likelihood,
+                                                 sample_certain_nuisance)
 
                 # Shift the constraint in the likelihood based on the background RMs we drew
                 likelihood.set_constraint_extra_args(**constraint_extra_args_SB)

From 954e58b2defcbe2fbd3357e762bdd1b5d1391abe Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Fri, 3 Oct 2025 04:23:25 -0700
Subject: [PATCH 41/48] Don't save fit results as tensors.

---
 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 74bb4746f..493f356be 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -32,9 +32,11 @@ def __call__(self, mu_test, signal_source_name, guess_dict):
         # Conditional fit
         bf_conditional = self.likelihood.bestfit(fix=fix_dict, guess=guess_dict_nuisance, suppress_warnings=True,
                                                  allow_failure=True)
+        bf_conditional = {k: v.numpy() for k, v in bf_conditional.items()}
         # Uncnditional fit
         bf_unconditional = self.likelihood.bestfit(guess=guess_dict, suppress_warnings=True,
                                                    allow_failure=True)
+        bf_unconditional = {k: v.numpy() for k, v in bf_unconditional.items()}
 
         # Return the test statistic, unconditional fit and conditional fit
         return self.evaluate(bf_unconditional, bf_conditional), bf_unconditional, bf_conditional

From 2e74990139a3ea5bddde2702f45a08a1ac49d36b Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Mon, 6 Oct 2025 03:50:19 -0700
Subject: [PATCH 42/48] Option to return the toys (toy indices) for the toys
 setting ULs closest to the band quantiles.

---
 flamedisx/non_asymptotic_inference.py | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 493f356be..4698d0ef3 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -718,10 +718,11 @@ def critical_disco_value(self, disco_pot_curve, mus, discovery_sigma):
                                   rising_edge=True, inverse=True)
 
     def get_bands(self, conf_level=0.1, quantiles=[0, 1, -1, 2, -2],
-                  use_CLs=False):
+                  use_CLs=False, return_toy_indices=False):
         """
         """
         bands = dict()
+        toy_indices = dict()
 
         # Loop over signal sources
         for signal_source in self.signal_source_names:
@@ -747,16 +748,26 @@ def get_bands(self, conf_level=0.1, quantiles=[0, 1, -1, 2, -2],
             p_val_curves = np.transpose(np.stack(p_val_curves, axis=0))
             upper_lims_bands = np.apply_along_axis(self.upper_lims_bands, 1, p_val_curves, mus, conf_level)
 
+            upper_lims_bands_all = upper_lims_bands
             if len(upper_lims_bands[upper_lims_bands == 0.]) > 0.:
                 print(f'Found {len(upper_lims_bands[upper_lims_bands == 0.])} failed toy for {signal_source}; removing...')
                 upper_lims_bands = upper_lims_bands[upper_lims_bands > 0.]
 
             these_bands = dict()
+            these_toy_indices = dict()
             for quantile in quantiles:
                 these_bands[quantile] = np.quantile(np.sort(upper_lims_bands), stats.norm.cdf(quantile))
+
+                nearest_index = np.argmin(np.abs(upper_lims_bands_all - these_bands[quantile]))
+                these_toy_indices[quantile] = nearest_index
+
             bands[signal_source] = these_bands
+            toy_indices[signal_source] = these_toy_indices
 
-        return bands
+        if return_toy_indices:
+            return bands, toy_indices
+        else:
+            return bands
 
     def get_disco_sig(self):
         """

From f18493190271bbe022b9f2b59dcb9682042cb5ca Mon Sep 17 00:00:00 2001
From: Robert James <robert.james.19@ucl.ac.uk>
Date: Tue, 7 Oct 2025 22:28:56 -0700
Subject: [PATCH 43/48] Was saving the wrong thing when saving toy fit results,
 for the background-only case.

---
 flamedisx/non_asymptotic_inference.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 4698d0ef3..6a0c7b110 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -487,8 +487,8 @@ def toy_test_statistic_dist(self,
                         unconditional_bfs_B_disco.append(ts_result_B_disco[1])
                         conditional_bfs_B_disco.append(ts_result_B_disco[2])
                     else:
-                        unconditional_bfs_B.append(ts_result_SB[1])
-                        conditional_bfs_B.append(ts_result_SB[2])
+                        unconditional_bfs_B.append(ts_result_B[1])
+                        conditional_bfs_B.append(ts_result_B[2])
 
         # Add to the test statistic distributions
         test_stat_dists_SB.add_ts_dist(mu_test, ts_values_SB)

From 89737c8a83baf383a7f2d09b2885f2d6a9b51b8c Mon Sep 17 00:00:00 2001
From: josh0-jrg <joshua.robert.green@gmail.com>
Date: Thu, 20 Nov 2025 07:38:04 -0800
Subject: [PATCH 44/48] Feat: allows evaluation of expected median discovery
 significance

---
 flamedisx/non_asymptotic_inference.py | 22 +++++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 6a0c7b110..146fc1ba8 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -391,7 +391,23 @@ def toy_test_statistic_dist(self,
                                 save_fits=False,
                                 SB_toys=False, B_toys=False, discovery_TS=False,
                                 sample_certain_nuisance=False):
-        """Internal function to get test statistic distribution.
+        """
+            Internal function to get test statistic distribution given a signal and POI value. |
+            test_stat_dists_SB:        TestStatisticDistributions, t(mu_test|mu=mu_test) * |
+            test_stat_dists_SB_disco:  TestStatisticDistributions, t(0.|mu=mu_test) * |
+            test_stat_dists_B:         TestStatisticDistributions, t(mu_test|mu=0.) * |
+            test_stat_dists_B_disco:   TestStatisticDistributions, t(0.|mu=0.) * |
+            mu_test:                   float, POI test value (usually signal counts). |
+            signal_source_name:        string, the source that takes the POI. |
+            likelihood:                LogLikelihood,the likelihood object. |
+            save_fits:                 bool, whether or not to save cond/uncond fits, stored in 
+                                             TestStatisticDistributions. |
+            SB_toys:                   bool, whether or not to simulate S+B toys. |
+            B_toys:                    bool, whether or not to simulate B toys. |
+            discovery_TS:              bool, wether to **only** evaluate test_stat_dists_SB_disco 
+                                             and not test_stat_dists_SB. |
+            return:                    None, updates flamedisx TestStatisticDistributions objects in first
+                                             inputs (*). |
         """
         ts_values_SB = []
         ts_values_SB_disco = []
@@ -436,9 +452,13 @@ def toy_test_statistic_dist(self,
                         guess_dict_SB[key] = 0.1
                 # Evaluate and save test statistics
                 if discovery_TS:
+                    # If we're doing significance, lots of toys
                     ts_result_SB_disco = test_statistic_SB(0., signal_source_name, guess_dict_SB)
                     ts_values_SB_disco.append(ts_result_SB_disco[0])
                 else:
+                    # If we're doing significance limits, can afford the extra eval.
+                    ts_result_SB_disco = test_statistic_SB(0., signal_source_name, guess_dict_SB)
+                    ts_values_SB_disco.append(ts_result_SB_disco[0])
                     ts_result_SB = test_statistic_SB(mu_test, signal_source_name, guess_dict_SB)
                     ts_values_SB.append(ts_result_SB[0])
                 # Possibly save fits

From e4cc6885d0d1b1b2a8e28ef5a1a53e61f6e5dfce Mon Sep 17 00:00:00 2001
From: Sam Eriksen <sam.eriksen@bristol.ac.uk>
Date: Tue, 1 Jul 2025 10:43:57 +0200
Subject: [PATCH 45/48] fix: get_median_disc_asymptotic to pass sigma_level

---
 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 146fc1ba8..4729d3907 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -834,7 +834,7 @@ def get_median_disco_asymptotic(self, sigma_level=3):
             disco_sig_curves = np.stack(disco_sig_curves, axis=0)
             median_disco_sigs = [np.median(disco_sigs) for disco_sigs in disco_sig_curves]
 
-            median_crossing_point = self.critical_disco_value(median_disco_sigs, mus, 3)
+            median_crossing_point = self.critical_disco_value(median_disco_sigs, mus, sigma_level)
             medians[signal_source] = median_crossing_point
 
         return medians

From 1db1cc3f67e030be26f8edaa445bb8f1f4803d6c Mon Sep 17 00:00:00 2001
From: Sam Eriksen <sam.eriksen@bristol.ac.uk>
Date: Tue, 1 Jul 2025 10:55:09 +0200
Subject: [PATCH 46/48] debug: adding print statements whilst debugging

---
 flamedisx/non_asymptotic_inference.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 4729d3907..01939b0a9 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -833,7 +833,9 @@ def get_median_disco_asymptotic(self, sigma_level=3):
 
             disco_sig_curves = np.stack(disco_sig_curves, axis=0)
             median_disco_sigs = [np.median(disco_sigs) for disco_sigs in disco_sig_curves]
-
+            print("median_disco_sigs")
+            print(median_disco_sigs)
+            print(f"sigma_level: {sigma_level}")
             median_crossing_point = self.critical_disco_value(median_disco_sigs, mus, sigma_level)
             medians[signal_source] = median_crossing_point
 

From 40bf3d7c49af892553ee574b61e246a777e0fc39 Mon Sep 17 00:00:00 2001
From: Sam Eriksen <sam.eriksen@bristol.ac.uk>
Date: Tue, 25 Nov 2025 11:39:00 +0000
Subject: [PATCH 47/48] chore: remove debug prints from
 non_asymptotic_inference.py

Removed debug print statements for median disco signals.
---
 flamedisx/non_asymptotic_inference.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/flamedisx/non_asymptotic_inference.py b/flamedisx/non_asymptotic_inference.py
index 01939b0a9..8fc741aa8 100644
--- a/flamedisx/non_asymptotic_inference.py
+++ b/flamedisx/non_asymptotic_inference.py
@@ -833,9 +833,6 @@ def get_median_disco_asymptotic(self, sigma_level=3):
 
             disco_sig_curves = np.stack(disco_sig_curves, axis=0)
             median_disco_sigs = [np.median(disco_sigs) for disco_sigs in disco_sig_curves]
-            print("median_disco_sigs")
-            print(median_disco_sigs)
-            print(f"sigma_level: {sigma_level}")
             median_crossing_point = self.critical_disco_value(median_disco_sigs, mus, sigma_level)
             medians[signal_source] = median_crossing_point
 

From 8e46e7de8b5f6a0c80145cca6fe7e3aee9dfdb53 Mon Sep 17 00:00:00 2001
From: josh0-jrg <joshua.robert.green@gmail.com>
Date: Thu, 4 Dec 2025 04:48:40 -0800
Subject: [PATCH 48/48] Fix: return zero gradients instead of none on rate
 multipliers

This was causing a serious bug where-in if a sub-likelihood had 0 events but hessians were used in the combined likelihood it would break
---
 flamedisx/likelihood.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/flamedisx/likelihood.py b/flamedisx/likelihood.py
index d8d8a0717..ff8c8067d 100644
--- a/flamedisx/likelihood.py
+++ b/flamedisx/likelihood.py
@@ -469,7 +469,9 @@ def mu(self, *,
         :param dataset_name: ... for just this dataset
         :param source_name: ... for just this source.
         You must provide either dsetname or source, since it makes no sense to
-        add events from multiple datasets
+        add events from multiple datasets.
+        For rate multipliers (always linear) add a 0 x r.m**2 term to give a 0
+        hessian instead of None.
         """
         kwargs = {**self.param_defaults, **kwargs}
         if dataset_name is None and source_name is None:
@@ -482,8 +484,10 @@ def mu(self, *,
             if source_name is not None and sname != source_name:
                 continue
             filtered_params = self._filter_source_kwargs(kwargs, sname)
-            mu += (self._get_rate_mult(sname, kwargs)
-                   * self.mu_estimators[sname](**filtered_params))
+            _rate_multiplier = self._get_rate_mult(sname, kwargs)
+            mu += (_rate_multiplier
+                   * self.mu_estimators[sname](**filtered_params) 
+                   + tf.constant(0.,fd.float_type())*_rate_multiplier**2)
         return mu
 
     @tf.function