From ec0bc272c60c086d0a8bd6ce86a830d196659cd4 Mon Sep 17 00:00:00 2001
From: dachengx <dx2227@columbia.edu>
Date: Sun, 30 Jul 2023 20:45:23 -0500
Subject: [PATCH 01/21] Simplify TemplateSource

---
 alea/template_source.py | 554 ++++++++++++++++------------------------
 1 file changed, 221 insertions(+), 333 deletions(-)

diff --git a/alea/template_source.py b/alea/template_source.py
index f7dac572..b2dcc2de 100644
--- a/alea/template_source.py
+++ b/alea/template_source.py
@@ -2,48 +2,93 @@
 import warnings
 import logging
 
-import blueice
 import numpy as np
 from scipy.interpolate import interp1d
+from blueice import HistogramPdfSource
 from inference_interface import template_to_multihist
 
 logging.basicConfig(level=logging.INFO)
 can_check_binning = True
 
 
-class TemplateSource(blueice.HistogramPdfSource):
+class TemplateSource(HistogramPdfSource):
     """
-    A source that constructs a  from a template histogram
-    :param templatename: root file to open.
+    A source that constructs a from a template histogram.
+    The parameters are set in self.config.
+    'templatename', 'histname', 'analysis_space' must be in self.config
+
+    :ivar config: The configuration of the source.
+    :vartype config: dict
+    :ivar dtype: The data type of the source.
+    :vartype dtype: list
+    :ivar _bin_volumes: The bin volumes of the source.
+    :vartype _bin_volumes: numpy.ndarray
+    :ivar _n_events_histogram: The histogram of the number of events of the source.
+    :vartype _n_events_histogram: multihist.MultiHistBase
+    :ivar events_per_day: The number of events per day of the source.
+    :vartype events_per_day: float
+    :ivar _pdf_histogram: The histogram of the probability density function of the source.
+    :vartype _pdf_histogram: multihist.MultiHistBase
+
+    :param templatename: Hdf5 file to open.
     :param histname: Histogram name.
-    :param named_parameters: list of config setting names to pass to .format on histname and filename.
+    :param named_parameters:
+        List of config setting names to pass to .format on histname and filename.
+    :type named_parameters: list
+    :param normalise_template: Normalise the template histogram.
+    :type normalise_template: bool
     :param in_events_per_bin:
-    if True, histogram is taken to be in events per day / bin,
-    if False or absent, taken to be events per day / bin volume
-    :param histogram_multiplier: multiply histogram by this number
-    :param log10_bins: List of axis numbers.
-    If True, bin edges on this axis in the root file are log10() of the actual bin edges.
+        If True, histogram is in events per day / bin.
+        If False or absent, histogram is already pdf.
+    :param histogram_scale_factor: Multiply histogram by this number
+    :param convert_to_uniform: Convert the histogram to a uniform per bin distribution.
+    :param log10_bins:
+        List of axis numbers.
+        If True, bin edges on this axis in the hdf5 file are log10() of the actual bin edges.
     """
 
     def build_histogram(self):
-        format_dict = {
-            k: self.config[k]
-            for k in self.config.get('named_parameters', [])
-        }
-        logging.debug("loading template")
-        templatename = self.config['templatename'].format(**format_dict)
-        histname = self.config['histname'].format(**format_dict)
-
-        slice_args = self.config.get("slice_args", {})
-        if type(slice_args) is dict:
-            slice_args = [slice_args]
-
+        templatename = self.config['templatename'].format(**self.format_named_parameters)
+        histname = self.config['histname'].format(**self.format_named_parameters)
         h = template_to_multihist(templatename, histname)
+        if np.min(h.histogram) < 0:
+            raise AssertionError(
+                f"There are bins for source {templatename} with negative entries.")
 
         if self.config.get('normalise_template', False):
             h /= h.n
-            if not self.config.get('in_events_per_bin', True):
-                h.histogram *= h.bin_volumes()
+        if not self.config.get('in_events_per_bin', True):
+            h.histogram *= h.bin_volumes()
+
+        h = self.apply_slice_args(h)
+
+        # Fix the bin sizes
+        if can_check_binning:
+            self._check_binning(h, f"{histname:s} in hdf5 file {templatename:s}")
+
+        self.set_dtype()
+        self.set_pdf_histogram(h)
+
+    @property
+    def format_named_parameters(self):
+        """Get the named parameters in the config to dictionary format."""
+        format_named_parameters = {
+            k: self.config[k] for k in self.config.get('named_parameters', [])}
+        return format_named_parameters
+
+    def apply_slice_args(self, h, slice_args=None):
+        """
+        Apply slice arguments to the histogram.
+
+        :param h: The histogram to apply the slice arguments to.
+        :type h: multihist.MultiHistBase
+        :param slice_args: The slice arguments to apply.
+        :type slice_args: dict
+        """
+        if slice_args is None:
+            slice_args = self.config.get("slice_args", {})
+            if type(slice_args) is dict:
+                slice_args = [slice_args]
 
         for sa in slice_args:
             slice_axis = sa.get("slice_axis", None)
@@ -57,7 +102,9 @@ def build_histogram(self):
                 slice_axis_limits = sa.get(
                     "slice_axis_limits", [bes[0], bes[-1]])
                 if sum_axis:
