update to lyotmask (adding init function)

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

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)
 

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

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')

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:

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__":
     '''