This repository has been archived on 2024-11-23. You can view files and clone it, but cannot push or open issues or pull requests.
GBM-data-tools/background/background.py

588 lines
23 KiB
Python
Raw Normal View History

2022-07-15 15:36:07 +08:00
# background.py: Module containing background fitting and model classes
#
# Authors: William Cleveland (USRA),
# Adam Goldstein (USRA) and
# Daniel Kocevski (NASA)
#
# Portions of the code are Copyright 2020 William Cleveland and
# Adam Goldstein, Universities Space Research Association
# All rights reserved.
#
# Written for the Fermi Gamma-ray Burst Monitor (Fermi-GBM)
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
#
import numpy as np
from ..data.primitives import EventList, TimeEnergyBins, EnergyBins
from ..data import PHAII, TTE
class BackgroundFitter:
"""Class for fitting a background, given a fitting algorithm,
to time-energy data (e.g. :class:`~gbm.data.phaii.PHAII`,
:class:`~gbm.data.TTE`).
When a BackgroundFitter is created, an algorithm must be specified. In
particular, the algorithm must be a class that has two public methods:
``fit()`` and ``interpolate()``.
For PHAII data, the class, upon initialization, must take the following as
arguments:
1. a 2D array of counts with shape (``numtimes``, ``numchans``);
2. a 1D array of time bin start times, shape (``numtimes``,);
3. a 1D array of time bin end times, shape (``numtimes``,);
4. a 1D array of exposures for each time bin, shape (``numtimes``,).
While for TTE data, the class, upon initialization, must take the following
as an argument:
1. a list, of length ``numchans``, where each item is a np.array
of event times.
The ``fit()`` method takes no arguments, hoewever, parameters that are
required for the algorithm may be specified as keywords.
The interpolate() method must take in the following as arguments:
1. a 1D array of time bin start times, shape (``numtimes``,);
2. a 1D array of time bin end times, shape (``numtimes``,).
Any additional parameters required can be specified as keywords.
The ``interpolate()`` method must return:
1. a 2D rates array, shape (``numtimes``, ``numchans``);
2. a 2D rate uncertainty array, shape (``numtimes``, ``numchans``).
Additionally, the class can provide the following public attributes that
will be exposed by BackgroundFitter:
- ``dof``: The degrees of freedom of the fits, array shape (``numchans``,)
- ``statistic``: The fit statistic for each fit, array shape (``numchans``,)
- ``statistic_name``: A string of the fit statistic used
Attributes:
dof (np.array): If available, the degrees-of-freedom of the fit for
each energy channel
livetime (float): The total livetime of the data used for the background
method (str): The name of the fitting algorithm class
parameters (dict): All parameters passed to the fitting algorithm
statistic (np.array): If available, the fit statistic for each energy
channel
statistic_name (str): If available, the name of the fit statistic
type (str): The type of background algorithm, either 'binned' or 'unbinned'
"""
def __init__(self):
self._data_obj = None
self._method = None
self._type = None
self._statistic = None
self._dof = None
self._livetime = None
self._parameters = None
@property
def method(self):
return self._method.__class__.__name__
@property
def type(self):
return self._type
@property
def statistic_name(self):
return getattr(self._method, 'statistic_name', None)
@property
def statistic(self):
return getattr(self._method, 'statistic', None)
@property
def dof(self):
return getattr(self._method, 'dof', None)
@property
def livetime(self):
return self._livetime
@property
def parameters(self):
return self._parameters
@classmethod
def from_phaii(cls, phaii, method, time_ranges=None):
"""Create a background fitter from a PHAII object
Args:
phaii (:class:`~gbm.data.phaii.PHAII`): A PHAII data object
method (<class>): A background fitting/estimation class for binned data
time_ranges ([(float, float), ...]):
The time range or time ranges over which to fit the background.
If omitted, uses the full time range of the data
Returns:
:class:`BackgroundFitter`: The initialized background fitting object
"""
if not isinstance(phaii, PHAII):
raise TypeError('Input data must be a PHAII object')
obj = cls()
obj._data_obj = phaii
obj._validate_method(method)
time_ranges = obj._validate_time_ranges(time_ranges)
# Slice the PHAII data and merge if multiple slices
data = [phaii.data.slice_time(trange[0], trange[1]) for trange in
time_ranges]
data = TimeEnergyBins.merge_time(data)
obj._method = method(data.counts, data.tstart, data.tstop,
data.exposure)
obj._type = 'binned'
obj._livetime = np.sum(data.exposure)
return obj
@classmethod
def from_tte(cls, tte, method, time_ranges=None):
"""Create a background fitter from a TTE object
Args:
tte (:class:`~gbm.data.TTE`): A TTE data object
method (<class>): A background fitting/estimation class for unbinned data
time_ranges ([(float, float), ...]):
The time range or time ranges over which to fit the background.
If omitted, uses the full time range of the data
Returns:
:class:`BackgroundFitter`: The initialized background fitting object
"""
if not isinstance(tte, TTE):
raise TypeError('Input data must be a TTE object')
obj = cls()
obj._data_obj = tte
obj._validate_method(method)
time_ranges = obj._validate_time_ranges(time_ranges)
# Slice the TTE data and merge if multiple slices
data = [tte.data.time_slice(trange[0], trange[1]) for trange in
time_ranges]
data = EventList.merge(data)
data.sort('TIME')
# pull out the events in each channel
events = [data.channel_slice(i, i).time for i in
range(tte.numchans)]
obj._method = method(events)
obj._type = 'unbinned'
obj._livetime = data.get_exposure(time_ranges=time_ranges)
return obj
def fit(self, **kwargs):
"""Perform a background fit of the data
Args:
**kwargs: Options to be passed as parameters to the fitting class
"""
self._parameters = kwargs
self._method.fit(**kwargs)
def interpolate_bins(self, tstart, tstop, **kwargs):
"""Interpolate the fitted background model over a set of bins.
The exposure is calculated for each bin of the background model
in case the background model counts is needed.
Args:
tstart (np.array): The starting times
tstop (np.array): The ending times
**kwargs: Options to be passed as parameters to the interpolation method
Returns:
:class:`BackgroundRates`: The background rates/uncertainty object
"""
tstart_, tstop_ = (tstart.copy(), tstop.copy())
# do the interpolation
rate, rate_uncert = self._method.interpolate(tstart_, tstop_, **kwargs)
# get the exposure
numtimes = tstart_.shape[0]
exposure = np.array([self._data_obj.get_exposure((tstart_[i], tstop_[i])) \
for i in range(numtimes)])
# create the rates object
det = self._data_obj.detector
dtype = self._data_obj.datatype
trigtime = self._data_obj.trigtime
rates = BackgroundRates(rate, rate_uncert, tstart_, tstop_,
self._data_obj.data.emin.copy(),
self._data_obj.data.emax.copy(),
exposure=exposure,
detector=det, datatype=dtype,
trigtime=trigtime)
return rates
def interpolate_times(self, times, **kwargs):
"""Interpolate the fitted background model over a set of times.
Does not calculate an exposure since this returns a set of point
estimates of the background rates.
Args:
tstart (np.array): The sampling times
**kwargs: Options to be passed as parameters to the interpolation method
Returns:
:class:`BackgroundRates`: The background rates/uncertainty object
"""
# do the interpolation
rate, rate_uncert = self._method.interpolate(times, times, **kwargs)
# create the rates object
det = self._data_obj.detector
dtype = self._data_obj.datatype
trigtime = self._data_obj.trigtime
rates = BackgroundRates(rate, rate_uncert, times, times,
self._data_obj.data.emin.copy(),
self._data_obj.data.emax.copy(), detector=det,
datatype=dtype, trigtime=trigtime)
return rates
def _validate_method(self, method):
try:
method
except:
raise NameError('Input method is not a known function')
has_fit = callable(method.fit)
has_interp = callable(method.interpolate)
if (not has_fit) or (not has_interp):
raise NotImplementedError(
"User-defined Background class must have "
"both fit() and an interpolate() methods")
def _validate_time_ranges(self, time_ranges):
if time_ranges is None:
time_ranges = [self._data_obj.time_range]
try:
iter(time_ranges[0])
except:
raise TypeError('time_ranges must be a list of tuples')
return time_ranges
class BackgroundRates(TimeEnergyBins):
"""Class containing the background rate data.
Parameters:
rates (np.array): The array of background rates in each bin
rate_uncertainty (np.array): The array of background rate uncertainties
in each bin
tstart (np.array): The low-value edges of the time bins
tstop (np.array): The high-value edges of the time bins
emin (np.array): The low-value edges of the energy bins
emax (np.array): The high-value edges of the energy bins
exposure (np.array, optional): The exposure of each bin
detector (str, optional): The associated detector
datatype (str, optional): The datatype associated with the background
trigtime (float, optional): The trigger time
Attributes:
chanwidths (np.array): The bin widths along the energy axis
count_uncertainty (np.array): The counts uncertainty in each bin
energy_centroids (np.array): The bin centroids along the energy axis
energy_range (float, float): The range of the data along the energy axis
numchans (int): The number of energy channels along the energy axis
numtimes (int): The number of bins along the time axis
rates (np.array): The rates in each Time-Energy Bin
rates_per_kev (np.array): The differential rates in units of counts/s/keV
rate_uncertainty (np.array): The rate uncertainty in each bin
rate_uncertainty_per_kev (np.array):
The differential rate uncertainty in units of counts/s/keV
size (int, int): The number of bins along both axes (numtimes, numchans)
time_centroids (np.array): The bin centroids along the time axis
time_range (float, float): The range of the data along the time axis
time_widths (np.array): The bin widths along the time axis
"""
def __init__(self, rates, rate_uncertainty, tstart, tstop, emin, emax,
exposure=None, detector=None, datatype=None, trigtime=None):
if exposure is None:
exposure = np.zeros_like(tstart)
counts = np.squeeze(rates * exposure[:, np.newaxis])
super().__init__(counts, tstart, tstop, exposure, emin, emax)
self._count_uncertainty = np.squeeze(rate_uncertainty * exposure[:, np.newaxis])
self._rates = rates.squeeze()
self._rate_uncertainty = rate_uncertainty.squeeze()
self._detector = detector
self._datatype = datatype
self._trigtime = trigtime
@property
def count_uncertainty(self):
return self._count_uncertainty
@property
def rates(self):
return self._rates
@property
def rate_uncertainty(self):
return self._rate_uncertainty
def integrate_energy(self, emin=None, emax=None):
"""Integrate the over the energy axis.
Limits on the integration smaller than the full range can be set.
Args:
emin (float, optional): The low end of the integration range. If not
set, uses the lowest energy edge of the histogram
emax (float, optional): The high end of the integration range. If not
set, uses the highest energy edge of the histogram
Returns:
:class:`gbm.data.primitives.TimeBins`: A TimeBins object containing \
the lightcurve histogram
"""
if emin is None:
emin = self.energy_range[0]
if emax is None:
emax = self.energy_range[1]
mask = self._slice_energy_mask(emin, emax)
emin = self.emin[mask][0]
emax = self.emax[mask][-1]
rates = np.nansum(self.rates[:, mask], axis=1).reshape(-1,1)
rate_uncert = np.sqrt(
np.nansum(self.rate_uncertainty[:, mask] ** 2, axis=1)).reshape(-1,1)
obj = BackgroundRates(rates, rate_uncert, self.tstart.copy(),
self.tstop.copy(), np.array([emin]),
np.array([emax]), exposure=self.exposure.copy())
return obj
def integrate_time(self, tstart=None, tstop=None):
"""Integrate the background over the time axis (producing a count rate
spectrum). Limits on the integration smaller than the full range can
be set.
Args:
tstart (float, optional): The low end of the integration range.
If not set, uses the lowest time edge of the histogram
tstop (float, optional): The high end of the integration range.
If not set, uses the highest time edge of the histogram
Returns:
:class:`BackgroundSpectrum`: A BackgroundSpectrum object containing \
the count rate spectrum
"""
if tstart is None:
tstart = self.time_range[0]
if tstop is None:
tstop = self.time_range[1]
mask = self._slice_time_mask(tstart, tstop)
exposure = np.nansum(self.exposure[mask])
rates = np.nansum(self.counts[mask, :], axis=0) / exposure
rate_uncert = np.sqrt(np.nansum(self.count_uncertainty[mask, :] ** 2,
axis=0)) / exposure
exposure = np.full(rates.size, exposure)
obj = BackgroundSpectrum(rates, rate_uncert, self.emin.copy(),
self.emax.copy(), exposure)
return obj
def to_bak(self, time_range=None, **kwargs):
"""Integrate over the time axis and produce a BAK object
Args:
time_range ((float, float), optional):
The time range to integrate over
**kwargs: Options to pass to BAK.from_data()
Returns:
:class:`gbm.data.BAK`: The background BAK object
"""
from gbm.data.pha import BAK
if time_range is None:
time_range = self.time_range
back_spec = self.integrate_time(*time_range)
dtype = self._datatype.lower()
bak = BAK.from_data(back_spec, *time_range, datatype=dtype,
detector=self._detector, trigtime=self._trigtime,
**kwargs)
return bak
def rebin_time(self, method, *args, tstart=None, tstop=None):
"""Not Implemented
"""
raise NotImplementedError
def rebin_energy(self, method, *args, emin=None, emax=None):
"""Not Implemented
"""
raise NotImplementedError
@classmethod
def merge_time(cls, histos):
"""Merge multiple BackroundRates together along the time axis.
Args:
histos (list of :class:`BackgroundRates`):
A list containing the BackgroundRates to be merged
Returns:
:class:`BackgroundRates`: A new BackgroundRates object containing \
the merged BackgroundRates
"""
rates = np.vstack((histo.rates for histo in histos))
rate_uncertainty = np.vstack(
(histo.rate_uncertainty for histo in histos))
bins = TimeEnergyBins.merge_time(histos)
obj = cls(rates, rate_uncertainty, bins.tstart, bins.tstop,
bins.emin, bins.emax, exposure=bins.exposure)
return obj
@classmethod
def sum_time(cls, bkgds):
"""Sum multiple TimeBins together if they have the same time range.
Example use would be summing two backgrounds from two detectors.
Args:
bkgds (list of :class:`BackgroundRates):
A list containing the BackgroundRates to be summed
Returns:
:class:`BackgroundRates`: A new BackgroundRates object containing \
the summed BackgroundRates
"""
rates = np.zeros_like(bkgds[0].rates)
rates_var = np.zeros_like(bkgds[0].rates)
for bkgd in bkgds:
assert bkgd.numtimes == bkgds[0].numtimes, \
"The backgrounds must all have the same support"
rates += bkgd.rates
rates_var += bkgd.rate_uncertainty ** 2
# averaged exposure, sampling times
exposure = np.mean([bkgd.exposure for bkgd in bkgds], axis=0)
tstart = np.mean([bkgd.tstart for bkgd in bkgds], axis=0)
tstop = np.mean([bkgd.tstop for bkgd in bkgds], axis=0)
emin = np.array([np.min([bkgd.emin for bkgd in bkgds])])
emax = np.array([np.min([bkgd.emax for bkgd in bkgds])])
sum_bkgd = cls(rates, np.sqrt(rates_var), tstart, tstop, emin, emax,
exposure=exposure, datatype=bkgds[0]._datatype,
trigtime=bkgds[0]._trigtime)
return sum_bkgd
class BackgroundSpectrum(EnergyBins):
"""A class defining a Background Spectrum.
Parameters:
rates (np.array): The array of background rates in each bin
rate_uncertainty (np.array): The array of background rate uncertainties
in each bin
lo_edges (np.array): The low-value edges of the bins
hi_edges (np.array): The high-value edges of the bins
exposure (np.array): The exposure of each bin
Attributes:
centroids (np.array): The geometric centroids of the bins
counts (np.array): The counts in each bin
count_uncertainty (np.array): The count uncertainty in each bin
exposure (np.array): The exposure of each bin
hi_edges (np.array): The high-value edges of the bins
lo_edges (np.array): The low-value edges of the bins
range (float, float): The range of the bin edges
rates (np.array): count rate of each bin
rate_uncertainty (np.array): The count rate uncertainty of each bin
size (int): Number of bins
widths (np.array): The widths of the bins
"""
def __init__(self, rates, rate_uncertainty, lo_edges, hi_edges, exposure):
counts = rates * exposure
super().__init__(counts, lo_edges, hi_edges, exposure)
self._count_uncertainty = rate_uncertainty * exposure
self._rates = rates
self._rate_uncertainty = rate_uncertainty
@property
def count_uncertainty(self):
return self._count_uncertainty
@property
def rates(self):
return self._rates
@property
def rate_uncertainty(self):
return self._rate_uncertainty
@classmethod
def sum(cls, histos):
"""Sum multiple BackgroundSpectrums together if they have the same
energy range (support).
Args:
histos (list of :class:`BackgroundSpectrum`):
A list containing the background spectra to be summed
Returns:
:class:`BackgroundSpectrum`: A new object containing the \
summed BackgroundSpectrum
"""
counts = np.zeros(histos[0].size)
count_variance = np.zeros(histos[0].size)
exposure = 0.0
for histo in histos:
assert histo.size == histos[0].size, \
"The histograms must all have the same size"
assert np.all(histo.lo_edges == histos[0].lo_edges), \
"The histograms must all have the same support"
counts += histo.counts
count_variance += histo.count_uncertainty**2
exposure += histo.exposure
rates = counts/exposure
rate_uncertainty = np.sqrt(count_variance)/exposure
sum_bins = cls(rates, rate_uncertainty, histos[0].lo_edges,
histos[0].hi_edges, exposure)
return sum_bins
def slice(self, emin, emax):
"""Perform a slice over an energy range and return a new
BackgroundSpectrum object. Note that the emin and emax values that fall
inside a bin will result in that bin being included.
Args:
emin (float): The low energy edge of the slice
emax (float): The high energy of the slice
Returns:
:class:`BackgroundSpectrum`: A new object containing the energy slice
"""
emin_snap = self.closest_edge(emin, which='low')
emax_snap = self.closest_edge(emax, which='high')
mask = (self.lo_edges < emax_snap) & (self.hi_edges > emin_snap)
obj = self.__class__(self.rates[mask], self.rate_uncertainty[mask],
self.lo_edges[mask], self.hi_edges[mask],
self.exposure[mask])
return obj