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


[     pynrc:INFO]   jwst_backgrounds is not installed and will not be used for bg estimates.

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]

V2/V3 Limits


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()

Interpolate within bounds of CV3 measurements


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()


Extrapolate using RegularGridInterpolator


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()


Do All Aperture


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 [ ]: