From 1324a30fdcb4c81c769ebf2da4b891e09bbcb706 Mon Sep 17 00:00:00 2001
From: Robert Hammann <robert.hammann@outlook.com>
Date: Thu, 10 Oct 2024 11:32:33 -0500
Subject: [PATCH 01/14] first idea of source-wise interpolation

---
 blueice/likelihood.py | 127 ++++++++++++++++++++++++++++++++++--------
 1 file changed, 105 insertions(+), 22 deletions(-)

diff --git a/blueice/likelihood.py b/blueice/likelihood.py
index 7fe2aa1..5eee25c 100644
--- a/blueice/likelihood.py
+++ b/blueice/likelihood.py
@@ -73,6 +73,7 @@ def __init__(self, pdf_base_config, likelihood_config=None, **kwargs):
             likelihood_config = {}
         self.config = likelihood_config
         self.config.setdefault('morpher', 'GridInterpolator')
+        self.source_wise_interpolation = self.config.get('source_wise_interpolation', False)
 
         # Base model: no variations of any settings
         self.base_model = Model(self.pdf_base_config)
@@ -101,6 +102,7 @@ def __init__(self, pdf_base_config, likelihood_config=None, **kwargs):
 
         # If there are shape parameters:
         self.anchor_models = OrderedDict()  # dictionary mapping model zs -> actual model
+        self.anchor_sources = OrderedDict()  # dictionary mapping source_name -> zs -> source
         # Interpolators created by morphers. These map zs to...
         self.mus_interpolator = None                          # rates for each source
         self.ps_interpolator = None                          # (source, event) p-values (unbinned), or pmf grid (binned)
@@ -108,14 +110,57 @@ def __init__(self, pdf_base_config, likelihood_config=None, **kwargs):
         self.n_model_events_interpolator = lambda x: None
         self.n_model_events = None
 
+    @property
+    def source_shape_parameters(self):
+        source_shape_parameters = OrderedDict()
+        for sn, source_config in zip(self.source_name_list, self.pdf_base_config["sources"]):
+            parameter_names = source_config["parameters"]
+            shape_parameters = OrderedDict({k: v for k, v in self.shape_parameters.items() if k in parameter_names})
+            if shape_parameters:
+                source_shape_parameters[sn] = shape_parameters
+        return source_shape_parameters
+
+    def _get_shape_indices(self, source_name):
+        """Return the indices of the shape parameters used by the source."""
+        shape_keys = self.source_shape_parameters[source_name].keys()
+        return [i for i, k in enumerate(self.shape_parameters.keys()) if k in shape_keys]
+
+    def _get_model_anchor(self, anchor, source_name):
+        """Return the shape anchors of the full model, given the shape anchors of a signle source.
+        All values of shape parameters not used by this source will be set to None.
+        """
+        shape_keys = self.source_shape_parameters[source_name].keys()
+        shape_indices = self.get_shape_indices(shape_keys)
+        model_anchor = [None] * len(self.shape_parameters)
+        for i, idx in enumerate(shape_indices):
+            model_anchor[idx] = anchor[i]
+        return tuple(model_anchor)
+
     def prepare(self, n_cores=1, ipp_client=None):
         """Prepares a likelihood function with shape parameters for use.
         This will compute the models for each shape parameter anchor value combination.
         """
         if len(self.shape_parameters):
-            self.morpher = MORPHERS[self.config['morpher']](self.config.get('morpher_config', {}),
-                                                            self.shape_parameters)
-            zs_list = self.morpher.get_anchor_points(bounds=self.get_bounds())
+            if self.source_wise_interpolation:
+                # Create morphers for each source individually
+                self.source_morphers = OrderedDict()
+                for sn, shape_parameters in self.source_shape_parameters.items():
+                    self.source_morphers[sn] = MORPHERS[self.config['morpher']](
+                        self.config.get('morpher_config', {}),
+                        shape_parameters
+                    )
+                zs_list = set()
+                for source_name, morpher in self.source_morphers.items():
+                    anchor_points = morpher.get_anchor_points(bounds=None)
+                    for anchor in anchor_points:
+                        zs = self._get_model_anchor(anchor, source_name)
+                        zs_list.add(zs)
+                zs_list = list(zs_list)
+                
+            else:
+                self.morpher = MORPHERS[self.config['morpher']](self.config.get('morpher_config', {}),
+                                                                self.shape_parameters)
+                zs_list = self.morpher.get_anchor_points(bounds=self.get_bounds())
 
             # Create the configs for each new model
             configs = []
@@ -123,7 +168,8 @@ def prepare(self, n_cores=1, ipp_client=None):
                 config = deepcopy(self.pdf_base_config)
                 for i, (setting_name, (anchors, _, _)) in enumerate(self.shape_parameters.items()):
                     # Translate from zs to settings using the anchors dict. Maybe not all settings are numerical.
-                    config[setting_name] = anchors[zs[i]]
+                    if zs[i] is not None:
+                        config[setting_name] = anchors[zs[i]]
                 if ipp_client is None and n_cores != 1:
                     # We have to compute in parallel: must have delayed computation on
                     config['delay_pdf_computation'] = True
