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

waveorder/io/metadata_reader.py

@@ -0,0 +1,195 @@
+import json
+import os
+
+from natsort import natsorted
+
+
+def load_json(path):
+    with open(path, "r") as f:
+        data = json.load(f)
+
+    return data
+
+
+def get_last_metadata_file(path):
+    last_metadata_file = natsorted(
+        [
+            file
+            for file in os.listdir(path)
+            if file.startswith("calibration_metadata")
+        ]
+    )[-1]
+    return os.path.join(path, last_metadata_file)
+
+
+class MetadataReader:
+    """
+    Calibration metadata reader class. Helps load metadata from different metadata formats and naming conventions
+    """
+
+    def __init__(self, path: str):
+        """
+
+        Parameters
+        ----------
+        path: full path to calibration metadata
+        """
+        self.metadata_path = path
+        self.json_metadata = load_json(self.metadata_path)
+
+        self.Timestamp = self.get_summary_calibration_attr("Timestamp")
+        self.waveorder_version = self.get_summary_calibration_attr(
+            "waveorder version"
+        )
+        self.Calibration_scheme = self.get_calibration_scheme()
+        self.Swing = self.get_swing()
+        self.Wavelength = self.get_summary_calibration_attr("Wavelength (nm)")
+        self.Black_level = self.get_black_level()
+        self.Extinction_ratio = self.get_extinction_ratio()
+        self.Channel_names = self.get_channel_names()
+        self.LCA_retardance = self.get_lc_retardance("LCA")
+        self.LCB_retardance = self.get_lc_retardance("LCB")
+        self.LCA_voltage = self.get_lc_voltage("LCA")
+        self.LCB_voltage = self.get_lc_voltage("LCB")
+        self.Swing_measured = self.get_swing_measured()
+        self.Notes = self.get_notes()
+
+    def get_summary_calibration_attr(self, attr):
+        try:
+            val = self.json_metadata["Summary"][attr]
+        except KeyError:
+            try:
+                val = self.json_metadata["Calibration"][attr]
+            except KeyError:
+                val = None
+        return val
+
+    def get_cal_states(self):
+        if self.Calibration_scheme == "4-State":
+            states = ["ext", "0", "60", "120"]
+        elif self.Calibration_scheme == "5-State":
+            states = ["ext", "0", "45", "90", "135"]
+        return states
+
+    def get_lc_retardance(self, lc):
+        """
+
+        Parameters
+        ----------
+        lc: 'LCA' or 'LCB'
+
+        Returns
+        -------
+
+        """
+        states = self.get_cal_states()
+
+        val = None
+        try:
+            val = [
+                self.json_metadata["Calibration"]["LC retardance"][
+                    f"{lc}_{state}"
+                ]
+                for state in states
+            ]
+        except KeyError:
+            states[0] = "Ext"
+            if lc == "LCA":
+                val = [
+                    self.json_metadata["Summary"][
+                        f"[LCA_{state}, LCB_{state}]"
+                    ][0]
+                    for state in states
+                ]
+            elif lc == "LCB":
+                val = [
+                    self.json_metadata["Summary"][
+                        f"[LCA_{state}, LCB_{state}]"
+                    ][1]
+                    for state in states
+                ]
+
+        return val
+
+    def get_lc_voltage(self, lc):
+        """
+
+        Parameters
+        ----------
+        lc: 'LCA' or 'LCB'
+
+        Returns
+        -------
+
+        """
+        states = self.get_cal_states()
+
+        val = None
+        if "Calibration" in self.json_metadata:
+            lc_voltage = self.json_metadata["Calibration"]["LC voltage"]
+            if lc_voltage:
+                val = [
+                    self.json_metadata["Calibration"]["LC voltage"][
+                        f"{lc}_{state}"
+                    ]
+                    for state in states
+                ]
+
+        return val
+
+    def get_swing(self):
+        try:
+            val = self.json_metadata["Calibration"]["Swing (waves)"]
+        except KeyError:
+            val = self.json_metadata["Summary"]["Swing (fraction)"]
+        return val
+
+    def get_swing_measured(self):
+        states = self.get_cal_states()
+        try:
+            val = [
+                self.json_metadata["Calibration"][f"Swing_{state}"]
+                for state in states[1:]
+            ]
+        except KeyError:
+            val = [
+                self.json_metadata["Summary"][f"Swing{state}"]
+                for state in states[1:]
+            ]
+
+        return val
+
+    def get_calibration_scheme(self):
+        try:
+            val = self.json_metadata["Calibration"]["Calibration scheme"]
+        except KeyError:
+            val = self.json_metadata["Summary"]["Acquired Using"]
+        return val
+
+    def get_black_level(self):
+        try:
+            val = self.json_metadata["Calibration"]["Black level"]
+        except KeyError:
+            val = self.json_metadata["Summary"]["BlackLevel"]
+        return val
+
+    def get_extinction_ratio(self):
+        try:
+            val = self.json_metadata["Calibration"]["Extinction ratio"]
+        except KeyError:
+            val = self.json_metadata["Summary"]["Extinction Ratio"]
+        return val
+
+    def get_channel_names(self):
+        try:
+            val = self.json_metadata["Calibration"]["Channel names"]
+        except KeyError:
+            val = self.json_metadata["Summary"]["ChNames"]
+        return val
+
+    def get_notes(self):
+        try:
+            val = self.json_metadata["Notes"]
+        except KeyError:
+            val = None
+        return val

