Oblique illumination reconstructions

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

@@ -25,6 +25,7 @@ phase:
     numerical_aperture_detection: 1.2
     numerical_aperture_illumination: 0.5
     invert_phase_contrast: false
+    illumination_sector_angles: null
   apply_inverse:
     reconstruction_algorithm: Tikhonov
     regularization_strength: 0.001

docs/examples/configs/phase.yml

@@ -13,6 +13,7 @@ phase:
     numerical_aperture_detection: 1.2
     numerical_aperture_illumination: 0.5
     invert_phase_contrast: false
+    illumination_sector_angles: null
   apply_inverse:
     reconstruction_algorithm: Tikhonov
     regularization_strength: 0.001

docs/examples/models/phase_thick_3d.py

@@ -34,15 +34,15 @@
     **simulation_arguments, **phantom_arguments
 )
 
-# Calculate transfer function
+# Calculate transfer function (returns shape (C, Z, Y, X))
 (
     real_potential_transfer_function,
     imag_potential_transfer_function,
 ) = phase_thick_3d.calculate_transfer_function(
     **simulation_arguments, **transfer_function_arguments
 )
 
-# Display transfer function
+# Display transfer function (extract single channel for visualization)
 viewer = napari.Viewer()
 zyx_scale = np.array(
     [
@@ -53,25 +53,25 @@
 )
 phase_thick_3d.visualize_transfer_function(
     viewer,
-    real_potential_transfer_function,
-    imag_potential_transfer_function,
+    real_potential_transfer_function[0],
+    imag_potential_transfer_function[0],
     zyx_scale,
 )
 input("Showing OTFs. Press <enter> to continue...")
 viewer.layers.select_all()
 viewer.layers.remove_selected()
 
-# Simulate
+# Simulate (extract single channel for forward model)
 zyx_data = phase_thick_3d.apply_transfer_function(
     zyx_phase,
-    real_potential_transfer_function,
+    real_potential_transfer_function[0],
     transfer_function_arguments["z_padding"],
     brightness=1e3,
 )
 
-# Reconstruct
+# Reconstruct (wrap data in channel dimension for inverse)
 zyx_recon = phase_thick_3d.apply_inverse_transfer_function(
-    zyx_data,
+    zyx_data[None],  # Add channel dimension: (Z, Y, X) -> (1, Z, Y, X)
     real_potential_transfer_function,
     imag_potential_transfer_function,
     transfer_function_arguments["z_padding"],

docs/examples/models/phase_thick_3d_sector_illumination.py

@@ -0,0 +1,182 @@
+"""
+phase thick 3d with sector illumination
+========================================
+
+# 3D phase reconstruction with oblique sector illumination
+# This example demonstrates multi-channel phase reconstruction where each channel
+# corresponds to a different illumination sector angle.
+"""
+
+import napari
+import numpy as np
+import torch
+
+from waveorder.models import phase_thick_3d
+
+# Parameters
+# all lengths must use consistent units e.g. um
+simulation_arguments = {
+    "zyx_shape": (100, 256, 256),
+    "yx_pixel_size": 6.5 / 63,
+    "z_pixel_size": 0.25,
+    "index_of_refraction_media": 1.3,
+}
+phantom_arguments = {"index_of_refraction_sample": 1.50, "sphere_radius": 5}
+transfer_function_arguments = {
+    "z_padding": 0,
+    "wavelength_illumination": 0.532,
+    "numerical_aperture_illumination": 0.9,
+    "numerical_aperture_detection": 1.2,
+}
+
+# Define 9 sector illumination angles
+# 8 sectors at 45-degree intervals + 1 full aperture
+sector_angle = 45
+sector_angle_offset = -22.5
+illumination_sector_angles = [
+    (
+        (i * sector_angle + sector_angle_offset) % 360,
+        ((i + 1) * sector_angle + sector_angle_offset) % 360,
+    )
+    for i in range(8)
+] + [(0, 360)]
+
+print(f"Using {illumination_sector_angles} illumination sectors")
+
+# Create a phantom
+zyx_phase = phase_thick_3d.generate_test_phantom(
+    **simulation_arguments, **phantom_arguments
+)
+
+# Calculate multi-channel transfer function (one for each sector)
+(
+    real_potential_transfer_function,
+    imag_potential_transfer_function,
+) = phase_thick_3d.calculate_transfer_function(
+    **simulation_arguments,
+    **transfer_function_arguments,
+    illumination_sector_angles=illumination_sector_angles,
+)
+
+print(
+    f"Transfer function shape: {real_potential_transfer_function.shape}"
+)  # Should be (C, Z, Y, X)
+
+# Display complete multi-channel transfer function
+viewer = napari.Viewer()
+zyx_scale = np.array(
+    [
+        simulation_arguments["z_pixel_size"],
+        simulation_arguments["yx_pixel_size"],
+        simulation_arguments["yx_pixel_size"],
+    ]
+)
+
+# Add full CZYX transfer function (imaginary part) as single 4D layer
+# Match the visualization style from add_transfer_function_to_viewer
+czyx_shape = imag_potential_transfer_function.shape
+voxel_scale = np.array(
+    [
+        czyx_shape[1] * zyx_scale[0],  # Z extent
+        czyx_shape[2] * zyx_scale[1],  # Y extent
+        czyx_shape[3] * zyx_scale[2],  # X extent
+    ]
+)
+lim = 0.5 * torch.max(torch.abs(imag_potential_transfer_function)).item()
+
+viewer.add_image(
+    torch.fft.ifftshift(
+        torch.imag(imag_potential_transfer_function), dim=(-3, -2, -1)
+    )
+    .cpu()
+    .numpy(),
+    name="Imag pot. TF (CZYX)",
+    colormap="bwr",
+    contrast_limits=(-lim, lim),
+    scale=(1,) + tuple(1 / voxel_scale),  # No scaling on C dimension
+)
+
+# Set up XZ view with C and Y as sliders
+viewer.dims.order = [0, 2, 1, 3]  # (C, Y, Z, X) for XZ display
+viewer.dims.current_step = (
+    0,
+    czyx_shape[1] // 2,
+    czyx_shape[2] // 2,
+    czyx_shape[3] // 2,
+)
+
+input(
+    "Showing CZYX OTF in XZ view (use C and Y sliders). Press <enter> to continue..."
+)
+viewer.layers.select_all()
+viewer.layers.remove_selected()
+
+# Simulate multi-channel data (one channel per sector)
+# In practice, these would come from your microscope as separate acquisitions
+zyx_data_multi_channel = []
+for c in range(len(illumination_sector_angles)):
+    zyx_data_channel = phase_thick_3d.apply_transfer_function(
+        zyx_phase,
+        real_potential_transfer_function[c],
+        transfer_function_arguments["z_padding"],
+        brightness=1e3,
+    )
+    zyx_data_multi_channel.append(zyx_data_channel)
+
+# Stack into (C, Z, Y, X) tensor
+zyx_data_multi_channel = torch.stack(zyx_data_multi_channel, dim=0)
+print(f"Multi-channel data shape: {zyx_data_multi_channel.shape}")
+
+# Reconstruct phase from all channels combined
+zyx_recon = phase_thick_3d.apply_inverse_transfer_function(
+    zyx_data_multi_channel,
+    real_potential_transfer_function,
+    imag_potential_transfer_function,
+    transfer_function_arguments["z_padding"],
+)
+
+# Display
+viewer.add_image(zyx_phase.numpy(), name="Phantom", scale=zyx_scale)
+viewer.add_image(
+    zyx_data_multi_channel.numpy(),
+    name="Data (CZYX)",
+    scale=zyx_scale,
+)
+viewer.add_image(zyx_recon.numpy(), name="Reconstruction", scale=zyx_scale)
+
+# Show comparison with single channel (full aperture) for reference
+print("\nComparing with single-channel (full aperture) reconstruction...")
+(
+    real_tf_single,
+    imag_tf_single,
+) = phase_thick_3d.calculate_transfer_function(
+    **simulation_arguments,
+    **transfer_function_arguments,
+    illumination_sector_angles=None,  # Full aperture
+)
+zyx_data_single = phase_thick_3d.apply_transfer_function(
+    zyx_phase,
+    real_tf_single[0],  # Single channel
+    transfer_function_arguments["z_padding"],
+    brightness=1e3,
+)
+zyx_recon_single = phase_thick_3d.apply_inverse_transfer_function(
+    zyx_data_single[None, ...],  # Add channel dimension
+    real_tf_single,
+    imag_tf_single,
+    transfer_function_arguments["z_padding"],
+)
+viewer.add_image(
+    zyx_recon_single.numpy(),
+    name="Reconstruction (single channel)",
+    scale=zyx_scale,
+)
+
+print(
+    f"\nReconstruction error (multi-channel): {torch.mean(torch.abs(zyx_recon - zyx_phase)).item():.6f}"
+)
+print(
+    f"Reconstruction error (single channel): {torch.mean(torch.abs(zyx_recon_single - zyx_phase)).item():.6f}"
+)
+
+input("\nShowing phantom, data, and reconstructions. Press <enter> to quit...")

tests/models/test_isotropic_thin_3d.py

@@ -17,5 +17,6 @@ def test_calculate_transfer_function(invert_phase_contrast):
         invert_phase_contrast=invert_phase_contrast,
     )
 
-    assert Hu.shape == (3, 100, 101)
-    assert Hp.shape == (3, 100, 101)
+    # Transfer functions now have shape (C, Z, Y, X) where C is number of channels
+    assert Hu.shape == (1, 3, 100, 101)
+    assert Hp.shape == (1, 3, 100, 101)

tests/models/test_phase_thick_3d.py

@@ -19,8 +19,9 @@ def test_calculate_transfer_function(invert_phase_contrast):
         invert_phase_contrast=invert_phase_contrast,
     )
 