@@ -150,14 +196,25 @@ def prepare(self, n_cores=1, ipp_client=None):
                 # Reload models so computation takes effect
                 models = [Model(c) for c in tqdm(configs, desc="Loading computed models")]
 
-            # Add the new models to the anchor_models dict
-            for zs, model in zip(zs_list, models):
-                self.anchor_models[tuple(zs)] = model
+            if self.source_wise_interpolation:
+                for i,(source_name, morpher) in enumerate(source_morphers.items()):
+                    anchors = morpher.get_anchor_points(bounds=None)
+                    self.anchor_sources[source_name] = OrderedDict()
+                    for anchor  in anchors:
+                        model_anchor = self._get_model_anchor(anchor, source_name)
+                        model_index = zs_list.index(model_anchor)
+                        self.anchor_sources[source_name][anchor] = models[model_index].sources[i]
+                # TODO: Implement self.mus_interpolator for source-wise interpolation
+            
+            else:
+                # Add the new models to the anchor_models dict
+                for zs, model in zip(zs_list, models):
+                    self.anchor_models[tuple(zs)] = model
 
-            # Build the interpolator for the rates of each source.
-            self.mus_interpolator = self.morpher.make_interpolator(f=lambda m: m.expected_events(),
-                                                                   extra_dims=[len(self.source_name_list)],
-                                                                   anchor_models=self.anchor_models)
+                # Build the interpolator for the rates of each source.
+                self.mus_interpolator = self.morpher.make_interpolator(f=lambda m: m.expected_events(),
+                                                                    extra_dims=[len(self.source_name_list)],
+                                                                    anchor_models=self.anchor_models)
 
         self.is_data_set = False
         self.is_prepared = True
@@ -440,9 +497,32 @@ class UnbinnedLogLikelihood(LogLikelihoodBase):
     def set_data(self, d):
         LogLikelihoodBase.set_data(self, d)
         if len(self.shape_parameters):
-            self.ps_interpolator = self.morpher.make_interpolator(f=lambda m: m.score_events(d),
-                                                                  extra_dims=[len(self.source_name_list), len(d)],
-                                                                  anchor_models=self.anchor_models)
+            if self.source_wise_interpolation:
+                ps_interpolators = OrderedDict()
+                for sn, base_source in zip(self.source_name_list, self.base_model.sources):
+                    if sn in self.source_morphers:
+                        ps_interpolators[sn] = self.source_morphers[sn].make_interpolator(
+                            f=lambda s: s.pdf(*self.base_model.to_analysis_dimensions(d)),
+                            extra_dims=[len(d)],
+                            anchor_models=self.anchor_sources[sn])
+                    else:
+                        ps_interpolators[sn] = base_source.pdf(*self.base_model.to_analysis_dimensions(d))
+                def ps_interpolator(*args):
+                    # take zs, convert to values for each source's interpolator call the respective interpolator
+                    ps = []
+                    for sn in self.source_name_list:
+                        if sn in self.source_shape_parameters:
+                            shape_indices = self._get_shape_indices(sn)
+                            these_args = [args[i] for i in shape_indices]
+                            ps.append(ps_interpolators[sn](np.asarray(these_args)))
+                        else:
+                            ps.append(ps_interpolators[sn])
+                    return ps
+                self.ps_interpolator = ps_interpolator
+            else:
+                self.ps_interpolator = self.morpher.make_interpolator(f=lambda m: m.score_events(d),
+                                                                    extra_dims=[len(self.source_name_list), len(d)],
+                                                                    anchor_models=self.anchor_models)
         else:
             self.ps = self.base_model.score_events(d)
 
@@ -472,14 +552,17 @@ def prepare(self, *args):
         self.ps, self.n_model_events = self.base_model.pmf_grids()
 
         if len(self.shape_parameters):
-            self.ps_interpolator = self.morpher.make_interpolator(f=lambda m: m.pmf_grids()[0],
-                                                                  extra_dims=list(self.ps.shape),
-                                                                  anchor_models=self.anchor_models)
-
-            if self.model_statistical_uncertainty_handling is not None:
-                self.n_model_events_interpolator = self.morpher.make_interpolator(f=lambda m: m.pmf_grids()[1],
-                                                                                  extra_dims=list(self.ps.shape),
-                                                                                  anchor_models=self.anchor_models)
+            if self.source_wise_interpolation:
+                raise NotImplementedError("Source-wise interpolation not implemented for binned likelihoods")
+            else:
+                self.ps_interpolator = self.morpher.make_interpolator(f=lambda m: m.pmf_grids()[0],
+                                                                    extra_dims=list(self.ps.shape),
+                                                                    anchor_models=self.anchor_models)
+
+                if self.model_statistical_uncertainty_handling is not None:
+                    self.n_model_events_interpolator = self.morpher.make_interpolator(f=lambda m: m.pmf_grids()[1],
+                                                                                    extra_dims=list(self.ps.shape),
+                                                                                    anchor_models=self.anchor_models)
 
     @inherit_docstring_from(LogLikelihoodBase)
     def set_data(self, d):

From d92246a846c09dc59a40a0b1090802e6b4832e6a Mon Sep 17 00:00:00 2001
From: Robert Hammann <robert.hammann@outlook.com>
Date: Thu, 10 Oct 2024 11:55:06 -0500
Subject: [PATCH 02/14] move expected_events to source

---
 blueice/model.py  | 3 +--
 blueice/source.py | 4 ++++
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/blueice/model.py b/blueice/model.py
index 82fdf4d..38804c8 100644
--- a/blueice/model.py
+++ b/blueice/model.py
@@ -106,11 +106,10 @@ def pmf_grids(self):
     def expected_events(self, s=None):
         """Return the total number of events expected in the analysis range for the source s.
         If no source specified, return an array of results for all sources.
-        # TODO: Why is this not a method of source?
         """
         if s is None:
             return np.array([self.expected_events(s) for s in self.sources])
-        return s.events_per_day * self.config['livetime_days'] * s.fraction_in_range * s.config['rate_multiplier']
+        return s.expected_events(self.config['livetime_days'])
 
     def show(self, d, ax=None, dims=None, **kwargs):
         """Plot the events from dataset d in the analysis range
diff --git a/blueice/source.py b/blueice/source.py
index d31fccb..f07b228 100644
--- a/blueice/source.py
+++ b/blueice/source.py
@@ -261,6 +261,10 @@ def simulate(self, n_events):
     def get_pmf_grid(self):
         return self._pdf_histogram.histogram * self._bin_volumes, self._n_events_histogram.histogram
 
+    def expected_events(self, livetime_days):
+        """Return the total number of events expected in the analysis range for the source."""
+        return self.events_per_day * livetime_days * self.fraction_in_range * self.config['rate_multiplier']
+
 
 class DensityEstimatingSource(HistogramPdfSource):
     """A source which estimates its PDF by some events you give to it.

From 0fdb5b28133067c1ea226de2ecfcdfa10b31f0f4 Mon Sep 17 00:00:00 2001
From: Robert Hammann <robert.hammann@outlook.com>
Date: Thu, 10 Oct 2024 13:06:48 -0500
Subject: [PATCH 03/14] fix expected_events property

---
 blueice/model.py  | 2 +-
 blueice/source.py | 9 +++++----
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/blueice/model.py b/blueice/model.py
index 38804c8..9e73a81 100644
--- a/blueice/model.py
+++ b/blueice/model.py
@@ -109,7 +109,7 @@ def expected_events(self, s=None):
         """
         if s is None:
             return np.array([self.expected_events(s) for s in self.sources])
-        return s.expected_events(self.config['livetime_days'])
+        return s.expected_events
 
     def show(self, d, ax=None, dims=None, **kwargs):
         """Plot the events from dataset d in the analysis range
diff --git a/blueice/source.py b/blueice/source.py
index f07b228..9537dc3 100644
--- a/blueice/source.py
+++ b/blueice/source.py
@@ -182,6 +182,11 @@ def simulate(self, n_events):
         if you decide some events are not detectable.
         """
         raise NotImplementedError
+    
+    @property
+    def expected_events(self):
+        """Return the default total number of events expected in the analysis range for the source."""
+        return self.events_per_day * self.config['livetime_days'] * self.fraction_in_range * self.config['rate_multiplier']
 
 
 class HistogramPdfSource(Source):
@@ -261,10 +266,6 @@ def simulate(self, n_events):
     def get_pmf_grid(self):
         return self._pdf_histogram.histogram * self._bin_volumes, self._n_events_histogram.histogram
 
-    def expected_events(self, livetime_days):
-        """Return the total number of events expected in the analysis range for the source."""
-        return self.events_per_day * livetime_days * self.fraction_in_range * self.config['rate_multiplier']
-
 
 class DensityEstimatingSource(HistogramPdfSource):
     """A source which estimates its PDF by some events you give to it.

From 5baa97a03e85ad89b1fa376276c7a9c45a7a428b Mon Sep 17 00:00:00 2001
From: Robert Hammann <robert.hammann@outlook.com>
Date: Thu, 10 Oct 2024 13:07:03 -0500
Subject: [PATCH 04/14] add source-wise mus_interpolator

---
 blueice/likelihood.py | 36 ++++++++++++++++++++++++++++--------
 1 file changed, 28 insertions(+), 8 deletions(-)

diff --git a/blueice/likelihood.py b/blueice/likelihood.py
index 5eee25c..aa7de7d 100644
--- a/blueice/likelihood.py
+++ b/blueice/likelihood.py
@@ -73,7 +73,7 @@ def __init__(self, pdf_base_config, likelihood_config=None, **kwargs):
             likelihood_config = {}
         self.config = likelihood_config
         self.config.setdefault('morpher', 'GridInterpolator')
-        self.source_wise_interpolation = self.config.get('source_wise_interpolation', False)
+        self.source_wise_interpolation = self.pdf_base_config.get('source_wise_interpolation', False)
 
         # Base model: no variations of any settings
         self.base_model = Model(self.pdf_base_config)
