Fix an issue with the modulus projector, update test to catch that issue

This commit is contained in:
Abe Levitan
2019-03-28 15:56:48 -04:00
parent 7f31d402ff
commit 89f96b459e
2 changed files with 14 additions and 9 deletions
+2 -2
View File
@@ -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
+12 -7
View File
@@ -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():