Repository: mehta-lab/waveorder

Issue:
Focus-finding features
@ieivanov and I just discussed some desired features for the focus-finding algorithms:

- return a normalized "SNR" or "confidence" metric alongside the in-focus slice. We're concerned that low-SNR images may make the focus-finding fail and cause erratic stage motion. I need to think more and test exactly what metric we'll use, but @ieivanov and I agreed that we want the metric to be invariant to (1) the number of pixels, (2) the overall brightness of the sample, and (3) the size of the sample. 
- a `plot=True` option to save a small figure of the focus-finding run, likely mid-band power vs. slice, a dot for the peak, and an annotation for the confidence metric. 

Discussion:

Patch:
diff --git a/examples/2D_QLIPP_simulation/2D_QLIPP_forward.py b/examples/2D_QLIPP_simulation/2D_QLIPP_forward.py
index 84a958a..70e14c8 100644
--- a/examples/2D_QLIPP_simulation/2D_QLIPP_forward.py
+++ b/examples/2D_QLIPP_simulation/2D_QLIPP_forward.py
@@ -12,9 +12,13 @@
 
 import numpy as np
 import matplotlib.pyplot as plt
-from numpy.fft import fft2, ifft2, fftshift, ifftshift
-
-import waveorder as wo
+from numpy.fft import fftshift
+from waveorder import (
+    optics,
+    waveorder_simulator,
+    visual,
+    util,
+)
 
 # Key parameters
 N = 256  # number of pixel in y dimension
@@ -31,8 +35,8 @@ chi = 0.03 * 2 * np.pi  # swing of Polscope analyzer
 
 # Sample : star with uniform phase, uniform retardance, and radial orientation
 # generate Siemens star pattern
-star, theta, _ = wo.genStarTarget(N, M)
-wo.plot_multicolumn(np.array([star, theta]), num_col=2, size=5)
+star, theta, _ = util.genStarTarget(N, M)
+visual.plot_multicolumn(np.array([star, theta]), num_col=2, size=5)
 
 # Assign uniform phase, uniform retardance, and radial slow axes to the star pattern
 phase_value = 1  # average phase in radians (optical path length)
@@ -44,7 +48,7 @@ t_eigen = np.zeros((2, N, M), complex)  # complex specimen transmission
 t_eigen[0] = np.exp(-mu_s + 1j * phi_s)
 t_eigen[1] = np.exp(-mu_f + 1j * phi_f)
 sa = theta % np.pi  # slow axes.
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     np.array([phi_s, phi_f, mu_s, sa]),
     num_col=2,
     size=5,
@@ -58,9 +62,9 @@ plt.show()
 # Source pupil
 # Subsample source pattern for speed
 
