From 668f81ab380647ab5ab8bd38fb0237e88ba0bf42 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Mon, 1 Apr 2019 00:18:32 -0400 Subject: [PATCH] Update the interactions, propagators, and measurements to accept both single wavefields and stacks of wavefields, update the tests as well --- CDTools/datasets.py | 4 +-- CDTools/tools/interactions.py | 21 ++++++++++++-- CDTools/tools/measurements.py | 22 +++++++++++---- CDTools/tools/propagators.py | 1 - tests/test_datasets.py | 2 +- tests/tools/test_interactions.py | 36 ++++++++++++++++++++++-- tests/tools/test_measurements.py | 47 +++++++++++++++++++++----------- 7 files changed, 102 insertions(+), 31 deletions(-) diff --git a/CDTools/datasets.py b/CDTools/datasets.py index d4d364b..c2e83f3 100644 --- a/CDTools/datasets.py +++ b/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): diff --git a/CDTools/tools/interactions.py b/CDTools/tools/interactions.py index ef5ad0d..8593d6e 100644 --- a/CDTools/tools/interactions.py +++ b/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) diff --git a/CDTools/tools/measurements.py b/CDTools/tools/measurements.py index fdcbbf1..9b21cd6 100644 --- a/CDTools/tools/measurements.py +++ b/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 + diff --git a/CDTools/tools/propagators.py b/CDTools/tools/propagators.py index d3eed3c..d13fe23 100644 --- a/CDTools/tools/propagators.py +++ b/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) diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 02dc473..52620c7 100644 --- a/tests/test_datasets.py +++ b/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,:,:])) diff --git a/tests/tools/test_interactions.py b/tests/tools/test_interactions.py index 0c6db8e..810d532 100644 --- a/tests/tools/test_interactions.py +++ b/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 + diff --git a/tests/tools/test_measurements.py b/tests/tools/test_measurements.py index 93e45dd..16db152 100644 --- a/tests/tools/test_measurements.py +++ b/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