From 5c588da8ee37dc241ffdc5b9ed603da204608f2f Mon Sep 17 00:00:00 2001
From: Maddie Cain <maddiecain10@gmail.com>
Date: Wed, 27 Mar 2019 23:53:03 -0400
Subject: [PATCH] dealt with masking in modulus

---
 CDTools/tools/projectors.py    | 12 ++++++++----
 tests/tools/test_projectors.py |  7 ++++++-
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/CDTools/tools/projectors.py b/CDTools/tools/projectors.py
index b2aa8cf..e66f334 100644
--- a/CDTools/tools/projectors.py
+++ b/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):
diff --git a/tests/tools/test_projectors.py b/tests/tools/test_projectors.py
index 36ce13b..49d67c3 100644
--- a/tests/tools/test_projectors.py
+++ b/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():