# Licensed under a 3-clause BSD style license - see LICENSE.rst
# -*- coding: utf-8 -*-
"""
================
desiutil.bitmask
================
Mask bits for the spectro pipeline.
Individual packages will define their own mask bits and use this as a utility
access wrapper. Typical users will get their bitmasks pre-made from those
packages, not from here.
Stephen Bailey, Lawrence Berkeley National Lab
Fall 2015
Examples
--------
desispec_ could create a ccdmask like this:
>>> from desiutil.bitmask import BitMask
>>> import yaml
>>> _bitdefs = yaml.safe_load('''
... ccdmask:
... - [BAD, 0, "Pre-determined bad pixel (any reason)"]
... - [HOT, 1, "Hot pixel"]
... - [DEAD, 2, "Dead pixel"]
... - [SATURATED, 3, "Saturated pixel from object"]
... - [COSMIC, 4, "Cosmic ray"]
... ''')
...
>>> ccdmask = BitMask('ccdmask', _bitdefs)
Users would then access this mask with:
>>> from desispec.bitmasks import ccdmask
>>> ccdmask.COSMIC | ccdmask.SATURATED #- 2**4 + 2**3
24
>>> ccdmask.mask('COSMIC') # 2**4, same as ccdmask.COSMIC
16
>>> ccdmask.mask(4) # 2**4, same as ccdmask.COSMIC
16
>>> ccdmask.COSMIC # 2**4, same as ccdmask.mask('COSMIC')
16
>>> ccdmask.bitnum('COSMIC')
4
>>> ccdmask.bitname(4)
'COSMIC'
>>> ccdmask.names()
['BAD', 'HOT', 'DEAD', 'SATURATED', 'COSMIC']
>>> ccdmask.names(3)
['BAD', 'HOT']
>>> ccdmask.comment(0)
'Pre-determined bad pixel (any reason)'
>>> ccdmask.comment('COSMIC')
'Cosmic ray'
.. _desispec: https://desispec.readthedocs.io/en/latest/
"""
import numpy as np
[docs]
class _MaskBit(int):
"""A single mask bit.
Subclasses :class:`int` to act like an :class:`int`, but allows the
ability to extend with blat.name, blat.comment, blat.mask, blat.bitnum.
Attributes
----------
name : :class:`str`
The name of the bit.
bitnum : :class:`int`
The number of the bit. The value of the bit is ``2**bitnum``.
mask : :class:`int`
The value of the bit, ``2**bitnum``.
comment : :class:`str`
A comment explaining the meaning of the bit.
"""
def __new__(cls, name, bitnum, comment, extra=dict()):
self = super(_MaskBit, cls).__new__(cls, 2**bitnum)
self.name = name
self.bitnum = bitnum
self.mask = 2**bitnum
self.comment = comment
self._extra = extra
for key, value in extra.items():
if hasattr(self, key):
raise AttributeError(
"Bit {0} extra key '{1}' is already in use by int objects.".format(name, key))
self.__dict__[key] = value
return self
def __str__(self):
return ('{0.name:16s} bit {0.bitnum} mask 0x{0.mask:X} - ' +
'{0.comment}').format(self)
# We do not define __repr__(self) since the superclass has an arguably more useful __repr__().
def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
"""This function makes :class:`~desiutil.bitmask._MaskBit` objects "NumPy-aware",
so that they can be handled by NumPy `ufunc <https://numpy.org/doc/stable/glossary.html#term-ufunc>`_ methods.
Parameters
----------
ufunc : :class:`numpy.ufunc`
The NumPy function, *e.g.* :func:`~numpy.bitwise_or`.
method : :class:`str`
How `ufunc` was called. This is meant to be passed to another operand.
inputs : :class:`list`
The positional arguments.
kwargs : :class:`dict`
The keyword arguments.
Returns
-------
array-like
This function passes the actual computation to another entry in `inputs`.
Notes
-----
* In NumPy 1.x, subclasses of :class:`int` could be used exactly as if the
object was a plain Python :class:`int`.
* In Numpy 2.x, that is no longer the case, so we must explicitly use the
array API to tell NumPy how to handle the subclass.
* The solution is to modify the inputs so that a :class:`~desiutil.bitmask._MaskBit`
object is replaced by its ``.mask`` attribute, which is a plain Python :class:`int`.
* See GitHub issue `#218`_ for examples and discussion.
* This method is based on `this discussion`_.
.. _`#218`: https://github.com/desihub/desiutil/issues/218
.. _`this discussion`: https://numpy.org/doc/stable/user/basics.subclassing.html#array-ufunc-for-ufuncs
"""
#
# Filter inputs and replace any _MaskBit with a Python int.
# np.generic is the superclass of all NumPy scalars.
#
args = []
for input_ in inputs:
if isinstance(input_, _MaskBit):
args.append(input_.mask)
elif isinstance(input_, np.generic):
args.append(np.asarray(input_))
else:
args.append(input_)
#
# Identify another argument that we can pass the computation to.
#
sup = None
for arg in args:
if hasattr(arg, '__array_ufunc__'):
sup = arg
if sup is None:
return NotImplemented
return sup.__array_ufunc__(ufunc, method, *args, **kwargs)
# Class to provide mask bit utility functions
[docs]
class BitMask(object):
"""BitMask object to represent bit names, masks, and comments.
Typical users are not expected to create BitMask objects directly;
other packages like desispec and desitarget will have used this
to pre-create the bitmasks for them using definition files in those
packages.
Parameters
----------
name : :class:`str`
Name of this mask, must be key in `bitdefs`.
bitdefs : :class:`dict`
Dictionary of different mask bit definitions;
each value is a list of ``[bitname, bitnum, comment]``.
A 4th entry is optional, which must be a dictionary.
"""
def __init__(self, name, bitdefs):
"""Init.
"""
self._bits = dict()
self._name = name
for x in bitdefs[name]:
bitname, bitnum, comment = x[0:3]
if len(x) == 4:
extra = x[3]
if not isinstance(extra, dict):
raise ValueError(
'{} extra values should be a dict'.format(bitname))
else:
extra = dict()
self._bits[bitname] = _MaskBit(bitname, bitnum, comment, extra)
self._bits[bitnum] = self._bits[bitname]
def __getitem__(self, bitname):
"""Return mask for individual bitname.
"""
return self._bits[bitname]
[docs]
def bitnum(self, bitname):
"""Return bit number (int) for this `bitname` (string).
Parameters
----------
bitname : :class:`str`
The bit name.
Returns
-------
:class:`int`
The bit value.
"""
return self._bits[bitname].bitnum
[docs]
def bitname(self, bitnum):
"""Return bit name (string) for this `bitnum` (integer).
Parameters
----------
bitnum : :class:`int`
The number of the bit.
Returns
-------
:class:`str`
The name of the bit.
"""
return self._bits[bitnum].name
[docs]
def mask(self, name_or_num):
"""Return mask value.
Parameters
----------
name_or_num : :class:`int` or :class:`str`
Name of number of the mask.
Returns
-------
:class:`int`
The value of the mask.
Examples
--------
>>> bitmask.mask(3) # 2**3
8
>>> bitmask.mask('BLAT')
>>> bitmask.mask('BLAT|FOO')
"""
if isinstance(name_or_num, int):
return self._bits[name_or_num].mask
else:
mask = 0
for name in name_or_num.split('|'):
mask |= self._bits[name].mask
return mask
[docs]
def names(self, mask=None):
"""Return list of names of masked bits.
Parameters
----------
mask : :class:`int`, optional
The mask integer to convert to names. If not supplied,
return names of all known bits. `mask` can also be a numpy scalar
or a numpy array of length 1.
Returns
-------
:class:`list`
The list of names contained in the mask.
"""
names = list()
if mask is None:
# return names in sorted order of bitnum
bitnums = [x for x in self._bits.keys() if isinstance(x, int)]
for bitnum in sorted(bitnums):
names.append(self._bits[bitnum].name)
else:
if isinstance(mask, int) or isinstance(mask, np.integer):
pass
elif isinstance(mask, np.ndarray) and mask.shape == (1, ):
mask = mask[0]
else:
raise ValueError('Unknown type or invalid shape for mask!')
if mask < 0:
mask = np.int64(mask).astype(np.uint64)
# For Numpy < 2, the operation (int & np.uint64) was not defined.
# For Numpy >= 2, this operation should be harmless.
mask = int(mask)
bitnum = 0
while 2**bitnum <= mask:
if (2**bitnum & mask):
if bitnum in self._bits.keys():
names.append(self._bits[bitnum].name)
else:
names.append('UNKNOWN' + str(bitnum))
bitnum += 1
return names
def __getattr__(self, name):
"""Enable ``mask.BITNAME`` equivalent to ``mask['BITNAME']``.
"""
if name in self._bits:
return self._bits[name]
else:
raise AttributeError('Unknown mask bit name ' + name)
def __repr__(self):
'''Return yaml representation defining the bits of this bitmask.
'''
result = list()
result.append(self._name + ':')
# return names in sorted order of bitnum
bitnums = [x for x in self._bits.keys() if isinstance(x, int)]
for bitnum in sorted(bitnums):
bit = self._bits[bitnum]
# format the line for single bit, with or without extra keys
line = ' - [{:16s} {:2d}, "{}"'.format(
bit.name+',', bit.bitnum, bit.comment)
if len(bit._extra) > 0:
line = line + ', '+str(bit._extra)+']'
else:
line = line + ']'
result.append(line)
return "\n".join(result)