@@ -129,8 +129,7 @@ def _get_model_anchor(self, anchor, source_name):
         """Return the shape anchors of the full model, given the shape anchors of a signle source.
         All values of shape parameters not used by this source will be set to None.
         """
-        shape_keys = self.source_shape_parameters[source_name].keys()
-        shape_indices = self.get_shape_indices(shape_keys)
+        shape_indices = self._get_shape_indices(source_name)
         model_anchor = [None] * len(self.shape_parameters)
         for i, idx in enumerate(shape_indices):
             model_anchor[idx] = anchor[i]
@@ -197,15 +196,36 @@ def prepare(self, n_cores=1, ipp_client=None):
                 models = [Model(c) for c in tqdm(configs, desc="Loading computed models")]
 
             if self.source_wise_interpolation:
-                for i,(source_name, morpher) in enumerate(source_morphers.items()):
+                print("USING SOURCE-WISE INTERPOLATION")
+                for i,(source_name, morpher) in enumerate(self.source_morphers.items()):
                     anchors = morpher.get_anchor_points(bounds=None)
                     self.anchor_sources[source_name] = OrderedDict()
                     for anchor  in anchors:
                         model_anchor = self._get_model_anchor(anchor, source_name)
                         model_index = zs_list.index(model_anchor)
                         self.anchor_sources[source_name][anchor] = models[model_index].sources[i]
-                # TODO: Implement self.mus_interpolator for source-wise interpolation
-            
+                mus_interpolators = OrderedDict()
+                for sn, base_source in zip(self.source_name_list, self.base_model.sources):
+                    if sn in self.source_morphers:
+                        mus_interpolators[sn] = self.source_morphers[sn].make_interpolator(
+                            f=lambda s: s.expected_events,
+                            extra_dims=[1],
+                            anchor_models=self.anchor_sources[sn])
+                    else:
+                        mus_interpolators[sn] = base_source.expected_events
+                def mus_interpolator(*args):
+                    # take zs, convert to values for each source's interpolator call the respective interpolator
+                    mus = []
+                    for sn in self.source_name_list:
+                        if sn in self.source_shape_parameters:
+                            shape_indices = self._get_shape_indices(sn)
+                            these_args = [args[0][i] for i in shape_indices]
+                            mus.append(mus_interpolators[sn](np.asarray(these_args))[0])
+                        else:
+                            mus.append(mus_interpolators[sn])
+                    return np.array(mus)
+                self.mus_interpolator = mus_interpolator
+
             else:
                 # Add the new models to the anchor_models dict
                 for zs, model in zip(zs_list, models):
@@ -513,11 +533,11 @@ def ps_interpolator(*args):
                     for sn in self.source_name_list:
                         if sn in self.source_shape_parameters:
                             shape_indices = self._get_shape_indices(sn)
-                            these_args = [args[i] for i in shape_indices]
+                            these_args = [args[0][i] for i in shape_indices]
                             ps.append(ps_interpolators[sn](np.asarray(these_args)))
                         else:
                             ps.append(ps_interpolators[sn])
-                    return ps
+                    return np.array(ps)
                 self.ps_interpolator = ps_interpolator
             else:
                 self.ps_interpolator = self.morpher.make_interpolator(f=lambda m: m.score_events(d),

From 324b3e0ebfb4c19cc9c93bf9c46ddb89b5d7cff4 Mon Sep 17 00:00:00 2001
From: Robert Hammann <robert.hammann@outlook.com>
Date: Fri, 11 Oct 2024 12:01:39 +0200
Subject: [PATCH 05/14] some cleanup..

---
 blueice/likelihood.py | 30 +++++++++++++++++-------------
 blueice/source.py     |  2 +-
 2 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/blueice/likelihood.py b/blueice/likelihood.py
index aa7de7d..c791d8e 100644
--- a/blueice/likelihood.py
+++ b/blueice/likelihood.py
@@ -155,7 +155,7 @@ def prepare(self, n_cores=1, ipp_client=None):
                         zs = self._get_model_anchor(anchor, source_name)
                         zs_list.add(zs)
                 zs_list = list(zs_list)
-                
+
             else:
                 self.morpher = MORPHERS[self.config['morpher']](self.config.get('morpher_config', {}),
                                                                 self.shape_parameters)
@@ -196,11 +196,10 @@ def prepare(self, n_cores=1, ipp_client=None):
                 models = [Model(c) for c in tqdm(configs, desc="Loading computed models")]
 
             if self.source_wise_interpolation:
-                print("USING SOURCE-WISE INTERPOLATION")
-                for i,(source_name, morpher) in enumerate(self.source_morphers.items()):
+                for i, (source_name, morpher) in enumerate(self.source_morphers.items()):
                     anchors = morpher.get_anchor_points(bounds=None)
                     self.anchor_sources[source_name] = OrderedDict()
-                    for anchor  in anchors:
+                    for anchor in anchors:
                         model_anchor = self._get_model_anchor(anchor, source_name)
                         model_index = zs_list.index(model_anchor)
                         self.anchor_sources[source_name][anchor] = models[model_index].sources[i]
@@ -213,6 +212,7 @@ def prepare(self, n_cores=1, ipp_client=None):
                             anchor_models=self.anchor_sources[sn])
                     else:
                         mus_interpolators[sn] = base_source.expected_events
