378 lines
15 KiB
Python
378 lines
15 KiB
Python
# 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
|
|
|
|
|
|
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
|
|
|
|
@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
|
|
|
|
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]
|
|
|
|
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
|
|
|
|
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)
|
|
|
|
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
|
|
|
|
|
|
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 = []
|
|
|
|
@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
|
|
|
|
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)
|
|
|
|
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
|