def test_move_sur(plot=False):
""" Test we can move mirrors using Segment Update Requests
"""
import webbpsf
import os
import glob
surdir = os.path.join(webbpsf.__path__[0], 'tests', 'surs')
surs = glob.glob(surdir+'/*sur.xml')
nrc = webbpsf.NIRCam()
nrc.filter='F212N'
nrc, ote = webbpsf.enable_adjustable_ote(nrc)
ote.zero(zero_original=True)
for s in surs:
print("Testing "+s)
ote.reset()
ote.move_sur(s)
# the coarse phasing SUR is a no-op after 3 groups; all others have some effect
if 'coarse_phasing' not in s:
assert not np.allclose(ote.segment_state, 0), "Expected some segments to be moved"
ote.move_sur(s, reverse=True)
assert np.allclose(ote.segment_state, 0), "Reversing moves didn't bring us back to zero"
# Test every DOF on A1-1 and SM and check the OTE state updated accordingly
s = glob.glob(surdir+'/example_alldof_A1-SM_sur.xml')[0]
print("Testing "+s)
ote.reset()
ote.move_sur(s)
assert np.allclose(ote.segment_state[0], [1, 2, 3, 4, 5, 6])
assert np.allclose(ote.segment_state[-1], [1, 2, 3, 4, 5, 0])
# Test moving one at a time. This test relies on specifics of what's in the image stacking SUR.
s = glob.glob(surdir+'/example_image_stacking*sur.xml')[0]
print("Testing moving one group at a time with "+s)
ote.reset()
sur = webbpsf.surs.SUR(s)
ngroups = len(sur.groups)
oldstate = ote.segment_state.copy()
for igrp in range(1, ngroups+1):
print("Group {} should move segment {}".format(igrp, 2*igrp+6))
ote.move_sur(s, group=igrp)
movedsegs = np.abs((ote.segment_state - oldstate).sum(axis=1))
assert (movedsegs!=0).sum()==1, "Only expected one segment to move"
whichmoved = np.argmax(movedsegs)+1
print ("Moved segment", whichmoved)
assert whichmoved == 2*igrp+6, "An unexpected segment moved"
oldstate = ote.segment_state.copy()
if plot:
psf = nrc.calc_psf(fov_pixels=256, add_distortion=False)
plt.figure()
ote.display_opd(title="After Group {}".format(igrp))
plt.figure()
webbpsf.display_psf(psf, ext=1, title="After Group {}".format(igrp))
def test_defocus(fov_arcsec=1, display=False):
"""Test we can apply a defocus to a PSF
via either a weak lens, or via the options dict,
and we get consistent results either way.
Test for #59 among other things
"""
nrc = webbpsf_core.NIRCam()
nrc.pupilopd=None
nrc.include_si_wfe=False
# Calculate defocus with a weak lens
nrc.pupil_mask = 'WEAK LENS +4'
psf = nrc.calc_psf(nlambda=1, fov_arcsec=fov_arcsec, oversample=1, display=False, add_distortion=False)
# Calculate equivalent via the options structure
nrc.pupil_mask = None
nrc.options['defocus_waves']=3.9024 # as measured
nrc.options['defocus_wavelength']=2.12e-6
psf_2 = nrc.calc_psf(nlambda=1, fov_arcsec=fov_arcsec, oversample=1, display=False, add_distortion=False)
assert np.allclose(psf[0].data, psf_2[0].data), "Defocused PSFs calculated two ways don't agree"
if display:
import webbpsf
plt.figure()
webbpsf.display_psf(psf)
plt.figure()
webbpsf.display_psf(psf_2)
def test_defocus(fov_arcsec=1, display=False):
"""Test we can apply a defocus to a PSF
via either a weak lens, or via the options dict,
and we get consistent results either way.
Note this is now an *inexact* comparison, because the weak lenses now include non-ideal effects, in particular field dependent astigmatism
Test for #59 among other things
"""
nrc = webbpsf_core.NIRCam()
nrc.set_position_from_aperture_name('NRCA3_FP1')
nrc.pupilopd = None
nrc.include_si_wfe = False
# Calculate defocus with a weak lens
nrc.pupil_mask = 'WEAK LENS +4'
psf = nrc.calc_psf(nlambda=1,
fov_arcsec=fov_arcsec,
oversample=1,
display=False,
add_distortion=False)
# Calculate equivalent via the options structure
nrc.pupil_mask = None
nrc.options['defocus_waves'] = 3.9024 # as measured
nrc.options['defocus_wavelength'] = 2.12e-6
psf_2 = nrc.calc_psf(nlambda=1,
fov_arcsec=fov_arcsec,
oversample=1,
display=False,
add_distortion=False)
assert np.allclose(
psf[0].data, psf_2[0].data, atol=1e-4
), "Defocused PSFs calculated two ways don't agree as precisely as expected"
if display:
import webbpsf
plt.figure()
webbpsf.display_psf(psf)
plt.figure()
webbpsf.display_psf(psf_2)
def display_ote_and_psf(inst, ote, opd_vmax=500, psf_vmax=0.1, title="OPD and PSF", **kwargs):
"""
Display OTE and PSF of a JWST instrument next to each other.
Adapted from:
https://github.com/spacetelescope/webbpsf/blob/develop/notebooks/Simulated%20OTE%20Mirror%20Move%20Demo.ipynb
:param inst: WebbPSF instrument instance, e.g. webbpsf.NIRCam()
:param ote: OTE of inst, usually obtained with: instrument, ote = webbpsf.enable_adjustable_ote(instrument)
:param opd_vmax: float, max display value for the OPD
:param psf_vmax: float, max display valued for PSF
:param title: string, plot title
:param kwargs:
"""
psf = inst.calc_psf(nlambda=1)
plt.figure(figsize=(12, 8))
ax1 = plt.subplot(121)
ote.display_opd(ax=ax1, vmax=opd_vmax, colorbar_orientation='horizontal', title='OPD with aberrated segments')
ax2 = plt.subplot(122)
webbpsf.display_psf(psf, ext=2, vmax=psf_vmax, vmin=psf_vmax/1e4, colorbar_orientation='horizontal', title="PSF simulation")
plt.suptitle(title, fontsize=16)
def display_ote_and_psf(inst,
ote,
opd_vmax=500,
psf_vmax=0.1,
title="OPD and PSF",
**kwargs):
import matplotlib.pyplot as plt
import webbpsf
psf = inst.calc_psf(monochromatic=2e-6, )
plt.figure(figsize=(12, 8))
ax1 = plt.subplot(121)
ote.display_opd(ax=ax1,
vmax=opd_vmax,
colorbar_orientation='horizontal',
title='OPD modified for mirror moves') #, cbpad=0.05)
ax2 = plt.subplot(122)
webbpsf.display_psf(psf,
ext=1,
vmax=psf_vmax,
vmin=psf_vmax / 1e4,
colorbar_orientation='horizontal',
title="PSF sim, 2 microns") #, cbpad=0.05)
plt.suptitle(title, fontsize=16)
long = 6
wide = 2
fig, axs = plt.subplots(wide, long, figsize=(12, 6), sharey=True)
for i, filter in enumerate(sorted(all_filters)):
r = int(np.floor(i / long))
c = (i % long if i < long else
(i % long) + 1) # remove else for left-justified bottom row
ax = axs[r][c]
wfi.filter = filter
psf = wfi.calc_psf(oversample=4)
display_psf(psf, ax=ax, colorbar=False, title=filter)
ax.title.set_fontsize(20)
ax.tick_params(axis='both', labelsize=10)
ax.xaxis.label.set_visible(False)
ax.yaxis.label.set_visible(False)
axs[-1][0].remove() # change last index to -1 to justify left
fig.tight_layout(w_pad=.1, h_pad=0)
fig.tight_layout(w_pad=.1, h_pad=0) # calling twice somehow tightens h_pad
# fig.savefig('webbpsf-roman_page_header.png', dpi=100, facecolor='w')
#### Create compare_wfi_sca09_sca17.png
wfi2 = roman.WFI()
wfi2.filter = 'F129'
wfi2.detector = 'SCA09'
# Write model image with point sources to fits
fits.writeto('model.fits', data, overwrite=True)
# Create custom convolution Kernel from PSF
try:
psf_kern = CustomKernel(psf_data)
except:
print('Adding dummy row and column to psf data for CustomKernel to work.')
psfnpix_y = psf_data.shape[0] + 1
psfnpix_x = psf_data.shape[1] + 1
new_psf = np.zeros((psfnpix_y, psfnpix_x), dtype=np.float64)
new_psf[:psfnpix_y - 1, :psfnpix_x - 1] = psf_data
#hdu_new_psf = fits.PrimaryHDU(new_psf)
psf_header['NAXIS1'] = psfnpix_y
psf_header['NAXIS2'] = psfnpix_x
fits.writeto('new_psf.fits', new_psf, psf_header, overwrite=True)
psf_kern = CustomKernel(new_psf)
#Convolve model image with PSF and save FITS file
convl_image = convolve_fft(data, psf_kern, allow_huge=True)
im_head = psf_header
im_head['NAXIS1'] = convl_image.shape[0]
im_head['NAXIS2'] = convl_image.shape[1]
fits.writeto('convl_model.fits', convl_image, im_head, overwrite=True)
#Display on screen and save png
webbpsf.display_psf('convl_model.fits')
plt.savefig('convl_model.png')
plt.show()
# Script that uses WebbPSF to create PSF to use in simulations
# Needs to run in webbpsf-env to avoid conflicts
import webbpsf
import matplotlib.pyplot as plt
filter = 'F200W'
nc = webbpsf.NIRCam()
nc.filter = filter
psf = nc.calc_psf(outfile='{}_psf.fits'.format(filter),oversample=2,fov_arcsec=5.0)
webbpsf.display_psf(psf)
plt.show()