Merge pull request #252 from hiyoneda/line_background_estimation

Line background estimation

cosipy/__init__.py

@@ -15,3 +15,5 @@
 from .ts_map import FastTSMap
 
 from .source_injector import SourceInjector
+
+from .background_estimation import LineBackgroundEstimation

cosipy/background_estimation/LineBackgroundEstimation.py

@@ -0,0 +1,257 @@
+import logging
+logger = logging.getLogger(__name__)
+
+from histpy import Histogram, Axis, Axes
+import astropy.units as u
+import numpy as np
+import matplotlib.pyplot as plt
+from scipy.optimize import curve_fit
+from scipy import integrate
+from iminuit import Minuit
+
+class LineBackgroundEstimation:
+    """
+    A class for estimating and modeling background in line spectra.
+
+    This class provides methods for setting up a background model,
+    fitting it to data, and generating background model histograms.
+
+    Attributes
+    ----------
+    event_histogram : Histogram
+        The input event histogram.
+    energy_axis : Axis
+        The energy axis of the event histogram.
+    energy_spectrum : Histogram
+        The projected energy spectrum.
+    bkg_spectrum_model : callable
+        The background spectrum model function.
+    bkg_spectrum_model_parameter : list
+        The parameters of the background spectrum model.
+    mask : ndarray
+        Boolean mask for excluding regions from the fit.
+    """
+
+    def __init__(self, event_histogram):
+        """
+        Initialize the LineBackgroundEstimation object.
+
+        Parameters
+        ----------
+        event_histogram : Histogram
+            The input event histogram.
+        """
+        # event histogram
+        self.event_histogram = event_histogram
+
+        # projected histogram onto the energy axis
+        self.energy_axis = self.event_histogram.axes['Em']
+        self.energy_spectrum = self.event_histogram.project('Em')
+        if self.energy_spectrum.is_sparse:
+            self.energy_spectrum = self.energy_spectrum.to_dense()
+
+        self.energy_spectrum.clear_underflow_and_overflow()
+
+        # background fitting model
+        self.bkg_spectrum_model = None
+        self.bkg_spectrum_model_parameter = None
+
+        # bins to be masked
+        self.mask = np.zeros(self.energy_axis.nbins, dtype=bool)
+
+    def set_bkg_energy_spectrum_model(self, bkg_spectrum_model, bkg_spectrum_model_parameter):
+        """
+        Set the background energy spectrum model and its initial parameters.
+
+        Parameters
+        ----------
+        bkg_spectrum_model : callable
+            The background spectrum model function.
+        bkg_spectrum_model_parameter : list
+            Initial parameters for the background spectrum model.
+        """
+        self.bkg_spectrum_model = bkg_spectrum_model
+        self.bkg_spectrum_model_parameter = bkg_spectrum_model_parameter
+
+    def set_mask(self, *mask_energy_ranges):
+        """
+        Set mask for excluding energy ranges from the fit.
+
+        Parameters
+        ----------
+        *mask_energy_ranges : tuple
+            Variable number of energy range tuples to be masked.
+        """
+        self.mask = np.zeros(self.energy_axis.nbins, dtype=bool)
+        for mask_energy_range in mask_energy_ranges:
+            this_mask = (mask_energy_range[0] <= self.energy_axis.bounds[:, 1]) & (self.energy_axis.bounds[:, 0] <= mask_energy_range[1])
+            self.mask = self.mask | this_mask
+
+    def _calc_expected_spectrum(self, *args):
+        """
+        Calculate the expected spectrum based on the current model and parameters.
+
+        Parameters
+        ----------
+        *args : float
+            Model parameters.
+
+        Returns
+        -------
+        ndarray
+            The calculated expected spectrum.
+        """
+        return np.array([integrate.quad(lambda x: self.bkg_spectrum_model(x, *args), *energy_range)[0] for energy_range in self.energy_axis.bounds.value])
+
+    def _negative_log_likelihood(self, *args):
+        """
+        Calculate the negative log-likelihood for the current model and parameters.
+
+        Parameters
+        ----------
+        *args : float
+            Model parameters.
+
+        Returns
+        -------
+        float
+            The calculated negative log-likelihood.
+        """
+        expected_spectrum = self._calc_expected_spectrum(*args)
+        return -np.sum(self.energy_spectrum.contents[~self.mask] * np.log(expected_spectrum)[~self.mask]) + np.sum(expected_spectrum[~self.mask])
+
+    def plot_energy_spectrum(self):
+        """
+        Plot the energy spectrum and the fitted model if available.
+
+        Returns
+        -------
+        tuple
+            A tuple containing the matplotlib axis object and any additional objects returned by the plotting function.
+        """
+        ax, _ = self.energy_spectrum.draw(label='input data')
+
+        # plot background model
+        if self.bkg_spectrum_model is not None:
+            expected_spectrum = self._calc_expected_spectrum(*self.bkg_spectrum_model_parameter)
+            ax.plot(self.energy_axis.centers, expected_spectrum, label='model')
+
+        # shade mask regions
+        start, end = None, None
+        for i, this_mask in enumerate(self.mask):
+            if this_mask:
+                if start is None:
+                    start, end = self.energy_axis.bounds[i]
+                else:
+                    _, end = self.energy_axis.bounds[i]
+            else:
+                if start is not None:
+                    ax.axvspan(start.value, end.value, color='lightgrey', alpha=0.5)
+                    start, end = None, None
+
+        if start is not None:
+            ax.axvspan(start.value, end.value, color='lightgrey', alpha=0.5)
+
+        # legend and grid
+        ax.legend()
+        ax.grid()
+
+        return ax, _
+
+    def fit_energy_spectrum(self):
+        """
+        Fit the background energy spectrum model to the data.
+
+        Returns
+        -------
+        Minuit
+            The Minuit object containing the fit results.
+        """
+        m = Minuit(self._negative_log_likelihood, *self.bkg_spectrum_model_parameter)
+        m.errordef = Minuit.LIKELIHOOD
+
+        m.migrad()
+        m.hesse()
+
+        # update the background model parameters
+        self.bkg_spectrum_model_parameter = list(m.values)
+        self.bkg_spectrum_model_parameter_errors = list(m.errors)
+
+        return m
+
+    def _get_weight_indices(self, energy_range):
+        """
+        Get the weight and indices for a given energy range.
+
+        Parameters
+        ----------
+        energy_range : tuple
+            The energy range to calculate the weight for.
+
+        Returns
+        -------
+        tuple
+            A tuple containing the calculated weight and the corresponding energy indices.
+        """
+        energy_indices = np.where((energy_range[0] <= self.energy_axis.lower_bounds) & (self.energy_axis.upper_bounds <= energy_range[1]))[0]
+
+        if len(energy_indices) == 0:
+            raise ValueError("The input energy range is too narrow to find a corresponding energy bin.")
+
+        integrate_energy_range = [self.energy_axis.lower_bounds[energy_indices[0]].value, self.energy_axis.upper_bounds[energy_indices[-1]].value]
+
+        if integrate_energy_range[0] != energy_range[0].value or integrate_energy_range[1] != energy_range[1].value:
+            logger.info(f"The energy range {energy_range.value} is modified to {integrate_energy_range}")
+        weight = integrate.quad(lambda x: self.bkg_spectrum_model(x, *self.bkg_spectrum_model_parameter), *integrate_energy_range)[0]
+        return weight, energy_indices
+
+    def generate_bkg_model_histogram(self, source_energy_range, bkg_estimation_energy_ranges):
+        """
+        Generate a background model histogram based on the fitted model.
+
+        Parameters
+        ----------
+        bkg_estimation_energy_ranges : list of tuple
+            List of energy ranges for background estimation.
+        smoothing_fwhm : float, optional
+            Full width at half maximum for smoothing, by default None.
+
+        Returns
+        -------
+        Histogram
+            The generated background model histogram.
+        """
+        # intergrated spectrum in the background estimation energy ranges
+        weights = []
+        energy_indices_list = []
+        for bkg_estimation_energy_range in bkg_estimation_energy_ranges:
+            weight, energy_indices = self._get_weight_indices(bkg_estimation_energy_range)
+            weights.append(weight)
+            energy_indices_list.append(energy_indices)
+
+        # intergrated spectrum in the source region
+        source_weight = integrate.quad(lambda x: self.bkg_spectrum_model(x, *self.bkg_spectrum_model_parameter), *source_energy_range.value)[0]
+
+        # prepare a new histogram
+        new_axes = []
+        for axis in self.event_histogram.axes:
+            if axis.label != "Em":
+                new_axes.append(axis)
+            else:
+                new_axes.append(Axis(source_energy_range, label = "Em"))
+
+        bkg_model_histogram = Histogram(new_axes)
+
+        # fill contents
+        for energy_indices in energy_indices_list:
+            for energy_index in energy_indices:
+                if new_axes[0].label != "Em":
+                    bkg_model_histogram[:] += self.event_histogram[:,energy_index].todense()
+                else:
+                    bkg_model_histogram[:] += self.event_histogram[energy_index].todense()
+
+        # normalization
+        corr_factor = source_weight / np.sum(weights)
+        bkg_model_histogram[:] *= corr_factor
+
+        return bkg_model_histogram

