Fix 3D to 2D phase reconstruction

docs/examples/configs/birefringence-and-phase.yml

@@ -20,6 +20,7 @@ phase:
     yx_pixel_size: 0.325
     z_pixel_size: 2.0
     z_padding: 0
+    z_focus_offset: 0
     index_of_refraction_media: 1.3
     numerical_aperture_detection: 1.2
     numerical_aperture_illumination: 0.5

docs/examples/configs/fluorescence.yml

@@ -7,6 +7,7 @@ fluorescence:
     yx_pixel_size: 0.325
     z_pixel_size: 2.0
     z_padding: 0
+    z_focus_offset: 0
     index_of_refraction_media: 1.3
     numerical_aperture_detection: 1.2
     wavelength_emission: 0.507

docs/examples/configs/phase.yml

@@ -8,6 +8,7 @@ phase:
     yx_pixel_size: 0.325
     z_pixel_size: 2.0
     z_padding: 0
+    z_focus_offset: 0
     index_of_refraction_media: 1.3
     numerical_aperture_detection: 1.2
     numerical_aperture_illumination: 0.5

docs/examples/models/isotropic_thin_3d.py

@@ -19,16 +19,25 @@
 phantom_arguments = {"index_of_refraction_sample": 1.33, "sphere_radius": 5}
 z_shape = 100
 z_pixel_size = 0.25
+zyx_scale = np.array(
+    [
+        z_pixel_size,
+        simulation_arguments["yx_pixel_size"],
+        simulation_arguments["yx_pixel_size"],
+    ]
+)
 transfer_function_arguments = {
     "z_position_list": (np.arange(z_shape) - z_shape // 2) * z_pixel_size,
     "numerical_aperture_illumination": 0.9,
     "numerical_aperture_detection": 1.2,
 }
 
-# Create a phantom
+# Create a disk phantom
 yx_absorption, yx_phase = isotropic_thin_3d.generate_test_phantom(
     **simulation_arguments, **phantom_arguments
 )
+yx_absorption[:, 128:] = 0  # half absorbing
+yx_phase[128:] = 0  # half phase
 
 # Calculate transfer function
 (
@@ -38,6 +47,12 @@
     **simulation_arguments, **transfer_function_arguments
 )
 
+# Calculate singular system
+singular_system = isotropic_thin_3d.calculate_singular_system(
+    absorption_2d_to_3d_transfer_function,
+    phase_2d_to_3d_transfer_function,
+)
+
 # Display transfer function
 viewer = napari.Viewer()
 zyx_scale = np.array(
@@ -70,8 +85,8 @@
     yx_phase_recon,
 ) = isotropic_thin_3d.apply_inverse_transfer_function(
     zyx_data,
-    absorption_2d_to_3d_transfer_function,
-    phase_2d_to_3d_transfer_function,
+    singular_system,
+    regularization_strength=1e-2,
 )
 
 # Display
@@ -84,5 +99,10 @@
 ]
 
 for array in arrays:
