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

waveorder/cli/compute_transfer_function.py

@@ -0,0 +1,432 @@
+from pathlib import Path
+
+import click
+import numpy as np
+from iohub.ngff import Position, open_ome_zarr
+
+from waveorder import focus
+from waveorder.cli.parsing import (
+    config_filepath,
+    input_position_dirpaths,
+    output_dirpath,
+)
+from waveorder.cli.printing import echo_headline, echo_settings
+from waveorder.cli.settings import ReconstructionSettings
+from waveorder.io import utils
+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,
+)
+
+
+def _position_list_from_shape_scale_offset(
+    shape: int, scale: float, offset: float
+) -> list:
+    """
+    Generates a list of positions based on the given array shape, pixel size (scale), and offset.
+
+    Examples
+    --------
+    >>> _position_list_from_shape_scale_offset(5, 1.0, 0.0)
+    [2.0, 1.0, 0.0, -1.0, -2.0]
+    >>> _position_list_from_shape_scale_offset(4, 0.5, 1.0)
+    [1.5, 1.0, 0.5, 0.0]
+    """
+    return list((-np.arange(shape) + (shape // 2) + offset) * scale)
+
+
+def generate_and_save_vector_birefringence_transfer_function(
+    settings: ReconstructionSettings, dataset: Position, zyx_shape: tuple
+):
+    """Generates and saves the vector birefringence transfer function
+    to the dataset, based on the settings.
+
+    Parameters
+    ----------
+    settings : ReconstructionSettings
+    dataset : NGFF Node
+        The dataset that will be updated.
+    zyx_shape : tuple
+        A tuple of integers specifying the input data's shape in (Z, Y, X) order
+    """
+    echo_headline(
+        "Generating vector birefringence transfer function with settings:"
+    )
+    echo_settings(settings.birefringence.transfer_function)
+    echo_settings(settings.phase.transfer_function)
+
+    num_elements = np.array(zyx_shape).prod()
+    max_tf_elements = 1e7  # empirical, based on memory usage
+    transverse_downsample_factor = np.ceil(
+        np.sqrt(num_elements / max_tf_elements)
+    )
+    echo_headline(
+        f"Downsampling transfer function in X and Y by {transverse_downsample_factor}x"
+    )
+    phase_settings_dict = settings.phase.transfer_function.model_dump()
+    phase_settings_dict.pop("z_focus_offset")  # not used in 3D
+
+    sfZYX_transfer_function, _, singular_system = (
+        inplane_oriented_thick_pol3d_vector.calculate_transfer_function(
+            zyx_shape=zyx_shape,
+            scheme=str(len(settings.input_channel_names)) + "-State",
+            **settings.birefringence.transfer_function.model_dump(),
+            **phase_settings_dict,
+            fourier_oversample_factor=int(transverse_downsample_factor),
+        )
+    )
+
+    U, S, Vh = singular_system
+    chunks = (1, 1, 1, zyx_shape[1], zyx_shape[2])
+
+    # Add dummy channels
+    for i in range(3):
+        dataset.append_channel(f"ch{i}")
+
+    dataset.create_image(
+        "vector_transfer_function",
+        sfZYX_transfer_function.cpu().numpy(),
+        chunks=chunks,
+    )
+    dataset.create_image(
+        "vector_singular_system_U",
+        U.cpu().numpy(),
+        chunks=chunks,
+    )
+    dataset.create_image(
+        "vector_singular_system_S",
+        S[None].cpu().numpy(),
+        chunks=chunks,
+    )
+    dataset.create_image(
+        "vector_singular_system_Vh",
+        Vh.cpu().numpy(),
+        chunks=chunks,
+    )
+
+
+def generate_and_save_birefringence_transfer_function(
+    settings: ReconstructionSettings, dataset
+):
+    """Generates and saves the birefringence transfer function to the dataset, based on the settings.
+
+    Parameters
+    ----------
+    settings: ReconstructionSettings
+    dataset: NGFF Node
+        The dataset that will be updated.
+    """
+    echo_headline("Generating birefringence transfer function with settings:")
+    echo_settings(settings.birefringence.transfer_function)
+
+    # Calculate transfer functions
+    intensity_to_stokes_matrix = (
+        inplane_oriented_thick_pol3d.calculate_transfer_function(
+            scheme=str(len(settings.input_channel_names)) + "-State",
+            **settings.birefringence.transfer_function.model_dump(),
+        )
+    )
+    # Save
+    dataset["intensity_to_stokes_matrix"] = (
+        intensity_to_stokes_matrix.cpu().numpy()[None, None, None, ...]
+    )
+
+
+def generate_and_save_phase_transfer_function(
+    settings: ReconstructionSettings,
+    dataset: Position,
+    zyx_shape: tuple,
+):
+    """Generates and saves the phase transfer function to the dataset, based on the settings.
+
+    Parameters
+    ----------
+    settings: ReconstructionSettings
+    dataset: Position
+        The dataset that will be updated.
+    zyx_shape : tuple
+        A tuple of integers specifying the input data's shape in (Z, Y, X) order
+    """
+    echo_headline("Generating phase transfer function with settings:")
+    echo_settings(settings.phase.transfer_function)
+
+    settings_dict = settings.phase.transfer_function.model_dump()
+    if settings.reconstruction_dimension == 2:
+        # Convert zyx_shape and z_pixel_size into yx_shape and z_position_list
+        settings_dict["yx_shape"] = [zyx_shape[1], zyx_shape[2]]
+        settings_dict["z_position_list"] = (
+            _position_list_from_shape_scale_offset(
+                shape=zyx_shape[0],
+                scale=settings_dict["z_pixel_size"],
+                offset=settings_dict["z_focus_offset"],
+            )
+        )
+
+        # Remove unused parameters
+        settings_dict.pop("z_pixel_size")
+        settings_dict.pop("z_padding")
+        settings_dict.pop("z_focus_offset")
+
+        # Calculate transfer functions
+        (
+            absorption_transfer_function,
+            phase_transfer_function,
+        ) = isotropic_thin_3d.calculate_transfer_function(
+            **settings_dict,
+        )
+
+        # Calculate singular system
+        U, S, Vh = isotropic_thin_3d.calculate_singular_system(
+            absorption_transfer_function,
+            phase_transfer_function,
+        )
+
+        # Save
+        dataset.create_image(
+            "singular_system_U",
+            U.cpu().numpy()[None],
+        )
+        dataset.create_image(
+            "singular_system_S",
+            S.cpu().numpy()[None, None],
+        )
+        dataset.create_image(
+            "singular_system_Vh",
+            Vh.cpu().numpy()[None],
+        )
+
+    elif settings.reconstruction_dimension == 3:
+        settings_dict.pop("z_focus_offset")  # not used in 3D
+
+        # Calculate transfer functions
+        (
+            real_potential_transfer_function,
+            imaginary_potential_transfer_function,
+        ) = phase_thick_3d.calculate_transfer_function(
+            zyx_shape=zyx_shape,
+            **settings_dict,
+        )
+        # Save
+        dataset.create_image(
+            "real_potential_transfer_function",
+            real_potential_transfer_function.cpu().numpy()[None, None, ...],
+            chunks=(1, 1, 1, zyx_shape[1], zyx_shape[2]),
+        )
+        dataset.create_image(
+            "imaginary_potential_transfer_function",
+            imaginary_potential_transfer_function.cpu().numpy()[
+                None, None, ...
+            ],
+            chunks=(1, 1, 1, zyx_shape[1], zyx_shape[2]),
+        )
+
+
+def generate_and_save_fluorescence_transfer_function(
+    settings: ReconstructionSettings,
+    dataset: Position,
+    zyx_shape: tuple,
+):
+    """Generates and saves the fluorescence transfer function to the dataset, based on the settings.
+
+    Parameters
+    ----------
+    settings: ReconstructionSettings
+    dataset: Position
+        The dataset that will be updated.
+    zyx_shape : tuple
+        A tuple of integers specifying the input data's shape in (Z, Y, X) order
+    """
+    echo_headline("Generating fluorescence transfer function with settings:")
+    echo_settings(settings.fluorescence.transfer_function)
+    settings_dict = settings.fluorescence.transfer_function.model_dump()
+
+    if settings.reconstruction_dimension == 2:
+        # Convert zyx_shape and z_pixel_size into yx_shape and z_position_list
+        settings_dict["yx_shape"] = [zyx_shape[1], zyx_shape[2]]
+        settings_dict["z_position_list"] = (
+            _position_list_from_shape_scale_offset(
+                shape=zyx_shape[0],
+                scale=settings_dict["z_pixel_size"],
+                offset=settings_dict["z_focus_offset"],
+            )
+        )
+
+        # Remove unused parameters
+        settings_dict.pop("z_pixel_size")
+        settings_dict.pop("z_padding")
+        settings_dict.pop("z_focus_offset")
+
+        # Calculate 2D fluorescence transfer functions
+        fluorescent_2d_to_3d_transfer_function = (
+            isotropic_fluorescent_thin_3d.calculate_transfer_function(
+                **settings_dict,
+            )
+        )
+
+        # Calculate singular system for 2D reconstruction
+        U, S, Vh = isotropic_fluorescent_thin_3d.calculate_singular_system(
+            fluorescent_2d_to_3d_transfer_function
+        )
+
+        # Get yx_shape for chunk sizes
+        yx_shape = zyx_shape[1:]
+
+        # Save singular system components
+        dataset.create_image(
+            "singular_system_U",
+            U.cpu().numpy()[None, ...],
+            chunks=(1, 1, 1, yx_shape[0], yx_shape[1]),
+        )
+        dataset.create_image(
+            "singular_system_S",
+            S.cpu().numpy()[None, None, ...],
+            chunks=(1, 1, 1, yx_shape[0], yx_shape[1]),
+        )
+        dataset.create_image(
+            "singular_system_Vh",
+            Vh.cpu().numpy()[None, ...],
+            chunks=(1, 1, zyx_shape[0], yx_shape[0], yx_shape[1]),
+        )
+
+    elif settings.reconstruction_dimension == 3:
+        # Remove unused parameters for 3D
+        settings_dict.pop("z_focus_offset")
+
+        # Calculate transfer functions
+        optical_transfer_function = (
+            isotropic_fluorescent_thick_3d.calculate_transfer_function(
+                zyx_shape=zyx_shape,
+                **settings_dict,
+            )
+        )
+        # Save
+        dataset.create_image(
+            "optical_transfer_function",
+            optical_transfer_function.cpu().numpy()[None, None, ...],
+            chunks=(1, 1, 1, zyx_shape[1], zyx_shape[2]),
+        )
+
+
+def compute_transfer_function_cli(
+    input_position_dirpath: Path,
+    config_filepath: Path,
+    output_dirpath: Path,
+) -> None:
+    """CLI command to compute the transfer function given a configuration file path
+    and a desired output path.
+    """
+
+    # Load config file
+    settings = utils.yaml_to_model(config_filepath, ReconstructionSettings)
+
+    echo_headline(
+        f"Generating transfer functions and storing in {output_dirpath}\n"
+    )
+
+    # Read shape from input dataset
+    input_dataset = open_ome_zarr(
+        input_position_dirpath, layout="fov", mode="r"
+    )
+    zyx_shape = input_dataset.data.shape[
+        2:
+    ]  # only loads a single position "0"
+
+    # 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_dirpaths[0]} which currently contains channel_names = {input_dataset.channel_names}."
+        )
+
+    # Find in-focus slices for 2D reconstruction in "auto" mode
+    if (
+        settings.phase is not None
+        and settings.reconstruction_dimension == 2
+        and settings.phase.transfer_function.z_focus_offset == "auto"
+    ):
+
+        c_idx = input_dataset.get_channel_index(
+            settings.input_channel_names[0]
+        )
+        zyx_array = input_dataset["0"][0, c_idx]
+
+        in_focus_index = focus.focus_from_transverse_band(
+            zyx_array,
+            NA_det=settings.phase.transfer_function.numerical_aperture_detection,
+            lambda_ill=settings.phase.transfer_function.wavelength_illumination,
+            pixel_size=settings.phase.transfer_function.yx_pixel_size,
+            mode="min",
+            polynomial_fit_order=4,
+        )
+
+        z_focus_offset = in_focus_index - (zyx_shape[0] // 2)
+        settings.phase.transfer_function.z_focus_offset = z_focus_offset
+        print("Found z_focus_offset:", z_focus_offset)
+
+    # Prepare output dataset
+    num_channels = (
+        2 if settings.reconstruction_dimension == 2 else 1
+    )  # space for SVD
+    output_dataset = open_ome_zarr(
+        output_dirpath,
+        layout="fov",
+        mode="w",
+        channel_names=num_channels * ["None"],
+    )
+
+    # Pass settings to appropriate calculate_transfer_function and save
+    if settings.birefringence is not None:
+        generate_and_save_birefringence_transfer_function(
+            settings, output_dataset
+        )
+    if settings.phase is not None:
+        generate_and_save_phase_transfer_function(
+            settings, output_dataset, zyx_shape
+        )
+    if settings.fluorescence is not None:
+        generate_and_save_fluorescence_transfer_function(
+            settings, output_dataset, zyx_shape
+        )
+    if settings.birefringence is not None and settings.phase is not None:
+        generate_and_save_vector_birefringence_transfer_function(
+            settings, output_dataset, zyx_shape
+        )
+
+    # Write settings to metadata
+    output_dataset.zattrs["settings"] = settings.model_dump()
+
+    echo_headline(f"Closing {output_dirpath}\n")
+    output_dataset.close()
+
+    echo_headline(
+        f"Recreate this transfer function with:\n$ waveorder compute-tf {input_position_dirpaths} -c {config_filepath} -o {output_dirpath}"
+    )
+
+
+@click.command("compute-tf")
+@input_position_dirpaths()
+@config_filepath()
+@output_dirpath()
+def _compute_transfer_function_cli(
+    input_position_dirpaths: list[Path],
+    config_filepath: Path,
+    output_dirpath: Path,
+) -> None:
+    """
+    Compute a transfer function using a dataset and configuration file.
+
+    Calculates the transfer function based on the shape of the first position
+    in the list `input-position-dirpaths`.
+
+    See /examples for example configuration files.
+
+    >> waveorder compute-tf -i ./input.zarr/0/0/0 -c ./examples/birefringence.yml -o ./transfer_function.zarr
+    """
+    compute_transfer_function_cli(
+        input_position_dirpaths[0], config_filepath, output_dirpath
+    )

waveorder/cli/gui_widget.py

@@ -0,0 +1,58 @@
+import sys
+
+import click
+
+try:
+    from waveorder.plugin import tab_recon
+except:
+    pass
+
+try:
+    from qtpy.QtWidgets import QApplication, QStyle, QVBoxLayout, QWidget
+except:
+    pass
+
+try:
+    import qdarktheme  # pip install pyqtdarktheme==2.1.0 --ignore-requires-python
+except:
+    pass
+
+PLUGIN_NAME = "waveorder: Computational Toolkit for Label-Free Imaging"
+PLUGIN_ICON = "🔬"
+
+
+@click.command()
+def gui():
+    """GUI for waveorder: Computational Toolkit for Label-Free Imaging"""
+
+    app = QApplication(sys.argv)
+    app.setStyle(
+        "Fusion"
+    )  # Other options: "Fusion", "Windows", "macOS", "WindowsVista"
+    try:
+        qdarktheme.setup_theme("dark")
+    except Exception as e:
+        print(e.args)
+        pass
+    window = MainWindow()
+    window.setWindowTitle(PLUGIN_ICON + " " + PLUGIN_NAME + " " + PLUGIN_ICON)
+
+    pixmapi = getattr(QStyle.StandardPixmap, "SP_TitleBarMenuButton")
+    icon = app.style().standardIcon(pixmapi)
+    window.setWindowIcon(icon)
+
+    window.show()
+    sys.exit(app.exec())
+
+
+class MainWindow(QWidget):
+    def __init__(self):
+        super().__init__()
+        recon_tab = tab_recon.Ui_ReconTab_Form(stand_alone=True)
+        layout = QVBoxLayout()
+        self.setLayout(layout)
+        layout.addWidget(recon_tab.recon_tab_mainScrollArea)
+
+
+if __name__ == "__main__":
+    gui()

waveorder/cli/main.py

@@ -0,0 +1,39 @@
+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.reconstruct import _reconstruct_cli
+
+try:
+    from waveorder.cli.gui_widget import gui
+except:
+    pass
+
+CONTEXT = {"help_option_names": ["-h", "--help"]}
+
+
+# `waveorder -h` will show subcommands in the order they are added
+class NaturalOrderGroup(click.Group):
+    def list_commands(self, ctx):
+        return self.commands.keys()
+
+
+@click.group(context_settings=CONTEXT, cls=NaturalOrderGroup)
+def cli():
+    """\033[92mwaveorder: Computational Toolkit for Label-Free Imaging\033[0m\n"""
+
+
+cli.add_command(_reconstruct_cli)
+cli.add_command(_compute_transfer_function_cli)
+cli.add_command(_apply_inverse_transfer_function_cli)
+try:
+    cli.add_command(gui)
+except:
+    pass
+
+if __name__ == "__main__":
+    cli()

waveorder/cli/monitor.py

@@ -0,0 +1,163 @@
+import shutil
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import submitit
+
+
+def _move_cursor_up(n_lines, do_print=True):
+    if do_print:
+        sys.stdout.write("\033[F" * n_lines)
+
+
+def _print_status(
+    jobs, position_dirpaths, elapsed_list, print_indices=None, do_print=True
+):
+
+    columns = [15, 30, 40, 50]
+
+    # header
+    if do_print:
+        sys.stdout.write(
+            "\033[K"  # clear line
+            "\033[96mID"  # cyan
+            f"\033[{columns[0]}G WELL "
+            f"\033[{columns[1]}G STATUS "
+            f"\033[{columns[2]}G NODE "
+            f"\033[{columns[2]}G ELAPSED\n"
+        )
+
+    if print_indices is None:
+        print_indices = range(len(jobs))
+
+    complete_count = 0
+
+    for i, (job, position_dirpath) in enumerate(zip(jobs, position_dirpaths)):
+        try:
+            node_name = job.get_info()["NodeList"]  # slowest, so do this first
+        except:
+            node_name = "SUBMITTED"
+
+        if job.state == "COMPLETED":
+            color = "\033[32m"  # green
+            complete_count += 1
+        elif job.state == "RUNNING":
+            color = "\033[93m"  # yellow
+            elapsed_list[i] += 1  # inexact timing
+        else:
+            color = "\033[91m"  # red
+
+        if i in print_indices:
+            if do_print:
+                sys.stdout.write(
+                    f"\033[K"  # clear line
+                    f"{color}{job.job_id}"
+                    f"\033[{columns[0]}G {'/'.join(position_dirpath.parts[-3:])}"
+                    f"\033[{columns[1]}G {job.state}"
+                    f"\033[{columns[2]}G {node_name}"
+                    f"\033[{columns[3]}G {elapsed_list[i]} s\n"
+                )
+    sys.stdout.flush()
+    if do_print:
+        print(
+            f"\033[32m{complete_count}/{len(jobs)} jobs complete. "
+            "<ctrl+z> to move monitor to background. "
+            "<ctrl+c> twice to cancel jobs."
+        )
+
+    return elapsed_list
+
+
+def _get_jobs_to_print(jobs, num_to_print):
+    job_indices_to_print = []
+
+    # if number of jobs is smaller than termanal size, print all
+    if len(jobs) <= num_to_print:
+        return list(range(len(jobs)))
+
+    # prioritize incomplete jobs
+    for i, job in enumerate(jobs):
+        if not job.done():
+            job_indices_to_print.append(i)
+        if len(job_indices_to_print) == num_to_print:
+            return job_indices_to_print
+
+    # fill in the rest with complete jobs
+    for i, job in enumerate(jobs):
+        job_indices_to_print.append(i)
+        if len(job_indices_to_print) == num_to_print:
+            return job_indices_to_print
+
+    # shouldn't reach here
+    return job_indices_to_print
+
+
+def monitor_jobs(
+    jobs: list[submitit.Job], position_dirpaths: list[Path], do_print=True
+):
+    """Displays the status of a list of submitit jobs with corresponding paths.
+
+    Parameters
+    ----------
+    jobs : list[submitit.Job]
+        List of submitit jobs
+    position_dirpaths : list[Path]
+        List of corresponding position paths
+    """
+    NON_JOB_LINES = 3
+
+    if not len(jobs) == len(position_dirpaths):
+        raise ValueError(
+            "The number of jobs and position_dirpaths should be the same."
+        )
+
+    elapsed_list = [0] * len(jobs)  # timer for each job
+
+    # print all jobs once if terminal is too small
+    if shutil.get_terminal_size().lines - NON_JOB_LINES < len(jobs):
+        _print_status(jobs, position_dirpaths, elapsed_list, do_print=do_print)
+
+    # main monitor loop
+    try:
+        while not all(job.done() for job in jobs):
+            terminal_lines = shutil.get_terminal_size().lines
+            num_jobs_to_print = np.min(
+                [terminal_lines - NON_JOB_LINES, len(jobs)]
+            )
+
+            job_indices_to_print = _get_jobs_to_print(jobs, num_jobs_to_print)
+
+            elapsed_list = _print_status(
+                jobs,
+                position_dirpaths,
+                elapsed_list,
+                job_indices_to_print,
+                do_print,
+            )
+
+            time.sleep(1)
+            _move_cursor_up(num_jobs_to_print + 2, do_print)
+
+        # Print final status
+        time.sleep(1)
+        _print_status(jobs, position_dirpaths, elapsed_list, do_print=do_print)
+
+    # cancel jobs if ctrl+c
+    except KeyboardInterrupt:
+        for job in jobs:
+            job.cancel()
+        print("All jobs cancelled.\033[97m")
+
+    # Print STDOUT and STDERR for first incomplete job
+    incomplete_count = 0
+    for job in jobs:
+        if not job.done():
+            if incomplete_count == 0:
+                print("\033[32mSTDOUT")
+                print(job.stdout())
+                print("\033[91mSTDERR")
+                print(job.stderr())
+
+    print("\033[97m")  # print white