<jupyter_start><jupyter_text>Testing<jupyter_code># Generate list of science OPDs and residuals for use in reference drift.
pool = mp.Pool(nproc)
%time output = pool.map(opd_sci_gen, opds_all)
opd_sci_list, opd_resid_list = zip(*output)
pool.close()
pool.join()
#wfe_drift = 10.0
#%time opd_ref_list = [opd_drift_nogood(opd,wfe_drift) for opd in opd_sci_list]
diff = np.asarray(opd_sci_list) - np.asarray(opd_ref_list)
diff.tolist()
wfe_drift = 10.0

#i=0
#test = opd_ref_gen((opds_all[i], wfe_drift, pup_cf_std, seg_cf_std, opd_resid_list[i]))

args = (opds_all, pup_cf_std, seg_cf_std, opd_resid_list)
%time opd_ref_list = ODP_drift_all(wfe_drift, *args)
vlim = 50
fig, axes = plt.subplots(2,5,figsize=(15,6.2))

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.925, bottom=0.18)

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');

outdir = '/Users/jwstnircam/Desktop/NRC_Coronagraph/WFE_models/'
fig.suptitle('WFE Drift Maps (Lebreton Method)')
#fig.savefig(outdir+'wfe_diff_10nm_lebreton.pdf')<jupyter_output><empty_output><jupyter_text>Science OPDs and PSFs<jupyter_code># Generate list of science OPDs and residuals for use in reference drift.
pool = mp.Pool(nproc)
%time output = pool.map(opd_sci_gen, opds_all)
opd_sci_list, opd_resid_list = zip(*output)
pool.close()
pool.join()
# Filters and masks
filt_coron, mask, pupil = ('F335M', 'MASK335R','CIRCLYOT') # Coronagraphic observations
filt_direct = 'F323N' # Direct Imaging
# Get all planet and stellar PSFs for coronagraphic observations
%time psf_sci_all = [get_psf_sci(opd, filt_coron, mask, pupil) for opd in opd_sci_list]
# Get all planet and stellar PSFs for direct imaging
%time psf_sci_direct = [get_psf_sci(opd, filt_direct, None, None) for opd in opd_sci_list]<jupyter_output>CPU times: user 6.97 s, sys: 2.73 s, total: 9.7 s
Wall time: 25.4 s<jupyter_text>Drift OPDs and Reference PSFs<jupyter_code># For a series of WFE drift values:
#   - Generate a new set of OPDs
#   - Generate a new set of reference PSFs
#   - Calculate the contrast
drift_list = [1.0,2.0,5.0,10.0]
args = (opds_all, pup_cf_std, seg_cf_std, opd_resid_list) # Arguments to pass
# OPDs for all four drift values (10x4)
%time opd_ref_list_all = [ODP_drift_all(wfe_drift, *args) for wfe_drift in drift_list]
# Coronagraphic contrast
%time contrast_spot = [contrast_drift(wfe_drift, psf_sci_all, filt_coron, mask, pupil, opd_ref_list_all[i], \
                                      *args) for i,wfe_drift in enumerate(drift_list)]
# Direct imaging contrast
%time contrast_direct = [contrast_drift(wfe_drift, psf_sci_direct, filt_direct, None, None, opd_ref_list_all[i], \
                                        *args) for i,wfe_drift in enumerate(drift_list)]<jupyter_output>Finished: 1 nm
Finished: 2 nm
Finished: 5 nm
Finished: 10 nm
CPU times: user 1min 10s, sys: 14.1 s, total: 1min 24s
Wall time: 2min 25s<jupyter_text>Plot Contrast Curves<jupyter_code>#f, (ax1,ax2,ax3) = plt.subplots(1,3,figsize=(14,3))
import matplotlib.gridspec as gridspec

fig = plt.figure(figsize=(12, 8))
gs = gridspec.GridSpec(2, 2, height_ratios=[2,3])
ax1 = plt.subplot(gs[0, 0])
ax2 = plt.subplot(gs[0, 1])
ax3 = plt.subplot(gs[1, :])

current_palette = sns.color_palette()
pal1 = sns.color_palette("deep")
pal2 = sns.color_palette("muted")

