# collection.py: Data collection classes
#
# Authors: William Cleveland (USRA),
# Adam Goldstein (USRA) and
# Daniel Kocevski (NASA)
#
# Portions of the code are Copyright 2020 William Cleveland and
# Adam Goldstein, Universities Space Research Association
# All rights reserved.
#
# Written for the Fermi Gamma-ray Burst Monitor (Fermi-GBM)
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
#
import re
from functools import partial
from collections import OrderedDict
[docs]class DataCollection:
"""A container for a collection of like data objects, such as a collection
of :class:`Ctime` objects. This class exposes the individual objects'
attributes, and it exposes the methods of the object class so that methods
can be called on the collection as a whole. For that reason, each object in
the DataCollection must be of the same type, otherwise an error is raised.
The type of the collection is set by the first object inserted into the
collection and the collection type is immutable thereafter.
Objects are stored in the collection in the order they are added.
The number of items in the collection can be retrieved by ``len()`` and
one can iterate over the items::
[data_item for data_item in DataCollection]
In addition to the DataCollection methods, all of the individual object
attributes and methods are exposed, and they become methods of the
DataCollection. Note that individual object attributes become *methods*
i.e. if you have an item attribute called item.name, then the corresponding
DataCollection method would be item.name().
Attributes:
items (list): The names of the items in the DataCollection
types (str): The type of the objects in the DataCollection
"""
def __init__(self):
self._data_dict = OrderedDict()
self._type = None
def __iter__(self):
for item in self._data_dict.values():
yield item
def __len__(self):
return len(self._data_dict)
def _enforce_type(self, data_item):
if not isinstance(data_item, self._type) and self._type is not None:
raise TypeError(
'Incorrect data item for {}'.format(self.__class__.__name__))
@property
def items(self):
return list(self._data_dict.keys())
@property
def types(self):
return self._type
[docs] @classmethod
def from_list(cls, data_list, names=None):
"""Given a list of objects and optionally a list of corresponding names,
create a new DataCollection.
Args:
data_list (list of :obj:`objects`):
The list of objects to be in the collection
names (list of :obj:`str`, optional):
The list of corresponding names to the objects. If not set,
will try to retrieve a name from object.filename (assuming it's
a data object). If that fails, each item will be named
ambiguously 'item1', 'item2', etc.
Returns
:py:class:`DataCollection`: The newly created collection
"""
obj = cls()
# set the names
if names is not None:
if len(names) != len(data_list):
raise ValueError('Names list must be same size as data list')
else:
names = [None] * len(data_list)
# include the objects
for data_item, name in zip(data_list, names):
obj.include(data_item, name=name)
return obj
[docs] def to_list(self):
"""Return the objects contained in the DataCollection as a list.
Returns:
(list of :obj:`objects`):
The list of objects, in the order that they were inserted
"""
return [self.get_item(name) for name in self.items]
[docs] def include(self, data_item, name=None):
"""Insert an object into the collection. The first item inserted will
set the immutable type.
Args:
data_item (:obj:`object`): A data object to include
name (str, optional):
An optional corresponding name. If not set, will try to
retrieve a name from object.filename (assuming it's a data
object). If that fails, each item will be named ambiguously
'item1', 'item2', etc.
"""
# if this is the first item inserted, set the type of the Collection
# and expose the attributes and methods of the object
if len(self) == 0:
self._type = type(data_item)
dir = [key for key in data_item.__dir__() if
not re.match('_.', key)]
for key in dir:
setattr(self, key, partial(self._method_call, key))
else:
# otherwise, ensure that each object inserted is of the same type
if type(data_item) != self._type:
raise TypeError('A DataCollection must contain like objects')
# insert with user-defined name
if name is not None:
self._data_dict[name] = data_item
else:
# or try to insert using filename attribute
try:
self._data_dict[data_item.filename] = data_item
# otherwise default to ambiguity
except AttributeError:
self._data_dict['item{}'.format(len(self) + 1)] = data_item
[docs] def remove(self, item_name):
"""Remove an object from the collection given the name
Args:
item_name (str): The name of the item to remove
"""
self._data_dict.pop(item_name)
[docs] def get_item(self, item_name):
"""Retrieve an object from the DataCollection by name
Args:
item_name (str): The name of the item to retrieve
Returns:
:obj:`object`: The retrieved data item
"""
return self._data_dict[item_name]
def _method_call(self, method_name, *args, **kwargs):
"""This is the wrapper for the exposde attribute and method calls.
Applies method_name over all items in the DataCollection
Args:
method_name (str): The name of the method or attribute
*args: Additional arguments to be passed to the method
**kwargs: Additional keyword arguments to be passed to the method
Returns:
None or list: If not None, will return the results from all
objects in the list
"""
# get the attributes/methods for each item
refs = [getattr(obj, method_name) for obj in self._data_dict.values()]
# if method_name is a method, then it will be callable
if callable(refs[0]):
res = [getattr(obj, method_name)(*args, **kwargs)
for obj in self._data_dict.values()]
# otherwise, method_name will not be callable if it is an attribute
else:
# we are setting an attribute
if len(args) != 0:
res = [setattr(obj, method_name, *args)
for obj in self._data_dict.values()]
# we are retrieving an attribute
else:
res = refs
if res[0] is not None:
return res
[docs]class GbmDetectorCollection(DataCollection):
"""A container for a collection of GBM-specific data objects, such as a
collection of ``Ctime`` objects from different detectors.
The special behavior of this class is to provide a way to interact with
a collection of detector data that may contain a mix of different *types*
of detectors. For example, many times we want a collection of GBM NaI
and GBM BGO detectors. These detectors have very different energy ranges,
and so may require different inputs for a variety of functions. This
collection allows one to specify the different arguments for NaI and BGO
data without having to implement many ugly and space-wasting loops and
``if...else`` decisions.
In addition to the GbmDetectorCollection methods, all of the individual
object attributes and methods are exposed, and they become methods of the
DataCollection. Note that individual object attributes become *methods*
i.e. if you have an item attribute called item.name, then the corresponding
DataCollection method would be item.name().
Attributes:
items (list): The names of the items in the DataCollection
types (str): The type of the objects in the DataCollection
"""
def __init__(self):
super().__init__()
self._dets = []
[docs] @classmethod
def from_list(cls, data_list, names=None, dets=None):
"""Given a list of objects and optionally a list of corresponding names
and corresponding detector names, create a new GbmDetectorCollection.
Args:
data_list (list of :obj:`objects`):
The list of objects to be in the collection
names (list of :obj:`str`, optional):
The list of corresponding names to the objects. If not set,
will try to retrieve a name from object.filename (assuming it's
a data object). If that fails, each item will be named
ambiguously 'item1', 'item2', etc.
dets (list of :obj:`str`, optional):
The detector names for each object. If not set, will try to
retrieve from the object.detector attribute. If that attribute
doesn't exist, an error will be raised, and the user will need
to specify this list.
Returns
:py:class:`GbmDetectorCollection`: The newly created collection
"""
obj = cls()
# set the detector names
if dets is not None:
if len(dets) != len(data_list):
raise ValueError(
'Detector list must be same size as data list')
else:
try:
dets = [data_item.detector for data_item in data_list]
except:
raise AttributeError('Cannot find detector information. '
'Need to manually set')
# set the names
if names is not None:
if len(names) != len(data_list):
raise ValueError('Names list must be same size as data list')
else:
names = [None] * len(data_list)
# include the objects
[obj.include(data_item, det, name=name) for (data_item, det, name)
in zip(data_list, dets, names)]
return obj
[docs] def remove(self, item_name):
"""Remove an object from the collection given the name
Args:
item_name (str): The name of the item to remove
"""
index = [item == item_name for item in self.items].index(True)
self._dets.pop(index)
self._data_dict.pop(item_name)
[docs] def include(self, data_item, det, name=None):
"""Insert an object into the GbmDetectorCollection. The first item
inserted will set the immutable type.
Args:
data_item (:obj:`object`): A data object to include
det (str): The corresponding detector for the item
name (str, optional):
An optional corresponding name. If not set, will try to
retrieve a name from object.filename (assuming it's a data
object). If that fails, each item will be named ambiguously
'item1', 'item2', etc.
"""
super().include(data_item, name=None)
self._dets.append(det)
def _method_call(self, method_name, *args, nai_args=(), nai_kwargs=None,
bgo_args=(), bgo_kwargs=None, **kwargs):
"""This is the wrapper for the attribute and method calls. Applies
method_name over all items in the GbmDetectorCollection.
Args:
method_name (str): The name of the method or attribute
*args: Additional arguments to be passed to the method
nai_args: Arguments to be applied only to the NaI objects
bgo_args: Arguments to be applied only to the BGO objects
nai_kwargs: Keywords to be applied only to the NaI objects
bgo_kwargs: Keywords to be applied only to the BGO objects
**kwargs: Additional keyword arguments to be passed to the
method. Will be applied to both NaI and BGO objects
and will be appended to any existing keywords from
nai_kwargs or bgo_kwargs
Returns:
None or list: If not None, will return the results from all objects
in the list
"""
if nai_kwargs is None:
nai_kwargs = {}
if bgo_kwargs is None:
bgo_kwargs = {}
if len(args) > 0:
nai_args = args
bgo_args = args
if len(kwargs) > 0:
nai_kwargs.update(kwargs)
bgo_kwargs.update(kwargs)
# get the attributes/methods for each item
refs = [getattr(obj, method_name) for obj in self._data_dict.values()]
# if method_name is a method, then it will be callable
if callable(refs[0]):
res = []
for obj, det in zip(self._data_dict.values(), self._dets):
if 'n' in det:
our_args = nai_args
our_kwargs = nai_kwargs
elif 'b' in det:
our_args = bgo_args
our_kwargs = bgo_kwargs
res.append(getattr(obj, method_name)(*our_args, **our_kwargs))
# otherwise, method_name will not be callable if it is an attribute
else:
# we are setting an attribute
if len(nai_args) != 0 or len(bgo_args) != 0:
res = []
for obj, det in zip(self._data_dict.values(), self._dets):
if 'n' in obj.detector:
our_args = nai_args
elif 'b' in obj.detector:
our_args = bgo_args
res.append(setattr(obj, method_name, *args))
# we are retrieving an attribute
else:
res = refs
if res[0] is not None:
return res