Update the interactions, propagators, and measurements to accept both single wavefields and stacks of wavefields, update the tests as well

CDTools/datasets.py

@@ -177,8 +177,8 @@ class Ptycho_2D_Dataset(CDataset):
     def __len__(self):
         return self.patterns.shape[0]
 
-    def __getitem__(self,index):
-        return index, self.translations[index], self.patterns[index]
+    def __getitem__(self, index):
+        return (index, self.translations[index]), self.patterns[index]
 
 
     def to(self, *args, **kwargs):

CDTools/tools/interactions.py

@@ -26,11 +26,19 @@ def ptycho_2D_round(probe, obj, translations):
     Returns:
         torch.Tensor : An NxMxL tensor of the calculated exit waves
     """
+    single_translation = False
+    if translations.dim() == 1:
+        translations = translations[None,:]
+        single_translation = True
+        
     integer_translations = t.round(translations).to(dtype=t.int32)
     selections = [obj[tr[0]:tr[0]+probe.shape[0],
                       tr[1]:tr[1]+probe.shape[1]]
                   for tr in integer_translations]
-    return t.stack([cmult(probe,selection) for selection in selections])
+    if single_translation:
+        return [cmult(probe,selection) for selection in selections][0]
+    else:
+        return t.stack([cmult(probe,selection) for selection in selections])
 
 
 
@@ -55,7 +63,11 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
     Returns:
         torch.Tensor : An NxMxL tensor of the calculated exit waves
     """
-
+    single_translation = False
+    if translations.dim() == 1:
+        translations = translations[None,:]
+        single_translation = True
+        
     # Separate the translations into a part that chooses the window
     # And a part that defines the windowing function
     integer_translations = t.floor(translations)
@@ -112,7 +124,10 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
 
             exit_waves.append(cmult(probe,selection))
 
-    return t.stack(exit_waves)
+    if single_translation:
+        return exit_waves[0]
+    else:
+        return t.stack(exit_waves)
 
 
 

CDTools/tools/measurements.py

@@ -20,7 +20,7 @@ def intensity(wavefield, detector_slice=None):
     will only include that slice from the simulated wavefront.
     
     Args:
-        wavefield (torch.Tensor) : A complex MxNx2 wavefield
+        wavefield (torch.Tensor) : A JxMxNx2 stack of complex wavefields
         detector_slice (slice) : Optional, a slice or tuple of slices defining a section of the simulation to return
 
     Returns:
@@ -29,7 +29,10 @@ def intensity(wavefield, detector_slice=None):
     if detector_slice is None:
         return cmath.cabssq(wavefield)
     else:
-        return cmath.cabssq(wavefield[detector_slice])
+        if wavefield.dim() == 3:
+            return cmath.cabssq(wavefield[detector_slice])
+        else:
+            return cmath.cabssq(wavefield[(np.s_[:],) + detector_slice])
 
 
 def incoherent_sum(wavefields, detector_slice=None):
@@ -38,9 +41,14 @@ def incoherent_sum(wavefields, detector_slice=None):
     The intensity is defined as the sum of the magnitudes squared of
     the wavefields. If a detector slice is given, the returned array
     will only include that slice from the simulated wavefronts.
+
+    The first index is the index of the diffraction pattern to measure,
+    the second index is the set of incoherently adding patterns, and
+    the next two indices index the wavefield. The final index is the complex
+    index.
     
     Args:
-        wavefields (torch.Tensor) : A JxMxNx2 stack of complex wavefields
+        wavefields (torch.Tensor) : A JxLxMxNx2 stack of complex wavefields
         detector_slice (slice) : Optional, a slice or tuple of slices defining a section of the simulation to return
 
     Returns:
@@ -50,7 +58,10 @@ def incoherent_sum(wavefields, detector_slice=None):
     if detector_slice is None:
         return t.sum(cmath.cabssq(wavefields),dim=-3)
     else:
-        return t.sum(cmath.cabssq(wavefields[(np.s_[:],)+detector_slice]),dim=-3)
+        if wavefields.dim() == 4:
+            return t.sum(cmath.cabssq(wavefields[(np.s_[:],)+detector_slice]),dim=-3)
+        else:
+            return t.sum(cmath.cabssq(wavefields[(np.s_[:],np.s_[:])+detector_slice]),dim=-3)
                  
 
 
