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

waveorder/cli/option_eat_all.py

@@ -0,0 +1,47 @@
+import click
+
+
+# Copied directly from https://stackoverflow.com/a/48394004
+# Enables `-i ./input.zarr/*/*/*`
+class OptionEatAll(click.Option):
+    def __init__(self, *args, **kwargs):
+        self.save_other_options = kwargs.pop("save_other_options", True)
+        nargs = kwargs.pop("nargs", -1)
+        assert nargs == -1, "nargs, if set, must be -1 not {}".format(nargs)
+        super(OptionEatAll, self).__init__(*args, **kwargs)
+        self._previous_parser_process = None
+        self._eat_all_parser = None
+
+    def add_to_parser(self, parser, ctx):
+        def parser_process(value, state):
+            # method to hook to the parser.process
+            done = False
+            value = [value]
+            if self.save_other_options:
+                # grab everything up to the next option
+                while state.rargs and not done:
+                    for prefix in self._eat_all_parser.prefixes:
+                        if state.rargs[0].startswith(prefix):
+                            done = True
+                    if not done:
+                        value.append(state.rargs.pop(0))
+            else:
+                # grab everything remaining
+                value += state.rargs
+                state.rargs[:] = []
+            value = tuple(value)
+
+            # call the actual process
+            self._previous_parser_process(value, state)
+
+        retval = super(OptionEatAll, self).add_to_parser(parser, ctx)
+        for name in self.opts:
+            our_parser = parser._long_opt.get(name) or parser._short_opt.get(
+                name
+            )
+            if our_parser:
+                self._eat_all_parser = our_parser
+                self._previous_parser_process = our_parser.process
+                our_parser.process = parser_process
+                break
+        return retval

waveorder/cli/parsing.py

@@ -0,0 +1,122 @@
+from pathlib import Path
+from typing import Callable
+
+import click
+import torch.multiprocessing as mp
+from iohub.ngff import Plate, open_ome_zarr
+from natsort import natsorted
+
+from waveorder.cli.option_eat_all import OptionEatAll
+
+
+def _validate_and_process_paths(
+    ctx: click.Context, opt: click.Option, value: str
+) -> list[Path]:
+    # Sort and validate the input paths, expanding plates into lists of positions
+    input_paths = [Path(path) for path in natsorted(value)]
+    for path in input_paths:
+        with open_ome_zarr(path, mode="r") as dataset:
+            if isinstance(dataset, Plate):
+                plate_path = input_paths.pop()
+                for position in dataset.positions():
+                    input_paths.append(plate_path / position[0])
+
+    return input_paths
+
+
+def _str_to_path(ctx: click.Context, opt: click.Option, value: str) -> Path:
+    return Path(value)
+
+
+def input_position_dirpaths() -> Callable:
+    def decorator(f: Callable) -> Callable:
+        return click.option(
+            "--input-position-dirpaths",
+            "-i",
+            cls=OptionEatAll,
+            type=tuple,
+            required=True,
+            callback=_validate_and_process_paths,
+            help="List of paths to input positions, each with the same TCZYX shape. Supports wildcards e.g. 'input.zarr/*/*/*'.",
+        )(f)
+
+    return decorator
+
+
+def config_filepath() -> Callable:
+    def decorator(f: Callable) -> Callable:
+        return click.option(
+            "--config-filepath",
+            "-c",
+            required=True,
+            type=click.Path(exists=True, file_okay=True, dir_okay=False),
+            callback=_str_to_path,
+            help="Path to YAML configuration file.",
+        )(f)
+
+    return decorator
+
+
+def transfer_function_dirpath() -> Callable:
+    def decorator(f: Callable) -> Callable:
+        return click.option(
+            "--transfer-function-dirpath",
+            "-t",
+            required=True,
+            type=click.Path(exists=False),
+            callback=_str_to_path,
+            help="Path to transfer function .zarr.",
+        )(f)
+
+    return decorator
+
+
+def output_dirpath() -> Callable:
+    def decorator(f: Callable) -> Callable:
+        return click.option(
+            "--output-dirpath",
+            "-o",
+            required=True,
+            type=click.Path(exists=False),
+            callback=_str_to_path,
+            help="Path to output directory.",
+        )(f)
+
+    return decorator
+
+
+# TODO: this setting will have to be collected from SLURM?
+def processes_option(default: int = None) -> Callable:
+    def check_processes_option(ctx, param, value):
+        max_processes = mp.cpu_count()
+        if value > max_processes:
+            raise click.BadParameter(
+                f"Maximum number of processes is {max_processes}"
+            )
+        return value
+
+    def decorator(f: Callable) -> Callable:
+        return click.option(
+            "--num_processes",
+            "-j",
+            default=default or mp.cpu_count(),
+            type=int,
+            help="Number of processes to run in parallel.",
+            callback=check_processes_option,
+        )(f)
+
+    return decorator
+
+
+def unique_id() -> Callable:
+    def decorator(f: Callable) -> Callable:
+        return click.option(
+            "--unique-id",
+            "-uid",
+            default="",
+            required=False,
+            type=str,
+            help="Unique ID.",
+        )(f)
+
+    return decorator