for j,drift in enumerate(drift_list):
    contrast = contrast_spot[j]
    r,c = zip(*contrast_spot[j])
    med = np.median(c, axis=0)
   
    for rc in contrast:
        ax1.semilogy(rc[0], rc[1], color=current_palette[j], alpha=0.2)
    ax1.semilogy(r[0], med, color=pal1[j], label='Spot - {:.0f} nm'.format(drift))
    ax3.semilogy(r[0], med, color=pal1[j], label='Spot - {:.0f} nm'.format(drift), lw=3)

for j,drift in enumerate(drift_list):
    contrast2 = contrast_direct[j]
    r2,c2 = zip(*contrast_direct[j])
    med2 = np.median(c2, axis=0)

    for rc in contrast2:
        ax2.semilogy(rc[0], rc[1], color=current_palette[j], alpha=0.2)
    ax2.semilogy(r2[0], med2, color=pal1[j], label='Direct - {:.0f} nm'.format(drift))
    ax3.semilogy(r2[0], med2, color=pal1[j], label='Direct - {:.0f} nm'.format(drift), lw=3, ls='--')

for ax in (ax1,ax2,ax3):
    ax.legend()
    ax.set_xlim([0,5]);
    ax.set_ylim([1e-8,1e-3]);
    ax.set_ylabel('Contrast Ratio')
    ax.minorticks_on()
    #ax.set_axis_bgcolor('blue')
ax3.legend(ncol=2)
ax3.set_xlabel("Radius ('')")

ax1.set_title('Coronagraphic Filter ' + filt_coron)
ax2.set_title('Direct Imaging Filter ' + filt_direct)
ax3.set_title('Contrast Curves (Lebreton Method)')

fig.tight_layout()
outdir = '/Users/jwstnircam/Desktop/NRC_Coronagraph/WFE_models/'
#fig.savefig(outdir+filt_coron+'_contrast_lebreton.pdf', facecolor='none')
# WFE drift for reference star

# Pupil
#coeff_pupil_new = opd_object.coeff_pupil + delta_coeff[:pup_var.size]
#opd_pupil_new = opd_object._opd_from_coeff(coeff_pupil_new, opd_object.basis_zernike, mask_pupil) * mask_pupil

# Segments
#coeff_segs_new = [cf + delta_coeff[pup_var.size:] for cf in opd_object.coeff_segs]
#opd_segs_new_list = opd_object._get_opd_new_segs(coeff_segs_new)
#opd_segs_new = opd_object.combine_opd_segs(opd_segs_new_list) * mask_pupil

# Total
#opd_ref = opd_pupil_new + opd_segs_new + opd_resid
rad_all, con_all = zip(*res_all)
f, axes= plt.subplots(1,3, figsize=[16,3])
(ax1,ax2,ax3) = axes

# PSF profiles for star and planet PSFs (occulted and unocculted)
# Plot both resolutions (pixel-sampled and 4x oversampled)
ax1.semilogy(rr0, mn0*16, label='Stellar PSf')
#ax1.plot(rr1, mn1)
ax1.plot(rr_psf0, mn_psf0*16, label='Planet PSF')
#ax1.plot(rr_psf1, mn_psf1)
ax1.set_ylim([1e-8,1e-2])
ax1.set_title('Radial Profile')

# Encircled energy for star and planet PSFs
ax2.semilogy(rr0, ee0, label='Stellar PSF')
#ax2.plot(rr1, ee1)
ax2.plot(rr_psf0, ee_psf0, label='Planet PSF')
#ax2.plot(rr_psf1, ee_psf1)
ax2.set_title('Encircled Energy')

ax3.semilogy(rr1,contrast1, label='Method 1')
ax3.plot(rr1,contrast2, label='Method 2')
ax3.plot(rr0,contrast3, label='Method 3')
ax3.set_ylim([1e-8,1e-3])
ax3.set_title('Contrast')

for ax in axes:
    ax.set_xlim([0,5])
    ax.set_xlabel('Arcsec')
    ax.legend(loc=0)
mask0 = (opd0==0) & (opd_pupil==0) & (opd_segs==0)
mask1 = ~mask0
print('OPD RMS Variations: \n\tOriginal: {:.3f} \n\tPupil: {:.3f} \n\tSegments: {:.3f} \n\tResiduals: {:.3f}'
    .format(opd0[mask1].std(), opd_pupil[mask1].std(), opd_segs[mask1].std(), opd_resid[mask1].std()))
