visual cleanup

waveorder/visuals/matplotlib_visuals.py

@@ -13,7 +13,7 @@ def plot_5d_ortho(
     filename,
     voxel_size,
     zyx_slice,
-    color_func,
+    color_funcs,
     row_labels=None,
     column_labels=None,
     rose_path=None,
@@ -45,6 +45,13 @@ def plot_5d_ortho(
     assert zyx_slice[0] < Z and zyx_slice[1] < Y and zyx_slice[2] < X
     assert zyx_slice[0] >= 0 and zyx_slice[1] >= 0 and zyx_slice[2] >= 0
 
+    assert R == len(color_funcs)
+    for color_func_row in color_funcs:
+        if isinstance(color_func_row, list):
+            assert len(color_func_row) == C
+        else:
+            color_func_row = [color_func_row] * C
+
     n_rows = 1 + (2 * R)
     n_cols = 1 + (2 * C)
 
@@ -100,6 +107,8 @@ def plot_5d_ortho(
 
             # Add data
             if i > 0 and j > 0:
+                color_func = color_funcs[int((i - 1) / 2)][int((j - 1) / 2)]
+
                 Cyx_data = rcCzyx_data[
                     int((i - 1) / 2), int((j - 1) / 2), :, zyx_slice[0]
                 ]

waveorder/visuals/napari_visuals.py

@@ -25,7 +25,7 @@ def add_transfer_function_to_viewer(
 
     if complex_rgb:
         rgb_transfer_function = complex_tensor_to_rgb(
-            torch.fft.ifftshift(transfer_function, dim=shift_dims),
+            np.array(torch.fft.ifftshift(transfer_function, dim=shift_dims)),
             saturate_clim_fraction=clim_factor,
         )
         viewer.add_image(

waveorder/visuals/utils.py

@@ -15,7 +15,10 @@ def complex_tensor_to_rgb(array, saturate_clim_fraction=1.0):
     hue = np.mod(hue + 0.5, 1)
 
     # Normalize magnitude to [0, 1] for saturation
-    max_abs_val = np.amax(magnitude) * saturate_clim_fraction
+    if saturate_clim_fraction is not None:
+        max_abs_val = np.amax(magnitude) * saturate_clim_fraction
+    else:
+        max_abs_val = 1.0
 
     sat = magnitude / max_abs_val if max_abs_val != 0 else magnitude