waveorder/cli/printing.py

@@ -0,0 +1,16 @@
+import click
+import yaml
+
+from waveorder.cli.settings import MyBaseModel
+
+
+def echo_settings(settings: MyBaseModel):
+    click.echo(
+        yaml.dump(
+            settings.model_dump(), default_flow_style=False, sort_keys=False
+        )
+    )
+
+
+def echo_headline(headline):
+    click.echo(click.style(headline, fg="green"))

waveorder/cli/reconstruct.py

@@ -0,0 +1,67 @@
+from pathlib import Path
+
+import click
+
+from waveorder.cli.apply_inverse_transfer_function import (
+    apply_inverse_transfer_function_cli,
+)
+from waveorder.cli.compute_transfer_function import (
+    compute_transfer_function_cli,
+)
+from waveorder.cli.parsing import (
+    config_filepath,
+    input_position_dirpaths,
+    output_dirpath,
+    processes_option,
+    unique_id,
+)
+
+
+@click.command("reconstruct")
+@input_position_dirpaths()
+@config_filepath()
+@output_dirpath()
+@processes_option(default=1)
+@unique_id()
+def _reconstruct_cli(
+    input_position_dirpaths,
+    config_filepath,
+    output_dirpath,
+    num_processes,
+    unique_id,
+):
+    """
+    Reconstruct a dataset using a configuration file. This is a
+    convenience function for a `compute-tf` call followed by a `apply-inv-tf`
+    call.
+
+    Calculates the transfer function based on the shape of the first position
+    in the list `input-position-dirpaths`, then applies that transfer function
+    to all positions in the list `input-position-dirpaths`, so all positions
+    must have the same TCZYX shape.
+
+    See /examples for example configuration files.
+
+    >> waveorder reconstruct -i ./input.zarr/*/*/* -c ./examples/birefringence.yml -o ./output.zarr
+    """
+
+    # Handle transfer function path
+    transfer_function_path = output_dirpath.parent / Path(
+        "transfer_function_" + config_filepath.stem + ".zarr"
+    )
+
+    # Compute transfer function
+    compute_transfer_function_cli(
+        input_position_dirpaths[0],
+        config_filepath,
+        transfer_function_path,
+    )
+
+    # Apply inverse transfer function
+    apply_inverse_transfer_function_cli(
+        input_position_dirpaths,
+        transfer_function_path,
+        config_filepath,
+        output_dirpath,
+        num_processes,
+    )

waveorder/cli/settings.py