cosipy/background_estimation/__init__.py

@@ -0,0 +1 @@
+from .LineBackgroundEstimation import LineBackgroundEstimation

docs/tutorials/background_estimation/line_background/inputs_bkg_estimation.yaml

@@ -0,0 +1,39 @@
+#----------#
+# Data I/O:
+
+# data files available on the COSI Sharepoint: https://drive.google.com/drive/folders/1UdLfuLp9Fyk4dNussn1wt7WEOsTWrlQ6
+data_file: "/scratch/astrohome/smittal/bkg_data_files/total_bg_3months_unbinned_data.fits.gz" #"GalacticScan.inc1.id1.crab2hr.extracted.tra.gz" # full path
+ori_file: "/scratch/astrohome/smittal/20280301_3_month.ori" # full path
+unbinned_output: 'fits' # 'fits' or 'hdf5'
+time_bins: 7776000 # time bin size in seconds. Takes int, float, or list of bin edges.
+energy_bins: [1500.        , 1505.05050505, 1510.1010101 , 1515.15151515,
+       1520.2020202 , 1525.25252525, 1530.3030303 , 1535.35353535,
+       1540.4040404 , 1545.45454545, 1550.50505051, 1555.55555556,
+       1560.60606061, 1565.65656566, 1570.70707071, 1575.75757576,
+       1580.80808081, 1585.85858586, 1590.90909091, 1595.95959596,
+       1601.01010101, 1606.06060606, 1611.11111111, 1616.16161616,
+       1621.21212121, 1626.26262626, 1631.31313131, 1636.36363636,
+       1641.41414141, 1646.46464646, 1651.51515152, 1656.56565657,
+       1661.61616162, 1666.66666667, 1671.71717172, 1676.76767677,
+       1681.81818182, 1686.86868687, 1691.91919192, 1696.96969697,
+       1702.02020202, 1707.07070707, 1712.12121212, 1717.17171717,
+       1722.22222222, 1727.27272727, 1732.32323232, 1737.37373737,
+       1742.42424242, 1747.47474747, 1752.52525253, 1757.57575758,
+       1762.62626263, 1767.67676768, 1772.72727273, 1777.77777778,
+       1782.82828283, 1787.87878788, 1792.92929293, 1797.97979798,
+       1803.03030303, 1808.08080808, 1813.13131313, 1818.18181818,
+       1823.23232323, 1828.28282828, 1833.33333333, 1838.38383838,
+       1843.43434343, 1848.48484848, 1853.53535354, 1858.58585859,
+       1863.63636364, 1868.68686869, 1873.73737374, 1878.78787879,
+       1883.83838384, 1888.88888889, 1893.93939394, 1898.98989899,
+       1904.04040404, 1909.09090909, 1914.14141414, 1919.19191919,
+       1924.24242424, 1929.29292929, 1934.34343434, 1939.39393939,
+       1944.44444444, 1949.49494949, 1954.54545455, 1959.5959596 ,
+       1964.64646465, 1969.6969697 , 1974.74747475, 1979.7979798 ,
+       1984.84848485, 1989.8989899 , 1994.94949495, 2000.        ] #[1500., 1550., 1600., 1650., 1700., 1750., 1800., 1850., 1900., 1950., 2000.] #[100.,  200.,  500., 1000., 2000., 5000.] # Takes list. Needs to match response.
+phi_pix_size: 3 # binning of Compton scattering anlge [deg]
+nside: 16 # healpix binning of psi chi local
+scheme: 'ring' # healpix binning of psi chi local
+tmin: 1835478000.0 # Min time cut in seconds.
+tmax: 1843254000.0 # Max time cut in seconds.
+#----------#

docs/tutorials/background_estimation/line_background/inputs_check_results.yaml

@@ -0,0 +1,15 @@
+#----------#
+# Data I/O:
+
+# data files available on the COSI Sharepoint: https://drive.google.com/drive/folders/1UdLfuLp9Fyk4dNussn1wt7WEOsTWrlQ6
+data_file: "/scratch/astrohome/smittal/bkg_data_files/total_bg_3months_unbinned_data.fits.gz" #"GalacticScan.inc1.id1.crab2hr.extracted.tra.gz" # full path
+ori_file: "/scratch/astrohome/smittal/20280301_3_month.ori" # full path
+unbinned_output: 'fits' # 'fits' or 'hdf5'
+time_bins: 7776000 # time bin size in seconds. Takes int, float, or list of bin edges.
+energy_bins: [1805.0, 1812.0]
+phi_pix_size: 3 # binning of Compton scattering anlge [deg]
+nside: 16 # healpix binning of psi chi local
+scheme: 'ring' # healpix binning of psi chi local
+tmin: 1835478000.0 # Min time cut in seconds.
+tmax: 1843254000.0 # Max time cut in seconds.
+#----------#