def propagate_cube(wf,
phase_screens,
tiptilts,
misaligns,
cpu_count=1,
send_to=None,
tag=None,
onaxis=True,
savefits=False,
verbose=False,
**conf):
# 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, 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, tiptilt, misalign) \
in enumerate(zip(phase_screens, tiptilts, misaligns)):
psf = propagate_one(wf, phase_screen, tiptilt, misalign, \
onaxis=onaxis, verbose=False, **conf)
psfs = psf if i == 0 else np.dstack((psfs.T, psf.T)).T
if verbose is True:
print('%s: finished, elapsed %.3f seconds\n'\
%(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
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 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