Adding odak.learn.wave.incoherent_angular_spectrum() test case in tes…

…t/test_learn_wave_compare_beam_propagations.py.
kaanaksit · Sep 28, 2024 · 3eb6ec5 · 3eb6ec5
1 parent 12d36a1
commit 3eb6ec5
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 5 deletions.
diff --git a/odak/learn/wave/classical.py b/odak/learn/wave/classical.py
@@ -126,13 +126,52 @@ def get_propagation_kernel(
                          Complex kernel for the given propagation type.
     """
     if propagation_type == 'Bandlimited Angular Spectrum':
-        kernel = get_band_limited_angular_spectrum_kernel(nu, nv, dx, wavelength, distance, device)
+        kernel = get_band_limited_angular_spectrum_kernel(
+                                                          nu = nu,
+                                                          nv = nv,
+                                                          dx = dx,
+                                                          wavelength = wavelength,
+                                                          distance = distance,
+                                                          device = device
+                                                         )
     elif propagation_type == 'Angular Spectrum':
-        kernel = get_angular_spectrum_kernel(nu, nv, dx, wavelength, distance, device)
+        kernel = get_angular_spectrum_kernel(
+                                             nu = nu,
+                                             nv = nv,
+                                             dx = dx,
+                                             wavelength = wavelength,
+                                             distance = distance,
+                                             device = device
+                                            )
     elif propagation_type == 'Transfer Function Fresnel':
-        kernel = get_transfer_function_fresnel_kernel(nu, nv, dx, wavelength, distance, device)
+        kernel = get_transfer_function_fresnel_kernel(
+                                                      nu = nu,
+                                                      nv = nv,
+                                                      dx = dx,
+                                                      wavelength = wavelength,
+                                                      distance = distance,
+                                                      device = device
+                                                     )
     elif propagation_type == 'Impulse Response Fresnel':
-        kernel = get_impulse_response_fresnel_kernel(nu, nv, dx, wavelength, distance, device, scale = scale, aperture_samples = samples)
+        kernel = get_impulse_response_fresnel_kernel(
+                                                     nu = nu, 
+                                                     nv = nv, 
+                                                     dx = dx, 
+                                                     wavelength = wavelength,
+                                                     distance = distance,
+                                                     device =  device,
+                                                     scale = scale,
+                                                     aperture_samples = samples
+                                                    )
+    elif propagation_type == 'Incoherent Angular Spectrum':
+        kernel = get_incoherent_angular_spectrum_kernel(
+                                                        nu = nu,
+                                                        nv = nv, 
+                                                        dx = dx, 
+                                                        wavelength = wavelength, 
+                                                        distance = distance,
+                                                        device = device
+                                                       )
     else:
         logging.warning('Propagation type not recognized')
         assert True == False

diff --git a/test/test_learn_wave_compare_beam_propagations.py b/test/test_learn_wave_compare_beam_propagations.py
@@ -10,12 +10,13 @@ def test(device = torch.device('cpu'), output_directory = 'test_output'):
     wavelength = 532e-9
     pixel_pitch = 3.74e-6
     distance = 1e-3
-    aperture_samples = [5, 5, 1, 1] # Replace it with this: [50, 50, 5, 5]
+    aperture_samples = [35, 35, 1, 1] # Replace it with this: [50, 50, 5, 5]
     propagation_types = [
                          'Impulse Response Fresnel',
                          'Transfer Function Fresnel',
                          'Angular Spectrum',
                          'Bandlimited Angular Spectrum',
+                         'Incoherent Angular Spectrum',
                         ]
     k = odak.learn.wave.wavenumber(wavelength)