mask.py

import numpy as np
import numpy.typing as npt
from typing import Optional


def spherical_mask(
    box_size: int,
    radius: float,
    smooth: Optional[float] = None,
    cutoff_sd: int = 3,
    center: Optional[float] = None,
) -> npt.NDArray[float]:
    """Wrapper around ellipsoidal_mask() to create a spherical mask with just a single
    radius.

    Parameters
    ----------
    box_size: int
        box size of the mask, equal in each dimension
    radius: float
        radius of sphere
    smooth: Optional[float], default None
        sigma (float relative to number of pixels) of gaussian falloff around mask
    cutoff_sd: int, default 3
        how many standard deviations of the Gaussian falloff to include, default of 3 is
        a good choice
    center: Optional[float], default None
        alternative center for the mask, default is (size - 1) / 2

    Returns
    -------
    mask: npt.NDArray[float]
        3D numpy array with mask at center
    """
    return ellipsoidal_mask(
        box_size, radius, radius, radius, smooth, cutoff_sd=cutoff_sd, center=center
    )


def ellipsoidal_mask(
    box_size: int,
    major: float,
    minor1: float,
    minor2: float,
    smooth: Optional[float] = None,
    cutoff_sd: int = 3,
    center: Optional[float] = None,
) -> npt.NDArray[float]:
    """Create an ellipsoidal mask in the specified square box. Ellipsoid is defined by 3
    radius on x,y, and z axis.

    Parameters
    ----------
    box_size: int
        box size of the mask, equal in each dimension
    major: float
        radius of ellipsoid in x
    minor1: float
        radius of ellipsoid in y
    minor2: float
        radius of ellipsoid in z
    smooth: Optional[float], default None
        sigma (float relative to number of pixels) of gaussian falloff around mask
    cutoff_sd: int, default 3
        how many standard deviations of the Gaussian falloff to include, default of 3 is
        a good choice
    center: Optional[float], default None
        alternative center for the mask, default is (size - 1) / 2

    Returns
    -------
    mask: npt.NDArray[float]
        3D numpy array with mask at center
    """
    if not all([box_size > 0, major > 0, minor1 > 0, minor2 > 0]):
        raise ValueError("Invalid input for mask creation: box_size or radii are <= 0")

    center = (box_size - 1) / 2 if center is None else center
    x, y, z = (
        np.arange(box_size) - center,
        np.arange(box_size) - center,
        np.arange(box_size) - center,
    )

    # use broadcasting
    r = np.sqrt(
        ((x / major) ** 2)[:, np.newaxis, np.newaxis]
        + ((y / minor1) ** 2)[:, np.newaxis]
        + (z / minor2) ** 2
    ).astype(np.float32)

    if smooth is not None:
        if not all([smooth >= 0, cutoff_sd >= 0]):
            raise ValueError(
                "Invalid input for mask smoothing: smooth or sd cutoff are <= 0"
            )

        r[r <= 1] = 1
        sigma = smooth / ((major + minor1 + minor2) / 3)
        mask = np.exp(-1 * ((r - 1) / sigma) ** 2)
        mask[mask <= np.exp(-(cutoff_sd**2) / 2.0)] = 0

    else:
        mask = np.zeros_like(r)
        mask[r <= 1] = 1.0

    return mask