change imports in examples

mehta-lab · Jul 12, 2023 · 2a0d724 · 2a0d724
1 parent 4f75a5e
commit 2a0d724
Show file tree

Hide file tree

Showing 8 changed files with 122 additions and 82 deletions.
diff --git a/examples/2D_QLIPP_simulation/2D_QLIPP_forward.py b/examples/2D_QLIPP_simulation/2D_QLIPP_forward.py
@@ -13,8 +13,13 @@
 import numpy as np
 import matplotlib.pyplot as plt
 from numpy.fft import fft2, ifft2, fftshift, ifftshift
-
-import waveorder as wo
+from waveorder import (
+    optics,
+    waveorder_simulator,
+    waveorder_reconstructor,
+    visual,
+    util,
+)
 
 # Key parameters
 N = 256  # number of pixel in y dimension
@@ -31,8 +36,8 @@
 
 # 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 +49,7 @@
 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 +63,9 @@
 # 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 +75,7 @@
 # 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
@@ -13,7 +13,14 @@
 import numpy as np
 import matplotlib.pyplot as plt
 from numpy.fft import fft2, ifft2, fftshift, ifftshift
-import waveorder as wo
+from waveorder import (
+    optics,
+    waveorder_simulator,
+    waveorder_reconstructor,
+    visual,
+    util,
+)
+
 
 # ### Load simulated data
 # Load simulations
@@ -34,7 +41,7 @@
 cali = False
 bg_option = "global"
 
-setup = wo.waveorder_microscopy(
+setup = waveorder_reconstructor.waveorder_microscopy(
     (N, M),
     lambda_illu,
     ps,
@@ -53,7 +60,7 @@
     S_image_tm
 )  # Without accounting for diffraction
 
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     np.array(
         [
             Recon_para[0, :, :, L // 2],
@@ -69,7 +76,7 @@
     origin="lower",
 )
 
-wo.plot_hsv(
+visual.plot_hsv(
     [Recon_para[1, :, :, L // 2], Recon_para[0, :, :, L // 2]],
     max_val=1,
     origin="lower",
@@ -86,15 +93,15 @@
     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,
     set_title=True,
     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,15 +117,15 @@
     verbose=True,
 )
 
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     np.array([retardance_TV, azimuth_TV]),
     num_col=2,
     size=10,
     set_title=True,
     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 +137,7 @@
 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 +155,7 @@
 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
@@ -13,7 +13,13 @@
 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 waveorder import (
+    optics,
+    waveorder_simulator,
+    waveorder_reconstructor,
+    visual,
+    util,
+)
 
 # Experiment parameters
 N = 256  # number of pixel in y dimension
@@ -30,8 +36,8 @@
 # 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 @@
 # 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 @@
 
 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 @@
     Source=Source_cont,
 )
 
-simulator = wo.waveorder_microscopy_simulator(
+simulator = waveorder_simulator.waveorder_microscopy_simulator(
     (N, M),
     lambda_illu,
     ps,
@@ -99,7 +105,7 @@
 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 @@
 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
@@ -15,8 +15,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 waveorder import (
+    optics,
+    waveorder_simulator,
+    waveorder_reconstructor,
+    visual,
+    util,
+)
 
 # Load data
 # Load simulations
@@ -47,7 +52,7 @@
 )
 
 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 +73,7 @@
 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
@@ -14,7 +14,13 @@
 from numpy.fft import fft, ifft, fft2, ifft2, fftshift, ifftshift, fftn, ifftn
 
 import pickle
-import waveorder as wo
+from waveorder import (
+    optics,
+    waveorder_simulator,
+    waveorder_reconstructor,
+    visual,
+    util,
+)
 
 #####################################################################
 # Initialization - imaging system and sample                        #
@@ -68,7 +74,7 @@
 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 +86,7 @@
     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 +144,7 @@
 ### Visualize sample properties
 
 #### XY sections
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     [
         target[:, :, z_layer],
         azimuth[:, :, z_layer] % (2 * np.pi),
@@ -190,7 +196,7 @@
     ),
     (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 +205,7 @@
 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 +215,7 @@
 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 +257,7 @@
 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 +333,7 @@
 )
 
 
-wo.plot_multicolumn(
+visual.plot_multicolumn(
     [
         del_f_component[0, :, :, z_layer],
         del_f_component[1, :, :, z_layer],
@@ -361,9 +367,9 @@
 
 # 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
 )
 
@@ -380,7 +386,7 @@
 )
 
 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
 )
 
@@ -414,11 +420,11 @@
 
 #### 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 +436,7 @@
 #### 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,