From 89f96b459ee95fb77224cf7366ebe80a46617023 Mon Sep 17 00:00:00 2001
From: Abe Levitan <alevitan@mit.edu>
Date: Thu, 28 Mar 2019 15:56:48 -0400
Subject: [PATCH] Fix an issue with the modulus projector, update test to catch
 that issue

---
 CDTools/tools/projectors.py    |  4 ++--
 tests/tools/test_projectors.py | 19 ++++++++++++-------
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/CDTools/tools/projectors.py b/CDTools/tools/projectors.py
index 117225b..a4bb10a 100644
--- a/CDTools/tools/projectors.py
+++ b/CDTools/tools/projectors.py
@@ -32,9 +32,9 @@ def modulus(wavefront, intensities, mask = None):
     projected = wavefront * (amplitudes / wavefront_mag)[...,None]
     # Replace amplitude of wavefront with measured amplitude
     if mask is not None:
-        selection = (mask == 0)
+        selection = mask == 0
         # Apply the mask to replace unmasked pixels in the original wavefront
-        projected[selection] = wavefront[selection]
+        projected = projected.masked_scatter(selection, wavefront.masked_select(selection))
 
     return projected
 
diff --git a/tests/tools/test_projectors.py b/tests/tools/test_projectors.py
index 0605567..2b3da30 100644
--- a/tests/tools/test_projectors.py
+++ b/tests/tools/test_projectors.py
@@ -7,18 +7,23 @@ import torch as t
 from scipy.fftpack import fftshift, ifftshift
 
 def test_modulus():
-    # Create a complex array with modulus 12 and phase pi/4
-    np_result = np.sqrt(6) * (1 + 1j) * np.ones((10,10))
+    # Create a complex array with random modulus and known phase
+    np_result = np.sqrt(6) * (1 + 1j) * np.random.rand(10,10)
+    projection_intensity =  t.from_numpy(np.abs(np_result)**2).to(t.float32)
+    original_wavefront = cmath.complex_to_torch((1+1j) * np.random.rand(10,10)).to(t.float32)
     # Test without masks
-    assert np.allclose(cmath.torch_to_complex(projectors.modulus(t.ones((10,10,2)), 12*t.ones((10,10)))), np_result)
+    torch_result = projectors.modulus(original_wavefront,projection_intensity)
+    assert np.allclose(cmath.torch_to_complex(torch_result),np_result)
     
     # Test with mask
     mask = t.ones((10,10,2), dtype = t.uint8)
     mask[5]*=0
-    np_result[5] = 1+1j
-    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)
+    np_result[5] = cmath.torch_to_complex(original_wavefront[5])
+    torch_result = projectors.modulus(original_wavefront,projection_intensity, mask=mask)
+    print(np_result[5])
+    print(cmath.torch_to_complex(torch_result)[5])
+    assert np.allclose(cmath.torch_to_complex(torch_result),np_result)
+
 
 
 def test_support():