In [1]:
# Import the usual libraries
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
#import matplotlib.patches as mpatches
# Enable inline plotting
%matplotlib inline
#from IPython.display import display, Latex, clear_output
In [2]:
# WebbPSF
import pynrc, webbpsf, os
# Interpolation and extrapolation
from scipy.interpolate import griddata, RegularGridInterpolator
from scipy.ndimage import rotate
# pySIAF stuff for plotting
from pysiaf.siaf import Siaf
from pysiaf.siaf import plot_main_apertures
# Astropy Tables and FITS
from astropy.table import Table, join
from astropy.io import fits
In [3]:
# Directory and file information
outdir = pynrc.conf.path + 'opd_mod/'
# Read in measured SI Zernike data
data_dir = webbpsf.utils.get_webbpsf_data_path() + '/'
zernike_file = data_dir + 'si_zernikes_isim_cv3.fits'
# Read in table data
ztable_full = Table.read(zernike_file)
# Zernike names
keys = np.array(ztable_full.keys())
ind_z = ['Zernike' in k for k in keys]
zkeys = keys[ind_z]
In [4]:
def get_v2v3_limits():
"""
V2/V3 Limits for a given module stored within an dictionary
"""
names = ['SW', 'LW', 'SWA', 'LWA', 'SWB', 'LWB']
apnames = ['NRCALL_FULL', 'NRCALL_FULL',
['NRCA1_FULL', 'NRCA2_FULL', 'NRCA3_FULL', 'NRCA4_FULL'], 'NRCA5_FULL',
['NRCB1_FULL', 'NRCB2_FULL', 'NRCB3_FULL', 'NRCB4_FULL'], 'NRCB5_FULL']
siaf = Siaf('NIRCam')
siaf.generate_toc()
v2v3_limits = {}
for i, name in enumerate(names):
# Do all four apertures for each SWA & SWB
if (name=='SWA') or (name=='SWB'):
v2_vals = []
v3_vals = []
for apsw_name in apnames[i]:
ap = siaf[apsw_name]
v2_ref, v3_ref = ap.closed_polygon_points('tel', rederive=False)
v2_vals.append(v2_ref)
v3_vals.append(v3_ref)
v2_ref = np.array(v2_vals).flatten()
v3_ref = np.array(v3_vals).flatten()
else:
ap = siaf[apnames[i]]
v2_ref, v3_ref = ap.closed_polygon_points('tel', rederive=False)
# Add a 10" border margin
v2_min = int(v2_ref.min()-10)
v2_max = int(v2_ref.max()+10)
v3_min = int(v3_ref.min()-10)
v3_max = int(v3_ref.max()+10)
v2v3_limits[name] = {'V2': [v2_min, v2_max],
'V3': [v3_min, v3_max]}
return v2v3_limits
In [5]:
# V2/V3 Limits stored in a dictionary
v2v3_limits = get_v2v3_limits()
In [6]:
def gen_cv3_zgrid(ztable_full, v2v3_limits, mod, zkey):
"""Gridded data from CV3 only"""
ind_nrc = ['NIRCam'+mod in row['instrument'] for row in ztable_full]
ind_nrc = np.where(ind_nrc)
# Grab V2/V3 coordinates
# In units of arcmin
v2 = ztable_full[ind_nrc]['V2']
v3 = ztable_full[ind_nrc]['V3']
# Create finer mesh grid
v2_lims = np.array(v2v3_limits[mod]['V2']) / 60.
v3_lims = np.array(v2v3_limits[mod]['V3']) / 60.
dstep = 1. / 60. # 1" steps
xgrid = np.arange(v2_lims[0], v2_lims[1]+dstep, dstep)
ygrid = np.arange(v3_lims[0], v3_lims[1]+dstep, dstep)
X, Y = np.meshgrid(xgrid,ygrid)
ztable = ztable_full[ind_nrc]
zvals = ztable[zkey]
# There will be some NaNs along the outer borders
zgrid = griddata((v2, v3), zvals, (X, Y), method='cubic')
return xgrid, ygrid, zgrid
In [7]:
mod = 'LW'
xgrid, ygrid, zgrid = gen_cv3_zgrid(ztable_full, v2v3_limits, mod, 'Zernike_4')
In [8]:
fig, ax = plt.subplots(1,1, figsize=(8,5))
zmin, zmax = np.nanmin(zgrid), np.nanmax(zgrid)
extent = np.array([xgrid.min(),xgrid.max(),ygrid.min(),ygrid.max()]) * 60
ax.imshow(zgrid, extent=extent, vmin=zmin, vmax=zmax)
siaf = Siaf('NIRCam')
ap = siaf['NRCALL_FULL']
ap.plot(ax=ax, fill=False, color='C1')
fig.tight_layout()
In [9]:
# To extrapolate outside the measured field points, we proceed
# in two steps. This first creates a fine-meshed cubic fit
# over the known field points, fixes any NaN's using
# RegularGridInterpolator, then again uses RegularGridInterpolator
# on the fixed data to extrapolate the requested field point.
# In principle, the first call of RegularGridInterpolator can be
# used to extrapolate the requested field point to eliminate
# the intermediate step, but this method enables use of all the
# real data rather than the trimmed data set.
def _fix_zgrid_NaNs(xgrid, ygrid, zgrid_input, rot_ang=0):
"""Fix NaN's in Zernike Grid
We trim NaN's within `zgrid`, then generate an extrapolation function
using `RegularGridInterpolator`. A rotation angle can also be specified
to maximize the number of remaining data points due to irregular
polygons of the real `zgrid` data.
Returns `zgrid` with the NaN's fixed using the extrapolation function.
Parameter
=========
xgrid : ndarray
1D V2 regular grid information
ygrid : ndarray
1D V3 regular grid information
zgrid : ndarray
2D Zernike grid
rot_ang : float
Option to rotate grid data for more optimal
trimming of NaN's.
"""
# Prevent this function from replacing original data
zgrid = zgrid_input.copy()
# Rotate zgrid
if rot_ang != 0:
zgrid = rotate(zgrid, rot_ang, reshape=False, order=1, cval=np.nan)
# There will be some NaN's along the border that need to be replaced
ind_nan = np.isnan(zgrid)
# Remove rows/cols with NaN's
xgrid2, ygrid2, zgrid2 = _trim_nan_image(xgrid, ygrid, zgrid)
# Create regular grid interpolator function for extrapolation of NaN's
func = RegularGridInterpolator((ygrid2,xgrid2), zgrid2, method='linear',
bounds_error=False, fill_value=None)
# Replace NaNs
X, Y = np.meshgrid(xgrid,ygrid)
pts = np.array([Y[ind_nan], X[ind_nan]]).transpose()
zgrid[ind_nan] = func(pts)
# De-rotate clipped zgrid image and redo RegularGridInterpolator
if rot_ang != 0:
# De-rotate
zgrid = rotate(zgrid, -rot_ang, reshape=False, order=1, cval=np.nan)
# There will be some NaNs along the border that need to be replaced
ind_nan = np.isnan(zgrid)
# Remove rows/cols 1 by 1 until no NaNs
xgrid2, ygrid2, zgrid2 = _trim_nan_image(xgrid, ygrid, zgrid)
# Create regular grid interpolator function for extrapolation of NaN's
func = RegularGridInterpolator((ygrid2,xgrid2), zgrid2, method='linear',
bounds_error=False, fill_value=None)
# Replace NaNs
pts = np.array([Y[ind_nan], X[ind_nan]]).transpose()
zgrid[ind_nan] = func(pts)
return zgrid
def _trim_nan_image(xgrid, ygrid, zgrid):
"""NaN Trimming of Image
Remove rows/cols with NaN's while trying to preserve
the maximum footprint of real data.
"""
xgrid2, ygrid2, zgrid2 = xgrid, ygrid, zgrid
# Create a mask of NaN'ed values
nan_mask = np.isnan(zgrid2)
nrows, ncols = nan_mask.shape
# Determine number of NaN's along each row and col
num_nans_cols = nan_mask.sum(axis=0)
num_nans_rows = nan_mask.sum(axis=1)
# First, crop all rows/cols that are only NaN's
xind_good = np.where(num_nans_cols < nrows)[0]
yind_good = np.where(num_nans_rows < ncols)[0]
# get border limits
x1, x2 = (xind_good.min(), xind_good.max()+1)
y1, y2 = (yind_good.min(), yind_good.max()+1)
# Trim of NaN borders
xgrid2 = xgrid2[x1:x2]
ygrid2 = ygrid2[y1:y2]
zgrid2 = zgrid2[y1:y2,x1:x2]
# Find a optimal rectangule subsection free of NaN's
# Iterative cropping
ndiff = 5
while np.isnan(zgrid2.sum()):
# Make sure ndiff is not negative
if ndiff<0:
break
npix = zgrid2.size
# Create a mask of NaN'ed values
nan_mask = np.isnan(zgrid2)
nrows, ncols = nan_mask.shape
# Determine number of NaN's along each row and col
num_nans_cols = nan_mask.sum(axis=0)
num_nans_rows = nan_mask.sum(axis=1)
# Look for any appreciable diff row-to-row/col-to-col
col_diff = num_nans_cols - np.roll(num_nans_cols,-1)
row_diff = num_nans_rows - np.roll(num_nans_rows,-1)
# For edge wrapping, just use last minus previous
col_diff[-1] = col_diff[-2]
row_diff[-1] = row_diff[-2]
#print(npix, zgrid2.shape, num_nans_cols, num_nans_rows)
# Keep rows/cols composed mostly of real data
# and where number of NaN's don't change dramatically
xind_good = np.where( ( np.abs(col_diff) <= ndiff ) &
( num_nans_cols < 0.5*nrows ) )[0]
yind_good = np.where( ( np.abs(row_diff) <= ndiff ) &
( num_nans_rows < 0.5*ncols ) )[0]
# get border limits
x1, x2 = (xind_good.min(), xind_good.max()+1)
y1, y2 = (yind_good.min(), yind_good.max()+1)
# Trim of NaN borders
xgrid2 = xgrid2[x1:x2]
ygrid2 = ygrid2[y1:y2]
zgrid2 = zgrid2[y1:y2,x1:x2]
# Check for convergence
# If we've converged, reduce
if npix==zgrid2.size:
ndiff -= 1
# Last ditch effort in case there are still NaNs
# If so, remove rows/cols 1 by 1 until no NaNs
while np.isnan(zgrid2.sum()):
xgrid2 = xgrid2[1:-1]
ygrid2 = ygrid2[1:-1]
zgrid2 = zgrid2[1:-1,1:-1]
return xgrid2, ygrid2, zgrid2
In [10]:
# Fix NaN's via extrapolation
zgrid_fix = _fix_zgrid_NaNs(xgrid, ygrid, zgrid)
In [11]:
fig, ax = plt.subplots(1,1, figsize=(8,5))
#zmin, zmax = np.nanmin(zgrid), np.nanmax(zgrid)
extent = np.array([xgrid.min(),xgrid.max(),ygrid.min(),ygrid.max()]) * 60
ax.imshow(zgrid_fix, extent=extent, vmin=zmin, vmax=zmax)
siaf = Siaf('NIRCam')
ap = siaf['NRCALL_FULL']
ap.plot(ax=ax, fill=False, color='C1')
fig.tight_layout()
In [12]:
# for mod in ['SW', 'LW', 'SWA', 'LWA', 'SWB', 'LWB']:
for mod in ['SWA', 'LWA', 'SWB', 'LWB']:
# Create a Zernike cube
zcube = []
for k in zkeys:
xgrid, ygrid, zgrid = gen_cv3_zgrid(ztable_full, v2v3_limits, mod, k)
zgrid_fix = _fix_zgrid_NaNs(xgrid, ygrid, zgrid)
zcube.append(zgrid_fix)
zube = np.array(zcube)
hdu = fits.PrimaryHDU(zcube)
hdr = hdu.header
dstep = 1. / 60.
hdr['units'] = 'meters'
hdr['xunits'] = 'Arcmin'
hdr['xmin'] = xgrid.min()
hdr['xmax'] = xgrid.max()
hdr['xdel'] = dstep
hdr['yunits'] = 'Arcmin'
hdr['ymin'] = ygrid.min()
hdr['ymax'] = ygrid.max()
hdr['ydel'] = dstep
hdr['wave'] = 2.10 if 'SW' in mod else 3.23
hdr['comment'] = 'X and Y values correspond to V2 and V3 coordinates (arcmin).'
hdr['comment'] = 'Slices in the cube correspond to Zernikes 1 to 36.'
hdr['comment'] = 'Zernike values calculated using 2D cubic interpolation'
hdr['comment'] = 'and linear extrapolation outside gridded data.'
fname = 'NIRCam{}_zernikes_isim_cv3.fits'.format(mod)
hdu.writeto(outdir + fname, overwrite=True)
In [ ]: