waveorder 2.2.1b0__py3-none-any.whl → 3.0.0__py3-none-any.whl

waveorder/focus.py

@@ -3,11 +3,53 @@ from typing import Literal, Optional
 
 import matplotlib.pyplot as plt
 import numpy as np
+import torch
 from scipy.signal import peak_widths
 
 from waveorder import util
 
 
+def compute_midband_power(
+    yx_array: torch.Tensor,
+    NA_det: float,
+    lambda_ill: float,
+    pixel_size: float,
+    midband_fractions: tuple[float, float] = (0.125, 0.25),
+) -> torch.Tensor:
+    """Compute midband spatial frequency power by summing over a 2D midband donut.
+
+    Parameters
+    ----------
+    yx_array : torch.Tensor
+        2D tensor in (Y, X) order.
+    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.
+        Default is (0.125, 0.25).
+
+    Returns
+    -------
+    torch.Tensor
+        Sum of absolute FFT values in the midband region.
+    """
+    _, _, fxx, fyy = util.gen_coordinate(yx_array.shape, pixel_size)
+    frr = torch.tensor(np.sqrt(fxx**2 + fyy**2))
+    xy_abs_fft = torch.abs(torch.fft.fftn(yx_array))
+    cutoff = 2 * NA_det / lambda_ill
+    mask = torch.logical_and(
+        frr > cutoff * midband_fractions[0],
+        frr < cutoff * midband_fractions[1],
+    )
+    return torch.sum(xy_abs_fft[mask])
+
+
 def focus_from_transverse_band(
     zyx_array,
     NA_det,
@@ -15,8 +57,11 @@ def focus_from_transverse_band(
     pixel_size,
     midband_fractions=(0.125, 0.25),
     mode: Literal["min", "max"] = "max",
+    polynomial_fit_order: Optional[int] = None,
     plot_path: Optional[str] = None,
     threshold_FWHM: float = 0,
+    return_statistics: bool = False,
+    enable_subpixel_precision: bool = False,
 ):
     """Estimates the in-focus slice from a 3D stack by optimizing a transverse spatial frequency band.
 
@@ -36,8 +81,10 @@ def focus_from_transverse_band(
     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.
+    mode: {'min', 'max'}, optional
+        Option to choose the in-focus slice by minimizing or maximizing the midband power. By default 'max'.
+    polynomial_fit_order: int, optional
+        Default None is no fit. If integer, a polynomial of that degree is fit to the midband power before choosing the extreme point as the in-focus slice.
     plot_path: str or None, optional
         File name for a diagnostic plot (supports matplotlib filetypes .png, .pdf, .svg, etc.).
         Use None to skip.
@@ -46,14 +93,22 @@ def focus_from_transverse_band(
         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.
+    return_statistics: bool, optional
+        If True, returns a tuple (in_focus_index, peak_stats) instead of just in_focus_index.
+        Default is False for backward compatibility.
+    enable_subpixel_precision: bool, optional
+        If True and polynomial_fit_order is provided, enables sub-pixel precision focus detection
+        by finding the continuous extremum of the polynomial fit. Default is False for backward compatibility.
 
     Returns
-    ------
-    slice : int or None
-        If peak's FWHM > peak_width_threshold:
-            return the index of the in-focus slice
-        else:
-            return None
+    -------
+    slice : int, float, None, or tuple
+        If return_statistics is False (default):
+            Returns in_focus_index (int if enable_subpixel_precision=False,
+            float if enable_subpixel_precision=True and polynomial_fit_order is not None, or None).
+        If return_statistics is True:
+            Returns tuple (in_focus_index, peak_stats) where peak_stats is a dict
+            containing 'peak_index' and 'peak_FWHM'.
 
     Example
     ------
@@ -64,6 +119,7 @@ def focus_from_transverse_band(
     >>> in_focus_data = data[slice,:,:]
     """
     minmaxfunc = _mode_to_minmaxfunc(mode)
+    peak_stats = {"peak_index": None, "peak_FWHM": None}
 
     _check_focus_inputs(
         zyx_array, NA_det, lambda_ill, pixel_size, midband_fractions
@@ -74,29 +130,71 @@ def focus_from_transverse_band(
         warnings.warn(
             "The dataset only contained a single slice. Returning trivial slice index = 0."
         )
+        if return_statistics:
+            return 0, peak_stats
         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 midband power for each slice
+    midband_sum = np.array(
+        [
+            compute_midband_power(
+                torch.from_numpy(zyx_array[z]),
+                NA_det,
+                lambda_ill,
+                pixel_size,
+                midband_fractions,
+            ).numpy()
+            for z in range(zyx_array.shape[0])
+        ]
+    )
 
-    # Calculate fft
-    xy_abs_fft = np.abs(np.fft.fftn(zyx_array, axes=(1, 2)))
+    if polynomial_fit_order is None:
+        peak_index = minmaxfunc(midband_sum)
+    else:
+        x = np.arange(len(midband_sum))
+        coeffs = np.polyfit(x, midband_sum, polynomial_fit_order)
+        poly_func = np.poly1d(coeffs)
 
-    # Calculate midband mask
-    cutoff = 2 * NA_det / lambda_ill
-    midband_mask = np.logical_and(
-        frr > cutoff * midband_fractions[0],
-        frr < cutoff * midband_fractions[1],
-    )
+        if enable_subpixel_precision:
+            # Find the continuous extremum using derivative
+            poly_deriv = np.polyder(coeffs)
+            # Find roots of the derivative (critical points)
+            critical_points = np.roots(poly_deriv)
+
+            # Filter for real roots within the data range
+            real_critical_points = []
+            for cp in critical_points:
+                if np.isreal(cp) and 0 <= cp.real < len(midband_sum):
+                    real_critical_points.append(cp.real)
+
+            if real_critical_points:
+                # Evaluate the polynomial at critical points to find extremum
+                critical_values = [
+                    poly_func(cp) for cp in real_critical_points
+                ]
+                if mode == "max":
+                    best_idx = np.argmax(critical_values)
+                else:  # mode == "min"
+                    best_idx = np.argmin(critical_values)
+                peak_index = real_critical_points[best_idx]
+            else:
+                # Fall back to discrete maximum if no valid critical points
+                peak_index = float(minmaxfunc(poly_func(x)))
+        else:
+            peak_index = minmaxfunc(poly_func(x))
 
-    # 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)
+    # For peak width calculation, use integer peak index
+    if enable_subpixel_precision and polynomial_fit_order is not None:
+        # Use the closest integer index for peak width calculation
+        integer_peak_index = int(np.round(peak_index))
+    else:
+        integer_peak_index = int(peak_index)
 
-    peak_results = peak_widths(midband_sum, [peak_index])
+    peak_results = peak_widths(midband_sum, [integer_peak_index])
     peak_FWHM = peak_results[0][0]
+    peak_stats.update(
+        {"peak_index": int(peak_index), "peak_FWHM": float(peak_FWHM)}
+    )
 
     if peak_FWHM >= threshold_FWHM:
         in_focus_index = peak_index
@@ -114,6 +212,9 @@ def focus_from_transverse_band(
             threshold_FWHM,
         )
 
+    if return_statistics:
+        return in_focus_index, peak_stats
+
     return in_focus_index
 
 
@@ -168,9 +269,19 @@ def _plot_focus_metric(
 ):
     _, ax = plt.subplots(1, 1, figsize=(4, 4))
     ax.plot(midband_sum, "-k")
+
+    # Handle floating-point peak_index for plotting
+    if isinstance(peak_index, float) and not peak_index.is_integer():
+        # Use interpolation to get the y-value at the floating-point x-position
+        peak_y_value = np.interp(
+            peak_index, np.arange(len(midband_sum)), midband_sum
+        )
+    else:
+        peak_y_value = midband_sum[int(peak_index)]
+
     ax.plot(
         peak_index,
-        midband_sum[peak_index],
+        peak_y_value,
         "go" if in_focus_index is not None else "ro",
     )
     ax.hlines(*peak_results[1:], color="k", linestyles="dashed")

waveorder/io/_reader.py

@@ -0,0 +1,61 @@
+from typing import Dict, List, Tuple, Union
+
+import zarr
+from iohub import read_micromanager
+from napari_ome_zarr._reader import napari_get_reader as fallback_reader
+
+
+def napari_get_reader(path):
+    if isinstance(path, str):
+        if ".zarr" in path:
+            with zarr.open(path) as root:
+                if "plate" in root.attrs:
+                    return hcs_zarr_reader
+                else:
+                    return fallback_reader(path)
+        else:
+            return ome_tif_reader
+    else:
+        return None
+
+
+def hcs_zarr_reader(
+    path: Union[str, List[str]],
+) -> List[Tuple[zarr.Array, Dict]]:
+    reader = read_micromanager(path)
+    results = list()
+
+    zs = zarr.open(path, "r")
+    names = []
+
+    dict_ = zs.attrs.asdict()
+    wells = dict_["plate"]["wells"]
+    for well in wells:
+        path = well["path"]
+        well_dict = zs[path].attrs.asdict()
+        for name in well_dict["well"]["images"]:
+            names.append(name["path"])
+    for pos in range(reader.get_num_positions()):
+        meta = dict()
+        name = names[pos]
+        meta["name"] = name
+        results.append((reader.get_zarr(pos), meta))
+    return results
+
+
+def ome_tif_reader(
+    path: Union[str, List[str]],
+) -> List[Tuple[zarr.Array, Dict]]:
+    reader = read_micromanager(path)
+    results = list()
+
+    npos = reader.get_num_positions()
+    for pos in range(npos):
+        meta = dict()
+        if npos == 1:
+            meta["name"] = "Pos000_000"
+        else:
+            meta["name"] = reader.stage_positions[pos]["Label"][2:]
+        results.append((reader.get_zarr(pos), meta))
+
+    return results

waveorder/io/core_functions.py

@@ -0,0 +1,272 @@
+import time
+from contextlib import contextmanager
+
+import numpy as np
+
+
+@contextmanager
+def suspend_live_sm(snap_manager):
+    """Context manager that suspends/unsuspends MM live mode for `SnapLiveManager`.
+
+    Parameters
+    ----------
+    snap_manager : object
+        `org.micromanager.internal.SnapLiveManager` object via pycromanager
+
+    Yields
+    ------
+    object
+        `org.micromanager.internal.SnapLiveManager` object via pycromanager
+
+    Usage
+    -----
+    ```py
+    with suspend_live_sm(snap_manager) as sm:
+        pass # do something with MM that can't be done in live mode
+    ```
+    """
+    snap_manager.setSuspended(True)
+    try:
+        yield snap_manager
+    finally:
+        snap_manager.setSuspended(False)
+
+
+def snap_and_get_image(snap_manager):
+    """
+    Snap and get image using Snap Live Window Manager + transfer of ZMQ
+
+    Parameters
+    ----------
+    snap_manager:   (object) MM Snap Live Window object
+
+    Returns
+    -------
+    image:          (array) 2D array of size (Y, X)
+
+    """
+    snap_manager.snap(True)
+    time.sleep(
+        0.3
+    )  # sleep after snap to make sure the image we grab is the correct one
+
+    # get pixels + dimensions
+    height = snap_manager.getDisplay().getDisplayedImages().get(0).getHeight()
+    width = snap_manager.getDisplay().getDisplayedImages().get(0).getWidth()
+    array = (
+        snap_manager.getDisplay().getDisplayedImages().get(0).getRawPixels()
+    )
+
+    return np.reshape(array, (height, width))
+
+
+def snap_and_average(snap_manager, display=True):
+    """
+    Snap an image with Snap Live manager + grab only the mean (computed in java)
+
+    Parameters
+    ----------
+    snap_manager:   (object) MM Snap Live Window object
+    display:        (bool) Whether to show the snap on the Snap Live Window in MM
+
+    Returns
+    -------
+    mean:           (float) mean of snapped image
+
+    """
+
+    snap_manager.snap(display)
+    time.sleep(
+        0.3
+    )  # sleep after snap to make sure the image we grab is the correct one
+
+    return snap_manager.getDisplay().getImagePlus().getStatistics().umean
+
+
+def set_lc_waves(mmc, device_property: tuple, value: float):
+    """
+    Set retardance in waves for LC in device_property
+
+    Parameters
+    ----------
+    mmc : object
+        MM Core object
+    device_property : tuple
+        (device_name, property_name) set
+    value : float
+        Retardance to set as fraction of a wavelength
+
+    Returns
+    -------
+
+    """
+    device_name = device_property[0]
+    prop_name = device_property[1]
+
+    if value > 1.6 or value < 0.001:
+        raise ValueError(
+            f"Requested retardance value is {value} waves. "
+            f"Retardance must be greater than 0.001 and less than 1.6 waves."
+        )
+
+    mmc.setProperty(device_name, prop_name, str(value))
+    time.sleep(20 / 1000)
+
+
+def set_lc_voltage(mmc, device_property: tuple, value: float):
+    """
+    Set LC retardance by specifying LC voltage
+
+    Parameters
+    ----------
+    mmc : object
+        MM Core object
+    device_property : tuple
+        (device_name, property_name) set
+    value : float
+        LC voltage in volts. Applied voltage is limited to 20V
+
+    Returns
+    -------
+
+    """
+    device_name = device_property[0]
+    prop_name = device_property[1]
+
+    if value > 20.0 or value < 0.0:
+        raise ValueError(
+            f"Requested LC voltage is {value} V. "
+            f"LC voltage must be greater than 0.0 and less than 20.0 V."
+        )
+
+    mmc.setProperty(device_name, prop_name, str(value))
+    time.sleep(20 / 1000)
+
+
+def set_lc_daq(mmc, device_property: tuple, value: float):
+    """
+    Set LC retardance based on DAQ output
+
+    Parameters
+    ----------
+    mmc : object
+        MM Core object
+    device_property : tuple
+        (device_name, property_name) set
+    value : float
+        DAQ output voltage in volts. DAQ output must be in 0-5V range
+
+    Returns
+    -------
+
+    """
+    device_name = device_property[0]
+    prop_name = device_property[1]
+
+    if value > 5.0 or value < 0.0:
+        raise ValueError(
+            "DAC voltage must be greater than 0.0 and less than 5.0"
+        )
+
+    mmc.setProperty(device_name, prop_name, str(value))
+    time.sleep(20 / 1000)
+
+
+def get_lc(mmc, device_property: tuple):
+    """
+    Get LC state in the native units of the device property
+
+    Parameters
+    ----------
+    mmc : object
+        MM Core object
+    device_property : tuple
+        (device_name, property_name) set
+
+    Returns
+    -------
+
+    """
+
+    device_name = device_property[0]
+    prop_name = device_property[1]
+
+    val = float(mmc.getProperty(device_name, prop_name))
+    return val
+
+
+def define_meadowlark_state(mmc, device_property: tuple):
+    """
+    Defines pallet element in the Meadowlark device adapter for the given state.
+    Make sure LC values for this state are set before calling this function
+
+    Parameters
+    ----------
+    mmc : object
+        MM Core object
+    device_property : tuple
+        (device_name, property_name) set, e.g.
+        ('MeadowlarkLC', 'Pal. elem. 00; enter 0 to define; 1 to activate')
+
+    Returns
+    -------
+
+    """
+
+    device_name = device_property[0]
+    prop_name = device_property[1]
+
+    # define LC state
+    # setting pallet elements to 0 defines LC state
+    mmc.setProperty(device_name, prop_name, 0)
+    mmc.waitForDevice(device_name)
+
+
+def define_config_state(
+    mmc, group: str, config: str, device_properties: list, values: list
+):
+    """
+    Define config state by specifying the values for all device properties in this config
+
+    Parameters
+    ----------
+    mmc : object
+        MM Core object
+    group : str
+        Name of config group
+    config : str
+        Name of config, e.g. State0
+    device_properties: list
+        List of (device_name, property_name) tuples in config
+    values: list
+        List of matching device property values
+
+    Returns
+    -------
+
+    """
+
+    for device_property, value in zip(device_properties, values):
+        device_name = device_property[0]
+        prop_name = device_property[1]
+        mmc.defineConfig(group, config, device_name, prop_name, str(value))
+    mmc.waitForConfig(group, config)
+
+
+def set_lc_state(mmc, group: str, config: str):
+    """
+    Change to the specific LC State
+
+    Parameters
+    ----------
+    mmc : object
+        MM Core object
+    group : str
+        Name of config group
+    config : str
+        Name of config, e.g. State0
+
+    """
+
+    mmc.setConfig(group, config)
+    time.sleep(20 / 1000)  # delay for LC settle time