@@ -0,0 +1,187 @@
+import os
+import warnings
+from pathlib import Path
+from typing import List, Literal, Optional, Union
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Extra,
+    NonNegativeFloat,
+    NonNegativeInt,
+    PositiveFloat,
+    field_validator,
+    model_validator,
+)
+
+# This file defines the configuration settings for the CLI.
+
+# Example settings files in `/docs/examples/settings/` are autmatically generated
+# by the tests in `/tests/cli_tests/test_settings.py` - `test_generate_example_settings`.
+
+# To keep the example settings up to date, run `pytest` locally when this file changes.
+
+
+# All settings classes inherit from MyBaseModel, which forbids extra parameters to guard against typos
+class MyBaseModel(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+# Bottom level settings
+class WavelengthIllumination(MyBaseModel):
+    wavelength_illumination: PositiveFloat = 0.532
+
+
+class BirefringenceTransferFunctionSettings(MyBaseModel):
+    swing: float = 0.1
+
+    @field_validator("swing")
+    @classmethod
+    def swing_range(cls, v):
+        if v <= 0 or v >= 1.0:
+            raise ValueError(f"swing = {v} should be between 0 and 1.")
+        return v
+
+
+class BirefringenceApplyInverseSettings(WavelengthIllumination):
+    background_path: Union[str, Path] = ""
+
+    @field_validator("background_path")
+    @classmethod
+    def check_background_path(cls, v):
+        if v == "":
+            return v
+
+        raw_dir = r"{}".format(v)
+        if not os.path.isdir(raw_dir):
+            raise ValueError(f"{v} is not a existing directory")
+        return raw_dir
+
+    remove_estimated_background: bool = False
+    flip_orientation: bool = False
+    rotate_orientation: bool = False
+
+
+class FourierTransferFunctionSettings(MyBaseModel):
+    yx_pixel_size: PositiveFloat = 6.5 / 20
+    z_pixel_size: PositiveFloat = 2.0
+    z_padding: NonNegativeInt = 0
+    z_focus_offset: Union[float, Literal["auto"]] = 0
+    index_of_refraction_media: PositiveFloat = 1.3
+    numerical_aperture_detection: PositiveFloat = 1.2
+
+    @model_validator(mode="after")
+    def validate_numerical_aperture_detection(self):
+        if self.numerical_aperture_detection > self.index_of_refraction_media:
+            raise ValueError(
+                f"numerical_aperture_detection = {self.numerical_aperture_detection} must be less than or equal to index_of_refraction_media = {self.index_of_refraction_media}"
+            )
+        return self
+
+    @model_validator(mode="after")
+    def warn_unit_consistency(self):
+        ratio = self.yx_pixel_size / self.z_pixel_size
+        if ratio < 1.0 / 20 or ratio > 20:
+            warnings.warn(
+                f"yx_pixel_size ({self.yx_pixel_size}) / z_pixel_size ({self.z_pixel_size}) = {ratio}. Did you use consistent units?",
+                UserWarning,
+            )
+        return self
+
+
+class FourierApplyInverseSettings(MyBaseModel):
+    reconstruction_algorithm: Literal["Tikhonov", "TV"] = "Tikhonov"
+    regularization_strength: NonNegativeFloat = 1e-3
+    TV_rho_strength: PositiveFloat = 1e-3
+    TV_iterations: NonNegativeInt = 1
+
+
+class PhaseTransferFunctionSettings(
+    FourierTransferFunctionSettings,
+    WavelengthIllumination,
+):
+    numerical_aperture_illumination: NonNegativeFloat = 0.5
+    invert_phase_contrast: bool = False
+
+    @model_validator(mode="after")
+    def validate_numerical_aperture_illumination(self):
+        if (
+            self.numerical_aperture_illumination
+            > self.index_of_refraction_media
+        ):
+            raise ValueError(
+                f"numerical_aperture_illumination = {self.numerical_aperture_illumination} must be less than or equal to index_of_refraction_media = {self.index_of_refraction_media}"
+            )
+        return self
+
+
+class FluorescenceTransferFunctionSettings(FourierTransferFunctionSettings):
+    wavelength_emission: PositiveFloat = 0.507
+    confocal_pinhole_diameter: Optional[PositiveFloat] = None
+
+    @model_validator(mode="after")
+    def warn_unit_consistency(self):
+        ratio = self.yx_pixel_size / self.wavelength_emission
+        if ratio < 1.0 / 20 or ratio > 20:
+            warnings.warn(
+                f"yx_pixel_size ({self.yx_pixel_size}) / wavelength_illumination ({self.wavelength_emission}) = {ratio}. Did you use consistent units?",
+                UserWarning,
+            )
+        return self
+
+
+# Second level settings
+class BirefringenceSettings(MyBaseModel):
+    transfer_function: BirefringenceTransferFunctionSettings = (
+        BirefringenceTransferFunctionSettings()
+    )
+    apply_inverse: BirefringenceApplyInverseSettings = (
+        BirefringenceApplyInverseSettings()
+    )
+
+
+class PhaseSettings(MyBaseModel):
+    transfer_function: PhaseTransferFunctionSettings = (
+        PhaseTransferFunctionSettings()
+    )
+    apply_inverse: FourierApplyInverseSettings = FourierApplyInverseSettings()
+
+
+class FluorescenceSettings(MyBaseModel):
+    transfer_function: FluorescenceTransferFunctionSettings = (
+        FluorescenceTransferFunctionSettings()
+    )
+    apply_inverse: FourierApplyInverseSettings = FourierApplyInverseSettings()
+
+
+# Top level settings
+class ReconstructionSettings(MyBaseModel):
+    input_channel_names: List[str] = [f"State{i}" for i in range(4)]
+    time_indices: Union[
+        NonNegativeInt, List[NonNegativeInt], Literal["all"]
+    ] = "all"
+    reconstruction_dimension: Literal[2, 3] = 3
+    birefringence: Optional[BirefringenceSettings] = None
+    phase: Optional[PhaseSettings] = None
+    fluorescence: Optional[FluorescenceSettings] = None
+
+    @model_validator(mode="after")
+    def validate_reconstruction_types(self):
+        if (
+            self.birefringence or self.phase
+        ) and self.fluorescence is not None:
+            raise ValueError(
+                '"fluorescence" cannot be present alongside "birefringence" or "phase". Please use one configuration file for a "fluorescence" reconstruction and another configuration file for a "birefringence" and/or "phase" reconstructions.'
+            )
+        num_channel_names = len(self.input_channel_names)
+        if self.birefringence is None:
+            if self.phase is None and self.fluorescence is None:
+                raise ValueError(
+                    "Provide settings for either birefringence, phase, birefringence + phase, or fluorescence."
+                )
+            if num_channel_names != 1:
+                raise ValueError(
+                    f"{num_channel_names} channels names provided. Please provide a single channel for fluorescence/phase reconstructions."
+                )
+
+        return self

waveorder/cli/utils.py

@@ -0,0 +1,175 @@
+from pathlib import Path
+from typing import Tuple
+
+import click
+import numpy as np
+import torch
+from iohub.ngff import Position, open_ome_zarr
+from iohub.ngff.models import TransformationMeta
+from numpy.typing import DTypeLike
+
+
+def generate_valid_position_key(index: int) -> tuple[str, str, str]:
+    """Generate a valid HCS position key for single-position stores.
+
+    Args:
+        index: Position index (0-based)
+
+    Returns:
+        Tuple of (row, column, field) with alphanumeric characters only
+    """
+    row = chr(65 + (index // 10))  # A, B, C, etc.
+    column = str((index % 10) + 1)  # 1, 2, 3, etc.
+    field = "0"  # Always 0 for single positions
+    return (row, column, field)
+
+
+def is_single_position_store(position_path: Path) -> bool:
+    """Check if a position path is from a single-position store (not HCS plate).
+
+    Args:
+        position_path: Path to the position directory
+
+    Returns:
+        True if it's a single-position store, False if it's part of an HCS plate
+    """
+    try:
+        # Try to open as HCS plate 3 levels up
+        open_ome_zarr(position_path.parent.parent.parent, mode="r")
+        return False  # Successfully opened as plate
+    except (RuntimeError, FileNotFoundError):
+        return True  # Not a plate structure
+
+
+def create_empty_hcs_zarr(
+    store_path: Path,
+    position_keys: list[Tuple[str]],
+    shape: Tuple[int],
+    chunks: Tuple[int],
+    scale: Tuple[float],
+    channel_names: list[str],
+    dtype: DTypeLike,
+    plate_metadata: dict = {},
+) -> None:
+    """If the plate does not exist, create an empty zarr plate.
+
+    If the plate exists, append positions and channels if they are not
+    already in the plate.
+
+    Parameters
+    ----------
+    store_path : Path
+        hcs plate path
+    position_keys : list[Tuple[str]]
+        Position keys, will append if not present in the plate.
+        e.g. [("A", "1", "0"), ("A", "1", "1")]
+    shape : Tuple[int]
+    chunks : Tuple[int]
+    scale : Tuple[float]
+    channel_names : list[str]
+        Channel names, will append if not present in metadata.
+    dtype : DTypeLike
+    plate_metadata : dict
+    """
+
+    # Create plate
+    output_plate = open_ome_zarr(
+        str(store_path), layout="hcs", mode="a", channel_names=channel_names
+    )
+
+    # Pass metadata
+    output_plate.zattrs.update(plate_metadata)
+
+    # Create positions
+    for position_key in position_keys:
+        position_key_string = "/".join(position_key)
+        # Check if position is already in the store, if not create it
+        if position_key_string not in output_plate.zgroup:
+            position = output_plate.create_position(*position_key)
+
+            _ = position.create_zeros(
+                name="0",
+                shape=shape,
+                chunks=chunks,
+                dtype=dtype,
+                transform=[TransformationMeta(type="scale", scale=scale)],
+            )
+        else:
+            position = output_plate[position_key_string]
+
+        # Check if channel_names are already in the store, if not append them
+        for channel_name in channel_names:
+            # Read channel names directly from metadata to avoid race conditions
+            metadata_channel_names = [
+                channel.label for channel in position.metadata.omero.channels
+            ]
+            if channel_name not in metadata_channel_names:
+                position.append_channel(channel_name, resize_arrays=True)
+
+
+def apply_inverse_to_zyx_and_save(
+    func,
+    position: Position,
+    output_path: Path,
+    input_channel_indices: list[int],
+    output_channel_indices: list[int],
+    t_idx: int = 0,
+    **kwargs,
+) -> None:
+    """Load a zyx array from a Position object, apply a transformation and save the result to file"""
+    click.echo(f"Reconstructing t={t_idx}")
+
+    # Load data
+    czyx_uint16_numpy = position.data.oindex[t_idx, input_channel_indices]
+
+    # Check if all values in czyx_uint16_numpy are not zeros or Nan
+    if _check_nan_n_zeros(czyx_uint16_numpy):
+        click.echo(
+            f"All values at t={t_idx} are zero or Nan, skipping reconstruction."
+        )
+        return
+
+    # convert to np.int32 (torch doesn't accept np.uint16), then convert to tensor float32
+    czyx_data = torch.tensor(np.int32(czyx_uint16_numpy), dtype=torch.float32)
+
+    # Apply transformation
+    reconstruction_czyx = func(czyx_data, **kwargs)
+
+    # Write to file
+    # for c, recon_zyx in enumerate(reconstruction_zyx):
+    with open_ome_zarr(output_path, mode="r+") as output_dataset:
+        output_dataset[0].oindex[
+            t_idx, output_channel_indices
+        ] = reconstruction_czyx
+    click.echo(f"Finished Writing.. t={t_idx}")
+
+
+def estimate_resources(shape, settings, num_processes):
+    T, C, Z, Y, X = shape
+
+    gb_ram_per_cpu = 0
+    gb_per_element = 4 / 2**30  # bytes_per_float32 / bytes_per_gb
+    voxel_resource_multiplier = 4
+    fourier_resource_multiplier = 32
+    input_memory = Z * Y * X * gb_per_element
+
+    if settings.birefringence is not None:
+        gb_ram_per_cpu += input_memory * voxel_resource_multiplier
+    if settings.phase is not None:
+        gb_ram_per_cpu += input_memory * fourier_resource_multiplier
+    if settings.fluorescence is not None:
+        gb_ram_per_cpu += input_memory * fourier_resource_multiplier
+    ram_multiplier = 1
+    gb_ram_per_cpu = np.ceil(
+        np.max([1, ram_multiplier * gb_ram_per_cpu])
+    ).astype(int)
+    num_cpus = np.min([32, num_processes])
+
+    return num_cpus, gb_ram_per_cpu
+
+
+def _check_nan_n_zeros(input_array):
+    """
+    Checks if data are all zeros or nan
+    """
+    return np.all(np.isnan(input_array)) or np.all(input_array == 0)

waveorder/filter.py

@@ -74,7 +74,6 @@ def apply_filter_bank(
         )
 
     num_input_channels, num_output_channels = io_filter_bank.shape[:2]
-    spatial_dims = io_filter_bank.shape[2:]
 
     # Pad input_array until each dimension is divisible by transfer_function
     pad_sizes = [
@@ -98,7 +97,7 @@ def apply_filter_bank(
     # Further optimization is likely with a combination of
     # torch.baddbmm, torch.pixel_shuffle, torch.pixel_unshuffle.
     padded_output_spectrum = torch.zeros(
-        (num_output_channels,) + spatial_dims,
+        (num_output_channels,) + padded_input_spectrum.shape[1:],
         dtype=padded_input_spectrum.dtype,
         device=padded_input_spectrum.device,
     )