waveorder/io/utils.py

@@ -0,0 +1,175 @@
+import os
+import textwrap
+from pathlib import Path
+
+import psutil
+import torch
+import yaml
+from iohub import open_ome_zarr
+
+from waveorder.cli.settings import MyBaseModel
+
+
+def add_index_to_path(path: Path):
+    """Takes a path to a file or folder and appends the smallest index that does
+    not already exist in that folder.
+
+    For example:
+    './output.txt' -> './output_0.txt' if no other files named './output*.txt' exist.
+    './output.txt' -> './output_2.txt' if './output_0.txt' and './output_1.txt' already exist.
+
+    Parameters
+    ----------
+    path: Path
+        Base path to add index to
+
+    Returns
+    -------
+    Path
+    """
+    index = 0
+    new_stem = f"{path.stem}_{index}"
+
+    while (path.parent / (new_stem + path.suffix)).exists():
+        index += 1
+        new_stem = f"{path.stem}_{index}"
+
+    return path.parent / (new_stem + path.suffix)
+
+
+def load_background(background_path):
+    with open_ome_zarr(
+        os.path.join(background_path, "background.zarr", "0", "0", "0")
+    ) as dataset:
+        cyx_data = dataset["0"][0, :, 0]
+        return torch.tensor(cyx_data, dtype=torch.float32)
+
+
+class MockEmitter:
+    def emit(self, value):
+        pass
+
+
+def ram_message():
+    """
+    Determine if the system's RAM capacity is sufficient for running reconstruction.
+    The message should be treated as a warning if the RAM detected is less than 32 GB.
+
+    Returns
+    -------
+    ram_report    (is_warning, message)
+    """
+    BYTES_PER_GB = 2**30
+    gb_available = psutil.virtual_memory().total / BYTES_PER_GB
+    is_warning = gb_available < 32
+
+    if is_warning:
+        message = " \n".join(
+            textwrap.wrap(
+                f"waveorder reconstructions often require more than the {gb_available:.1f} "
+                f"GB of RAM that this computer is equipped with. We recommend starting with reconstructions of small "
+                f"volumes ~1000 x 1000 x 10 and working up to larger volumes while monitoring your RAM usage with "
+                f"Task Manager or htop.",
+            )
+        )
+    else:
+        message = f"{gb_available:.1f} GB of RAM is available."
+
+    return (is_warning, message)
+
+
+def model_to_yaml(model: MyBaseModel, yaml_path: Path) -> None:
+    """
+    Save a model's dictionary representation to a YAML file.
+
+    Parameters
+    ----------
+    model : MyBaseModel
+        The model object to convert to YAML.
+    yaml_path : Path
+        The path to the output YAML file.
+
+    Raises
+    ------
+    TypeError
+        If the `model` object does not have a `dict()` method.
+
+    Notes
+    -----
+    This function converts a model object into a dictionary representation
+    using the `dict()` method. It removes any fields with None values before
+    writing the dictionary to a YAML file.
+
+    Examples
+    --------
+    >>> from my_model import MyModel
+    >>> model = MyModel()
+    >>> model_to_yaml(model, 'model.yaml')
+
+    """
+    yaml_path = Path(yaml_path)
+
+    if not hasattr(model, "dict"):
+        raise TypeError("The 'model' object does not have a 'dict()' method.")
+
+    model_dict = model.model_dump()
+
+    # Remove None-valued fields
+    clean_model_dict = {
+        key: value for key, value in model_dict.items() if value is not None
+    }
+
+    with open(yaml_path, "w+") as f:
+        yaml.dump(
+            clean_model_dict, f, default_flow_style=False, sort_keys=False
+        )
+
+
+def yaml_to_model(yaml_path: Path, model):
+    """
+    Load model settings from a YAML file and create a model instance.
+
+    Parameters
+    ----------
+    yaml_path : Path
+        The path to the YAML file containing the model settings.
+    model : class
+        The model class used to create an instance with the loaded settings.
+
+    Returns
+    -------
+    object
+        An instance of the model class with the loaded settings.
+
+    Raises
+    ------
+    TypeError
+        If the provided model is not a class or does not have a callable constructor.
+    FileNotFoundError
+        If the YAML file specified by `yaml_path` does not exist.
+
+    Notes
+    -----
+    This function loads model settings from a YAML file using `yaml.safe_load()`.
+    It then creates an instance of the provided `model` class using the loaded settings.
+
+    Examples
+    --------
+    # >>> from my_model import MyModel
+    # >>> model = yaml_to_model('model.yaml', MyModel)
+
+    """
+    yaml_path = Path(yaml_path)
+
+    if not callable(getattr(model, "__init__", None)):
+        raise TypeError(
+            "The provided model must be a class with a callable constructor."
+        )
+
+    try:
+        with open(yaml_path, "r") as file:
+            raw_settings = yaml.safe_load(file)
+    except FileNotFoundError:
+        raise FileNotFoundError(f"The YAML file '{yaml_path}' does not exist.")
+
+    return model(**raw_settings)

