Update WFE_interpolation.py

pynrc/WFE_interpolation.py

@@ -220,5 +220,82 @@
 #     plt.ylim([Y.min(),Y.max()])
 #     plt.axes().set_aspect('equal', 'datalim')
 
-# names = ['Guider1', 'Guider2', 'MIRI', 'NIRCamLWA', 'NIRCamLWB', 
-# 'NIRCamSWA', 'NIRCamSWB', 'NIRISS', 'NIRSpec']
+
+
+
+names = ['Guider1', 'Guider2', 'MIRI', 'NIRISS', 'NIRSpec',
+         'NIRCamLWA', 'NIRCamLWB', 'NIRCamSWA', 'NIRCamSWB']
+
+from pysiaf.siaf import plot_main_apertures
+
+fig, axes = plt.subplots(3,1, figsize=(8,10))
+axes = axes.flatten()
+
+lookup_name = 'NIRSpec'
+
+for lookup_name in names:
+    ztable = ztable_full[ztable_full['instrument'] == lookup_name]
+
+    v2 = ztable['V2']
+    v3 = ztable['V3']
+
+    i=4
+    zkey = 'Zernike_{}'.format(i)
+    zvals = ztable[zkey]
+
+
+    dstep = 1. / 60. # 1" steps
+    xgrid = np.arange(v2.min(), v2.max()+dstep, dstep)
+    ygrid = np.arange(v3.min(), v3.max()+dstep, dstep)
+    X, Y = np.meshgrid(xgrid,ygrid)
+
+    # Cubic interpolation of all points
+    zgrid = griddata((v2, v3), zvals, (X, Y), method='cubic')
+    zmin, zmax = np.nanmin(zgrid), np.nanmax(zgrid)
+
+    ax = axes[0]
+    extent = np.array([xgrid.min(),xgrid.max(),ygrid.min(),ygrid.max()]) * 60
+    ax.imshow(zgrid, extent=extent, vmin=zmin, vmax=zmax)
+    ax.plot(v2*60, v3*60, ls='none', marker='x', color='C1')
+    ax.set_title("griddata with method='cubic'")
+
+
+    # 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
+    zgrid2 = zgrid
+    ygrid2 = ygrid
+    xgrid2 = xgrid
+    while np.isnan(zgrid2.sum()):
+        zgrid2 = zgrid2[1:-1,1:-1]
+        ygrid2 = ygrid2[1:-1]
+        xgrid2 = xgrid2[1:-1]
+
+    ax = axes[1]
+    extent = np.array([xgrid2.min(),xgrid2.max(),ygrid2.min(),ygrid2.max()]) * 60
+    ax.imshow(zgrid2, extent=extent, vmin=zmin, vmax=zmax)
+    ax.plot(v2*60, v3*60, ls='none', marker='x', color='C1')
+    ax.set_title("Trimmed for RegularGridInterpolator")
+
+
+    # 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)
+
+    ax = axes[2]
+    extent = np.array([xgrid.min(),xgrid.max(),ygrid.min(),ygrid.max()]) * 60
+    ax.imshow(zgrid, extent=extent, vmin=zmin, vmax=zmax)
+    ax.plot(v2*60, v3*60, ls='none', marker='x', color='C1')
+    ax.set_title("NaN replacement")
+
+for ax in axes:
+    plot_main_apertures(ax=ax, fill=False, mark_ref=False)
+    ax.set_xlim([600,-600])
+    ax.set_ylim([-800,-250])
+
+fig.tight_layout()
+