def spatial(allSF, kernel, npupil=None, norm=False, verbose=False):
# mask with nans
mask_nan = np.isnan(allSF) + (allSF == 0)
allSF[mask_nan] = np.nan
# get low and high spatial frequencies
with warnings.catch_warnings():
warnings.simplefilter("ignore") # NANs
LSF = astroconv.convolve(allSF, kernel, boundary='extend')
LSF[mask_nan] = np.nan
HSF = allSF - LSF
# print rms
if verbose is True:
print('rms(all SF) = %3.2f' % (np.nanstd(allSF)))
print('rms(LSF) = %3.2f' % (np.nanstd(LSF)))
print('rms(HSF) = %3.2f' % (np.nanstd(HSF)))
# normalize
if norm is True:
allSF /= np.nanstd(allSF)
LSF /= np.nanstd(LSF)
HSF /= np.nanstd(HSF)
allSF -= np.nanmean(allSF)
LSF -= np.nanmean(LSF)
HSF -= np.nanmean(HSF)
# remove nans
allSF = np.nan_to_num(allSF)
LSF = np.nan_to_num(LSF)
HSF = np.nan_to_num(HSF)
# resize outputs
if npupil is not None:
allSF = impro.resize_img(allSF, npupil)
LSF = impro.resize_img(LSF, npupil)
HSF = impro.resize_img(HSF, npupil)
return allSF, LSF, HSF
def spatial(allSF, kernel, new_size=None):
# low spatial frequencies
LSF = astroconv.convolve(allSF, kernel, boundary='extend')
LSF[allSF != allSF] = np.nan
# high spatial frequencies
HSF = allSF - LSF
HSF[allSF != allSF] = np.nan
# resize outputs
if new_size is not None:
allSF = impro.resize_img(allSF, new_size)
LSF = impro.resize_img(LSF, new_size)
HSF = impro.resize_img(HSF, new_size)
# print rms
if False:
print('rms(all SF) = %3.2f' % (np.nanstd(allSF)))
print('rms(LSF) = %3.2f' % (np.nanstd(LSF)))
print('rms(HSF) = %3.2f' % (np.nanstd(HSF)))
# normalize
if False:
allSF /= np.nanstd(allSF)
LSF /= np.nanstd(LSF)
HSF /= np.nanstd(HSF)
return allSF, LSF, HSF
def lyot_stop(wf, mode='RAVC', ravc_r=0.6, ls_dRext=0.03, ls_dRint=0.05,
ls_dRspi=0.04, spi_width=0.5, spi_angles=[0,60,120], diam_ext=37,
diam_int=11, ls_misalign=None, file_app_phase='', file_app_amp='',
ngrid=1024, npupil=285, margin=50, get_amp=False,
get_phase=False, verbose=False, **conf):
"""Add a Lyot stop, or an APP."""
# case 1: Lyot stop
if mode in ['CVC', 'RAVC']:
# LS parameters
r_obstr = ravc_r if mode in ['RAVC'] else diam_int/diam_ext
ls_int = r_obstr + ls_dRint
ls_ext = 1 - ls_dRext
ls_spi = spi_width/diam_ext + ls_dRspi
# LS misalignments
ls_misalign = [0,0,0,0,0,0] if ls_misalign is None else list(ls_misalign)
dx_amp, dy_amp, dz_amp = ls_misalign[0:3]
dx_phase, dy_phase, dz_phase = ls_misalign[3:6]
# create Lyot stop
proper.prop_circular_aperture(wf, ls_ext, dx_amp, dy_amp, NORM=True)
if diam_int > 0:
proper.prop_circular_obscuration(wf, ls_int, dx_amp, dy_amp, NORM=True)
if spi_width > 0:
for angle in spi_angles:
proper.prop_rectangular_obscuration(wf, ls_spi, 2, \
dx_amp, dy_amp, ROTATION=angle, NORM=True)
if verbose is True:
print('Create Lyot stop')
print(' ls_int=%3.4f, ls_ext=%3.4f, ls_spi=%3.4f'\
%(ls_int, ls_ext, ls_spi))
print('')
# case 2: APP
elif mode in ['APP']:
if verbose is True:
print('Load APP from files\n')
# get amplitude and phase data
APP_amp = fits.getdata(file_app_amp) if os.path.isfile(file_app_amp) \
else np.ones((npupil, npupil))
APP_phase = fits.getdata(file_app_phase) if os.path.isfile(file_app_phase) \
else np.zeros((npupil, npupil))
# resize to npupil
APP_amp = impro.resize_img(APP_amp, npupil)
APP_phase = impro.resize_img(APP_phase, npupil)
# pad with zeros to match PROPER ngrid
APP_amp = impro.pad_img(APP_amp, ngrid, 1)
APP_phase = impro.pad_img(APP_phase, ngrid, 0)
# multiply the loaded APP
proper.prop_multiply(wf, APP_amp*np.exp(1j*APP_phase))
# get the LS amplitude and phase for output
LS_amp = impro.crop_img(proper.prop_get_amplitude(wf), npupil, margin)\
if get_amp is True else None
LS_phase = impro.crop_img(proper.prop_get_phase(wf), npupil, margin)\
if get_phase is True else None
return wf, LS_amp, LS_phase
def remove_piston(filename):
data = np.float32(fits.getdata(filename))
data = crop_img(data, nimg)
data -= np.mean(data[mask != 0]) # remove piston
data[mask == 0] = 0
data = resize_img(data, npupil)
data = np.rot90(data) * 1e-6 # rotate, convert to meters
return data
def conv_kernel(Npup, cpp, HR=2**11):
ker_range = np.arange(-HR, HR) / HR
Xs, Ys = np.meshgrid(ker_range, ker_range) # high res kernel XY grid
Rs = np.abs(Xs + 1j * Ys) # high res kernel radii
kernel = np.ones((2 * HR, 2 * HR))
kernel[Rs > 1] = 0
# kernel must have odd dimensions
nkernel = int(Npup / cpp)
nkernel = nkernel + 1 if nkernel % 2 == 0 else nkernel
# resize kernel
kernel = impro.resize_img(kernel, nkernel)
kernel /= np.sum(kernel) # need to normalize the kernel
return kernel
def lyot_stop(wf, mode='RAVC', ravc_r=0.6, ls_dRext=0.03, ls_dRint=0.05,
ls_dRspi=0.04, spi_width=0.5, spi_angles=[0,60,120], diam_ext=37,
diam_int=11, diam_nominal=37, ls_misalign=None, ngrid=1024, npupil=285,
file_lyot_stop='', verbose=False, **conf):
""" Add a Lyot stop for a focal plane mask """
if mode in ['CVC', 'RAVC', 'CLC']:
# load lyot stop from file if provided
if os.path.isfile(file_lyot_stop):
if verbose is True:
print(" apply lyot stop from '%s'"%os.path.basename(file_lyot_stop))
# get amplitude and phase data
ls_mask = fits.getdata(file_lyot_stop)
# resize to npupil
ls_mask = resize_img(ls_mask, npupil)
# pad with zeros and add to wavefront
proper.prop_multiply(wf, pad_img(ls_mask, ngrid))
# if no lyot stop, create one
else:
# scale nominal values to pupil external diameter
scaling = diam_nominal/diam_ext
# LS parameters
r_obstr = ravc_r if mode in ['RAVC'] else diam_int/diam_ext
ls_int = r_obstr + ls_dRint*scaling
ls_ext = 1 - ls_dRext*scaling
ls_spi = spi_width/diam_ext + ls_dRspi*scaling
# LS misalignments
ls_misalign = [0,0,0,0,0,0] if ls_misalign is None else list(ls_misalign)
dx_amp, dy_amp, dz_amp = ls_misalign[0:3]
dx_phase, dy_phase, dz_phase = ls_misalign[3:6]
# create Lyot stop
proper.prop_circular_aperture(wf, ls_ext, dx_amp, dy_amp, NORM=True)
if diam_int > 0:
proper.prop_circular_obscuration(wf, ls_int, dx_amp, dy_amp, NORM=True)
if spi_width > 0:
for angle in spi_angles:
proper.prop_rectangular_obscuration(wf, 2*ls_spi, 2, \
dx_amp, dy_amp, ROTATION=angle, NORM=True)
if verbose is True:
print(' apply Lyot stop: ls_int=%s, ls_ext=%s, ls_spi=%s'\
%(round(ls_int, 4), round(ls_ext, 4), round(ls_spi, 4)))
return wf
def propagate_one(wf, phase_screen=None, tiptilt=None, misalign=None,
ngrid=1024, npupil=285, tag=None, onaxis=True, savefits=False, verbose=False,
**conf):
"""
Propagate one single wavefront.
An off-axis PSF can be obtained by switching onaxis to False,
thereby decentering the focal plane mask (if any).
"""
# update conf
conf.update(ngrid=ngrid, npupil=npupil)
# keep a copy of the input wavefront
wf1 = deepcopy(wf)
# apply phase screen (scao residuals, ncpa, petal piston)
if phase_screen is not None:
phase_screen = np.nan_to_num(phase_screen)
phase_screen = pad_img(resize_img(phase_screen, npupil), ngrid)
proper.prop_add_phase(wf1, phase_screen)
# apply tip-tilt (Zernike 2,3)
if tiptilt is not None:
# translate the tip/tilt from lambda/D into RMS phase errors
# RMS = x/2 = ((lam/D)*(D/2))/2 = lam/4
tiptilt = np.array(tiptilt, ndmin=1)*conf['lam']/4
proper.prop_zernikes(wf1, [2,3], tiptilt)
# update RAVC misalignment
conf.update(ravc_misalign=misalign)
# pupil-plane apodization
wf1, apo_amp, apo_phase = apodizer(wf1, get_amp=True, verbose=verbose, **conf)
# focal-plane mask, only in 'on-axis' configuration
if onaxis == True:
wf1 = fp_mask(wf1, verbose=verbose, **conf)
# Lyot-stop or APP
wf1, ls_amp, ls_phase = lyot_stop(wf1, get_amp=True, verbose=verbose, **conf)
# detector
psf = detector(wf1, verbose=verbose, **conf)
# save psf as fits file
if savefits == True:
tag = '' if tag is None else '%s_'%tag
name = '%s%s_PSF'%(tag, {True: 'onaxis', False: 'offaxis'}[onaxis])
save2fits(psf, name, **conf)
return psf
pscale = band_specs[band]['pscale']
modes = band_specs[band]['modes']
colors = band_specs[band]['colors']
markers = band_specs[band]['markers']
fig = plt.figure(figsize=figsize)
for mode, color, marker in zip(modes, colors, markers):
# if APP vertical band was replaced by horizontal one
replaced = '_replaced' if mode == 'APP' and APP_replaced is True else ''
# off-axis PSF
psf_OFF = fits.getdata(os.path.join(path_offaxis, 'offaxis_PSF_%s_%s.fits' \
%(band, mode)))
# save PSF size
npsf = psf_OFF.shape[1]
# resample
psf_OFF_rim = impro.resize_img(psf_OFF, 2 * rim)
# on-axis PSFs (cube)
psf_ON = fits.getdata(os.path.join(path_onaxis, 'onaxis_PSF_%s_%s%s.fits' \
%(band, mode, replaced)))
if psf_ON.ndim != 3:
psf_ON = np.array(psf_ON, ndmin=3)
# average
psf_ON_avg = np.mean(psf_ON, 0)
# resample
psf_ON_rim = impro.resize_img(psf_ON_avg, 2 * rim)
# radial profiles
y1 = impro.get_radial_profile(psf_OFF_rim, (xo, yo), xbin)[:-1]
y2 = impro.get_radial_profile(psf_ON_rim, (xo, yo), xbin)[:-1]
# normalize by the peak of the off-axis PSF
peak = np.max(y1)
y1 /= peak
def fp_mask(wf, mode='RAVC', focal=660, verbose=False, **conf):
# case 1: vortex coronagraphs
if mode in ['CVC', 'RAVC']:
if verbose is True:
print('Apply Vortex phase mask')
# update conf
conf.update(focal=focal)
# load vortex calibration files:
# conf['psf_num'], conf['vvc'], conf['perf_num']
conf = vortex_init(verbose=verbose, **conf)
# propagate to vortex
lens(wf, focal)
# apply vortex
scale_psf = wf._wfarr[0,0]/conf['psf_num'][0,0]
wf_corr = (conf['psf_num']*conf['vvc'] - conf['perf_num'])*scale_psf
wf._wfarr = wf._wfarr*conf['vvc'] - wf_corr
# propagate to lyot stop
lens(wf, focal)
# TODO: cleanup CLC
# case 2: classical Lyot
elif mode in ['CLC']:
if verbose is True:
print('Apply Classical Lyot mask\n')
f_lens = conf['focal']
beam_ratio = conf['beam_ratio']
CLC_diam = conf['CLC_diam'] # classical lyot diam in lam/D (default to 4)
gridsize = conf['gridsize']
tmp_dir = conf['temp_dir']
# propagate to lyot mask
lens(wf, conf['focal'])
# create or load the classical Lyot mask
calib = str(CLC_diam)+str('_')+str(int(beam_ratio*100))+str('_')+str(gridsize)
my_file = os.path.join(tmp_dir, 'clc_'+calib+'.fits')
if not os.path.isfile(my_file):
# calculate exact size of Lyot mask diameter, in pixels
Dmask = CLC_diam/beam_ratio
# oversample the Lyot mask (round up)
samp = 100
ndisk = int(samp*np.ceil(Dmask))
ndisk = ndisk + 1 if not ndisk % 2 else ndisk # must be odd
# find center
cdisk = int((ndisk - 1)/2)
# calculate the distances to center
xy = range(-cdisk, cdisk + 1)
x,y = np.meshgrid(xy, xy)
dist = np.sqrt(x**2 + y**2)
# create the Lyot mask
mask = np.zeros((ndisk, ndisk))
mask[np.where(dist > samp*Dmask/2)] = 1
# resize to Lyot mask real size, and pad with ones
mask = resize_img(mask, int(ndisk/samp))
mask = pad_img(mask, gridsize, 1)
# write mask
fits.writeto(my_file, mask)
else:
mask = fits.getdata(my_file)
# apply lyot mask
mask = proper.prop_shift_center(mask)
wf._wfarr.real *= mask
# propagate to lyot stop
lens(wf, conf['focal'])
return wf
def create_pupil(nhr=2**10,
npupil=285,
pupil_img_size=40,
diam_ext=37,
diam_int=11,
spi_width=0.5,
spi_angles=[0, 60, 120],
seg_width=0,
seg_gap=0,
seg_rms=0,
seg_ny=[
10, 13, 16, 19, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31,
30, 31, 30, 31, 30, 31, 30, 31, 30, 29, 28, 27, 26, 25,
24, 23, 22, 19, 16, 13, 10
],
seg_missing=[],
seed=123456,
**conf):
''' Create a pupil.
Args:
nhr: int
high resolution grid
npupil: int
number of pixels of the pupil
pupil_img_size: float
pupil image (for PROPER) in m
diam_ext: float
outer circular aperture in m
diam_int: float
central obscuration in m
spi_width: float
spider width in m
spi_angles: list of float
spider angles in deg
seg_width: float
segment width in m
seg_gap: float
gap between segments in m
seg_rms: float
rms of the reflectivity of all segments
seg_ny: list of int
number of hexagonal segments per column (from left to right)
seg_missing: list of tupples
coordinates of missing segments
'''
# create a high res pupil with PROPER of even size (nhr)
nhr_size = pupil_img_size * nhr / (nhr - 1)
wf_tmp = proper.prop_begin(nhr_size, 1, nhr, diam_ext / nhr_size)
if diam_ext > 0:
proper.prop_circular_aperture(wf_tmp, 1, NORM=True)
if diam_int > 0:
proper.prop_circular_obscuration(wf_tmp,
diam_int / diam_ext,
NORM=True)
if spi_width > 0:
for angle in spi_angles:
proper.prop_rectangular_obscuration(wf_tmp, spi_width/nhr_size, 2, \
ROTATION=angle, NORM=True)
pup = proper.prop_get_amplitude(wf_tmp)
# crop the pupil to odd size (nhr-1), and resize to npupil
pup = pup[1:, 1:]
pup = resize_img(pup, npupil)
# add segments
if seg_width > 0:
segments = np.zeros((nhr, nhr))
# sampling in meters/pixel
sampling = pupil_img_size / nhr
# dist between center of two segments, side by side
seg_d = seg_width * np.cos(np.pi / 6) + seg_gap
# segment radius
seg_r = seg_width / 2
# segment radial distance wrt x and y axis
seg_ny = np.array(seg_ny)
seg_nx = len(seg_ny)
seg_rx = np.arange(seg_nx) - (seg_nx - 1) / 2
seg_ry = (seg_ny - 1) / 2
# loop through segments
np.random.seed(seed)
for i in range(seg_nx):
seg_x = seg_rx[i] * seg_d * np.cos(np.pi / 6)
seg_y = -seg_ry[i] * seg_d
for j in range(1, seg_ny[i] + 1):
# removes secondary and if any missing segment is present
if (np.sqrt(seg_x**2 + seg_y**2) <= 4.01*seg_d) \
or ((seg_rx[i], j) in seg_missing):
seg_y += seg_d
else:
# creates one hexagonal segment at x, y position in meters
segment = create_hexagon(nhr, seg_r, seg_y, seg_x,
sampling)
# multiply by segment reflectivity and add to segments
seg_refl = np.random.normal(1, seg_rms)
segments += segment * seg_refl
seg_y += seg_d
# need to transpose, due to the orientation of hexagons in create_hexagon
segments = segments.T
# resize to npupil, and add to pupil
segments = resize_img(segments, npupil)
pup *= segments
return pup
def lyotmask_init(lyotmask_calib='',
dir_temp='',
clc_diam=80,
pscale=5.47,
magnification=100,
ngrid=1024,
beam_ratio=0.26,
verbose=False,
**conf):
'''
Creates/writes classical lyot masks, or loads them if files
already exist.
The following parameters will be added to conf:
lyotmask_calib, lyotmask
Returns: conf (updated and sorted)
'''
# update conf with local variables (remove unnecessary)
conf.update(locals())
[conf.pop(key) for key in ['conf', 'verbose'] if key in conf]
# check if lyot mask already loaded for this calib
calib = 'lyotmask_%s_%s_%3.4f' % (clc_diam, ngrid, beam_ratio)
if lyotmask_calib == calib:
return conf
else:
# check for existing file
filename = os.path.join(dir_temp, '%s.fits' % calib)
if os.path.isfile(filename):
if verbose is True:
print(' loading lyot mask')
lyotmask = fits.getdata(os.path.join(dir_temp, filename))
# create file
else:
if verbose is True:
print(" writing lyot mask")
# calculate the size (in pixels) of the lyot mask,
# rounded up to an odd value
nmask = int(np.ceil(clc_diam / pscale))
nmask += 1 - nmask % 2
# create a magnified lyot mask
rmask = clc_diam / pscale / nmask
r, t = polar_coord(nmask * magnification)
lyotmask = np.zeros(np.shape(r))
lyotmask[r > rmask] = 1
# resize and pad with ones to amtch ngrid
lyotmask = pad_img(resize_img(lyotmask, nmask), ngrid, 1)
# save as fits file
fits.writeto(os.path.join(dir_temp, filename),
np.float32(lyotmask),
overwrite=True)
# shift the lyotmask amplitude
lyotmask = proper.prop_shift_center(lyotmask)
# add lyotmask amplitude at the end of conf
conf = {k: v for k, v in sorted(conf.items())}
conf.update(lyotmask_calib=calib, lyotmask=lyotmask)
if verbose is True:
print(' clc_diam=%s, ngrid=%s, beam_ratio=%3.4f'%\
(clc_diam, ngrid, beam_ratio))
return conf
def propagate_cube(wf,
phase_screens,
amp_screens,
tiptilts,
misaligns,
cpu_count=1,
vc_chrom_leak=2e-3,
add_cl_det=False,
add_cl_vort=False,
tag=None,
onaxis=True,
send_to=None,
savefits=False,
verbose=False,
**conf):
# update conf
conf.update(cpu_count=cpu_count, vc_chrom_leak=vc_chrom_leak, \
add_cl_det=add_cl_det, add_cl_vort=add_cl_vort, tag=tag, onaxis=onaxis)
# preload amp screen if only one frame
if len(amp_screens) == 1 and np.any(amp_screens) != None:
import proper
from heeps.util.img_processing import pad_img, resize_img
amp_screens = np.nan_to_num(amp_screens[0])
amp_screens = pad_img(resize_img(amp_screens, conf['npupil']),
conf['ngrid'])
proper.prop_multiply(wf, amp_screens)
# then create a cube of None values
amp_screens = [None] * int((conf['nframes'] / conf['nstep']) + 0.5)
# preload apodizer when no drift
if np.all(misaligns) == None or 'APP' in conf['mode']:
wf = apodizer(wf, verbose=False, **conf)
if verbose == True:
print('Create %s-axis PSF cube' % {True: 'on', False: 'off'}[onaxis])
if add_cl_det is True:
print(' adding chromatic leakage at detector plane: %s' %
vc_chrom_leak)
if add_cl_vort is True:
print(' adding chromatic leakage at vortex plane: %s' %
vc_chrom_leak)
# run simulation
posvars = [phase_screens, amp_screens, tiptilts, misaligns]
kwargs = dict(verbose=False, **conf)
psfs = multiCPU(propagate_one, posargs=[wf], posvars=posvars, kwargs=kwargs, \
case='e2e simulation', cpu_count=cpu_count)
# if only one wavefront, make dim = 2
if len(psfs) == 1:
psfs = psfs[0]
# save cube of PSFs to fits file, and notify by email
if savefits == True:
tag = '' if tag is None else '%s_' % tag
name = '%s%s_PSF' % (tag, {True: 'onaxis', False: 'offaxis'}[onaxis])
filename = save2fits(psfs, name, **conf)
notify('saved to %s' % filename, send_to)
return psfs
def propagate_cube(wf,
phase_screens,
amp_screens,
tiptilts,
misaligns,
cpu_count=1,
send_to=None,
tag=None,
onaxis=True,
savefits=False,
verbose=False,
**conf):
# preload amp screen if only one frame
if len(amp_screens) == 1 and np.any(amp_screens) != None:
import proper
from heeps.util.img_processing import pad_img, resize_img
amp_screens = np.nan_to_num(amp_screens[0])
amp_screens = pad_img(resize_img(amp_screens, conf['npupil']),
conf['ngrid'])
proper.prop_multiply(wf, amp_screens)
# then create a cube of None values
amp_screens = [None] * int((conf['nframes'] / conf['nstep']) + 0.5)
# preload apodizer when no drift
if np.all(misaligns) == None or 'APP' in conf['mode']:
wf = apodizer(wf, verbose=False, **conf)
if verbose == True:
print('Create %s-axis PSF cube' % {True: 'on', False: 'off'}[onaxis])
# run simulation
t0 = time.time()
if cpu_count != 1 and platform in ['linux', 'linux2', 'darwin']:
if cpu_count == None:
cpu_count = mpro.cpu_count() - 1
if verbose is True:
print(' %s: e2e simulation starts, using %s cores'\
%(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), cpu_count))
p = mpro.Pool(cpu_count)
func = partial(propagate_one, wf, onaxis=onaxis, verbose=False, **conf)
psfs = np.array(
p.starmap(func, zip(phase_screens, amp_screens, tiptilts,
misaligns)))
p.close()
p.join()
else:
if verbose is True:
print(' %s: e2e simulation starts, using 1 core'\
%(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())))
for i, (phase_screen, amp_screen, tiptilt, misalign) \
in enumerate(zip(phase_screens, amp_screens, tiptilts, misaligns)):
psf = propagate_one(wf, phase_screen, amp_screen, tiptilt, misalign, \
onaxis=onaxis, verbose=True, **conf)
psfs = psf if i == 0 else np.dstack((psfs.T, psf.T)).T
if verbose is True:
print(' %s: finished, elapsed %.3f seconds'\
%(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), time.time() - t0))
# if only one wavefront, make dim = 2
if len(psfs) == 1:
psfs = psfs[0]
# save cube of PSFs to fits file, and notify by email
if savefits == True:
tag = '' if tag is None else '%s_' % tag
name = '%s%s_PSF' % (tag, {True: 'onaxis', False: 'offaxis'}[onaxis])
filename = save2fits(psfs, name, **conf)
notify('saved to %s' % filename, send_to)
return psfs
for band in bands:
print(' %s' % band, end=', ')
pscale = band_specs[band]['pscale']
modes = band_specs[band]['modes']
for mode in modes:
print(' %s' % mode)
# if APP vertical band was replaced by horizontal one
replaced = '_replaced' if mode == 'APP' and APP_replaced is True else ''
# off-axis PSF
psf_OFF = fits.getdata(os.path.join(folder, 'offaxis_PSF_%s_%s.fits' \
%(band, mode)))
# save PSF size
npsf = psf_OFF.shape[1]
# resample
psf_OFF_ndet = impro.resize_img(psf_OFF, ndet)
# on-axis PSFs (cube)
psf_ON = fits.getdata(os.path.join(folder, 'onaxis_PSF_%s_%s%s.fits' \
%(band, mode, replaced)))
if psf_ON.ndim != 3:
psf_ON = np.array(psf_ON, ndmin=3)
# average
psf_ON_avg = np.mean(psf_ON, 0)
# resample
psf_ON_ndet = impro.resize_img(psf_ON_avg, ndet)
# radial profiles
y1 = impro.get_radial_profile(psf_OFF_ndet, (rim, rim), xbin)[:-1]
y2 = impro.get_radial_profile(psf_ON_ndet, (rim, rim), xbin)[:-1]
# normalize by the peak of the off-axis PSF
peak = np.max(y1)
y1 /= peak
def apodizer(wf,
mode='RAVC',
ravc_t=0.8,
ravc_r=0.6,
ravc_misalign=None,
ngrid=1024,
npupil=285,
file_ravc_amp='',
file_ravc_phase='',
margin=50,
get_amp=False,
get_phase=False,
verbose=False,
**conf):
''' Create a wavefront object at the entrance pupil plane.
The pupil is either loaded from a fits file, or created using
pupil parameters.
Can also select only one petal and mask the others.
wf: WaveFront
PROPER wavefront object
mode: str
HCI mode
ravc_t: float
RA transmittance
ravc_r: float
RA radius
ravc_misalign: list of float
RA misalignment
ngrid: int
number of pixels of the wavefront array
npupil: int
number of pixels of the pupil
file_ravc_amp: str
file_ravc_phase: str
ring apodizer files (optional)
'''
if mode in ['RAVC']:
# load apodizer from files if provided
if os.path.isfile(file_ravc_amp) and os.path.isfile(file_ravc_phase):
if verbose is True:
print('Load ring apodizer from files\n')
# get amplitude and phase data
RAVC_amp = fits.getdata(file_ravc_amp)
RAVC_phase = fits.getdata(file_ravc_phase)
# resize to npupil
RAVC_amp = impro.resize_img(RAVC_amp, npupil)
RAVC_phase = impro.resize_img(RAVC_phase, npupil)
# pad with zeros to match PROPER gridsize
RAVC_amp = impro.pad_img(RAVC_amp, ngrid)
RAVC_phase = impro.pad_img(RAVC_phase, ngrid)
# build complex apodizer
apo = RAVC_amp * np.exp(1j * RAVC_phase)
# or else, define the apodizer as a ring (with % misalignments)
else:
# RAVC misalignments
ravc_misalign = [
0, 0, 0, 0, 0, 0
] if ravc_misalign is None else list(ravc_misalign)
dx_amp, dy_amp, dz_amp = ravc_misalign[0:3]
dx_phase, dy_phase, dz_phase = ravc_misalign[3:6]
# create apodizer
apo = circular_apodization(wf, ravc_r, 1., ravc_t, xc=dx_amp, \
yc=dy_amp, NORM=True)
apo = proper.prop_shift_center(apo)
if verbose is True:
print('Create ring apodizer')
print(' ravc_t=%3.4f, ravc_r=%3.4f'\
%(ravc_t, ravc_r))
print(' ravc_misalign=%s' % ravc_misalign)
print('')
# multiply the loaded apodizer
proper.prop_multiply(wf, apo)
# get the apodizer amplitude and phase for output
apo_amp = impro.crop_img(proper.prop_get_amplitude(wf), npupil,\
margin) if get_amp is True else None
apo_phase = impro.crop_img(proper.prop_get_phase(wf), npupil,\
margin) if get_phase is True else None
return wf, apo_amp, apo_phase
else: # no ring apodizer
return wf, None, None
def apodizer(wf,
mode='RAVC',
ravc_t=0.8,
ravc_r=0.6,
ngrid=1024,
npupil=285,
f_app_amp='',
f_app_phase='',
f_ravc_amp='',
f_ravc_phase='',
apo_misalign=None,
onaxis=True,
verbose=False,
**conf):
''' Create a wavefront object at the entrance pupil plane.
The pupil is either loaded from a fits file, or created using
pupil parameters.
Can also select only one petal and mask the others.
wf: WaveFront
PROPER wavefront object
mode: str
HCI mode
ravc_t: float
RA transmittance
ravc_r: float
RA radius
ngrid: int
number of pixels of the wavefront array
npupil: int
number of pixels of the pupil
f_app_amp: str
f_app_phase: str
apodizing phase plate files
f_ravc_amp: str
f_ravc_phase: str
ring apodizer files (optional)
apo_misalign: list of float
apodizer misalignment
'''
# case 1: Ring Apodizer
if 'RAVC' in mode and ravc_r > 0:
# load apodizer from files if provided
if os.path.isfile(f_ravc_amp) and os.path.isfile(f_ravc_phase):
if verbose is True:
print(' apply ring apodizer from files')
# get amplitude and phase data
RAVC_amp = fits.getdata(f_ravc_amp)
RAVC_phase = fits.getdata(f_ravc_phase)
# resize to npupil
RAVC_amp = impro.resize_img(RAVC_amp, npupil)
RAVC_phase = impro.resize_img(RAVC_phase, npupil)
# pad with zeros to match PROPER gridsize
RAVC_amp = impro.pad_img(RAVC_amp, ngrid)
RAVC_phase = impro.pad_img(RAVC_phase, ngrid)
# build complex apodizer
ring = RAVC_amp * np.exp(1j * RAVC_phase)
# else, define the apodizer as a ring (with % misalignments)
else:
# RAVC misalignments
dx, dy = [0, 0
] if apo_misalign is None else list(apo_misalign)[0:2]
# create apodizer
ring = circular_apodization(wf,
ravc_r,
1,
ravc_t,
xc=dx,
yc=dy,
NORM=True)
if verbose is True:
print(' apply ring apodizer: ravc_t=%s, ravc_r=%s' %
(round(ravc_t, 4), round(ravc_r, 4)))
# multiply the loaded apodizer
proper.prop_multiply(wf, ring)
# case 2: Apodizing Phase Plate
elif 'APP' in mode:
# get amplitude and phase data
if os.path.isfile(f_app_amp):
if verbose is True:
print(' apply APP stop (amplitude)')
APP_amp = fits.getdata(f_app_amp)
else:
APP_amp = np.ones((npupil, npupil))
if os.path.isfile(f_app_phase) and onaxis == True:
if verbose is True:
print(' apply APP phase')
APP_phase = fits.getdata(f_app_phase)
else:
APP_phase = np.zeros((npupil, npupil))
# resize to npupil
APP_amp = impro.resize_img(APP_amp, npupil)
APP_phase = impro.resize_img(APP_phase, npupil)
# rotate for negative PSF
if 'neg' in mode:
APP_amp = np.rot90(APP_amp, 2)
APP_phase = np.rot90(APP_phase, 2)
# pad with zeros to match PROPER ngrid
APP_amp = impro.pad_img(APP_amp, ngrid, 0)
APP_phase = impro.pad_img(APP_phase, ngrid, 0)
# multiply the loaded APP
proper.prop_multiply(wf, APP_amp * np.exp(1j * APP_phase))
return wf
savename = 'cube_%s_%ss_%sms_0piston_meters_scao_only_%s_%s.fits' % (
tag, duration, samp, band, npupil)
#savename = 'cube_%s_%ss_%sms_0piston_meters_scao_only_%s_WVseeing.fits'%(tag, duration, samp, npupil)
#input_folder = '/mnt/disk4tb/METIS/METIS_COMPASS_RAW_PRODUCTS/gorban_metis_baseline_Cbasic_2020-10-16T10:25:14/residualPhaseScreens'
#input_folder = '/mnt/disk4tb/METIS/METIS_COMPASS_RAW_PRODUCTS/gorban_metis_baseline_Cbasic_2020-11-05T12:40:27/residualPhaseScreens'
#input_folder = '/mnt/disk4tb/METIS/METIS_COMPASS_RAW_PRODUCTS/gorban_metis_baseline_Cbasic_2020-11-30T20:52:24/residualPhaseScreens'
#input_folder = '/mnt/disk4tb/METIS/METIS_COMPASS_RAW_PRODUCTS/gorban_metis_baseline_Cbasic_uncorrected_2021-06-01T12:02:36/residualPhaseScreens'
input_folder = '/mnt/disk12tb/Users/gorban/METIS/METIS_COMPASS/gorban_metis_baseline_Cfull_noWtt_2021-10-07T09:00:32/residualPhaseScreens'
output_folder = 'wavefront/cfull'
cpu_count = None
# mask
mask = fits.getdata(os.path.join(input_folder, 'Telescope_Pupil.fits'))
mask = crop_img(mask, nimg)
mask_pupil = np.rot90(resize_img(mask, npupil))
fits.writeto(os.path.join(output_folder,
'mask_%s_%s_%s.fits' % (tag, band, npupil)),
np.float32(mask_pupil),
overwrite=True)
# filenames
nframes = len(
[name for name in os.listdir(input_folder) if name.startswith(prefix)])
nframes = 12000
print('%s frames' % nframes)
frames = [str(frame).zfill(6) if pad_frame is True else str(frame) \
for frame in range(start, start + nframes*samp, samp)]
filenames = np.array([os.path.join(input_folder, '%s%s%s.fits'%(prefix, frame, suffix)) \
for frame in frames])
def pupil(file_pupil='',
lam=3.8e-6,
ngrid=1024,
npupil=285,
pupil_img_size=40,
diam_ext=37,
diam_int=11,
spi_width=0.5,
spi_angles=[0, 60, 120],
npetals=6,
seg_width=0,
seg_gap=0,
seg_rms=0,
seg_ny=[
10, 13, 16, 19, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 30, 31,
30, 31, 30, 31, 30, 31, 30, 29, 28, 27, 26, 25, 24, 23, 22, 19,
16, 13, 10
],
seg_missing=[],
select_petal=None,
savefits=False,
verbose=False,
**conf):
''' Create a wavefront object at the entrance pupil plane.
The pupil is either loaded from a fits file, or created using
pupil parameters.
Can also select only one petal and mask the others.
Args:
dir_output (str):
path to saved pupil file
band (str):
spectral band (e.g. 'L', 'M', 'N1', 'N2')
mode (str):
HCI mode: RAVC, CVC, APP, CLC
file_pupil: str
path to a pupil fits file
lam: float
wavelength in m
ngrid: int
number of pixels of the wavefront array
npupil: int
number of pixels of the pupil
pupil_img_size: float
pupil image (for PROPER) in m
diam_ext: float
outer circular aperture in m
diam_int: float
central obscuration in m
spi_width: float
spider width in m
spi_angles: list of float
spider angles in deg
seg_width: float
segment width in m
seg_gap: float
gap between segments in m
seg_rms: float
rms of the reflectivity of all segments
seg_ny: list of int
number of hexagonal segments per column (from left to right)
seg_missing: list of tupples
coordinates of missing segments
npetals: int
number of petals
select_petal: int
selected petal in range npetals, default to None
'''
# initialize wavefront using PROPER
beam_ratio = npupil / ngrid * (diam_ext / pupil_img_size)
wf = proper.prop_begin(pupil_img_size, lam, ngrid, beam_ratio)
# load pupil file
if os.path.isfile(file_pupil):
if verbose is True:
print("Load pupil from '%s'" % os.path.basename(file_pupil))
#conf.update()
pup = fits.getdata(file_pupil).astype(np.float32)
# resize to npupil
pup = resize_img(pup, npupil)
# if no pupil file, create a pupil
else:
if verbose is True:
print("Create pupil: spi_width=%s m, seg_width=%s m, seg_gap=%s m, seg_rms=%s"\
%(spi_width, seg_width, seg_gap, seg_rms))
conf.update(npupil=npupil,
pupil_img_size=pupil_img_size,
diam_ext=diam_ext,
diam_int=diam_int,
spi_width=spi_width,
spi_angles=spi_angles,
seg_width=seg_width,
seg_gap=seg_gap,
seg_ny=seg_ny,
seg_missing=seg_missing,
seg_rms=seg_rms)
pup = create_pupil(**conf)
# normalize the entrance pupil intensity (total flux = 1)
I_pup = pup**2
pup = np.sqrt(I_pup / np.sum(I_pup))
# pad with zeros and add to wavefront
proper.prop_multiply(wf, pad_img(pup, ngrid))
# select one petal (optional)
if select_petal in range(npetals) and npetals > 1:
if verbose is True:
print(" select_petal=%s" % select_petal)
# petal start and end angles
pet_angle = 2 * np.pi / npetals
pet_start = pet_angle / 2 + (select_petal - 1) * pet_angle
pet_end = pet_start + pet_angle
# petal angles must be 0-2pi
ti %= (2 * np.pi)
pet_start %= (2 * np.pi)
pet_end %= (2 * np.pi)
# check if petal crosses 0 angle
if pet_end - pet_start < 0:
pet_start = (pet_start + np.pi) % (2 * np.pi) - np.pi
pet_end = (pet_end + np.pi) % (2 * np.pi) - np.pi
ti = (ti + np.pi) % (2 * np.pi) - np.pi
# create petal and add to grid
petal = ((ti >= pet_start) * (ti <= pet_end)).astype(int)
petal = resize_img(petal, ngrid)
proper.prop_multiply(wf, petal)
# save pupil as fits file
if savefits == True:
save2fits(pup, 'pupil', **conf)
if verbose is True:
print(' diam=%s m, resize to %s pix, zero-pad to %s pix\n'\
%(round(diam_ext, 2), npupil, ngrid))
return wf
def vortex_init(vortex_calib='',
dir_temp='',
diam_ext=37,
lam=3.8,
ngrid=1024,
beam_ratio=0.26,
focal=660,
vc_charge=2,
verbose=False,
**conf):
'''
Creates/writes vortex back-propagation fitsfiles, or loads them if files
already exist.
The following parameters will be added to conf:
vortex_calib, psf_num, perf_num, vvc
Returns: conf (updated and sorted)
'''
# update conf with local variables (remove unnecessary)
conf.update(locals())
[conf.pop(key) for key in ['conf', 'verbose'] if key in conf]
# check if back-propagation params already loaded for this calib
calib = 'vortex_%s_%s_%3.4f' % (vc_charge, ngrid, beam_ratio)
if vortex_calib == calib:
return conf
else:
# check for existing file
filename = os.path.join(dir_temp, '%s.fits' % calib)
if os.path.isfile(filename):
if verbose is True:
print(' loading vortex back-propagation params')
data = fits.getdata(os.path.join(dir_temp, filename))
# read the pre-vortex field
psf_num = data[0] + 1j * data[1]
# read the theoretical vortex field
vvc = data[2] + 1j * data[3]
# read the perfect-result vortex field
perf_num = data[4] + 1j * data[5]
# create files
else:
if verbose is True:
print(" writing vortex back-propagation params")
# create circular pupil
wf_tmp = proper.prop_begin(diam_ext, lam, ngrid, beam_ratio)
proper.prop_circular_aperture(wf_tmp, 1, NORM=True)
# propagate to vortex
lens(wf_tmp, focal)
# pre-vortex field
psf_num = deepcopy(wf_tmp.wfarr)
# vortex phase ramp is oversampled for a better discretization
ramp_oversamp = 11.
nramp = int(ngrid * ramp_oversamp)
start = -nramp / 2 - int(ramp_oversamp) / 2 + 0.5
end = nramp / 2 - int(ramp_oversamp) / 2 + 0.5
Vp = np.arange(start, end, 1.)
# Pancharatnam Phase = arg<Vref,Vp> (horizontal input polarization)
Vref = np.ones(Vp.shape)
prod = np.outer(Vref, Vp)
phiPan = np.angle(prod + 1j * prod.T)
# vortex phase ramp exp(ilphi)
ofst = 0
ramp_sign = 1
vvc_tmp = np.exp(1j * (ramp_sign * vc_charge * phiPan + ofst))
vvc = np.array(impro.resize_img(vvc_tmp.real, ngrid),
dtype=complex)
vvc.imag = impro.resize_img(vvc_tmp.imag, ngrid)
phase_ramp = np.angle(vvc)
# theoretical vortex field
vvc_complex = np.array(np.zeros((ngrid, ngrid)), dtype=complex)
vvc_complex.imag = phase_ramp
vvc = np.exp(vvc_complex)
# apply vortex
proper.prop_multiply(wf_tmp, vvc)
# null the amplitude inside the Lyot Stop, and back propagate
lens(wf_tmp, focal)
proper.prop_circular_obscuration(wf_tmp, 1., NORM=True)
lens(wf_tmp, -focal)
# perfect-result vortex field
perf_num = deepcopy(wf_tmp.wfarr)
# write all fields
data = np.dstack((psf_num.real.T, psf_num.imag.T, vvc.real.T, vvc.imag.T,\
perf_num.real.T, perf_num.imag.T)).T
fits.writeto(os.path.join(dir_temp, filename),
np.float32(data),
overwrite=True)
# shift the phase ramp
vvc = proper.prop_shift_center(vvc)
# add vortex back-propagation parameters at the end of conf
conf = {k: v for k, v in sorted(conf.items())}
conf.update(vortex_calib=calib,
psf_num=psf_num,
vvc=vvc,
perf_num=perf_num)
if verbose is True:
print(' vc_charge=%s, ngrid=%s, beam_ratio=%3.4f'%\
(vc_charge, ngrid, beam_ratio))
return conf
def pupil(pup=None,
f_pupil='',
lam=3.8e-6,
ngrid=1024,
npupil=285,
pupil_img_size=40,
diam_ext=37,
diam_int=11,
spi_width=0.5,
spi_angles=[0, 60, 120],
npetals=6,
seg_width=0,
seg_gap=0,
seg_rms=0,
seg_ny=[
10, 13, 16, 19, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 30, 31,
30, 31, 30, 31, 30, 31, 30, 29, 28, 27, 26, 25, 24, 23, 22, 19,
16, 13, 10
],
seg_missing=[],
select_petal=None,
norm_I=True,
savefits=False,
verbose=False,
**conf):
''' Create a wavefront object at the entrance pupil plane.
The pupil is either loaded from a fits file, or created using
pupil parameters.
Can also select only one petal and mask the others.
Args:
dir_output (str):
path to saved pupil file
band (str):
spectral band (e.g. 'L', 'M', 'N1', 'N2')
mode (str):
HCI mode: RAVC, CVC, APP, CLC
f_pupil: str
path to a pupil fits file
lam: float
wavelength in m
ngrid: int
number of pixels of the wavefront array
npupil: int
number of pixels of the pupil
pupil_img_size: float
pupil image (for PROPER) in m
diam_ext: float
outer circular aperture in m
diam_int: float
central obscuration in m
spi_width: float
spider width in m
spi_angles: list of float
spider angles in deg
seg_width: float
segment width in m
seg_gap: float
gap between segments in m
seg_rms: float
rms of the reflectivity of all segments
seg_ny: list of int
number of hexagonal segments per column (from left to right)
seg_missing: list of tupples
coordinates of missing segments
npetals: int
number of petals
select_petal: int
selected petal in range npetals, default to None
'''
# initialize wavefront using PROPER
beam_ratio = npupil / ngrid * (diam_ext / pupil_img_size)
wf = proper.prop_begin(diam_ext, lam, ngrid, beam_ratio)
# case 1: load pupil from data
if pup is not None:
if verbose is True:
print("Load pupil data from 'pup'")
pup = resize_img(pup, npupil)
# case 2: load pupil from file
elif os.path.isfile(f_pupil):
if verbose is True:
print("Load pupil from '%s'" % os.path.basename(f_pupil))
pup = resize_img(fits.getdata(f_pupil), npupil)
# case 3: create a pupil
else:
if verbose is True:
print("Create pupil: spi_width=%s m, seg_width=%s m, seg_gap=%s m, seg_rms=%s"\
%(spi_width, seg_width, seg_gap, seg_rms))
conf.update(npupil=npupil,
pupil_img_size=pupil_img_size,
diam_ext=diam_ext,
diam_int=diam_int,
spi_width=spi_width,
spi_angles=spi_angles,
seg_width=seg_width,
seg_gap=seg_gap,
seg_ny=seg_ny,
seg_missing=seg_missing,
seg_rms=seg_rms)
pup = create_pupil(**conf)
# select one petal (optional)
if select_petal in range(npetals) and npetals > 1:
if verbose is True:
print(" select_petal=%s" % select_petal)
petal = create_petal(select_petal, npetals, npupil)
pup *= petal
# normalize the entrance pupil intensity (total flux = 1)
if norm_I is True:
I_pup = pup**2
pup = np.sqrt(I_pup / np.sum(I_pup))
# save pupil as fits file
if savefits == True:
save2fits(pup, 'pupil', **conf)
# pad with zeros and add to wavefront
proper.prop_multiply(wf, pad_img(pup, ngrid))
return wf
