More closely hit requested RMS

The power spectrum-based WFE method accepts an RMS attribute. While the
resulting phase map was getting somewhat close to the requested RMS, it
regularly coming in quite a bit lower than the requested value. This
commit does not guarantee that the requested RMS is hit exactly, but it
gets a lot closer than before.

lentil/wfe.py

@@ -59,7 +59,7 @@ def power_spectrum(mask, pixelscale, rms, half_power_freq, exp, seed=None):
     # Generate the PSD function and noise filter
     psd = 1 / (1 + (dr / half_power_freq) ** exp)
     psd[dr == 0] = 0
-    psd = psd * rms ** 2 / np.sum(psd)
+    psd = psd / np.sum(psd)
 
     H = np.fft.fftshift(np.sqrt(psd))
 
@@ -68,9 +68,12 @@ def power_spectrum(mask, pixelscale, rms, half_power_freq, exp, seed=None):
     # Generate noise, filter it to realize the requested PSD, and enforce
     # the pupil mask
     noise = rng.normal(size=[n, m])
-    opd = np.real(np.fft.ifft2(np.fft.fft2(noise) * H)) * np.sqrt(m * n)
+    phase = np.real(np.fft.ifft2(np.fft.fft2(noise) * H)) * np.sqrt(m * n)
 
-    return opd * mask
+    phase *= mask
+    phase = phase * np.sqrt(np.count_nonzero(phase)/np.sum(np.abs(phase)**2)) * rms
+
+    return phase
 
 
 def translation_defocus(mask, f_number, translation):

tests/wfe/test_wfe.py

@@ -2,6 +2,15 @@
 import lentil
 
 
+def test_power_spectrum_wfe():
+    mask = lentil.util.circlemask((256, 256), 128)
+    rms = 50e-9
+    phase = lentil.wfe.power_spectrum(mask, pixelscale=1/256, rms=rms,
+                                      half_power_freq=5, exp=3)
+
+    assert np.std(phase[np.nonzero(phase)])/rms >= 0.8
+
+
 def test_translation_defocus():
     mask = lentil.util.circlemask((256, 256), 128)
     f_number = np.random.normal(10)