Add sources and example notebook for 1D CES inference

.gitignore

@@ -29,3 +29,4 @@ debug.py
 docs/source/reference/release_notes.rst
 .vscode
 .hypothesis
+.venv

alea/ces_source.py

@@ -0,0 +1,234 @@
+from typing import Dict, Literal
+import numpy as np
+from scipy.interpolate import interp1d
+
+from inference_interface import template_to_multihist
+from blueice import HistogramPdfSource, Source
+from blueice.exceptions import PDFNotComputedException
+
+from multihist import Hist1d
+from alea.ces_transformation import Transformation
+
+MINIMAL_ENERGY_RESOLUTION = 0.05
+
+
+class CESTemplateSource(HistogramPdfSource):
+    def __init__(self, config: Dict, *args, **kwargs):
+        """Initialize the TemplateSource."""
+        # override the default interpolation method
+        if "pdf_interpolation_method" not in config:
+            config["pdf_interpolation_method"] = "piecewise"
+        super().__init__(config, *args, **kwargs)
+
+    def _load_inputs(self):
+        self.ces_space = self.config["analysis_space"][0][1]
+        self.max_e = np.max(self.ces_space)
+        self.min_e = np.min(self.ces_space)
+        self.templatename = self.config["templatename"]
+        self.histname = self.config["histname"]
+
+    def _load_true_histogram(self):
+        h = template_to_multihist(self.templatename, self.histname, hist_to_read=Hist1d)
+        return h
+
+    def _check_histogram(self, h: Hist1d):
+        """Check if the histogram has expected binning."""
+        # We only take 1d histogram in the ces axes
+        if not isinstance(h, Hist1d):
+            raise ValueError("Only Hist1d object is supported")
+        if self.ces_space.ndim != 1:
+            raise ValueError("Only 1d analysis space is supported")
+        if np.min(h.histogram) < 0:
+            raise AssertionError(
+                f"There are bins for source {self.templatename} with negative entries."
+            )
+
+        # check if the histogram contains the analysis space.
+        histogram_max = np.max(h.bin_edges)
+        histogram_min = np.min(h.bin_edges)
+        if self.min_e < histogram_min or self.max_e > histogram_max:
+            raise ValueError(
+                f"The histogram edge ({histogram_min},{histogram_max}) \
+                does not contain the analysis space ({self.min_e},{self.max_e})"
+            )
+
+    def _create_transformation(
+        self, transformation_type: Literal["efficiency", "smearing", "bias"]
+    ):
+        if self.config.get(f"apply_{transformation_type}", True):
+            parameters_key = f"{transformation_type}_parameters"
+            model_key = f"{transformation_type}_model"
+
+            if model_key not in self.config:
+                raise ValueError(f"{transformation_type.capitalize()} model is not provided")
+
+            if parameters_key not in self.config:
+                raise ValueError(f"{transformation_type.capitalize()} parameters are not provided")
+            else:
+                parameter_list = self.config[parameters_key]
+                # to get the values we need to iterate over the list and use self.config.get
+                combined_parameter_dict = {k: self.config.get(k) for k in parameter_list}
+
+            # Also take the peak_energy parameter if it is a mono smearing model
+            if "mono" in self.config[model_key]:
+                combined_parameter_dict["peak_energy"] = self.config["peak_energy"]
+
+            return Transformation(
+                parameters=combined_parameter_dict,
+                action=transformation_type,
+                model=self.config[model_key],
+            )
+        return None
+
+    def _transform_histogram(self, h: Hist1d):
+        # Create transformations for efficiency, smearing, and bias
+        efficiency_transformation = self._create_transformation("efficiency")
+        smearing_transformation = self._create_transformation("smearing")
+        bias_transformation = self._create_transformation("bias")
+
+        # Apply the transformations to the histogram
+        if efficiency_transformation is not None:
+            h = efficiency_transformation.apply_transformation(h)
+        if smearing_transformation is not None:
+            h = smearing_transformation.apply_transformation(h)
+        if bias_transformation is not None:
+            h = bias_transformation.apply_transformation(h)
+        return h
+
+    def _normalize_histogram(self, h: Hist1d):
+        # To avoid confusion, we always normalize the histogram, regardless of the bin volume
+        # So the unit is always events/year/keV, the rate multipliers are always in terms of that
+        total_integration = np.sum(h.histogram * h.bin_volumes())
+        h.histogram /= total_integration
+
+        # Apply the transformations to the histogram
+        h = self._transform_histogram(h)
+
+        # Calculate the integration of the histogram after all transformations
+        # Only from min_e to max_e
+        left_edges = h.bin_edges[:-1]
+        right_edges = h.bin_edges[1:]
+        outside_index = np.where((left_edges < self.min_e) | (right_edges > self.max_e))
+        h.histogram[outside_index] = 0
+
+        self._bin_volumes = h.bin_volumes()
+        self._n_events_histogram = h.similar_blank_histogram()
+
+        # Note that it already does what "fraction_in_roi" does in the old code
+        # So no need to do again
+        integration_after_transformation_in_roi = np.sum(h.histogram * h.bin_volumes())
+
+        self.events_per_year = (
+            integration_after_transformation_in_roi * self.config["rate_multiplier"]
+        )
+        self.events_per_day = self.events_per_year / 365
+
+        # For pdf, we need to normalize the histogram to 1 again
+        return h, integration_after_transformation_in_roi
+
+    def build_histogram(self):
+        """Build the histogram of the source.
+
+        It's always called during the initialization of the source. So the attributes are set here.
+
+        """
+        print("Building histogram")
+        self._load_inputs()
+        h = self._load_true_histogram()
+        self._check_histogram(h)
+        h, frac_in_roi = self._normalize_histogram(h)
+        h_pdf = h
+        h_pdf /= frac_in_roi
+        self._pdf_histogram = h_pdf
+        self.set_dtype()
+
+    def simulate(self, n_events: int):
+        dtype = [
+            ("ces", float),
+            ("source", int),
+        ]
+        ret = np.zeros(n_events, dtype=dtype)
+        ret["ces"] = self._pdf_histogram.get_random(n_events)
+        return ret
+
+    def compute_pdf(self):
+        self.build_histogram()
+        Source.compute_pdf(self)
+
+    def pdf(self, *args):
+        # override the default interpolation method in blueice (RegularGridInterpolator)
+        if not self.pdf_has_been_computed:
+            raise PDFNotComputedException(
+                "%s: Attempt to call a PDF that has not been computed" % self
+            )
+
+        method = self.config["pdf_interpolation_method"]
+
+        if method == "linear":
+            if not hasattr(self, "_pdf_interpolator"):
+                # First call:
+                # Construct a linear interpolator between the histogram bins
+                self._pdf_interpolator = interp1d(
+                    self._pdf_histogram.bin_centers,
+                    self._pdf_histogram.histogram,
+                )
+            bcs = self._pdf_histogram.bin_centers
+            clipped_data = np.clip(args, bcs.min(), bcs.max())
+
+            return self._pdf_interpolator(np.transpose(clipped_data))
+
+        elif method == "piecewise":
+            return self._pdf_histogram.lookup(*args)
+
+        else:
+            raise NotImplementedError("PDF Interpolation method %s not implemented" % method)
+
+    def set_dtype(self):
+        self.dtype = [
+            ("ces", float),
+            ("source", int),
+        ]
+
+
+class CESMonoenergySource(CESTemplateSource):
+    def _load_inputs(self):
+        self.ces_space = self.config["analysis_space"][0][1]
+        self.max_e = np.max(self.ces_space)
+        self.min_e = np.min(self.ces_space)
+        self.mu = self.config["peak_energy"]
+
+    def _load_true_histogram(self):
+        number_of_bins = int((self.max_e - self.min_e) / MINIMAL_ENERGY_RESOLUTION)
+        h = Hist1d(
+            data=np.repeat(self.mu, 1),
+            bins=number_of_bins,
+            range=(self.min_e, self.max_e),
+        )
+        h.histogram = h.histogram.astype(np.float64)
+        self.config["smearing_model"] = "mono_" + self.config["smearing_model"]
+        return h
+
+
+class CESFlatSource(CESTemplateSource):
+    def _load_inputs(self):
+        self.ces_space = self.config["analysis_space"][0][1]
+        self.max_e = np.max(self.ces_space)
+        self.min_e = np.min(self.ces_space)
+
+    def _load_true_histogram(self):
+        number_of_bins = int((self.max_e - self.min_e) / MINIMAL_ENERGY_RESOLUTION)
+        h = Hist1d(
+            data=np.linspace(self.min_e, self.max_e, number_of_bins),
+            bins=number_of_bins,
+            range=(self.min_e, self.max_e),
+        )
+        h.histogram = h.histogram.astype(np.float64)
+        return h
+
+    def _transform_histogram(self, h: Hist1d):
+        # Only efficiency is applicable for flat source
+        efficiency_transformation = self._create_transformation("efficiency")
+
+        if efficiency_transformation is not None:
+            h = efficiency_transformation.apply_transformation(h)
+        return h