-                    logging.debug(f"Slice and sum over axis {slice_axis} from {slice_axis_limits[0]} to {slice_axis_limits[1]}")
+                    logging.debug(
+                        f"Slice and sum over axis {slice_axis} from {slice_axis_limits[0]} "
+                        f"to {slice_axis_limits[1]}")
                     axis_names = h.axis_names_without(slice_axis)
                     h = h.slicesum(
                         axis=slice_axis,
@@ -66,10 +113,13 @@ def build_histogram(self):
                     h.axis_names = axis_names
                 else:
                     logging.debug(f"Normalization before slicing: {h.n}.")
-                    logging.debug(f"Slice over axis {slice_axis} from {slice_axis_limits[0]} to {slice_axis_limits[1]}")
-                    h = h.slice(axis=slice_axis,
-                                start=slice_axis_limits[0],
-                                stop=slice_axis_limits[1])
+                    logging.debug(
+                        f"Slice over axis {slice_axis} from {slice_axis_limits[0]} "
+                        f"to {slice_axis_limits[1]}")
+                    h = h.slice(
+                        axis=slice_axis,
+                        start=slice_axis_limits[0],
+                        stop=slice_axis_limits[1])
                     logging.debug(f"Normalization after slicing: {h.n}.")
 
             if collapse_axis is not None:
@@ -77,163 +127,156 @@ def build_histogram(self):
                     raise ValueError(
                         "To collapse you must supply collapse_slices")
                 h = h.collapse_bins(collapse_slices, axis=collapse_axis)
+        return h
+
+    def _check_binning(self, h, histogram_info):
+        """
+        Check if the histogram's bin edges are the same to analysis_space.
 
+        :param h: The histogram to check.
+        :type h: multihist.MultiHistBase
+        :param histogram_info: Information of the histogram.
+        :type histogram_info: str
+        """
+        # Deal with people who have log10'd their bins
+        for axis_i in self.config.get('log10_bins', []):
+            h.bin_edges[axis_i] = 10 ** h.bin_edges[axis_i]
+
+        # Check if the histogram bin edges are correct
+        analysis_space = self.config['analysis_space']
+        for axis_i, (_, expected_bin_edges) in enumerate(analysis_space):
+            expected_bin_edges = np.array(expected_bin_edges)
+            seen_bin_edges = h.bin_edges[axis_i]
+            # If 1D, hist1d returns bin_edges straight, not as list
+            if len(analysis_space) == 1:
+                seen_bin_edges = h.bin_edges
+            logging.debug("axis_i: " + str(axis_i))
+            logging.debug("expected_bin_edges: " + str(expected_bin_edges))
+            logging.debug("seen_bin_edges: " + str(seen_bin_edges))
+            logging.debug("h.bin_edges type" + str(h.bin_edges))
+            if len(seen_bin_edges) != len(expected_bin_edges):
+                raise ValueError(
+                    f"Axis {axis_i:d} of histogram {histogram_info} "
+                    f"has {len(seen_bin_edges)} bin edges, but expected {expected_bin_edges}.")
+            try:
+                np.testing.assert_almost_equal(
+                    seen_bin_edges,
+                    expected_bin_edges,
+                    decimal=2)
+            except AssertionError:
+                warnings.warn(
+                    f"Axis {axis_i:d} of histogram {histogram_info} "
+                    f"has bin edges {seen_bin_edges}, but expected {expected_bin_edges}. "
+                    "Since length matches, setting it expected values...")
+                h.bin_edges[axis_i] = expected_bin_edges
+
+    def set_dtype(self):
+        """Set the data type of the source."""
         self.dtype = []
         for n, _ in self.config['analysis_space']:
             self.dtype.append((n, float))
         self.dtype.append(('source', int))
 
-        # Fix the bin sizes
-        if can_check_binning:
-            # Deal with people who have log10'd their bins
-            for axis_i in self.config.get('log10_bins', []):
-                h.bin_edges[axis_i] = 10**h.bin_edges[axis_i]
-
-            # Check if the histogram bin edges are correct
-            for axis_i, (_, expected_bin_edges) in enumerate(
-                    self.config['analysis_space']):
-                expected_bin_edges = np.array(expected_bin_edges)
-                seen_bin_edges = h.bin_edges[axis_i]
-                # If 1D, hist1d returns bin_edges straight, not as list
-                if len(self.config['analysis_space']) == 1:
-                    seen_bin_edges = h.bin_edges
-                logging.debug("axis_i: " + str(axis_i))
-                logging.debug("expected_bin_edges: " + str(expected_bin_edges))
-                logging.debug("seen_bin_edges: " + str(seen_bin_edges))
-                logging.debug("h.bin_edges type" + str(h.bin_edges))
-                if len(seen_bin_edges) != len(expected_bin_edges):
-                    raise ValueError(
-                        "Axis %d of histogram %s in root file %s has %d bin edges, but expected %d"
-                        % (
-                            axis_i, histname, templatename, len(seen_bin_edges),
-                            len(expected_bin_edges)))
-                try:
-                    np.testing.assert_almost_equal(
-                        seen_bin_edges,
-                        expected_bin_edges,
-                        decimal=2)
-                except AssertionError:
-                    warnings.warn(
-                        "Axis %d of histogram %s in root file %s has bin edges %s, "
-                        "but expected %s. Since length matches, setting it expected values..."
-                        % (
-                            axis_i, histname, templatename, seen_bin_edges,
-                            expected_bin_edges))
-                    h.bin_edges[axis_i] = expected_bin_edges
-
+    def set_pdf_histogram(self, h):
+        """Set the histogram of the probability density function of the source."""
         self._bin_volumes = h.bin_volumes()  # TODO: make alias
         # Shouldn't be in HistogramSource... anyway
         self._n_events_histogram = h.similar_blank_histogram()
 
-        h *= self.config.get('histogram_scale_factor', 1)
-        logging.debug(f"Multiplying histogram with histogram_scale_factor {self.config.get('histogram_scale_factor', 1)}. Histogram is now normalised to {h.n}.")
+        histogram_scale_factor = self.config.get('histogram_scale_factor', 1)
+        h *= histogram_scale_factor
+        logging.debug(
+            f"Multiplying histogram with histogram_scale_factor {histogram_scale_factor}. "
+            f"Histogram is now normalised to {h.n}.")
 
-        # Convert h to density...
-        if self.config.get('in_events_per_bin'):
-            h.histogram /= h.bin_volumes()
-        self.events_per_day = (h.histogram * self._bin_volumes).sum()
+        self.events_per_day = h.n
         logging.debug(f"events_per_day: " + str(self.events_per_day))
 
+        # Convert h to density...
+        if self.config.get('in_events_per_bin', True):
+            h.histogram /= self._bin_volumes
+
         # ... and finally to probability density
-        if 0 < self.events_per_day:
+        if self.events_per_day > 0:
             h.histogram /= self.events_per_day
+        else:
+            raise ValueError("Sum of histogram is zero.")
 
         if self.config.get('convert_to_uniform', False):
-            h.histogram = 1./self._bin_volumes
+            h.histogram = 1. / self._bin_volumes
 
         self._pdf_histogram = h
         logging.debug(f"Setting _pdf_histogram normalised to {h.n}.")
 
-        if np.min(self._pdf_histogram.histogram) < 0:
-            raise AssertionError(
-                f"There are bins for source {templatename} with negative entries."
-                    )
-
     def simulate(self, n_events):
-        dtype = []
-        for n, _ in self.config['analysis_space']:
-            dtype.append((n, float))
-        dtype.append(('source', int))
-        ret = np.zeros(n_events, dtype=dtype)
+        """
+        Simulate events from the source.
+
+        :param n_events: The number of events to simulate.
+        :type n_events: int
+        :return: The simulated events.
+        :rtype: numpy.ndarray
+        """
+        ret = np.zeros(n_events, dtype=self.dtype)
         # t = self._pdf_histogram.get_random(n_events)
         h = self._pdf_histogram * self._bin_volumes
         t = h.get_random(n_events)
-        for i, (n, _) in enumerate(self.config.get('analysis_space')):
+        for i, (n, _) in enumerate(self.config['analysis_space']):
             ret[n] = t[:, i]
         return ret
 
 
-class CombinedSource(blueice.HistogramPdfSource):
+class CombinedSource(TemplateSource, HistogramPdfSource):
     """
     Source that inherits structure from TH2DSource by Jelle,
     but takes in lists of histogram names.
-    :param weights: weights
-    :param histnames: list of filenames containing the histograms
-    :param templatenames: list of names of histograms within the hdf5 files
-    :param named_parameters : list of names of weights to be applied to histograms.
-    Must be 1 shorter than histnames, templatenames
-    :param histogram_parameters: names of parameters that should be put in the hdf5/histogram names,
+    The first histogram is the base histogram, and the rest are added to it with weights.
+
+    :param weights: Weights of the 2nd to the last histograms.
+    :param histnames: List of filenames containing the histograms.
+    :param templatenames: List of names of histograms within the hdf5 files.
+    :param histogram_multiplier: Absolute rate of the combined template before multihist slicing.
     """
 
     def build_histogram(self):
-        weight_names = self.config.get("weight_names")
+        if not self.config.get('in_events_per_bin', True):
+            raise ValueError(
+                "CombinedSource does not support in_events_per_bin=False")
+        if "weight_names" not in self.config:
+            raise ValueError("weight_names must be specified")
         weights = [
-            self.config.get(weight_name, 0) for weight_name in weight_names
+            self.config.get(weight_name, 0) for weight_name in self.config["weight_names"]
         ]
         histograms = []
-        format_dict = {
-            k: self.config[k]
-            for k in self.config.get('named_parameters', [])
-        }
         histnames = self.config['histnames']
         templatenames = self.config['templatenames']
+        if len(histnames) != len(templatenames):
+            raise ValueError(
+                "histnames and templatenames must be the same length")
+        if len(weights) + 1 != len(histograms):
+            raise ValueError(
+                "weights must be 1 shorter than histnames, templatenames")
 
         slice_args = self.config.get("slice_args", {})
         if type(slice_args) is dict:
             slice_args = [slice_args]
-        if type(slice_args[0]) is dict:
-            slice_argss = []
-            for n in histnames:
-                slice_argss.append(slice_args)
+        if all(type(sa) is dict for sa in slice_args):
+            # Recognize as a single slice_args for all histograms
+            slice_argss = [slice_args] * len(histnames)
         else:
+            # Recognize as a list of slice_args for each histogram
             slice_argss = slice_args
 
         slice_fractions = []
-
         for histname, templatename, slice_args in zip(
                 histnames, templatenames, slice_argss):
-            templatename = templatename.format(**format_dict)
-            histname = histname.format(**format_dict)
+            templatename = templatename.format(**self.format_named_parameters)
+            histname = histname.format(**self.format_named_parameters)
             h = template_to_multihist(templatename, histname)
             slice_fraction = 1. / h.n
             if h.n == 0.:
                 slice_fraction = 0.
-            for sa in slice_args:
-                slice_axis = sa.get("slice_axis", None)
-                # Decide if you wish to sum the histogram into lower dimensions or
-                sum_axis = sa.get("sum_axis", False)
-
-                slice_axis_limits = sa.get("slice_axis_limits", [0, 0])
-                collapse_axis = sa.get('collapse_axis', None)
-                collapse_slices = sa.get('collapse_slices', None)
-                if (slice_axis is not None):
-                    if sum_axis:
-                        axis_names = h.axis_names_without(slice_axis)
-                        h = h.slicesum(
-                            axis=slice_axis,
-                            start=slice_axis_limits[0],
-                            stop=slice_axis_limits[1])
-                        h.axis_names = axis_names
-                    else:
-                        h = h.slice(axis=slice_axis,
-                                    start=slice_axis_limits[0],
-                                    stop=slice_axis_limits[1])
-
-                if collapse_axis is not None:
-                    if collapse_slices is None:
-                        raise ValueError(
-                            "To collapse you must supply collapse_slices")
-                    h = h.collapse_bins(collapse_slices, axis=collapse_axis)
+            h = self.apply_slice_args(h, slice_args)
             slice_fraction *= h.n
             slice_fractions.append(slice_fraction)
 
@@ -243,18 +286,12 @@ def build_histogram(self):
             else:
                 h /= h.n
             histograms.append(h)
-
-        assert len(weights) + 1 == len(histograms)
-
-        # Here a common normalisation is performed (normalisation to bin widths comes later)
-        # for histogram in histograms:
-        #     histogram/=histogram.n
+        logging.debug("slice fractions are", slice_fractions)
 
         for i in range(len(weights)):
-
             h_comp = histograms[0].similar_blank_histogram()
             h = histograms[i + 1]
-            base_part_norm = 0.
+            # base_part_norm = 0.
             # h_centers = h_comp.bin_centers()
             # h_inds = [range(len(hc)) for hc in h_centers]
             # for inds in product(*h_inds):
@@ -262,114 +299,56 @@ def build_histogram(self):
             #     inds_comp = h_comp.get_bin_indices(bincs)
             #     base_part_norm += histograms[0][inds]
             #     h_comp[inds] = h[inds_comp]
-            # print("i", i, "h_comp", h_comp.histogram.shape, h_comp.n)
             # h_comp *= base_part_norm
             hslices = []
             for j, bincs in enumerate(h.bin_centers()):
                 hsliced = h_comp.get_axis_bin_index(bincs[0], j)
                 hsliceu = h_comp.get_axis_bin_index(bincs[-1], j) + 1
                 hslices.append(slice(hsliced, hsliceu))
-            base_part_norm = histograms[0][hslices].sum()
-            h_comp[hslices] += h.histogram  # TODO check here what norm I want.
+            # TODO: check here what norm I want.
+            h_comp[hslices] += h.histogram
             histograms[0] += h_comp * weights[i]
+        h = histograms[0]
 
         # Set pdf values that are below 0 to zero:
-        histograms[0].histogram[np.isinf(histograms[0].histogram)] = 0.
-        histograms[0].histogram[np.isnan(histograms[0].histogram)] = 0.
-        histograms[0].histogram[histograms[0].histogram < 0] = 0.
-
-        logging.debug("composite, histograms.n", histograms[0].n)
-        if 0 < histograms[0].n:
-            h = histograms[0] / histograms[0].n
-        else:
-            h = histograms[0]
-            h.histogram = np.ones(h.histogram.shape)
-        self._bin_volumes = h.bin_volumes()  # TODO: make alias
-        # Shouldn't be in HistogramSource... anyway
-        self._n_events_histogram = h.similar_blank_histogram()
-
-        # Fix the bin sizes
-        if can_check_binning:
-            # Deal with people who have log10'd their bins
-            for axis_i in self.config.get('log10_bins', []):
-                h.bin_edges[axis_i] = 10**h.bin_edges[axis_i]
-
-            # Check if the histogram bin edges are correct
-            for axis_i, (_, expected_bin_edges) in enumerate(
-                    self.config['analysis_space']):
-                expected_bin_edges = np.array(expected_bin_edges)
-                seen_bin_edges = h.bin_edges[axis_i]
-                logging.debug("in axis_i check", axis_i)
-                logging.debug("expected", expected_bin_edges)
-                logging.debug("seen", seen_bin_edges)
-                logging.debug("h.bin_edges type", type(h.bin_edges))
-                if len(seen_bin_edges) != len(expected_bin_edges):
-                    raise ValueError(
-                        "Axis %d of histogram %s in hdf5 file %s has %d bin edges, but expected %d"
-                        % (
-                            axis_i, histname, templatename, len(seen_bin_edges),
-                            len(expected_bin_edges)))
-                try:
-                    np.testing.assert_almost_equal(
-                        seen_bin_edges,
-                        expected_bin_edges, decimal=4)
-                except AssertionError:
-                    warnings.warn(
-                        "Axis %d of histogram %s in hdf5 file %s has bin edges %s, "
-                        "but expected %s. Since length matches, setting it expected values..."
-                        % (
-                            axis_i, histname, templatename, seen_bin_edges,
-                            expected_bin_edges))
-                    h.bin_edges[axis_i] = expected_bin_edges
+        if np.min(h.histogram) < 0:
+            raise AssertionError(
+                f"There are bins for source {templatename} with negative entries.")
 
+        logging.info("Normalising combined template, the absolute rate is defined in histogram_multiplier")
+        h = h / h.n
         h *= self.config.get('histogram_multiplier', 1)
-        logging.debug("slice fractions are", slice_fractions)
+        logging.info("Applying slice fraction to combined template")
         h *= slice_fractions[0]
 
-        # Convert h to density...
-        logging.debug("hist sum", np.sum(h.histogram))
-        if self.config.get('in_events_per_bin'):
-            h.histogram /= h.bin_volumes()
-        self.events_per_day = (h.histogram * h.bin_volumes()).sum()
-        logging.debug("init events per day", self.events_per_day)
-
-        # ... and finally to probability density
-        if 0 < self.events_per_day:
-            h.histogram /= self.events_per_day
-        else:
-            h.histogram = np.ones(h.histogram.shape)
-        logging.debug("_pdf_histogram sum", h.n)
-        self._pdf_histogram = h
+        # Fix the bin sizes
+        if can_check_binning:
+            self._check_binning(h, f"combined template {histnames} in hdf5 file {templatenames}")
 
-    def simulate(self, n_events):
-        dtype = []
-        for n, _ in self.config['analysis_space']:
-            dtype.append((n, float))
-        dtype.append(('source', int))
-        ret = np.zeros(n_events, dtype=dtype)
-        # t = self._pdf_histogram.get_random(n_events)
-        h = self._pdf_histogram * self._bin_volumes
-        t = h.get_random(n_events)
-        for i, (n, _) in enumerate(self.config.get('analysis_space')):
-            ret[n] = t[:, i]
-        return ret
+        self.set_dtype()
+        self.set_pdf_histogram(h)
 
 
-class SpectrumTemplateSource(blueice.HistogramPdfSource):
+class SpectrumTemplateSource(TemplateSource, HistogramPdfSource):
     """
-    :param spectrum_name: name of bbf json-like spectrum _OR_ function that can be called
-    templatename #3D histogram (Etrue,S1,S2) to open
-    :param histname: histogram name
-    :param named_parameters: list of config settings to pass to .format on histname and filename
+    Reweighted template source by energy spectrum.
+    The first axis of the template is assumed to be energy.
+
+    :param spectrum_name:
+        Name of bbf json-like spectrum _OR_ function that can be called
+        templatename #3D histogram (Etrue, S1, S2) to open
     """
 
     @staticmethod
-    def _get_json_spectrum(fn):
+    def _get_json_spectrum(filename):
         """
         Translates bbf-style JSON files to spectra.
         units are keV and /kev*day*kg
+
+        :param filename: Name of the JSON file.
+        :type filename: str
         """
-        contents = json.load(open(fn, "r"))
+        contents = json.load(open(filename, "r"))
         logging.debug(contents["description"])
         esyst = contents["coordinate_system"][0][1]
         ret = interp1d(
@@ -378,126 +357,35 @@ def _get_json_spectrum(fn):
         return ret
 
     def build_histogram(self):
-        logging.debug("building a hist")
-        format_dict = {
-            k: self.config[k]
-            for k in self.config.get('named_parameters', [])
-        }
-        templatename = self.config['templatename'].format(**format_dict)
-        histname = self.config['histname'].format(**format_dict)
+        templatename = self.config['templatename'].format(**self.format_named_parameters)
+        histname = self.config['histname'].format(**self.format_named_parameters)
+        h = template_to_multihist(templatename, histname)
 
         spectrum = self.config["spectrum"]
         if type(spectrum) is str:
-            spectrum = self._get_json_spectrum(spectrum.format(**format_dict))
+            spectrum = self._get_json_spectrum(spectrum.format(**self.format_named_parameters))
 
-        slice_args = self.config.get("slice_args", {})
-        if type(slice_args) is dict:
-            slice_args = [slice_args]
-
-        h = template_to_multihist(templatename, histname)
-
-        # Perform E-scaling
-        ebins = h.bin_edges[0]
-        ecenters = 0.5 * (ebins[1::] + ebins[0:-1])
-        ediffs = (ebins[1::] - ebins[0:-1])
+        # Perform E-scaling, assume first axis is energy
+        ecenters = h.bin_centers[0]
+        ediffs = h.bin_volumes[0]
         h.histogram = h.histogram * (spectrum(ecenters) * ediffs)[:, None, None]
         h = h.sum(axis=0)  # Remove energy-axis
         logging.debug("source", "mu ", h.n)
 
-        for sa in slice_args:
-            slice_axis = sa.get("slice_axis", None)
-            # Decide if you wish to sum the histogram into lower dimensions or
-            sum_axis = sa.get("sum_axis", False)
-
-            slice_axis_limits = sa.get("slice_axis_limits", [0, 0])
-            collapse_axis = sa.get('collapse_axis', None)
-            collapse_slices = sa.get('collapse_slices', None)
-
-            if (slice_axis is not None):
-                if sum_axis:
-                    h = h.slicesum(
-                        axis=slice_axis,
-                        start=slice_axis_limits[0],
-                        stop=slice_axis_limits[1])
-                else:
-                    h = h.slice(
-                        axis=slice_axis,
-                        start=slice_axis_limits[0],
-                        stop=slice_axis_limits[1])
-
-            if collapse_axis is not None:
-                if collapse_slices is None:
-                    raise ValueError(
-                        "To collapse you must supply collapse_slices")
-                h = h.collapse_bins(collapse_slices, axis=collapse_axis)
-
-        self.dtype = []
-        for n, _ in self.config['analysis_space']:
-            self.dtype.append((n, float))
-        self.dtype.append(('source', int))
-
-        # Fix the bin sizes
-        if can_check_binning:
-            # Deal with people who have log10'd their bins
-            for axis_i in self.config.get('log10_bins', []):
-                h.bin_edges[axis_i] = 10 ** h.bin_edges[axis_i]
-
-            # Check if the histogram bin edges are correct
-            for axis_i, (_, expected_bin_edges) in enumerate(
-                    self.config['analysis_space']):
-                expected_bin_edges = np.array(expected_bin_edges)
-                seen_bin_edges = h.bin_edges[axis_i]
-                # If 1D, hist1d returns bin_edges straight, not as list
-                if len(self.config['analysis_space']) == 1:
-                    seen_bin_edges = h.bin_edges
-                if len(seen_bin_edges) != len(expected_bin_edges):
-                    raise ValueError(
-                        "Axis %d of histogram %s in root file %s has %d bin edges, but expected %d"
-                        % (
-                            axis_i, histname, templatename, len(seen_bin_edges),
-                            len(expected_bin_edges)))
-                try:
-                    np.testing.assert_almost_equal(
-                        seen_bin_edges / expected_bin_edges,
-                        np.ones_like(seen_bin_edges),
-                        decimal=4)
-                except AssertionError:
-                    warnings.warn(
-                        "Axis %d of histogram %s in root file %s has bin edges %s, "
-                        "but expected %s. Since length matches, setting it expected values..."
-                        % (
-                            axis_i, histname, templatename, seen_bin_edges,
-                            expected_bin_edges))
-                    h.bin_edges[axis_i] = expected_bin_edges
-
-        self._bin_volumes = h.bin_volumes()  # TODO: make alias
-        # Shouldn't be in HistogramSource... anyway
-        self._n_events_histogram = h.similar_blank_histogram()
+        if np.min(h.histogram) < 0:
+            raise AssertionError(
+                f"There are bins for source {templatename} with negative entries.")
 
         if self.config.get('normalise_template', False):
             h /= h.n
+        if not self.config.get('in_events_per_bin', True):
+            h.histogram *= h.bin_volumes()
 
-        h *= self.config.get('histogram_multiplier', 1)
-
-        # Convert h to density...
-        if self.config.get('in_events_per_bin'):
-            h.histogram /= h.bin_volumes()
-        self.events_per_day = (h.histogram * self._bin_volumes).sum()
-        logging.debug("source evt per day is", self.events_per_day)
+        h = self.apply_slice_args(h)
 
-        # ... and finally to probability density
-        h.histogram /= self.events_per_day
-        self._pdf_histogram = h
+        # Fix the bin sizes
+        if can_check_binning:
+            self._check_binning(h, f"reweighted {histname:s} in hdf5 file {templatename:s}")
 
-    def simulate(self, n_events):
-        dtype = []
-        for n, _ in self.config['analysis_space']:
-            dtype.append((n, float))
-        dtype.append(('source', int))
-        ret = np.zeros(n_events, dtype=dtype)
-        # t = self._pdf_histogram.get_random(n_events)
-        h = self._pdf_histogram * self._bin_volumes
-        t = h.get_random(n_events)
-        for i, (n, _) in enumerate(self.config.get('analysis_space')):
-            ret[n] = t[:, i]
-        return ret
+        self.set_dtype()
+        self.set_pdf_histogram(h)

From cd0a5942b96834e126bd978b5d12bcd0fc757da0 Mon Sep 17 00:00:00 2001
From: dachengx <dx2227@columbia.edu>
Date: Sun, 30 Jul 2023 23:46:18 -0500
Subject: [PATCH 02/21] Update settings of utils, add SR2 to demo
 CombinedSource

---
 .../unbinned_wimp_statistical_model.yaml      |  38 +++++
 alea/template_source.py                       | 148 +++++++++---------
 alea/utils.py                                 |  16 +-
 3 files changed, 125 insertions(+), 77 deletions(-)

diff --git a/alea/examples/configs/unbinned_wimp_statistical_model.yaml b/alea/examples/configs/unbinned_wimp_statistical_model.yaml
index cd4f3f37..fbaa4821 100644
--- a/alea/examples/configs/unbinned_wimp_statistical_model.yaml
+++ b/alea/examples/configs/unbinned_wimp_statistical_model.yaml
@@ -14,6 +14,11 @@ parameter_definition:
     fittable: false
     description: Livetime of SR1 in years
 
+  livetime_sr2:
+    nominal_value: 0.5
+    fittable: false
+    description: Livetime of SR1 in years
+
   wimp_rate_multiplier:
     nominal_value: 1.0
     ptype: rate
@@ -127,3 +132,36 @@ likelihood_config:
         template_filename: wimp50gev_template.h5
         apply_efficiency: True
         efficiency_name: signal_efficiency
+
+    # SR2
+    - name: sr2
+      default_source_class: alea.template_source.TemplateSource
+      likelihood_type: blueice.likelihood.UnbinnedLogLikelihood
+      analysis_space:
+        - "cs1": 'np.arange(0, 102, 2)'
+        - "cs2": 'np.geomspace(100, 100000, 51)'
+      in_events_per_bin: true
+      livetime_parameter: livetime_sr2
+      slice_args: {}
+      sources:
+      - name: er
+        class: alea.template_source.CombinedSource
+        parameters:
+          - er_rate_multiplier
+        weight_1: 0.2
+        weight_-1: 0.3
+        weight_names: [weight_1, weight_-1]
+        histnames: [er_template, er_template, er_template,]
+        template_filenames: [er_template_0.h5, er_template_1.h5, er_template_-1.h5]
+        histogram_multiplier: 200
+        histogram_scale_factor: 3
+
+      - name: wimp
+        histname: wimp_template
+        parameters:
+          - wimp_rate_multiplier
+          - wimp_mass
+          - signal_efficiency
+        template_filename: wimp50gev_template.h5
+        apply_efficiency: True
+        efficiency_name: signal_efficiency
diff --git a/alea/template_source.py b/alea/template_source.py
index b2dcc2de..f3fd2fb5 100644
--- a/alea/template_source.py
+++ b/alea/template_source.py
@@ -15,7 +15,7 @@ class TemplateSource(HistogramPdfSource):
     """
     A source that constructs a from a template histogram.
     The parameters are set in self.config.
-    'templatename', 'histname', 'analysis_space' must be in self.config
+    "templatename", "histname", "analysis_space" must be in self.config
 
     :ivar config: The configuration of the source.
     :vartype config: dict
@@ -47,17 +47,65 @@ class TemplateSource(HistogramPdfSource):
         If True, bin edges on this axis in the hdf5 file are log10() of the actual bin edges.
     """
 
+    def _check_binning(self, h, histogram_info):
+        """
+        Check if the histogram"s bin edges are the same to analysis_space.
+
+        :param h: The histogram to check.
+        :type h: multihist.MultiHistBase
+        :param histogram_info: Information of the histogram.
+        :type histogram_info: str
+        """
+        # Deal with people who have log10"d their bins
+        for axis_i in self.config.get("log10_bins", []):
+            h.bin_edges[axis_i] = 10 ** h.bin_edges[axis_i]
+
+        # Check if the histogram bin edges are correct
+        analysis_space = self.config["analysis_space"]
+        for axis_i, (_, expected_bin_edges) in enumerate(analysis_space):
+            expected_bin_edges = np.array(expected_bin_edges)
+            seen_bin_edges = h.bin_edges[axis_i]
+            # If 1D, hist1d returns bin_edges straight, not as list
+            if len(analysis_space) == 1:
+                seen_bin_edges = h.bin_edges
+            logging.debug("axis_i: " + str(axis_i))
+            logging.debug("expected_bin_edges: " + str(expected_bin_edges))
+            logging.debug("seen_bin_edges: " + str(seen_bin_edges))
+            logging.debug("h.bin_edges type" + str(h.bin_edges))
+            if len(seen_bin_edges) != len(expected_bin_edges):
+                raise ValueError(
+                    f"Axis {axis_i:d} of histogram {histogram_info} "
+                    f"has {len(seen_bin_edges)} bin edges, but expected {expected_bin_edges}.")
+            try:
+                np.testing.assert_almost_equal(
+                    seen_bin_edges,
+                    expected_bin_edges,
+                    decimal=2)
+            except AssertionError:
+                warnings.warn(
+                    f"Axis {axis_i:d} of histogram {histogram_info} "
+                    f"has bin edges {seen_bin_edges}, but expected {expected_bin_edges}. "
+                    "Since length matches, setting it expected values...")
+                h.bin_edges[axis_i] = expected_bin_edges
+
+    @property
+    def format_named_parameters(self):
+        """Get the named parameters in the config to dictionary format."""
+        format_named_parameters = {
+            k: self.config[k] for k in self.config.get("named_parameters", [])}
+        return format_named_parameters
+
     def build_histogram(self):
-        templatename = self.config['templatename'].format(**self.format_named_parameters)
-        histname = self.config['histname'].format(**self.format_named_parameters)
+        templatename = self.config["templatename"].format(**self.format_named_parameters)
+        histname = self.config["histname"].format(**self.format_named_parameters)
         h = template_to_multihist(templatename, histname)
         if np.min(h.histogram) < 0:
             raise AssertionError(
                 f"There are bins for source {templatename} with negative entries.")
 
-        if self.config.get('normalise_template', False):
+        if self.config.get("normalise_template", False):
             h /= h.n
-        if not self.config.get('in_events_per_bin', True):
+        if not self.config.get("in_events_per_bin", True):
             h.histogram *= h.bin_volumes()
 
         h = self.apply_slice_args(h)
@@ -69,13 +117,6 @@ def build_histogram(self):
         self.set_dtype()
         self.set_pdf_histogram(h)
 
-    @property
-    def format_named_parameters(self):
-        """Get the named parameters in the config to dictionary format."""
-        format_named_parameters = {
-            k: self.config[k] for k in self.config.get('named_parameters', [])}
-        return format_named_parameters
-
     def apply_slice_args(self, h, slice_args=None):
         """
         Apply slice arguments to the histogram.
@@ -94,8 +135,8 @@ def apply_slice_args(self, h, slice_args=None):
             slice_axis = sa.get("slice_axis", None)
             # Decide if you wish to sum the histogram into lower dimensions or
             sum_axis = sa.get("sum_axis", False)
-            collapse_axis = sa.get('collapse_axis', None)
-            collapse_slices = sa.get('collapse_slices', None)
+            collapse_axis = sa.get("collapse_axis", None)
+            collapse_slices = sa.get("collapse_slices", None)
             if slice_axis is not None:
                 slice_axis_number = h.get_axis_number(slice_axis)
                 bes = h.bin_edges[slice_axis_number]
@@ -129,61 +170,20 @@ def apply_slice_args(self, h, slice_args=None):
                 h = h.collapse_bins(collapse_slices, axis=collapse_axis)
         return h
 
-    def _check_binning(self, h, histogram_info):
-        """
-        Check if the histogram's bin edges are the same to analysis_space.
-
-        :param h: The histogram to check.
-        :type h: multihist.MultiHistBase
-        :param histogram_info: Information of the histogram.
-        :type histogram_info: str
-        """
-        # Deal with people who have log10'd their bins
-        for axis_i in self.config.get('log10_bins', []):
-            h.bin_edges[axis_i] = 10 ** h.bin_edges[axis_i]
-
-        # Check if the histogram bin edges are correct
-        analysis_space = self.config['analysis_space']
-        for axis_i, (_, expected_bin_edges) in enumerate(analysis_space):
-            expected_bin_edges = np.array(expected_bin_edges)
-            seen_bin_edges = h.bin_edges[axis_i]
-            # If 1D, hist1d returns bin_edges straight, not as list
-            if len(analysis_space) == 1:
-                seen_bin_edges = h.bin_edges
-            logging.debug("axis_i: " + str(axis_i))
-            logging.debug("expected_bin_edges: " + str(expected_bin_edges))
-            logging.debug("seen_bin_edges: " + str(seen_bin_edges))
-            logging.debug("h.bin_edges type" + str(h.bin_edges))
-            if len(seen_bin_edges) != len(expected_bin_edges):
-                raise ValueError(
-                    f"Axis {axis_i:d} of histogram {histogram_info} "
-                    f"has {len(seen_bin_edges)} bin edges, but expected {expected_bin_edges}.")
-            try:
-                np.testing.assert_almost_equal(
-                    seen_bin_edges,
-                    expected_bin_edges,
-                    decimal=2)
-            except AssertionError:
-                warnings.warn(
-                    f"Axis {axis_i:d} of histogram {histogram_info} "
-                    f"has bin edges {seen_bin_edges}, but expected {expected_bin_edges}. "
-                    "Since length matches, setting it expected values...")
-                h.bin_edges[axis_i] = expected_bin_edges
-
     def set_dtype(self):
         """Set the data type of the source."""
         self.dtype = []
-        for n, _ in self.config['analysis_space']:
+        for n, _ in self.config["analysis_space"]:
             self.dtype.append((n, float))
-        self.dtype.append(('source', int))
+        self.dtype.append(("source", int))
 
     def set_pdf_histogram(self, h):
         """Set the histogram of the probability density function of the source."""
         self._bin_volumes = h.bin_volumes()  # TODO: make alias
-        # Shouldn't be in HistogramSource... anyway
+        # Should not be in HistogramSource... anyway
         self._n_events_histogram = h.similar_blank_histogram()
 
-        histogram_scale_factor = self.config.get('histogram_scale_factor', 1)
+        histogram_scale_factor = self.config.get("histogram_scale_factor", 1)
         h *= histogram_scale_factor
         logging.debug(
             f"Multiplying histogram with histogram_scale_factor {histogram_scale_factor}. "
@@ -193,7 +193,7 @@ def set_pdf_histogram(self, h):
         logging.debug(f"events_per_day: " + str(self.events_per_day))
 
         # Convert h to density...
-        if self.config.get('in_events_per_bin', True):
+        if self.config.get("in_events_per_bin", True):
             h.histogram /= self._bin_volumes
 
         # ... and finally to probability density
@@ -202,7 +202,7 @@ def set_pdf_histogram(self, h):
         else:
             raise ValueError("Sum of histogram is zero.")
 
-        if self.config.get('convert_to_uniform', False):
+        if self.config.get("convert_to_uniform", False):
             h.histogram = 1. / self._bin_volumes
 
         self._pdf_histogram = h
@@ -221,7 +221,7 @@ def simulate(self, n_events):
         # t = self._pdf_histogram.get_random(n_events)
         h = self._pdf_histogram * self._bin_volumes
         t = h.get_random(n_events)
-        for i, (n, _) in enumerate(self.config['analysis_space']):
+        for i, (n, _) in enumerate(self.config["analysis_space"]):
             ret[n] = t[:, i]
         return ret
 
@@ -239,7 +239,7 @@ class CombinedSource(TemplateSource, HistogramPdfSource):
     """
 
     def build_histogram(self):
-        if not self.config.get('in_events_per_bin', True):
+        if not self.config.get("in_events_per_bin", True):
             raise ValueError(
                 "CombinedSource does not support in_events_per_bin=False")
         if "weight_names" not in self.config:
@@ -247,13 +247,12 @@ def build_histogram(self):
         weights = [
             self.config.get(weight_name, 0) for weight_name in self.config["weight_names"]
         ]
-        histograms = []
-        histnames = self.config['histnames']
-        templatenames = self.config['templatenames']
+        histnames = self.config["histnames"]
+        templatenames = self.config["templatenames"]
         if len(histnames) != len(templatenames):
             raise ValueError(
                 "histnames and templatenames must be the same length")
-        if len(weights) + 1 != len(histograms):
+        if len(weights) + 1 != len(templatenames):
             raise ValueError(
                 "weights must be 1 shorter than histnames, templatenames")
 
@@ -267,6 +266,7 @@ def build_histogram(self):
             # Recognize as a list of slice_args for each histogram
             slice_argss = slice_args
 
+        histograms = []
         slice_fractions = []
         for histname, templatename, slice_args in zip(
                 histnames, templatenames, slice_argss):
@@ -315,9 +315,11 @@ def build_histogram(self):
             raise AssertionError(
                 f"There are bins for source {templatename} with negative entries.")
 
-        logging.info("Normalising combined template, the absolute rate is defined in histogram_multiplier")
+        logging.info(
+            "Normalising combined template, "
+            "the absolute rate is defined in histogram_multiplier")
         h = h / h.n
-        h *= self.config.get('histogram_multiplier', 1)
+        h *= self.config.get("histogram_multiplier", 1)
         logging.info("Applying slice fraction to combined template")
         h *= slice_fractions[0]
 
@@ -357,8 +359,8 @@ def _get_json_spectrum(filename):
         return ret
 
     def build_histogram(self):
-        templatename = self.config['templatename'].format(**self.format_named_parameters)
-        histname = self.config['histname'].format(**self.format_named_parameters)
+        templatename = self.config["templatename"].format(**self.format_named_parameters)
+        histname = self.config["histname"].format(**self.format_named_parameters)
         h = template_to_multihist(templatename, histname)
 
         spectrum = self.config["spectrum"]
@@ -376,9 +378,9 @@ def build_histogram(self):
             raise AssertionError(
                 f"There are bins for source {templatename} with negative entries.")
 
-        if self.config.get('normalise_template', False):
+        if self.config.get("normalise_template", False):
             h /= h.n
-        if not self.config.get('in_events_per_bin', True):
+        if not self.config.get("in_events_per_bin", True):
             h.histogram *= h.bin_volumes()
 
         h = self.apply_slice_args(h)
diff --git a/alea/utils.py b/alea/utils.py
index 3d24fe6b..e429a4d3 100644
--- a/alea/utils.py
+++ b/alea/utils.py
@@ -54,12 +54,20 @@ def adapt_likelihood_config_for_blueice(
     likelihood_config_copy["analysis_space"] = get_analysis_space(
         likelihood_config_copy["analysis_space"])
 
-    likelihood_config_copy["default_source_class"] = locate(
-        likelihood_config_copy["default_source_class"])
+    if "default_source_class" in likelihood_config_copy:
+        likelihood_config_copy["default_source_class"] = locate(
+            likelihood_config_copy["default_source_class"])
 
     for source in likelihood_config_copy["sources"]:
-        source["templatename"] = get_file_path(
-            source["template_filename"], template_folder_list)
+        if "template_filename" in source:
+            source["templatename"] = get_file_path(
+                source["template_filename"], template_folder_list)
+        if "class" in source:
+            source["class"] = locate(source["class"])
+        if "template_filenames" in source:
+            source["templatenames"] = [
+                get_file_path(template_filename, template_folder_list)
+                for template_filename in source["template_filenames"]]
     return likelihood_config_copy
 
 

From 1f9251b57ac5261992cd6eb8f346327fb37719f8 Mon Sep 17 00:00:00 2001
From: dachengx <dx2227@columbia.edu>
Date: Mon, 31 Jul 2023 00:21:40 -0500
Subject: [PATCH 03/21] Minor change

---
 .../unbinned_wimp_statistical_model.yaml      | 12 +++++------
 alea/template_source.py                       | 21 ++++++++++++-------
 2 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/alea/examples/configs/unbinned_wimp_statistical_model.yaml b/alea/examples/configs/unbinned_wimp_statistical_model.yaml
index fbaa4821..ab5c77c0 100644
--- a/alea/examples/configs/unbinned_wimp_statistical_model.yaml
+++ b/alea/examples/configs/unbinned_wimp_statistical_model.yaml
@@ -148,13 +148,13 @@ likelihood_config:
         class: alea.template_source.CombinedSource
         parameters:
           - er_rate_multiplier
-        weight_1: 0.2
-        weight_-1: 0.3
-        weight_names: [weight_1, weight_-1]
-        histnames: [er_template, er_template, er_template,]
+        weight_a: 0.2
+        weight_b: 0.3
+        weight_names: [weight_a, weight_b]  # not meaningful, just an example
+        histnames: [er_template, er_template, er_template]
         template_filenames: [er_template_0.h5, er_template_1.h5, er_template_-1.h5]
-        histogram_multiplier: 200
-        histogram_scale_factor: 3
+        histogram_multiplier: 100  # absolute rate, /year
+        histogram_scale_factor: 2
 
       - name: wimp
         histname: wimp_template
diff --git a/alea/template_source.py b/alea/template_source.py
index f3fd2fb5..71234ec7 100644
--- a/alea/template_source.py
+++ b/alea/template_source.py
@@ -7,7 +7,7 @@
 from blueice import HistogramPdfSource
 from inference_interface import template_to_multihist
 
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(level=logging.debug)
 can_check_binning = True
 
 
@@ -128,7 +128,7 @@ def apply_slice_args(self, h, slice_args=None):
         """
         if slice_args is None:
             slice_args = self.config.get("slice_args", {})
-            if type(slice_args) is dict:
+            if isinstance(slice_args, dict):
                 slice_args = [slice_args]
 
         for sa in slice_args:
@@ -231,6 +231,8 @@ class CombinedSource(TemplateSource, HistogramPdfSource):
     Source that inherits structure from TH2DSource by Jelle,
     but takes in lists of histogram names.
     The first histogram is the base histogram, and the rest are added to it with weights.
+    Currently the weights are hardcoded in the config, can not be changed in the fit.
+    In other words, the weight can not be shape parameter.
 
     :param weights: Weights of the 2nd to the last histograms.
     :param histnames: List of filenames containing the histograms.
@@ -257,9 +259,9 @@ def build_histogram(self):
                 "weights must be 1 shorter than histnames, templatenames")
 
         slice_args = self.config.get("slice_args", {})
-        if type(slice_args) is dict:
+        if isinstance(slice_args, dict):
             slice_args = [slice_args]
-        if all(type(sa) is dict for sa in slice_args):
+        if all(isinstance(sa, dict) for sa in slice_args):
             # Recognize as a single slice_args for all histograms
             slice_argss = [slice_args] * len(histnames)
         else:
@@ -310,17 +312,22 @@ def build_histogram(self):
             histograms[0] += h_comp * weights[i]
         h = histograms[0]
 
+        logging.debug("Setting zero for NaNs and Infs in combined template.")
+        h.histogram[np.isinf(h.histogram)] = 0.
+        h.histogram[np.isnan(h.histogram)] = 0.
+        h.histogram[h.histogram < 0] = 0.
+
         # Set pdf values that are below 0 to zero:
         if np.min(h.histogram) < 0:
             raise AssertionError(
                 f"There are bins for source {templatename} with negative entries.")
 
-        logging.info(
+        logging.debug(
             "Normalising combined template, "
             "the absolute rate is defined in histogram_multiplier")
         h = h / h.n
         h *= self.config.get("histogram_multiplier", 1)
-        logging.info("Applying slice fraction to combined template")
+        logging.debug("Applying slice fraction to combined template")
         h *= slice_fractions[0]
 
         # Fix the bin sizes
@@ -364,7 +371,7 @@ def build_histogram(self):
         h = template_to_multihist(templatename, histname)
 
         spectrum = self.config["spectrum"]
-        if type(spectrum) is str:
+        if isinstance(spectrum, str):
             spectrum = self._get_json_spectrum(spectrum.format(**self.format_named_parameters))
 
         # Perform E-scaling, assume first axis is energy

From 789e3f25950237e5f89b3eda1a42557684393d29 Mon Sep 17 00:00:00 2001
From: dachengx <dx2227@columbia.edu>
Date: Mon, 31 Jul 2023 00:58:27 -0500
Subject: [PATCH 04/21] Set numpy version to less than 1.23.0

Because of https://github.com/XENONnT/alea/issues/70
---
 alea/template_source.py | 28 +++++++++++++---------------
 requirements.txt        |  2 +-
 2 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/alea/template_source.py b/alea/template_source.py
index 71234ec7..7cb0cd3a 100644
--- a/alea/template_source.py
+++ b/alea/template_source.py
@@ -96,6 +96,7 @@ def format_named_parameters(self):
         return format_named_parameters
 
     def build_histogram(self):
+        """Build the histogram of the source."""
         templatename = self.config["templatename"].format(**self.format_named_parameters)
         histname = self.config["histname"].format(**self.format_named_parameters)
         h = template_to_multihist(templatename, histname)
@@ -112,7 +113,8 @@ def build_histogram(self):
 
         # Fix the bin sizes
         if can_check_binning:
-            self._check_binning(h, f"{histname:s} in hdf5 file {templatename:s}")
+            histogram_info = f"{histname:s} in hdf5 file {templatename:s}"
+            self._check_binning(h, histogram_info)
 
         self.set_dtype()
         self.set_pdf_histogram(h)
@@ -241,6 +243,7 @@ class CombinedSource(TemplateSource, HistogramPdfSource):
     """
 
     def build_histogram(self):
+        """Build the histogram of the source."""
         if not self.config.get("in_events_per_bin", True):
             raise ValueError(
                 "CombinedSource does not support in_events_per_bin=False")
@@ -293,21 +296,12 @@ def build_histogram(self):
         for i in range(len(weights)):
             h_comp = histograms[0].similar_blank_histogram()
             h = histograms[i + 1]
-            # base_part_norm = 0.
-            # h_centers = h_comp.bin_centers()
-            # h_inds = [range(len(hc)) for hc in h_centers]
-            # for inds in product(*h_inds):
-            #     bincs = [h_centers[j][k] for j, k in enumerate(inds)]
-            #     inds_comp = h_comp.get_bin_indices(bincs)
-            #     base_part_norm += histograms[0][inds]
-            #     h_comp[inds] = h[inds_comp]
-            # h_comp *= base_part_norm
             hslices = []
             for j, bincs in enumerate(h.bin_centers()):
                 hsliced = h_comp.get_axis_bin_index(bincs[0], j)
                 hsliceu = h_comp.get_axis_bin_index(bincs[-1], j) + 1
-                hslices.append(slice(hsliced, hsliceu))
-            # TODO: check here what norm I want.
+                hslices.append(slice(hsliced, hsliceu, 1))
+            # TODO: check the normalization here.
             h_comp[hslices] += h.histogram
             histograms[0] += h_comp * weights[i]
         h = histograms[0]
@@ -332,7 +326,8 @@ def build_histogram(self):
 
         # Fix the bin sizes
         if can_check_binning:
-            self._check_binning(h, f"combined template {histnames} in hdf5 file {templatenames}")
+            histogram_info = f"combined template {histnames} in hdf5 file {templatenames}"
+            self._check_binning(h, histogram_info)
 
         self.set_dtype()
         self.set_pdf_histogram(h)
@@ -357,7 +352,8 @@ def _get_json_spectrum(filename):
         :param filename: Name of the JSON file.
         :type filename: str
         """
-        contents = json.load(open(filename, "r"))
+        with open(filename, "r") as f:
+            contents = json.load(f)
         logging.debug(contents["description"])
         esyst = contents["coordinate_system"][0][1]
         ret = interp1d(
@@ -366,6 +362,7 @@ def _get_json_spectrum(filename):
         return ret
 
     def build_histogram(self):
+        """Build the histogram of the source."""
         templatename = self.config["templatename"].format(**self.format_named_parameters)
         histname = self.config["histname"].format(**self.format_named_parameters)
         h = template_to_multihist(templatename, histname)
@@ -394,7 +391,8 @@ def build_histogram(self):
 
         # Fix the bin sizes
         if can_check_binning:
-            self._check_binning(h, f"reweighted {histname:s} in hdf5 file {templatename:s}")
+            histogram_info = f"reweighted {histname:s} in hdf5 file {templatename:s}"
+            self._check_binning(h, histogram_info)
 
         self.set_dtype()
         self.set_pdf_histogram(h)
diff --git a/requirements.txt b/requirements.txt
index 27a03eb8..388c7650 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,7 +7,7 @@ iminuit>=2.21.0
 matplotlib
 mergedeep
 multihist
-numpy
+numpy<1.23.0
 PyYAML
 scipy
 setuptools

From e1aac2a348f3420fffdc33ab2e82e6215a581fc1 Mon Sep 17 00:00:00 2001
From: dachengx <dx2227@columbia.edu>
Date: Mon, 31 Jul 2023 14:26:54 -0500
Subject: [PATCH 05/21] Update docstring to Google style

---
 alea/template_source.py | 96 ++++++++++++++++++++---------------------
 docs/source/conf.py     | 15 ++++++-
 2 files changed, 61 insertions(+), 50 deletions(-)

diff --git a/alea/template_source.py b/alea/template_source.py
index 7cb0cd3a..303186fd 100644
--- a/alea/template_source.py
+++ b/alea/template_source.py
@@ -17,44 +17,40 @@ class TemplateSource(HistogramPdfSource):
     The parameters are set in self.config.
     "templatename", "histname", "analysis_space" must be in self.config
 
-    :ivar config: The configuration of the source.
-    :vartype config: dict
-    :ivar dtype: The data type of the source.
-    :vartype dtype: list
-    :ivar _bin_volumes: The bin volumes of the source.
-    :vartype _bin_volumes: numpy.ndarray
-    :ivar _n_events_histogram: The histogram of the number of events of the source.
-    :vartype _n_events_histogram: multihist.MultiHistBase
-    :ivar events_per_day: The number of events per day of the source.
-    :vartype events_per_day: float
-    :ivar _pdf_histogram: The histogram of the probability density function of the source.
-    :vartype _pdf_histogram: multihist.MultiHistBase
-
-    :param templatename: Hdf5 file to open.
-    :param histname: Histogram name.
-    :param named_parameters:
-        List of config setting names to pass to .format on histname and filename.
-    :type named_parameters: list
-    :param normalise_template: Normalise the template histogram.
-    :type normalise_template: bool
-    :param in_events_per_bin:
-        If True, histogram is in events per day / bin.
-        If False or absent, histogram is already pdf.
-    :param histogram_scale_factor: Multiply histogram by this number
-    :param convert_to_uniform: Convert the histogram to a uniform per bin distribution.
-    :param log10_bins:
-        List of axis numbers.
-        If True, bin edges on this axis in the hdf5 file are log10() of the actual bin edges.
+    Attributes:
+        config (dict): The configuration of the source.
+        dtype (list): The data type of the source.
+        _bin_volumes (numpy.ndarray): The bin volumes of the source.
+        _n_events_histogram (multihist.MultiHistBase):
+            The histogram of the number of events of the source.
+        events_per_day (float): The number of events per day of the source.
+        _pdf_histogram (multihist.MultiHistBase):
+            The histogram of the probability density function of the source.
+
+    Args:
+        config (dict): The configuration of the source.
+        templatename: Hdf5 file to open.
+        histname: Histogram name.
+        named_parameters (list):
+            List of config setting names to pass to .format on histname and filename.
+        normalise_template (bool): Normalise the template histogram.
+        in_events_per_bin (bool):
+            If True, histogram is in events per day / bin.
+            If False or absent, histogram is already pdf.
+        histogram_scale_factor (float): Multiply histogram by this number
+        convert_to_uniform (bool): Convert the histogram to a uniform per bin distribution.
+        log10_bins (list):
+            List of axis numbers.
+            If True, bin edges on this axis in the hdf5 file are log10() of the actual bin edges.
     """
 
     def _check_binning(self, h, histogram_info):
         """
         Check if the histogram"s bin edges are the same to analysis_space.
 
-        :param h: The histogram to check.
-        :type h: multihist.MultiHistBase
-        :param histogram_info: Information of the histogram.
-        :type histogram_info: str
+        Args:
+            h (multihist.MultiHistBase): The histogram to check.
+            histogram_info (str): Information of the histogram.
         """
         # Deal with people who have log10"d their bins
         for axis_i in self.config.get("log10_bins", []):
@@ -123,10 +119,9 @@ def apply_slice_args(self, h, slice_args=None):
         """
         Apply slice arguments to the histogram.
 
-        :param h: The histogram to apply the slice arguments to.
-        :type h: multihist.MultiHistBase
-        :param slice_args: The slice arguments to apply.
-        :type slice_args: dict
+        Args:
+            h (multihist.MultiHistBase): The histogram to apply the slice arguments to.
+            slice_args (dict): The slice arguments to apply.
         """
         if slice_args is None:
             slice_args = self.config.get("slice_args", {})
@@ -214,10 +209,11 @@ def simulate(self, n_events):
         """
         Simulate events from the source.
 
-        :param n_events: The number of events to simulate.
-        :type n_events: int
-        :return: The simulated events.
-        :rtype: numpy.ndarray
+        Args:
+            n_events (int): The number of events to simulate.
+
+        Returns:
+            numpy.ndarray: The simulated events.
         """
         ret = np.zeros(n_events, dtype=self.dtype)
         # t = self._pdf_histogram.get_random(n_events)
@@ -236,10 +232,11 @@ class CombinedSource(TemplateSource, HistogramPdfSource):
     Currently the weights are hardcoded in the config, can not be changed in the fit.
     In other words, the weight can not be shape parameter.
 
-    :param weights: Weights of the 2nd to the last histograms.
-    :param histnames: List of filenames containing the histograms.
-    :param templatenames: List of names of histograms within the hdf5 files.
-    :param histogram_multiplier: Absolute rate of the combined template before multihist slicing.
+    Args:
+        weights: Weights of the 2nd to the last histograms.
+        histnames: List of filenames containing the histograms.
+        templatenames: List of names of histograms within the hdf5 files.
+        histogram_multiplier: Absolute rate of the combined template before multihist slicing.
     """
 
     def build_histogram(self):
@@ -338,9 +335,10 @@ class SpectrumTemplateSource(TemplateSource, HistogramPdfSource):
     Reweighted template source by energy spectrum.
     The first axis of the template is assumed to be energy.
 
-    :param spectrum_name:
-        Name of bbf json-like spectrum _OR_ function that can be called
-        templatename #3D histogram (Etrue, S1, S2) to open
+    Args:
+        spectrum_name:
+            Name of bbf json-like spectrum _OR_ function that can be called
+            templatename #3D histogram (Etrue, S1, S2) to open
     """
 
     @staticmethod
@@ -349,8 +347,8 @@ def _get_json_spectrum(filename):
         Translates bbf-style JSON files to spectra.
         units are keV and /kev*day*kg
 
-        :param filename: Name of the JSON file.
-        :type filename: str
+        Args:
+            filename (str): Name of the JSON file.
         """
         with open(filename, "r") as f:
             contents = json.load(f)
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 02a92ac4..8421d809 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -13,10 +13,12 @@
 # -- General configuration
 
 extensions = [
+    'nbsphinx',
     'sphinx.ext.autodoc',
     'sphinx.ext.intersphinx',
     'sphinx.ext.viewcode',
-    'nbsphinx',
+    'sphinx.ext.napoleon',
+    'sphinx.ext.todo',
 ]
 
 intersphinx_mapping = {
@@ -27,13 +29,24 @@
 
 templates_path = ['_templates']
 
+# -- Options for NAPOLEON output
+
+napoleon_include_init_with_doc = True
+napoleon_include_private_with_doc = True
+napoleon_include_special_with_doc = True
+
 # -- Options for HTML output
 
 html_theme = 'sphinx_rtd_theme'
 
 # -- Options for EPUB output
+
 epub_show_urls = 'footnote'
 
+# -- Options for TODO output
+
+todo_include_todos = True
+
 def setup(app):
     # Hack to import something from this dir. Apparently we're in a weird
     # situation where you get a __name__  is not in globals KeyError

From ef53ff84313d3afa1b543fe0817119dd3555690d Mon Sep 17 00:00:00 2001
From: dachengx <dx2227@columbia.edu>
Date: Tue, 1 Aug 2023 00:38:17 -0500
Subject: [PATCH 06/21] Check needed weighted parameters of CombinedSource

---
 .../configs/unbinned_wimp_statistical_model.yaml       | 10 ++++++----
 alea/parameters.py                                     |  2 ++
 alea/template_source.py                                |  7 ++++++-
 3 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/alea/examples/configs/unbinned_wimp_statistical_model.yaml b/alea/examples/configs/unbinned_wimp_statistical_model.yaml
index ab5c77c0..0912194a 100644
--- a/alea/examples/configs/unbinned_wimp_statistical_model.yaml
+++ b/alea/examples/configs/unbinned_wimp_statistical_model.yaml
@@ -53,7 +53,7 @@ parameter_definition:
   er_band_shift:
     nominal_value: 0
     ptype: shape
-    uncertainty: null  # scipy.stats.uniform(loc=-1, scale=2)
+    uncertainty: null  # 'scipy.stats.uniform(loc=-1, scale=2)'
     # relative_uncertainty: false
     fittable: true
     blueice_anchors:
@@ -148,9 +148,11 @@ likelihood_config:
         class: alea.template_source.CombinedSource
         parameters:
           - er_rate_multiplier
-        weight_a: 0.2
-        weight_b: 0.3
-        weight_names: [weight_a, weight_b]  # not meaningful, just an example
+          - er_band_shift
+        named_parameters:
+          - er_band_shift
+        fixed_weight: 0.2
+        weight_names: [fixed_weight, er_band_shift]  # not meaningful, just an example
         histnames: [er_template, er_template, er_template]
         template_filenames: [er_template_0.h5, er_template_1.h5, er_template_-1.h5]
         histogram_multiplier: 100  # absolute rate, /year
diff --git a/alea/parameters.py b/alea/parameters.py
index e7f26da0..80540049 100644
--- a/alea/parameters.py
+++ b/alea/parameters.py
@@ -1,5 +1,7 @@
 from typing import Any, Dict, List, Optional, Tuple
 
+import scipy
+
 
 class Parameter:
     """
diff --git a/alea/template_source.py b/alea/template_source.py
index 303186fd..5715cea1 100644
--- a/alea/template_source.py
+++ b/alea/template_source.py
@@ -244,10 +244,15 @@ def build_histogram(self):
         if not self.config.get("in_events_per_bin", True):
             raise ValueError(
                 "CombinedSource does not support in_events_per_bin=False")
+        # check if all the necessary parameters, like weights are specified
         if "weight_names" not in self.config:
             raise ValueError("weight_names must be specified")
+        find_weight = [weight_name in self.config for weight_name in self.config["weight_names"]]
+        if not all(find_weight):
+            missing_weight = self.config["weight_names"][find_weight]
+            raise ValueError(f"All weight_names must be specified, but {missing_weight} are not. ")
         weights = [
-            self.config.get(weight_name, 0) for weight_name in self.config["weight_names"]
+            self.config[weight_name] for weight_name in self.config["weight_names"]
         ]
         histnames = self.config["histnames"]
         templatenames = self.config["templatenames"]

From 459c93b0be2d20e98a5e055b94494a2f17640fcb Mon Sep 17 00:00:00 2001
From: dachengx <dx2227@columbia.edu>
Date: Tue, 1 Aug 2023 00:45:25 -0500
Subject: [PATCH 07/21] Update docstring

---
 alea/template_source.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/alea/template_source.py b/alea/template_source.py
index 5715cea1..14867985 100644
--- a/alea/template_source.py
+++ b/alea/template_source.py
@@ -229,8 +229,7 @@ class CombinedSource(TemplateSource, HistogramPdfSource):
     Source that inherits structure from TH2DSource by Jelle,
     but takes in lists of histogram names.
     The first histogram is the base histogram, and the rest are added to it with weights.
-    Currently the weights are hardcoded in the config, can not be changed in the fit.
-    In other words, the weight can not be shape parameter.
+    The weights can be set as shape parameters in the config.
 
     Args:
         weights: Weights of the 2nd to the last histograms.

From 698e0ba2db6888b7a091b7e8804c821274054174 Mon Sep 17 00:00:00 2001
From: dachengx <dx2227@columbia.edu>
Date: Tue, 1 Aug 2023 13:47:08 -0500
Subject: [PATCH 08/21] No need to use multiple inheritance

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

diff --git a/alea/template_source.py b/alea/template_source.py
index 14867985..9a26b57a 100644
--- a/alea/template_source.py
+++ b/alea/template_source.py
@@ -224,7 +224,7 @@ def simulate(self, n_events):
         return ret
 
 
-class CombinedSource(TemplateSource, HistogramPdfSource):
+class CombinedSource(TemplateSource):
     """
     Source that inherits structure from TH2DSource by Jelle,
     but takes in lists of histogram names.
@@ -334,7 +334,7 @@ def build_histogram(self):
         self.set_pdf_histogram(h)
 
 
-class SpectrumTemplateSource(TemplateSource, HistogramPdfSource):
+class SpectrumTemplateSource(TemplateSource):
     """
     Reweighted template source by energy spectrum.
     The first axis of the template is assumed to be energy.

From 462fb5b6c0d94c5ca4865b05f3ca93afe2c81e95 Mon Sep 17 00:00:00 2001
From: dachengx <dx2227@columbia.edu>
Date: Wed, 2 Aug 2023 09:00:00 -0500
Subject: [PATCH 09/21] Raise error when there is no key in the json file

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

diff --git a/alea/template_source.py b/alea/template_source.py
index 9a26b57a..5c9dd001 100644
--- a/alea/template_source.py
+++ b/alea/template_source.py
@@ -353,10 +353,16 @@ def _get_json_spectrum(filename):
 
         Args:
             filename (str): Name of the JSON file.
+
+        Todo:
+            Define the format of the JSON file clearly.
         """
         with open(filename, "r") as f:
             contents = json.load(f)
-        logging.debug(contents["description"])
+        if "description" in contents:
+            logging.debug(contents["description"])
+        if "coordinate_system" not in contents:
+            raise ValueError("Coordinate system not in JSON file.")
         esyst = contents["coordinate_system"][0][1]
         ret = interp1d(
             np.linspace(*esyst), contents["map"],
@@ -369,9 +375,11 @@ def build_histogram(self):
         histname = self.config["histname"].format(**self.format_named_parameters)
         h = template_to_multihist(templatename, histname)
 
-        spectrum = self.config["spectrum"]
+        if "spectrum_name" not in self.config:
+            raise ValueError("spectrum_name not in config")
+        spectrum_name = self.config["spectrum_name"]
         if isinstance(spectrum, str):
-            spectrum = self._get_json_spectrum(spectrum.format(**self.format_named_parameters))
+            spectrum = self._get_json_spectrum(spectrum_name.format(**self.format_named_parameters))
 
         # Perform E-scaling, assume first axis is energy
         ecenters = h.bin_centers[0]

From de2044d93a852489a3d851b3f0e2aa2a93cd3820 Mon Sep 17 00:00:00 2001
From: dachengx <dx2227@columbia.edu>
Date: Fri, 4 Aug 2023 14:15:01 -0500
Subject: [PATCH 10/21] Recover logging, update docstring

---
 alea/template_source.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/alea/template_source.py b/alea/template_source.py
index 5c9dd001..17deed57 100644
--- a/alea/template_source.py
+++ b/alea/template_source.py
@@ -7,7 +7,7 @@
 from blueice import HistogramPdfSource
 from inference_interface import template_to_multihist
 
-logging.basicConfig(level=logging.debug)
+logging.basicConfig(level=logging.INFO)
 can_check_binning = True
 
 
@@ -226,8 +226,7 @@ def simulate(self, n_events):
 
 class CombinedSource(TemplateSource):
     """
-    Source that inherits structure from TH2DSource by Jelle,
-    but takes in lists of histogram names.
+    Source combines in lists of histogram names.
     The first histogram is the base histogram, and the rest are added to it with weights.
     The weights can be set as shape parameters in the config.
 

From 07fb1c0c66f088921dd6fcd1ffc5f4e986cf4a55 Mon Sep 17 00:00:00 2001
From: dachengx <dx2227@columbia.edu>
Date: Fri, 4 Aug 2023 14:16:10 -0500
Subject: [PATCH 11/21] Kill cache of coverage

Reference: https://github.com/lemurheavy/coveralls-public/issues/1065#issuecomment-435494495
---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index db7cfaf4..26ff03ca 100644
--- a/README.md
+++ b/README.md
@@ -3,7 +3,7 @@ A tool to perform toyMC-based inference constructions
 
 [![DOI](https://zenodo.org/badge/654100988.svg)](https://zenodo.org/badge/latestdoi/654100988)
 [![Test package](https://github.com/XENONnT/alea/actions/workflows/pytest.yml/badge.svg?branch=main)](https://github.com/XENONnT/alea/actions/workflows/pytest.yml)
-[![Coverage Status](https://coveralls.io/repos/github/XENONnT/alea/badge.svg?branch=main)](https://coveralls.io/github/XENONnT/alea?branch=main)
+[![Coverage Status](https://coveralls.io/repos/github/XENONnT/alea/badge.svg?branch=main&kill_cache=1)](https://coveralls.io/github/XENONnT/alea?branch=main&kill_cache=1)
 [![PyPI version shields.io](https://img.shields.io/pypi/v/alea-inference.svg)](https://pypi.python.org/pypi/alea-inference/)
 [![Readthedocs Badge](https://readthedocs.org/projects/alea/badge/?version=latest)](https://alea.readthedocs.io/en/latest/?badge=latest)
 [![CodeFactor](https://www.codefactor.io/repository/github/xenonnt/alea/badge)](https://www.codefactor.io/repository/github/xenonnt/alea)

From 6fd9c1972acfd1a1b4f7434fc1e94622f79baeb2 Mon Sep 17 00:00:00 2001
From: dachengx <dx2227@columbia.edu>
Date: Sat, 5 Aug 2023 21:51:16 -0500
Subject: [PATCH 12/21] Test SpectrumTemplateSource

---
 alea/examples/configs/cs1_spectrum.json       |  10 ++
 .../unbinned_wimp_statistical_model.yaml      |  53 +++++++--
 alea/parameters.py                            |  23 ++--
 alea/template_source.py                       | 110 +++++++++---------
 alea/utils.py                                 |  17 ++-
 5 files changed, 133 insertions(+), 80 deletions(-)
 create mode 100644 alea/examples/configs/cs1_spectrum.json

diff --git a/alea/examples/configs/cs1_spectrum.json b/alea/examples/configs/cs1_spectrum.json
new file mode 100644
index 00000000..5a9d7e95
--- /dev/null
+++ b/alea/examples/configs/cs1_spectrum.json
@@ -0,0 +1,10 @@
+{
+    "coordinate_system": [
+        0.0,
+        100.0
+    ],
+    "map": [
+        1.0,
+        1.0
+    ]
+}
\ No newline at end of file
diff --git a/alea/examples/configs/unbinned_wimp_statistical_model.yaml b/alea/examples/configs/unbinned_wimp_statistical_model.yaml
index d203f1e5..fe76c736 100644
--- a/alea/examples/configs/unbinned_wimp_statistical_model.yaml
+++ b/alea/examples/configs/unbinned_wimp_statistical_model.yaml
@@ -17,7 +17,12 @@ parameter_definition:
   livetime_sr2:
     nominal_value: 0.5
     fittable: false
-    description: Livetime of SR1 in years
+    description: Livetime of SR2 in years
+
+  livetime_sr3:
+    nominal_value: 1.0
+    fittable: false
+    description: Livetime of SR3 in years
 
   wimp_rate_multiplier:
     nominal_value: 1.0
@@ -56,7 +61,7 @@ parameter_definition:
   er_band_shift:
     nominal_value: 0
     ptype: shape
-    uncertainty: null  # 'scipy.stats.uniform(loc=-1, scale=2)'
+    uncertainty: null  # scipy.stats.uniform(loc=0, scale=2)
     # relative_uncertainty: false
     fittable: true
     blueice_anchors:
@@ -79,8 +84,8 @@ likelihood_config:
       default_source_class: alea.template_source.TemplateSource
       likelihood_type: blueice.likelihood.UnbinnedLogLikelihood
       analysis_space:
-        - "cs1": 'np.arange(0, 102, 2)'
-        - "cs2": 'np.geomspace(100, 100000, 51)'
+        - cs1: np.linspace(0, 100, 51)
+        - cs2: np.geomspace(100, 100000, 51)
       in_events_per_bin: true
       livetime_parameter: livetime_sr0
       slice_args: {}
@@ -110,8 +115,8 @@ likelihood_config:
       default_source_class: alea.template_source.TemplateSource
       likelihood_type: blueice.likelihood.UnbinnedLogLikelihood
       analysis_space:
-        - "cs1": 'np.arange(0, 102, 2)'
-        - "cs2": 'np.geomspace(100, 100000, 51)'
+        - cs1: np.linspace(0, 100, 51)
+        - cs2: np.geomspace(100, 100000, 51)
       in_events_per_bin: true
       livetime_parameter: livetime_sr1
       slice_args: {}
@@ -141,8 +146,8 @@ likelihood_config:
       default_source_class: alea.template_source.TemplateSource
       likelihood_type: blueice.likelihood.UnbinnedLogLikelihood
       analysis_space:
-        - "cs1": 'np.arange(0, 102, 2)'
-        - "cs2": 'np.geomspace(100, 100000, 51)'
+        - cs1: np.linspace(0, 100, 51)
+        - cs2: np.geomspace(100, 100000, 51)
       in_events_per_bin: true
       livetime_parameter: livetime_sr2
       slice_args: {}
@@ -170,3 +175,35 @@ likelihood_config:
         template_filename: wimp50gev_template.h5
         apply_efficiency: True
         efficiency_name: signal_efficiency
+
+    # SR3
+    - name: sr3
+      default_source_class: alea.template_source.TemplateSource
+      likelihood_type: blueice.likelihood.UnbinnedLogLikelihood
+      analysis_space:
+        - cs1: np.linspace(0, 100, 51)
+        - cs2: np.geomspace(100, 100000, 51)
+      in_events_per_bin: true
+      livetime_parameter: livetime_sr3
+      slice_args: {}
+      sources:
+      - name: er
+        histname: er_template
+        parameters:
+          - er_rate_multiplier
+          - er_band_shift
+        named_parameters:
+          - er_band_shift
+        template_filename: er_template_{er_band_shift}.h5
+
+      - name: wimp
+        class: alea.template_source.SpectrumTemplateSource
+        histname: wimp_template
+        parameters:
+          - wimp_rate_multiplier
+          - wimp_mass
+          - signal_efficiency
+        template_filename: wimp50gev_template.h5
+        spectrum_name: cs1_spectrum.json
+        apply_efficiency: True
+        efficiency_name: signal_efficiency
diff --git a/alea/parameters.py b/alea/parameters.py
index 5de3ef10..55d94e93 100644
--- a/alea/parameters.py
+++ b/alea/parameters.py
@@ -1,10 +1,6 @@
 import warnings
 from typing import Any, Dict, List, Optional, Tuple
 
-# These imports are needed to evaluate the uncertainty string
-import numpy  # noqa: F401
-import scipy  # noqa: F401
-
 from alea.utils import within_limits, clip_limits
 
 
@@ -50,8 +46,8 @@ def __init__(
         self.nominal_value = nominal_value
         self.fittable = fittable
         self.ptype = ptype
-        self.uncertainty = uncertainty
         self.relative_uncertainty = relative_uncertainty
+        self.uncertainty = uncertainty
         self.blueice_anchors = blueice_anchors
         self.fit_limits = fit_limits
         self.parameter_interval_bounds = parameter_interval_bounds
@@ -70,21 +66,22 @@ def __repr__(self) -> str:
     def uncertainty(self) -> float or Any:
         """
         Return the uncertainty of the parameter.
-        If the uncertainty is a string, it can be evaluated as a numpy or scipy function.
         """
         if isinstance(self._uncertainty, str):
-            # Evaluate the uncertainty if it's a string starting with "scipy." or "numpy."
-            if self._uncertainty.startswith("scipy.") or self._uncertainty.startswith("numpy."):
-                return eval(self._uncertainty)
-            else:
-                raise ValueError(
-                    f"Uncertainty string '{self._uncertainty}'"
-                    " must start with 'scipy.' or 'numpy.'")
+            NotImplementedError(
+                "Only float uncertainties are supported at the moment.")
         else:
             return self._uncertainty
 
     @uncertainty.setter
     def uncertainty(self, value: float or str) -> None:
+        """
+        Uncertainty can be a float or a string, but can only be a float if relative_uncertainty
+        """
+        if self.relative_uncertainty and isinstance(value, str):
+            raise ValueError(
+                f"relative_uncertainty is not supported for "
+                f"string uncertainties of {self.name}.")
         self._uncertainty = value
 
     @property
diff --git a/alea/template_source.py b/alea/template_source.py
index 17deed57..57be00db 100644
--- a/alea/template_source.py
+++ b/alea/template_source.py
@@ -7,6 +7,8 @@
 from blueice import HistogramPdfSource
 from inference_interface import template_to_multihist
 
+from alea.utils import get_file_path
+
 logging.basicConfig(level=logging.INFO)
 can_check_binning = True
 
@@ -122,6 +124,7 @@ def apply_slice_args(self, h, slice_args=None):
         Args:
             h (multihist.MultiHistBase): The histogram to apply the slice arguments to.
             slice_args (dict): The slice arguments to apply.
+                The sum_axis, slice_axis, and slice_axis_limits are supported.
         """
         if slice_args is None:
             slice_args = self.config.get("slice_args", {})
@@ -129,42 +132,36 @@ def apply_slice_args(self, h, slice_args=None):
                 slice_args = [slice_args]
 
         for sa in slice_args:
+            sum_axis = sa.get("sum_axis", None)
             slice_axis = sa.get("slice_axis", None)
-            # Decide if you wish to sum the histogram into lower dimensions or
-            sum_axis = sa.get("sum_axis", False)
-            collapse_axis = sa.get("collapse_axis", None)
-            collapse_slices = sa.get("collapse_slices", None)
             if slice_axis is not None:
-                slice_axis_number = h.get_axis_number(slice_axis)
-                bes = h.bin_edges[slice_axis_number]
+                # When slice_axis is None, slice_axis_limits is not used
+                bin_edges = h.bin_edges[h.get_axis_number(slice_axis)]
                 slice_axis_limits = sa.get(
-                    "slice_axis_limits", [bes[0], bes[-1]])
-                if sum_axis:
-                    logging.debug(
-                        f"Slice and sum over axis {slice_axis} from {slice_axis_limits[0]} "
-                        f"to {slice_axis_limits[1]}")
-                    axis_names = h.axis_names_without(slice_axis)
-                    h = h.slicesum(
-                        axis=slice_axis,
-                        start=slice_axis_limits[0],
-                        stop=slice_axis_limits[1])
-                    h.axis_names = axis_names
-                else:
-                    logging.debug(f"Normalization before slicing: {h.n}.")
-                    logging.debug(
-                        f"Slice over axis {slice_axis} from {slice_axis_limits[0]} "
-                        f"to {slice_axis_limits[1]}")
-                    h = h.slice(
-                        axis=slice_axis,
-                        start=slice_axis_limits[0],
-                        stop=slice_axis_limits[1])
-                    logging.debug(f"Normalization after slicing: {h.n}.")
-
-            if collapse_axis is not None:
-                if collapse_slices is None:
-                    raise ValueError(
-                        "To collapse you must supply collapse_slices")
-                h = h.collapse_bins(collapse_slices, axis=collapse_axis)
+                    "slice_axis_limits", [bin_edges[0], bin_edges[-1]])
+            # sum and/or slice the histogram
+            if (sum_axis is not None) and (slice_axis is not None):
+                logging.debug(
+                    f"Slice and sum over axis {slice_axis} from {slice_axis_limits[0]} "
+                    f"to {slice_axis_limits[1]}")
+                h = h.slicesum(
+                    start=slice_axis_limits[0],
+                    stop=slice_axis_limits[1],
+                    axis=slice_axis)
+            elif sum_axis is not None:
+                logging.debug(
+                    f"Sum over axis {sum_axis}")
+                h = h.sum(axis=sum_axis)
+            elif slice_axis is not None:
+                logging.debug(
+                    f"Slice over axis {slice_axis} from {slice_axis_limits[0]} "
+                    f"to {slice_axis_limits[1]}")
+                logging.debug(f"Normalization before slicing: {h.n}.")
+                h = h.slice(
+                    start=slice_axis_limits[0],
+                    stop=slice_axis_limits[1],
+                    axis=slice_axis)
+                logging.debug(f"Normalization after slicing: {h.n}.")
         return h
 
     def set_dtype(self):
@@ -176,7 +173,7 @@ def set_dtype(self):
 
     def set_pdf_histogram(self, h):
         """Set the histogram of the probability density function of the source."""
-        self._bin_volumes = h.bin_volumes()  # TODO: make alias
+        self._bin_volumes = h.bin_volumes()
         # Should not be in HistogramSource... anyway
         self._n_events_histogram = h.similar_blank_histogram()
 
@@ -216,7 +213,6 @@ def simulate(self, n_events):
             numpy.ndarray: The simulated events.
         """
         ret = np.zeros(n_events, dtype=self.dtype)
-        # t = self._pdf_histogram.get_random(n_events)
         h = self._pdf_histogram * self._bin_volumes
         t = h.get_random(n_events)
         for i, (n, _) in enumerate(self.config["analysis_space"]):
@@ -242,7 +238,7 @@ def build_histogram(self):
         if not self.config.get("in_events_per_bin", True):
             raise ValueError(
                 "CombinedSource does not support in_events_per_bin=False")
-        # check if all the necessary parameters, like weights are specified
+        # Check if all the necessary parameters, like weights are specified
         if "weight_names" not in self.config:
             raise ValueError("weight_names must be specified")
         find_weight = [weight_name in self.config for weight_name in self.config["weight_names"]]
@@ -301,8 +297,7 @@ def build_histogram(self):
                 hsliced = h_comp.get_axis_bin_index(bincs[0], j)
                 hsliceu = h_comp.get_axis_bin_index(bincs[-1], j) + 1
                 hslices.append(slice(hsliced, hsliceu, 1))
-            # TODO: check the normalization here.
-            h_comp[hslices] += h.histogram
+            h_comp[tuple(hslices)] += h.histogram
             histograms[0] += h_comp * weights[i]
         h = histograms[0]
 
@@ -335,20 +330,18 @@ def build_histogram(self):
 
 class SpectrumTemplateSource(TemplateSource):
     """
-    Reweighted template source by energy spectrum.
-    The first axis of the template is assumed to be energy.
+    Reweighted template source by 1D spectrum.
+    The first axis of the template is assumed to be reweighted.
 
     Args:
         spectrum_name:
-            Name of bbf json-like spectrum _OR_ function that can be called
-            templatename #3D histogram (Etrue, S1, S2) to open
+            Name of bbf json-like spectrum file
     """
 
     @staticmethod
     def _get_json_spectrum(filename):
         """
-        Translates bbf-style JSON files to spectra.
-        units are keV and /kev*day*kg
+        Translates bbf-style JSON files to spectra
 
         Args:
             filename (str): Name of the JSON file.
@@ -356,15 +349,16 @@ def _get_json_spectrum(filename):
         Todo:
             Define the format of the JSON file clearly.
         """
-        with open(filename, "r") as f:
+        with open(get_file_path(filename), "r") as f:
             contents = json.load(f)
         if "description" in contents:
             logging.debug(contents["description"])
         if "coordinate_system" not in contents:
             raise ValueError("Coordinate system not in JSON file.")
-        esyst = contents["coordinate_system"][0][1]
+        if "map" not in contents:
+            raise ValueError("Map not in JSON file.")
         ret = interp1d(
-            np.linspace(*esyst), contents["map"],
+            contents["coordinate_system"], contents["map"],
             bounds_error=False, fill_value=0.)
         return ret
 
@@ -376,16 +370,18 @@ def build_histogram(self):
 
         if "spectrum_name" not in self.config:
             raise ValueError("spectrum_name not in config")
-        spectrum_name = self.config["spectrum_name"]
-        if isinstance(spectrum, str):
-            spectrum = self._get_json_spectrum(spectrum_name.format(**self.format_named_parameters))
-
-        # Perform E-scaling, assume first axis is energy
-        ecenters = h.bin_centers[0]
-        ediffs = h.bin_volumes[0]
-        h.histogram = h.histogram * (spectrum(ecenters) * ediffs)[:, None, None]
-        h = h.sum(axis=0)  # Remove energy-axis
-        logging.debug("source", "mu ", h.n)
+        spectrum = self._get_json_spectrum(
+            self.config["spectrum_name"].format(**self.format_named_parameters))
+
+        # Perform scaling, the first axis is assumed to be reweighted
+        # The spectrum is assumed to be probability density (in per the unit of first axis).
+        axis = 0
+        # h = h.normalize(axis=axis)
+        bin_edges = h.bin_edges[axis]
+        bin_centers = h.bin_centers(axis=axis)
+        slices = [None] * h.histogram.ndim
+        slices[axis] = slice(None)
+        h.histogram = h.histogram * (spectrum(bin_centers) * np.diff(bin_edges))[tuple(slices)]
 
         if np.min(h.histogram) < 0:
             raise AssertionError(
diff --git a/alea/utils.py b/alea/utils.py
index 9e404d69..05b056ba 100644
--- a/alea/utils.py
+++ b/alea/utils.py
@@ -7,7 +7,10 @@
 from pydoc import locate
 import logging
 
-import numpy as np
+# These imports are needed to evaluate strings
+import numpy  # noqa: F401
+import numpy as np  # noqa: F401
+from scipy import stats  # noqa: F401
 
 logging.basicConfig(level=logging.INFO)
 
@@ -15,6 +18,16 @@
 MAX_FLOAT = np.sqrt(np.finfo(np.float32).max)
 
 
+def evaluate_numpy_scipy_expression(value: str):
+    """Evaluate numpy(np) and scipy.stats expression."""
+    if value.startswith("stats."):
+        return eval(value)
+    elif value.startswith("np.") or value.startswith("numpy."):
+        return eval(value)
+    else:
+        raise ValueError(f"Expression {value} not understood.")
+
+
 def get_analysis_space(analysis_space: dict) -> list:
     """Convert analysis_space to a list of tuples with evaluated values."""
     eval_analysis_space = []
@@ -22,7 +35,7 @@ def get_analysis_space(analysis_space: dict) -> list:
     for element in analysis_space:
         for key, value in element.items():
             if isinstance(value, str) and value.startswith("np."):
-                eval_element = (key, eval(value))
+                eval_element = (key, evaluate_numpy_scipy_expression(value))
             elif isinstance(value, str):
                 eval_element = (
                     key,

From 1356abb6db6b8b535f7ec9aba8825cbb414a1af0 Mon Sep 17 00:00:00 2001
From: dachengx <dx2227@columbia.edu>
Date: Sat, 5 Aug 2023 21:52:04 -0500
Subject: [PATCH 13/21] Demo future callable uncertainty

---
 alea/examples/configs/unbinned_wimp_statistical_model.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/alea/examples/configs/unbinned_wimp_statistical_model.yaml b/alea/examples/configs/unbinned_wimp_statistical_model.yaml
index fe76c736..62269892 100644
--- a/alea/examples/configs/unbinned_wimp_statistical_model.yaml
+++ b/alea/examples/configs/unbinned_wimp_statistical_model.yaml
@@ -61,7 +61,7 @@ parameter_definition:
   er_band_shift:
     nominal_value: 0
     ptype: shape
-    uncertainty: null  # scipy.stats.uniform(loc=0, scale=2)
+    uncertainty: null  # stats.uniform(loc=0, scale=2).logpdf
     # relative_uncertainty: false
     fittable: true
     blueice_anchors:

From 52c2681fc810f4332b9ded2e8d08ac9ece815595 Mon Sep 17 00:00:00 2001
From: dachengx <dx2227@columbia.edu>
Date: Tue, 8 Aug 2023 10:54:25 -0500
Subject: [PATCH 14/21] Use latest numpy

---
 requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/requirements.txt b/requirements.txt
index 388c7650..27a03eb8 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,7 +7,7 @@ iminuit>=2.21.0
 matplotlib
 mergedeep
 multihist
-numpy<1.23.0
+numpy
 PyYAML
 scipy
 setuptools

From bf29d490b8fe77488e256e51aec5b457b2a18bf4 Mon Sep 17 00:00:00 2001
From: dachengx <dx2227@columbia.edu>
Date: Thu, 10 Aug 2023 13:25:43 -0500
Subject: [PATCH 15/21] Rename histogram_multiplier to histogram_scale_factor

---
 alea/examples/configs/unbinned_wimp_statistical_model.yaml | 3 +--
 alea/template_source.py                                    | 5 -----
 2 files changed, 1 insertion(+), 7 deletions(-)

diff --git a/alea/examples/configs/unbinned_wimp_statistical_model.yaml b/alea/examples/configs/unbinned_wimp_statistical_model.yaml
index 62269892..8bc9b4b2 100644
--- a/alea/examples/configs/unbinned_wimp_statistical_model.yaml
+++ b/alea/examples/configs/unbinned_wimp_statistical_model.yaml
@@ -163,8 +163,7 @@ likelihood_config:
         weight_names: [fixed_weight, er_band_shift]  # not meaningful, just an example
         histnames: [er_template, er_template, er_template]
         template_filenames: [er_template_0.h5, er_template_1.h5, er_template_-1.h5]
-        histogram_multiplier: 100  # absolute rate, /year
-        histogram_scale_factor: 2
+        histogram_scale_factor: 100  # absolute rate, /year
 
       - name: wimp
         histname: wimp_template
diff --git a/alea/template_source.py b/alea/template_source.py
index 57be00db..48624f82 100644
--- a/alea/template_source.py
+++ b/alea/template_source.py
@@ -230,7 +230,6 @@ class CombinedSource(TemplateSource):
         weights: Weights of the 2nd to the last histograms.
         histnames: List of filenames containing the histograms.
         templatenames: List of names of histograms within the hdf5 files.
-        histogram_multiplier: Absolute rate of the combined template before multihist slicing.
     """
 
     def build_histogram(self):
@@ -311,11 +310,7 @@ def build_histogram(self):
             raise AssertionError(
                 f"There are bins for source {templatename} with negative entries.")
 
-        logging.debug(
-            "Normalising combined template, "
-            "the absolute rate is defined in histogram_multiplier")
         h = h / h.n
-        h *= self.config.get("histogram_multiplier", 1)
         logging.debug("Applying slice fraction to combined template")
         h *= slice_fractions[0]
 

From b3955f3cbade1b034105ce716641c806609909c2 Mon Sep 17 00:00:00 2001
From: dachengx <dx2227@columbia.edu>
Date: Wed, 16 Aug 2023 22:55:22 +0800
Subject: [PATCH 16/21] Test template source separately in another test module

---
 ...1_spectrum.json => test_cs1_spectrum.json} |  2 +-
 .../configs/test_template_source.yaml         | 68 +++++++++++++++++++
 .../unbinned_wimp_statistical_model.yaml      | 66 ------------------
 tests/test_template_source.py                 | 15 ++--
 4 files changed, 80 insertions(+), 71 deletions(-)
 rename alea/examples/configs/{cs1_spectrum.json => test_cs1_spectrum.json} (98%)
 create mode 100644 alea/examples/configs/test_template_source.yaml

diff --git a/alea/examples/configs/cs1_spectrum.json b/alea/examples/configs/test_cs1_spectrum.json
similarity index 98%
rename from alea/examples/configs/cs1_spectrum.json
rename to alea/examples/configs/test_cs1_spectrum.json
index 5a9d7e95..679a5a20 100644
--- a/alea/examples/configs/cs1_spectrum.json
+++ b/alea/examples/configs/test_cs1_spectrum.json
@@ -7,4 +7,4 @@
         1.0,
         1.0
     ]
-}
\ No newline at end of file
+}
diff --git a/alea/examples/configs/test_template_source.yaml b/alea/examples/configs/test_template_source.yaml
new file mode 100644
index 00000000..f35df827
--- /dev/null
+++ b/alea/examples/configs/test_template_source.yaml
@@ -0,0 +1,68 @@
+likelihood_config:
+  template_folder: null  # will try to find the templates in alea
+  likelihood_terms:
+    # SR2
+    - name: sr2
+      default_source_class: alea.template_source.TemplateSource
+      likelihood_type: blueice.likelihood.UnbinnedLogLikelihood
+      analysis_space:
+        - cs1: np.linspace(0, 100, 51)
+        - cs2: np.geomspace(100, 100000, 51)
+      in_events_per_bin: true
+      livetime_parameter: livetime_sr2
+      slice_args: {}
+      sources:
+      - name: er
+        class: alea.template_source.CombinedSource
+        parameters:
+          - er_rate_multiplier
+          - er_band_shift
+        named_parameters:
+          - er_band_shift
+        fixed_weight: 0.2
+        weight_names: [fixed_weight, er_band_shift]  # not meaningful, just an example
+        histnames: [er_template, er_template, er_template]
+        template_filenames: [er_template_0.h5, er_template_1.h5, er_template_-1.h5]
+        histogram_scale_factor: 100  # absolute rate, /year
+
+      - name: wimp
+        histname: wimp_template
+        parameters:
+          - wimp_rate_multiplier
+          - wimp_mass
+          - signal_efficiency
+        template_filename: wimp50gev_template.h5
+        apply_efficiency: True
+        efficiency_name: signal_efficiency
+
+    # SR3
+    - name: sr3
+      default_source_class: alea.template_source.TemplateSource
+      likelihood_type: blueice.likelihood.UnbinnedLogLikelihood
+      analysis_space:
+        - cs1: np.linspace(0, 100, 51)
+        - cs2: np.geomspace(100, 100000, 51)
+      in_events_per_bin: true
+      livetime_parameter: livetime_sr3
+      slice_args: {}
+      sources:
+      - name: er
+        histname: er_template
+        parameters:
+          - er_rate_multiplier
+          - er_band_shift
+        named_parameters:
+          - er_band_shift
+        template_filename: er_template_{er_band_shift}.h5
+
+      - name: wimp
+        class: alea.template_source.SpectrumTemplateSource
+        histname: wimp_template
+        parameters:
+          - wimp_rate_multiplier
+          - wimp_mass
+          - signal_efficiency
+        template_filename: wimp50gev_template.h5
+        spectrum_name: test_cs1_spectrum.json
+        apply_efficiency: True
+        efficiency_name: signal_efficiency
diff --git a/alea/examples/configs/unbinned_wimp_statistical_model.yaml b/alea/examples/configs/unbinned_wimp_statistical_model.yaml
index d8413402..c929305f 100644
--- a/alea/examples/configs/unbinned_wimp_statistical_model.yaml
+++ b/alea/examples/configs/unbinned_wimp_statistical_model.yaml
@@ -142,69 +142,3 @@ likelihood_config:
         template_filename: wimp50gev_template.h5
         apply_efficiency: True
         efficiency_name: signal_efficiency
-
-    # SR2
-    - name: sr2
-      default_source_class: alea.template_source.TemplateSource
-      likelihood_type: blueice.likelihood.UnbinnedLogLikelihood
-      analysis_space:
-        - cs1: np.linspace(0, 100, 51)
-        - cs2: np.geomspace(100, 100000, 51)
-      in_events_per_bin: true
-      livetime_parameter: livetime_sr2
-      slice_args: {}
-      sources:
-      - name: er
-        class: alea.template_source.CombinedSource
-        parameters:
-          - er_rate_multiplier
-          - er_band_shift
-        named_parameters:
-          - er_band_shift
-        fixed_weight: 0.2
-        weight_names: [fixed_weight, er_band_shift]  # not meaningful, just an example
-        histnames: [er_template, er_template, er_template]
-        template_filenames: [er_template_0.h5, er_template_1.h5, er_template_-1.h5]
-        histogram_scale_factor: 100  # absolute rate, /year
-
-      - name: wimp
-        histname: wimp_template
-        parameters:
-          - wimp_rate_multiplier
-          - wimp_mass
-          - signal_efficiency
-        template_filename: wimp50gev_template.h5
-        apply_efficiency: True
-        efficiency_name: signal_efficiency
-
-    # SR3
-    - name: sr3
-      default_source_class: alea.template_source.TemplateSource
-      likelihood_type: blueice.likelihood.UnbinnedLogLikelihood
-      analysis_space:
-        - cs1: np.linspace(0, 100, 51)
-        - cs2: np.geomspace(100, 100000, 51)
-      in_events_per_bin: true
-      livetime_parameter: livetime_sr3
-      slice_args: {}
-      sources:
-      - name: er
-        histname: er_template
-        parameters:
-          - er_rate_multiplier
-          - er_band_shift
-        named_parameters:
-          - er_band_shift
-        template_filename: er_template_{er_band_shift}.h5
-
-      - name: wimp
-        class: alea.template_source.SpectrumTemplateSource
-        histname: wimp_template
-        parameters:
-          - wimp_rate_multiplier
-          - wimp_mass
-          - signal_efficiency
-        template_filename: wimp50gev_template.h5
-        spectrum_name: cs1_spectrum.json
-        apply_efficiency: True
-        efficiency_name: signal_efficiency
diff --git a/tests/test_template_source.py b/tests/test_template_source.py
index 1d73a0f0..05021690 100644
--- a/tests/test_template_source.py
+++ b/tests/test_template_source.py
@@ -1,10 +1,17 @@
 from unittest import TestCase
 
+from alea.models import BlueiceExtendedModel
+from alea.utils import load_yaml
+
 
 class TestTemplateSource(TestCase):
     """Test of the TemplateSource class."""
 
-    def test_load_templates(self):
-        # TODO: not sure whether we want to test TemplateSource in alea
-        # not sure TemplateSource.build_histogram and TemplateSource.simulate are called
-        pass
+    def test_init_templates(self):
+        """Test whether we can initialize template sources."""
+        parameter_definition = load_yaml("unbinned_wimp_statistical_model.yaml")[
+            "parameter_definition"
+        ]
+        likelihood_config = load_yaml("test_template_source.yaml")["likelihood_config"]
+        model = BlueiceExtendedModel(parameter_definition, likelihood_config)
+        model.nominal_expectation_values

From b063d6ef1f3b21ea496cfb304b079b545726b05b Mon Sep 17 00:00:00 2001
From: dachengx <dx2227@columbia.edu>
Date: Wed, 16 Aug 2023 23:07:01 +0800
Subject: [PATCH 17/21] Remove double quotes if not necessary

---
 .../unbinned_wimp_statistical_model.yaml      |  2 +-
 ...nbinned_wimp_statistical_model_simple.yaml |  4 +-
 ...atistical_model_template_source_test.yaml} | 67 +++++++++++++++++++
 3 files changed, 70 insertions(+), 3 deletions(-)
 rename alea/examples/configs/{test_template_source.yaml => unbinned_wimp_statistical_model_template_source_test.yaml} (62%)

diff --git a/alea/examples/configs/unbinned_wimp_statistical_model.yaml b/alea/examples/configs/unbinned_wimp_statistical_model.yaml
index c929305f..fee3a78c 100644
--- a/alea/examples/configs/unbinned_wimp_statistical_model.yaml
+++ b/alea/examples/configs/unbinned_wimp_statistical_model.yaml
@@ -63,7 +63,7 @@ parameter_definition:
   er_band_shift:
     nominal_value: 0
     ptype: shape
-    uncertainty: null  # stats.uniform(loc=0, scale=2).logpdf
+    uncertainty: null  # stats.uniform(loc=0, scale=2)
     # relative_uncertainty: false
     fittable: true
     blueice_anchors:
diff --git a/alea/examples/configs/unbinned_wimp_statistical_model_simple.yaml b/alea/examples/configs/unbinned_wimp_statistical_model_simple.yaml
index 1226c5aa..f99ff47f 100644
--- a/alea/examples/configs/unbinned_wimp_statistical_model_simple.yaml
+++ b/alea/examples/configs/unbinned_wimp_statistical_model_simple.yaml
@@ -55,8 +55,8 @@ likelihood_config:
       default_source_class: alea.template_source.TemplateSource
       likelihood_type: blueice.likelihood.UnbinnedLogLikelihood
       analysis_space:
-        - "cs1": 'np.arange(0, 102, 2)'
-        - "cs2": 'np.geomspace(100, 100000, 51)'
+        - cs1: np.arange(0, 102, 2)
+        - cs2: np.geomspace(100, 100000, 51)
       in_events_per_bin: true
       livetime_parameter: livetime
       slice_args: {}
diff --git a/alea/examples/configs/test_template_source.yaml b/alea/examples/configs/unbinned_wimp_statistical_model_template_source_test.yaml
similarity index 62%
rename from alea/examples/configs/test_template_source.yaml
rename to alea/examples/configs/unbinned_wimp_statistical_model_template_source_test.yaml
index f35df827..d1dfb68f 100644
--- a/alea/examples/configs/test_template_source.yaml
+++ b/alea/examples/configs/unbinned_wimp_statistical_model_template_source_test.yaml
@@ -1,3 +1,70 @@
+parameter_definition:
+  wimp_mass:
+    nominal_value: 50
+    fittable: false
+    description: WIMP mass in GeV/c^2
+
+  livetime_sr2:
+    nominal_value: 0.5
+    fittable: false
+    description: Livetime of SR2 in years
+
+  livetime_sr3:
+    nominal_value: 1.0
+    fittable: false
+    description: Livetime of SR3 in years
+
+  wimp_rate_multiplier:
+    nominal_value: 1.0
+    ptype: rate
+    fittable: true
+    fit_limits:
+      - 0
+      - null
+    parameter_interval_bounds:
+      - 0
+      - null
+
+  er_rate_multiplier:
+    nominal_value: 1.0
+    ptype: rate
+    uncertainty: 0.2
+    relative_uncertainty: true
+    fittable: true
+    fit_limits:
+      - 0
+      - null
+    fit_guess: 1.0
+
+  signal_efficiency:
+    nominal_value: 1.0
+    ptype: efficiency
+    uncertainty: 0.1
+    relative_uncertainty: true
+    fittable: true
+    fit_limits:
+      - 0
+      - 10.
+    fit_guess: 1.0
+    description: Parameter to account for the uncertain signal expectation given a certain cross-section
+
+  er_band_shift:
+    nominal_value: 0
+    ptype: shape
+    uncertainty: null  # stats.uniform(loc=0, scale=2)
+    # relative_uncertainty: false
+    fittable: true
+    blueice_anchors:
+      - -2
+      - -1
+      - 0
+      - 1
+      - 2
+    fit_limits:
+      - -2
+      - 2
+    description: ER band shape parameter (shifts the ER band up and down)
+
 likelihood_config:
   template_folder: null  # will try to find the templates in alea
   likelihood_terms:

From 4f5508a1158aa45f534285fae5724c06b978b137 Mon Sep 17 00:00:00 2001
From: dachengx <dx2227@columbia.edu>
Date: Wed, 16 Aug 2023 23:08:50 +0800
Subject: [PATCH 18/21] Forget to upload this file

---
 tests/test_template_source.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/tests/test_template_source.py b/tests/test_template_source.py
index 05021690..47e70393 100644
--- a/tests/test_template_source.py
+++ b/tests/test_template_source.py
@@ -9,9 +9,8 @@ class TestTemplateSource(TestCase):
 
     def test_init_templates(self):
         """Test whether we can initialize template sources."""
-        parameter_definition = load_yaml("unbinned_wimp_statistical_model.yaml")[
-            "parameter_definition"
-        ]
-        likelihood_config = load_yaml("test_template_source.yaml")["likelihood_config"]
+        model_configs = load_yaml("unbinned_wimp_statistical_model_template_source_test.yaml")
+        parameter_definition = model_configs["parameter_definition"]
+        likelihood_config = model_configs["likelihood_config"]
         model = BlueiceExtendedModel(parameter_definition, likelihood_config)
         model.nominal_expectation_values

From 68d2a2a97e88655c490ffb33e25cea8408beaa28 Mon Sep 17 00:00:00 2001
From: Knut Dundas Moraa <knutdundasmoraa@gmail.com>
Date: Wed, 16 Aug 2023 15:45:14 -0400
Subject: [PATCH 19/21] tiny typo

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index d5174f0c..0b929e10 100644
--- a/README.md
+++ b/README.md
@@ -28,7 +28,7 @@ You are now ready to use alea!
 ## Getting started
 The best way to get started is to check out the [documentation](https://alea.readthedocs.io/en/latest/) and have a look at our [tutorial notebooks](https://github.com/XENONnT/alea/tree/main/notebooks). To explore the notebooks interactively, you can use [Binder](https://mybinder.org/v2/gh/XENONnT/alea/HEAD?labpath=notebooks).
 
-## Ackgnowledgements
+## Acknowledgements
 `alea` is a public package inherited the spirits of previously private XENON likelihood definition and inference construction code `binference` that based on the blueice repo https://github.com/JelleAalbers/blueice.
 
 Binference was developed for XENON1T WIMP searches by Knut Dundas Morå, and for the first XENONnT results by Robert Hammann, Knut Dundas Morå and Tim Wolf.

From 72b4e1788d82b7ad3e3ee9dd20155ebb02c5a779 Mon Sep 17 00:00:00 2001
From: Knut Dundas Moraa <knutdundasmoraa@gmail.com>
Date: Wed, 16 Aug 2023 15:52:25 -0400
Subject: [PATCH 20/21] tiny docstring change

---
 alea/template_source.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/alea/template_source.py b/alea/template_source.py
index 2a726e4d..50ec6cfa 100644
--- a/alea/template_source.py
+++ b/alea/template_source.py
@@ -14,7 +14,7 @@
 
 
 class TemplateSource(HistogramPdfSource):
-    """A source that constructs a from a template histogram. The parameters are set in self.config.
+    """A source defined with a template histogram. The parameters are set in self.config.
     "templatename", "histname", "analysis_space" must be in self.config.
 
     Attributes:
@@ -216,8 +216,9 @@ def simulate(self, n_events):
 
 
 class CombinedSource(TemplateSource):
-    """Source combines in lists of histogram names. The first histogram is the base histogram, and
-    the rest are added to it with weights. The weights can be set as shape parameters in the config.
+    """Source that is a weighted sums of histograms. Useful e.g. for safeguard. The first histogram
+    is the base histogram, and the rest are added to it with weights. The weights can be set as
+    shape parameters in the config.
 
     Args:
         weights: Weights of the 2nd to the last histograms.

From acbd7db13af539885911a700782ed1775c0ae2be Mon Sep 17 00:00:00 2001
From: dachengx <dx2227@columbia.edu>
Date: Thu, 17 Aug 2023 04:59:20 +0800
Subject: [PATCH 21/21] Fix little typo

---
 alea/examples/configs/unbinned_wimp_statistical_model.yaml      | 2 +-
 .../unbinned_wimp_statistical_model_template_source_test.yaml   | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/alea/examples/configs/unbinned_wimp_statistical_model.yaml b/alea/examples/configs/unbinned_wimp_statistical_model.yaml
index fee3a78c..f1357c13 100644
--- a/alea/examples/configs/unbinned_wimp_statistical_model.yaml
+++ b/alea/examples/configs/unbinned_wimp_statistical_model.yaml
@@ -63,7 +63,7 @@ parameter_definition:
   er_band_shift:
     nominal_value: 0
     ptype: shape
-    uncertainty: null  # stats.uniform(loc=0, scale=2)
+    uncertainty: null  # stats.uniform(loc=-2, scale=4)
     # relative_uncertainty: false
     fittable: true
     blueice_anchors:
diff --git a/alea/examples/configs/unbinned_wimp_statistical_model_template_source_test.yaml b/alea/examples/configs/unbinned_wimp_statistical_model_template_source_test.yaml
index d1dfb68f..52a9eaec 100644
--- a/alea/examples/configs/unbinned_wimp_statistical_model_template_source_test.yaml
+++ b/alea/examples/configs/unbinned_wimp_statistical_model_template_source_test.yaml
@@ -51,7 +51,7 @@ parameter_definition:
   er_band_shift:
     nominal_value: 0
     ptype: shape
-    uncertainty: null  # stats.uniform(loc=0, scale=2)
+    uncertainty: null  # stats.uniform(loc=-2, scale=4)
     # relative_uncertainty: false
     fittable: true
     blueice_anchors: