In [1]:
# Makes print and division act like Python 3
from __future__ import print_function, division
# Import the usual libraries
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
# Enable inline plotting at lower left
%matplotlib inline
matplotlib.rcParams['image.origin'] = 'lower'
matplotlib.rcParams['image.interpolation'] = 'none'
# seaborn package for making pretty plots, but not necessary
try:
import seaborn as sns
params = {'xtick.direction': 'in', 'ytick.direction': 'in', 'font.family': ['serif'],
'text.usetex': True, 'text.latex.preamble': ['\usepackage{gensymb}']}
sns.set_style("ticks", params)
except ImportError:
print('Seaborn module is not installed.')
from IPython.display import display, Latex
In [2]:
import pynrc
from pynrc import nrc_utils
from pynrc.nrc_utils import (webbpsf, poppy, pix_noise)
from pynrc import speckle_noise as sn
import astropy.io.fits as fits
import multiprocessing as mp
In [3]:
pynrc.setup_logging('WARNING', verbose=False)
In [4]:
data_path = webbpsf.utils.get_webbpsf_data_path() + '/'
opd_path = data_path + 'NIRCam/OPD/'
opd_file = 'OPD_RevV_nircam_132.fits'
opds, header = fits.getdata(opd_path + opd_file, header=True)
In [5]:
nopd = opds.shape[0]
nproc = int(np.min([nopd,mp.cpu_count()*0.75]))
In [6]:
# Multiprocessing for each OPD
%time opds_all = sn.opd_extract_mp(opds, header)
In [7]:
# For the pupil OPD and each segment OPD, find the stdev of each Zern/Hexike coefficient
pup_cf_std = np.array([opds_all[i].coeff_pupil for i in range(9)]).std(axis=0)
nseg = 18
seg_cf_std_all = []
for j in range(nseg):
std = np.array([opds_all[i].coeff_segs[j] for i in range(9)]).std(axis=0)
seg_cf_std_all.append(std)
seg_cf_std = np.median(seg_cf_std_all, axis=0)
# These values will be used to vary RMS WFE
# Set the piston values (Z1) to 0
# John Krist also says tip/tilt (Z2-3) are 0
pup_cf_std[0:3] = 0.0
seg_cf_std[0:3] = 0.0
# Zern/Hexikes to vary: Z = 3-14 (indices 3-7)
pup_cf_std[15:] = 0
seg_cf_std[15:] = 0
znum_pup = np.arange(len(pup_cf_std))+1
znum_seg = np.arange(len(seg_cf_std))+1
f, (ax1,ax2) = plt.subplots(1,2,figsize=(16,5))
ax1.plot(znum_pup, pup_cf_std*1000, color='steelblue', marker='o')
ax1.plot(znum_seg, seg_cf_std*1000, color='indianred', marker='o')
ax1.set_title('Variations of Pupil and Segment Coefficients')
ax2.set_title('Median of Variations in Segment Coefficients')
for seg in seg_cf_std_all:
ax2.plot(znum_seg, seg*1000, color='indianred', alpha=0.15)
ax2.plot(znum_seg, seg_cf_std*1000, color='indianred', marker='o', lw=3)
for ax in (ax1,ax2):
ax.set_xlabel('Zern/Hexike #')
ax.set_ylabel('Coefficient Sigma (nm)')
In [8]:
# Generate list of science OPDs and residuals for use in reference drift.
%time opd_sci_list, opd_resid_list = sn.opd_sci_gen_mp(opds_all)
In [9]:
# Generate a new set of OPDs for a series of WFE drift values
drift_list = [1.0,5.0,10.0,20.0]
args = (opds_all, pup_cf_std, seg_cf_std, opd_resid_list) # Arguments to pass
# OPDs for all four drift values (10x4)
opd_ref_list_all = [sn.ODP_drift_all(wfe_drift, *args) for wfe_drift in drift_list]
In [10]:
vlim = 50
fig, axes = plt.subplots(2,5,figsize=(11.5,6))
ind = 2
opd_ref_list = opd_ref_list_all[ind]
for i,ax in enumerate(axes.flat):
im = ax.imshow((opd_sci_list[i]-opd_ref_list[i])*1000, cmap='RdBu', vmin=-vlim, vmax=vlim)
ax.set_aspect('equal')
if i % 5 > 0: ax.set_yticklabels([])
if i < 5: ax.set_xticklabels([])
#fig.tight_layout()
fig.subplots_adjust(wspace=0.05, hspace=0.05, top=0.94, bottom=0.18, left=0, right=1)
cbar_ax = fig.add_axes([0.15, 0.05, 0.7, 0.025])
fig.colorbar(im, cax=cbar_ax, orientation = 'horizontal')
cbar_ax.set_xlabel('WFE Difference (nm)')
cbar_ax.xaxis.set_label_position('top');
fig.suptitle('{:.0f} nm WFE Drift Maps (Pupil only)'.format(drift_list[ind]));
#outdir = '/Users/jwstnircam/Desktop/NRC_Coronagraph/WFE_models/'
#fig.savefig(outdir+'wfe_diff_10nm_lebreton.pdf')
In [11]:
psf_list = [
('F323N', None , None ), # Direct Imaging
('F335M', 'MASKLWB' , 'WEDGELYOT')] # Coronagraphic observations
# Filters and masks
filt_direct, _, _ = psf_list[0]
filt_coron, mask, pupil = psf_list[1]
In [12]:
# Get planet PSF
psf_planet_coron = sn.get_psf(opd_sci_list[0], header, filt_coron, None, pupil)
In [13]:
# Coronagraphic PSF at bar's center positin
%time psf_star0 = [sn.get_psf(opd, header, filt_coron, mask, pupil, r=0) for opd in opd_sci_list]
In [14]:
# Coronagraphic PSF shifted to the right
%time psf_star1 = [sn.get_psf(opd, header, filt_coron, mask, pupil, theta=-90) for opd in opd_sci_list]
In [15]:
# shifted to the left
%time psf_star2 = [sn.get_psf(opd, header, filt_coron, mask, pupil, theta=90) for opd in opd_sci_list]
In [26]:
sn.offset_bar(filt_coron,mask)
Out[26]:
In [16]:
# Get all reference PSFs
# This includes the multiple drift values
%time psf_ref0_all = sn.gen_psf_ref_all(opd_ref_list_all, header, filt_coron, mask, pupil, r=0)
In [17]:
%time psf_ref1_all = sn.gen_psf_ref_all(opd_ref_list_all, header, filt_coron, mask, pupil, theta=-90)
In [18]:
%time psf_ref2_all = sn.gen_psf_ref_all(opd_ref_list_all, header, filt_coron, mask, pupil, theta=90)
In [19]:
vlim = 5e-8
fig, axes = plt.subplots(2,5,figsize=(11.5,6))
ext = 0
if ext==0: vlim = vlim/16
ind = 2
psf_star = psf_star0
psf_ref = psf_ref0_all[ind]
for i,ax in enumerate(axes.flat):
hdu1 = psf_star[i]
hdu2 = psf_ref[i]
im = webbpsf.display_psf_difference(hdu1, hdu2, ext, ext, -vlim, vlim, title='', imagecrop=5,
colorbar=False, cmap='RdBu', ax=ax, return_ax=True)
if i % 5 > 0: ax.set_yticklabels([])
if i < 5: ax.set_xticklabels([])
#fig.tight_layout()
fig.subplots_adjust(wspace=0.05, hspace=0.05, top=0.94, bottom=0.18, left=0, right=1)
cbar_ax = fig.add_axes([0.15, 0.05, 0.7, 0.025])
fig.colorbar(im.images[0], cax=cbar_ax, orientation = 'horizontal')
cbar_ax.set_xlabel('PSF Difference')
cbar_ax.xaxis.set_label_position('top');
fig.suptitle('Speckles ({:.0f} nm WFE)'.format(drift_list[ind]));
In [71]:
# Create a series of speckle noise images for each drift amount
speckle0 = [sn.speckle_noise_image(psf_star0, psf_ref) for psf_ref in psf_ref0_all]
speckle1 = [sn.speckle_noise_image(psf_star1, psf_ref) for psf_ref in psf_ref1_all]
speckle2 = [sn.speckle_noise_image(psf_star2, psf_ref) for psf_ref in psf_ref2_all]
In [21]:
contrast0 = [sn.get_contrast(speckle_image, psf_planet_coron) for speckle_image in speckle0]
contrast1 = [sn.get_contrast(speckle_image, psf_planet_coron) for speckle_image in speckle1]
contrast2 = [sn.get_contrast(speckle_image, psf_planet_coron) for speckle_image in speckle2]
In [25]:
#plt.imshow(speckle_coron[2][1].data[50:100,50:100],vmin=1e-10,vmax=1e-8)
#speckle_coron[2][1].data.max()
ext = 1
vmin = 1e-12; vmax = 1e-5
if ext==0: vmin /= 16; vmax /= 16
fig, axes = plt.subplots(2,4,figsize=(10,6.2))
for i, ax in enumerate(axes[0]):
hdu = speckle1[i]
im = webbpsf.display_psf(hdu, ext, vmin=vmin, vmax=vmax, title='', colorbar=False, imagecrop=20,
ax=ax, return_ax=True)
ax.set_xticklabels([])
if i > 0: ax.set_yticklabels([])
for i, ax in enumerate(axes[1]):
hdu = speckle2[i]
im = webbpsf.display_psf(hdu, ext, vmin=vmin, vmax=vmax, title='', colorbar=False, imagecrop=20,
ax=ax, return_ax=True)
if i > 0: ax.set_yticklabels([])
#fig.tight_layout()
fig.subplots_adjust(wspace=0.05, hspace=0.05, top=0.94, bottom=0.15, left=0, right=1)
cbar_ax = fig.add_axes([0.15, 0.05, 0.7, 0.025])
fig.colorbar(im.images[0], cax=cbar_ax, orientation = 'horizontal')
cbar_ax.set_xlabel('Fractional Noise')
cbar_ax.xaxis.set_label_position('top');
fig.suptitle('Speckle Noise Distribution');
In [52]:
ext=0
ndrift = len(drift_list)
ratio = [np.median(speckle_coron[i][ext].data / speckle_coron[0][ext].data) for i in range(ndrift)]
ratio2 = [np.median(speckle_direct[i][ext].data / speckle_direct[0][ext].data) for i in range(ndrift)]
#plt.plot(drift_list, ratio)
#plt.plot(drift_list, ratio2)
#plt.imshow(speckle_coron[2][ext].data / speckle_coron[0][ext].data)
vmin = 1e-13; vmax = 1e-8
norm = matplotlib.colors.LogNorm(vmin=vmin, vmax=vmax)
speckle_scaled = []
for i in range(ndrift):
im = speckle_coron[i][ext].data / ratio[i]
speckle_scaled.append(im)
print(np.median(speckle_coron[0][ext].data / im))
speckle_master = np.mean(np.array(speckle_scaled), axis=0)
plt.imshow(speckle_master, norm=norm)
Out[52]:
In [48]:
fig,ax = plt.subplots(1,1, figsize=(10, 5))
current_palette = sns.color_palette()
pal1 = sns.color_palette("deep")
pal2 = sns.color_palette("muted")
for j,drift in enumerate(drift_list):
r, c = contrast_spot[j]
ax.semilogy(r, c, color=pal1[j], label='Spot - {:.0f} nm'.format(drift))
for j,drift in enumerate(drift_list):
r, c = contrast_direct[j]
ax.semilogy(r, c, color=pal1[j], label='Direct - {:.0f} nm'.format(drift), ls='--')
ax.legend()
ax.set_xlim([0,5]);
ax.set_ylim([1e-8,1e-3]);
ax.set_ylabel('Contrast Ratio')
ax.minorticks_on()
ax.legend(ncol=2)
ax.set_xlabel("Radius ('')")
ax.set_title('Contrast Curves')
fig.tight_layout()
outdir = '/Users/jwstnircam/Desktop/NRC_Coronagraph/WFE_models/'
#fig.savefig(outdir+filt_coron+'_contrast_lebreton.pdf', facecolor='none')
In [94]:
fig,ax = plt.subplots(1,1, figsize=(10, 5))
current_palette = sns.color_palette()
pal1 = sns.color_palette("deep")
pal2 = sns.color_palette("muted")
for j,drift in enumerate(drift_list):
r, c = contrast1[j]
ax.semilogy(r, c, color=pal1[j], label='Bar0 - {:.0f} nm'.format(drift))
for j,drift in enumerate(drift_list):
r, c = contrast2[j]
ax.semilogy(r, c, color=pal1[j], label='Bar1 - {:.0f} nm'.format(drift), ls='--')
ax.legend()
ax.set_xlim([0,5]);
ax.set_ylim([1e-8,1e-3]);
ax.set_ylabel('Contrast Ratio')
ax.minorticks_on()
ax.legend(ncol=2)
ax.set_xlabel("Radius ('')")
ax.set_title('Contrast Curves')
fig.tight_layout()
outdir = '/Users/jwstnircam/Desktop/NRC_Coronagraph/WFE_models/'
#fig.savefig(outdir+filt_coron+'_contrast_lebreton.pdf', facecolor='none')
In [ ]: