"""
General utility for bit-based masking.
Class usage examples
====================
.. include:: ../include/bitmaskarray_usage.rst
.. include common links, assuming primary doc root is up one directory
.. include:: ../include/links.rst
"""
from IPython import embed
import numpy as np
from astropy.io import fits
from pypeit.datamodel import DataContainer
from pypeit.core.bitmask import BitMask
from pypeit import log
from pypeit import PypeItBitMaskError
from pypeit import PypeItError
[docs]
class BitMaskArray(DataContainer):
"""
Object that holds both the mask data and the mask interpretation object.
This is an abstract class that should not be directly instantiated.
Args:
shape (:obj:`tuple`):
Shape of the mask to create.
asuint (:obj:`bool`, optional):
When setting the data-type for the mask array (see
:func:`~pypeit.core.bitmask.BitMask.minimum_dtype`), use an *unsigned*
integer instead of a signed integer (e.g., ``uint16`` instead of
``int16``).
"""
version = None
"""
DataContainer version. Must be defined by the subclass.
"""
datamodel = {'mask': dict(otype=np.ndarray, atype=np.integer, descr='Bitmask values')}
"""
Datamodel is simple, containing only the mask array.
"""
internals = ['lower_keys']
bitmask = None
"""
:class:`~pypeit.core.bitmask.BitMask` object used to interpret the bit array.
Must be defined by the subclass. When defining subclasses, note that the
bitmask flags *must* be case-insensitive strings.
"""
def __init__(self, shape, asuint=False):
# Instantiate as an empty DataContainer
super().__init__()
self._set_keys()
self.mask = np.zeros(shape, dtype=self.bitmask.minimum_dtype(asuint=asuint))
[docs]
@classmethod
def from_array(cls, arr):
# Instantiate using the shape of the provided array
self = cls(arr.shape, asuint=np.issubdtype(arr.dtype, np.unsignedinteger))
self.mask[...] = arr[...]
return self
[docs]
def _set_keys(self):
"""
Set :attr:`lower_keys`, which are needed for the bit access convenience
method.
"""
# Check the bitmask
keys = self.bit_keys()
if any([not isinstance(k, str) for k in keys]):
raise PypeItError(f'CODING ERROR: {self.bitmask.__class__.__name__} must only contain '
'string bit flags.')
self.lower_keys = [k.lower() for k in keys]
if len(np.unique(self.lower_keys)) != len(keys):
raise PypeItError('CODING ERROR: All bitmask keys must be case-insensitive and unique: '
f'{keys}')
[docs]
def __getattr__(self, item):
"""
Override the attribute access to allow for immediate construction and
access to boolean arrays that select array elements with the provided flag.
For example, if ``'BPM'`` is a flag in the :attr:`bitmask`, and ``mask``
is an instance of the the relevant subclass, ``mask.bpm`` is identical
to calling ``mask.flagged(flag='BPM')``.
.. warning::
Using this functionality creates a new numpy array *every time it is
called*. This can be slow for large arrays. This means that, if
you need to access the boolean array for a given flag multiple
times, you should set it to a new object (e.g., ``maskbpm =
mask.bpm``)!
If the attribute (``item``) requested is *not* one of the bitmask flags,
the :class:`~pypeit.datamodel.DataContainer` base class function is
called.
Args:
item (object):
The attribute being accessed.
"""
try:
i = self.lower_keys.index(item.lower())
except ValueError:
return super().__getattr__(item)
# TODO: This creates a new numpy array every time it is called, which
# could be slow for large arrays. Instead, we might want to do some
# clever lazy-loading to make this faster.
return self.flagged(flag=list(self.bit_keys())[i])
[docs]
def __getitem__(self, item):
"""Allow direct access to the mask."""
try:
return self.mask[item]
except:
return super().__getitem__(item)
def __setitem__(self, item, value):
try:
self.mask[item] = value
except:
super().__setitem__(item, value)
[docs]
def __or__(self, other):
"""Override or operation for mask."""
_self = self.copy()
_self.mask |= other.mask
return _self
[docs]
def __and__(self, other):
"""Override and operation for mask."""
_self = self.copy()
_self.mask &= other.mask
return _self
# TODO: This loses the description of the bits. Might be better to override
# to_hdu; although this gets sticky trying to figure out which hdu has the
# mask... Leaving it this way for now.
[docs]
def _bundle(self):
"""
Override the base-class bundle method so that the bitmask keys can be
added to the header.
Returns:
:obj:`list`: List of dictionaries indicating what should be written
to a file.
"""
d = super()._bundle()
d[0].update(self.bitmask.to_dict())
return d
[docs]
@classmethod
def from_hdu(cls, hdu, chk_version=True, **kwargs):
"""
Instantiate the object from an HDU extension.
This overrides the base-class method, only to add checks (or not) for
the bitmask.
Args:
hdu (`astropy.io.fits.HDUList`_, `astropy.io.fits.ImageHDU`_, `astropy.io.fits.BinTableHDU`_):
The HDU(s) with the data to use for instantiation.
chk_version (:obj:`bool`, optional):
If True, raise an error if the datamodel version or
type check failed. If False, throw a warning only.
**kwargs:
Passed directly to :func:`_parse`.
"""
# Run the default parser
d, version_passed, type_passed, parsed_hdus = cls._parse(hdu, **kwargs)
# Check
cls._check_parsed(version_passed, type_passed, chk_version=chk_version)
# Instantiate
self = super().from_dict(d=d)
# Check the bitmasks. Bits should have been written to *any* header
# associated with the object
hdr = hdu[parsed_hdus[0]].header if isinstance(hdu, fits.HDUList) else hdu.header
hdr_bitmask = BitMask.from_header(hdr)
if chk_version and hdr_bitmask.bits != self.bitmask.bits:
raise PypeItBitMaskError(
'The bitmask in this fits file appear to be out of date! Recreate this file by '
're-running the relevant script or set chk_version=False.'
)
return self
[docs]
def copy(self):
"""Create a deep copy."""
_self = super().__new__(self.__class__)
DataContainer.__init__(_self)
_self._set_keys()
_self.mask = self.mask.copy()
return _self
# NOTE: This function cannot be called "keys" because that would override
# the DataContainer base-class function!
[docs]
def bit_keys(self):
"""
Return a list of the bit keywords.
"""
return self.bitmask.keys()
@property
def bits(self):
"""Return the bit dictionary."""
return self.bitmask.bits
@property
def shape(self):
"""Return the shape of the internal array."""
return self.mask.shape
[docs]
def info(self):
"""
Print the list of bits and, if available, their descriptions.
"""
self.bitmask.info()
[docs]
def flagged(self, flag=None, invert=False):
"""
Determine if a bit is on in the provided bitmask value. The
function can be used to determine if any individual bit is on or
any one of many bits is on.
Args:
flag (:obj:`str`, array-like, optional):
One or more bit names to check. If None, then it checks
if *any* bit is on.
invert (:obj:`bool`, optional):
Invert the boolean such that unflagged pixels are True and
flagged pixels are False.
Returns:
`numpy.ndarray`_: Boolean array indicating where the internal
bitmask is flagged by the selected bits. Flagged values are True,
unflagged values are False. If ``invert`` is True, this is
reversed.
"""
indx = self.bitmask.flagged(self.mask, flag=flag)
return np.logical_not(indx) if invert else indx
[docs]
def flagged_bits(self, index):
"""
Return the list of flagged bit names for a single bit value.
Args:
index (:obj:`tuple`):
Tuple with the indices in the array.
Returns:
:obj:`list`: List of flagged bit value keywords.
"""
return self.bitmask.flagged_bits(self.mask[index])
[docs]
def toggle(self, flag, select=None):
"""
Toggle bits for selected array elements.
Args:
flag (:obj:`str`, array-like):
Bit name(s) to toggle.
select (:obj:`tuple`, :obj:`slice`, `numpy.ndarray`_, optional):
Object used to select elements of the mask array to at which to
toggle the provided bit flags. I.e., for the internal
:attr:`mask`, ``mask[select]`` must be a valid (fancy indexing)
operation. If None, the bit is toggled for the full mask!
"""
if select is None:
select = np.s_[...]
self.mask[select] = self.bitmask.toggle(self.mask[select], flag)
[docs]
def turn_on(self, flag, select=None):
"""
Ensure that a bit is turned on for the selected elements.
Args:
flag (:obj:`str`, array-like):
Bit name(s) to turn on.
select (:obj:`tuple`, :obj:`slice`, `numpy.ndarray`_, optional):
Object used to select elements of the mask array to at which to
turn on the provided bit flags. I.e., for the internal
:attr:`mask`, ``mask[select]`` must be a valid (fancy indexing)
operation. If None, the bit is turned on for the full mask!
"""
if select is None:
select = np.s_[...]
self.mask[select] = self.bitmask.turn_on(self.mask[select], flag)
[docs]
def turn_off(self, flag, select=None):
"""
Ensure that a bit is turned off in the provided bitmask value.
Args:
flag (:obj:`str`, array-like):
Bit name(s) to turn off.
select (:obj:`tuple`, :obj:`slice`, `numpy.ndarray`_, optional):
Object used to select elements of the mask array to at which to
turn off the provided bit flags. I.e., for the internal
:attr:`mask`, ``mask[select]`` must be a valid (fancy indexing)
operation. If None, the bit is turned off for the full mask!
"""
if select is None:
select = np.s_[...]
self.mask[select] = self.bitmask.turn_off(self.mask[select], flag)
[docs]
def consolidate(self, flag_set, consolidated_flag):
"""
Consolidate a set of flags into a single flag.
That is, any bit flagged with any of the flags provided by ``flag_set``
will also be flagged by ``consolidate_flag`` after executing this
function.
Args:
flag_set (:obj:`str`, array-like):
List of flags that are consolidated into a single flag.
consolidated_flag (:obj:`str`):
Consolidated flag name.
"""
self.turn_on(self.flagged(flag=flag_set), consolidated_flag)