Update the projectors to be compatible with AD

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]
+

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))