waveorder 0.2.2rc0__py3-none-any.whl

waveorder/_version.py

@@ -0,0 +1,16 @@
+# file generated by setuptools_scm
+# don't change, don't track in version control
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple, Union
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.2.2rc0'
+__version_tuple__ = version_tuple = (0, 2, 2)

waveorder/background_estimator.py

@@ -0,0 +1,319 @@
+"""Estimate flat field images"""
+
+import numpy as np
+import itertools
+
+
+"""
+
+This script is adopted from 
+
+https://github.com/mehta-lab/reconstruct-order
+
+
+"""
+
+
+class BackgroundEstimator2D:
+    """Estimates flat field image"""
+
+    def __init__(self, block_size=32):
+        """
+        Background images are estimated once per channel for 2D data
+        :param int block_size: Size of blocks image will be divided into
+        """
+
+        if block_size is None:
+            block_size = 32
+        self.block_size = block_size
+
+    def sample_block_medians(self, im):
+        """Subdivide a 2D image in smaller blocks of size block_size and
+        compute the median intensity value for each block. Any incomplete
+        blocks (remainders of modulo operation) will be ignored.
+
+        :param np.array im:         2D image
+        :return np.array(float) sample_coords: Image coordinates for block
+                                               centers
+        :return np.array(float) sample_values: Median intensity values for
+                                               blocks
+        """
+
+        im_shape = im.shape
+        assert (
+            self.block_size < im_shape[0]
+        ), "Block size larger than image height"
+        assert (
+            self.block_size < im_shape[1]
+        ), "Block size larger than image width"
+
+        nbr_blocks_x = im_shape[0] // self.block_size
+        nbr_blocks_y = im_shape[1] // self.block_size
+        sample_coords = np.zeros(
+            (nbr_blocks_x * nbr_blocks_y, 2), dtype=np.float64
+        )
+        sample_values = np.zeros(
+            (nbr_blocks_x * nbr_blocks_y,), dtype=np.float64
+        )
+        for x in range(nbr_blocks_x):
+            for y in range(nbr_blocks_y):
+                idx = y * nbr_blocks_x + x
+                sample_coords[idx, :] = [
+                    x * self.block_size + (self.block_size - 1) / 2,
+                    y * self.block_size + (self.block_size - 1) / 2,
+                ]
+                sample_values[idx] = np.median(
+                    im[
+                        x * self.block_size : (x + 1) * self.block_size,
+                        y * self.block_size : (y + 1) * self.block_size,
+                    ]
+                )
+        return sample_coords, sample_values
+
+    @staticmethod
+    def fit_polynomial_surface_2D(
+        sample_coords, sample_values, im_shape, order=2, normalize=True
+    ):
+        """
+        Given coordinates and corresponding values, this function will fit a
+        2D polynomial of given order, then create a surface of given shape.
+
+        :param np.array sample_coords: 2D sample coords (nbr of points, 2)
+        :param np.array sample_values: Corresponding intensity values (nbr points,)
+        :param tuple im_shape:         Shape of desired output surface (height, width)
+        :param int order:              Order of polynomial (default 2)
+        :param bool normalize:         Normalize surface by dividing by its mean
+                                       for background correction (default True)
+
+        :return np.array poly_surface: 2D surface of shape im_shape
+        """
+        assert (order + 1) * (order + 2) / 2 <= len(
+            sample_values
+        ), "Can't fit a higher degree polynomial than there are sampled values"
+        # Number of coefficients is determined by (order + 1)*(order + 2)/2
+        orders = np.arange(order + 1)
+        variable_matrix = np.zeros(
+            (sample_coords.shape[0], int((order + 1) * (order + 2) / 2))
+        )
+        order_pairs = list(itertools.product(orders, orders))
+        # sum of orders of x,y <= order of the polynomial
+        variable_iterator = itertools.filterfalse(
+            lambda x: sum(x) > order, order_pairs
+        )
+        for idx, (m, n) in enumerate(variable_iterator):
+            variable_matrix[:, idx] = (
+                sample_coords[:, 0] ** n * sample_coords[:, 1] ** m
+            )
+        # Least squares fit of the points to the polynomial
+        coeffs, _, _, _ = np.linalg.lstsq(
+            variable_matrix, sample_values, rcond=-1
+        )
+        # Create a grid of image (x, y) coordinates
+        x_mesh, y_mesh = np.meshgrid(
+            np.linspace(0, im_shape[1] - 1, im_shape[1]),
+            np.linspace(0, im_shape[0] - 1, im_shape[0]),
+        )
+        # Reconstruct the surface from the coefficients
+        poly_surface = np.zeros(im_shape, np.float64)
+        order_pairs = list(itertools.product(orders, orders))
+        # sum of orders of x,y <= order of the polynomial
+        variable_iterator = itertools.filterfalse(
+            lambda x: sum(x) > order, order_pairs
+        )
+        for coeff, (m, n) in zip(coeffs, variable_iterator):
+            poly_surface += coeff * x_mesh**m * y_mesh**n
+
+        if normalize:
+            poly_surface /= np.mean(poly_surface)
+        return poly_surface
+
+    def get_background(self, im, order=2, normalize=True):
+        """
+        Combine sampling and polynomial surface fit for background estimation.
+        To background correct an image, divide it by background.
+
+        :param np.array im:        2D image
+        :param int order:          Order of polynomial (default 2)
+        :param bool normalize:     Normalize surface by dividing by its mean
+                                   for background correction (default True)
+
+        :return np.array background:    Background image
+        """
+
+        coords, values = self.sample_block_medians(im=im)
+        background = self.fit_polynomial_surface_2D(
+            sample_coords=coords,
+            sample_values=values,
+            im_shape=im.shape,
+            order=order,
+            normalize=normalize,
+        )
+        # Backgrounds can't contain zeros or negative values
+        # if background.min() <= 0:
+        #     raise ValueError(
+        #         "The generated background was not strictly positive {}.".format(
+        #             background.min()),
+        #     )
+        return background
+
+
+class BackgroundEstimator2D_GPU:
+    """Estimates flat field image"""
+
+    def __init__(self, block_size=32, gpu_id=0):
+        """
+        Background images are estimated once per channel for 2D data
+        :param int block_size: Size of blocks image will be divided into
+        """
+        globals()["cp"] = __import__("cupy")
+        self.gpu_id = gpu_id
+        cp.cuda.Device(self.gpu_id).use()
+
+        if block_size is None:
+            block_size = 32
+        self.block_size = block_size
+
+    def sample_block_medians(self, im):
+        """Subdivide a 2D image in smaller blocks of size block_size and
+        compute the median intensity value for each block. Any incomplete
+        blocks (remainders of modulo operation) will be ignored.
+
+        :param np.array im:         2D image
+        :return np.array(float) sample_coords: Image coordinates for block
+                                               centers
+        :return np.array(float) sample_values: Median intensity values for
+                                               blocks
+        """
+
+        im_shape = im.shape
+        assert (
+            self.block_size < im_shape[0]
+        ), "Block size larger than image height"
+        assert (
+            self.block_size < im_shape[1]
+        ), "Block size larger than image width"
+
+        nbr_blocks_x = im_shape[0] // self.block_size
+        nbr_blocks_y = im_shape[1] // self.block_size
+        sample_coords = np.zeros(
+            (nbr_blocks_x * nbr_blocks_y, 2), dtype=cp.float64
+        )
+        sample_values = np.zeros(
+            (nbr_blocks_x * nbr_blocks_y,), dtype=cp.float64
+        )
+        for x in range(nbr_blocks_x):
+            for y in range(nbr_blocks_y):
+                idx = y * nbr_blocks_x + x
+                sample_coords[idx, :] = [
+                    x * self.block_size + (self.block_size - 1) / 2,
+                    y * self.block_size + (self.block_size - 1) / 2,
+                ]
+                sample_values[idx] = np.median(
+                    im[
+                        x * self.block_size : (x + 1) * self.block_size,
+                        y * self.block_size : (y + 1) * self.block_size,
+                    ]
+                )
+        return sample_coords, sample_values
+
+    @staticmethod
+    def median_cp(x):
+        x = x.flatten()
+        n = x.shape[0]
+        s = cp.sort(x)
+        m_odd = cp.take(s, n // 2)
+        if n % 2 == 1:
+            return m_odd
+        else:
+            m_even = cp.take(s, n // 2 - 1)
+            return (m_odd + m_even) / 2
+
+    @staticmethod
+    def fit_polynomial_surface_2D(
+        sample_coords, sample_values, im_shape, order=2, normalize=True
+    ):
+        """
+        Given coordinates and corresponding values, this function will fit a
+        2D polynomial of given order, then create a surface of given shape.
+
+        :param np.array sample_coords: 2D sample coords (nbr of points, 2)
+        :param np.array sample_values: Corresponding intensity values (nbr points,)
+        :param tuple im_shape:         Shape of desired output surface (height, width)
+        :param int order:              Order of polynomial (default 2)
+        :param bool normalize:         Normalize surface by dividing by its mean
+                                       for background correction (default True)
+
+        :return np.array poly_surface: 2D surface of shape im_shape
+        """
+        assert (order + 1) * (order + 2) / 2 <= len(
+            sample_values
+        ), "Can't fit a higher degree polynomial than there are sampled values"
+        # Number of coefficients is determined by (order + 1)*(order + 2)/2
+        orders = np.arange(order + 1)
+        variable_matrix = np.zeros(
+            (sample_coords.shape[0], int((order + 1) * (order + 2) / 2))
+        )
+        order_pairs = list(itertools.product(orders, orders))
+        # sum of orders of x,y <= order of the polynomial
+        variable_iterator = itertools.filterfalse(
+            lambda x: sum(x) > order, order_pairs
+        )
+        for idx, (m, n) in enumerate(variable_iterator):
+            variable_matrix[:, idx] = (
+                sample_coords[:, 0] ** n * sample_coords[:, 1] ** m
+            )
+        # Least squares fit of the points to the polynomial
+        coeffs, _, _, _ = np.linalg.lstsq(
+            variable_matrix, sample_values, rcond=-1
+        )
+
+        # Create a grid of image (x, y) coordinates
+        x_mesh, y_mesh = cp.meshgrid(
+            cp.linspace(0, im_shape[1] - 1, im_shape[1]),
+            cp.linspace(0, im_shape[0] - 1, im_shape[0]),
+        )
+        # Reconstruct the surface from the coefficients
+        poly_surface = cp.zeros(im_shape, cp.float)
+        coeffs = cp.array(coeffs)
+        order_pairs = list(itertools.product(orders, orders))
+        # sum of orders of x,y <= order of the polynomial
+        variable_iterator = itertools.filterfalse(
+            lambda x: sum(x) > order, order_pairs
+        )
+        for coeff, (m, n) in zip(coeffs, variable_iterator):
+            poly_surface += coeff * x_mesh**m * y_mesh**n
+
+        if normalize:
+            poly_surface /= cp.mean(poly_surface)
+        return poly_surface
+
+    def get_background(self, im, order=2, normalize=True):
+        """
+        Combine sampling and polynomial surface fit for background estimation.
+        To background correct an image, divide it by background.
+
+        :param np.array im:        2D image
+        :param int order:          Order of polynomial (default 2)
+        :param bool normalize:     Normalize surface by dividing by its mean
+                                   for background correction (default True)
+
+        :return np.array background:    Background image
+        """
+        cp.cuda.Device(self.gpu_id).use()
+        im = cp.asnumpy(im)
+
+        coords, values = self.sample_block_medians(im=im)
+        background = self.fit_polynomial_surface_2D(
+            sample_coords=coords,
+            sample_values=values,
+            im_shape=im.shape,
+            order=order,
+            normalize=normalize,
+        )
+        # Backgrounds can't contain zeros or negative values
+        # if background.min() <= 0:
+        #     raise ValueError(
+        #         "The generated background was not strictly positive {}.".format(
+        #             background.min()),
+        #     )
+        return background

waveorder/correction.py

@@ -0,0 +1,107 @@
+"""Background correction methods"""
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, Size
+
+
+def _sample_block_medians(image: Tensor, block_size) -> Tensor:
+    """
+    Sample densely tiled square blocks from a 2D image and return their medians.
+    Incomplete blocks (overhangs) will be ignored.
+
+    Parameters
+    ----------
+    image : Tensor
+        2D image
+    block_size : int, optional
+        Width and height of the blocks
+
+    Returns
+    -------
+    Tensor
+        Median intensity values for each block, flattened
+    """
+    if not image.dtype.is_floating_point:
+        image.to(torch.float)
+    blocks = F.unfold(image[None, None], block_size, stride=block_size)[0]
+    return blocks.median(0)[0]
+
+
+def _grid_coordinates(image: Tensor, block_size: int) -> Tensor:
+    """Build image coordinates from the center points of square blocks"""
+    coords = torch.meshgrid(
+        [
+            torch.arange(
+                0 + block_size / 2,
+                boundary - block_size / 2 + 1,
+                block_size,
+                device=image.device,
+            )
+            for boundary in image.shape
+        ]
+    )
+    return torch.stack(coords, dim=-1).reshape(-1, 2)
+
+
+def _fit_2d_polynomial_surface(
+    coords: Tensor, values: Tensor, order: int, surface_shape: Size
+) -> Tensor:
+    """Fit a 2D polynomial to a set of coordinates and their values,
+    and return the surface evaluated at every point."""
+    n_coeffs = int((order + 1) * (order + 2) / 2)
+    if n_coeffs >= len(values):
+        raise ValueError(
+            f"Cannot fit a {order} degree 2D polynomial "
+            f"with {len(values)} sampled values"
+        )
+    orders = torch.arange(order + 1, device=coords.device)
+    order_pairs = torch.stack(torch.meshgrid(orders, orders), -1)
+    order_pairs = order_pairs[order_pairs.sum(-1) <= order].reshape(-1, 2)
+    terms = torch.stack(
+        [coords[:, 0] ** i * coords[:, 1] ** j for i, j in order_pairs], -1
+    )
+    # use "gels" driver for precision and GPU consistency
+    coeffs = torch.linalg.lstsq(terms, values, driver="gels").solution
+    dense_coords = torch.meshgrid(
+        [
+            torch.arange(s, dtype=values.dtype, device=values.device)
+            for s in surface_shape
+        ]
+    )
+    dense_terms = torch.stack(
+        [dense_coords[0] ** i * dense_coords[1] ** j for i, j in order_pairs],
+        -1,
+    )
+    return torch.matmul(dense_terms, coeffs)
+
+
+def estimate_background(image: Tensor, order: int = 2, block_size: int = 32):
+    """
+    Combine sampling and polynomial surface fit for background estimation.
+    To background correct an image, divide it by the background.
+
+    Parameters
+    ----------
+    image : Tensor
+        2D image
+    order : int, optional
+        Order of polynomial, by default 2
+    block_size : int, optional
+        Width and height of the blocks, by default 32
+
+    Returns
+    -------
+    Tensor
+        Background image
+    """
+    if image.ndim != 2:
+        raise ValueError(f"Image must be 2D, got shape {image.shape}")
+    height, width = image.shape
+    if block_size > width:
+        raise ValueError("Block size larger than image height")
+    if block_size > height:
+        raise ValueError("Block size larger than image width")
+    medians = _sample_block_medians(image, block_size)
+    coords = _grid_coordinates(image, block_size)
+    return _fit_2d_polynomial_surface(coords, medians, order, image.shape)

waveorder/focus.py

@@ -0,0 +1,198 @@
+from scipy.signal import peak_widths
+from typing import Literal, Optional
+from waveorder import util
+import matplotlib.pyplot as plt
+import numpy as np
+import warnings
+
+
+def focus_from_transverse_band(
+    zyx_array,
+    NA_det,
+    lambda_ill,
+    pixel_size,
+    midband_fractions=(0.125, 0.25),
+    mode: Literal["min", "max"] = "max",
+    plot_path: Optional[str] = None,
+    threshold_FWHM: float = 0,
+):
+    """Estimates the in-focus slice from a 3D stack by optimizing a transverse spatial frequency band.
+
+    Parameters
+    ----------
+    zyx_array: np.array
+        Data stack in (Z, Y, X) order.
+        Requires len(zyx_array.shape) == 3.
+    NA_det: float
+        Detection NA.
+    lambda_ill: float
+        Illumination wavelength
+        Units are arbitrary, but must match [pixel_size]
+    pixel_size: float
+        Object-space pixel size = camera pixel size / magnification.
+        Units are arbitrary, but must match [lambda_ill]
+    midband_fractions: Tuple[float, float], optional
+        The minimum and maximum fraction of the cutoff frequency that define the midband.
+        Requires: 0 <= midband_fractions[0] < midband_fractions[1] <= 1.
+    mode: {'max', 'min'}, optional
+        Option to choose the in-focus slice by minimizing or maximizing the midband frequency.
+    plot_path: str or None, optional
+        File name for a diagnostic plot (supports matplotlib filetypes .png, .pdf, .svg, etc.).
+        Use None to skip.
+    threshold_FWHM: float, optional
+        Threshold full-width half max for a peak to be considered in focus.
+        The default value, 0, applies no threshold, and the maximum midband power is always considered in focus.
+        For values > 0, the peak's FWHM must be greater than the threshold for the slice to be considered in focus.
+        If the peak does not meet this threshold, the function returns None.
+
+    Returns
+    ------
+    slice : int or None
+        If peak's FWHM > peak_width_threshold:
+            return the index of the in-focus slice
+        else:
+            return None
+
+    Example
+    ------
+    >>> zyx_array.shape
+    (11, 2048, 2048)
+    >>> from waveorder.focus import focus_from_transverse_band
+    >>> slice = focus_from_transverse_band(zyx_array, NA_det=0.55, lambda_ill=0.532, pixel_size=6.5/20)
+    >>> in_focus_data = data[slice,:,:]
+    """
+    minmaxfunc = _mode_to_minmaxfunc(mode)
+
+    _check_focus_inputs(
+        zyx_array, NA_det, lambda_ill, pixel_size, midband_fractions
+    )
+
+    # Check for single slice
+    if zyx_array.shape[0] == 1:
+        warnings.warn(
+            "The dataset only contained a single slice. Returning trivial slice index = 0."
+        )
+        return 0
+
+    # Calculate coordinates
+    _, Y, X = zyx_array.shape
+    _, _, fxx, fyy = util.gen_coordinate((Y, X), pixel_size)
+    frr = np.sqrt(fxx**2 + fyy**2)
+
+    # Calculate fft
+    xy_abs_fft = np.abs(np.fft.fftn(zyx_array, axes=(1, 2)))
+
+    # Calculate midband mask
+    cutoff = 2 * NA_det / lambda_ill
+    midband_mask = np.logical_and(
+        frr > cutoff * midband_fractions[0],
+        frr < cutoff * midband_fractions[1],
+    )
+
+    # Find slice index with min/max power in midband
+    midband_sum = np.sum(xy_abs_fft[:, midband_mask], axis=1)
+    peak_index = minmaxfunc(midband_sum)
+
+    peak_results = peak_widths(midband_sum, [peak_index])
+    peak_FWHM = peak_results[0][0]
+
+    if peak_FWHM >= threshold_FWHM:
+        in_focus_index = peak_index
+    else:
+        in_focus_index = None
+
+    # Plot
+    if plot_path is not None:
+        _plot_focus_metric(
+            plot_path,
+            midband_sum,
+            peak_index,
+            in_focus_index,
+            peak_results,
+            threshold_FWHM,
+        )
+
+    return in_focus_index
+
+
+def _mode_to_minmaxfunc(mode):
+    if mode == "min":
+        minmaxfunc = np.argmin
+    elif mode == "max":
+        minmaxfunc = np.argmax
+    else:
+        raise ValueError("mode must be either `min` or `max`")
+    return minmaxfunc
+
+
+def _check_focus_inputs(
+    zyx_array, NA_det, lambda_ill, pixel_size, midband_fractions
+):
+    N = len(zyx_array.shape)
+    if N != 3:
+        raise ValueError(
+            f"{N}D array supplied. `focus_from_transverse_band` only accepts 3D arrays."
+        )
+
+    if NA_det < 0:
+        raise ValueError("NA must be > 0")
+    if lambda_ill < 0:
+        raise ValueError("lambda_ill must be > 0")
+    if pixel_size < 0:
+        raise ValueError("pixel_size must be > 0")
+    if not 0.4 < lambda_ill / pixel_size < 10:
+        warnings.warn(
+            f"WARNING: lambda_ill/pixel_size = {lambda_ill/pixel_size}."
+            f"Did you use the same units?"
+            f"Did you enter the pixel size in (demagnified) object-space units?"
+        )
+    if not midband_fractions[0] < midband_fractions[1]:
+        raise ValueError(
+            "midband_fractions[0] must be less than midband_fractions[1]"
+        )
+    if not (0 <= midband_fractions[0] <= 1):
+        raise ValueError("midband_fractions[0] must be between 0 and 1")
+    if not (0 <= midband_fractions[1] <= 1):
+        raise ValueError("midband_fractions[1] must be between 0 and 1")
+
+
+def _plot_focus_metric(
+    plot_path,
+    midband_sum,
+    peak_index,
+    in_focus_index,
+    peak_results,
+    threshold_FWHM,
+):
+    _, ax = plt.subplots(1, 1, figsize=(4, 4))
+    ax.plot(midband_sum, "-k")
+    ax.plot(
+        peak_index,
+        midband_sum[peak_index],
+        "go" if in_focus_index is not None else "ro",
+    )
+    ax.hlines(*peak_results[1:], color="k", linestyles="dashed")
+
+    ax.set_xlabel("Slice index")
+    ax.set_ylabel("Midband power")
+
+    ax.annotate(
+        f"In-focus slice = {in_focus_index}\n Peak width = {peak_results[0][0]:.2f}\n Peak width threshold = {threshold_FWHM}",
+        xy=(1, 1),
+        xytext=(1.0, 1.1),
+        textcoords="axes fraction",
+        xycoords="axes fraction",
+        ha="right",
+        va="center",
+        annotation_clip=False,
+    )
+
+    ax.spines["right"].set_visible(False)
+    ax.spines["top"].set_visible(False)
+    ax.spines["left"].set_position(("outward", 10))
+    ax.spines["bottom"].set_position(("outward", 10))
+    ax.ticklabel_format(style="sci", scilimits=(-2, 2))
+
+    print(f"Saving plot to {plot_path}")
+    plt.savefig(plot_path, bbox_inches="tight", dpi=300)
+    plt.close()