vlim = 0.5
f, (ax1,ax2,ax3) = plt.subplots(1,3, figsize=[16,5])
ax1.imshow(opd_pupil, vmin=-vlim, vmax=vlim, cmap='gist_heat') #seg*seg_mask+
ax2.imshow(opd_segs, vmin=-vlim, vmax=vlim, cmap='gist_heat') #seg*seg_mask+
ax3.imshow(opd_resid, vmin=-vlim, vmax=vlim, cmap='gist_heat') #seg*seg_mask+
jopd = '/Users/jwstnircam/Desktop/JWST_SIMULATOR/PSFs/' + \
    'OPD_RevV_nircam_132_jitter_01_025_05_10_20_50_100_zernike.fits'
hdulist = fits.open(jopd)
noise = [0.1, 0.25, 0.5, 1.0, 2.0, 5.0, 10.]
i=6
opd_new = pynrc.nrc_utils.opd_drift(opds[0,:,:], header, noise[i])
diff = opd_new[0].data - hdulist[0].data[i,:,:]
nterms = 8
diam = header['PUPLDIAM']
pix_m = header['PUPLSCAL']

# Create x/y carteesian grid
sh = opds.shape
if len(sh) == 3: 
    nz,ny,nx = sh
    opd_sum = opds.sum(axis=0)
    mask0 = (opd_sum == 0)
    mask1 = (opd_sum != 0)
else:
    nz = 1
    ny,nx = sh
    # Masks for those values equal (and not) to 0
    mask0 = (opds == 0)
    mask1 = (opds != 0)

y,x = np.indices((ny,nx))
center = tuple((a - 1) / 2.0 for a in [nx,ny])
y = y - center[1]; x = x - center[0]
y *= pix_m; x *= pix_m

# Convert to polar coordinates
rho = np.sqrt(x**2 + y**2) / diam
theta = np.arctan2(y,x)

# Generate Zernike maps
zall = (poppy.zernike.zernike_basis(nterms, rho=rho, theta=theta))[1:,:,:]

# Mask out defocus
zmasked = np.ma.array(zall, mask=False)
zmasked.mask[2,:,:] = True

ztot = zmasked.sum(axis=0)
#ztot = (1.0-defocus_frac)*ztot + defocus_frac*zall[2,:,:]
#ztot /= ztot.sum()
defocus_frac = 0
print(ztot.sum(),zall[2,:,:].sum())
ztot = (1.0-defocus_frac)*ztot + defocus_frac*(zall[2,:,:]/zall[2,:,:].sum())
ztot[mask0] = 0
ztot /= ztot[mask1].std()
ztot[mask1] -= ztot[mask1].mean()
print(ztot[mask1].mean(), ztot[mask1].sum(), ztot[mask1].min(), ztot[mask1].max())
ztot = poppy.zernike.opd_from_zernikes([0,0,0,1,0,0,0,0], rho=rho, theta=theta, basis=poppy.zernike.zernike_basis)
ztot[mask0]=0
#ztot /= ztot[mask1].std()
print(ztot[mask1].mean(), ztot[mask1].sum(), ztot[mask1].min(), ztot[mask1].max())
for i in range(7):
    print('{:.3f} {:.3f} {:.3f} {:.3f}'.\
            format(zall[i,:,:].mean(), zall[i,:,:].sum(), zall[i,:,:].min(), zall[i,:,:].max()))
