diff --git a/heeps/config/read_config.py b/heeps/config/read_config.py index dda535a..dea2519 100644 --- a/heeps/config/read_config.py +++ b/heeps/config/read_config.py @@ -86,7 +86,7 @@ def read_config(verbose=False, **update_conf): ravc_t = 0.7903, # (calc=False) mean-M1 RA trans ravc_r = 0.6033, # (calc=False) mean-M1 RA radius wrt allglass ravc_misalign = [0,0,0,0,0,0], # RA misalignment - clc_diam = 4, # CLC occulter diam in lam/D + clc_diam = 80, # CLC occulter diam in mas file_app_amp = 'apodizers/APP_stop_L_285_v2.fits', # apodizing phase plate files file_app_phase = 'apodizers/vAPP_Dshape_Lband_asymmetric.fits', app_strehl = 0.64, # APP Strehl ratio diff --git a/heeps/optics/fp_mask.py b/heeps/optics/fp_mask.py index f3a3202..e2164af 100644 --- a/heeps/optics/fp_mask.py +++ b/heeps/optics/fp_mask.py @@ -1,5 +1,6 @@ from .lens import lens from .vortex_init import vortex_init +from .lyotmask_init import lyotmask_init import os.path import numpy as np from heeps.util.img_processing import resize_img, pad_img @@ -29,39 +30,12 @@ def fp_mask(wf, mode='RAVC', focal=660, verbose=False, **conf): elif mode in ['CLC']: if verbose is True: print(' apply classical lyot mask') + # load lyotmask amplitude file + conf = lyotmask_init(verbose=verbose, **conf) # propagate to lyot mask - lens(wf, focal) - # create or load the classical Lyot mask - calib = str(conf['clc_diam'])+str('_')+str(int(conf['beam_ratio']*100))+str('_')+str(conf['ngrid']) - my_file = os.path.join(conf['dir_temp'], 'clc_'+calib+'.fits') - if not os.path.isfile(my_file): - # calculate exact size of Lyot mask diameter, in pixels - Dmask = conf['clc_diam']/conf['beam_ratio'] - # oversample the Lyot mask (round up) - samp = 100 - ndisk = int(samp*np.ceil(Dmask)) - ndisk = ndisk if ndisk % 2 else ndisk + 1# 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) - r = abs(x + 1j*y) - # create the Lyot mask - mask = np.zeros((ndisk, ndisk)) - mask[np.where(r > 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, conf['ngrid'], 1) - # write mask - fits.writeto(my_file, mask, overwrite=True) - else: - mask = fits.getdata(my_file) - + lens(wf, focal) # apply lyot mask - mask = proper.prop_shift_center(mask) - wf._wfarr.real *= mask - + wf._wfarr.real *= conf['lyotmask'] # propagate to lyot stop lens(wf, focal) diff --git a/heeps/optics/lyotmask_init.py b/heeps/optics/lyotmask_init.py new file mode 100644 index 0000000..fd39632 --- /dev/null +++ b/heeps/optics/lyotmask_init.py @@ -0,0 +1,67 @@ +from heeps.util.coord import polar_coord +from heeps.util.img_processing import resize_img, pad_img +import numpy as np +import os.path +from astropy.io import fits +import proper + +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 \ No newline at end of file diff --git a/heeps/optics/vortex_init.py b/heeps/optics/vortex_init.py index adddf2e..3582afb 100644 --- a/heeps/optics/vortex_init.py +++ b/heeps/optics/vortex_init.py @@ -25,13 +25,13 @@ def vortex_init(vortex_calib='', dir_temp='', diam_ext=37, lam=3.8, ngrid=1024, [conf.pop(key) for key in ['conf', 'verbose'] if key in conf] # check if back-propagation params already loaded for this calib - calib = '%s_%s_%3.4f'%(vc_charge, ngrid, beam_ratio) + 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, 'calib_vvc_%s.fits'%calib) + filename = os.path.join(dir_temp, '%s.fits'%calib) if os.path.isfile(filename): if verbose is True: print(' loading vortex back-propagation params') diff --git a/heeps/util/save2pkl.py b/heeps/util/save2pkl.py index e537b1e..f6dbab1 100644 --- a/heeps/util/save2pkl.py +++ b/heeps/util/save2pkl.py @@ -5,6 +5,12 @@ def save2pkl(data, name, dir_output='output_files', band='L', mode='RAVC', **con ''' Save data to a pickle file. ''' + # update dict + conf.update(dir_output=dir_output) + ignore = ['perf_num','psf_num','vvc','lyotmask'] + for key in ignore : + if key in conf.keys(): + del conf[key] os.makedirs(dir_output, exist_ok=True) filename = os.path.join(dir_output, '%s_%s_%s.pkl'%(name, band, mode)) with open(filename, 'wb') as file: diff --git a/run_heeps.py b/run_heeps.py index db8c886..fbe3dba 100644 --- a/run_heeps.py +++ b/run_heeps.py @@ -50,7 +50,7 @@ def run_heeps(savename='ADI_contrast_curve.png'): plt.xticks([0.02, 0.1, 0.5, 1]) plt.xlim(0.02,1) plt.ylim(1e-6,1e-2) - plt.savefig(savename, dpi=300, transparent=True) + plt.savefig('%s/%s'%(conf['dir_output'], savename), dpi=300, transparent=True) if __name__ == "__main__": '''