+
                 def mus_interpolator(*args):
                     # take zs, convert to values for each source's interpolator call the respective interpolator
                     mus = []
@@ -232,9 +232,10 @@ def mus_interpolator(*args):
                     self.anchor_models[tuple(zs)] = model
 
                 # Build the interpolator for the rates of each source.
-                self.mus_interpolator = self.morpher.make_interpolator(f=lambda m: m.expected_events(),
-                                                                    extra_dims=[len(self.source_name_list)],
-                                                                    anchor_models=self.anchor_models)
+                self.mus_interpolator = self.morpher.make_interpolator(
+                    f=lambda m: m.expected_events(),
+                    extra_dims=[len(self.source_name_list)],
+                    anchor_models=self.anchor_models)
 
         self.is_data_set = False
         self.is_prepared = True
@@ -527,6 +528,7 @@ def set_data(self, d):
                             anchor_models=self.anchor_sources[sn])
                     else:
                         ps_interpolators[sn] = base_source.pdf(*self.base_model.to_analysis_dimensions(d))
+
                 def ps_interpolator(*args):
                     # take zs, convert to values for each source's interpolator call the respective interpolator
                     ps = []
@@ -540,9 +542,10 @@ def ps_interpolator(*args):
                     return np.array(ps)
                 self.ps_interpolator = ps_interpolator
             else:
-                self.ps_interpolator = self.morpher.make_interpolator(f=lambda m: m.score_events(d),
-                                                                    extra_dims=[len(self.source_name_list), len(d)],
-                                                                    anchor_models=self.anchor_models)
+                self.ps_interpolator = self.morpher.make_interpolator(
+                    f=lambda m: m.score_events(d),
+                    extra_dims=[len(self.source_name_list), len(d)],
+                    anchor_models=self.anchor_models)
         else:
             self.ps = self.base_model.score_events(d)
 
@@ -575,9 +578,10 @@ def prepare(self, *args):
             if self.source_wise_interpolation:
                 raise NotImplementedError("Source-wise interpolation not implemented for binned likelihoods")
             else:
