def test_apply_field_dependence_model():
''' Test to make sure the field dependence model is giving sensible output
Checks cases comparing master chief ray, center of NIRCam, and center of NIRISS.
'''
rms = lambda array, mask: np.sqrt((array[mask]**2).mean())
# Get the OPD without any sort of field dependence
ote = webbpsf.opds.OTE_Linear_Model_WSS(v2v3=None)
opd_no_field_model = ote.opd.copy()
mask = ote.get_transmission(0) != 0
# Get the OPD at the zero field point of v2 = 0, v3 = -468 arcsec
# Center of NIRCAM fields, but not physically on a detector.
ote.v2v3 = (0, -468) * u.arcsec
ote._apply_field_dependence_model(assume_si_focus=False) # Do this test directly on the OTE OPD, without any implicit focus adjustments
opd_zero_field = ote.opd.copy() * ote.get_transmission(0)
rms1 = rms(opd_no_field_model - opd_zero_field, mask)
assert(rms1 < 7e-9), "OPDs expected to match didn't, at center field (zero field dependence)"
# Get the OPD at some arbitrary nonzero field point (Center of NRC A)
ote.v2v3 = (1.4, -8.2) * u.arcmin
ote._apply_field_dependence_model(assume_si_focus=False)
opd_arb_field = ote.opd.copy() * ote.get_transmission(0)
rms2 = rms(opd_no_field_model - opd_arb_field, mask)
assert(rms2 > 7e-9), "OPDs expected to differ didn't"
assert np.isclose(rms2, 26.1e-9, atol=1e-9), "field dep OPD at center of NIRCam A was not the expected value"
# Now we invoke this via an SI class, to show that works too:
# Get the OPD at the center of NIRISS
nis = webbpsf.NIRISS()
nis.pupilopd = None # disable any global WFE, so we just look at the field dependent part
nis.detector_position = (1024, 1024)
nis, ote_nis = webbpsf.enable_adjustable_ote(nis)
# Test if we directly invoke the OTE model, in this case also disabling SI focus implicit optimization
ote_nis._apply_field_dependence_model(assume_si_focus=False)
opd_nis_cen = ote_nis.opd.copy()
rms3 = rms(opd_nis_cen, mask)
# The value used in the following test is derived from this model itself, so it's a bit circular;
# but at least this test should suffice to detect any unintended significant change in the
# outputs of this model
assert np.isclose(rms3, 36.0e-9, atol=1e-9), "Field-dependent OTE WFE at selected field point (NIRISS center) didn't match expected value (test case: explicit call, assume_si_focus=False)"
# Now test as usd in a webbpsf calculation, implicitly, and with the defocus backout ON
# The WFE here is slightly less, due to the focus optimization
nis = webbpsf.NIRISS()
nis.pupilopd = None # disable any global WFE, so we just look at the field dependent part
nis.detector_position = (1024, 1024)
osys = nis.get_optical_system()
opd_nis_cen_v2 = osys.planes[0].opd.copy()
rms4 = rms(opd_nis_cen_v2, mask)
assert np.isclose(rms4, 28.0e-9, atol=1e-9), "Field-dependent OTE WFE at selected field point (NIRISS center) didn't match expected value(test case: implicit call, assume_si_focus=True."
def main(args):
niriss = webbpsf.NIRISS()
niriss.pupil_mask = "MASK_NRM"
if args.filter in FILTERS:
niriss.filter = args.filter
for i in tqdm(range(10)):
niriss.pupilopd = ('OPD_RevW_ote_for_NIRISS_requirements.fits.gz',
i)
for spectral in tqdm(spectral_types):
src = webbpsf.specFromSpectralType(spectral, catalog="ck04")
niriss.calc_psf(
oversample=args.oversample,
outfile=
f"../data/psf/jwst_{args.filter}_{args.oversample}_psf_OPD{i+1:02}_{spectral}.fits",
display=args.display,
overwrite=True,
normalize='last',
source=src)
if args.display:
plt.show()
elif args.filter == "all":
for f in FILTERS:
niriss.filers = f
niriss.calc_psf(
oversample=args.oversample,
outfile=f"../data/psf/jwst_{f}_{args.oversample}_psf.fits",
display=False,
overwrite=True)
def loicpsf(wavelist=None, wfe_real=None, filepath=''):
''' Utility function which calls the WebbPSF package to create
monochromatic PSFs for NIRISS SOSS obserations and save them to disk.
Parameters
----------
wavelist : list
List of wavelengths (in meters) for which to generate PSFs.
wfe_real : int
Index of wavefront realization to use for the PSF (if non-default
WFE realization is desired).
filepath : str
Path to the directory to which the PSFs will be written.
Defaults to the current directory.
Returns
-------
None : NoneType
PSFs are written to disk.
'''
if wavelist is None:
# List of wavelengths to generate PSFs for
wavelist = np.linspace(0.5, 5.2, 95) * 1e-6
# Dimension of the PSF in native pixels
pixel = 128
# Pixel oversampling factor
oversampling = 10
# Select the NIRISS instrument
niriss = webbpsf.NIRISS()
# Override the default minimum wavelength of 0.6 microns
niriss.SHORT_WAVELENGTH_MIN = 0.5e-6
# Set correct filter and pupil wheel components
niriss.filter = 'CLEAR'
niriss.pupil_mask = 'GR700XD'
# Change the WFE realization if desired
if wfe_real is not None:
niriss.pupilopd = ('OPD_RevW_ote_for_NIRISS_predicted.fits.gz',
wfe_real)
# Loop through all wavelengths to generate PSFs
for wave in wavelist:
print('Calculating PSF at wavelength ', round(wave / 1e-6, 2),
' microns')
psf = niriss.calc_psf(monochromatic=wave,
fov_pixels=pixel,
oversample=oversampling,
display=False)
# Save psf realization to disk
text = '{0:5f}'.format(wave * 1e+6)
psf.writeto(str(filepath) + 'SOSS_os' + str(oversampling) + '_' +
str(pixel) + 'x' + str(pixel) + '_' + text + '_' +
str(wfe_real) + '.fits',
overwrite=True)
return None
def generate_starPSF(FILTER=None, fov=None, osample=None, spectraltype="A0V"):
niriss = wp.NIRISS()
niriss.filter = FILTER
niriss.pupil_mask = 'MASK_NRM'
# set the WFE file to use...
#iriss.pupilopd = ("OPD_RevV_niriss_162.fits", 3) old webbpsf
path_to_webbpsf_data = wp.utils.get_webbpsf_data_path()
opdfile = "OPD_RevW_ote_for_NIRISS_requirements.fits.gz"
opd = os.path.join(path_to_webbpsf_data, "NIRISS", "OPD", opdfile)
opdslice = 3 # anywhere between 0 and 9: 10 realizations...
niriss.pupilopd = (opd, opdslice)
fov_pixels = fov # handoff no refactor
oversample = osample # handoff no refactor
niriss.pixelscale = U.pixscl # handoff no refactor
src = specFromSpectralType(spectraltype)
#Create an oversized array for star PSF.
#sf_fits = niriss.calcPSF(fov_pixels=fov + 4,oversample=osample,source=src,rebin=False,clobber=True) old call webbpsf
psf_fits = niriss.calc_psf(oversample=oversample,
source=src,
fov_pixels=fov_pixels +
4) # +4 because of jittering?
psf_array = psf_fits[0].data
psf_header = psf_fits[0].header
print(psf_array.sum(), 'sum of star psf')
return psf_array, psf_header
def ami_sim_recompute_psf(_filter: str, filename: Union[str, Path],
fov_pixels: int, oversample: int, pupil_mask: str):
"""
Recompute the PSF using webbpsf
:param _filter: str, the filter to use
:param filename: str, the output filename
:param fov_pixels: int, the fov in pixels
:param oversample: int, the oversampling factor
:param pupil_mask: str, the pupil mask to use
:return: None
"""
# get niriss instance from webb psf
niriss = webbpsf.NIRISS()
# set the filter name
niriss.filter = _filter
# set the pupil mask
niriss.pupil_mask = pupil_mask
# TODO: This shouldn't be needed but without it calc_psf breaks?
# TODO: Error --> AttributeError: 'NIRISS' object has no attribute
# TODO: '_extra_keywords'
niriss._extra_keywords = []
# run the psf calculation
niriss.calc_psf(str(filename),
fov_pixels=fov_pixels,
oversample=oversample)
def generate_SOSS_psfs(filt):
"""
Gnerate a cube of the psf at 100 wavelengths from the min to the max wavelength
Parameters
----------
filt: str
The filter to use, ['CLEAR', 'F277W']
"""
try:
import webbpsf
# Get the file
file = os.path.join(PSF_DIR, 'SOSS_{}_PSF.fits'.format(filt))
# Get the NIRISS class from webbpsf and set the filter
ns = webbpsf.NIRISS()
ns.filter = filt
ns.pupil_mask = 'GR700XD'
# Get the min and max wavelengths
wavelengths = utils.wave_solutions('SUBSTRIP256').flatten()
wave_min = np.max([
ns.SHORT_WAVELENGTH_MIN * 1E6,
np.min(wavelengths[wavelengths > 0])
])
wave_max = np.min([
ns.LONG_WAVELENGTH_MAX * 1E6,
np.max(wavelengths[wavelengths > 0])
])
# webbpsf.calc_datacube can only handle 100 but that's sufficient
W = np.linspace(wave_min, wave_max, 100) * 1E-6
# Calculate the psfs
print("Generating SOSS psfs. This takes about 8 minutes...")
start = time.time()
PSF = ns.calc_datacube(W, oversample=1)[0].data
print("Finished in", time.time() - start)
# Make the HDUList
psfhdu = fits.PrimaryHDU(data=PSF)
wavhdu = fits.ImageHDU(data=W * 1E6, name='WAV')
hdulist = fits.HDUList([psfhdu, wavhdu])
# Write the file
hdulist.writeto(file, overwrite=True)
hdulist.close()
except (ImportError, OSError, IOError):
print(
"Could not import `webbpsf` package. Functionality limited. Generating dummy file."
)
def generate_SOSS_psfs(filt):
"""
Gnerate a cube of the psf at 100 wavelengths from the min to the max wavelength
Parameters
----------
filt: str
The filter to use, ['CLEAR','F277W']
"""
# Get the file
file = pkg_resources.resource_filename(
'awesimsoss', 'files/SOSS_{}_PSF.fits'.format(filt))
# Get the NIRISS class from webbpsf and set the filter
ns = webbpsf.NIRISS()
ns.filter = filt
ns.pupil_mask = 'GR700XD'
# Get the min and max wavelengths
wavelengths = wave_solutions(256).flatten()
wave_min = np.max(
[ns.SHORT_WAVELENGTH_MIN * 1E6,
np.min(wavelengths[wavelengths > 0])])
wave_max = np.min(
[ns.LONG_WAVELENGTH_MAX * 1E6,
np.max(wavelengths[wavelengths > 0])])
# webbpsf.calc_datacube can only handle 100 but that's sufficient
W = np.linspace(wave_min, wave_max, 100) * 1E-6
# Calculate the psfs
print("Generating SOSS psfs. This takes about 8 minutes...")
start = time.time()
PSF = ns.calc_datacube(W, oversample=1)[0].data
print("Finished in", time.time() - start)
# Make the HDUList
psfhdu = fits.PrimaryHDU(data=PSF)
wavhdu = fits.ImageHDU(data=W * 1E6, name='WAV')
hdulist = fits.HDUList([psfhdu, wavhdu])
# Write the file
hdulist.writeto(file, overwrite=True)
hdulist.close()
def test_one_psf():
"""Check that setting num_psfs = 1 produces the PSF in the expected location"""
oversample = 2
fov_pixels = 101
# Create 2 cases with different locations: the default center and a set location
inst1 = CreatePSFLibrary(instrument="NIRISS",
filters="F090W",
detectors="NIS",
num_psfs=1,
add_distortion=True,
oversample=oversample,
fov_pixels=fov_pixels,
save=False)
grid1 = inst1.create_files()
inst2 = CreatePSFLibrary(instrument="NIRISS",
filters="F090W",
detectors="NIS",
num_psfs=1,
add_distortion=True,
oversample=oversample,
fov_pixels=fov_pixels,
psf_location=(0, 10),
save=False)
grid2 = inst2.create_files()
assert grid1[0][0].header[
"DET_YX0"] == "(1023.0, 1023.0)" # the default is the integer center of the NIS aperture
assert grid2[0][0].header["DET_YX0"] == "(0.0, 10.0)" # (y,x)
# Compare to the WebbPSF calc_psf output to make sure it's placing the PSF in the right location
nis = webbpsf.NIRISS()
nis.filter = "F090W"
nis.detector_position = (10, 0) # (x,y)
calc = nis.calc_psf(add_distortion=True,
oversample=oversample,
fov_pixels=fov_pixels)
kernel = astropy.convolution.Box2DKernel(width=oversample)
convpsf = astropy.convolution.convolve(calc["OVERDIST"].data, kernel)
assert np.array_equal(convpsf, grid2[0][0].data[0, 0, 0, :, :])
def psf(outputprefix='myPSF_', filter='F430M'):
outputname = outputprefix + filter + '.fits'
nis = webbpsf.NIRISS()
nis.filter = filter
nis.pupil_mask = 'MASK_NRM'
nis.calc_psf(outputname, fov_pixels=77, oversample=11)
def loicpsf(wavelist=None, wfe_real=None, save_to_disk=True):
'''Utility function which calls the WebbPSF package to create monochromatic
PSFs for NIRISS SOSS mode obserations.
Parameters
----------
wavelist : list
List of wavelengths (in meters) for which to generate PSFs.
wfe_real : int
Index of wavefront realization to use for the PSF (if non-default
WFE realization is desired).
save_to_disk : bool
Whether to save PSFs to disk.
Returns
-------
None : NoneType
If PSFs are written to disk.
psf-list : list
List of np.ndarrays with the PSF data.
'''
if wavelist is None:
# List of wavelengths to generate PSFs for
wavelist = np.linspace(0.5, 5.2, 95) * 1e-6
# Dimension of the PSF in native pixels
pixel = 128
# Pixel oversampling factor
oversampling = 10
# Select the NIRISS instrument
niriss = webbpsf.NIRISS()
# Override the default minimum wavelength of 0.6 microns
niriss.SHORT_WAVELENGTH_MIN = 0.5e-6
# Set correct filter and pupil wheel components
niriss.filter = 'CLEAR'
niriss.pupil_mask = 'GR700XD'
# Change the WFE realization if desired
if wfe_real is not None:
niriss.pupilopd = ('OPD_RevW_ote_for_NIRISS_predicted.fits.gz',
wfe_real)
# Loop through all wavelengths to generate PSFs
if save_to_disk is False:
psf_list = [] # Create running list of PSF realizations
for wave in wavelist:
print('Calculating PSF at wavelength ', wave / 1e-6, ' microns')
psf = niriss.calc_psf(monochromatic=wave,
fov_pixels=pixel,
oversample=oversampling,
display=False)
# Save psf realization to disk if desired
if save_to_disk is True:
text = '{0:5f}'.format(wave * 1e+6)
psf.writeto('SOSS_os' + str(oversampling) + '_' + str(pixel) +
'x' + str(pixel) + '_' + text + '.fits',
overwrite=True)
else:
psf_list.append(psf[0].data)
if save_to_disk is False:
return psf_list
else:
return None
wmax=6.0,
instrument=NIRSpec,
outname='NIRSpec_SLIT',
fov_arcsec=3.0,
aperture='Shutter,A200,A400,A1600')
NIRSpec = wp.NIRSpec()
NIRSpec.pixelscale = 0.105
psf_suite(nw=30,
wmin=0.5,
wmax=6.0,
instrument=NIRSpec,
outname='NIRSpec_IFU',
fov_arcsec=3.0,
aperture='IFU')
if doNIRISS:
NIRISS = wp.NIRISS()
NIRISS.pixelscale = 0.0656
psf_suite(nw=30,
wmin=0.5,
wmax=6.0,
instrument=NIRISS,
outname='NIRISS',
fov_arcsec=2.,
aperture='Imager')
# NIRISS = wp.NIRISS()
# NIRISS.pupil_mask = 'GR700XD'
# NIRISS.pixelscale = 0.0656
# psf_suite(nw=20, wmin=0.6, wmax=2.6, instrument=NIRISS, outname='NIRISS_SOSS', fov_arcsec=2., aperture='GR700XD')
def psf(outputname='myPSF.fits'):
nis = webbpsf.NIRISS()
nis.filter = 'F430M'
nis.pupil_mask = 'MASK_NRM'
nis.calc_psf(outputname, fov_pixels=77, oversample=11)
def make_niriss_image(plot=False):
"""
Create a direct image with sources that have been convolved with the
NIRISS PSF from webbpsf as well as the segmentation (ID) image.
Args:
plot (Bool): True if plots should be saved.
"""
# Add only ra-dec, rotation, and wcs to add the sources at the right places
# `size` is image dimension in arcsec
hdu = grizli.utils.make_wcsheader(size=NAXIS[0] * 0.0656,
pixscale=0.0656,
get_hdu=True)
nis_header = hdu.header
nis_header["CRVAL1"] = RA_CENTER
nis_header["CRVAL2"] = DEC_CENTER
# Rotate image to desired PA
nis_wcs = pywcs.WCS(nis_header)
cd_rot = rotate_CD_matrix(nis_wcs.wcs.cd, PA_APER)
for i in range(2):
for j in range(2):
nis_header["CD{0}_{1}".format(i + 1, j + 1)] = cd_rot[i][j]
nis_wcs = pywcs.WCS(nis_header)
nis_wcs.pscale = 0.0656 #grizli.utils.get_wcs_pscale(nis_wcs)
nis_header["PA_APER"] = PA_APER
# Load the catalog and compute detector flux
gfit = Table.read(SOURCE_CATALOG, format='ascii.commented_header')
object_mag = gfit[MAG_COL]
ZP = ZPs[NIS_FILTER.lower()]
object_flux = 10**(-0.4 * (object_mag - ZP))
# Determine which objects are within the NIRISS FoV
xc, yc = nis_wcs.all_world2pix(gfit['ra'], gfit['dec'], 0)
obj_in_img = (xc > 1) & (yc > 1) & (xc < NAXIS[0] - 1) & (yc <
NAXIS[1] - 1)
obj_in_img &= object_flux > 0
# Point sources and faint sources
stars = obj_in_img & ((gfit['star_flag'] == 1) | (object_mag > MAX_MAG))
star_idx = np.arange(len(gfit))[stars]
# Galaxies
if 're' not in gfit.colnames:
gals = obj_in_img < -100 # False
else:
gals = obj_in_img & (~stars) & (object_mag <= MAX_MAG) & (gfit['re'] >
0)
gal_idx = np.arange(len(gfit))[gals]
# Put sources in the image
# Initialize model and segmentation images with zeros
nis_model = np.zeros(NAXIS[::-1], dtype=np.float32)
nis_seg = np.zeros(NAXIS[::-1], dtype=int)
# pixel indices
yp, xp = np.indices(NAXIS[::-1])
# Add star point sources
xpix = np.cast[int](np.round(xc))
ypix = np.cast[int](np.round(yc))
for ix in star_idx:
nis_model[ypix[ix], xpix[ix]] = object_flux[ix]
if SEG_THRESHOLD > 0:
Rseg = 5
for i, ix in enumerate(star_idx):
print('seg: {0} ({1}/{2})'.format(ix, i, stars.sum()))
R = np.sqrt((xp - xpix[ix])**2 + (yp - ypix[ix])**2)
clip_seg = (R <= Rseg) & (nis_seg == 0)
nis_seg[clip_seg] = gfit['id'][ix]
# Add Sersic sources
for i, ix in enumerate(gal_idx):
print('ID: {0} ({1}/{2})'.format(ix, i, gals.sum()))
# Effective radius, in pixels
re = gfit['re'][ix] / nis_wcs.pscale
se = models.Sersic2D(amplitude=1.,
r_eff=re,
n=gfit['n'][ix],
x_0=xc[ix],
y_0=yc[ix],
ellip=1 - gfit['q'][ix],
theta=(gfit['pa'][ix] + 90 - PA_APER) / 180 *
np.pi)
# Normalize to catalog flux
m = se(xp, yp)
renorm = object_flux[ix] / m.sum()
if not np.isfinite(renorm):
continue
# Add to model image
nis_model += m * renorm
if SEG_THRESHOLD > 0:
# Add to segmentation image
clip_seg = (m * renorm > SEG_THRESHOLD) & (nis_seg == 0)
nis_seg[clip_seg] = gfit['id'][ix]
# Check image
if plot is True:
plt.figure(figsize=(10, 10))
plt.subplot(projection=pywcs.WCS(nis_header))
plt.imshow(np.log10(nis_model))
#plt.imshow(nis_seg)
plt.grid(color='black', ls='solid')
plt.xlabel('Galactic Longitude')
plt.ylabel('Galactic Latitude')
plt.title("Sources (logscale)")
figname = "{}_sources.png".format(OUTROOT)
plt.savefig(figname)
print("Wrote {}".format(figname))
#plt.show()
# Convolve with NIRISS PSF and save image and segmentation map.
nis = webbpsf.NIRISS()
# Add -0.5, -0.5 pixel offset to get convolved image in correct place ??
nis.options[
'source_offset_r'] = 0.5 * nis_wcs.pscale # offset in arcseconds
nis.options['source_offset_theta'] = 135. # degrees CCW from +Y
# Get the PSF
nis.filter = NIS_FILTER.upper()
psf_hdu = nis.calcPSF(fov_pixels=64) # power of 2 for fast FFT convolution
psf = psf_hdu[1].data
psf /= psf.sum()
# Convolve the model image
nis_fullsim = stsci.convolve.convolve2d(nis_model,
psf,
output=None,
mode='nearest',
cval=0.0,
fft=1)
# Mask low fluxes to make images gzip smaller
#mask = nis_fullsim > CLIP_FLUX
#nis_fullsim[~mask] = 0
# Save the output image
filename = '{0}-{1}.fits'.format(OUTROOT, nis.filter.lower())
fits.writeto(filename,
data=nis_fullsim,
header=nis_header,
overwrite=True,
output_verify='fix')
print("Wrote {}".format(filename))
# Save the segmentation image
if SEG_THRESHOLD > 0:
filename = '{0}-{1}_seg.fits'.format(OUTROOT, nis.filter.lower())
fits.writeto(filename,
data=nis_seg,
header=nis_header,
overwrite=True,
output_verify='fix')
print("Wrote {}".format(filename))
# Check image.
if plot is True:
plt.subplot(projection=pywcs.WCS(nis_header))
# plt.imshow(nis_fullsim)
plt.imshow(np.log10(nis_fullsim))
#plt.imshow(nis_seg)
plt.grid(color='black', ls='solid')
plt.xlabel('Galactic Longitude')
plt.ylabel('Galactic Latitude')
plt.title("Distorted sources convolved with PSF (logscale)")
figname = "{}_sourcepsf.png".format(OUTROOT)
plt.savefig(figname)
print("Wrote {}".format(figname))
