# primitives.py: Primitive data classes for pre-binned and unbinned data
#
# 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 .
#
import numpy as np
class EventList:
"""A primitive class defining a TTE event list.
Attributes:
channel_range (int, int): The range of the channels in the list
count_spectrum (:class:`EnergyBins`): The count spectrum histogram
emax (np.array): The PHA upper energy edges
emin (np.array): The PHA lower energy edges
energy_range (float, float): The range of the energy edges in the list
numchans (int): The number of energy channels. Note that not all
channels will necessarily have events, especially if a
slice is made over energy.
pha (np.array): The PHA channel array
size (int): The number of events in the list
time (np.array): The time array
time_range (float, float): The range of the times in the list
"""
def __init__(self):
self._events = np.array([], dtype=[('TIME', '>f8'), ('PHA', '>i2')])
self._ebounds = np.array([], dtype=[('CHANNEL', '>i2'),
('E_MIN', '>f4'),
('E_MAX', '>f4')])
def _assert_range(self, valrange):
assert valrange[0] <= valrange[1], \
'Range must be in increasing order: (lo, hi)'
return valrange
def sort(self, attrib):
"""In-place sort by attribute. Either by 'TIME' or 'PHA'
Args:
attrib (str): The name of the EventList attribute,
either 'TIME' or 'PHA'
"""
if attrib not in self._keys:
raise ValueError('{} is not a valid attribute.'.format(attrib) + \
' {} are valid attributes'.format(self._keys))
idx = np.argsort(self._events[attrib])
self._events = self._events[:][idx]
def time_slice(self, tstart, tstop):
"""Perform a slice in time of the EventList and return a new EventList
Args:
tstart (float): The start of the time slice
tstop (float): The end of the time slice
Returns:
:class:`EventList`: A new EventList object containing the time slice
"""
mask = (self.time >= tstart) & (self.time <= tstop)
return EventList.from_fits_array(self._events[mask], self._ebounds)
def channel_slice(self, chanlo, chanhi):
"""Perform a slice in energy channels of the EventList and return a
new EventList
Args:
chanlo (int): The start of the channel slice
chanhi (int): The end of the channel slice
Returns:
:class:`EventList`: A new EventList object containing the channel slice
"""
mask = (self.pha >= chanlo) & (self.pha <= chanhi)
return EventList.from_fits_array(self._events[mask],
self._ebounds)
def energy_slice(self, emin, emax):
"""Perform a slice in energy of the EventList and return a new EventList.
Since energies are binned, an emin or emax falling inside of an energy
channel bin will include that bin in the slice.
Args:
emin (float): The start of the energy slice
emax (float): The end of the energy slice
Returns:
:class:`EventList`: A new EventList object containing the energy slice
"""
mask = (self.emin < emax) & (self.emax > emin)
ebounds = self._ebounds[mask]
return self.channel_slice(ebounds['CHANNEL'][0],
ebounds['CHANNEL'][-1])
def bin(self, method, *args, tstart=None, tstop=None,
event_deadtime=2.6e-6, overflow_deadtime=1.0e-5, **kwargs):
"""Bin the EventList in time given a binning function and return a
2D time-energy channel histogram.
The binning function should take as input an array of times as well
as a tstart and tstop keywords for partial list binning. Additional
arguments and keyword arguments specific to the function are allowed.
The function should return an array of time edges for the bins, such
that, for `n` bins, there are `n` + 1 edges.
Args:
method (): A binning function
*args: Arguments to be passed to the binning function
tstart (float, optional):
If set, defines the start time of the EventList to be binned,
otherwise binning will begin at the time of the first event.
tstop (float, optional):
If set, defines the end time of the EventList to be binned,
otherwise binning will end at the time of the last event.
event_deadtime (float, optional): The deadtime per event in seconds.
Default is 2.6e-6.
overflow_deadtime (float, optional):
The deadtime per event in the overflow channel in seconds.
Default is 1e-5.
**kwargs: Options to be passed to the binning function
Returns:
:class:`TimeEnergyBins`: A Time-Energy Channel histogram
"""
if tstart is None:
tstart = self.time_range[0]
if tstop is None:
tstop = self.time_range[1]
# set the start and stop of the rebinning segment
mask = (self.time >= tstart) & (self.time <= tstop)
bin_times = self.time[mask]
kwargs['tstart'] = tstart
kwargs['tstop'] = tstop
# get the time edges from the binning function and then do the 2d histo
time_edges = method(bin_times, *args, **kwargs)
counts = np.histogram2d(bin_times, self.pha[mask],
[time_edges, np.arange(self.numchans + 1)])[0]
# calculate exposure
# for gbm, 2.6 microsec per count in channels < 127;
# 10 microsec per count for overflow
lo_edges = time_edges[:-1]
hi_edges = time_edges[1:]
overflow_counts = counts[:, -1]
deadtime = np.sum(counts, axis=1) * event_deadtime + \
overflow_counts * (overflow_deadtime - event_deadtime)
exposure = (hi_edges - lo_edges) - deadtime
bins = TimeEnergyBins(counts, lo_edges, hi_edges, exposure,
self.emin, self.emax)
return bins
def rebin_energy(self, method, *args):
"""Rebin the PHA channels using the specified binning algorithm. This
does not change the number of events in the EventList, but changes their
assignment to a PHA channel and bins the energy bounds mapping to those
channels. A new EventList object is returned.
Args:
method (): A binning function
*args: Arguments to be passed to the binning function
Returns:
:class:`EventList: A new EventList object with the rebinned PHA channels
"""
# exposure and counts; not really used other than for some specific
# binning algorithms
exposure = np.full(self.numchans, self.get_exposure())
chans, counts = np.unique(self.pha, return_counts=True)
if chans.size != self.numchans:
counts_fill = np.zeros(self.numchans)
counts_fill[chans] = counts
counts = counts_fill
edges = np.arange(self.numchans + 1)
# call the binning algorithm and get the new edges
_, _, new_edges = method(counts, exposure, edges, *args)
# re-assign the pha channels based on the new edges
# and also rebin the ebounds
new_pha = np.zeros_like(self.pha)
new_ebounds = []
for i in range(len(new_edges) - 1):
emin = new_edges[i]
emax = new_edges[i + 1]
mask = (self.pha >= emin) & (self.pha < emax)
new_pha[mask] = i
new_ebounds.append((i, self._ebounds[emin]['E_MIN'],
self._ebounds[emax - 1]['E_MAX']))
# create the new EventList object with the rebinned ebounds
new_events = self._events
new_events['PHA'] = new_pha
new_ebounds = np.array(new_ebounds, dtype=[('CHANNEL', '>i2'),
('E_MIN', '>f4'),
('E_MAX', '>f4')])
obj = EventList.from_fits_array(new_events, new_ebounds)
return obj
def get_exposure(self, time_ranges=None, event_deadtime=2.6e-6,
overflow_deadtime=1.0e-5):
"""Calculate the total exposure of a time range or time ranges of data
Args:
time_ranges ([(float, float), ...], optional):
The time range or time ranges over which to calculate the
exposure. If omitted, calculates the total exposure of the data.
event_deadtime (float, optional): The deadtime per event in seconds.
Default is 2.6e-6.
overflow_deadtime (float, optional):
The deadtime per event in the overflow channel in seconds.
Default is 1e-5.
Returns:
float: The exposure of the time selections
"""
if time_ranges is None:
time_ranges = [self.time_range]
try:
iter(time_ranges[0])
except:
time_ranges = [time_ranges]
exposure = 0.0
for i in range(len(time_ranges)):
tstart, tstop = self._assert_range(time_ranges[i])
tcent = (tstop + tstart) / 2.0
dt = (tstop - tstart) / 2.0
mask = (np.abs(self.time - tcent) <= dt)
# mask = (self.time >= tstart) & (self.time < tstop)
deadtime = np.sum(mask) * event_deadtime
deadtime += np.sum(self.pha[mask] == self.channel_range[1]) * \
(overflow_deadtime - event_deadtime)
exposure += (tstop - tstart) - deadtime
return exposure
@property
def time(self):
return self._events['TIME']
@property
def pha(self):
return self._events['PHA']
@property
def emin(self):
return self._ebounds['E_MIN']
@property
def emax(self):
return self._ebounds['E_MAX']
@property
def size(self):
return self._events.shape[0]
@property
def numchans(self):
return self._ebounds.shape[0]
@property
def time_range(self):
if self.size > 0:
return np.min(self.time), np.max(self.time)
@property
def channel_range(self):
if self.size > 0:
return np.min(self.pha), np.max(self.pha)
@property
def energy_range(self):
if self.size > 0:
emin = self._ebounds[self.pha.min()]['E_MIN']
emax = self._ebounds[self.pha.max()]['E_MAX']
return (emin, emax)
@property
def count_spectrum(self):
counts = np.histogram(self.pha, bins=np.arange(self.numchans + 1))[0]
exposure = np.full(self.numchans, self.get_exposure())
bins = EnergyBins(counts, self.emin, self.emax, exposure)
return bins
@classmethod
def from_fits_array(cls, events_arr, ebounds):
"""Create an EventList object from TTE FITS arrays
Args:
events_arr (np.recarray or astropy.io.fits.fitsrec.FITS_rec):
The TTE events array
ebounds (np.recarray or astropy.io.fits.fitsrec.FITS_rec):
The TTE Ebounds array, mapping channel numbers to energy bins
Returns:
:class:`EventList`: The new EventList
"""
cls = EventList()
try:
cls._keys = events_arr.names
except:
cls._keys = events_arr.dtype.names
cls._events = events_arr
cls._ebounds = ebounds
return cls
@classmethod
def from_lists(cls, times_list, pha_list, chan_lo, chan_hi):
"""Create an EventList object from lists of times, channels,
and the channel bounds. The list of times and channels must be the
same length, and the list of channel boundaries are used to map the
PHA index number in `pha_list` to energy channels.
Args:
times_list (list of float): A list of event times
pha_list (list of int): A list of PHA channels. Must be same length
as `times_list`.
chan_lo (list of float): A list of the lower edges for the energy
channels.
chan_hi (list of float): A list of the upper edges for the energy
channels. Must be same length as `chan_hi`.
Returns:
:class:`EventList`: The new EventList
"""
num_events = len(times_list)
num_chans = len(chan_lo)
if num_events != len(pha_list):
raise ValueError(
'The length of times_list and pha_list must be the same')
if num_chans != len(chan_hi):
raise ValueError(
'The length of chan_lo and chan_hi must be the same')
# events array
events_arr = np.array(list(zip(*(times_list, pha_list))),
dtype=[('TIME', '>f8'), ('PHA', '>i2')])
# ebounds array
chan_idx = np.arange(num_chans).tolist()
ebounds = np.array(list(zip(*(chan_idx, chan_lo, chan_hi))),
dtype=[('CHANNEL', '>i2'), ('E_MIN', '>f4'),
('E_MAX', '>f4')])
obj = cls.from_fits_array(events_arr, ebounds)
return obj
@classmethod
def merge(cls, eventlists, sort_attrib=None, force_unique=False):
"""Merge multiple EventLists together in time and optionally sort.
Args:
eventlist (list of :class:`EventList`):
A list containing the EventLists to be merged
sort_attrib (str, optional):
The name of the EventList attribute to sort, either 'TIME' or 'PHA'
force_unique (bool, optional):
If True, force all events to be unique via brute force sorting.
If False, the EventLists will only be checked and masked for
overlapping time ranges. Events can potentially be lost if the
merged EventLists contain overlapping times (but not necessarily
duplicate events), however this method is much faster.
Default is False.
Returns:
:class:`EventList`: A new EventList object containing the merged \
EventLists
"""
# put in time order
idx = np.argsort([np.min(eventlist.time) for eventlist in eventlists])
eventlists = [eventlists[i] for i in idx]
new_events = np.array([], dtype=[('TIME', '>f8'), ('PHA', '>i2')])
for eventlist in eventlists:
# skip empty EventLists
if eventlist.size == 0:
continue
# fix for time-shifted data
ref_time = eventlist.time[0]
temp_events = np.array(eventlist._events)
if temp_events['TIME'][0] != ref_time:
offset = ref_time - temp_events['TIME'][0]
temp_events['TIME'] += offset
# if not forcing to be unique, just make sure there is no time overlap
if (not force_unique) and (new_events.size > 0):
mask = (temp_events['TIME'] > np.max(new_events['TIME']))
temp_events = temp_events[mask]
new_events = np.concatenate((new_events, temp_events))
# force unique: make sure that we only keep unique events (slower)
if force_unique:
new_events = np.unique(new_events)
# mark: TODO add check for ebounds consistency
cls = EventList.from_fits_array(new_events, eventlists[0]._ebounds)
# do a sort
if sort_attrib is not None:
cls.sort(sort_attrib)
return cls
class Bins:
"""A primitive class defining a set of histogram bins
Parameters:
counts (np.array): The array of counts 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
Attributes:
centroids (np.array): The centroids of the bins
counts (np.array): The counts in each bin
count_uncertainty (np.array): The count uncertainty in 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): The count rate of each bin: counts/width
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, counts, lo_edges, hi_edges):
self.counts = counts
self.lo_edges = lo_edges
self.hi_edges = hi_edges
self._good_segments = self._calculate_good_segments()
@property
def size(self):
return len(self.lo_edges)
@property
def counts(self):
return self._counts
@counts.setter
def counts(self, val):
try:
iter(val)
self._counts = np.array(val)
except:
raise TypeError('Bins.counts must be an iterable')
@property
def lo_edges(self):
return self._lo_edges
@lo_edges.setter
def lo_edges(self, val):
try:
iter(val)
self._lo_edges = np.array(val)
except:
raise TypeError('Bins.lo_edges must be an iterable')
@property
def hi_edges(self):
return self._hi_edges
@hi_edges.setter
def hi_edges(self, val):
try:
iter(val)
self._hi_edges = np.array(val)
except:
raise TypeError('Bins.hi_edges must be an iterable')
@property
def widths(self):
return self.hi_edges - self.lo_edges
@property
def centroids(self):
return (self.hi_edges + self.lo_edges) / 2.0
@property
def range(self):
return (self.lo_edges[0], self.hi_edges[-1])
@property
def count_uncertainty(self):
return np.sqrt(self._counts)
@property
def rates(self):
return self.counts / self.widths
@property
def rate_uncertainty(self):
return self.count_uncertainty / self.width
def closest_edge(self, val, which='either'):
"""Return the closest bin edge
Args:
val (float): Input value
which (str, optional): Options are:
* 'either' - closest edge to val;
* 'low' - closest edge lower than val; or
* 'high' - closest edge higher than val. Default is 'either'
Returns:
float: The closest bin edge to the input value
"""
edges = np.concatenate((self.lo_edges, [self.hi_edges[-1]]))
idx = np.argmin(np.abs(val - edges))
if which == 'low' and (idx - 1) >= 0:
if edges[idx] > val:
idx -= 1
elif (which == 'high') and (idx + 1) < edges.size:
if edges[idx] < val:
idx += 1
else:
pass
return edges[idx]
def slice(self, lo_edge, hi_edge):
"""Perform a slice over the range of the bins and return a new Bins
object. Note that lo_edge and hi_edge values that fall inside a bin will
result in that bin being included.
Args:
lo_edge (float): The start of the slice
hi_edge (float): The end of the slice
Returns:
:class:`Bins`: A new Bins object containing the slice
"""
lo_snap = self.closest_edge(lo_edge, which='low')
hi_snap = self.closest_edge(hi_edge, which='high')
if lo_snap == hi_snap:
mask = (self.lo_edges < hi_snap) & (self.hi_edges >= lo_snap)
else:
mask = (self.lo_edges < hi_snap) & (self.hi_edges > lo_snap)
obj = Bins(self.counts[mask], self.lo_edges[mask], self.hi_edges[mask])
return obj
def contiguous_bins(self):
"""Return a list of Bins, each one containing a continuous segment of
data. This is done by comparing the edges of each bin, and if there
is a gap between edges, the data is split into separate Bin objects,
each containing a contiguous set of data.
Returns
list of :class:`Bins`: A list of Bins
"""
bins = [self.slice(seg[0], seg[1]) for seg in self._good_segments]
return bins
def _calculate_good_segments(self):
"""Calculates the ranges of data that are contiguous segments
Returns:
[(float, float), ...]: A list of tuples, each containing the edges \
of a contiguous segment.
"""
mask = (self.lo_edges[1:] != self.hi_edges[:-1])
if mask.sum() == 0:
return [(self.lo_edges[0], self.hi_edges[-1])]
times = np.concatenate(([self.lo_edges[0]], self.hi_edges[:-1][mask],
self.lo_edges[1:][mask], [self.hi_edges[-1]]))
times.sort()
return times.reshape(-1, 2).tolist()
class TimeBins(Bins):
"""A class defining a set of Time History (lightcurve) bins.
Parameters:
counts (np.array): The array of counts 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 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): The count rate of each bin: counts/exposure
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, counts, lo_edges, hi_edges, exposure):
super().__init__(counts, lo_edges, hi_edges)
self.exposure = exposure
@property
def exposure(self):
return self._exposure
@exposure.setter
def exposure(self, val):
try:
iter(val)
self._exposure = np.array(val)
except:
raise TypeError('TimeBins.exposure must be an iterable')
@property
def rates(self):
r = np.zeros_like(self.exposure)
mask = (self.exposure > 0.0)
r[mask] = self.counts[mask] / self.exposure[mask]
return r
@property
def rate_uncertainty(self):
r = np.zeros_like(self.exposure)
mask = (self.exposure > 0.0)
r[mask] = self.count_uncertainty[mask] / self.exposure[mask]
return r
@classmethod
def from_fits_array(cls, arr):
"""Create an TimeBins object from a FITS counts array
Args:
arr (np.recarray or astropy.io.fits.fitsrec.FITS_rec):
The FITS counts array
Returns:
:class:`TimeBins`: The new TimeBins object
"""
keys = arr.dtype.names
if 'COUNTS' not in keys or 'TIME' not in keys or \
'ENDTIME' not in keys or 'EXPOSURE' not in keys:
raise ValueError('Array has missing or invalid columns')
counts = arr['COUNTS']
if len(counts.shape) == 2:
counts = np.sum(arr['COUNTS'], axis=1)
obj = TimeBins(counts, arr['TIME'], arr['ENDTIME'], arr['EXPOSURE'])
return obj
def slice(self, tstart, tstop):
"""Perform a slice over a time range and return a new Bins object. Note
that tstart and tstop values that fall inside a bin will result in
that bin being included.
Args:
lo_edge (float): The start of the slice
hi_edge (float): The end of the slice
Returns:
:class:`TimeBins`: A new TimeBins object containing the time slice
"""
tstart_snap = self.closest_edge(tstart, which='low')
tstop_snap = self.closest_edge(tstop, which='high')
mask = (self.lo_edges < tstop_snap) & (self.hi_edges > tstart_snap)
obj = self.__class__(self.counts[mask], self.lo_edges[mask],
self.hi_edges[mask], self.exposure[mask])
return obj
@classmethod
def merge(cls, histos):
"""Merge multiple TimeBins together.
Args:
histos (list of :class:`TimeBins`):
A list containing the TimeBins to be merged
Returns:
:class:`TimeBins`: A new TimeBins object containing the merged TimeBins
"""
num = len(histos)
# sort by start time
tstarts = np.concatenate([histo.lo_edges for histo in histos])
idx = np.argsort(tstarts)
# concatenate the histos in order
counts = histos[idx[0]].counts
lo_edges = histos[idx[0]].lo_edges
hi_edges = histos[idx[0]].hi_edges
exposure = histos[idx[0]].exposure
for i in range(1, num):
bin_starts = histos[idx[i]].lo_edges
# make sure there is no overlap
mask = (bin_starts >= hi_edges[-1])
counts = np.concatenate((counts, histos[idx[i]].counts[mask]))
lo_edges = np.concatenate(
(lo_edges, histos[idx[i]].lo_edges[mask]))
hi_edges = np.concatenate(
(hi_edges, histos[idx[i]].hi_edges[mask]))
exposure = np.concatenate(
(exposure, histos[idx[i]].exposure[mask]))
# new TimeBins object
merged_bins = cls(counts, lo_edges, hi_edges, exposure)
return merged_bins
@classmethod
def sum(cls, histos):
"""Sum multiple TimeBins together if they have the same time range
(support) and the same bin widths. Example use would be summing
two histograms at different energies.
Args:
histos (list of :class:`TimeBins`):
A list containing the TimeBins to be summed
Returns:
:class:`TimeBins`: A new TimeBins object containing the summed TimeBins
"""
counts = np.zeros(histos[0].size)
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"
assert np.all(histo.widths == histos[0].widths), \
"The histograms must all have the same exposure"
counts += histo.counts
# averaged exposure
exposure = np.mean([histo.exposure for histo in histos], axis=0)
sum_bins = cls(counts, histos[0].lo_edges, histos[0].hi_edges,
exposure)
return sum_bins
def rebin(self, method, *args, tstart=None, tstop=None):
"""Rebin the TimeBins object in given a binning function and return a
a new TimeBins object
The binning function should take as input an array of counts,
array of exposures, and an array of bin edges. Additional arguments
specific to the function are allowed. The function should return an
array of the new counts, new exposure, and new edges.
Args:
method (): A binning function
*args: Arguments to be passed to the binning function
tstart (float, optional):
If set, defines the start time of the TimeBins to be binned,
otherwise binning will begin at the time of the first bin edge.
tstop (float, optional):
If set, defines the end time of the TimeBins to be binned,
otherwise binning will end at the time of the last bin edge.
Returns:
:class:`TimeBins`: The rebinned TimeBins object
"""
# set the start and stop of the rebinning segment
trange = self.range
if tstart is None:
tstart = trange[0]
if tstop is None:
tstop = trange[1]
if tstart < trange[0]:
tstart = trange[0]
if tstop > trange[1]:
tstop = trange[1]
bins = self.contiguous_bins()
new_histos = []
for bin in bins:
trange = bin.range
# split the histogram into pieces so that we only rebin the piece
# that needs to be rebinned
pre = None
post = None
if (tstop < trange[0]) or (tstart > trange[1]):
histo = None
elif tstop == trange[1]:
if tstart > trange[0]:
pre = bin.slice(trange[0], tstart)
histo = bin.slice(self.closest_edge(tstart, which='high'),
trange[1])
elif tstart == trange[0]:
histo = bin.slice(trange[0],
self.closest_edge(tstop, which='low'))
if tstop < trange[1]:
post = bin.slice(tstop, trange[1])
elif (tstart > trange[0]) and (tstop < trange[1]):
pre = bin.slice(trange[0], tstart)
histo = bin.slice(self.closest_edge(tstart, which='high'),
self.closest_edge(tstop, which='low'))
post = bin.slice(tstop, trange[1])
else:
histo = bin
# perform the rebinning and create a new histo with the rebinned rates
if histo is not None:
edges = np.append(histo.lo_edges, histo.hi_edges[-1])
new_counts, new_exposure, new_edges = method(histo.counts,
histo.exposure,
edges, *args)
new_histo = self.__class__(new_counts, new_edges[:-1],
new_edges[1:],
new_exposure)
else:
new_histo = None
# now merge the split histo back together again
histos_to_merge = [i for i in (pre, new_histo, post) if
i is not None]
new_histos.append(self.__class__.merge(histos_to_merge))
new_histo = self.__class__.merge(new_histos)
return new_histo
class EnergyBins(TimeBins):
"""A class defining a set of Energy (count spectra) bins.
Parameters:
counts (np.array): The array of counts 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): The differential count rate of each bin:
counts/(exposure*widths)
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, counts, lo_edges, hi_edges, exposure):
super().__init__(counts, lo_edges, hi_edges, exposure)
@property
def centroids(self):
return np.sqrt(self.hi_edges * self.lo_edges)
@property
def rates(self):
return self.counts / (self.exposure * self.widths)
@property
def rate_uncertainty(self):
return self.count_uncertainty / (self.exposure * self.widths)
@classmethod
def from_fits_array(cls, arr, ebounds):
"""Create an EnergyBins object from a FITS counts array
Args:
arr (np.recarray or astropy.io.fits.fitsrec.FITS_rec):
The FITS counts array
ebounds (np.recarray or astropy.io.fits.fitsrec.FITS_rec):
The ebounds array, mapping channel numbers to energy bins
Returns:
:class:`EnergyBins`: The new EnergyBins object
"""
keys = arr.dtype.names
if 'COUNTS' not in keys or 'EXPOSURE' not in keys:
raise ValueError('Array has missing or invalid columns')
keys = ebounds.dtype.names
if 'E_MIN' not in keys or 'E_MAX' not in keys:
raise ValueError('ebounds has missing or invalid columns')
counts = arr['COUNTS']
if len(counts.shape) == 2:
counts = np.sum(arr['COUNTS'], axis=0)
exposure = np.full(counts.size, np.sum(arr['EXPOSURE']))
obj = EnergyBins(counts, ebounds['E_MIN'], ebounds['E_MAX'], exposure)
return obj
@classmethod
def sum(cls, histos):
"""Sum multiple EnergyBins together if they have the same energy range
(support). Example use would be summing two count spectra.
Args:
histos (list of :class:`EnergyBins`):
A list containing the EnergyBins to be summed
Returns:
:class:`EnergyBins`: A new EnergyBins object containing the \
summed EnergyBins
"""
counts = 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
exposure += histo.exposure
sum_bins = cls(counts, histos[0].lo_edges, histos[0].hi_edges,
exposure)
return sum_bins
def rebin(self, method, *args, emin=None, emax=None):
"""Rebin the EnergyBins object in given a binning function and return a
a new EnergyBins object
The binning function should take as input an array of counts,
array of exposures, and an array of bin edges. Additional arguments
specific to the function are allowed. The function should return an
array of the new counts, new exposure, and new edges.
Args:
method (): A binning function
*args: Arguments to be passed to the binning function
emin (float, optional):
If set, defines the starting energy of the EnergyBins to be
binned, otherwise binning will begin at the first bin edge.
emax (float, optional):
If set, defines the ending energy of the EnergyBins to be binned,
otherwise binning will end at the last bin edge.
Returns:
:class:`EnergyBins`: The rebinned EnergyBins object
"""
histo = super().rebin(method, *args, tstart=emin, tstop=emax)
histo._exposure = self.exposure[:histo.size]
return histo
class TimeEnergyBins():
"""A class defining a set of 2D Time-Energy bins.
Parameters:
counts (np.array): The array of counts 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
exposure (np.array): The exposure of each bin
emin (np.array): The low-value edges of the energy bins
emax (np.array): The high-value edges of the energy bins
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, counts, tstart, tstop, exposure, emin, emax):
self.counts = counts
self.tstart = tstart
self.tstop = tstop
self.exposure = exposure
self.emin = emin
self.emax = emax
if self.numtimes > 0:
self._good_time_segments = self._calculate_good_segments(
self.tstart,
self.tstop)
else:
self._good_time_segments = [(None, None)]
if self.numchans > 0:
self._good_energy_segments = self._calculate_good_segments(
self.emin,
self.emax)
else:
self._good_energy_segments = [(None, None)]
@property
def numtimes(self):
return len(self.exposure)
@property
def numchans(self):
return len(self.emin)
@property
def size(self):
return (self.numtimes, self.numchans)
@property
def time_widths(self):
return (self.tstop - self.tstart)
@property
def chan_widths(self):
return (self.emax - self.emin)
@property
def time_centroids(self):
return (self.tstop + self.tstart) / 2.0
@property
def energy_centroids(self):
return np.sqrt(self.emin * self.emax)
@property
def time_range(self):
return (self.tstart[0], self.tstop[-1])
@property
def energy_range(self):
return (self.emin[0], self.emax[-1])
@property
def count_uncertainty(self):
return np.sqrt(self.counts)
@property
def rates(self):
return self.counts / (self.exposure[:, np.newaxis])
@property
def rate_uncertainty(self):
return self.count_uncertainty / (self.exposure[:, np.newaxis])
@property
def rates_per_kev(self):
return self.rates / self.chan_widths[:, np.newaxis]
@property
def rate_uncertainty_per_kev(self):
return self.rate_uncertainty / self.chan_widths[:, np.newaxis]
def _assert_range(self, valrange):
assert valrange[0] <= valrange[1], \
'Range must be in increasing order: (lo, hi)'
return valrange
def _slice_time_mask(self, tstart, tstop):
tstart_snap = self.closest_time_edge(tstart, which='low')
tstop_snap = self.closest_time_edge(tstop, which='high')
if tstart_snap == tstop_snap:
mask = (self.tstart < tstop_snap) & (self.tstop >= tstart_snap)
else:
mask = (self.tstart < tstop_snap) & (self.tstop > tstart_snap)
return mask
def _slice_energy_mask(self, emin, emax):
emin_snap = self.closest_energy_edge(emin, which='low')
emax_snap = self.closest_energy_edge(emax, which='high')
if emin_snap == emax_snap:
mask = (self.emin < emax_snap) & (self.emax >= emin_snap)
else:
mask = (self.emin < emax_snap) & (self.emax > emin_snap)
return mask
def _calculate_good_segments(self, lo_edges, hi_edges):
"""Calculates the ranges of data that are contiguous segments
Args:
lo_edges (np.array): The lower bin edges
hi_edges (np.array): The upper bin edges
Returns:
[(float, float), ...]: A list of tuples, each containing the edges \
of a contiguous segment
"""
mask = (lo_edges[1:] != hi_edges[:-1])
if mask.sum() == 0:
return [(lo_edges[0], hi_edges[-1])]
edges = np.concatenate(([lo_edges[0]], hi_edges[:-1][mask],
lo_edges[1:][mask], [hi_edges[-1]]))
edges.sort()
return edges.reshape(-1, 2).tolist()
def closest_time_edge(self, val, which='either'):
"""Return the closest time bin edge
Args:
val (float): Input value
which (str, optional): Options are:
* 'either' - closest edge to val;
* 'low' - closest edge lower than val;
* 'high' - closest edge higher than val. Default is 'either'
Returns:
float: The closest time bin edge to the input value
"""
edges = np.concatenate((self.tstart, [self.tstop[-1]]))
idx = np.argmin(np.abs(val - edges))
if which == 'low':
if (edges[idx] > val) and (idx - 1) >= 0:
idx -= 1
elif (which == 'high') and (idx + 1) < edges.size:
if edges[idx] < val:
idx += 1
else:
pass
return edges[idx]
def closest_energy_edge(self, val, which='either'):
"""Return the closest energy bin edge
Args:
val (float): Input value
which (str, optional): Options are:
* 'either' - closest edge to val;
* 'low' - closest edge lower than val;
* 'high' - closest edge higher than val. Default is 'either'
Returns:
float: The closest energy bin edge to the input value
"""
edges = np.concatenate((self.emin, [self.emax[-1]]))
idx = np.argmin(np.abs(val - edges))
if which == 'low':
if (edges[idx] > val) and (idx - 1) >= 0:
idx -= 1
elif (which == 'high') and (idx + 1) < edges.size:
if edges[idx] < val:
idx += 1
else:
pass
return edges[idx]
def slice_time(self, tstart, tstop):
"""Perform a slice over a time range and return a new TimeEnergyBins
object. Note that tstart and tstop values that fall inside a bin will
result in that bin being included.
Args:
tstart (float): The start of the slice
tstop (float): The end of the slice
Returns:
:class:`TimeEnergyBins`: A new TimeEnergyBins object containing \
the time slice
"""
mask = self._slice_time_mask(tstart, tstop)
cls = type(self)
obj = cls(self.counts[mask, :], self.tstart[mask], self.tstop[mask],
self.exposure[mask], self.emin, self.emax)
return obj
def slice_energy(self, emin, emax):
"""Perform a slice over an energy range and return a new TimeEnergyBins
object. Note that emin and emax values that fall inside a bin will
result in that bin being included.
Args:
emin (float): The start of the slice
emax (float): The end of the slice
Returns:
:class:`TimeEnergyBins`: A new TimeEnergyBins object containing \
the energy slice
"""
mask = self._slice_energy_mask(emin, emax)
obj = TimeEnergyBins(self.counts[:, mask], self.tstart, self.tstop,
self.exposure, self.emin[mask], self.emax[mask])
return obj
def integrate_energy(self, emin=None, emax=None):
"""Integrate the histogram over the energy axis (producing a lightcurve).
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:`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)
counts = np.sum(self.counts[:, mask], axis=1)
obj = TimeBins(counts, self.tstart, self.tstop, self.exposure)
return obj
def integrate_time(self, tstart=None, tstop=None):
"""Integrate the histogram 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:`EnergyBins`: A EnergyBins object containing the count \
rate spectrum histogram
"""
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)
counts = np.sum(self.counts[mask, :], axis=0)
exposure = np.sum(self.exposure[mask])
exposure = np.full(counts.size, exposure)
obj = EnergyBins(counts, self.emin, self.emax, exposure)
return obj
def contiguous_time_bins(self):
"""Return a list of TimeEnergyBins, each one containing a contiguous
time segment of data. This is done by comparing the edges of each time
bin, and if thereis a gap between edges, the data is split into
separate TimeEnergyBin objects, each containing a time-contiguous set
of data.
Returns:
list of :class:`TimeEnergyBins`: A list of TimeEnergyBins
"""
bins = [self.slice_time(seg[0], seg[1]) for seg in
self._good_time_segments]
return bins
def contiguous_energy_bins(self):
"""Return a list of TimeEnergyBins, each one containing a contiguous
energy segment of data. This is done by comparing the edges of each
energy bin, and if thereis a gap between edges, the data is split into
separate TimeEnergyBin objects, each containing an energy-contiguous set
of data.
Returns:
list of :class:`TimeEnergyBins`: A list of TimeEnergyBins
"""
bins = [self.slice_energy(seg[0], seg[1]) for seg in
self._good_energy_segments]
return bins
def rebin_time(self, method, *args, tstart=None, tstop=None):
"""Rebin the TimeEnergyBins object along the time axis given a binning
function and return a new TimeEnergyBins object
The binning function should take as input an array of counts,
array of exposures, and an array of bin edges. Additional arguments
specific to the function are allowed. The function should return an
array of the new counts, new exposure, and new edges.
Args:
method (): A binning function
*args: Arguments to be passed to the binning function
tstart (float, optional):
If set, defines the start time of the TimeEnergyBins to be
binned, otherwise binning will begin at the time of the first
bin edge.
tstop (float, optional):
If set, defines the end time of the TimeEnergyBins to be
binned, otherwise binning will end at the time of the last
bin edge.
Returns:
:class:`TimeEnergyBins`: The rebinned TimeEnergyBins object
"""
# set the start and stop of the rebinning segment
trange = self.time_range
if tstart is None:
tstart = trange[0]
if tstop is None:
tstop = trange[1]
if tstart < trange[0]:
tstart = trange[0]
if tstop > trange[1]:
tstop = trange[1]
bins = self.contiguous_time_bins()
new_histos = []
for bin in bins:
trange = bin.time_range
# split the histogram into pieces so that we only rebin the piece
# that needs to be rebinned
pre = None
post = None
if (tstop < trange[0]) or (tstart > trange[1]):
histo = None
elif tstop == trange[1]:
if tstart > trange[0]:
pre = bin.slice_time(trange[0], tstart)
histo = bin.slice_time(self.closest_time_edge(tstart,
which='high'),
trange[1])
elif tstart == trange[0]:
histo = bin.slice_time(trange[0],
self.closest_time_edge(tstop,
which='low'))
if tstop < trange[1]:
post = bin.slice_time(tstop, trange[1])
elif (tstart > trange[0]) and (tstop < trange[1]):
pre = bin.slice_time(trange[0], tstart)
histo = bin.slice_time(
self.closest_time_edge(tstart, which='high'),
self.closest_time_edge(tstop, which='low'))
post = bin.slice_time(tstop, trange[1])
else:
histo = bin
# perform the rebinning and create a new histo with the rebinned rates
if histo is not None:
edges = np.append(histo.tstart, histo.tstop[-1])
new_counts = []
for i in range(bin.numchans):
new_cts, new_exposure, new_edges = method(
histo.counts[:, i],
histo.exposure,
edges, *args)
new_counts.append(new_cts)
new_counts = np.array(new_counts).T
new_histo = TimeEnergyBins(new_counts, new_edges[:-1],
new_edges[1:],
new_exposure, bin.emin, bin.emax)
else:
new_histo = bin
if len(new_histo.tstart) == 0:
new_histo = None
# now merge the split histo back together again
histos_to_merge = [i for i in (pre, new_histo, post) if
i is not None]
new_histos.append(TimeEnergyBins.merge_time(histos_to_merge))
new_histo = TimeEnergyBins.merge_time(new_histos)
return new_histo
def rebin_energy(self, method, *args, emin=None, emax=None):
"""Rebin the TimeEnergyBins object along the energy axis given a binning
function and return a new TimeEnergyBins object
The binning function should take as input an array of counts,
array of exposures, and an array of bin edges. Additional arguments
specific to the function are allowed. The function should return an
array of the new counts, new exposure, and new edges.
Args:
method (): A binning function
*args: Arguments to be passed to the binning function
emin (float, optional):
If set, defines the starting edge of the TimeEnergyBins to be
binned, otherwise binning will begin at the the first bin edge.
emax (float, optional):
If set, defines the ending edge of the TimeEnergyBins to be
binned, otherwise binning will end at the last bin edge.
Returns:
:class:`TimeEnergyBins`: The rebinned TimeEnergyBins object
"""
# set the start and stop of the rebinning segment
erange = self.energy_range
if emin is None:
emin = erange[0]
if emax is None:
emax = erange[1]
if emin < erange[0]:
emin = erange[0]
if emax > erange[1]:
emax = erange[1]
bins = self.contiguous_energy_bins()
new_histos = []
for bin in bins:
# split the histogram into pieces so that we only rebin the piece
# that needs to be rebinned
pre = None
post = None
if (emax < erange[0]) or (emin > erange[1]):
histo = None
elif emax == erange[1]:
if emin > erange[0]:
pre = bin.slice_energy(erange[0], emin)
histo = bin.slice_energy(self.closest_energy_edge(emin,
which='high'),
erange[1])
elif emin == erange[0]:
histo = bin.slice_energy(erange[0],
self.closest_energy_edge(emax,
which='low'))
if emax < erange[1]:
post = bin.slice_energy(emax, erange[1])
elif (emin > erange[0]) and (emax < erange[1]):
pre = bin.slice_energy(erange[0], emin)
histo = bin.slice_energy(
self.closest_energy_edge(emin, which='high'),
self.closest_energy_edge(emax, which='low'))
post = bin.slice_energy(emax, erange[1])
else:
histo = bin
# perform the rebinning and create a new histo with the rebinned rates
if histo is not None:
edges = np.append(histo.emin, histo.emax[-1])
numtimes, numchans = histo.size
new_counts = []
for i in range(numtimes):
exposure = np.full(numchans, histo.exposure[i])
new_cts, _, new_edges = method(histo.counts[i, :],
exposure,
edges, *args)
new_counts.append(new_cts)
new_counts = np.array(new_counts)
new_histo = TimeEnergyBins(new_counts, bin.tstart, bin.tstop,
bin.exposure, new_edges[:-1],
new_edges[1:])
else:
new_histo = bin
# now merge the split histo back together again
histos_to_merge = [i for i in (pre, new_histo, post) if
i is not None]
new_histos.append(TimeEnergyBins.merge_energy(histos_to_merge))
new_histo = TimeEnergyBins.merge_energy(new_histos)
return new_histo
def get_exposure(self, time_ranges=None, scale=False):
"""Calculate the total exposure of a time range or time ranges of data
Args:
time_ranges ([(float, float), ...], optional):
The time range or time ranges over which to calculate the
exposure. If omitted, calculates the total exposure of the data
scale (bool, optional):
If True and the time ranges don't match up with the data binning,
will scale the exposure based on the requested time range.
Default is False.
Returns:
float: The exposure of the time selections
"""
if time_ranges is None:
time_ranges = [self.time_range]
try:
iter(time_ranges[0])
except:
time_ranges = [time_ranges]
exposure = 0.0
for i in range(len(time_ranges)):
mask = self._slice_time_mask(*self._assert_range(time_ranges[i]))
dt = (time_ranges[i][1] - time_ranges[i][0])
data_exp = np.sum(self.exposure[mask])
dts = np.sum(self.tstop[mask] - self.tstart[mask])
if dts > 0:
if scale:
exposure += data_exp * (dt / ds)
else:
exposure += data_exp
return exposure
@classmethod
def merge_time(cls, histos):
"""Merge multiple TimeEnergyBins together along the time axis.
Args:
histos (list of :class:`TimeEnergyBins`):
A list containing the TimeEnergyBins to be merged
Returns
:class:`TimeEnergyBins`: A new TimEnergyBins object containing the \
merged TimeEnergyBins
"""
num = len(histos)
# sort by start time
tstarts = np.array([histo.tstart[0] for histo in histos])
idx = np.argsort(tstarts)
# concatenate the histos in order
counts = histos[idx[0]].counts
tstart = histos[idx[0]].tstart
tstop = histos[idx[0]].tstop
exposure = histos[idx[0]].exposure
emin = histos[idx[0]].emin
emax = histos[idx[0]].emax
for i in range(1, num):
bin_starts = histos[idx[i]].tstart
# make sure there is no overlap
mask = (bin_starts >= tstop[-1])
counts = np.vstack((counts, histos[idx[i]].counts[mask, :]))
tstart = np.concatenate((tstart, histos[idx[i]].tstart[mask]))
tstop = np.concatenate((tstop, histos[idx[i]].tstop[mask]))
exposure = np.concatenate(
(exposure, histos[idx[i]].exposure[mask]))
# new TimeEnergyBins object
merged_bins = cls(counts, tstart, tstop, exposure, emin, emax)
return merged_bins
@classmethod
def merge_energy(cls, histos):
"""Merge multiple TimeEnergyBins together along the energy axis.
Args:
histos (list of :class:`TimeEnergyBins`):
A list containing the TimeEnergyBins to be merged
Returns:
:class:`TimeEnergyBins`: A new TimEnergyBins object containing \
the merged TimeEnergyBins
"""
num = len(histos)
# sort by channel edge
emins = np.concatenate([histo.emin for histo in histos])
idx = np.argsort(emins)
# concatenate the histos in order
counts = histos[idx[0]].counts
tstart = histos[idx[0]].tstart
tstop = histos[idx[0]].tstop
exposure = histos[idx[0]].exposure
emin = histos[idx[0]].emin
emax = histos[idx[0]].emax
for i in range(1, num):
bin_starts = histos[idx[i]].emin
# make sure there is no overlap
mask = (bin_starts >= emax[-1])
counts = np.hstack((counts, histos[idx[i]].counts[:, mask]))
emin = np.concatenate((emin, histos[idx[i]].emin[mask]))
emax = np.concatenate((emax, histos[idx[i]].emax[mask]))
# new TimeEnergyBins object
merged_bins = cls(counts, tstart, tstop, exposure, emin, emax)
return merged_bins
class TimeRange():
"""A primitive class defining a time range
Parameters:
tstart (float): The start time of the range
tstop (float): The end time of the range
Attributes:
center (float): The center of the time range
duration (float): The duration of the time range
tstart (float): The start time of the range
tstop (float): The end time of the range
"""
def __init__(self, tstart, tstop):
tstart = self._assert_time(tstart)
tstop = self._assert_time(tstop)
if tstop >= tstart:
self._tstart = tstart
self._tstop = tstop
else:
self._tstart = tstop
self._tstop = tstart
def __str__(self):
return '({0}, {1})'.format(self.tstart, self.tstop)
def _assert_time(self, atime):
try:
atime = float(atime)
except:
raise TypeError('time must be a float')
return atime
@property
def tstart(self):
return self._tstart
@property
def tstop(self):
return self._tstop
@property
def duration(self):
return self.tstop - self.tstart
@property
def center(self):
return (self.tstart + self.tstop) / 2.0
def as_tuple(self):
"""Return the time range as a tuple.
Returns:
(float, float): The starting and ending time of the time range
"""
return (self.tstart, self.tstop)
def contains(self, a_time, inclusive=True):
"""Determine if the time range contains a time.
Args:
a_time (float): The input time to check
inclusive (bool, optional):
If True, then includes the edges of the range for the check,
otherwise it is edge-exclusive. Default is True.
Returns:
bool: True if the time is in the time range, False otherwise
"""
a_time = self._assert_time(a_time)
if inclusive:
test = (a_time <= self.tstop) and (a_time >= self.tstart)
else:
test = (a_time < self.tstop) and (a_time > self.tstart)
if test:
return True
else:
return False
@classmethod
def union(cls, range1, range2):
"""Return a new TimeRange that is the union of two input TimeRanges
Args:
range1 (:class:`TimeRange`): A time range used to calculate the union
range2 (:class:`TimeRange`): Another time range used to calculate
the union
Returns:
:class:`TimeRange`: The unionized time range
"""
tstart = np.min((range1.tstart, range2.tstart))
tstop = np.max((range1.tstop, range2.tstop))
obj = cls(tstart, tstop)
return obj
@classmethod
def intersection(cls, range1, range2):
"""Return a new TimeRange that is the intersection of two input
TimeRanges. If the input TimeRanges do not intersect, then None is
returned.
Args:
range1 (:class:`TimeRange`): A time range used to calculate the
intersection
range2 (:class:`TimeRange`): Another time range used to calculate
the intersection
Returns:
:class:`TimeRange`: The intersected time range
"""
# test if one tstart is inside the other time range
if range1.contains(range2.tstart):
lower = range1
upper = range2
elif range2.contains(range1.tstart):
lower = range2
upper = range1
else:
return None
# do the merge
tstart = upper.tstart
tstop = np.min((lower.tstop, upper.tstop))
obj = cls(tstart, tstop)
return obj
class GTI:
"""A primitive class defining a set of Good Time Intervals (GTIs)
Parameters:
tstart (float): The start time of the GTI
tstop (float): The end time of the GTI
Attributes:
num_intervals (int): The number of intervals in the GTI
range (float, float): The full range of the GTI
"""
def __init__(self, tstart=None, tstop=None):
self._gti = None
if (tstart is not None) and (tstop is not None):
self._gti = [TimeRange(tstart, tstop)]
@property
def num_intervals(self):
return len(self._gti)
@property
def range(self):
"""The full range of the GTI
:type: (float, float)"""
return (self._gti[0].tstart, self._gti[-1].tstop)
def as_list(self):
"""Return the GTI as a list of tuples.
Returns:
[(float, float), ...]: The list of GTI interval tuples
"""
gti_list = [one_gti.as_tuple() for one_gti in self._gti]
return gti_list
def insert(self, tstart, tstop):
"""Insert a new interval into the GTI
Args:
tstart (float): The start time of the new interval
tstop (float): The end time of the new interval
"""
# where the new time range should be inserted
time_range = TimeRange(tstart, tstop)
idx = [i for i, j in enumerate(self._gti) if
j.tstart <= time_range.tstart]
# determine if there is overlap with the lower bounding range, and if so
# then merge
if len(idx) != 0:
idx = idx[-1]
if self._gti[idx].contains(time_range.tstart, inclusive=True):
the_range = self._gti.pop(idx)
time_range = TimeRange.union(the_range, time_range)
else:
idx += 1
else:
idx = 0
# determine if there is overlap with the upper bounding range, and if so
# then merge
if idx < len(self._gti):
if self._gti[idx].contains(time_range.tstop):
the_range = self._gti.pop(idx)
time_range = TimeRange.union(the_range, time_range)
self._gti.insert(idx, time_range)
def contains(self, a_time, inclusive=True):
"""Determine if the GTI contains a time.
Args:
a_time (float): The input time to check
inclusive (bool, optional):
If True, then includes the edges of the range for the check,
otherwise it is edge-exclusive. Default is True.
Returns:
bool: True if the time is in the GTI, False otherwise
"""
test = [gti.contains(a_time, inclusive=inclusive) for gti in self._gti]
return any(test)
@classmethod
def from_list(cls, gti_list):
"""Create a new GTI object from a list of tuples.
Args:
gti_list ([(float, float), ...]): A list of interval tuples
Returns:
:class:`GTI`: The new GTI object
"""
gtis = [TimeRange(*one_gti) for one_gti in gti_list]
obj = cls()
obj._gti = gtis
return obj
@classmethod
def from_boolean_mask(cls, times, mask):
"""Create a new GTI object from a list of times and a Boolean mask
Splits the boolean mask into segments of contiguous values and applies
to array of times to create a GTI object.
Args:
times (np.array): An array of times
mask (np.array(dtype=bool)): The boolean array. Must be the same
size as times.
Returns:
:class:`GTI`: The new GTI object
"""
# split a boolean mask array into segments based on True/False
indices = np.nonzero(mask[1:] != mask[:-1])[0] + 1
time_segs = np.split(times, indices)
mask_segs = np.split(mask, indices)
# retrieve the start and stop times for the "off" intervals
segs = []
numsegs = len(indices) + 1
for i in range(numsegs):
if mask_segs[i][0] == 0:
segs.append((time_segs[i][0], time_segs[i][-1]))
# mark FIXME: null if mask is all True or all False
# currently assuming that it must be all True
if len(segs) == 0:
segs = [self.time_range]
return cls.from_list(segs)
@classmethod
def merge(cls, gti1, gti2):
"""Return a new GTI object that is a merge of two existing GTI objects.
Args:
gti1 (:class:`GTI`): A GTI to be merged
gti2 (:class:`GTI`): A GTI to be merged
Returns:
:class:`GTI`: The new merged GTI object
"""
time_list = gti1.as_list()
time_list.extend(gti2.as_list())
time_list = sorted(time_list)
obj = cls(*time_list.pop(0))
[obj.insert(*interval) for interval in time_list]
return obj