ax = plt.imshow(ztot, cmap='gist_heat')
plt.colorbar(ax)
def pix_noise_contrast(ngroup=2, nf=1, nd2=0, tf=10.737, rn=15.0, ktc=29.0, p_excess=(0,0),
    fsrc=0.0, idark=0.003, fzodi=0, fbg=0, **kwargs):
    """
    Similar to pix_noise(), except 

    Parameters
    ===========
    n (int) : Number of groups in integration ramp
    m (int) : Number of frames in each group
    s (int) : Number of dropped frames in each group
    tf (float) : Frame time
    rn (float) : Read Noise per pixel
    ktc (float) : kTC noise only valid for single frame (n=1)
    p_excess: An array or list of two elements that holding the
        parameters that describe the excess variance observed in
        effective noise plots. By default these are both 0.
        Recommended values are [1.0,5.0] or SW and [1.5,10.0] for LW.

    fsrc  (float) : Flux of source in e-/sec/pix
    idark (float) : Dark current in e-/sec/pix
    fzodi (float) : Zodiacal light emission in e-/sec/pix
    fbg   (float) : Any additional background (telescope emission or scattered light?)

    Various parameters can either be single values or numpy arrays.
    If multiple inputs are arrays, make sure their array sizes match.
    Variables that need to have the same array sizes (or a single value):
        - n, m, s, & tf
        - rn, idark, ktc, fsrc, fzodi, & fbg

    Array broadcasting also works:
        For Example
        n = np.arange(50)+1 # An array of groups to test out

        # Create 2D Gaussian PSF with FWHM = 3 pix
        npix = 20 # Number of pixels in x and y direction
        x = np.arange(0, npix, 1, dtype=float)
        y = x[:,np.newaxis]
        x0 = y0 = npix // 2 # Center position
        fwhm = 3.0
        fsrc = np.exp(-4*np.log(2.) * ((x-x0)**2 + (y-y0)**2) / fwhm**2)
        fsrc /= fsrc.max()
        fsrc *= 10 # Total source counts/sec (arbitrarily scaled)
        fsrc = fsrc.reshape(npix,npix,1) # Necessary for broadcasting

        # Represents pixel array w/ different RN/pix
        rn = np.ones([npix,npix,1])*15. 
        # Results is a (20x20)x50 showing the noise in e-/sec/pix at each group
        noise = pix_noise(ngroup=n, rn=rn, fsrc=fsrc) 
    """

    n = np.array(ngroup)
    m = np.array(nf)
    s = np.array(nd2)
    tf = np.array(tf)
    max_size = np.max([n.size,m.size,s.size,tf.size])
    if n.size  != max_size: n  = n.repeat(max_size)
    if m.size  != max_size: m  = m.repeat(max_size)
    if s.size  != max_size: s  = s.repeat(max_size)
    if tf.size != max_size: tf = tf.repeat(max_size)

    # Total flux (e-/sec/pix)
    ftot = fsrc + idark + fzodi + fbg

    # Special case if n=1
    if (n==1).any():
        # Variance after averaging m frames
        var = ktc**2 + (rn**2 + ftot*tf) / m
        noise = np.sqrt(var) 
        noise /= tf # In terms of e-/sec

        if (n==1).all(): return noise
        noise_n1 = noise

    ind_n1 = (n==1)
    temp = np.array(rn+ktc+ftot)
    temp_bool = np.zeros(temp.shape, dtype=bool)
    ind_n1_all = (temp_bool | ind_n1)

    # Group time
    tg = tf * (m + s)

    # Read noise, group time, and frame time variances
    var_rn = rn**2      * 12.               * (n - 1.) / (m * n * (n + 1.))
    var_gp = ftot  * tg * 6. * (n**2. + 1.) * (n - 1.) / (5 * n * (n + 1.))
    var_fm = ftot  * tf * 2. * (m**2. - 1.) * (n - 1.) / (m * n * (n + 1.))

    # Functional form for excess variance above theoretical
    var_ex = 12. * (n - 1.)/(n + 1.) * p_excess[0]**2 - p_excess[1] / m**0.5

    # Variance of total signal
    var = var_rn + var_gp - var_fm + var_ex
    sig = np.sqrt(var)

    # Noise in e-/sec
    noise = sig / (tg * (n - 1))
    #print(ind_n1_all.shape,noise.shape,noise_n1.shape)
    if (n==1).any():
        noise[ind_n1_all] = noise_n1[ind_n1_all]

    # Include flat field noise
    # JWST-CALC-003894
    noise_ff = 1E-4 # Uncertainty in the flat field
    factor = 1 + noise_ff*np.sqrt(ftot)
    noise *= factor

    return noise

image = np.zeros([500,1000])
center = tuple((a - 1) / 2.0 for a in image.shape[::-1])

y,x = np.indices(image.shape)
y = y - center[1]
x = x - center[0]

r = np.sqrt(x**2 + y**2)
ang = np.arctan2(-x,y) * (180/np.pi)
print(ang.min(), ang.max())

plt.imshow(ang)
print(ang[0,999])<jupyter_output>-179.885178952 179.885178952
-116.542105331
{'fraction_non_alphanumeric': 0.16105151835624718, 'fraction_numerical': 0.04169814171324974, 'mean_word_length': 3.822950819672131, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '41700', 'license': 'MIT', 'n_tokens_mistral': 6050, 'n_tokens_neox': 5668, 'n_words': 1224, 'path': 'dev_utils/contrast_develop.ipynb', 'repo_name': 'JarronL/pyNRC', 'size': '13238'}