alea/ces_transformation.py

@@ -0,0 +1,155 @@
+from pydantic import BaseModel, validator
+from typing import Dict, Optional, Any, Literal, Callable
+import numpy as np
+from scipy import stats
+from copy import deepcopy
+from multihist import Hist1d
+
+
+def energy_res(energy, a=25.8, b=1.429):
+    """Return energy resolution in keV.
+
+    :param energy: true energy in keV :return: energy resolution in keV
+
+    """
+    # Reference for the values of a,b:
+    # xenon:xenonnt:analysis:ntsciencerun0:g1g2_update#standard_gaussian_vs_skew-gaussian_yue
+    return (np.sqrt(energy) * a + energy * b) / 100
+
+
+def smearing_mono_gaussian(
+    hist: Any,
+    smearing_a: float,
+    smearing_b: float,
+    peak_energy: float,
+    bins: Optional[np.ndarray] = None,
+):
+
+    if bins is None:
+        # create an emptyzero histogram with the same binning as the input histogram
+        data = stats.norm.pdf(
+            hist.bin_centers,
+            loc=peak_energy,
+            scale=energy_res(peak_energy, smearing_a, smearing_b),
+        )
+        hist_smeared = Hist1d(data=np.zeros_like(data), bins=hist.bin_edges)
+        hist_smeared.histogram = data
+    else:
+        # use the bins that set by the user
+        bins = np.array(bins)
+        if bins.size <= 1:
+            raise ValueError("bins must have at least 2 elements")
+        bin_centers = 0.5 * (bins[1:] + bins[:-1])
+        data = stats.norm.pdf(
+            bin_centers,
+            loc=peak_energy,
+            scale=energy_res(peak_energy, smearing_a, smearing_b),
+        )
+        # create an empty histogram with the user-defined binning
+        hist_smeared = Hist1d(data=np.zeros_like(data), bins=bins)
+        hist_smeared.histogram = data
+
+    return hist_smeared
+
+
+def smearing_hist_gaussian(
+    hist: Any,
+    smearing_a: float,
+    smearing_b: float,
+    bins: Optional[np.ndarray] = None,
+):
+    """Smear a histogram. This allows for non-uniform histogram binning.
+
+    :param hist: the spectrum we want to smear :param bins: bin edges of the returned spectrum
+    :return: smeared histogram in the same unit as input spectrum
+
+    """
+    assert isinstance(hist, Hist1d), "Only Hist1d object is supported"
+    if bins is None:
+        # set the bins to the bin edges of the input histogram
+        bins = hist.bin_edges
+    elif bins.size <= 1:
+        raise ValueError("bins must have at least 2 elements")
+    bins = np.array(bins)
+
+    e_true_s, rates, bin_volumes = hist.bin_centers, hist.histogram, hist.bin_volumes()
+    mask = np.where(e_true_s > 0)
+    e_true_s = e_true_s[mask]
+    rates = rates[mask]
+    bin_volumes = bin_volumes[mask]
+
+    e_smeared_s = 0.5 * (bins[1:] + bins[:-1])
+    smeared = np.zeros_like(e_smeared_s)
+
+    for idx, e_smeared in enumerate(e_smeared_s):
+        probs = (
+            stats.norm.pdf(
+                e_smeared,
+                loc=e_true_s,
+                scale=energy_res(e_true_s, smearing_a, smearing_b),
+            )
+            * bin_volumes
+        )
+        smeared[idx] = np.sum(probs * rates)
+
+    hist_smeared = Hist1d.from_histogram(smeared, bins)
+
+    return hist_smeared
+
+
+def biasing_hist_arctan(hist: Any, A: float = 0.01977, k: float = 0.01707):
+    """Apply a constant bias to a histogram.
+
+    :param hist: the spectrum we want to apply the bias to :param bias: the bias to apply to the
+    spectrum :return: the spectrum with the bias applied
+
+    """
+    assert isinstance(hist, Hist1d), "Only Hist1d object is supported"
+    true_energy = hist.bin_centers
+    h_bias = deepcopy(hist)
+    bias_derivative = A * k / (1 + k**2 * true_energy**2)
+    h_bias.histogram *= 1 / (1 + bias_derivative)
+    return h_bias
+
+
+def efficiency_hist_constant(hist: Any, efficiency: float):
+    """Apply a constant efficiency to a histogram.
+
+    :param hist: the spectrum we want to apply the efficiency to :param efficiency: the efficiency
+    to apply to the spectrum :return: the spectrum with the efficiency applied
+
+    """
+    assert isinstance(hist, Hist1d), "Only Hist1d object is supported"
+    assert 0 <= efficiency <= 1, "Efficiency must be between 0 and 1"
+    hist.histogram = hist.histogram * efficiency
+    return hist
+
+
+MODELS: Dict[str, Dict[str, Callable]] = {
+    "smearing": {
+        "gaussian": smearing_hist_gaussian,
+        "mono_gaussian": smearing_mono_gaussian,
+    },
+    "bias": {"arctan": biasing_hist_arctan},
+    "efficiency": {
+        "constant": efficiency_hist_constant,
+    },
+}
+
+
+# input: model name, parameters, transformation mode
+class Transformation(BaseModel):
+    parameters: Dict[str, float]
+    action: Literal["smearing", "bias", "efficiency"]
+    model: str
+
+    @validator("model")
+    @classmethod
+    def check_model(cls, v, values):
+        if v not in MODELS[values["action"]]:
+            raise ValueError(f"Model {v} not found for action {values['action']}")
+        return v
+
+    def apply_transformation(self, histogram: Hist1d):
+        chosen_model = MODELS[self.action][self.model]
+        return chosen_model(histogram, **self.parameters)