# 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