def psd_spatial_zernike(cube_name, pup, zpols, nzer, ncube):
spsd_name = cube_name[:-5] + '_%s' + '_%s.fits'%nzer
try:
spsd = fits.getdata(spsd_name%'spsd')
print('getdata ' + spsd_name%'spsd')
except FileNotFoundError:
print('writeto ' + spsd_name%'spsd')
cube = fits.getdata(cube_name)[:ncube]
nimg = cube.shape[-1]
pup = resize_cube(pup, nimg)
zpols = zpols[:ncube,:nzer]
wf = proper.prop_begin(1, 1, nimg, 1) # initial wavefront
LSFs = np.empty((nzer, ncube, nimg, nimg))
HSFs = np.empty((nzer, ncube, nimg, nimg))
HSFs_rms = []
for z in np.arange(nzer) + 1:
verbose = True if z == 1 else False
LSF, HSF = multiCPU(remove_zernike, posargs=[deepcopy(wf), pup],
posvars=[cube, zpols[:,:z]], case='remove zernike',
nout=2, verbose=verbose)
LSFs[z-1], HSFs[z-1] = LSF, HSF
HSFs_rms.append(get_rms(HSF, verbose=verbose))
print(z, end=', ')
spsd = [rms**2 for rms in HSFs_rms]
spsd = [spsd[0]] + spsd
fits.writeto(spsd_name%'spsd', np.float32(spsd))
fits.writeto(spsd_name%'LSFs', np.float32(LSFs))
fits.writeto(spsd_name%'HSFs', np.float32(HSFs))
return spsd
def crop_cube(cube, new_size, margin=0, cpu_count=None, verbose=False):
posvars = [cube, [new_size] * len(cube)]
kwargs = dict(margin=margin, verbose=False)
if cube.ndim < 3:
return crop_img(cube, new_size, **kwargs)
else:
return multiCPU(crop_img, posvars=posvars, kwargs=kwargs, \
case='crop cube', cpu_count=cpu_count, verbose=verbose)
def get_piston(cube, verbose=False):
if cube.ndim < 3:
return np.mean(cube[cube != 0])
else:
return np.mean(
multiCPU(get_piston,
posvars=[cube],
case='get piston',
verbose=verbose))
def get_zernike(cube_name, pup, nzer):
zpols_name = cube_name[:-5] + '_zpols_%s.fits'%nzer
try:
zpols = fits.getdata(zpols_name)
print('getdata ' + zpols_name)
except FileNotFoundError:
print('writeto ' + zpols_name)
cube = fits.getdata(cube_name)
nimg = cube.shape[-1]
pup = resize_cube(pup, nimg)
zpols = multiCPU(fit_zer, posargs=[pup, nimg/2, nzer],
posvars=[cube], case='get zpols')
fits.writeto(zpols_name, np.float32(zpols))
return zpols
def resize_cube(cube,
new_size,
preserve_range=True,
mode='reflect',
anti_aliasing=True,
cpu_count=None,
verbose=False):
posvars = [cube, [new_size] * len(cube)]
kwargs = dict(preserve_range=preserve_range,
mode=mode,
anti_aliasing=anti_aliasing)
if cube.ndim < 3:
return resize_img(cube, new_size, **kwargs)
else:
return multiCPU(resize_img, posvars=posvars, kwargs=kwargs, \
case='resize cube', cpu_count=cpu_count, verbose=verbose)
# 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 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
# create cube
cube = multiCPU(remove_piston,
posvars=[filenames],
case='create cube SCAO',
cpu_count=cpu_count)
print(cube.shape)
fits.writeto(os.path.join(output_folder, savename),
np.float32(cube),
overwrite=True)
notify(conf['send_message'], conf['send_to'])
def propagate_cube(wf,
phase_screens,
amp_screens,
tiptilts,
apo_misaligns,
ls_misaligns,
nframes=10,
nstep=1,
mode='RAVC',
ngrid=1024,
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(mode=mode,
ngrid=ngrid,
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)
# keep a copy of the input wavefront
wf1 = deepcopy(wf)
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)
# preload amp screen if only one frame
if len(amp_screens) == 1 and np.any(amp_screens) != None:
proper.prop_multiply(wf1, pad_img(amp_screens, ngrid))
# then create a cube of None values
amp_screens = [None] * int(nframes / nstep + 0.5)
if verbose == True:
print(' preloading amplitude screen')
# preload apodizer when no drift
if ('APP' in mode) or ('RAVC' in mode and np.all(apo_misaligns) == None):
wf1 = apodizer(wf1, verbose=False, **conf)
conf['apo_loaded'] = True
if verbose == True:
print(' preloading %s apodizer' % mode)
else:
conf['apo_loaded'] = False
# preload Lyot stop when no drift
if ('VC' in mode or 'LC' in mode) and np.all(ls_misaligns) == None:
conf['ls_mask'] = lyot_stop(wf1, apply_ls=False, verbose=False, **conf)
if verbose == True:
print(' preloading Lyot stop')
# run simulation
posvars = [
phase_screens, amp_screens, tiptilts, apo_misaligns, ls_misaligns
]
kwargs = dict(verbose=False, **conf)
psfs = multiCPU(propagate_one,
posargs=[wf1],
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,
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 get_rms(cube, verbose=False):
if cube.ndim < 3:
return np.std(cube[cube != 0])
else:
return np.mean(
multiCPU(get_rms, posvars=[cube], case='get rms', verbose=verbose))