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

waveorder/cli/apply_inverse_models.py

@@ -0,0 +1,328 @@
+"""
+This module converts GUI-level reconstruction calls into library calls
+"""
+
+import numpy as np
+import torch
+
+from waveorder.cli.settings import FluorescenceSettings, PhaseSettings
+from waveorder.models import (
+    inplane_oriented_thick_pol3d,
+    inplane_oriented_thick_pol3d_vector,
+    isotropic_fluorescent_thick_3d,
+    isotropic_fluorescent_thin_3d,
+    isotropic_thin_3d,
+    phase_thick_3d,
+)
+from waveorder.stokes import _s12_to_orientation, stokes_after_adr
+
+
+def radians_to_nanometers(retardance_rad, wavelength_illumination_um):
+    """
+    waveorder returns retardance in radians, while waveorder displays and saves
+    retardance in nanometers. This function converts from radians to nanometers
+    using the illumination wavelength (which is internally handled in um
+    in waveorder).
+    """
+    return retardance_rad * wavelength_illumination_um * 1e3 / (2 * np.pi)
+
+
+def birefringence(
+    czyx_data,
+    cyx_no_sample_data,
+    wavelength_illumination,
+    recon_dim,
+    biref_inverse_dict,
+    transfer_function_dataset,
+):
+    # Load transfer function
+    intensity_to_stokes_matrix = torch.tensor(
+        transfer_function_dataset["intensity_to_stokes_matrix"][0, 0, 0]
+    )
+
+    # Apply reconstruction
+    # (retardance, orientation, transmittance, depolarization)
+    reconstructed_parameters = (
+        inplane_oriented_thick_pol3d.apply_inverse_transfer_function(
+            czyx_data,
+            intensity_to_stokes_matrix,
+            cyx_no_sample_data=cyx_no_sample_data,
+            project_stokes_to_2d=(recon_dim == 2),
+            **biref_inverse_dict,
+        )
+    )
+
+    # Convert retardance
+    retardance = radians_to_nanometers(
+        reconstructed_parameters[0], wavelength_illumination
+    )
+
+    return torch.stack((retardance,) + reconstructed_parameters[1:])
+
+
+def phase(
+    czyx_data,
+    recon_dim,
+    settings_phase: PhaseSettings,
+    transfer_function_dataset,
+):
+    # [phase only, 2]
+    if recon_dim == 2:
+        # Load transfer functions
+        U = torch.from_numpy(transfer_function_dataset["singular_system_U"][0])
+        S = torch.from_numpy(
+            transfer_function_dataset["singular_system_S"][0, 0]
+        )
+        Vh = torch.from_numpy(
+            transfer_function_dataset["singular_system_Vh"][0]
+        )
+
+        # Apply
+        (
+            absorption_yx,
+            phase_yx,
+        ) = isotropic_thin_3d.apply_inverse_transfer_function(
+            czyx_data[0],
+            (U, S, Vh),
+            **settings_phase.apply_inverse.model_dump(),
+        )
+        # Stack to C1YX
+        output = phase_yx[None, None]
+        # TODO: Write phase and absorption to CZYX
+        # torch.stack((phase_yx[None], absorption_yx[None]))
+
+    # [phase only, 3]
+    elif recon_dim == 3:
+        # Load transfer functions
+        real_potential_transfer_function = torch.tensor(
+            transfer_function_dataset["real_potential_transfer_function"][0, 0]
+        )
+        imaginary_potential_transfer_function = torch.tensor(
+            transfer_function_dataset["imaginary_potential_transfer_function"][
+                0, 0
+            ]
+        )
+
+        # Apply
+        output = phase_thick_3d.apply_inverse_transfer_function(
+            czyx_data[0],
+            real_potential_transfer_function,
+            imaginary_potential_transfer_function,
+            z_padding=settings_phase.transfer_function.z_padding,
+            **settings_phase.apply_inverse.model_dump(),
+        )
+
+    # Pad to CZYX
+    while output.ndim != 4:
+        output = torch.unsqueeze(output, 0)
+
+    return output
+
+
+def birefringence_and_phase(
+    czyx_data,
+    cyx_no_sample_data,
+    wavelength_illumination,
+    recon_dim,
+    biref_inverse_dict,
+    settings_phase: PhaseSettings,
+    transfer_function_dataset,
+):
+    # Load birefringence transfer function
+    intensity_to_stokes_matrix = torch.tensor(
+        transfer_function_dataset["intensity_to_stokes_matrix"][0, 0, 0]
+    )
+
+    # [biref and phase, 2]
+    if recon_dim == 2:
+        # Load transfer functions
+        U = torch.from_numpy(
+            transfer_function_dataset["vector_singular_system_U"][0]
+        )
+        S = torch.from_numpy(
+            transfer_function_dataset["vector_singular_system_S"][0, 0]
+        )
+        Vh = torch.from_numpy(
+            transfer_function_dataset["vector_singular_system_Vh"][0]
+        )
+
+        # Apply
+        reconstructed_parameters_2d = (
+            inplane_oriented_thick_pol3d.apply_inverse_transfer_function(
+                czyx_data,
+                intensity_to_stokes_matrix,
+                cyx_no_sample_data=cyx_no_sample_data,
+                project_stokes_to_2d=True,
+                **biref_inverse_dict,
+            )
+        )
+
+        reconstructed_parameters_3d = (
+            inplane_oriented_thick_pol3d.apply_inverse_transfer_function(
+                czyx_data,
+                intensity_to_stokes_matrix,
+                cyx_no_sample_data=cyx_no_sample_data,
+                project_stokes_to_2d=False,
+                **biref_inverse_dict,
+            )
+        )
+
+        brightfield_3d = reconstructed_parameters_3d[2]
+
+        (
+            _,
+            yx_phase,
+        ) = isotropic_thin_3d.apply_inverse_transfer_function(
+            brightfield_3d,
+            (U, S, Vh),
+            **settings_phase.apply_inverse.model_dump(),
+        )
+
+        # Convert retardance
+        retardance = radians_to_nanometers(
+            reconstructed_parameters_2d[0], wavelength_illumination
+        )
+
+        output = torch.stack(
+            (retardance,)
+            + reconstructed_parameters_2d[1:]
+            + (torch.unsqueeze(yx_phase, 0),)
+        )  # CZYX
+
+    # [biref and phase, 3]
+    elif recon_dim == 3:
+        # Load phase transfer functions
+        intensity_to_stokes_matrix = torch.tensor(
+            transfer_function_dataset["intensity_to_stokes_matrix"][0, 0, 0]
+        )
+        # Load transfer functions
+        real_potential_transfer_function = torch.tensor(
+            transfer_function_dataset["real_potential_transfer_function"][0, 0]
+        )
+        imaginary_potential_transfer_function = torch.tensor(
+            transfer_function_dataset["imaginary_potential_transfer_function"][
+                0, 0
+            ]
+        )
+
+        # Apply
+        reconstructed_parameters_3d = (
+            inplane_oriented_thick_pol3d.apply_inverse_transfer_function(
+                czyx_data,
+                intensity_to_stokes_matrix,
+                cyx_no_sample_data=cyx_no_sample_data,
+                project_stokes_to_2d=False,
+                **biref_inverse_dict,
+            )
+        )
+
+        brightfield_3d = reconstructed_parameters_3d[2]
+
+        zyx_phase = phase_thick_3d.apply_inverse_transfer_function(
+            brightfield_3d,
+            real_potential_transfer_function,
+            imaginary_potential_transfer_function,
+            z_padding=settings_phase.transfer_function.z_padding,
+            **settings_phase.apply_inverse.model_dump(),
+        )
+
+        # Convert retardance
+        retardance = radians_to_nanometers(
+            reconstructed_parameters_3d[0], wavelength_illumination
+        )
+        # Load singular system
+        U = torch.tensor(
+            np.array(transfer_function_dataset["vector_singular_system_U"])
+        )
+        S = torch.tensor(
+            np.array(transfer_function_dataset["vector_singular_system_S"][0])
+        )
+        Vh = torch.tensor(
+            np.array(transfer_function_dataset["vector_singular_system_Vh"])
+        )
+        singular_system = (U, S, Vh)
+
+        # Convert retardance and orientation to stokes
+        stokes = stokes_after_adr(*reconstructed_parameters_3d)
+
+        stokes = torch.nan_to_num_(
+            torch.stack(stokes), nan=0.0
+        )  # very rare nans from previous like
+
+        # Apply reconstruction
+        joint_recon_params = inplane_oriented_thick_pol3d_vector.apply_inverse_transfer_function(
+            szyx_data=stokes,
+            singular_system=singular_system,
+            intensity_to_stokes_matrix=None,
+            **settings_phase.apply_inverse.model_dump(),
+        )
+
+        new_ret = (
+            joint_recon_params[1] ** 2 + joint_recon_params[2] ** 2
+        ) ** (0.5)
+        new_ori = _s12_to_orientation(
+            joint_recon_params[1], -joint_recon_params[2]
+        )
+
+        # Convert stokes to retardance and orientation
+        # new_ret, new_ori, _ = estimate_ar_from_stokes012(*joint_recon_params)
+
+        # Convert retardance
+        new_ret_nm = radians_to_nanometers(new_ret, wavelength_illumination)
+
+        # Save
+        output = torch.stack(
+            (retardance,)
+            + reconstructed_parameters_3d[1:]
+            + (zyx_phase,)
+            + (new_ret_nm,)
+            + (new_ori,)
+            + (joint_recon_params[0],)
+        )
+    return output
+
+
+def fluorescence(
+    czyx_data,
+    recon_dim,
+    settings_fluorescence: FluorescenceSettings,
+    transfer_function_dataset,
+):
+    # [fluo, 2]
+    if recon_dim == 2:
+        # Load transfer functions for 2D thin fluorescence reconstruction
+        U = torch.from_numpy(transfer_function_dataset["singular_system_U"][0])
+        S = torch.from_numpy(
+            transfer_function_dataset["singular_system_S"][0, 0]
+        )
+        Vh = torch.from_numpy(
+            transfer_function_dataset["singular_system_Vh"][0]
+        )
+
+        # Apply 2D fluorescence reconstruction
+        output = isotropic_fluorescent_thin_3d.apply_inverse_transfer_function(
+            czyx_data[0],
+            (U, S, Vh),
+            **settings_fluorescence.apply_inverse.model_dump(),
+        )
+    # [fluo, 3]
+    elif recon_dim == 3:
+        # Load transfer functions
+        optical_transfer_function = torch.tensor(
+            transfer_function_dataset["optical_transfer_function"][0, 0]
+        )
+
+        # Apply
+        output = (
+            isotropic_fluorescent_thick_3d.apply_inverse_transfer_function(
+                czyx_data[0],
+                optical_transfer_function,
+                settings_fluorescence.transfer_function.z_padding,
+                **settings_fluorescence.apply_inverse.model_dump(),
+            )
+        )
+        # Pad to CZYX
+    while output.ndim != 4:
+        output = torch.unsqueeze(output, 0)
+
+    return output