-xx, yy, fxx, fyy = wo.gen_coordinate((N, M), ps)
-Source_cont = wo.gen_Pupil(fxx, fyy, NA_illu, lambda_illu)
-Source_discrete = wo.Source_subsample(
+xx, yy, fxx, fyy = util.gen_coordinate((N, M), ps)
+Source_cont = optics.gen_Pupil(fxx, fyy, NA_illu, lambda_illu)
+Source_discrete = optics.Source_subsample(
     Source_cont, lambda_illu * fxx, lambda_illu * fyy, subsampled_NA=0.1
 )
 plt.figure(figsize=(10, 10))
@@ -70,7 +74,7 @@ plt.show()
 # Initialize microscope simulator with above source pattern and uniform imaging pupil
 # Microscope object generation
 
-simulator = wo.waveorder_microscopy_simulator(
+simulator = waveorder_simulator.waveorder_microscopy_simulator(
     (N, M),
     lambda_illu,
     ps,
diff --git a/examples/2D_QLIPP_simulation/2D_QLIPP_recon.py b/examples/2D_QLIPP_simulation/2D_QLIPP_recon.py
index 2d406a3..b1f64e7 100644
--- a/examples/2D_QLIPP_simulation/2D_QLIPP_recon.py
+++ b/examples/2D_QLIPP_simulation/2D_QLIPP_recon.py
@@ -12,8 +12,11 @@
 
 import numpy as np
 import matplotlib.pyplot as plt
-from numpy.fft import fft2, ifft2, fftshift, ifftshift
-import waveorder as wo
+from waveorder import (
+    waveorder_reconstructor,
+    visual,
+)
+
 
 # ### Load simulated data
 # Load simulations
@@ -34,7 +37,7 @@ _, N, M, L = I_meas.shape
 cali = False
 bg_option = "global"
 
-setup = wo.waveorder_microscopy(
+setup = waveorder_reconstructor.waveorder_microscopy(
     (N, M),
     lambda_illu,
     ps,
@@ -53,7 +56,7 @@ Recon_para = setup.Polarization_recon(
     S_image_tm
 )  # Without accounting for diffraction
 
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     np.array(
         [
             Recon_para[0, :, :, L // 2],
@@ -69,7 +72,7 @@ wo.plot_multicolumn(
     origin="lower",
 )
 
-wo.plot_hsv(
+visual.plot_hsv(
     [Recon_para[1, :, :, L // 2], Recon_para[0, :, :, L // 2]],
     max_val=1,
     origin="lower",
@@ -86,7 +89,7 @@ retardance, azimuth = setup.Birefringence_recon_2D(
     S1_stack, S2_stack, method="Tikhonov", reg_br=1e-3
 )
 
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     np.array([retardance, azimuth]),
     num_col=2,
     size=10,
@@ -94,7 +97,7 @@ wo.plot_multicolumn(
     titles=["Reconstructed retardance", "Reconstructed orientation"],
     origin="lower",
 )
-wo.plot_hsv([azimuth, retardance], size=10, origin="lower")
+visual.plot_hsv([azimuth, retardance], size=10, origin="lower")
 plt.show()
 
 
@@ -110,7 +113,7 @@ retardance_TV, azimuth_TV = setup.Birefringence_recon_2D(
     verbose=True,
 )
 
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     np.array([retardance_TV, azimuth_TV]),
     num_col=2,
     size=10,
@@ -118,7 +121,7 @@ wo.plot_multicolumn(
     titles=["Reconstructed retardance", "Reconstructed orientation"],
     origin="lower",
 )
-wo.plot_hsv([azimuth_TV, retardance_TV], size=10, origin="lower")
+visual.plot_hsv([azimuth_TV, retardance_TV], size=10, origin="lower")
 plt.show()
 
 
@@ -130,7 +133,7 @@ S0_stack = S_image_recon[0].copy()
 mu_sample, phi_sample = setup.Phase_recon(
     S0_stack, method="Tikhonov", reg_u=reg_u, reg_p=reg_p
 )
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     np.array([mu_sample, phi_sample]),
     num_col=2,
     size=10,
@@ -148,7 +151,7 @@ S0_stack = S_image_recon[0].copy()
 mu_sample_TV, phi_sample_TV = setup.Phase_recon(
     S0_stack, method="TV", lambda_u=lambda_u, lambda_p=lambda_p, itr=10, rho=1
 )
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     np.array([mu_sample_TV, phi_sample_TV]),
     num_col=2,
     size=10,
diff --git a/examples/3D_PODT_phase_simulation/3D_PODT_Phase_forward.py b/examples/3D_PODT_phase_simulation/3D_PODT_Phase_forward.py
index da7fd55..28fee46 100644
--- a/examples/3D_PODT_phase_simulation/3D_PODT_Phase_forward.py
+++ b/examples/3D_PODT_phase_simulation/3D_PODT_Phase_forward.py
@@ -12,8 +12,14 @@
 
 import numpy as np
 import matplotlib.pyplot as plt
-from numpy.fft import fft, ifft, fft2, ifft2, fftshift, ifftshift, fftn, ifftn
-import waveorder as wo
+from numpy.fft import fftshift
+from waveorder import (
+    optics,
+    waveorder_simulator,
+    waveorder_reconstructor,
+    visual,
+    util,
+)
 
 # Experiment parameters
 N = 256  # number of pixel in y dimension
@@ -30,8 +36,8 @@ NA_illu = 0.9  # illumination NA
 # Sample creation
 radius = 5
 blur_size = 2 * ps
-sphere, _, _ = wo.gen_sphere_target((N, M, L), ps, psz, radius, blur_size)
-wo.image_stack_viewer(np.transpose(sphere, (2, 0, 1)))
+sphere, _, _ = util.gen_sphere_target((N, M, L), ps, psz, radius, blur_size)
+visual.image_stack_viewer(np.transpose(sphere, (2, 0, 1)))
 
 # Physical value assignment
 
@@ -56,9 +62,9 @@ plt.show()
 # Setup acquisition
 # Subsampled Source pattern
 
-xx, yy, fxx, fyy = wo.gen_coordinate((N, M), ps)
-Source_cont = wo.gen_Pupil(fxx, fyy, NA_illu, lambda_illu)
-Source_discrete = wo.Source_subsample(
+xx, yy, fxx, fyy = util.gen_coordinate((N, M), ps)
+Source_cont = optics.gen_Pupil(fxx, fyy, NA_illu, lambda_illu)
+Source_discrete = optics.Source_subsample(
     Source_cont, lambda_illu * fxx, lambda_illu * fyy, subsampled_NA=0.1
 )
 plt.figure(figsize=(10, 10))
@@ -68,7 +74,7 @@ print(np.sum(Source_discrete))
 
 z_defocus = (np.r_[:L] - L // 2) * psz
 chi = 0.1 * 2 * np.pi
-setup = wo.waveorder_microscopy(
+setup = waveorder_reconstructor.waveorder_microscopy(
     (N, M),
     lambda_illu,
     ps,
@@ -82,7 +88,7 @@ setup = wo.waveorder_microscopy(
     Source=Source_cont,
 )
 
-simulator = wo.waveorder_microscopy_simulator(
+simulator = waveorder_simulator.waveorder_microscopy_simulator(
     (N, M),
     lambda_illu,
     ps,
@@ -99,7 +105,7 @@ plt.figure(figsize=(5, 5))
 plt.imshow(fftshift(setup.Source), cmap="gray")
 plt.colorbar()
 H_re_vis = fftshift(setup.H_re)
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     [
         np.real(H_re_vis)[:, :, L // 2],
         np.transpose(np.real(H_re_vis)[N // 2, :, :]),
@@ -120,7 +126,7 @@ wo.plot_multicolumn(
 plt.show()
 
 H_im_vis = fftshift(setup.H_im)
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     [
         np.real(H_im_vis)[:, :, L // 2],
         np.transpose(np.real(H_im_vis)[N // 2, :, :]),
diff --git a/examples/3D_PODT_phase_simulation/3D_PODT_Phase_recon.py b/examples/3D_PODT_phase_simulation/3D_PODT_Phase_recon.py
index 4fbbb6f..d152fd0 100644
--- a/examples/3D_PODT_phase_simulation/3D_PODT_Phase_recon.py
+++ b/examples/3D_PODT_phase_simulation/3D_PODT_Phase_recon.py
@@ -14,9 +14,11 @@
 
 import numpy as np
 import matplotlib.pyplot as plt
-from numpy.fft import fft, ifft, fft2, ifft2, fftn, ifftn, fftshift, ifftshift
-
-import waveorder as wo
+from numpy.fft import fftshift
+from waveorder import (
+    waveorder_reconstructor,
+    visual,
+)
 
 # Load data
 # Load simulations
@@ -33,7 +35,7 @@ N, M, L = I_meas.shape
 # Refractive index reconstruction
 z_defocus = (np.r_[:L] - L // 2) * psz
 chi = 0.1 * 2 * np.pi
-setup = wo.waveorder_microscopy(
+setup = waveorder_reconstructor.waveorder_microscopy(
     (N, M),
     lambda_illu,
     ps,
@@ -47,7 +49,7 @@ setup = wo.waveorder_microscopy(
 )
 
 H_re_vis = fftshift(setup.H_re)
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     [
         np.real(H_re_vis)[:, :, L // 2],
         np.transpose(np.real(H_re_vis)[N // 2, :, :]),
@@ -68,7 +70,7 @@ wo.plot_multicolumn(
 plt.show()
 
 H_im_vis = fftshift(setup.H_im)
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     [
         np.real(H_im_vis)[:, :, L // 2],
         np.transpose(np.real(H_im_vis)[N // 2, :, :]),
diff --git a/examples/PTI_simulation/PTI_Simulation_Forward_2D3D.py b/examples/PTI_simulation/PTI_Simulation_Forward_2D3D.py
index 105e629..eda1f42 100644
--- a/examples/PTI_simulation/PTI_Simulation_Forward_2D3D.py
+++ b/examples/PTI_simulation/PTI_Simulation_Forward_2D3D.py
@@ -11,10 +11,13 @@
 
 import numpy as np
 import matplotlib.pyplot as plt
-from numpy.fft import fft, ifft, fft2, ifft2, fftshift, ifftshift, fftn, ifftn
-
-import pickle
-import waveorder as wo
+from numpy.fft import fftshift
+from waveorder import (
+    optics,
+    waveorder_simulator,
+    visual,
+    util,
+)
 
 #####################################################################
 # Initialization - imaging system and sample                        #
@@ -68,7 +71,7 @@ gpu_id = 0  # gpu to be used
 if sample_type == "3D":
     # 3D spoke pattern whose spokes are inclined 60 degrees relative to the z-axis. The principal retardance varies with depth and 3D orientation of permittivity tensor is aligned with the structural orientation of the spoke.
     blur_size = 1 * ps
-    target, azimuth, inclination = wo.genStarTarget_3D(
+    target, azimuth, inclination = util.genStarTarget_3D(
         (N, M, L),
         ps,
         psz,
@@ -80,7 +83,7 @@ if sample_type == "3D":
     azimuth = np.round(azimuth / np.pi / 2 * 16) / 16 * np.pi * 2
 elif sample_type == "2D":
     ## 2D spoke pattern, azimuth aligned with spokes, and the inclination set to 60 degrees ##
-    target, azimuth, _ = wo.genStarTarget(N, M, blur_px=1 * ps, margin=10)
+    target, azimuth, _ = util.genStarTarget(N, M, blur_px=1 * ps, margin=10)
     inclination = np.ones_like(target) * np.pi / 3
     azimuth = azimuth % (np.pi * 2)
     azimuth = np.round(azimuth / np.pi / 2 * 16) / 16 * np.pi * 2
@@ -138,7 +141,7 @@ biref_map = ne_map_copy - no_map_copy
 ### Visualize sample properties
 
 #### XY sections
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     [
         target[:, :, z_layer],
         azimuth[:, :, z_layer] % (2 * np.pi),
@@ -151,7 +154,7 @@ wo.plot_multicolumn(
     set_title=True,
 )
 #### XZ sections
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     [
         np.transpose(target[y_layer, :, :]),
         np.transpose(azimuth[y_layer, :, :]) % (2 * np.pi),
@@ -190,7 +193,7 @@ orientation_3D_image = np.transpose(
     ),
     (3, 1, 2, 0),
 )
-orientation_3D_image_RGB = wo.orientation_3D_to_rgb(
+orientation_3D_image_RGB = visual.orientation_3D_to_rgb(
     orientation_3D_image, interp_belt=20 / 180 * np.pi, sat_factor=1
 )
 
@@ -199,7 +202,7 @@ plt.imshow(orientation_3D_image_RGB[z_layer], origin="lower")
 plt.figure(figsize=(10, 10))
 plt.imshow(orientation_3D_image_RGB[:, y_layer], origin="lower")
 plt.figure(figsize=(3, 3))
-wo.orientation_3D_colorwheel(
+visual.orientation_3D_colorwheel(
     wheelsize=128,
     circ_size=50,
     interp_belt=20 / 180 * np.pi,
@@ -209,7 +212,7 @@ wo.orientation_3D_colorwheel(
 plt.show()
 
 #### Angular histogram of 3D orientation
-wo.orientation_3D_hist(
+visual.orientation_3D_hist(
     azimuth.flatten(),
     inclination.flatten(),
     np.abs(target).flatten(),
@@ -251,7 +254,7 @@ epsilon_tensor[2, 1] = epsilon_del * np.sin(2 * inclination) * np.sin(azimuth)
 epsilon_tensor[2, 2] = epsilon_mean + epsilon_del * np.cos(2 * inclination)
 
 
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     [
         epsilon_tensor[0, 0, :, :, z_layer],
         epsilon_tensor[0, 1, :, :, z_layer],
@@ -327,7 +330,7 @@ del_f_component[6] = (
 )
 
 
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     [
         del_f_component[0, :, :, z_layer],
         del_f_component[1, :, :, z_layer],
@@ -361,9 +364,9 @@ plt.show()
 
 # DPC + BF illumination + PolState (sector illumination)
 
-xx, yy, fxx, fyy = wo.gen_coordinate((N, M), ps)
-Pupil_obj = wo.gen_Pupil(fxx, fyy, NA_obj / n_media, lambda_illu / n_media)
-Source_support = wo.gen_Pupil(
+xx, yy, fxx, fyy = util.gen_coordinate((N, M), ps)
+Pupil_obj = optics.gen_Pupil(fxx, fyy, NA_obj / n_media, lambda_illu / n_media)
+Source_support = optics.gen_Pupil(
     fxx, fyy, NA_illu / n_media, lambda_illu / n_media
 )
 
@@ -375,12 +378,12 @@ rotation_angle = [0, 45, 90, 135, 180, 225, 270, 315]
 Source = np.zeros((len(rotation_angle) + 1, N, M))
 Source_cont = np.zeros_like(Source)
 
-Source_BF = wo.gen_Pupil(
+Source_BF = optics.gen_Pupil(
     fxx, fyy, NA_illu / n_media / 2, lambda_illu / n_media
 )
 
 Source_cont[-1] = Source_BF.copy()
-Source[-1] = wo.Source_subsample(
+Source[-1] = optics.Source_subsample(
     Source_BF, NAx_coord, NAy_coord, subsampled_NA=0.1 / n_media
 )
 
@@ -401,7 +404,7 @@ for i in range(len(rotation_angle)):
 
     Source_cont[i] = Source_temp * Source_temp2 * Source_support
 
-    Source_discrete = wo.Source_subsample(
+    Source_discrete = optics.Source_subsample(
         Source_cont[i], NAx_coord, NAy_coord, subsampled_NA=0.1 / n_media
     )
     Source[i] = np.maximum(0, Source_discrete.copy())
@@ -414,11 +417,11 @@ for i in range(len(Source)):
 
 #### Circularly polarized illumination patterns
 
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     fftshift(Source_cont, axes=(1, 2)), origin="lower", num_col=5, size=5
 )
 # discretized illumination patterns used in simulation (faster forward model)
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     fftshift(Source, axes=(1, 2)), origin="lower", num_col=5, size=5
 )
 print(Source_PolState)
@@ -430,7 +433,7 @@ print(np.sum(Source, axis=(1, 2)))
 #### Initialize microscope simulator with above source pattern and uniform imaging pupil
 
 ## initiate the simulator
-simulator = wo.waveorder_microscopy_simulator(
+simulator = waveorder_simulator.waveorder_microscopy_simulator(
     (N, M),
     lambda_illu,
     ps,
diff --git a/examples/PTI_simulation/PTI_Simulation_Recon2D.py b/examples/PTI_simulation/PTI_Simulation_Recon2D.py
index 9419673..9fb4518 100644
--- a/examples/PTI_simulation/PTI_Simulation_Recon2D.py
+++ b/examples/PTI_simulation/PTI_Simulation_Recon2D.py
@@ -11,10 +11,13 @@
 
 import numpy as np
 import matplotlib.pyplot as plt
-from numpy.fft import fft, ifft, fft2, ifft2, fftn, ifftn, fftshift, ifftshift
-
-import waveorder as wo
+from numpy.fft import fftshift
 
+from waveorder import (
+    optics,
+    waveorder_reconstructor,
+    visual,
+)
 
 ## Initialization
 ## Load simulated images and parameters
@@ -50,7 +53,7 @@ gpu_id = 0
 A_matrix = 0.5 * np.array([[1, 1, 0], [1, 0, 1], [1, -1, 0], [1, 0, -1]])
 
 
-setup = wo.waveorder_microscopy(
+setup = waveorder_reconstructor.waveorder_microscopy(
     (N, M),
     lambda_illu,
     ps,
@@ -70,10 +73,10 @@ setup = wo.waveorder_microscopy(
     gpu_id=gpu_id,
 )
 
-## Visualize 2D transfer functions as a function of illumination pattern
+## Visualize 2  D transfer functions as a function of illumination pattern
 
 # illumination patterns used
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     fftshift(Source_cont, axes=(1, 2)), origin="lower", num_col=5, size=5
 )
 plt.show()
@@ -115,7 +118,7 @@ f_tensor = setup.scattering_potential_tensor_recon_2D_vec(
     S_image_tm, reg_inc=reg_inc, cupy_det=True
 )
 
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     f_tensor,
     num_col=4,
     origin="lower",
@@ -153,13 +156,13 @@ plt.show()
 # scaling to the physical properties of the material
 
 # optic sign probability
-p_mat_map = wo.optic_sign_probability(mat_map, mat_map_thres=0.1)
+p_mat_map = optics.optic_sign_probability(mat_map, mat_map_thres=0.1)
 
 # absorption and phase
-phase = wo.phase_inc_correction(f_tensor[0], retardance_pr[0], theta[0])
+phase = optics.phase_inc_correction(f_tensor[0], retardance_pr[0], theta[0])
 absorption = f_tensor[1].copy()
 phase_nm, absorption_nm, retardance_pr_nm = [
-    wo.unit_conversion_from_scattering_potential_to_permittivity(
+    optics.unit_conversion_from_scattering_potential_to_permittivity(
         SP_array, lambda_illu, n_media=n_media, imaging_mode=img_mode
     )
     for img_mode, SP_array in zip(
@@ -252,14 +255,14 @@ orientation_3D_image = np.transpose(
     ),
     (1, 2, 0),
 )
-orientation_3D_image_RGB = wo.orientation_3D_to_rgb(
+orientation_3D_image_RGB = visual.orientation_3D_to_rgb(
     orientation_3D_image, interp_belt=20 / 180 * np.pi, sat_factor=1
 )
 
 plt.figure(figsize=(5, 5))
 plt.imshow(orientation_3D_image_RGB, origin="lower")
 plt.figure(figsize=(3, 3))
-wo.orientation_3D_colorwheel(
+visual.orientation_3D_colorwheel(
     wheelsize=256, circ_size=50, interp_belt=20 / 180 * np.pi, sat_factor=1
 )
 plt.show()
@@ -294,7 +297,7 @@ in_plane_orientation = hsv_to_rgb(I_hsv.copy())
 plt.figure(figsize=(5, 5))
 plt.imshow(in_plane_orientation, origin="lower")
 plt.figure(figsize=(3, 3))
-wo.orientation_2D_colorwheel()
+visual.orientation_2D_colorwheel()
 plt.show()
 
 # out-of-plane tilt
@@ -336,7 +339,7 @@ spacing = 4
 plt.figure(figsize=(10, 10))
 
 fig, ax = plt.subplots(1, 1, figsize=(20, 10))
-wo.plot3DVectorField(
+visual.plot3DVectorField(
     np.abs(retardance_pr_nm[0]),
     azimuth[0],
     theta[0],
@@ -359,7 +362,7 @@ ret_mask[ret_mask < 0.5] = 0
 
 plt.figure(figsize=(10, 10))
 plt.imshow(ret_mask, cmap="gray", origin="lower")
-wo.orientation_3D_hist(
+visual.orientation_3D_hist(
     azimuth[0].flatten(),
     theta[0].flatten(),
     ret_mask.flatten(),
diff --git a/examples/PTI_simulation/PTI_Simulation_Recon3D.py b/examples/PTI_simulation/PTI_Simulation_Recon3D.py
index 4a3a276..e1edbb3 100644
--- a/examples/PTI_simulation/PTI_Simulation_Recon3D.py
+++ b/examples/PTI_simulation/PTI_Simulation_Recon3D.py
@@ -10,8 +10,12 @@
 ####################################################################
 import numpy as np
 import matplotlib.pyplot as plt
-from numpy.fft import fft, ifft, fft2, ifft2, fftn, ifftn, fftshift, ifftshift
-import waveorder as wo
+from numpy.fft import fftshift
+from waveorder import (
+    optics,
+    waveorder_reconstructor,
+    visual,
+)
 
 ## Initialization
 ## Load simulated images and parameters
@@ -39,7 +43,7 @@ gpu_id = 0
 A_matrix = 0.5 * np.array([[1, 1, 0], [1, 0, 1], [1, -1, 0], [1, 0, -1]])
 
 
-setup = wo.waveorder_microscopy(
+setup = waveorder_reconstructor.waveorder_microscopy(
     (N, M),
     lambda_illu,
     ps,
@@ -61,7 +65,7 @@ setup = wo.waveorder_microscopy(
 
 
 ### Illumination patterns used
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     fftshift(Source_cont, axes=(1, 2)), origin="lower", num_col=5, size=5
 )
 plt.show()
@@ -109,7 +113,7 @@ f_tensor = setup.scattering_potential_tensor_recon_3D_vec(
     S_image_tm, reg_inc=reg_inc, cupy_det=True
 )
 
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     f_tensor[..., L // 2],
     num_col=4,
     origin="lower",
@@ -154,14 +158,14 @@ plt.show()
 # scaling to the physical properties of the material
 
 # optic sign probability
-p_mat_map = wo.optic_sign_probability(mat_map, mat_map_thres=0.2)
+p_mat_map = optics.optic_sign_probability(mat_map, mat_map_thres=0.2)
 
 
 # absorption and phase
-phase = wo.phase_inc_correction(f_tensor[0], retardance_pr[0], theta[0])
+phase = optics.phase_inc_correction(f_tensor[0], retardance_pr[0], theta[0])
 absorption = f_tensor[1].copy()
 phase_PT, absorption_PT, retardance_pr_PT = [
-    wo.unit_conversion_from_scattering_potential_to_permittivity(
+    optics.unit_conversion_from_scattering_potential_to_permittivity(
         SP_array, lambda_illu, n_media=n_media, imaging_mode="3D"
     )
     for SP_array in [phase, absorption, retardance_pr]
@@ -179,7 +183,7 @@ phase_PT, absorption_PT, retardance_pr_PT = [
 ### Reconstructed phase, absorption, principal retardance, azimuth, and inclination assuming (+) and (-) optic sign
 
 # browse the reconstructed physical properties
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     np.stack(
         [
             phase_PT[..., L // 2],
@@ -385,7 +389,7 @@ orientation_3D_image = np.transpose(
     ),
     (3, 1, 2, 0),
 )
-orientation_3D_image_RGB = wo.orientation_3D_to_rgb(
+orientation_3D_image_RGB = visual.orientation_3D_to_rgb(
     orientation_3D_image, interp_belt=20 / 180 * np.pi, sat_factor=1
 )
 
@@ -398,7 +402,7 @@ plt.imshow(
 )
 # plot the top view of 3D orientation colorsphere
 plt.figure(figsize=(3, 3))
-wo.orientation_3D_colorwheel(
+visual.orientation_3D_colorwheel(
     wheelsize=128,
     circ_size=50,
     interp_belt=20 / 180 * np.pi,
@@ -436,7 +440,7 @@ plt.imshow(in_plane_orientation[z_layer], origin="lower")
 plt.figure(figsize=(10, 10))
 plt.imshow(in_plane_orientation[:, y_layer], origin="lower", aspect=psz / ps)
 plt.figure(figsize=(3, 3))
-wo.orientation_2D_colorwheel()
+visual.orientation_2D_colorwheel()
 plt.show()
 
 
@@ -507,7 +511,7 @@ linelength_scale = 20
 
 
 fig, ax = plt.subplots(2, 2, figsize=(10, 10))
-wo.plot3DVectorField(
+visual.plot3DVectorField(
     np.abs(retardance_pr_PT[0, :, :, z_layer]),
     azimuth[0, :, :, z_layer],
     theta[0, :, :, z_layer],
@@ -525,7 +529,7 @@ wo.plot3DVectorField(
 )
 ax[0, 0].set_title(f"XY section (z= {z_layer})")
 
-wo.plot3DVectorField(
+visual.plot3DVectorField(
     np.transpose(np.abs(retardance_pr_PT[0, :, x_layer, :])),
     np.transpose(azimuth_x[0, :, x_layer, :]),
     np.transpose(theta_x[0, :, x_layer, :]),
@@ -543,7 +547,7 @@ wo.plot3DVectorField(
 )
 ax[0, 1].set_title(f"YZ section (x = {x_layer})")
 
-wo.plot3DVectorField(
+visual.plot3DVectorField(
     np.transpose(np.abs(retardance_pr_PT[0, y_layer, :, :])),
     np.transpose(azimuth_y[0, y_layer, :, :]),
     np.transpose(theta_y[0, y_layer, :, :]),
@@ -580,7 +584,7 @@ ret_mask[ret_mask < 0.00125] = 0
 
 plt.figure(figsize=(10, 10))
 plt.imshow(ret_mask[:, :, z_layer], cmap="gray", origin="lower")
-wo.orientation_3D_hist(
+visual.orientation_3D_hist(
     azimuth[0].flatten(),
     theta[0].flatten(),
     ret_mask.flatten(),
diff --git a/waveorder/__init__.py b/waveorder/__init__.py
index b04566e..e69de29 100644
--- a/waveorder/__init__.py
+++ b/waveorder/__init__.py
@@ -1,6 +0,0 @@
-from .visual import *
-from .waveorder_reconstructor import *
-from .waveorder_simulator import *
-from .util import *
-from .optics import *
-from .background_estimator import *
diff --git a/waveorder/focus.py b/waveorder/focus.py
index 3afb772..76fcce0 100644
--- a/waveorder/focus.py
+++ b/waveorder/focus.py
@@ -1,5 +1,7 @@
-from typing import Literal
+from scipy.signal import peak_widths
+from typing import Literal, Optional
 from waveorder import util
+import matplotlib.pyplot as plt
 import numpy as np
 
 
@@ -9,33 +11,46 @@ def focus_from_transverse_band(
     lambda_ill,
     pixel_size,
     midband_fractions=(0.125, 0.25),
-    mode: Literal["min" "max"] = "max",
+    mode: Literal["min", "max"] = "max",
+    plot_path: Optional[str] = None,
+    threshold_FWHM: float = 0,
 ):
     """Estimates the in-focus slice from a 3D stack by optimizing a transverse spatial frequency band.
 
     Parameters
     ----------
-    zyx_array : np.array
+    zyx_array: np.array
         Data stack in (Z, Y, X) order.
-        Requires len(3d_array.shape) == 3.
-    NA_det : float
+        Requires len(zyx_array.shape) == 3.
+    NA_det: float
         Detection NA.
-    lambda_ill : float
+    lambda_ill: float
         Illumination wavelength
         Units are arbitrary, but must match [pixel_size]
-    pixel_size : float
+    pixel_size: float
         Object-space pixel size = camera pixel size / magnification.
         Units are arbitrary, but must match [lambda_ill]
     midband_fractions: Tuple[float, float], optional
         The minimum and maximum fraction of the cutoff frequency that define the midband.
         Requires: 0 <= midband_fractions[0] < midband_fractions[1] <= 1.
     mode: {'max', 'min'}, optional
-        Option to choose the in-focus slice my minimizing or maximizing the midband frequency.
+        Option to choose the in-focus slice by minimizing or maximizing the midband frequency.
+    plot_path: str or None, optional
+        File name for a diagnostic plot (supports matplotlib filetypes .png, .pdf, .svg, etc.).
+        Use None to skip.
+    threshold_FWHM: float, optional
+        Threshold full-width half max for a peak to be considered in focus.
+        The default value, 0, applies no threshold, and the maximum midband power is always considered in focus.
+        For values > 0, the peak's FWHM must be greater than the threshold for the slice to be considered in focus.
+        If the peak does not meet this threshold, the function returns None.
 
     Returns
     ------
-    slice : int
-        The index of the in-focus slice.
+    slice : int or None
+        If peak's FWHM > peak_width_threshold:
+            return the index of the in-focus slice
+        else:
+            return None
 
     Example
     ------
@@ -45,12 +60,53 @@ def focus_from_transverse_band(
     >>> slice = focus_from_transverse_band(zyx_array, NA_det=0.55, lambda_ill=0.532, pixel_size=6.5/20)
     >>> in_focus_data = data[slice,:,:]
     """
+    minmaxfunc = _check_focus_inputs(
+        zyx_array, NA_det, lambda_ill, pixel_size, midband_fractions, mode
+    )
+
+    # Calculate coordinates
+    _, Y, X = zyx_array.shape
+    _, _, fxx, fyy = util.gen_coordinate((Y, X), pixel_size)
+    frr = np.sqrt(fxx**2 + fyy**2)
+
+    # Calculate fft
+    xy_abs_fft = np.abs(np.fft.fftn(zyx_array, axes=(1, 2)))
+
+    # Calculate midband mask
+    cutoff = 2 * NA_det / lambda_ill
+    midband_mask = np.logical_and(
+        frr > cutoff * midband_fractions[0],
+        frr < cutoff * midband_fractions[1],
+    )
+
+    # Find slice index with min/max power in midband
+    midband_sum = np.sum(xy_abs_fft[:, midband_mask], axis=1)
+    peak_index = minmaxfunc(midband_sum)
+
+    peak_results = peak_widths(midband_sum, [peak_index])
+    peak_FWHM = peak_results[0][0]
+
+    if peak_FWHM > threshold_FWHM:
+        in_focus_index = peak_index
+    else:
+        in_focus_index = None
 
-    # Check inputs
+    # Plot
+    if plot_path is not None:
+        _plot_focus_metric(
+            plot_path, midband_sum, peak_index, in_focus_index, peak_results, threshold_FWHM
+        )
+
+    return in_focus_index
+
+
+def _check_focus_inputs(
+    zyx_array, NA_det, lambda_ill, pixel_size, midband_fractions, mode
+):
     N = len(zyx_array.shape)
     if N != 3:
         raise ValueError(
-            f"{N}D array supplied. `estimate_brightfield_focus` only accepts 3D arrays."
+            f"{N}D array supplied. `focus_from_transverse_band` only accepts 3D arrays."
         )
     if zyx_array.shape[0] == 1:
         print(
@@ -84,22 +140,40 @@ def focus_from_transverse_band(
         minmaxfunc = np.argmax
     else:
         raise ValueError("mode must be either `min` or `max`")
+    return minmaxfunc
 
-    # Calculate coordinates
-    _, Y, X = zyx_array.shape
-    _, _, fxx, fyy = util.gen_coordinate((Y, X), pixel_size)
-    frr = np.sqrt(fxx**2 + fyy**2)
 
-    # Calculate fft
-    xy_abs_fft = np.abs(np.fft.fftn(zyx_array, axes=(1, 2)))
+def _plot_focus_metric(
+    plot_path, midband_sum, peak_index, in_focus_index, peak_results, threshold_FWHM
+):
+    _, ax = plt.subplots(1, 1, figsize=(4, 4))
+    ax.plot(midband_sum, "-k")
+    ax.plot(
+        peak_index,
+        midband_sum[peak_index],
+        "go" if in_focus_index is not None else "ro",
+    )
+    ax.hlines(*peak_results[1:], color="k", linestyles="dashed")
 
-    # Calculate midband mask
-    cutoff = 2 * NA_det / lambda_ill
-    midband_mask = np.logical_and(
-        frr > cutoff * midband_fractions[0],
-        frr < cutoff * midband_fractions[1],
+    ax.set_xlabel("Slice index")
+    ax.set_ylabel("Midband power")
+
+    ax.annotate(
+        f"In-focus slice = {in_focus_index}\n Peak width = {peak_results[0][0]:.2f}\n Peak width threshold = {threshold_FWHM}",
+        xy=(1, 1),
+        xytext=(1.0, 1.1),
+        textcoords="axes fraction",
+        xycoords="axes fraction",
+        ha="right",
+        va="center",
+        annotation_clip=False,
     )
 
-    # Return slice index with min/max power in midband
-    midband_sum = np.sum(xy_abs_fft[:, midband_mask], axis=1)
-    return minmaxfunc(midband_sum)
+    ax.spines["right"].set_visible(False)
+    ax.spines["top"].set_visible(False)
+    ax.spines["left"].set_position(("outward", 10))
+    ax.spines["bottom"].set_position(("outward", 10))
+    ax.ticklabel_format(style="sci", scilimits=(-2, 2))
+
+    print(f"Saving plot to {plot_path}")
+    plt.savefig(plot_path, bbox_inches="tight", dpi=300)
diff --git a/waveorder/optics.py b/waveorder/optics.py
index 7c6a69f..22dbfc0 100644
--- a/waveorder/optics.py
+++ b/waveorder/optics.py
@@ -2,11 +2,10 @@ import numpy as np
 import matplotlib.pyplot as plt
 import gc
 import itertools
-from numpy.fft import fft, ifft, fft2, ifft2, fftn, ifftn, fftshift, ifftshift
+from numpy.fft import fft, fft2, ifft2, fftn, ifftn, fftshift, ifftshift
 
 
 def Jones_sample(Ein, t, sa):
-
     """
     compute output electric field after the interaction between the transparent retardant sample and incident electric field.
 
@@ -39,7 +38,6 @@ def Jones_sample(Ein, t, sa):
 
 
 def Jones_to_Stokes(Ein, use_gpu=False, gpu_id=0):
-
     """
 
     given a coherent electric field, compute the corresponding Stokes vector.
@@ -63,7 +61,6 @@ def Jones_to_Stokes(Ein, use_gpu=False, gpu_id=0):
     """
 
     if use_gpu:
-
         globals()["cp"] = __import__("cupy")
         cp.cuda.Device(gpu_id).use()
 
@@ -93,7 +90,6 @@ def Jones_to_Stokes(Ein, use_gpu=False, gpu_id=0):
 
 
 def analyzer_output(Ein, alpha, beta):
-
     """
 
     compute output electric field after passing through an universal analyzer.
@@ -123,7 +119,6 @@ def analyzer_output(Ein, alpha, beta):
 
 
 def gen_Pupil(fxx, fyy, NA, lambda_in):
-
     """
 
     compute pupil function given spatial frequency, NA, wavelength.
@@ -159,7 +154,6 @@ def gen_Pupil(fxx, fyy, NA, lambda_in):
 
 
 def gen_sector_Pupil(fxx, fyy, NA, lambda_in, sector_angle, rotation_angle):
-
     """
 
     compute sector pupil functions given spatial frequency, NA, wavelength, sector angle, and the rotational angles.
@@ -236,7 +230,6 @@ def gen_sector_Pupil(fxx, fyy, NA, lambda_in, sector_angle, rotation_angle):
 
 
 def Source_subsample(Source_cont, NAx_coord, NAy_coord, subsampled_NA=0.1):
-
     """
 
     compute the sub-sampled source function with the specified sampling spacing in NA coordinate
@@ -277,7 +270,6 @@ def Source_subsample(Source_cont, NAx_coord, NAy_coord, subsampled_NA=0.1):
     first_idx = True
 
     for i in NA_idx:
-
         if first_idx:
             illu_list.append(i)
             first_idx = False
@@ -298,7 +290,6 @@ def Source_subsample(Source_cont, NAx_coord, NAy_coord, subsampled_NA=0.1):
 
 
 def gen_Hz_stack(fxx, fyy, Pupil_support, lambda_in, z_stack):
-
     """
 
     generate propagation kernel
@@ -349,7 +340,6 @@ def gen_Hz_stack(fxx, fyy, Pupil_support, lambda_in, z_stack):
 
 
 def gen_Greens_function_z(fxx, fyy, Pupil_support, lambda_in, z_stack):
-
     """
 
     generate Green's function in u_x, u_y, z space
@@ -406,7 +396,6 @@ def gen_Greens_function_z(fxx, fyy, Pupil_support, lambda_in, z_stack):
 
 
 def gen_dyadic_Greens_tensor_z(fxx, fyy, G_fun_z, Pupil_support, lambda_in):
-
     """
 
     generate forward dyadic Green's function in u_x, u_y, z space
@@ -462,7 +451,6 @@ def gen_dyadic_Greens_tensor_z(fxx, fyy, G_fun_z, Pupil_support, lambda_in):
 
 
 def gen_Greens_function_real(img_size, ps, psz, lambda_in):
-
     """
 
     generate Green's function in real space
@@ -526,7 +514,6 @@ def gen_Greens_function_real(img_size, ps, psz, lambda_in):
 
 
 def gen_dyadic_Greens_tensor(G_real, ps, psz, lambda_in, space="real"):
-
     """
 
     generate dyadic Green's function tensor in real space or in Fourier space
@@ -582,11 +569,9 @@ def gen_dyadic_Greens_tensor(G_real, ps, psz, lambda_in, space="real"):
                 G_tensor[i, i] += G_real_f
 
     if space == "Fourier":
-
         return G_tensor
 
     elif space == "real":
-
         return (
             fftshift(ifftn(G_tensor, axes=(2, 3, 4)), axes=(2, 3, 4))
             / ps
@@ -596,7 +581,6 @@ def gen_dyadic_Greens_tensor(G_real, ps, psz, lambda_in, space="real"):
 
 
 def WOTF_2D_compute(Source, Pupil, use_gpu=False, gpu_id=0):
-
     """
 
     compute 2D weak object transfer function (2D WOTF)
@@ -646,7 +630,6 @@ def WOTF_2D_compute(Source, Pupil, use_gpu=False, gpu_id=0):
         Hp = cp.asnumpy(Hp)
 
     else:
-
         H1 = ifft2(fft2(Source * Pupil).conj() * fft2(Pupil))
         H2 = ifft2(fft2(Source * Pupil) * fft2(Pupil).conj())
         I_norm = np.sum(Source * Pupil * Pupil.conj())
@@ -659,7 +642,6 @@ def WOTF_2D_compute(Source, Pupil, use_gpu=False, gpu_id=0):
 def WOTF_semi_3D_compute(
     Source_support, Source, Pupil, Hz_det, G_fun_z, use_gpu=False, gpu_id=0
 ):
-
     """
 
     compute semi-3D weak object transfer function (semi-3D WOTF)
@@ -723,7 +705,6 @@ def WOTF_semi_3D_compute(
         Hp = cp.asnumpy(Hp)
 
     else:
-
         H1 = ifft2(
             fft2(Source * Pupil * Hz_det).conj() * fft2(Pupil * G_fun_z)
         )
@@ -747,7 +728,6 @@ def WOTF_3D_compute(
     use_gpu=False,
     gpu_id=0,
 ):
-
     """
 
     compute 3D weak object transfer function (2D WOTF)
@@ -834,7 +814,6 @@ def WOTF_3D_compute(
         H_im = cp.asnumpy(H_im)
 
     else:
-
         H1 = ifft2(
             fft2(
                 (Source * Pupil)[:, :, np.newaxis] * Hz_det, axes=(0, 1)
@@ -860,7 +839,6 @@ def WOTF_3D_compute(
 
 
 def gen_geometric_inc_matrix(incident_theta, incident_phi, Source):
-
     """
 
     compute forward and backward matrix mapping from inclination coefficients to retardance with geometric model
@@ -891,7 +869,6 @@ def gen_geometric_inc_matrix(incident_theta, incident_phi, Source):
     geometric_inc_matrix = []
 
     for i in range(N_pattern):
-
         idx_y, idx_x = np.where(Source[i])
 
         geometric_inc_matrix.append(
@@ -930,7 +907,6 @@ def gen_geometric_inc_matrix(incident_theta, incident_phi, Source):
 def SEAGLE_vec_forward(
     E_tot, f_scat_tensor, G_tensor, use_gpu=False, gpu_id=0
 ):
-
     """
 
     compute vectorial SEAGLE forward model
@@ -992,7 +968,6 @@ def SEAGLE_vec_forward(
                 E_in_est[p] += E_tot[p]
 
     else:
-
         pad_convolve_G = lambda x, y, z: ifftn(
             fftn(
                 np.pad(
@@ -1023,7 +998,6 @@ def SEAGLE_vec_forward(
 def SEAGLE_vec_backward(
     E_diff, f_scat_tensor, G_tensor, use_gpu=False, gpu_id=0
 ):
-
     """
 
     compute the adjoint of vectorial SEAGLE forward model
@@ -1087,7 +1061,6 @@ def SEAGLE_vec_backward(
             grad_E[p] = E_diff[p] + E_interact
 
     else:
-
         pad_convolve_G = lambda x, y, z: ifftn(
             fftn(
                 np.pad(
@@ -1120,7 +1093,6 @@ def SEAGLE_vec_backward(
 def scattering_potential_tensor_to_3D_orientation_PN(
     f_tensor, material_type="positive", reg_ret_pr=1e-1
 ):
-
     """
 
     compute principal retardance and 3D orientation from scattering potential tensor components
@@ -1157,7 +1129,6 @@ def scattering_potential_tensor_to_3D_orientation_PN(
     """
 
     if material_type == "positive":
-
         # Positive uniaxial material
 
         azimuth_p = (np.arctan2(-f_tensor[3], -f_tensor[2]) / 2) % np.pi
@@ -1177,7 +1148,6 @@ def scattering_potential_tensor_to_3D_orientation_PN(
         return retardance_pr_p, azimuth_p, theta_p
 
     elif material_type == "negative":
-
         # Negative uniaxial material
 
         azimuth_n = (np.arctan2(f_tensor[3], f_tensor[2]) / 2) % np.pi
@@ -1198,7 +1168,6 @@ def scattering_potential_tensor_to_3D_orientation_PN(
 
 
 def phase_inc_correction(f_tensor0, retardance_pr, theta):
-
     """
 
     compute the inclination-corrected phase from the principal retardance, inclination, and the 0-th component of the scattering potential tensor
@@ -1227,7 +1196,6 @@ def phase_inc_correction(f_tensor0, retardance_pr, theta):
 
 
 def optic_sign_probability(mat_map, mat_map_thres=0.1):
-
     """
 
     compute the optic sign probability from the reconstructed material tendancy
@@ -1259,7 +1227,6 @@ def optic_sign_probability(mat_map, mat_map_thres=0.1):
 def unit_conversion_from_scattering_potential_to_permittivity(
     SP_array, lambda_0, n_media=1, imaging_mode="3D"
 ):
-
     """
 
     compute the optic sign probability from the reconstructed material tendancy
diff --git a/waveorder/waveorder_reconstructor.py b/waveorder/waveorder_reconstructor.py
index a9f3c2c..2908c6e 100644
--- a/waveorder/waveorder_reconstructor.py
+++ b/waveorder/waveorder_reconstructor.py
@@ -6,7 +6,6 @@ import os
 from numpy.fft import fft, ifft, fft2, ifft2, fftn, ifftn, fftshift, ifftshift
 from IPython import display
 from scipy.ndimage import uniform_filter
-from concurrent.futures import ProcessPoolExecutor
 from .util import *
 from .optics import *
 from .background_estimator import *


Test patch:
diff --git a/tests/reconstructor/test_2D_QLIPP_recon.py b/tests/reconstructor/test_2D_QLIPP_recon.py
index 7b5f3a9..c4ea759 100644
--- a/tests/reconstructor/test_2D_QLIPP_recon.py
+++ b/tests/reconstructor/test_2D_QLIPP_recon.py
@@ -1,131 +1,224 @@
 import numpy as np
-import matplotlib.pyplot as plt
-from numpy.fft import fft2, ifft2, fftshift, ifftshift
-
-import pickle
-import waveorder as wo
-
+from numpy.fft import fft2, ifft2
+from waveorder import (
+    optics,
+    waveorder_simulator,
+    waveorder_reconstructor,
+    util,
+)
 
 
 def test_2D_QLIPP_recon():
-    
     """
     Test the whole pipeline of 2D QLIPP simulation and reconstruction using waveorder
-    
+
     """
-    
+
     # Simulation parameters
 
-    N           = 256                    # number of pixel in y dimension
-    M           = 256                    # number of pixel in x dimension
-    mag         = 40                     # magnification
-    ps          = 6.5/mag                # effective pixel size
-    lambda_illu = 0.532                  # wavelength
-    n_media     = 1                      # refractive index in the media
-    NA_obj      = 0.55                   # objective NA
-    NA_illu     = 0.4                    # illumination NA (condenser)
-    NA_illu_in  = 0.4                    # illumination NA (phase contrast inner ring)
-    z_defocus   = (np.r_[:5]-2)*1.757    # a set of defocus plane
-    chi         = 0.03*2*np.pi           # swing of Polscope analyzer
+    N = 256  # number of pixel in y dimension
+    M = 256  # number of pixel in x dimension
+    mag = 40  # magnification
+    ps = 6.5 / mag  # effective pixel size
+    lambda_illu = 0.532  # wavelength
+    n_media = 1  # refractive index in the media
+    NA_obj = 0.55  # objective NA
+    NA_illu = 0.4  # illumination NA (condenser)
+    NA_illu_in = 0.4  # illumination NA (phase contrast inner ring)
+    z_defocus = (np.r_[:5] - 2) * 1.757  # a set of defocus plane
+    chi = 0.03 * 2 * np.pi  # swing of Polscope analyzer
 
-    star, theta, _ = wo.genStarTarget(N,M)
+    star, theta, _ = util.genStarTarget(N, M)
 
     # Assign uniform phase, uniform retardance, and radial slow axes to the star pattern
 
-    phase_value = 1 # average phase in radians (optical path length)
+    phase_value = 1  # average phase in radians (optical path length)
 
-    phi_s = star*(phase_value + 0.15) # slower OPL across target
-    phi_f = star*(phase_value - 0.15) # faster OPL across target
+    phi_s = star * (phase_value + 0.15)  # slower OPL across target
+    phi_f = star * (phase_value - 0.15)  # faster OPL across target
 
-    mu_s = np.zeros((N,M)) # absorption
+    mu_s = np.zeros((N, M))  # absorption
     mu_f = mu_s.copy()
 
-    t_eigen = np.zeros((2, N, M), complex) # complex specimen transmission
+    t_eigen = np.zeros((2, N, M), complex)  # complex specimen transmission
 
-    t_eigen[0] = np.exp(-mu_s + 1j*phi_s) 
-    t_eigen[1] = np.exp(-mu_f + 1j*phi_f) 
+    t_eigen[0] = np.exp(-mu_s + 1j * phi_s)
+    t_eigen[1] = np.exp(-mu_f + 1j * phi_f)
 
-    sa = theta%np.pi #slow axes.
+    sa = theta % np.pi  # slow axes.
 
     # Subsample source pattern for speed
 
-    xx, yy, fxx, fyy = wo.gen_coordinate((N, M), ps)
-    Source_cont = wo.gen_Pupil(fxx, fyy, NA_illu, lambda_illu)
+    xx, yy, fxx, fyy = util.gen_coordinate((N, M), ps)
+    Source_cont = optics.gen_Pupil(fxx, fyy, NA_illu, lambda_illu)
 
+    Source_discrete = optics.Source_subsample(
+        Source_cont, lambda_illu * fxx, lambda_illu * fyy, subsampled_NA=0.1
+    )
 
-    Source_discrete = wo.Source_subsample(Source_cont, lambda_illu*fxx, lambda_illu*fyy, subsampled_NA = 0.1)
-    
     # initiate simulator
-    simulator = wo.waveorder_microscopy_simulator((N,M), lambda_illu, ps, NA_obj, NA_illu, z_defocus, chi, n_media=n_media,\
-                                                  illu_mode='Arbitrary', Source=Source_discrete)
-
-    I_meas, _ = simulator.simulate_waveorder_measurements(t_eigen, sa, multiprocess=False)
-
+    simulator = waveorder_simulator.waveorder_microscopy_simulator(
+        (N, M),
+        lambda_illu,
+        ps,
+        NA_obj,
+        NA_illu,
+        z_defocus,
+        chi,
+        n_media=n_media,
+        illu_mode="Arbitrary",
+        Source=Source_discrete,
+    )
+
+    I_meas, _ = simulator.simulate_waveorder_measurements(
+        t_eigen, sa, multiprocess=False
+    )
 
     # initiate reconstructor
-    setup = wo.waveorder_microscopy((N,M), lambda_illu, ps, NA_obj, NA_illu, z_defocus, chi, n_media=n_media, 
-                                    phase_deconv='2D', bire_in_plane_deconv='2D', illu_mode='BF')
-
+    setup = waveorder_reconstructor.waveorder_microscopy(
+        (N, M),
+        lambda_illu,
+        ps,
+        NA_obj,
+        NA_illu,
+        z_defocus,
+        chi,
+        n_media=n_media,
+        phase_deconv="2D",
+        bire_in_plane_deconv="2D",
+        illu_mode="BF",
+    )
 
     S_image_recon = setup.Stokes_recon(I_meas)
     S_image_tm = setup.Stokes_transform(S_image_recon)
-    Recon_para = setup.Polarization_recon(S_image_tm) # Without accounting for diffraction
-
+    Recon_para = setup.Polarization_recon(
+        S_image_tm
+    )  # Without accounting for diffraction
 
     # Tikhonov regularizer for phase
     reg_u = 1e-5
     reg_p = 1e-5
     S0_stack = S_image_recon[0].copy()
 
-    mu_sample, phi_sample = setup.Phase_recon(S0_stack, method='Tikhonov', reg_u = reg_u, reg_p = reg_p)
-
+    mu_sample, phi_sample = setup.Phase_recon(
+        S0_stack, method="Tikhonov", reg_u=reg_u, reg_p=reg_p
+    )
 
     # TV regularizer for phase
     lambda_u = 3e-3
     lambda_p = 1e-3
     S0_stack = S_image_recon[0].copy()
 
-    mu_sample_TV, phi_sample_TV = setup.Phase_recon(S0_stack, method='TV', lambda_u = lambda_u, lambda_p = lambda_p, itr = 10, rho=1)
-
+    mu_sample_TV, phi_sample_TV = setup.Phase_recon(
+        S0_stack,
+        method="TV",
+        lambda_u=lambda_u,
+        lambda_p=lambda_p,
+        itr=10,
+        rho=1,
+    )
 
     # Diffraction aware reconstruction assuming slowly varying transmission
-    S1_stack = S_image_recon[1].copy()/S_image_recon[0].mean()
-    S2_stack = S_image_recon[2].copy()/S_image_recon[0].mean()
+    S1_stack = S_image_recon[1].copy() / S_image_recon[0].mean()
+    S2_stack = S_image_recon[2].copy() / S_image_recon[0].mean()
 
     # Tikhonov
-    retardance, azimuth = setup.Birefringence_recon_2D(S1_stack, S2_stack, method='Tikhonov', reg_br = 1e-3)
+    retardance, azimuth = setup.Birefringence_recon_2D(
+        S1_stack, S2_stack, method="Tikhonov", reg_br=1e-3
+    )
 
     # TV
-    retardance_TV, azimuth_TV = setup.Birefringence_recon_2D(S1_stack, S2_stack, method='TV', reg_br = 1e-1,\
-                                                          rho = 1e-5, lambda_br=1e-3, itr = 20, verbose=True)
-
+    retardance_TV, azimuth_TV = setup.Birefringence_recon_2D(
+        S1_stack,
+        S2_stack,
+        method="TV",
+        reg_br=1e-1,
+        rho=1e-5,
+        lambda_br=1e-3,
+        itr=20,
+        verbose=True,
+    )
 
     # Compute ground truth images
 
-    phase_gt = (phi_s + phi_f)/2
-    phase_gt = np.real(ifft2(np.mean(fft2(phase_gt)[:,:,np.newaxis]*np.abs(setup.Hp)**2 / (np.abs(setup.Hp)**2+reg_p),axis=2)))
+    phase_gt = (phi_s + phi_f) / 2
+    phase_gt = np.real(
+        ifft2(
+            np.mean(
+                fft2(phase_gt)[:, :, np.newaxis]
+                * np.abs(setup.Hp) ** 2
+                / (np.abs(setup.Hp) ** 2 + reg_p),
+                axis=2,
+            )
+        )
+    )
     phase_gt -= np.mean(phase_gt)
 
-    S1_gt = (phi_s - phi_f)*np.sin(2*sa)
-    S2_gt = (phi_s - phi_f)*np.cos(2*sa)
-
+    S1_gt = (phi_s - phi_f) * np.sin(2 * sa)
+    S2_gt = (phi_s - phi_f) * np.cos(2 * sa)
 
     # test the accuracy of reconstruction using the relative error
 
     # phase_Tikhonov (sensitive to regularization tuning, might cause error if the regularization is not tuned properly)
-    assert(np.sum(np.abs(phase_gt-(phi_sample-np.mean(phi_sample)))**2)/np.sum(np.abs(phase_gt)**2) < 0.1)
+    assert (
+        np.sum(np.abs(phase_gt - (phi_sample - np.mean(phi_sample))) ** 2)
+        / np.sum(np.abs(phase_gt) ** 2)
+        < 0.1
+    )
 
     # phase_TV
-    assert(np.sum(np.abs(phase_gt-(phi_sample_TV-np.mean(phi_sample_TV)))**2)/np.sum(np.abs(phase_gt)**2) < 0.1)
+    assert (
+        np.sum(
+            np.abs(phase_gt - (phi_sample_TV - np.mean(phi_sample_TV))) ** 2
+        )
+        / np.sum(np.abs(phase_gt) ** 2)
+        < 0.1
+    )
 
     # retardance + orientation (no deconvolution)
-    assert(np.sum(np.abs(S1_gt-Recon_para[0,:,:,2]*np.sin(2*Recon_para[1,:,:,2]))**2)/np.sum(np.abs(S1_gt)**2) < 0.1)
-    assert(np.sum(np.abs(S2_gt-Recon_para[0,:,:,2]*np.cos(2*Recon_para[1,:,:,2]))**2)/np.sum(np.abs(S2_gt)**2) < 0.1)
+    assert (
+        np.sum(
+            np.abs(
+                S1_gt
+                - Recon_para[0, :, :, 2] * np.sin(2 * Recon_para[1, :, :, 2])
+            )
+            ** 2
+        )
+        / np.sum(np.abs(S1_gt) ** 2)
+        < 0.1
+    )
+    assert (
+        np.sum(
+            np.abs(
+                S2_gt
+                - Recon_para[0, :, :, 2] * np.cos(2 * Recon_para[1, :, :, 2])
+            )
+            ** 2
+        )
+        / np.sum(np.abs(S2_gt) ** 2)
+        < 0.1
+    )
 
     # retardance + orientation (Tikhonov)
-    assert(np.sum(np.abs(S1_gt-retardance*np.sin(2*azimuth))**2)/np.sum(np.abs(S1_gt)**2) < 0.1)
-    assert(np.sum(np.abs(S2_gt-retardance*np.cos(2*azimuth))**2)/np.sum(np.abs(S2_gt)**2) < 0.1)
+    assert (
+        np.sum(np.abs(S1_gt - retardance * np.sin(2 * azimuth)) ** 2)
+        / np.sum(np.abs(S1_gt) ** 2)
+        < 0.1
+    )
+    assert (
+        np.sum(np.abs(S2_gt - retardance * np.cos(2 * azimuth)) ** 2)
+        / np.sum(np.abs(S2_gt) ** 2)
+        < 0.1
+    )
 
     # retardance + orientation (TV)
-    assert(np.sum(np.abs(S1_gt-retardance_TV*np.sin(2*azimuth_TV))**2)/np.sum(np.abs(S1_gt)**2) < 0.1)
-    assert(np.sum(np.abs(S2_gt-retardance_TV*np.cos(2*azimuth_TV))**2)/np.sum(np.abs(S2_gt)**2) < 0.1)
\ No newline at end of file
+    assert (
+        np.sum(np.abs(S1_gt - retardance_TV * np.sin(2 * azimuth_TV)) ** 2)
+        / np.sum(np.abs(S1_gt) ** 2)
+        < 0.1
+    )
+    assert (
+        np.sum(np.abs(S2_gt - retardance_TV * np.cos(2 * azimuth_TV)) ** 2)
+        / np.sum(np.abs(S2_gt) ** 2)
+        < 0.1
+    )
diff --git a/tests/test_examples.py b/tests/test_examples.py
index ccddd06..1b9f5a8 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -1,6 +1,5 @@
 import subprocess
 import os
-from pathlib import Path
 
 
 def _run_scripts(scripts):
diff --git a/tests/test_focus_estimator.py b/tests/test_focus_estimator.py
index 31d7f7f..560dc58 100644
--- a/tests/test_focus_estimator.py
+++ b/tests/test_focus_estimator.py
@@ -3,8 +3,7 @@ import numpy as np
 from waveorder import focus
 
 
-def test_focus_estimator():
-
+def test_focus_estimator(tmp_path):
     ps = 6.5 / 100
     lambda_ill = 0.532
     NA_det = 1.4
@@ -37,7 +36,42 @@ def test_focus_estimator():
             np.zeros((2, 3, 4)), NA_det, lambda_ill, ps, mode="maxx"
         )
 
+    plot_path = tmp_path.joinpath("test.pdf")
     data3D = np.random.random((11, 256, 256))
-    slice = focus.focus_from_transverse_band(data3D, ps, lambda_ill, NA_det)
+    slice = focus.focus_from_transverse_band(
+        data3D, ps, lambda_ill, NA_det, plot_path=str(plot_path)
+    )
     assert slice >= 0
     assert slice <= data3D.shape[0]
+    assert plot_path.exists()
+
+
+def test_focus_estimator_snr(tmp_path):
+    ps = 6.5 / 100
+    lambda_ill = 0.532
+    NA_det = 1.4
+
+    x = np.linspace(-1, 1, 256)
+    y = np.linspace(-1, 1, 256)
+    z = np.linspace(-1, 1, 21)
+    zz, yy, xx = np.meshgrid(z, y, x, indexing="ij")
+    phantom = (np.sqrt(xx**2 + yy**2 + zz**2) < 0.5).astype(np.uint16)
+
+    for snr in [1000, 100, 10, 1, 0.1]:
+        np.random.seed(1)
+        data = np.random.poisson(
+            phantom * np.sqrt(snr), size=phantom.shape
+        ) + np.random.normal(loc=0, scale=3, size=phantom.shape)
+
+        plot_path = tmp_path / f"test-{snr}.pdf"
+        slice = focus.focus_from_transverse_band(
+            data,
+            ps,
+            lambda_ill,
+            NA_det,
+            plot_path=plot_path,
+            threshold_FWHM=5,
+        )
+        assert plot_path.exists()
+        if slice is not None:
+            assert np.abs(slice - 10) <= 2