waveorder/io/visualization.py

@@ -0,0 +1,160 @@
+from typing import Literal, Union
+
+import numpy as np
+from colorspacious import cspace_convert
+from matplotlib.colors import hsv_to_rgb
+from skimage.color import hsv2rgb
+from skimage.exposure import rescale_intensity
+
+
+def ret_ori_overlay(
+    czyx,
+    ret_min: float = 1,
+    ret_max: Union[float, Literal["auto"]] = 10,
+    cmap: Literal["JCh", "HSV"] = "HSV",
+):
+    """
+    Creates an overlay of retardance and orientation with two different colormap options.
+    "HSV" maps orientation to hue and retardance to value with maximum saturation.
+    "JCh" is a similar colormap but is perceptually uniform.
+
+    Parameters
+    ----------
+    czyx:                   (nd-array) czyx[0] is retardance in nanometers, czyx[1] is orientation in radians [0, pi],
+                            czyx.shape = (2, ...)
+
+    ret_min:                (float) minimum displayed retardance. Typically a noise floor.
+    ret_max:                (float) maximum displayed retardance. Typically used to adjust contrast limits.
+
+    cmap:                   (str) 'JCh' or 'HSV'
+
+    Returns
+    -------
+    overlay                 (nd-array) RGB image with shape (3, ...)
+
+    """
+    if czyx.shape[0] != 2:
+        raise ValueError(
+            f"Input must have shape (2, ...) instead of ({czyx.shape[0]}, ...)"
+        )
+
+    retardance = czyx[0]
+    orientation = czyx[1]
+
+    if ret_max == "auto":
+        ret_max = np.percentile(np.ravel(retardance), 99.99)
+
+    # Prepare input and output arrays
+    ret_ = np.clip(retardance, 0, ret_max)  # clip and copy
+    # Convert 180 degree range into 360 to match periodicity of hue.
+    ori_ = orientation * 360 / np.pi
+    overlay_final = np.zeros_like(retardance)
+
+    if cmap == "JCh":
+        J_MAX = 65
+        C_MAX = 60
+
+        J = (ret_ / ret_max) * J_MAX
+        C = np.ones_like(J) * C_MAX
+        C[ret_ < ret_min] = 0
+        h = ori_
+
+        JCh = np.stack((J, C, h), axis=-1)
+        JCh_rgb = cspace_convert(JCh, "JCh", "sRGB255")
+
+        JCh_rgb[JCh_rgb < 0] = 0
+        JCh_rgb[JCh_rgb > 255] = 255
+
+        overlay_final = JCh_rgb.astype(np.uint8)
+    elif cmap == "HSV":
+        I_hsv = np.moveaxis(
+            np.stack(
+                [
+                    ori_ / 360,
+                    np.ones_like(ori_),
+                    ret_ / np.max(ret_),
+                ]
+            ),
+            source=0,
+            destination=-1,
+        )
+        overlay_final = hsv_to_rgb(I_hsv)
+    else:
+        raise ValueError(f"Colormap {cmap} not understood")
+
+    return np.moveaxis(
+        overlay_final, source=-1, destination=0
+    )  # .shape = (3, ...)
+
+
+def ret_ori_phase_overlay(
+    czyx, max_val_V: float = 1.0, max_val_S: float = 1.0
+):
+    """
+    Creates an overlay of retardance, orientation, and phase.
+    Maps orientation to hue, retardance to saturation, and phase to value.
+
+    HSV encoding of retardance + orientation + phase image with hsv colormap
+    (orientation in h, retardance in s, phase in v)
+    Parameters
+    ----------
+        czyx        : numpy.ndarray
+                    czyx[0] corresponds to the retardance image
+                    czyx[1]is the orientation image (range from 0 to pi)
+                    czyx[2] is the the phase image
+
+        max_val_V   : float
+                      raise the brightness of the phase channel by 1/max_val_V
+
+        max_val_S   : float
+                      raise the brightness of the retardance channel by 1/max_val_S
+
+    Returns
+    -------
+    overlay                 (nd-array) RGB image with shape (3, ...)
+
+    Returns:
+        RGB with HSV
+    """
+
+    if czyx.shape[0] != 3:
+        raise ValueError(
+            f"Input must have shape (3, ...) instead of ({czyx.shape[0]}, ...)"
+        )
+
+    czyx_out = np.zeros_like(czyx, dtype=np.float32)
+
+    retardance = czyx[0]
+    orientation = czyx[1]
+    phase = czyx[2]
+
+    # Normalize the stack
+    ordered_stack = np.stack(
+        (
+            # Normalize the first channel by dividing by pi
+            orientation / np.pi,
+            # Normalize the second channel and rescale intensity
+            rescale_intensity(
+                retardance,
+                in_range=(
+                    np.min(retardance),
+                    np.max(retardance),
+                ),
+                out_range=(0, 1),
+            )
+            / max_val_S,
+            # Normalize the third channel and rescale intensity
+            rescale_intensity(
+                phase,
+                in_range=(
+                    np.min(phase),
+                    np.max(phase),
+                ),
+                out_range=(0, 1),
+            )
+            / max_val_V,
+        ),
+        axis=0,
+    )
+    czyx_out = hsv2rgb(ordered_stack, channel_axis=0)
+    return czyx_out

waveorder/models/inplane_oriented_thick_pol3d_vector.py

@@ -65,9 +65,9 @@ def calculate_transfer_function(
     print("Z factor:", z_factor)
 
     tf_calculation_shape = (
-        zyx_shape[0] * z_factor * fourier_oversample_factor,
-        int(np.ceil(zyx_shape[1] * yx_factor * fourier_oversample_factor)),
-        int(np.ceil(zyx_shape[2] * yx_factor * fourier_oversample_factor)),
+        int(zyx_shape[0] * z_factor / fourier_oversample_factor),
+        int(np.ceil(zyx_shape[1] * yx_factor / fourier_oversample_factor)),
+        int(np.ceil(zyx_shape[2] * yx_factor / fourier_oversample_factor)),
     )
 
     (

waveorder/models/isotropic_fluorescent_thick_3d.py

@@ -31,7 +31,40 @@ def calculate_transfer_function(
     z_padding: int,
     index_of_refraction_media: float,
     numerical_aperture_detection: float,
+    confocal_pinhole_diameter: float | None = None,
 ) -> Tensor:
+    """Calculate the optical transfer function for fluorescence imaging.
+
+    Supports both widefield and confocal microscopy modes. When
+    confocal_pinhole_diameter is None, computes widefield OTF. When specified,
+    computes confocal OTF by multiplying excitation and detection PSFs, where
+    the detection PSF is downweighted by the pinhole aperture function.
+
+    Parameters
+    ----------
+    zyx_shape : tuple[int, int, int]
+        Shape of the 3D volume
+    yx_pixel_size : float
+        Pixel size in YX plane
+    z_pixel_size : float
+        Pixel size in Z dimension
+    wavelength_emission : float
+        Emission wavelength
+    z_padding : int
+        Padding for axial dimension
+    index_of_refraction_media : float
+        Refractive index of imaging medium
+    numerical_aperture_detection : float
+        Numerical aperture of detection objective
+    confocal_pinhole_diameter : float | None, optional
+        Diameter of confocal pinhole in image space (demagnified). If None,
+        computes widefield OTF. If specified, computes confocal OTF.
+
+    Returns
+    -------
+    Tensor
+        3D optical transfer function
+    """
     transverse_nyquist = sampling.transverse_nyquist(
         wavelength_emission,
         numerical_aperture_detection,  # ill = det for fluorescence
@@ -43,6 +76,11 @@ def calculate_transfer_function(
         index_of_refraction_media,
     )
 
+    # For confocal, double the Nyquist range (half the sampling requirement)
+    if confocal_pinhole_diameter is not None:
+        transverse_nyquist = transverse_nyquist / 2
+        axial_nyquist = axial_nyquist / 2
+
     yx_factor = int(np.ceil(yx_pixel_size / transverse_nyquist))
     z_factor = int(np.ceil(z_pixel_size / axial_nyquist))
 
@@ -58,6 +96,7 @@ def calculate_transfer_function(
         z_padding,
         index_of_refraction_media,
         numerical_aperture_detection,
+        confocal_pinhole_diameter,
     )
     zyx_out_shape = (zyx_shape[0] + 2 * z_padding,) + zyx_shape[1:]
     return sampling.nd_fourier_central_cuboid(
@@ -65,6 +104,35 @@ def calculate_transfer_function(
     )
 
 
+def _calculate_pinhole_aperture_otf(
+    radial_frequencies: Tensor,
+    pinhole_diameter: float,
+) -> Tensor:
+    """Calculate the pinhole aperture OTF for confocal microscopy.
+
+    The pinhole acts as a spatial filter in the image plane. A smaller pinhole
+    (approaching a point) gives a broader OTF (approaching flat/ones).
+    A larger pinhole gives a narrower OTF (approaching a delta function).
+
+    Parameters
+    ----------
+    radial_frequencies : Tensor
+        Radial spatial frequencies (units of 1/length)
+    pinhole_diameter : float
+        Diameter (not radius) of the confocal pinhole (units of length, matching
+        radial_frequencies)
+
+    Returns
+    -------
+    Tensor
+        Pinhole aperture OTF (jinc^2 function)
+    """
+    argument = pinhole_diameter * radial_frequencies
+    j1_values = torch.special.bessel_j1(np.pi * argument)
+    jinc = torch.where(argument > 1e-10, j1_values / (2 * argument), 0.5)
+    return jinc**2
+
+
 def _calculate_wrap_unsafe_transfer_function(
     zyx_shape: tuple[int, int, int],
     yx_pixel_size: float,
@@ -73,6 +141,7 @@ def _calculate_wrap_unsafe_transfer_function(
     z_padding: int,
     index_of_refraction_media: float,
     numerical_aperture_detection: float,
+    confocal_pinhole_diameter: float | None = None,
 ) -> Tensor:
     radial_frequencies = util.generate_radial_frequencies(
         zyx_shape[1:], yx_pixel_size
@@ -102,6 +171,29 @@ def _calculate_wrap_unsafe_transfer_function(
     optical_transfer_function = torch.fft.fftn(
         point_spread_function, dim=(0, 1, 2)
     )
+
+    # Confocal: multiply excitation PSF with detection PSF (downweighted by pinhole)
+    if confocal_pinhole_diameter is not None:
+        pinhole_otf_2d = _calculate_pinhole_aperture_otf(
+            radial_frequencies, confocal_pinhole_diameter
+        )
+        # Detection OTF is downweighted by pinhole
+        otf_detection = optical_transfer_function * pinhole_otf_2d[None, :, :]
+
+        # Convert to PSFs
+        psf_excitation = torch.abs(
+            torch.fft.ifftn(optical_transfer_function, dim=(0, 1, 2))
+        )
+        psf_detection = torch.abs(
+            torch.fft.ifftn(otf_detection, dim=(0, 1, 2))
+        )
+
+        # Confocal PSF = excitation PSF * detection PSF (in real space)
+        psf_confocal = psf_excitation * psf_detection
+
+        # Convert back to OTF
+        optical_transfer_function = torch.fft.fftn(psf_confocal, dim=(0, 1, 2))
+
     optical_transfer_function /= torch.max(
         torch.abs(optical_transfer_function)
     )  # normalize