Add PixelMask class

eitprocessing/datahandling/pixelmap.py

@@ -51,6 +51,8 @@
     from matplotlib.axes import Axes
     from matplotlib.image import AxesImage
 
+    from eitprocessing.roi import PixelMask
+
 ColorType = str | tuple[float, float, float] | tuple[float, float, float, float] | float | Colormap
 
 T = TypeVar("T", bound="PixelMap")
@@ -327,46 +329,60 @@ def _check_normalization_reference(reference_: float) -> None:
         if reference_ < 0:
             warnings.warn("Normalization by a negative number may lead to unexpected results.", UserWarning)
 
-    def threshold(
+    def create_mask_from_threshold(
         self,
-        threshold: npt.ArrayLike,
+        threshold: float,
         *,
         comparator: Callable = np.greater_equal,
         absolute: bool = False,
-        keep_sign: bool = False,
-        fill_value: float = np.nan,
-        **return_attrs: dict | None,
-    ) -> Self:
-        """Threshold the pixel map values.
+    ) -> PixelMask:
+        """Create a pixel mask from the pixel map based on threshold values.
 
-        This method applies a threshold to the pixel map values, setting values that do not meet the threshold condition
-        to a specified fill value. The comparison is done using the provided comparator function (default is `>=`).
+        The values of the pixel map are compared to the threshold values. By default, the comparator is `>=`
+        (`np.greater_equal`), such that the resulting mask is 1.0 where the map values are at least the threshold
+        values, and NaN elsewhere. The comparator can be set to any comparison function, e.g.`np.less`, a function from
+        the `operator` module or custom function which takes pixel map values array and threshold as arguments, and
+        returns a boolean array with the same shape as the array.
 
-        If `absolute` is True, the threshold is applied to the absolute values of the pixel map. If `keep_sign` is True,
-        the sign of the original pixel values is retained when filling with the `fill_value`. Otherwise, the fill value
-        is applied uniformly.
+        If `absolute` is True, absolute values are compared to the threshold.
 
-        The `threshold` method returns a new instance of the same class with the modified values. Other attributes of
-        the returned object can be set using keyword arguments.
+        The shape of the pixel mask is the same as the shape of the pixel map.
 
         Args:
             threshold (float): The threshold value.
             comparator (Callable): A function that compares pixel values against the threshold.
             absolute (bool): If True, apply the threshold to the absolute values of the pixel map.
-            keep_sign (bool): If True, retain the sign of original values in the filled values.
-            fill_value (float): The value to set for pixels that do not meet the threshold condition.
-            **return_attrs (dict | None): Additional attributes to pass to the new PixelMap instance.
 
         Returns:
-            Self: A new object instance with the thresholded values.
+            PixelMask:
+                A PixelMask instance with values 1.0 where comparison is true, and NaN elsewhere.
+
+        Raises:
+            TypeError: If `threshold` is not a float or `comparator` is not callable.
+
+        Examples:
+        >>> pm = PixelMap([[0.1, 0.5, 0.9]])
+        >>> mask = pm.create_mask_from_threshold(0.5)
+        PixelMask(mask=array([[nan,  1.,  1.]]))
+        >>> mask.apply(pm)
+        PixelMap(values=array([[nan, 0.5, 0.9]]), ...)
+
+        >>> mask = pm.create_mask_from_threshold(0.5, comparator=np.less)
+        PixelMask(mask=array([[ 1., nan, nan]]))
         """
-        # TODO: add mode argument to allow for percentage/actual value thresholding
-        compare_values = np.abs(self.values) if absolute else self.values
-        sign = np.sign(self.values) if keep_sign else 1.0
-        new_values = np.where(comparator(compare_values, threshold), self.values, fill_value * sign)
+        if not isinstance(threshold, (float, np.floating, int, np.integer)):
+            msg = "`threshold` must be a number."
+            raise TypeError(msg)
+
+        if not callable(comparator):
+            msg = "`comparator` must be a callable function."
+            raise TypeError(msg)
+
+        from eitprocessing.roi import PixelMask
 
-        return_attrs = return_attrs or {}
-        return replace(self, values=new_values, **return_attrs)
+        compare_values = np.abs(self.values) if absolute else self.values
+        mask_values = comparator(compare_values, threshold)
+        return PixelMask(mask_values)
 
     def imshow(
         self,

eitprocessing/roi/__init__.py

@@ -0,0 +1,185 @@
+"""Region of Interest selection and pixel masking.
+
+This module contains tools for the selection of regions of interest and masking pixel data. The central class of this
+module is `PixelMask`. Any type of region of interest selection results in a `PixelMask` object. A mask can be applied
+to any pixel dataset (EITData, PixelMap) with the same shape.
+
+Several default masks have been predefined. NB: the right lung is to the left side of the EIT image and vice versa.
+
+- `VENTRAL_MASK` includes only the first 16 rows;
+- `DORSAL_MASK` includes only the last 16 rows;
+- `RIGHT_LUNG_MASK` includes only the first 16 columns;
+- `LEFT_LUNG_MASK` includes only the last 16 columns;
+- `QUADRANT_1_MASK` includes the top right quadrant;
+- `QUADRANT_2_MASK` includes the top left quadrant;
+- `QUADRANT_3_MASK` includes the bottom right quadrant;
+- `QUADRANT_4_MASK` includes the bottom left quadrant;
+- `LAYER_1_MASK` includes only the first 8 rows;
+- `LAYER_2_MASK` includes only the second set of 8 rows;
+- `LAYER_3_MASK` includes only the third set of 8 rows;
+- `LAYER_4_MASK` includes only the last 8 rows.
+"""
+
+import dataclasses
+import sys
+from dataclasses import InitVar, dataclass, field
+from dataclasses import replace as dataclass_replace
+from typing import Self, TypeVar, overload
+
+import numpy as np
+
+from eitprocessing.datahandling.eitdata import EITData
+from eitprocessing.datahandling.pixelmap import PixelMap
+
+T = TypeVar("T", np.ndarray, EITData, PixelMap)
+
+
+@dataclass(frozen=True)
+class PixelMask:
+    """Mask pixels by selecting or weighing them individually.
+
+    A mask is a 2D array with a value for each pixel. Most often, this value is NaN (`np.nan`, 'not a number') or 1, and
+    less commonly a value between 0 and 1. NaN values indicate the pixel is not part of the region of interest, e.g.,
+    falls outside the functional lung space, or is not part of the ventral region of the lung. A value of 1 indicates
+    the pixel is included in the region of interest. A value between 0 and 1 indicates that the pixel is part of the
+    region of interest, but is weighted, e.g., for a weighted summation of pixel values, or because the pixel is
+    considered part of multiple regions of interest.
+
+    You can initialize a mask using a nested list. At initialization, the mask is converted to a numpy array.
+
+    By default, 0-values are converted tot NaN. You can override this behaviour with `keep_zeros=True`. You can
+    therefore create a mask by supplying boolean values, where `True` indicates the pixel is part of the region of
+    interest (`True` equals 1), and `False` indicates it is not (`False` equals 0, and will be converted to NaN).
+
+    Since masking is not intended for other operations, masking values that are negative or higher than 1 will result in
+    a `ValueError`. You can override this check with `suppress_value_range_error=True`.
+
+    A mask can be applied to any pixel dataset, such as an `EITData` object or a `PixelMap` object. The mask is applied
+    to the last two dimensions of the data, which must match the shape of the mask. The mask is applied by multiplying
+    each pixel in the dataset by the corresponding masking value. Multiplication by NaN always results in NaN.
+
+    Masks can be combined by either adding or multiplying them. Adding masks results in a mask that includes all pixels
+    that are in either mask. Multiplying masks results in a mask that includes only pixels that are in both masks.
+
+    Example:
+    ```python
+    >>> assert VENTRAL_MASK + RIGHT_LUNG_MASK == QUADRANT_1_MASK
+    True  # quadrant 1 is the ventral part of the right lung
+    >>> assert DORSAL_MASK * LEFT_LUNG_MASK == QUADRANT_4_MASK
+    True  # quadrant 4 is the dorsal part of the left lung
+    ```
+
+    """
+
+    mask: np.ndarray
+    keep_zeros: InitVar[bool] = field(default=False, kw_only=True)
+    suppress_value_range_error: InitVar[bool] = field(default=False, kw_only=True)
+
+    def __init__(self, mask: list | np.ndarray, keep_zeros: bool = False, suppress_value_range_error: bool = False):
+        mask = np.array(mask, dtype=float)
+
+        if mask.ndim != 2:  # noqa: PLR2004
+            msg = f"Mask should be a 2D array, not {mask.ndim}D."
+            raise ValueError(msg)
+
+        if (not suppress_value_range_error) and (np.nanmax(mask) > 1 or np.nanmin(mask) < 0):
+            msg = "One or more mask values fall outside the range 0 to 1."
+            exc = ValueError(msg)
+            if sys.version_info >= (3, 11):
+                exc.add_note("Provided values should normally be a boolean value or a number from 0 to 1.")
+                exc.add_note(
+                    "In case you need a mask with values outside this range, "
+                    "provide `suppress_value_range_warning=True` when initializing a Mask."
+                )
+            raise exc
+
+        if not keep_zeros:
+            mask[mask == 0] = np.nan
+
+        mask.flags["WRITEABLE"] = False
+        object.__setattr__(self, "mask", mask)
+
+    @overload
+    def apply(self, data: np.ndarray) -> np.ndarray: ...
+
+    @overload
+    def apply(self, data: EITData, **kwargs) -> EITData: ...
+
+    @overload
+    def apply(self, data: "PixelMap", **kwargs) -> "PixelMap": ...
+
+    def apply(self, data, **kwargs):
+        """Apply pixel mask to data, returning a copy of the object with pixel values masked.
+
+        Data can be a numpy array, an EITData object or PixelMap object. In case of an EITData object, the mask will be
+        applied to the `pixel_impedance` attribute. In case of a PixelMap, the mask will be applied to the `values`
+        attribute.
+
+        The input data can have any dimension. The mask is applied to the last two dimensions. The size of the last two
+        dimensions must match the size of the dimensions of the mask, and will generally (but do not have to) have the
+        length 32.
+
+        The function returns the same data type as `data`. In case of `EITData` or `PixelMap` data, the object will have
+        the provided label, or the original data label if none is provided.
+        """
+
+        def transform_and_mask(data: np.ndarray) -> np.ndarray:
+            """Transform the mask to ensure it has the correct shape for the given data, and apply the mask.
+
+            The mask is transformed by adding new axes to the beginning of the array, such that the number of dimensions
+            match the number of dimensions of the data. The last two dimensions will contain the mask itself. This
+            allows the mask to be applied correctly, even if the data has more than two dimensions (e.g., a 3D array
+            with shape (time, channels, rows, cols)).
+            """
+            if self.mask.shape[-2:] != data.shape[-2:]:
+                msg = (
+                    f"Data shape {data.shape} does not match Mask shape {self.mask.shape}. "
+                    "The last two dimensions of the mask and data must match."
+                )
+                raise ValueError(msg)
+            mask = self.mask[*([np.newaxis] * (data.ndim - 2)), ...]
+            return data * mask
+
+        match data:
+            case np.ndarray():
+                return transform_and_mask(data)
+            case EITData():
+                return dataclass_replace(data, pixel_impedance=transform_and_mask(data.pixel_impedance), **kwargs)
+            case PixelMap():
+                return data.update(values=transform_and_mask(data.values), **kwargs)
+            case _:
+                msg = f"Data should be an array, or EITData or PixelMap object, not {type(data)}."
+                raise TypeError(msg)
+
+    @property
+    def is_weighted(self) -> bool:
+        """Whether the mask multiplies any pixels with a number other than NaN or 1."""
+        return not bool(np.all(np.isnan(self.mask) | (self.mask == 1.0)))
+
+    def __mul__(self, other: Self) -> Self:
+        """Combine masks by multiplying masking values."""
+        return dataclasses.replace(self, mask=self.mask * other.mask)
+
+    def __add__(self, other: Self) -> Self:
+        """Combine masks by adding masking values.
+
+        Values are clipped at 1, so that the resulting mask does not contain values higher than 1.
+        """
+        return dataclasses.replace(self, mask=np.clip(np.nansum([self.mask, other.mask], axis=0), a_min=None, a_max=1))
+
+
+LAYER_1_MASK = PixelMask(np.concat([np.ones((8, 32)), np.zeros((24, 32))], axis=0))
+LAYER_2_MASK = PixelMask(np.concat([np.zeros((8, 32)), np.ones((8, 32)), np.zeros((16, 32))], axis=0))
+LAYER_3_MASK = PixelMask(np.concat([np.zeros((16, 32)), np.ones((8, 32)), np.zeros((8, 32))], axis=0))
+LAYER_4_MASK = PixelMask(np.concat([np.zeros((24, 32)), np.ones((8, 32))], axis=0))
+
+VENTRAL_MASK = LAYER_1_MASK + LAYER_2_MASK
+DORSAL_MASK = LAYER_3_MASK + LAYER_4_MASK
+
+RIGHT_LUNG_MASK = PixelMask(np.concat([np.ones((32, 16)), np.zeros((32, 16))], axis=1))
+LEFT_LUNG_MASK = PixelMask(np.concat([np.zeros((32, 16)), np.ones((32, 16))], axis=1))
+
+QUADRANT_1_MASK = VENTRAL_MASK * RIGHT_LUNG_MASK
+QUADRANT_2_MASK = VENTRAL_MASK * LEFT_LUNG_MASK
+QUADRANT_3_MASK = DORSAL_MASK * RIGHT_LUNG_MASK
+QUADRANT_4_MASK = DORSAL_MASK * LEFT_LUNG_MASK

tests/test_pixelmap.py

@@ -135,29 +135,20 @@ def test_init_plot_parameters():
     assert pm11.plot_parameters.colorbar_kwargs == {"key": "another value"}
 
 
-def test_threshold():
-    pm = PixelMap(np.reshape(np.arange(-50, 50), (10, 10)))
-    pm_threshold = pm.threshold(10)
+def test_create_threshold_mask():
+    pm = PixelMap([[-2, -1, 0, 1, 2]])
 
-    assert type(pm) is type(pm_threshold)
+    mask = pm.create_mask_from_threshold(1)
+    assert np.array_equal(mask.mask, [[np.nan, np.nan, np.nan, 1.0, 1.0]], equal_nan=True)
 
-    pm = ODCLMap(np.reshape(np.arange(-50, 50), (10, 10)))
-    pm_threshold = pm.threshold(10)
+    mask = pm.create_mask_from_threshold(1, comparator=np.greater)
+    assert np.array_equal(mask.mask, [[np.nan, np.nan, np.nan, np.nan, 1.0]], equal_nan=True)
 
-    assert np.nanmin(pm_threshold.values) == 10
-    non_nan_values = pm_threshold.values[~np.isnan(pm_threshold.values)]
-    assert np.all(non_nan_values >= 10)
+    mask = pm.create_mask_from_threshold(1, absolute=True)
+    assert np.array_equal(mask.mask, [[1.0, 1.0, np.nan, 1.0, 1.0]], equal_nan=True)
 
-    pm_abs_threshold = pm.threshold(10, absolute=True)
-    assert np.nanmin(pm_abs_threshold.values) == -50
-
-    non_nan_values = pm_abs_threshold.values[~np.isnan(pm_abs_threshold.values)]
-    assert np.array_equal(non_nan_values, np.concatenate([np.arange(-50, -9), np.arange(10, 50)]))
-
-    pm_lower_threshold = pm.threshold(10, comparator=np.less)
-    assert np.nanmax(pm_lower_threshold.values) == 9
-    non_nan_values = pm_lower_threshold.values[~np.isnan(pm_lower_threshold.values)]
-    assert np.array_equal(non_nan_values, np.arange(-50, 10))
+    mask = pm.create_mask_from_threshold(1, comparator=np.less)
+    assert np.array_equal(mask.mask, [[1.0, 1.0, 1.0, np.nan, np.nan]], equal_nan=True)
 
 
 def test_convert():