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