dealt with masking in modulus

CDTools/tools/projectors.py

@@ -28,15 +28,19 @@ def modulus(wavefront, intensities, mask = None):
     amplitudes = intensities**.5
     # 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()
+    # Replace amplitude of wavefront with measured amplitude
     wavefront[...,0]/=abs
     wavefront[...,1]/=abs
+    wavefront[...,0]*=amplitudes
+    wavefront[...,1]*=amplitudes
     if mask is None:
-        # Replace amplitude of wavefront with measured amplitude
-        wavefront[...,0]*=amplitudes
-        wavefront[...,1]*=amplitudes
         return wavefront
     else:
-        return wavefront[mask != 0]
+        # Apply the mask to replace unmasked pixels in the original wavefront
+        return original_wavefront.masked_scatter_(mask, wavefront)
 
 
 def support(wavefront, support):

tests/tools/test_projectors.py

@@ -9,8 +9,13 @@ 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)))
-
+    # Test without masks
     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))
 
 
 def test_support():