src/libertem/udf/masks.py
from typing_extensions import Literal
from libertem.common.math import prod
import numpy as np
from libertem.common.udf import UDFMethod
from libertem.udf import UDF, UDFMeta
from libertem.common.buffers import AuxBufferWrapper
from libertem.common.container import MaskContainer
from libertem.common.numba import rmatmul
class ApplyMasksEngine:
def __init__(self, masks: MaskContainer, meta: UDFMeta, use_torch: bool):
self.masks = masks
self.meta = meta
try:
import torch
except ImportError:
torch = None
torch_incompatible = (
torch is None
or self.meta.input_dtype.kind != 'f'
or self.meta.input_dtype != self.masks.dtype
or self.meta.device_class != 'cpu'
or self.meta.array_backend != UDF.BACKEND_NUMPY
or self.masks.use_sparse
)
self.needs_transpose = True
if use_torch and (not torch_incompatible):
self.process_flat = self._process_flat_torch
elif (
self.meta.array_backend == UDF.BACKEND_NUMPY
and isinstance(self.masks.use_sparse, str)
and self.masks.use_sparse.startswith('scipy.sparse')
):
# Due to https://github.com/scipy/scipy/issues/13211
self.process_flat = self._process_flat_spsp
elif (
self.meta.array_backend in (
UDF.BACKEND_SCIPY_COO,
UDF.BACKEND_SCIPY_CSR,
UDF.BACKEND_SCIPY_CSC
) and isinstance(self.masks.use_sparse, str)
and self.masks.use_sparse.startswith('sparse.pydata')
):
self.process_flat = self._process_flat_sparsepyd
self.needs_transpose = False
else:
self.process_flat = self._process_flat_standard
def _get_masks(self):
return self.masks.get_for_sig_slice(
self.meta.sig_slice, transpose=self.needs_transpose
)
def _process_flat_torch(self, flat_tile, masks):
import torch
# CuPy back-end disables torch in get_task_data
# FIXME use GPU torch with CuPy array?
return torch.mm(
torch.from_numpy(flat_tile),
torch.from_numpy(masks),
).numpy()
def _process_flat_spsp(self, flat_tile, masks):
return rmatmul(flat_tile, masks)
def _process_flat_sparsepyd(self, flat_tile, masks):
# Make sure the sparse.pydata mask comes first
# to choose the right multiplication method
return (masks @ flat_tile.T).T
def _process_flat_standard(self, flat_tile, masks):
return flat_tile @ masks
def process_tile(self, tile):
flat_shape = (tile.shape[0], prod(tile.shape[1:]))
# Avoid reshape since older versions of scipy.sparse don't support it
flat_data = tile.reshape(flat_shape) if tile.shape != flat_shape else tile
return self.process_flat(flat_data, self._get_masks())
def process_frame_shifted(self, frame, shifts: tuple[int, ...]):
sig_shape = self.meta.dataset_shape.sig
masks = self._get_masks()
num_masks = len(self.masks)
shifted_slice = self.meta.sig_slice.shift_by(shifts)
inverse_shifted_slice = self.meta.sig_slice.shift_by(-1 * shifts)
left = self.meta.sig_slice.intersection_with(shifted_slice)
right = self.meta.sig_slice.intersection_with(inverse_shifted_slice)
if left.is_null():
# Zero overlap after shifts, shortcut return
return np.zeros((num_masks,), dtype=np.float32)
mask_slice = right.get()
if self.needs_transpose:
# expects masks in shape (sig_size, num_masks)
mask_slice = mask_slice + (slice(None), )
masks = masks.reshape((*sig_shape, num_masks))
final_mask_shape = (-1, num_masks)
else:
# expects masks in shape (num_masks, sig_size)
# NOTE unexpectedly don't need to reverse sig_shape or mask_slice ?
masks = masks.reshape((num_masks, *sig_shape))
mask_slice = (slice(None), ) + mask_slice
final_mask_shape = (num_masks, -1)
sliced_masks = masks[mask_slice].reshape(final_mask_shape)
# shift slicing requires sig_shape frames
# sparse array backend can provide flat frame
frame = frame.reshape(sig_shape)
try:
data = left.get(frame)
except TypeError as e:
# frame is in a form which doesn't support slicing
# the only recognized case is scipy.sparse.coo
if not hasattr(frame, 'getformat'):
raise e # pragma: no cover
assert frame.getformat() == 'coo'
frame = frame.tocsr()
data = left.get(frame)
flat_data = data.reshape((1, -1))
return self.process_flat(flat_data, sliced_masks).reshape((num_masks,))
class ApplyMasksUDF(UDF):
'''
Apply masks to signals/frames in the dataset. This can not only be used to integrate
over regions with a binary mask - the integration can be weighted by using
float or complex valued masks.
The result will be returned in a single sig-shaped buffer called intensity.
Its shape will be :code:`(*nav_shape, len(masks))`.
Parameters
----------
mask_factories : Union[Callable[[], array_like], Iterable[Callable[[], array_like]]]
Function or list of functions that take no arguments and create masks. The returned
masks can be
numpy arrays, scipy.sparse or sparse https://sparse.pydata.org/ matrices. The mask
factories should not reference large objects because they can create significant
overheads when they are pickled and unpickled. Each factory function should, when
called, return a numpy array with the same shape as frames in the dataset
(so dataset.shape.sig).
use_torch : bool, optional
Use pytorch back-end if available. Default True
use_sparse : Union[None, False, True, 'scipy.sparse', 'scipy.sparse.csc', \
'sparse.pydata'], optional
Which sparse back-end to use.
* None (default): Where possible use sparse matrix multiplication if all factory \
functions return a sparse mask, otherwise convert all masks to dense matrices \
and use dense matrix multiplication
* True: Convert all masks to sparse matrices and use default sparse back-end.
* False: Convert all masks to dense matrices.
* 'scipy.sparse': Use scipy.sparse.csr_matrix (default sparse)
* 'scipy.sparse.csc': Use scipy.sparse.csc_matrix
* 'sparse.pydata': Use sparse.pydata COO matrix
mask_count : int, optional
Specify the number of masks if a single factory function is used so that the
number of masks can be determined without calling the factory function.
mask_dtype : numpy.dtype, optional
Specify the dtype of the masks so that mask dtype
can be determined without calling the mask factory functions. This can be used to
override the mask dtype in the result dtype determination. As an example, setting
this to np.float32 means that masks of type float64 will not switch the calculation
and result dtype to float64 or complex128.
preferred_dtype : numpy.dtype, optional
Let :meth:`get_preferred_input_dtype` return the specified type instead of the
default `float32`. This can perform the calculation with integer types if both input
data and mask data are compatible with this.
backends : Iterable containing strings "numpy" and/or "cupy", or None
Control which back-ends are used. Default is numpy and cupy
shifts : Union[Tuple[int, int], AuxBufferWrapper], optional
(Y/X)-shifts to apply to all masks before multiplying with each frame. Can be either a
length-2 array-like for a constant :code:`(Y, X)` shift, or an
:class:`~libertem.common.buffers.AuxBufferWrapper` of
:code:`(kind='nav', extra_shape=(2,), dtype=int)` defining a per-frame shift to apply.
A positive y-shift moves the mask 'down' relative to the frame, while a positive
x-shift moves the mask 'right' relative to the frame. Elements of the mask and frame
which do not overlap after the shift are discarded. A shift resulting in no overlap
at all will return a sum of :code:`0.` for that frame.
.. note::
Float shift values are cast to integers internally; round values before
passing the shifts argument to better control the exact shifts used.
.. note::
The :code:`shifts` parameter requires frame-by-frame processing to function, and so
adds a performance penalty compared to unshifted mask application. If applying a
constant shift it may be worthwhile to manually create a new, pre-shifted mask
rather than relying on this feature.
Shifting is also currently incompatible with :code:`scipy.sparse` masks. If sparse
processing is required then where possible :code:`scipy.sparse` masks are converted to
:code:`sparse.pydata` equivalents. A consequence of this is that sparse processing
is not yet supported through CuPy when shifts are enabled, as :code:`sparse.pydata`
has no current CuPy implementation. If sparse masks are supplied on a CuPy backend
when :code:`use_sparse=None` (the default) they will be densified to allow the
calculation to take place.
Examples
--------
>>> dataset.shape
(16, 16, 32, 32)
>>> def my_masks():
... return [np.ones((32, 32)), np.zeros((32, 32))]
>>> udf = ApplyMasksUDF(mask_factories=my_masks)
>>> res = ctx.run_udf(dataset=dataset, udf=udf)['intensity']
>>> res.data.shape
(16, 16, 2)
>>> np.allclose(res.data[..., 1], 0) # same order as in the mask factory
True
Mask factories can also return all masks as a single array, stacked on the first axis:
>>> def my_masks_2():
... masks = np.zeros((2, 32, 32))
... masks[1, ...] = 1
... return masks
>>> udf = ApplyMasksUDF(mask_factories=my_masks_2)
>>> res_2 = ctx.run_udf(dataset=dataset, udf=udf)['intensity']
>>> np.allclose(res_2.data, res.data)
True
Masks can be shifted relative to the data using the :code:`shifts` parameter,
this can either be a constant shift for all frames:
>>> udf = ApplyMasksUDF(mask_factories=my_masks, shifts=(2, -5))
>>> res_shift_constant = ctx.run_udf(dataset=dataset, udf=udf)['intensity']
or a per-frame shift supplied using an :class:`~libertem.common.buffers.AuxBufferWrapper`
created using :meth:`~libertem.udf.base.UDF.aux_data`:
>>> udf = ApplyMasksUDF(
... mask_factories=my_masks,
... shifts=ApplyMasksUDF.aux_data(
... np.random.randint(-8, 8, size=(16, 16, 2)).ravel(),
... kind='nav',
... extra_shape=(2,),
... dtype=int,
... )
... )
>>> res_shift_aux = ctx.run_udf(dataset=dataset, udf=udf)['intensity']
.. versionadded:: 0.4.0
.. versionchanged:: 0.13.0
Added the :code:`shifts` parameter
'''
def __init__(self, mask_factories, use_torch=True, use_sparse=None, mask_count=None,
mask_dtype=None, preferred_dtype=None, backends=None, shifts=None, **kwargs):
_backends = backends
if backends is None:
backends = self.BACKEND_ALL
backends = tuple(b for b in backends if b in self.BACKEND_ALL)
if shifts is not None:
if isinstance(use_sparse, str) and use_sparse.startswith('scipy.sparse'):
# This is 'unsupported' because we need to slice the mask stack
# in the signal dimensions to shift it, and sig is flattened
# to give to 2D matrix in scipy.sparse
raise ValueError(
f'Sparse backend {use_sparse} not supported for '
'shifts, use sparse.pydata instead.'
)
if not isinstance(shifts, AuxBufferWrapper):
shifts = np.asarray(shifts)
backends = tuple(
b for b in backends
if b not in (
# Here we are doing frame-by-frame processing, so we can
# accept scipy.sparse-style frames, however we need to
# perform a reshape into sig-shaped frames which normally
# casts the frame into coo_matrix form, which then
# has to be re-cast into csr_matrix form to be sliced (i.e. shifted)
self.BACKEND_SCIPY_COO, # cannot be sliced
self.BACKEND_CUPY_SCIPY_COO, # cannot be sliced
)
)
if len(backends) == 0:
raise ValueError(f'No compatible backend found in {_backends}')
self._mask_container = None
super().__init__(
mask_factories=mask_factories,
use_torch=use_torch,
use_sparse=use_sparse,
mask_count=mask_count,
mask_dtype=mask_dtype,
preferred_dtype=preferred_dtype,
backends=backends,
shifts=shifts,
**kwargs
)
def get_preferred_input_dtype(self):
''
if self.params.preferred_dtype is None:
return super().get_preferred_input_dtype()
else:
return self.params.preferred_dtype
def get_mask_dtype(self):
if self.params.mask_dtype is None:
return self.masks.dtype
else:
return self.params.mask_dtype
def get_mask_count(self):
if self.params.mask_count is None:
return len(self.masks)
else:
return self.params.mask_count
@property
def masks(self):
if self._mask_container is None:
self._mask_container = self._make_mask_container()
return self._mask_container
def _make_mask_container(self):
p = self.params
if self.meta.array_backend in self.CUPY_BACKENDS:
backend = self.BACKEND_CUPY
else:
backend = self.BACKEND_NUMPY
# In the default case defer to default kwarg on MaskContainer
default_sparse = {}
if p.shifts is None:
default_sparse['default_sparse'] = 'scipy.sparse'
else:
default_sparse['default_sparse'] = 'sparse.pydata'
return MaskContainer(
p.mask_factories, dtype=p.mask_dtype, use_sparse=p.use_sparse, count=p.mask_count,
backend=backend, **default_sparse,
)
def get_task_data(self):
''
engine = ApplyMasksEngine(self.masks, self.meta, self.params.use_torch)
return {
'engine': engine,
}
def get_result_buffers(self):
''
dtype = np.result_type(self.meta.input_dtype, self.get_mask_dtype())
count = self.get_mask_count()
return {
'intensity': self.buffer(
kind='nav', extra_shape=(count, ), dtype=dtype, where='device'
)
}
def get_backends(self):
''
return self.params.backends
def get_method(self) -> Literal[UDFMethod.FRAME, UDFMethod.TILE]:
"""
:meta private:
"""
if self.params.get('shifts') is not None:
return UDFMethod.FRAME
else:
return UDFMethod.TILE
def process_tile(self, tile):
"""
Used for simple mask application, without shifts
:meta private:
"""
self.results.intensity[:] += self.forbuf(
self.task_data.engine.process_tile(tile),
self.results.intensity,
)
def process_frame(self, frame):
"""
Apply shifted masks to a frame
:meta private:
"""
shifts = self.params.shifts.astype(int)
self.results.intensity[:] += self.forbuf(
self.task_data.engine.process_frame_shifted(frame, shifts),
self.results.intensity,
)