Source code for libertem.udf.masks

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

            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
            self.process_flat = self._process_flat_spsp
        elif (
            self.meta.array_backend in (
            ) 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
            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?

    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)
            # 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)
            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,))

[docs] 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 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', '', \ '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) * '': 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'] >>> (16, 16, 2) >>> np.allclose([..., 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(, 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, )