waveorder/cli/apply_inverse_transfer_function.py

@@ -0,0 +1,379 @@
+import itertools
+import warnings
+from functools import partial
+from pathlib import Path
+
+import click
+import numpy as np
+import torch
+import torch.multiprocessing as mp
+from iohub import open_ome_zarr
+
+from waveorder.cli import apply_inverse_models
+from waveorder.cli.parsing import (
+    config_filepath,
+    input_position_dirpaths,
+    output_dirpath,
+    processes_option,
+    transfer_function_dirpath,
+)
+from waveorder.cli.printing import echo_headline, echo_settings
+from waveorder.cli.settings import ReconstructionSettings
+from waveorder.cli.utils import (
+    apply_inverse_to_zyx_and_save,
+    create_empty_hcs_zarr,
+    generate_valid_position_key,
+    is_single_position_store,
+)
+from waveorder.io import utils
+
+
+def _check_background_consistency(
+    background_shape, data_shape, input_channel_names
+):
+    data_cyx_shape = (len(input_channel_names),) + data_shape[3:]
+    if background_shape != data_cyx_shape:
+        raise ValueError(
+            f"Background shape {background_shape} does not match data shape {data_cyx_shape}"
+        )
+
+
+def get_reconstruction_output_metadata(position_path: Path, config_path: Path):
+    # Get non-OME-Zarr plate-level metadata if it's available
+    plate_metadata = {}
+    try:
+        input_plate = open_ome_zarr(
+            position_path.parent.parent.parent, mode="r"
+        )
+        plate_metadata = dict(input_plate.zattrs)
+        plate_metadata.pop("plate")
+    except (RuntimeError, FileNotFoundError):
+        warnings.warn(
+            "Position is not part of a plate...no plate metadata will be copied."
+        )
+
+    # Load the first position to infer dataset information
+    input_dataset = open_ome_zarr(str(position_path), mode="r")
+    T, _, Z, Y, X = input_dataset.data.shape
+
+    settings = utils.yaml_to_model(config_path, ReconstructionSettings)
+
+    # Simplify important settings names
+    recon_biref = settings.birefringence is not None
+    recon_phase = settings.phase is not None
+    recon_fluo = settings.fluorescence is not None
+    recon_dim = settings.reconstruction_dimension
+
+    # Prepare output dataset
+    channel_names = []
+    if recon_biref:
+        channel_names.append("Retardance")
+        channel_names.append("Orientation")
+        channel_names.append("BF")
+        channel_names.append("Pol")
+    if recon_phase:
+        if recon_dim == 2:
+            channel_names.append("Phase2D")
+            # channel_names.append("Absorption2D")
+        elif recon_dim == 3:
+            channel_names.append("Phase3D")
+    if recon_biref and recon_phase:
+        channel_names.append("Retardance_Joint_Decon")
+        channel_names.append("Orientation_Joint_Decon")
+        channel_names.append("Phase_Joint_Decon")
+    if recon_fluo:
+        fluor_name = settings.input_channel_names[0]
+        if recon_dim == 2:
+            channel_names.append(fluor_name + "_Density2D")
+        elif recon_dim == 3:
+            channel_names.append(fluor_name + "_Density3D")
+
+    if recon_dim == 2:
+        output_z_shape = 1
+    elif recon_dim == 3:
+        output_z_shape = input_dataset.data.shape[2]
+
+    return {
+        "shape": (T, len(channel_names), output_z_shape, Y, X),
+        "chunks": (1, 1, 1, Y, X),
+        "scale": input_dataset.scale,
+        "channel_names": channel_names,
+        "dtype": np.float32,
+        "plate_metadata": plate_metadata,
+    }
+
+
+def apply_inverse_transfer_function_single_position(
+    input_position_dirpath: Path,
+    transfer_function_dirpath: Path,
+    config_filepath: Path,
+    output_position_dirpath: Path,
+    num_processes,
+    output_channel_names: list[str],
+) -> None:
+
+    echo_headline("\nStarting reconstruction...")
+
+    # Load datasets
+    transfer_function_dataset = open_ome_zarr(transfer_function_dirpath)
+    input_dataset = open_ome_zarr(input_position_dirpath)
+    output_dataset = open_ome_zarr(output_position_dirpath, mode="r+")
+
+    # Load config file
+    settings = utils.yaml_to_model(config_filepath, ReconstructionSettings)
+
+    # Check input channel names
+    if not set(settings.input_channel_names).issubset(
+        input_dataset.channel_names
+    ):
+        raise ValueError(
+            f"Each of the input_channel_names = {settings.input_channel_names} in {config_filepath} must appear in the dataset {input_position_dirpath} which currently contains channel_names = {input_dataset.channel_names}."
+        )
+
+    # Find input channel indices
+    input_channel_indices = []
+    for input_channel_name in settings.input_channel_names:
+        input_channel_indices.append(
+            input_dataset.channel_names.index(input_channel_name)
+        )
+
+    # Find output channel indices
+    output_channel_indices = []
+    for output_channel_name in output_channel_names:
+        output_channel_indices.append(
+            output_dataset.channel_names.index(output_channel_name)
+        )
+
+    # Find time indices
+    if settings.time_indices == "all":
+        time_indices = range(input_dataset.data.shape[0])
+    elif isinstance(settings.time_indices, list):
+        time_indices = settings.time_indices
+    elif isinstance(settings.time_indices, int):
+        time_indices = [settings.time_indices]
+
+    # Check for invalid times
+    time_ubound = input_dataset.data.shape[0] - 1
+    if np.max(time_indices) > time_ubound:
+        raise ValueError(
+            f"time_indices = {time_indices} includes a time index beyond the maximum index of the dataset = {time_ubound}"
+        )
+
+    # Simplify important settings names
+    recon_biref = settings.birefringence is not None
+    recon_phase = settings.phase is not None
+    recon_fluo = settings.fluorescence is not None
+    recon_dim = settings.reconstruction_dimension
+
+    # Prepare birefringence parameters
+    if settings.birefringence is not None:
+        # settings.birefringence has more parameters than waveorder needs,
+        # so this section converts the settings to a dict and separates the
+        # waveorder parameters (biref_inverse_dict) from the waveorder
+        # parameters (cyx_no_sample_data, and wavelength_illumination)
+        biref_inverse_dict = settings.birefringence.apply_inverse.model_dump()
+
+        # Resolve background path into array
+        background_path = biref_inverse_dict.pop("background_path")
+        if background_path != "":
+            cyx_no_sample_data = utils.load_background(background_path)
+            _check_background_consistency(
+                cyx_no_sample_data.shape,
+                input_dataset.data.shape,
+                settings.input_channel_names,
+            )
+        else:
+            cyx_no_sample_data = None
+
+        # Get illumination wavelength for retardance radians -> nanometers conversion
+        biref_wavelength_illumination = biref_inverse_dict.pop(
+            "wavelength_illumination"
+        )
+
+    # Prepare the apply_inverse_model_function and its arguments
+
+    # [biref only]
+    if recon_biref and (not recon_phase):
+        echo_headline("Reconstructing birefringence with settings:")
+        echo_settings(settings.birefringence)
+
+        # Setup parameters for apply_inverse_to_zyx_and_save
+        apply_inverse_model_function = apply_inverse_models.birefringence
+        apply_inverse_args = {
+            "cyx_no_sample_data": cyx_no_sample_data,
+            "wavelength_illumination": biref_wavelength_illumination,
+            "recon_dim": recon_dim,
+            "biref_inverse_dict": biref_inverse_dict,
+            "transfer_function_dataset": transfer_function_dataset,
+        }
+
+    # [phase only]
+    if recon_phase and (not recon_biref):
+        echo_headline("Reconstructing phase with settings:")
+        echo_settings(settings.phase.apply_inverse)
+
+        # Setup parameters for apply_inverse_to_zyx_and_save
+        apply_inverse_model_function = apply_inverse_models.phase
+        apply_inverse_args = {
+            "recon_dim": recon_dim,
+            "settings_phase": settings.phase,
+            "transfer_function_dataset": transfer_function_dataset,
+        }
+
+    # [biref and phase]
+    if recon_biref and recon_phase:
+        echo_headline("Reconstructing birefringence and phase with settings:")
+        echo_settings(settings.birefringence.apply_inverse)
+        echo_settings(settings.phase.apply_inverse)
+
+        # Setup parameters for apply_inverse_to_zyx_and_save
+        apply_inverse_model_function = (
+            apply_inverse_models.birefringence_and_phase
+        )
+        apply_inverse_args = {
+            "cyx_no_sample_data": cyx_no_sample_data,
+            "wavelength_illumination": biref_wavelength_illumination,
+            "recon_dim": recon_dim,
+            "biref_inverse_dict": biref_inverse_dict,
+            "settings_phase": settings.phase,
+            "transfer_function_dataset": transfer_function_dataset,
+        }
+
+    # [fluo]
+    if recon_fluo:
+        echo_headline("Reconstructing fluorescence with settings:")
+        echo_settings(settings.fluorescence.apply_inverse)
+
+        # Setup parameters for apply_inverse_to_zyx_and_save
+        apply_inverse_model_function = apply_inverse_models.fluorescence
+        apply_inverse_args = {
+            "recon_dim": recon_dim,
+            "settings_fluorescence": settings.fluorescence,
+            "transfer_function_dataset": transfer_function_dataset,
+        }
+
+    # Make the partial function for apply inverse
+    partial_apply_inverse_to_zyx_and_save = partial(
+        apply_inverse_to_zyx_and_save,
+        apply_inverse_model_function,
+        input_dataset,
+        output_position_dirpath,
+        input_channel_indices,
+        output_channel_indices,
+        **apply_inverse_args,
+    )
+
+    # Multiprocessing logic
+    if num_processes > 1:
+        # Loop through T, processing and writing as we go
+        click.echo(
+            f"\nStarting multiprocess pool with {num_processes} processes"
+        )
+        with mp.Pool(num_processes) as p:
+            p.starmap(
+                partial_apply_inverse_to_zyx_and_save,
+                itertools.product(time_indices),
+            )
+    else:
+        for t_idx in time_indices:
+            partial_apply_inverse_to_zyx_and_save(t_idx)
+
+    # Save metadata at position level
+    output_dataset.zattrs["settings"] = settings.model_dump()
+
+    echo_headline(f"Closing {output_position_dirpath}\n")
+
+    output_dataset.close()
+    transfer_function_dataset.close()
+    input_dataset.close()
+
+    echo_headline(
+        f"Recreate this reconstruction with:\n$ waveorder apply-inv-tf {input_position_dirpath} {transfer_function_dirpath} -c {config_filepath} -o {output_position_dirpath}"
+    )
+
+
+def apply_inverse_transfer_function_cli(
+    input_position_dirpaths: list[Path],
+    transfer_function_dirpath: Path,
+    config_filepath: Path,
+    output_dirpath: Path,
+    num_processes,
+) -> None:
+    # Prepare output store
+    output_metadata = get_reconstruction_output_metadata(
+        input_position_dirpaths[0], config_filepath
+    )
+
+    # Generate position keys - use valid HCS keys for single-position stores
+    position_keys = []
+    for i, input_path in enumerate(input_position_dirpaths):
+        if is_single_position_store(input_path):
+            position_key = generate_valid_position_key(i)
+        else:
+            # Use original HCS plate structure
+            position_key = input_path.parts[-3:]
+        position_keys.append(position_key)
+
+    create_empty_hcs_zarr(
+        store_path=output_dirpath,
+        position_keys=position_keys,
+        **output_metadata,
+    )
+
+    # Initialize torch threads
+    if num_processes > 1:
+        torch.set_num_threads(1)
+        torch.set_num_interop_threads(1)
+
+    # Loop through positions
+    for i, input_position_dirpath in enumerate(input_position_dirpaths):
+        # Use the same position key generation logic
+        if is_single_position_store(input_position_dirpath):
+            position_key = generate_valid_position_key(i)
+        else:
+            position_key = input_position_dirpath.parts[-3:]
+
+        output_position_path = output_dirpath / Path(*position_key)
+
+        apply_inverse_transfer_function_single_position(
+            input_position_dirpath,
+            transfer_function_dirpath,
+            config_filepath,
+            output_position_path,
+            num_processes,
+            output_metadata["channel_names"],
+        )
+
+
+@click.command("apply-inv-tf")
+@input_position_dirpaths()
+@transfer_function_dirpath()
+@config_filepath()
+@output_dirpath()
+@processes_option(default=1)
+def _apply_inverse_transfer_function_cli(
+    input_position_dirpaths: list[Path],
+    transfer_function_dirpath: Path,
+    config_filepath: Path,
+    output_dirpath: Path,
+    num_processes,
+) -> None:
+    """
+    Apply an inverse transfer function to a dataset using a configuration file.
+
+    Applies a transfer function to all positions in the list `input-position-dirpaths`,
+    so all positions must have the same TCZYX shape.
+
+    Appends channels to ./output.zarr, so multiple reconstructions can fill a single store.
+
+    See /examples for example configuration files.
+
+    >> waveorder apply-inv-tf -i ./input.zarr/*/*/* -t ./transfer-function.zarr -c /examples/birefringence.yml -o ./output.zarr
+    """
+    apply_inverse_transfer_function_cli(
+        input_position_dirpaths,
+        transfer_function_dirpath,
+        config_filepath,
+        output_dirpath,
+        num_processes,
+    )