-    assert H_re.shape == (20 + 2 * z_padding, 100, 101)
-    assert H_im.shape == (20 + 2 * z_padding, 100, 101)
+    # Transfer functions now have shape (C, Z, Y, X) where C is number of channels
+    assert H_re.shape == (1, 20 + 2 * z_padding, 100, 101)
+    assert H_im.shape == (1, 20 + 2 * z_padding, 100, 101)
 
 
 # Helper function for testing reconstruction invariances
@@ -58,25 +59,25 @@ def simulate_phase_recon(
         **simulation_arguments, **phantom_arguments
     )
 
-    # Calculate transfer function
+    # Calculate transfer function (returns shape (C, Z, Y, X))
     (
         real_potential_transfer_function,
         imag_potential_transfer_function,
     ) = phase_thick_3d.calculate_transfer_function(
         **simulation_arguments, **transfer_function_arguments
     )
 
-    # Simulate
+    # Simulate (extract single channel for forward model)
     zyx_data = phase_thick_3d.apply_transfer_function(
         zyx_phase,
-        real_potential_transfer_function,
+        real_potential_transfer_function[0],
         transfer_function_arguments["z_padding"],
         brightness=1000,
     )
 
-    # Reconstruct
+    # Reconstruct (wrap data in channel dimension for inverse)
     zyx_recon = phase_thick_3d.apply_inverse_transfer_function(
-        zyx_data,
+        zyx_data[None],  # Add channel dimension: (Z, Y, X) -> (1, Z, Y, X)
         real_potential_transfer_function,
         imag_potential_transfer_function,
         transfer_function_arguments["z_padding"],

waveorder/cli/apply_inverse_models.py

@@ -93,19 +93,19 @@ def phase(
 
     # [phase only, 3]
     elif recon_dim == 3:
-        # Load transfer functions
+        # Load transfer functions (keep channel dimension)
         real_potential_transfer_function = torch.tensor(
-            transfer_function_dataset["real_potential_transfer_function"][0, 0]
+            transfer_function_dataset["real_potential_transfer_function"][0]
         )
         imaginary_potential_transfer_function = torch.tensor(
             transfer_function_dataset["imaginary_potential_transfer_function"][
-                0, 0
+                0
             ]
         )
 
-        # Apply
+        # Apply (pass full CZYX data)
         output = phase_thick_3d.apply_inverse_transfer_function(
-            czyx_data[0],
+            czyx_data,
             real_potential_transfer_function,
             imaginary_potential_transfer_function,
             z_padding=settings_phase.transfer_function.z_padding,

waveorder/cli/compute_transfer_function.py

@@ -69,6 +69,9 @@ def generate_and_save_vector_birefringence_transfer_function(
     )
     phase_settings_dict = settings.phase.transfer_function.model_dump()
     phase_settings_dict.pop("z_focus_offset")  # not used in 3D
+    phase_settings_dict.pop(
+        "illumination_sector_angles"
+    )  # not used in vector birefringence
 
     sfZYX_transfer_function, _, singular_system = (
         inplane_oriented_thick_pol3d_vector.calculate_transfer_function(
@@ -213,14 +216,12 @@ def generate_and_save_phase_transfer_function(
         # Save
         dataset.create_image(
             "real_potential_transfer_function",
-            real_potential_transfer_function.cpu().numpy()[None, None, ...],
+            real_potential_transfer_function.cpu().numpy()[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, ...
-            ],
+            imaginary_potential_transfer_function.cpu().numpy()[None, ...],
             chunks=(1, 1, 1, zyx_shape[1], zyx_shape[2]),
         )
 
@@ -369,14 +370,15 @@ def compute_transfer_function_cli(
         print("Found z_focus_offset:", z_focus_offset)
 
     # Prepare output dataset
-    num_channels = (
+    num_input_channel = len(settings.input_channel_names)
+    num_output_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"],
+        channel_names=num_input_channel * num_output_channels * ["None"],
     )
 
     # Pass settings to appropriate calculate_transfer_function and save