-                self.ps_interpolator = self.morpher.make_interpolator(f=lambda m: m.pmf_grids()[0],
-                                                                    extra_dims=list(self.ps.shape),
-                                                                    anchor_models=self.anchor_models)
+                self.ps_interpolator = self.morpher.make_interpolator(
+                    f=lambda m: m.pmf_grids()[0],
+                    extra_dims=list(self.ps.shape),
+                    anchor_models=self.anchor_models)
 
                 if self.model_statistical_uncertainty_handling is not None:
                     self.n_model_events_interpolator = self.morpher.make_interpolator(f=lambda m: m.pmf_grids()[1],
diff --git a/blueice/source.py b/blueice/source.py
index 9537dc3..145a260 100644
--- a/blueice/source.py
+++ b/blueice/source.py
@@ -182,7 +182,7 @@ def simulate(self, n_events):
         if you decide some events are not detectable.
         """
         raise NotImplementedError
-    
+
     @property
     def expected_events(self):
         """Return the default total number of events expected in the analysis range for the source."""

From 305fedc7b3d5fc8ed1252c0d40c4aadc39873fdf Mon Sep 17 00:00:00 2001
From: Robert Hammann <robert.hammann@outlook.com>
Date: Fri, 11 Oct 2024 13:57:19 +0200
Subject: [PATCH 06/14] add source parameters method

---
 blueice/likelihood.py | 7 +++++--
 blueice/source.py     | 2 ++
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/blueice/likelihood.py b/blueice/likelihood.py
index c791d8e..f3fdaba 100644
--- a/blueice/likelihood.py
+++ b/blueice/likelihood.py
@@ -112,9 +112,12 @@ def __init__(self, pdf_base_config, likelihood_config=None, **kwargs):
 
     @property
     def source_shape_parameters(self):
+        """Dict of sources with shape parameters. source name -> dict of shape parameters."""
         source_shape_parameters = OrderedDict()
-        for sn, source_config in zip(self.source_name_list, self.pdf_base_config["sources"]):
-            parameter_names = source_config["parameters"]
+        for sn, source in zip(self.source_name_list, self.base_model.sources):
+            parameter_names = source.parameters
+            if parameter_names is None:
+                raise ValueError("The `parameters` of each source need to be specified when using `source_wise_interpolation`.")
             shape_parameters = OrderedDict({k: v for k, v in self.shape_parameters.items() if k in parameter_names})
             if shape_parameters:
                 source_shape_parameters[sn] = shape_parameters
diff --git a/blueice/source.py b/blueice/source.py
index 145a260..9bcbeb9 100644
--- a/blueice/source.py
+++ b/blueice/source.py
@@ -86,6 +86,8 @@ def __init__(self, config, *args, **kwargs):
         self.name = c['name']
         del c['name']
 
+        self.parameters = c.pop('parameters', None)
+
         # events_per_day and fraction_in_range may be modified / set properly for the first time later (see comments
         # in 'defaults' above)
         if hasattr(self, 'events_per_day'):

From cb20f4a7a1ec566a27f8e9e0c178b8e464ddec5b Mon Sep 17 00:00:00 2001
From: Robert Hammann <robert.hammann@outlook.com>
Date: Fri, 11 Oct 2024 13:57:37 +0200
Subject: [PATCH 07/14] add unittest

---
 tests/test_likelihood.py | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/tests/test_likelihood.py b/tests/test_likelihood.py
index 7377d72..20dbf9b 100644
--- a/tests/test_likelihood.py
+++ b/tests/test_likelihood.py
@@ -92,6 +92,37 @@ def test_shape_uncertainty():
     assert lf(strlen_multiplier=2) == -2 + np.log(2 * stats.norm.pdf(0)) + log_prior(2)
 
 
+def test_source_wise_interpolation():
+    data = np.zeros(5, dtype=[('x', float), ('source', int)])
+    data['x'] = np.linspace(0, 1, 5)
+
+    config = conf_for_test(events_per_day=1)
+
+    lf = UnbinnedLogLikelihood(config)
+    lf.add_shape_parameter("mu", anchors={-2:-2, 0:0, 2:2})
+    lf.prepare()
+    lf.set_data(data)
+    ret_0 = lf(full_output=True)
+    ret_1 = lf(full_output=True, mu=1)
+
+    config["source_wise_interpolation"] = True
+    # We need to specify all parameters used by each source
+    config["sources"][0]["parameters"] = ["mu", "sigma", "strlen_multiplier", "some_multiplier", "s0_rate_multiplier"]
+    lf_source_wise = UnbinnedLogLikelihood(config)
+    lf_source_wise.add_shape_parameter("mu", anchors={-2:-2, 0:0, 2:2})
+    lf_source_wise.prepare()
+    lf_source_wise.set_data(data)
+    ret_source_wise_0 = lf_source_wise(full_output=True)
+    ret_source_wise_1 = lf_source_wise(full_output=True, mu=1)
+
+    assert ret_0[0] == ret_source_wise_0[0]
+    assert (ret_0[1] == ret_source_wise_0[1]).all()
+    assert (ret_0[2] == ret_source_wise_0[2]).all()
+    assert ret_1[0] == ret_source_wise_1[0]
+    assert (ret_1[1] == ret_source_wise_1[1]).all()
+    assert (ret_1[2] == ret_source_wise_1[2]).all()
+
+
 def test_multisource_likelihood():
     lf = UnbinnedLogLikelihood(conf_for_test(n_sources=2))
 

From 98a9a2bad29f890cc980af71d29f9b692b9c4d4c Mon Sep 17 00:00:00 2001
From: Robert Hammann <robert.hammann@outlook.com>
Date: Fri, 11 Oct 2024 14:35:06 +0200
Subject: [PATCH 08/14] base decision on dont_hash_settings

---
 blueice/likelihood.py    | 6 ++----
 blueice/source.py        | 2 --
 tests/test_likelihood.py | 2 --
 3 files changed, 2 insertions(+), 8 deletions(-)

diff --git a/blueice/likelihood.py b/blueice/likelihood.py
index f3fdaba..ba8eb76 100644
--- a/blueice/likelihood.py
+++ b/blueice/likelihood.py
@@ -115,10 +115,8 @@ def source_shape_parameters(self):
         """Dict of sources with shape parameters. source name -> dict of shape parameters."""
         source_shape_parameters = OrderedDict()
         for sn, source in zip(self.source_name_list, self.base_model.sources):
-            parameter_names = source.parameters
-            if parameter_names is None:
-                raise ValueError("The `parameters` of each source need to be specified when using `source_wise_interpolation`.")
-            shape_parameters = OrderedDict({k: v for k, v in self.shape_parameters.items() if k in parameter_names})
+            dont_hash_settings = source.config['dont_hash_settings']
+            shape_parameters = OrderedDict({k: v for k, v in self.shape_parameters.items() if k not in dont_hash_settings})
             if shape_parameters:
                 source_shape_parameters[sn] = shape_parameters
         return source_shape_parameters
diff --git a/blueice/source.py b/blueice/source.py
index 9bcbeb9..145a260 100644
--- a/blueice/source.py
+++ b/blueice/source.py
@@ -86,8 +86,6 @@ def __init__(self, config, *args, **kwargs):
         self.name = c['name']
         del c['name']
 
-        self.parameters = c.pop('parameters', None)
-
         # events_per_day and fraction_in_range may be modified / set properly for the first time later (see comments
         # in 'defaults' above)
         if hasattr(self, 'events_per_day'):
diff --git a/tests/test_likelihood.py b/tests/test_likelihood.py
index 20dbf9b..d67d79c 100644
--- a/tests/test_likelihood.py
+++ b/tests/test_likelihood.py
@@ -106,8 +106,6 @@ def test_source_wise_interpolation():
     ret_1 = lf(full_output=True, mu=1)
 
     config["source_wise_interpolation"] = True
-    # We need to specify all parameters used by each source
-    config["sources"][0]["parameters"] = ["mu", "sigma", "strlen_multiplier", "some_multiplier", "s0_rate_multiplier"]
     lf_source_wise = UnbinnedLogLikelihood(config)
     lf_source_wise.add_shape_parameter("mu", anchors={-2:-2, 0:0, 2:2})
     lf_source_wise.prepare()

From 7a5d7979c21adff772570a10247eda486b56041e Mon Sep 17 00:00:00 2001
From: Robert Hammann <robert.hammann@outlook.com>
Date: Fri, 11 Oct 2024 14:47:11 +0200
Subject: [PATCH 09/14] add efficiecnies back to source shape parameters

---
 blueice/likelihood.py | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/blueice/likelihood.py b/blueice/likelihood.py
index ba8eb76..efe9ef9 100644
--- a/blueice/likelihood.py
+++ b/blueice/likelihood.py
@@ -114,9 +114,17 @@ def __init__(self, pdf_base_config, likelihood_config=None, **kwargs):
     def source_shape_parameters(self):
         """Dict of sources with shape parameters. source name -> dict of shape parameters."""
         source_shape_parameters = OrderedDict()
-        for sn, source in zip(self.source_name_list, self.base_model.sources):
-            dont_hash_settings = source.config['dont_hash_settings']
-            shape_parameters = OrderedDict({k: v for k, v in self.shape_parameters.items() if k not in dont_hash_settings})
+        for sn, source, apply_eff, eff_name in zip(
+                self.source_name_list,
+                self.base_model.sources,
+                self.source_apply_efficiency,
+                self.source_efficiency_names
+                ):
+            ignore_parameters = source.config['dont_hash_settings']
+            # The efficiecny parameter doesn't need to be hashed but it needs to be passed to the morpher
+            if apply_eff:
+                ignore_parameters.pop(eff_name, None)
+            shape_parameters = OrderedDict({k: v for k, v in self.shape_parameters.items() if k not in ignore_parameters})
             if shape_parameters:
                 source_shape_parameters[sn] = shape_parameters
         return source_shape_parameters

From ff49044fa6e389aac7dd82e3a827cd2ce5e66bb3 Mon Sep 17 00:00:00 2001
From: Robert Hammann <robert.hammann@outlook.com>
Date: Fri, 11 Oct 2024 07:54:16 -0500
Subject: [PATCH 10/14] fix ignore parameters

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

diff --git a/blueice/likelihood.py b/blueice/likelihood.py
index efe9ef9..b4a81d1 100644
--- a/blueice/likelihood.py
+++ b/blueice/likelihood.py
@@ -120,10 +120,10 @@ def source_shape_parameters(self):
                 self.source_apply_efficiency,
                 self.source_efficiency_names
                 ):
-            ignore_parameters = source.config['dont_hash_settings']
+            ignore_parameters = set(source.config['dont_hash_settings'])
             # The efficiecny parameter doesn't need to be hashed but it needs to be passed to the morpher
             if apply_eff:
-                ignore_parameters.pop(eff_name, None)
+                ignore_parameters.discard(eff_name)
             shape_parameters = OrderedDict({k: v for k, v in self.shape_parameters.items() if k not in ignore_parameters})
             if shape_parameters:
                 source_shape_parameters[sn] = shape_parameters

From 0e99fa22a56ae04cf3a25fe80d564a5e1f357b6e Mon Sep 17 00:00:00 2001
From: Robert Hammann <robert.hammann@outlook.com>
Date: Thu, 17 Oct 2024 06:41:42 -0500
Subject: [PATCH 11/14] clean up ps_interpolator

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

diff --git a/blueice/likelihood.py b/blueice/likelihood.py
index b4a81d1..e948da8 100644
--- a/blueice/likelihood.py
+++ b/blueice/likelihood.py
@@ -528,27 +528,26 @@ def set_data(self, d):
         LogLikelihoodBase.set_data(self, d)
         if len(self.shape_parameters):
             if self.source_wise_interpolation:
-                ps_interpolators = OrderedDict()
+                self.ps_interpolators = OrderedDict()
                 for sn, base_source in zip(self.source_name_list, self.base_model.sources):
                     if sn in self.source_morphers:
-                        ps_interpolators[sn] = self.source_morphers[sn].make_interpolator(
+                        self.ps_interpolators[sn] = self.source_morphers[sn].make_interpolator(
                             f=lambda s: s.pdf(*self.base_model.to_analysis_dimensions(d)),
                             extra_dims=[len(d)],
                             anchor_models=self.anchor_sources[sn])
                     else:
-                        ps_interpolators[sn] = base_source.pdf(*self.base_model.to_analysis_dimensions(d))
+                        self.ps_interpolators[sn] = base_source.pdf(*self.base_model.to_analysis_dimensions(d))
 
                 def ps_interpolator(*args):
                     # take zs, convert to values for each source's interpolator call the respective interpolator
-                    ps = []
-                    for sn in self.source_name_list:
+                    ps = np.zeros((len(self.source_name_list), len(d)))
+                    for i, (sn, ps_interpolator) in enumerate(self.ps_interpolators.items()):
                         if sn in self.source_shape_parameters:
-                            shape_indices = self._get_shape_indices(sn)
-                            these_args = [args[0][i] for i in shape_indices]
-                            ps.append(ps_interpolators[sn](np.asarray(these_args)))
+                            these_args = np.asarray([args[0][j] for j in self._get_shape_indices(sn)])
+                            ps[i] = ps_interpolator(these_args)
                         else:
-                            ps.append(ps_interpolators[sn])
-                    return np.array(ps)
+                            ps[i] = ps_interpolator
+                    return ps
                 self.ps_interpolator = ps_interpolator
             else:
                 self.ps_interpolator = self.morpher.make_interpolator(

From 47e3ba13b5b4058ba762befaf8f8c6e385039aa8 Mon Sep 17 00:00:00 2001
From: Robert Hammann <robert.hammann@outlook.com>
Date: Tue, 29 Oct 2024 13:42:51 +0900
Subject: [PATCH 12/14] add create_source_morpher function for readability

---
 blueice/likelihood.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/blueice/likelihood.py b/blueice/likelihood.py
index e948da8..616d62a 100644
--- a/blueice/likelihood.py
+++ b/blueice/likelihood.py
@@ -152,11 +152,14 @@ def prepare(self, n_cores=1, ipp_client=None):
             if self.source_wise_interpolation:
                 # Create morphers for each source individually
                 self.source_morphers = OrderedDict()
-                for sn, shape_parameters in self.source_shape_parameters.items():
-                    self.source_morphers[sn] = MORPHERS[self.config['morpher']](
+
+                def create_source_morpher(shape_parameters):
+                    return MORPHERS[self.config['morpher']](
                         self.config.get('morpher_config', {}),
-                        shape_parameters
-                    )
+                        shape_parameters)
+
+                for sn, shape_parameters in self.source_shape_parameters.items():
+                    self.source_morphers[sn] = create_source_morpher(shape_parameters)
                 zs_list = set()
                 for source_name, morpher in self.source_morphers.items():
                     anchor_points = morpher.get_anchor_points(bounds=None)

From 9b6cad4ec7c2ac9b4399380c8a5e988d947640dd Mon Sep 17 00:00:00 2001
From: Robert Hammann <robert.hammann@outlook.com>
Date: Fri, 8 Nov 2024 15:35:14 +0100
Subject: [PATCH 13/14] fix unittest

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

diff --git a/tests/test_model.py b/tests/test_model.py
index 04d25c3..a59cdf2 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -5,9 +5,11 @@ def test_rates():
     m = Model(conf_for_test(n_sources=1))
     np.testing.assert_array_equal(m.expected_events(), np.array([1000]))
 
-    m.config['livetime_days'] = 2
+    for source in m.sources:
+        source.config['livetime_days'] = 2
     np.testing.assert_array_equal(m.expected_events(), np.array([2000]))
-    m.config['livetime_days'] = 1
+    for source in m.sources:
+        source.config['livetime_days'] = 1
 
     m.sources[0].fraction_in_range = 0.5
     np.testing.assert_array_equal(m.expected_events(), np.array([500]))

From 17b9abd6137ca1256f39f6b7d6f9f449c8a9e51b Mon Sep 17 00:00:00 2001
From: Robert Hammann <robert.hammann@outlook.com>
Date: Fri, 8 Nov 2024 15:41:19 +0100
Subject: [PATCH 14/14] improve readability

---
 blueice/likelihood.py | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/blueice/likelihood.py b/blueice/likelihood.py
index 616d62a..b938364 100644
--- a/blueice/likelihood.py
+++ b/blueice/likelihood.py
@@ -588,16 +588,16 @@ def prepare(self, *args):
         if len(self.shape_parameters):
             if self.source_wise_interpolation:
                 raise NotImplementedError("Source-wise interpolation not implemented for binned likelihoods")
-            else:
-                self.ps_interpolator = self.morpher.make_interpolator(
-                    f=lambda m: m.pmf_grids()[0],
-                    extra_dims=list(self.ps.shape),
-                    anchor_models=self.anchor_models)
 
-                if self.model_statistical_uncertainty_handling is not None:
-                    self.n_model_events_interpolator = self.morpher.make_interpolator(f=lambda m: m.pmf_grids()[1],
-                                                                                    extra_dims=list(self.ps.shape),
-                                                                                    anchor_models=self.anchor_models)
+            self.ps_interpolator = self.morpher.make_interpolator(
+                f=lambda m: m.pmf_grids()[0],
+                extra_dims=list(self.ps.shape),
+                anchor_models=self.anchor_models)
+
+            if self.model_statistical_uncertainty_handling is not None:
+                self.n_model_events_interpolator = self.morpher.make_interpolator(f=lambda m: m.pmf_grids()[1],
+                                                                                extra_dims=list(self.ps.shape),
+                                                                                anchor_models=self.anchor_models)
 
     @inherit_docstring_from(LogLikelihoodBase)
     def set_data(self, d):