def plot_planes(cpx_seq,
title=None,
logZ=[False],
use_axis=True,
vlim=[None, None],
subplt_cols=3,
dx=None,
first=False):
"""
view plot of intensity in each wavelength bin at a single (last) timestep
will pull out the plane(s) of sp.save_list at last tstep of cpx_sequence, convert to intensity, and sum over
wavelength and object
Currently, the atmosphere and enterance pupil are plotted in units of phase vs intensity. I think we should change
this later for user-specification
:param cpx_seq:
:param title: string, must be set or will error!
:param logZ: turn logscale plotting for z-axis on or off
:param use_axis: turn on/off using axis ticks, colorbar, etc
:param vlim: tuple of colorbar axis limits (min,max)
:param subplt_cols: number of subplots per row
:param dx: sampling of the image at each saved plane
:return:
"""
# Create figure & adjust subplot number, layout, size, whitespace
fig = plt.figure()
n_planes = len(sp.save_list)
n_rows = int(np.ceil(n_planes / float(subplt_cols)))
plt.axis('off')
gs = gridspec.GridSpec(n_rows, subplt_cols, wspace=0.08)
# Main Title
if title is None:
warnings.warn("Plots without titles: Don't Do It!")
title = input("Please Enter Title: ")
pass
fig.suptitle(title, fontweight='bold', fontsize=16)
# Small Hack to repeat axis if defaults used
np.array(vlim)
if len(logZ) == 1:
logZ = np.repeat(logZ, len(sp.save_list))
if len(vlim) == 2:
vlim = [
vlim,
] * len(sp.save_list)
if not first:
f = -1
else:
f = 0 # select first or last timestep
for p in range(n_planes):
ax = fig.add_subplot(gs[p])
###################
# Retreiving Data
##################
# Standard-Way
# [timestep, plane, wavelength, object, x, y]
# converts to intensity of last timestep, THEN sums over wavelength, then sums over object
plot_plane = sp.save_list[p]
plane = opx.extract_plane(cpx_seq, plot_plane)
# Distinguish plotting z-axis in phase units or intensity units
if plot_plane == "atmosphere" or plot_plane == "entrance_pupil":
plane = np.sum(np.angle(plane[f]), axis=(0, 1))
plane = opx.extract_center(plane,
new_size=sp.grid_size * sp.beam_ratio +
10)
logZ[p] = False
vlim[p] = [None, None]
phs = " phase"
elif plot_plane == "woofer" or plot_plane == "tweeter" or plot_plane == "DM":
# only show the star phase map since phase at other bodies just offsets to shift focal plane position
plane = np.sum(np.angle(plane[f]),
axis=(0)) # only sum over object
plane = plane[0] # plot the shortest wavelength
plane = opx.extract_center(plane,
new_size=sp.grid_size * sp.beam_ratio +
10) # zoom in on DM
logZ[p] = False
vlim[p] = [-np.pi, np.pi]
phs = " phase"
else:
plane = np.sum(opx.cpx_to_intensity(plane[f]), axis=(0, 1))
phs = ""
### Retreiving Data- Custom selection of plane ###
# plot_plane = sp.save_list[w]
# plane = opx.extract_plane(cpx_seq, plot_plane)
# # plane = opx.cpx_to_intensity(plane[-1])
# plane = opx.extract_center(np.angle(plane[0,1,1])) # wavelengths, objects
# Converting Sampling Units to Readable numbers
if dx[p, 0] < 1e-5:
dx[p, :] *= 1e6 # [convert to um]
axlabel = 'um'
elif dx[p, 0] < 1e-3:
dx[p, :] *= 1e3 # [convert to mm]
axlabel = 'mm'
elif 1e-2 > dx[p, 0] > 1e-3:
dx[p, :] *= 1e2 # [convert to cm]
axlabel = 'cm'
else:
axlabel = 'm'
# X,Y lables
if dx is not None:
tic_spacing = np.linspace(
0, plane.shape[0],
5) # 5 (# of ticks) is just set by hand, arbitrarily chosen
tic_lables = np.round(
np.linspace(-dx[p, 0] * plane.shape[0] / 2,
dx[p, 0] * plane.shape[0] / 2, 5)).astype(
int) # nsteps must be same as tic_spacing
tic_spacing[0] = tic_spacing[0] + 1 # hack for edge effects
tic_spacing[-1] = tic_spacing[-1] - 1 # hack for edge effects
plt.xticks(tic_spacing, tic_lables, fontsize=6)
plt.yticks(tic_spacing, tic_lables, fontsize=6)
plt.ylabel(axlabel, fontsize=8)
# Z-axis scale
if phs == ' phase':
cmap = sunlight
else:
cmap = "YlGnBu_r"
ax.set_title(f"{sp.save_list[p]}" + phs)
if logZ[p]:
if vlim[p][0] is not None and vlim[p][0] <= 0:
im = ax.imshow(plane,
interpolation='none',
origin='lower',
vmin=vlim[p][0],
vmax=vlim[p][1],
norm=SymLogNorm(linthresh=1e-5),
cmap=cmap)
add_colorbar(im)
# clabel = "Log Normalized Intensity"
else:
im = ax.imshow(plane,
interpolation='none',
origin='lower',
vmin=vlim[p][0],
vmax=vlim[p][1],
norm=LogNorm(),
cmap=cmap) #(1e-6,1e-3)
add_colorbar(im)
# clabel = "Log Normalized Intensity"
# cb.set_label(clabel)
else:
im = ax.imshow(plane,
interpolation='none',
origin='lower',
vmin=vlim[p][0],
vmax=vlim[p][1],
cmap=cmap) # "twilight"
add_colorbar(im)
# clabel = "Normalized Intensity"
# cb.set_label(clabel)
if use_axis:
warnings.simplefilter("ignore", category=UserWarning)
gs.tight_layout(fig, pad=1.08,
rect=(0, 0.02, 1,
0.9)) # rect = (left, bottom, right, top)
def cdi_postprocess(cpx_sequence, sampling, plot=False):
"""
this is the function that accepts the timeseries of intensity images from the simulation and returns the processed
single image. This function calculates the speckle amplitude phase, and then corrects for it to create the dark
hole over the specified region of the image.
From Give'on et al 2011, we have in eq 10 two different values: DeltaP-the change in the focal plane due to the
probe, and delta, the intensity difference measurements between pairs of phase probes. (It is unfortunate that
both these terms use delta, so I have done my best to distinguish them in the variable names)
Here I note that in the CDI phase stream generation, for n_probes there are n_pairs = n_probes/2 pairs of probes.
These get applied to the DM in a series such that the two probes that form the conjugate pair are separated by
n_pairs of probes. In other words, for n_probes = 6, the 0th and 3rd probes are a pair, the 1st and 4th are a pair,
and so on. This is a choice made when creating cdi.phase_series.
:param cpx_sequence: #timestream of 2D images (complex) from the focal plane complex field
:param sampling: focal plane sampling
:return:
"""
##
tic = time.time()
focal_plane = extract_plane(cpx_sequence, 'detector') # eliminates astro_body axis [tsteps,wvl,obj,x,y]
fp_seq = np.sum(focal_plane, axis=(1,2)) # sum over wavelength,object
n_pairs = cdi.n_probes//2 # number of deltas (probe differentials)
n_nulls = sp.numframes - cdi.n_probes
delta = np.zeros((n_pairs, sp.grid_size, sp.grid_size), dtype=float)
# absDelta = np.zeros((n_nulls, sp.grid_size, sp.grid_size))
# phsDelta = np.zeros((n_pairs, sp.grid_size, sp.grid_size), dtype=float)
E_est = np.zeros((n_nulls, sp.grid_size, sp.grid_size), dtype=complex)
I_processed = np.zeros((n_nulls, sp.grid_size, sp.grid_size))
H = np.zeros((n_pairs, 2), dtype=float)
b = np.zeros((n_pairs, 1))
# Get Masked Data
mask2D, imsk, jmsk, irng, jrng, imx, imn, jmx, jmn = get_fp_mask(cdi, thresh=1e-6)
if sp.debug:
fig, ax = plt.subplots(1,1)
fig.suptitle(f'Masked FP in CDI probe Region')
im = ax.imshow(cpx_to_intensity(fp_seq[0,:,:]*mask2D))
for ip in range(n_pairs):
# Compute deltas (I_ip+ - I_ip-)/2
delta[ip] = (np.abs(fp_seq[ip])**2 - np.abs(fp_seq[ip + n_pairs])**2) / 4
# for i,j in zip(imsk,jmsk):
# for i,j in zip(irng,jrng):
for i in irng:
for j in jrng:
for xn in range(n_nulls):
for ip in range(n_pairs):
# Amplitude DeltaP
Ip = np.abs(fp_seq[ip, i, j]) ** 2
Im = np.abs(fp_seq[ip + n_pairs, i, j]) ** 2
Io = np.abs(fp_seq[cdi.n_probes + xn, i, j]) ** 2
abs = (Ip + Im) / 2 - Io
if abs < 0:
abs = 0
absDeltaP = np.sqrt(abs)
# absDeltaP = np.sqrt(np.abs((Ip + Im) / 2 - Io))
# phsDeltaP = phsDelta[ip, i, j]
# Phase DeltaP
# The phase of the change in the focal plane of the probe applied to the DM
# First subtract Eo vector from each probe phase to make new field vectors dEa, dEb,
# then take the angle between the two
dEp = fp_seq[ip, i, j] - fp_seq[cdi.n_probes + xn, i, j]
dEm = fp_seq[ip + n_pairs, i, j] - fp_seq[cdi.n_probes + xn, i, j]
phsDeltaP = np.arctan2(dEp.imag - dEm.imag, dEp.real - dEm.real)
cpxDeltaP = absDeltaP * np.exp(1j * phsDeltaP)
# cpxDeltaP = absDeltaP * np.array((np.cos(phsDeltaP), np.sin(phsDeltaP)))
H[ip, :] = [-cpxDeltaP.imag, cpxDeltaP.real] # [n_pairs, 2]
# H[ip,:] = [cpxDeltaP[0], cpxDeltaP[1]]
b[ip] = delta[ip, i, j] # [n_pairs, 1]
a = 2 * H
Exy = linalg.lstsq(a, b)[0] # returns tuple, not array
E_est[xn, i, j] = Exy[0] + (1j * Exy[1])
toc = time.time()
dprint(f'CDI post-processing took {(toc-tic)/60:.2} minutes\n')
## ===========================
# Contrast Ratios
# ===========================
intensity_probe = np.zeros(n_nulls)
intensity_DM_FFT = np.zeros(n_nulls)
intensity_pre_process = np.zeros(n_nulls)
intensity_post_process = np.zeros(n_nulls)
for xn in range(n_nulls):
I_processed[xn] = np.abs(fp_seq[n_pairs+xn])**2 - np.abs(np.conj(E_est[xn])*mask2D)**2
# I_processed[xn] = np.sqrt(np.abs(np.abs(fp_seq[n_pairs+xn])**2 - np.abs(E_est[xn]*mask2D)**2))**2
# I_processed[xn] = np.abs(fp_seq[n_pairs+xn] - np.conj(E_est[xn]*mask2D))**2
# Contrast
intensity_probe[xn] = np.sum(np.abs(fp_seq[xn]*mask2D)**2)
intensity_pre_process[xn] = np.sum(np.abs(fp_seq[n_pairs + xn]*mask2D)**2)
intensity_post_process[xn] = np.sum(I_processed[xn]*mask2D) #np.sum(np.abs(E_processed[xn]*mask2D)**2)
print(f'\nIntensity in probed region for null step {xn} is '
f'\nprobe {intensity_probe[xn]}'
f'\npre-processed {intensity_pre_process[xn]} '
f'\npost-processed {intensity_post_process[xn]}'
f'\n difference = {intensity_post_process[xn] - intensity_pre_process[xn]}'
f'\n')
## if plot:
# ==================
# FFT of Tweeter Plane
# ==================
# fig, subplot = plt.subplots(1, n_pairs, figsize=(14, 5))
# fig.subplots_adjust(wspace=0.5, right=0.85)
# fig.suptitle('Tweeter DM Plane')
#
# tweet = extract_plane(cpx_sequence, 'tweeter') # eliminates astro_body axis [tsteps,wvl,obj,x,y]
# tweeter = np.sum(tweet, axis=(1, 2))
# tweeter_intensity = np.abs(tweeter) ** 2
# for ax, ix in zip(subplot.flatten(), range(n_pairs)):
# im = ax.imshow(extract_center(tweeter_intensity[ix], new_size=sp.grid_size*sp.beam_ratio+10),
# interpolation='none', norm=LogNorm(),
# vmin=5e-3, vmax=1e-2)
# ax.set_title(f'Probe Phase ' r'$\theta$' f'={cdi.phase_cycle[ix] / np.pi:.2f}' r'$\pi$')
#
# cax = fig.add_axes([0.9, 0.2, 0.03, 0.6]) # Add axes for colorbar @ position [left,bottom,width,height]
# cb = fig.colorbar(im, orientation='vertical', cax=cax) #
# cb.set_label('Intensity')
#
# # =================================================================================
# fig, subplot = plt.subplots(1, n_pairs, figsize=(14, 5))
# fig.subplots_adjust(wspace=0.5, right=0.85)
# fig.suptitle('FFT of Tweeter DM Plane')
#
# for ax, ix in zip(subplot.flatten(), range(n_pairs)):
# fft_tweeter = (1 / np.sqrt(2 * np.pi) *
# np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(tweeter[ix]))))
# im = ax.imshow(np.abs(fft_tweeter) ** 2,
# interpolation='none', norm=LogNorm(),
# vmin=1e-2, vmax=10)
# ax.set_title(f'Probe Phase ' r'$\theta$' f'={cdi.phase_cycle[ix] / np.pi:.2f}' r'$\pi$')
#
# cax = fig.add_axes([0.9, 0.2, 0.03, 0.6]) # Add axes for colorbar @ position [left,bottom,width,height]
# cb = fig.colorbar(im, orientation='vertical', cax=cax) #
# cb.set_label('Intensity')
# ==================
# Deltas
# ==================
# fig, subplot = plt.subplots(1, n_pairs, figsize=(14,5))
# fig.subplots_adjust(wspace=0.5, right=0.85)
# fig.suptitle('Deltas for CDI Probes')
#
# for ax, ix in zip(subplot.flatten(), range(n_pairs)):
# im = ax.imshow(delta[ix]*1e6*mask2D, interpolation='none',
# norm=SymLogNorm(linthresh=1),
# # vmin=-1, vmax=1
# ) #, norm=SymLogNorm(linthresh=1e-5))
# ax.set_title(f"Diff Probe\n" + r'$\theta$' + f'={cdi.phase_series[ix]/np.pi:.3f}' +
# r'$\pi$ -$\theta$' + f'={cdi.phase_series[ix+n_pairs]/np.pi:.3f}' + r'$\pi$')
#
# cax = fig.add_axes([0.9, 0.2, 0.03, 0.6]) # Add axes for colorbar @ position [left,bottom,width,height]
# cb = fig.colorbar(im, orientation='vertical', cax=cax) #
# cb.set_label('Intensity')
# ==================
# Null Probe E-Field
# ==================
fig, subplot = plt.subplots(1, n_nulls, figsize=(14, 5))
fig.subplots_adjust(wspace=0.5, right=0.85)
fig.suptitle('Original (Null-Probe) Image Plane Intensity')
for ax, ix in zip(subplot.flatten(), range(n_nulls)):
im = ax.imshow(np.abs(fp_seq[cdi.n_probes + ix, irng[0]:irng[-1], jrng[0]:jrng[-1]]) ** 2, # , 250:270, 150:170 *mask2D , irng[0]:irng[-1], jrng[0]:jrng[-1]
interpolation='none', norm=LogNorm(),
vmin=1e-7, vmax=1e-4)
ax.set_title(f'Null Step {ix}')
cax = fig.add_axes([0.9, 0.2, 0.03, 0.6]) # Add axes for colorbar @ position [left,bottom,width,height]
cb = fig.colorbar(im, orientation='vertical', cax=cax) #
cb.set_label('Intensity')
# ===============================
# E-filed Estimate
# ===============================
fig, subplot = plt.subplots(1, n_nulls, figsize=(14, 5))
fig.subplots_adjust(wspace=0.5, right=0.85)
fig.suptitle('Pairwise Probe Estimated Image Plane intensity')
for ax, ix in zip(subplot.flatten(), range(n_nulls)):
im = ax.imshow(np.abs(E_est[ix, irng[0]:irng[-1], jrng[0]:jrng[-1]])**2, # , 250:270, 150:170 *mask2D , irng[0]:irng[-1], jrng[0]:jrng[-1]
interpolation='none',
norm=LogNorm(),
vmin=1e-7, vmax=1e-4
) # , norm=SymLogNorm(linthresh=1e-5))
ax.set_title(f'Null Step {ix}')
cax = fig.add_axes([0.9, 0.2, 0.03, 0.6]) # Add axes for colorbar @ position [left,bottom,width,height]
cb = fig.colorbar(im, orientation='vertical', cax=cax) #
cb.set_label('Intensity')
# ==================================
# Processed E-field (CDI Subtracted)
# ==================================
fig, subplot = plt.subplots(1, n_nulls, figsize=(14, 5))
fig.subplots_adjust(wspace=0.5, right=0.85)
fig.suptitle('CDI Subtracted Image Plane Intensity')
for ax, ix in zip(subplot.flatten(), range(n_nulls)):
im = ax.imshow(I_processed[ix], # , 250:270, 150:170 *mask2D , irng[0]:irng[-1], jrng[0]:jrng[-1]
interpolation='none',
norm=LogNorm(),
vmin=1e-7, vmax=1e-4
) # , norm=SymLogNorm(linthresh=1e-5))
ax.set_title(f'Null Step {ix}')
cax = fig.add_axes([0.9, 0.2, 0.03, 0.6]) # Add axes for colorbar @ position [left,bottom,width,height]
cb = fig.colorbar(im, orientation='vertical', cax=cax) #
cb.set_label('Intensity')
# =======================================
# Compare Unprobed to Estimated E-fields
# ======================================
fig, subplot = plt.subplots(2, n_nulls, figsize=(16, 12))
fig.subplots_adjust(wspace=0.5, right=0.85)
fig.suptitle('Compare Null Probed to Pairwise Probe Estimated Image Plane Intensity')
ax1, ax2, ax3, ax4 = subplot.flatten()
ax1.imshow(np.abs(fp_seq[-2, irng[0]:irng[-1], jrng[0]:jrng[-1]]) ** 2, # , 250:270, 150:170 *mask2D
interpolation='none',
norm=LogNorm(),
# vmin=1e-8, vmax=1e-2
)
ax1.set_title(f'Original Null Step 0')
ax2.imshow(np.abs(fp_seq[-1, irng[0]:irng[-1], jrng[0]:jrng[-1]]) ** 2, # , 250:270, 150:170 *mask2D
interpolation='none',
norm=LogNorm(),
# vmin=1e-8, vmax=1e-2
)
ax2.set_title(f'Original Null Step 1')
ax3.imshow(np.abs(E_est[0, irng[0]:irng[-1], jrng[0]:jrng[-1]]) ** 2, # , 250:270, 150:170 *mask2D
interpolation='none',
norm=LogNorm(),
# vmin=1e-8, vmax=1e-2
)
ax3.set_title(f'PP Estimated Null Step 0')
ax4.imshow(np.abs(E_est[1, irng[0]:irng[-1], jrng[0]:jrng[-1]]) ** 2, # , 250:270, 150:170 *mask2D
interpolation='none',
norm=LogNorm(),
# vmin=1e-8, vmax=1e-2
)
ax4.set_title(f'PP Estimated Null Step 1')
cax = fig.add_axes([0.9, 0.2, 0.03, 0.6]) # Add axes for colorbar @ position [left,bottom,width,height]
cb = fig.colorbar(im, orientation='vertical', cax=cax) #
cb.set_label('Intensity')
# =======================================
# Compare Unprobed to Estimated E-fields
# ======================================
fig, subplot = plt.subplots(2, n_nulls, figsize=(16, 12))
fig.subplots_adjust(wspace=0.5, right=0.85)
fig.suptitle('Compare Null Probed to Pairwise Probe Estimated Image Plane Intensity')
ax1, ax2, ax3, ax4 = subplot.flatten()
ax1.imshow(np.abs(fp_seq[-2]) ** 2, # , 250:270, 150:170 *mask2D
interpolation='none',
norm=LogNorm(),
# vmin=1e-8, vmax=1e-2
)
ax1.set_title(f'Original Null Step 0')
ax2.imshow(np.abs(fp_seq[-1]) ** 2, # , 250:270, 150:170 *mask2D
interpolation='none',
norm=LogNorm(),
# vmin=1e-8, vmax=1e-2
)
ax2.set_title(f'Original Null Step 1')
ax3.imshow(I_processed[0], # , 250:270, 150:170 *mask2D
interpolation='none',
norm=LogNorm(),
# vmin=1e-8, vmax=1e-2
)
ax3.set_title(f'CDI Processed Null Step 0')
ax4.imshow(I_processed[1], # , 250:270, 150:170 *mask2D
interpolation='none',
norm=LogNorm(),
# vmin=1e-8, vmax=1e-2
)
ax4.set_title(f'CDI Processed Null Step 1')
cax = fig.add_axes([0.9, 0.2, 0.03, 0.6]) # Add axes for colorbar @ position [left,bottom,width,height]
cb = fig.colorbar(im, orientation='vertical', cax=cax) #
cb.set_label('Intensity')
# ==================
# View Time Series
# ==================
view_timeseries(cpx_to_intensity(fp_seq[:, 100:300, 300:500]), cdi, title=f"White Light Timeseries",
subplt_cols=sp.tseries_cols,
logZ=False,
vlim=(1e-7, 1e-4),
)
plt.show()
# Run it!!!!!!!!!!!!!!!!!
# =======================================================================
sim = mm.RunMedis(name=testname, product='fields')
observation = sim()
cpx_sequence = observation['fields']
sampling = observation['sampling']
# =======================================================================
# Focal Plane Processing
# =======================================================================
# obs_sequence = np.array(obs_sequence) # obs sequence is returned by gen_timeseries (called above)
# (n_timesteps ,n_planes, n_waves_init, n_astro_bodies, nx ,ny)
cpx_sequence = opx.interp_wavelength(cpx_sequence,
2) # interpolate over wavelength
focal_plane = opx.extract_plane(cpx_sequence,
'detector') # eliminates object axis
# convert to intensity THEN sum over object, keeping the dimension of tstep even if it's one
focal_plane = np.sum(opx.cpx_to_intensity(focal_plane), axis=2)
fp_sampling = np.copy(sampling[cpx_sequence.shape[1] - 1, :])
# numpy arrays have some weird effects that make copying the array necessary
# =======================================================================
# Plotting
# =======================================================================
# White Light, Last Timestep
if sp.show_wframe:
# vlim = (np.min(spectralcube) * 10, np.max(spectralcube)) # setting z-axis limits
img = np.sum(focal_plane[sp.numframes - 1],
axis=0) # sum over wavelength
quick2D(
opx.extract_center(img), #focal_plane[sp.numframes-1]),
sp.save_to_disk = False # save obs_sequence (timestep, wavelength, x, y)
sp.save_list = ['atmosphere', 'entrance_pupil', 'woofer',
'detector'] # list of locations in optics train to save
if __name__ == '__main__':
# =======================================================================
# Run it!!!!!!!!!!!!!!!!!
# =======================================================================
cpx_sequence, sampling = mm.RunMedis().telescope()
# =======================================================================
# Focal Plane Processing
# =======================================================================
# cpx_sequence = (n_timesteps ,n_planes, n_waves_init, n_astro_bodies, nx ,ny)
# cpx_sequence = mmu.interp_wavelength(cpx_sequence, 2) # interpolate over wavelength
focal_plane = opx.extract_plane(cpx_sequence,
'detector') # eliminates astro_body axis
# convert to intensity THEN sum over object, keeping the dimension of tstep even if it's one
focal_plane = np.sum(opx.cpx_to_intensity(focal_plane), axis=2)
fp_sampling = sampling[-1, :]
# =======================================================================
# Plotting
# =======================================================================
# White Light, Last Timestep
if sp.show_wframe:
# vlim = (np.min(spectralcube) * 10, np.max(spectralcube)) # setting z-axis limits
img = np.sum(focal_plane[sp.numframes - 1],
axis=0) # sum over wavelength
quick2D(
opx.extract_center(img), #focal_plane[sp.numframes-1]),
title=f"White light image at timestep {sp.numframes} \n" # img