From 7f31d402ff7bb122020eb04342c7044405338cf0 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Thu, 28 Mar 2019 00:49:33 -0400 Subject: [PATCH] Update the projectors to be compatible with AD --- CDTools/tools/projectors.py | 36 +++++++++++++++------------------- tests/tools/test_projectors.py | 22 ++++++++++++++------- 2 files changed, 31 insertions(+), 27 deletions(-) diff --git a/CDTools/tools/projectors.py b/CDTools/tools/projectors.py index e66f334..117225b 100644 --- a/CDTools/tools/projectors.py +++ b/CDTools/tools/projectors.py @@ -21,35 +21,31 @@ def modulus(wavefront, intensities, mask = None): wavefront (torch.Tensor) : The JxNxMx2 stack of complex propagated wavefronts intensities (torch.Tensor): The measured diffraction pattern(s) stored as an JxNxM stack of real tensors mask (torch.Tensor) : Mask for the intensities array with shape JxNxM, where bad detector pixels are set to 0 and usable pixels set to 1 + Returns: torch.Tensor : The JxNxMx2 propagated wavefield with corrected intensities """ # Calculate amplitudes from intensities - amplitudes = intensities**.5 + amplitudes = t.sqrt(intensities) # Normalize wavefront so the complex elements have modulus one - abs = cabs(wavefront) - if mask is not None: - # Record the original wavefront without amplitude replacements - original_wavefront = wavefront.clone() + wavefront_mag = cabs(wavefront) + projected = wavefront * (amplitudes / wavefront_mag)[...,None] # Replace amplitude of wavefront with measured amplitude - wavefront[...,0]/=abs - wavefront[...,1]/=abs - wavefront[...,0]*=amplitudes - wavefront[...,1]*=amplitudes - if mask is None: - return wavefront - else: + if mask is not None: + selection = (mask == 0) # Apply the mask to replace unmasked pixels in the original wavefront - return original_wavefront.masked_scatter_(mask, wavefront) + projected[selection] = wavefront[selection] + + return projected def support(wavefront, support): """Implements the support constraint in torch - This accepts a torch tensor representing the propagated simulated wavefront(s), + This accepts a torch tensor representing (a) simulated wavefield(s), where the last dimension represents the real and imaginary components of the propagated wavefield(s). It projects the support of the imaged object - onto the simulated wavefront via a mask. + onto the simulated wavefront via a support mask. It assumes that the wavefront is stored in an array [i,j] where i corresponds to the y-axis and j corresponds to the @@ -58,10 +54,10 @@ def support(wavefront, support): Args: wavefront (torch.Tensor) : The JxNxMx2 stack of complex propagated wavefronts - mask (torch.Tensor) : Mask for the intensities array with shape JxNxM, where bad detector pixels are set to 0 and usable pixels set to 1 + support (torch.Tensor) : An NxM support, with 1s within the support and 0s outside + Returns: - torch.Tensor : The JxNxMx2 wavefield with the mask applied + torch.Tensor : The JxNxMx2 wavefield with the support mask applied """ - wavefront[...,0] *= support - wavefront[...,1] *= support - return wavefront + return wavefront * support.to(wavefront.dtype)[...,None] + diff --git a/tests/tools/test_projectors.py b/tests/tools/test_projectors.py index 49d67c3..0605567 100644 --- a/tests/tools/test_projectors.py +++ b/tests/tools/test_projectors.py @@ -8,19 +8,27 @@ from scipy.fftpack import fftshift, ifftshift def test_modulus(): # Create a complex array with modulus 12 and phase pi/4 - np_result = 6**.5*(np.ones((10,10))+1j*np.ones((10,10))) + np_result = np.sqrt(6) * (1 + 1j) * np.ones((10,10)) # Test without masks - assert(np.allclose(cmath.torch_to_complex(projectors.modulus(t.ones((10,10,2)), 12*t.ones((10,10)))), np_result)) + assert np.allclose(cmath.torch_to_complex(projectors.modulus(t.ones((10,10,2)), 12*t.ones((10,10)))), np_result) + # Test with mask mask = t.ones((10,10,2), dtype = t.uint8) mask[5]*=0 np_result[5] = 1+1j - assert(np.allclose(cmath.torch_to_complex(projectors.modulus(t.ones((10,10,2)), 12*t.ones((10,10)), mask = mask)), np_result)) + print(mask) + print(cmath.torch_to_complex(projectors.modulus(t.ones((10,10,2)), 12*t.ones((10,10)), mask = mask))) + assert np.allclose(cmath.torch_to_complex(projectors.modulus(t.ones((10,10,2)), 12*t.ones((10,10)), mask = mask)), np_result) def test_support(): - # Test masking - support = t.zeros((10,10)) - np_result = np.zeros((10,10)) - + # Define a mask as uint8 to make sure it works when support is not + # the same type as the wavefield + support = t.zeros((10,10)).to(t.uint8) + # some masked, some unmasked + support[:3,:3] = 1 + + np_result = np.zeros((10,10)).astype(np.complex128) + np_result[:3,:3] = 1 + 1j + assert(np.allclose(cmath.torch_to_complex(projectors.support(t.ones((10,10,2)), support)), np_result))