-    viewer.add_image(array[0].cpu().numpy(), name=array[1])
+    scale = zyx_scale[1:] if array[0].ndim == 2 else zyx_scale
+    viewer.add_image(array[0].cpu().numpy(), name=array[1], scale=scale)
+
+viewer.grid.enabled = True
+viewer.dims.current_step = (z_shape // 2, 0, 0)
+
 input("Showing object, data, and recon. Press <enter> to quit...")

tests/cli_tests/test_compute_tf.py

@@ -1,7 +1,10 @@
+import numpy as np
+import pytest
 from click.testing import CliRunner
 
 from waveorder.cli import settings
 from waveorder.cli.compute_transfer_function import (
+    _position_list_from_shape_scale_offset,
     generate_and_save_birefringence_transfer_function,
     generate_and_save_fluorescence_transfer_function,
     generate_and_save_phase_transfer_function,
@@ -10,6 +13,18 @@
 from waveorder.io import utils
 
 
+@pytest.mark.parametrize(
+    "shape, scale, offset, expected",
+    [
+        (5, 1.0, 0.0, [2.0, 1.0, 0.0, -1.0, -2.0]),
+        (4, 0.5, 1.0, [1.5, 1.0, 0.5, 0.0]),
+    ],
+)
+def test_position_list_from_shape_scale_offset(shape, scale, offset, expected):
+    result = _position_list_from_shape_scale_offset(shape, scale, offset)
+    np.testing.assert_allclose(result, expected)
+
+
 def test_compute_transfer(tmp_path, example_plate):
     recon_settings = settings.ReconstructionSettings(
         input_channel_names=[f"State{i}" for i in range(4)],
@@ -116,9 +131,10 @@ def test_phase_3dim_write(birefringence_phase_recon_settings_function):
     settings, dataset = birefringence_phase_recon_settings_function
     settings.reconstruction_dimension = 2
     generate_and_save_phase_transfer_function(settings, dataset, (3, 4, 5))
-    assert dataset["absorption_transfer_function"]
-    assert dataset["phase_transfer_function"]
-    assert dataset["phase_transfer_function"].shape == (1, 1, 3, 4, 5)
+    assert dataset["singular_system_U"]
+    assert dataset["singular_system_U"].shape == (1, 2, 2, 4, 5)
+    assert dataset["singular_system_S"]
+    assert dataset["singular_system_Vh"]
     assert "real_potential_transfer_function" not in dataset
     assert "imaginary_potential_transfer_function" not in dataset
 

waveorder/cli/apply_inverse_models.py

@@ -65,23 +65,27 @@ def phase(
     # [phase only, 2]
     if recon_dim == 2:
         # Load transfer functions
-        absorption_transfer_function = torch.tensor(
-            transfer_function_dataset["absorption_transfer_function"][0, 0]
+        U = torch.from_numpy(transfer_function_dataset["singular_system_U"][0])
+        S = torch.from_numpy(
+            transfer_function_dataset["singular_system_S"][0, 0]
         )
-        phase_transfer_function = torch.tensor(
-            transfer_function_dataset["phase_transfer_function"][0, 0]
+        Vh = torch.from_numpy(
+            transfer_function_dataset["singular_system_Vh"][0]
         )
 
         # Apply
         (
-            _,
-            output,
+            absorption_yx,
+            phase_yx,
         ) = isotropic_thin_3d.apply_inverse_transfer_function(
             czyx_data[0],
-            absorption_transfer_function,
-            phase_transfer_function,
+            (U, S, Vh),
             **settings_phase.apply_inverse.dict(),
         )
+        # 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:
@@ -127,12 +131,13 @@ def birefringence_and_phase(
 
     # [biref and phase, 2]
     if recon_dim == 2:
-        # Load phase transfer functions
-        absorption_transfer_function = torch.tensor(
-            transfer_function_dataset["absorption_transfer_function"][0, 0]
+        # 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]
         )
-        phase_transfer_function = torch.tensor(
-            transfer_function_dataset["phase_transfer_function"][0, 0]
+        Vh = torch.from_numpy(
+            transfer_function_dataset["singular_system_Vh"][0]
         )
 
         # Apply
@@ -163,8 +168,7 @@ def birefringence_and_phase(
             yx_phase,
         ) = isotropic_thin_3d.apply_inverse_transfer_function(
             brightfield_3d,
-            absorption_transfer_function,
-            phase_transfer_function,
+            (U, S, Vh),
             **settings_phase.apply_inverse.dict(),
         )
 

waveorder/cli/apply_inverse_transfer_function.py

@@ -79,6 +79,7 @@ def get_reconstruction_output_metadata(position_path: Path, config_path: Path):
     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_fluo:

waveorder/cli/compute_transfer_function.py

@@ -4,6 +4,7 @@
 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,
@@ -20,6 +21,22 @@
 )
 
 
+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_birefringence_transfer_function(settings, dataset):
     """Generates and saves the birefringence transfer function to the dataset, based on the settings.
 
@@ -61,18 +78,22 @@ def generate_and_save_phase_transfer_function(
     echo_headline("Generating phase transfer function with settings:")
     echo_settings(settings.phase.transfer_function)
 
+    settings_dict = settings.phase.transfer_function.dict()
     if settings.reconstruction_dimension == 2:
         # Convert zyx_shape and z_pixel_size into yx_shape and z_position_list
-        settings_dict = settings.phase.transfer_function.dict()
         settings_dict["yx_shape"] = [zyx_shape[1], zyx_shape[2]]
-        settings_dict["z_position_list"] = list(
-            -(np.arange(zyx_shape[0]) - zyx_shape[0] // 2)
-            * settings_dict["z_pixel_size"]
+        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
         (
@@ -82,26 +103,36 @@ def generate_and_save_phase_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(
-            "absorption_transfer_function",
-            absorption_transfer_function.cpu().numpy()[None, None, ...],
-            chunks=(1, 1, 1, zyx_shape[1], zyx_shape[2]),
+            "singular_system_U",
+            U.cpu().numpy()[None],
         )
         dataset.create_image(
-            "phase_transfer_function",
-            phase_transfer_function.cpu().numpy()[None, None, ...],
-            chunks=(1, 1, 1, zyx_shape[1], zyx_shape[2]),
+            "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.phase.transfer_function.dict(),
+            **settings_dict,
         )
         # Save
         dataset.create_image(
@@ -133,6 +164,9 @@ def generate_and_save_fluorescence_transfer_function(
     """
     echo_headline("Generating fluorescence transfer function with settings:")
     echo_settings(settings.fluorescence.transfer_function)
+    # Remove unused parameters
+    settings_dict = settings.fluorescence.transfer_function.dict()
+    settings_dict.pop("z_focus_offset")
 
     if settings.reconstruction_dimension == 2:
         raise NotImplementedError
@@ -141,7 +175,7 @@ def generate_and_save_fluorescence_transfer_function(
         optical_transfer_function = (
             isotropic_fluorescent_thick_3d.calculate_transfer_function(
                 zyx_shape=zyx_shape,
-                **settings.fluorescence.transfer_function.dict(),
+                **settings_dict,
             )
         )
         # Save
@@ -182,9 +216,40 @@ def compute_transfer_function_cli(
             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=["None"]
+        output_dirpath,
+        layout="fov",
+        mode="w",
+        channel_names=num_channels * ["None"],
     )
 
     # Pass settings to appropriate calculate_transfer_function and save

waveorder/cli/settings.py

@@ -61,6 +61,7 @@ class FourierTransferFunctionSettings(MyBaseModel):
     yx_pixel_size: PositiveFloat = 6.5 / 20
     z_pixel_size: PositiveFloat = 2.0
     z_padding: NonNegativeInt = 0
+    z_focus_offset: Union[int, Literal["auto"]] = 0
     index_of_refraction_media: PositiveFloat = 1.3
     numerical_aperture_detection: PositiveFloat = 1.2