@@ -63,7 +74,7 @@ def quadratic_background(wavefield, background, detector_slice=None, measurement
     background model.
 
     Args:
-        wavefield (torch.Tensor) : A complex MxNx2 wavefield
+        wavefield (torch.Tensor) : A JxMxNx2 stack of complex wavefields
         background (torch.Tensor) : An tensor storing the square root of the detector background
         detector_slice (slice) : Optional, a slice or tuple of slices defining a section of the simulation to return
         measurement (function) : Optional, the measurement function to use. The default is measurements.intensity
@@ -76,3 +87,4 @@ def quadratic_background(wavefield, background, detector_slice=None, measurement
     else:
         return measurement(wavefield, detector_slice) \
             + background[detector_slice]**2
+

CDTools/tools/propagators.py

@@ -116,7 +116,6 @@ def near_field(wavefront, angular_spectrum_propagator):
     Returns:
         torch.Tensor : The propagated wavefront 
     """
-
     return t.ifft(cmult(angular_spectrum_propagator,t.fft(wavefront,2)), 2)
 
 

tests/test_datasets.py

@@ -237,7 +237,7 @@ def test_Ptycho_2D_Dataset_ops(ptycho_cxi_1):
     dataset = Ptycho_2D_Dataset.from_cxi(cxi)
 
     assert len(dataset) == expected['data'].shape[0]
-    idx, translation, pattern = dataset[3]
+    (idx, translation), pattern = dataset[3]
     assert idx == 3
     assert t.allclose(translation, t.tensor(expected['translations'][3,:]))
     assert t.allclose(pattern, t.tensor(expected['data'][3,:,:]))

tests/tools/test_interactions.py

@@ -29,35 +29,63 @@ def single_pixel_probe(scope='module'):
 
 
 def test_ptycho_2D_round(random_probe, random_obj):
+    # Test a stack of images
     translations = np.random.rand(10,2) * 500
     exit_waves_np = [random_probe * \
                      random_obj[tr[0]:tr[0]+random_probe.shape[0],
-                                tr[1]:tr[1]+random_probe.shape[1],] for
+                                tr[1]:tr[1]+random_probe.shape[1]] for
                      tr in np.round(translations).astype(int)]
     exit_waves_t = interactions.ptycho_2D_round(cmath.complex_to_torch(random_probe),
                                                 cmath.complex_to_torch(random_obj),
                                                 t.tensor(translations))
     assert np.allclose(cmath.torch_to_complex(exit_waves_t), exit_waves_np)
 
+    # Test the single wave case
+    exit_wave_t = interactions.ptycho_2D_round(cmath.complex_to_torch(random_probe),
+                                                cmath.complex_to_torch(random_obj),
+                                                t.tensor(translations[0]))
+    assert np.allclose(cmath.torch_to_complex(exit_wave_t), exit_waves_np[0])
+
 
 
 def test_ptycho_2D_linear(single_pixel_probe, random_obj):
 
-    # For this one, I just want to check one translation.
+    # For this one, I just want to check one translation, but
+    # I need to check both formats
     translations = np.array([[46.7,53.2]])
-    
+    translation = np.array([46.7,53.2])
+     
     exit_waves_probe = interactions.ptycho_2D_linear(
         cmath.complex_to_torch(single_pixel_probe),
         cmath.complex_to_torch(random_obj),
         t.tensor(translations),
         shift_probe=True)
 
+    exit_wave_probe = interactions.ptycho_2D_linear(
+        cmath.complex_to_torch(single_pixel_probe),
+        cmath.complex_to_torch(random_obj),
+        t.tensor(translation),
+        shift_probe=True)
+
+    # Check that the outputs match
+    assert t.allclose(exit_waves_probe[0],exit_wave_probe)
+
+    
     exit_waves_obj = interactions.ptycho_2D_linear(
         cmath.complex_to_torch(single_pixel_probe),
         cmath.complex_to_torch(random_obj),
         t.tensor(translations),
         shift_probe=False)
 
+    exit_wave_obj = interactions.ptycho_2D_linear(
+        cmath.complex_to_torch(single_pixel_probe),
+        cmath.complex_to_torch(random_obj),
+        t.tensor(translation),
+        shift_probe=False)
+
+    # Check that the outputs match
+    assert t.allclose(exit_waves_obj[0],exit_wave_obj)
+
     # For the shifted probe, we should find 4 pixels with intensity
     exit_waves_probe = cmath.torch_to_complex(exit_waves_probe)[0]
 
@@ -77,3 +105,5 @@ def test_ptycho_2D_linear(single_pixel_probe, random_obj):
     exit_pixel = exit_waves_obj[128,128]
     assert np.isclose(np.sum(obj_shift * obj_section),exit_pixel)
     
+    # Test for a single translation
+    

tests/tools/test_measurements.py

@@ -8,42 +8,57 @@ import pytest
 
 
 def test_intensity():
-    wavefield = t.rand((10,10,2))
-    np_result = np.abs(cmath.torch_to_complex(wavefield))**2
-    assert t.allclose(measurements.intensity(wavefield),
+    wavefields = t.rand((5,10,10,2))
+    np_result = np.abs(cmath.torch_to_complex(wavefields))**2
+    assert t.allclose(measurements.intensity(wavefields),
                       t.tensor(np_result))
 
-    det_slice = np.s_[3:,5:8]
-    assert t.allclose(measurements.intensity(wavefield,det_slice),
-                      t.tensor(np_result[det_slice]))
+    # Test single field case
+    assert t.allclose(measurements.intensity(wavefields[0]),
+                      t.tensor(np_result[0]))
 
+    
+    det_slice = np.s_[3:,5:8]
+    assert t.allclose(measurements.intensity(wavefields,det_slice),
+                      t.tensor(np_result[(np.s_[:],)+det_slice]))
+    
+    # Test single field case
+    assert t.allclose(measurements.intensity(wavefields[0],det_slice),
+                      t.tensor(np_result[0][det_slice]))
 
 
 def test_incoherent_sum():
-    wavefields = t.rand((4,10,10,2))
-    np_result = np.sum(np.abs(cmath.torch_to_complex(wavefields))**2,axis=0)
+    wavefields = t.rand((5,4,10,10,2))
+    np_result = np.sum(np.abs(cmath.torch_to_complex(wavefields))**2,axis=1)
     assert t.allclose(measurements.incoherent_sum(wavefields),
                       t.tensor(np_result))
+    # Test single field case
+    assert t.allclose(measurements.incoherent_sum(wavefields[0]),
+                      t.tensor(np_result[0]))
 
+    
     det_slice = np.s_[3:,5:8]
     assert t.allclose(measurements.incoherent_sum(wavefields,det_slice),
-                      t.tensor(np_result[det_slice]))
-
+                      t.tensor(np_result[(np.s_[:],)+det_slice]))
+    # Test single field case
+    assert t.allclose(measurements.incoherent_sum(wavefields[0],det_slice),
+                      t.tensor(np_result[0][det_slice]))
 
 
 def test_quadratic_background():
     # test with intensity
-    wavefield = t.rand((10,10,2))
+    wavefields = t.rand((5,10,10,2))
     background = t.rand((10,10))
-    np_result = np.abs(cmath.torch_to_complex(wavefield))**2 + background.numpy()**2
+    np_result = np.abs(cmath.torch_to_complex(wavefields))**2 + background.numpy()**2
     det_slice = np.s_[3:,5:8]
 
-    result = measurements.quadratic_background(wavefield,background,
+    result = measurements.quadratic_background(wavefields,background,
                                                detector_slice=det_slice,
                                                measurement=measurements.intensity)
-    assert t.allclose(result, t.tensor(np_result[det_slice]))
-   
-    # test with incoherent sum but no slice
+    assert t.allclose(result, t.tensor(np_result[(np.s_[:],)+det_slice]))
+    
+    
+    # test with incoherent sum but no slice and no stack
     wavefields = t.rand((4,10,10,2))
     np_result = np.sum(np.abs(cmath.torch_to_complex(wavefields))**2,axis=0)
     np_result += background.numpy()**2