diff --git a/.gitignore b/.gitignore index 464faa79..f61e66f8 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ *.pyc *.iml +*.sh .ipynb_checkpoints/ config_local.ini .idea/ diff --git a/pastis/config.py b/pastis/config.py index 5dd28476..a37c238c 100644 --- a/pastis/config.py +++ b/pastis/config.py @@ -20,9 +20,12 @@ def get_config_ini_path(): return os.path.join(os.path.dirname(os.path.realpath(__file__)), config_file_name) -def load_config_ini(): +def load_config_ini(dir=None): # Locate where on disk this file is. - code_directory = os.path.dirname(os.path.realpath(__file__)) + if dir == None: + code_directory = os.path.dirname(os.path.realpath(__file__)) + else: + code_directory = dir # Read config file once here. config = configparser.ConfigParser() diff --git a/pastis/config_pastis.ini b/pastis/config_pastis.ini index ceb385ce..4c866879 100644 --- a/pastis/config_pastis.ini +++ b/pastis/config_pastis.ini @@ -6,6 +6,9 @@ ; figure out webbpsf-data path with: webbpsf.utils.get_webbpsf_data_path() webbpsf_data_path = /Users//anaconda/envs/astroconda/share/webbpsf-data local_data_path = /Users//data_from_repos/pastis_data +analysis_name = last +; folder name to analysis or hockeystick curve, or 'last' to run last folder. + [telescope] name = LUVOIR @@ -131,6 +134,7 @@ valid_range_lower = -4 valid_range_upper = 4 ; telescope +design = small nb_subapertures = 120 diameter = 15. gaps = 0.02 @@ -141,6 +145,16 @@ lyot_stop_path_in_optics = inputs/LS_LUVOIR_ID0120_OD0982_no_struts_gy_ovsamp4_N ; absolute path to Harris spreadsheet, currently hosted off the repo harris_data_path = /Users/yourname/somewhere/Sensitivities2.xlsx +therm = False +mech = True +other = False + +; Possible: "seg_mirror", "harris_seg_mirror", "zernike_mirror", "default" +DM = zernike_mirror +; 1 will be set for harris_seg_mirror, "default" limit at 3 +DM_mode = 1 +; 0 to do all mode, accept multiple values but all must be lower than DM_mode +DM_mode_select = 0 ; coronagraph ; iwa and owa from dictionaries within files. could move that to util. @@ -191,6 +205,31 @@ IWA = 2 OWA = 32 lyot_stop_ratio = 0.95 +[generation] +run = True + +method = intensity +;Valid method: intensity, Efield + +[hockeystick] +run = True + +range_points = 30 +no_realizations = 1 + +[analysis] +run = True + +calculate_modes = True +calculate_sigmas = True +run_monte_carlo_modes = True +calc_cumulative_contrast = True +calculate_mus = True +run_monte_carlo_segments = True +calculate_covariance_matrices = True +analytical_statistics = True +calculate_segment_based = True + [numerical] ; size_seg used to be 100 in atlast case, 118 for JWST 512 px images, 239 for JWST 1024 px images size_seg = 239 @@ -206,6 +245,7 @@ z_pup_downsample = 10 ; this is not used automatically in the functions, it is always defined (or read from here) manually current_analysis = 2020-01-13T21-34-29_luvoir-small +; deprecated use [local] analysis_name [zernikes] ; Noll convention! @@ -234,3 +274,10 @@ number_of_low_order_modes = 15 number_of_mid_order_modes = 1 number_of_high_order_modes = 4 number_of_continuous_dm_actuators = 4 + +[save_data] +save_efields = True +save_psfs = True +save_opds = True +save_coro_floor = True +coro_simulator = True diff --git a/pastis/contrast_calculation_simple.py b/pastis/contrast_calculation_simple.py index bacfa199..30c58e7c 100644 --- a/pastis/contrast_calculation_simple.py +++ b/pastis/contrast_calculation_simple.py @@ -20,6 +20,7 @@ from pastis.config import CONFIG_PASTIS from pastis.e2e_simulators.hicat_imaging import set_up_hicat from pastis.e2e_simulators.luvoir_imaging import LuvoirAPLC +from pastis.launchers.parameters import parameters import pastis.e2e_simulators.webbpsf_imaging as webbpsf_imaging import pastis.analytical_pastis.image_pastis as impastis import pastis.util as util @@ -414,7 +415,7 @@ def contrast_rst_num(coro_floor, norm, matrix_dir, rms=50*u.nm): # Create random aberration coefficients on segments, scaled to total rms aber = util.create_random_rms_values(total_seg, rms) - ### E2E JWST sim + ### E2E RST sim start_e2e = time.time() rst_sim = webbpsf_imaging.set_up_cgi() @@ -422,7 +423,6 @@ def contrast_rst_num(coro_floor, norm, matrix_dir, rms=50*u.nm): nb_actu = rst_sim.nbactuator iwa = CONFIG_PASTIS.getfloat('RST', 'IWA') owa = CONFIG_PASTIS.getfloat('RST', 'OWA') - sampling = CONFIG_PASTIS.getfloat('RST', 'sampling') # Put aberration on OTE rst_sim.dm1.flatten() @@ -462,6 +462,67 @@ def contrast_rst_num(coro_floor, norm, matrix_dir, rms=50*u.nm): return contrast_rst, contrast_matrix +def contrast_general_num(matrix_dir, rms=50*u.nm): + """ + Compute the contrast for a random aberration over all DM actuators in the RST simulator. + + :param matrix_dir: str, directory of saved matrix + :param rms: astropy quantity (e.g. m or nm), WFE rms (OPD) to be put randomly over the entire continuous mirror + :return: 2x float, E2E and matrix contrast + """ + # Keep track of time + start_time = time.time() + + # Parameters + param = parameters() + telescope = param.def_telescope() + telescope.normalization_and_dark_hole() + telescope.calculate_unaberrated_contrast() + total_seg = telescope.number_all_modes + + # Import numerical PASTIS matrix + filename = 'pastis_matrix' + matrix_pastis = fits.getdata(os.path.join(matrix_dir, filename + '.fits')) + + # Create random aberration coefficients on segments, scaled to total rms + aber = util.create_random_rms_values(total_seg, rms) + + ### E2E RST sim + start_e2e = time.time() + + # Put aberration on OTE + telescope.flatten() + for nseg in range(total_seg): + telescope.push_mode(nseg, aber[nseg].value*u.nm) + + telescope.imaging_psf() + + # Get the mean contrast + contrast = telescope.contrast() + end_e2e = time.time() + + ## MATRIX PASTIS + log.info('Generating contrast from matrix-PASTIS') + start_matrixpastis = time.time() + # Get mean contrast from matrix PASTIS + contrast_matrix = util.pastis_contrast(aber, matrix_pastis) + telescope.contrast_floor # calculating contrast with PASTIS matrix model + end_matrixpastis = time.time() + + ## Outputs + log.info('\n--- CONTRASTS: ---') + log.info(f'Mean contrast from E2E: {contrast}') + log.info(f'Contrast from matrix PASTIS: {contrast_matrix}') + + log.info('\n--- RUNTIMES: ---') + log.info(f'E2E: {end_e2e-start_e2e}sec = {(end_e2e-start_e2e)/60}min') + log.info(f'Matrix PASTIS: {end_matrixpastis-start_matrixpastis}sec = {(end_matrixpastis-start_matrixpastis)/60}min') + + end_time = time.time() + runtime = end_time - start_time + log.info(f'Runtime for contrast_calculation_simple.py: {runtime} sec = {runtime/60} min') + + return contrast, contrast_matrix + if __name__ == '__main__': # Test JWST diff --git a/pastis/e2e_simulators/telescopes.py b/pastis/e2e_simulators/telescopes.py new file mode 100644 index 00000000..95a0985b --- /dev/null +++ b/pastis/e2e_simulators/telescopes.py @@ -0,0 +1,291 @@ +""" +This module contains class and all necessary requierement to work RST +""" + +import os +import logging +import matplotlib.pyplot as plt +import numpy as np +import hcipy +import re + +from pastis.config import CONFIG_PASTIS +import pastis.util as util +import pastis.launchers.parameters +import pastis.e2e_simulators.webbpsf_imaging as webbpsf_imaging +from pastis.e2e_simulators.luvoir_imaging import LuvoirAPLC + +log = logging.getLogger() + +class RST(): + + def __init__(self, initial_path=''): + self.sim = webbpsf_imaging.set_up_cgi() + self.nb_actu = self.sim.nbactuator + self.norm = 1 + + self.wfe_aber = CONFIG_PASTIS.getfloat('RST', 'calibration_aberration') * 1e-9 # m + self.number_all_modes = CONFIG_PASTIS.getint('RST', 'nb_subapertures') + + def normalization_and_dark_hole(self): + # Calculate direct reference images for contrast normalization + self.flatten() + rst_direct = self.sim.raw_coronagraph() + self.direct_psf = self.imaging_psf(inst=rst_direct) + self.norm = self.direct_psf.max() + + self.iwa = CONFIG_PASTIS.getfloat('RST', 'IWA') + self.owa = CONFIG_PASTIS.getfloat('RST', 'OWA') + self.sim.working_area(im=self.direct_psf, inner_rad=self.iwa, outer_rad=self.owa) + self.dh_mask = self.sim.WA + + def calculate_unaberrated_contrast(self): + '''Calcul an underrated contrast, usually needs one execution of normalization_and_dark_hole''' + self.flatten() + self.coro_simulator = self.sim + self.unaberrated = self.imaging_psf() + + # Calculate coronagraph floor in dark hole + self.contrast_floor = self.contrast() + + def flatten(self): + ''' Flat all actuators''' + self.sim.dm1.flatten() + return [] + + def push_mode(self, mode, amplitude): + actu_x, actu_y = util.seg_to_dm_xy(self.nb_actu, mode) + self.sim.dm1.set_actuator(actu_x, actu_y, amplitude) + + def imaging_psf(self, inst=None): + if inst == None : + inst = self.sim + fit_psf = inst.calc_psf(nlambda=1, fov_arcsec=1.6) + self.psf = fit_psf[0].data/self.norm + return self.psf + + def imaging_efield(self): + _psf , inter = self.sim.calc_psf(nlambda=1, fov_arcsec=1.6, return_intermediates=True) + self.efield = inter[-1].wavefront + return self.efield + + def contrast(self): + self.imaging_psf() + return util.dh_mean(self.psf, self.dh_mask) + + def setup_deformable_mirror(self): + pass + + def display_opd(self): + plt.figure(figsize=(self.nb_actu, self.nb_actu)) + self.sim.dm1.display(what='opd', opd_vmax=self.wfe_aber, colorbar_orientation='horizontal', + title='Aberrated actuator pair') + + + +class LUVOIRA(): + + instrument = 'LUVOIR' + + def __init__(self, initial_path=''): + # General telescope parameters + self.design = CONFIG_PASTIS.get('LUVOIR', 'design') + self.number_all_modes = CONFIG_PASTIS.getint(self.instrument, 'nb_subapertures') + self.seglist = util.get_segment_list(self.instrument) + self.wvln = CONFIG_PASTIS.getfloat(self.instrument, 'lambda') * 1e-9 # m + self.wfe_aber = CONFIG_PASTIS.getfloat(self.instrument, 'calibration_aberration') * 1e-9 # m + self.sampling = CONFIG_PASTIS.getfloat('LUVOIR', 'sampling') + self.optics_input = os.path.join(util.find_repo_location(), + CONFIG_PASTIS.get('LUVOIR', 'optics_path_in_repo')) + + self.dm_modes_list = [0] + self.dm_modes_max = 1 + + self.sim = LuvoirAPLC(self.optics_input, self.design, self.sampling) + self.parameters = pastis.launchers.parameters.parameters() + self.parameters.def_saves() + self.parameters.def_analysis() + + if self.design is None and self.parameters.hockeystick_curve : + raise ValueError( + 'Need to specify apodizer_choice when woant plot hockeystick curve with LUVOIR instrument.') + + def normalization_and_dark_hole(self): + # Calculate direct reference images for contrast normalization + self.flatten() + self.dh_mask = self.sim.dh_mask.shaped + + # Calculate contrast normalization factor from direct PSF (intensity) + psf, direct = self.sim.calc_psf(ref=True) + self.norm = direct.max() + self.direct_psf = psf.shaped + + def calculate_unaberrated_contrast(self): + '''Calcul an underrated contrast, usually needs one execution of normalization_and_dark_hole''' + self.flatten() + self.sim.calc_psf(ref=True) + self.unaberrated = self.imaging_psf() + self.coro_simulator = self.sim + + # Calculate coronagraph floor in dark hole + self.contrast_floor = self.contrast() + + def flatten(self): + ''' Flat all modes''' + self.sim.flatten() + return [] + + def push_mode(self, mode, amplitude): + mode_type = (mode % self.dm_modes_max)+1 + if mode_type in self.dm_modes_list or 0 in self.dm_modes_list: + if self.which_dm == 'default': + if mode_type == 1 : + self.dm_mode.set_segment(mode%self.dm_modes_max+1, amplitude /2, 0, 0) + elif mode_type == 2 : + self.dm_mode.set_segment(mode%self.dm_modes_max+1, 0, amplitude /2, 0) + elif mode_type == 3 : + self.dm_mode.set_segment(mode%self.dm_modes_max+1, 0, 0, amplitude /2) + else: + all_modes = np.zeros(self.number_all_modes) + all_modes[mode] = amplitude / 2 + self.dm_mode.actuators = all_modes + + def imaging_psf(self, inst=None): + + if inst == None : + inst = self.sim + if self.parameters.saveopds: + psf, inter = inst.calc_psf(ref=False, display_intermediate=False, + return_intermediate='intensity') + self.inter = inter['seg_mirror'] + else: + psf = inst.calc_psf(ref=False, display_intermediate=False, + return_intermediate=None) + self.psf = (psf.shaped/self.norm) + return self.psf + + def imaging_efield(self): + efield_focal_plane , inter = self.sim.calc_psf(return_intermediate='efield') + self.efield = efield_focal_plane.electric_field.shaped + self.inter = inter['seg_mirror'].phase + return self.efield + + def contrast(self): + self.imaging_psf() + return util.dh_mean(self.psf, self.dh_mask) + + def setup_deformable_mirror(self): + """ Set up the deformable mirror for the modes you're using and define the total number of mode actuators. """ + + #DM config + self.which_dm = CONFIG_PASTIS.get('LUVOIR', 'DM') + self.dm_modes_max = CONFIG_PASTIS.getint('LUVOIR', 'DM_mode') + + dm_modes_list = CONFIG_PASTIS.get('LUVOIR', 'DM_mode_select') + dm_modes_list = map(int, re.findall(r'\d+', dm_modes_list)) #Find integer + self.dm_modes_list = list(set(dm_modes_list)) #return unique items + + if np.max(self.dm_modes_list) > self.dm_modes_max : + error_msg = "one or more DM_mode_select in LUVOIR section inside config file is/are higher than DM_mode!" + log.error(error_msg) + raise ValueError(error_msg) + + log.info('Setting up deformable mirror...') + if self.which_dm == 'seg_mirror': + log.info(f'Creating segmented mirror with {self.dm_modes_max} local modes on each segment...') + if np.max(self.dm_modes_max) > 3: + error_msg = f"DM_mode={self.dm_modes_max} in LUVOIR section inside config file is higher than 3!" + log.error(error_msg) + raise ValueError(error_msg) + self.sim.create_segmented_mirror(self.dm_modes_max) + self.number_all_modes = self.sim.sm.num_actuators + self.dm_mode = self.sim.sm + elif self.which_dm == 'harris_seg_mirror': #TODO Test it + fpath = CONFIG_PASTIS.get('LUVOIR', 'harris_data_path') # path to Harris spreadsheet + pad_orientations = np.pi / 2 * np.ones(120) + therm = CONFIG_PASTIS.getboolean('LUVOIR', 'therm') + mech = CONFIG_PASTIS.getboolean('LUVOIR', 'mech') + other = CONFIG_PASTIS.getboolean('LUVOIR', 'other') + log.info(f'Reading Harris spreadsheet from {fpath}') + log.info(f'Using pad orientations: {pad_orientations}') + self.sim.create_segmented_harris_mirror(fpath, pad_orientations, therm, mech, other) + self.number_all_modes = self.sim.harris_sm.num_actuators + self.dm_mode = self.sim.harris_sm + self.dm_modes_max = 1 + elif self.which_dm == 'zernike_mirror': + log.info(f'Creating global Zernike mirror with {self.dm_modes_max} global modes...') + self.sim.create_global_zernike_mirror(self.dm_modes_max) + self.number_all_modes = self.sim.zernike_mirror.num_actuators + self.dm_mode = self.sim.zernike_mirror + else: + self.dm_mode = self.sim + if self.which_dm != 'default': + log.info(f'DM with name "{self.which_dm}" not recognized.') + self.which_dm = 'default' + log.info('Default mirror is on') + if np.max(self.dm_modes_max) > 3: + error_msg = f"DM_mode={self.dm_modes_max} in LUVOIR section inside config file is higher than 3!" + log.error(error_msg) + raise ValueError(error_msg) + self.number_all_modes *= self.dm_modes_max + + log.info(f'Total number of modes: {self.number_all_modes}') + + def display_opd(self): + hcipy.imshow_field(self.inter, grid=self.sim.aperture.grid, mask=self.sim.aperture, cmap='RdBu') + +class MODEL(): + '''This is minial template for a new telescope implementation all parameter inside are necessary replace every #message#''' + + def __init__(self, initial_path=''): + self.sim #= MODEL simulation start# + self.norm = 1 + + self.wfe_aber = CONFIG_PASTIS.getfloat('#MODEL#', 'calibration_aberration') * 1e-9 # m + self.number_all_modes = CONFIG_PASTIS.getint('#MODEL#', 'nb_subapertures') + + def normalization_and_dark_hole(self): + '''Calcul an underrated contrast, usually needs one execution of normalization_and_dark_hole''' + + # Calculate direct reference images for contrast normalization + self.flatten() + self.direct_psf + self.norm + + self.dh_mask + + def calculate_unaberrated_contrast(self): + self.flatten() + self.coro_simulator = self.sim + + # Calculate coronagraph floor in dark hole + self.contrast_floor = self.contrast() + + def flatten(self): + ''' Flat all modes''' + ## + return [] + + def push_mode(self, mode, amplitude): + pass + + def imaging_psf(self, inst=None): + + return self.psf + + def imaging_efield(self): + + return self.efield + + def contrast(self): + self.imaging_psf() + return util.dh_mean(self.psf, self.dh_mask) + + def setup_deformable_mirror(self): + pass + + def display_opd(self): + "" + pass + + diff --git a/pastis/hockeystick_contrast_curve.py b/pastis/hockeystick_contrast_curve.py index 5c9d4a99..9720b258 100644 --- a/pastis/hockeystick_contrast_curve.py +++ b/pastis/hockeystick_contrast_curve.py @@ -17,6 +17,7 @@ import matplotlib.pyplot as plt import numpy as np import pandas as pd +import pastis.config from pastis.config import CONFIG_PASTIS import pastis.contrast_calculation_simple as consim @@ -174,16 +175,7 @@ def hockeystick_curve(instrument, apodizer_choice=None, matrixdir='', resultdir= log.info(f"Random realization: {j+1}/{no_realizations}") log.info(f"Total: {(i*no_realizations)+(j+1)}/{len(rms_range)*no_realizations}") - # Chose correct contrast propagation function for chose instrument - if instrument == 'LUVOIR': - c_e2e, c_matrix = consim.contrast_luvoir_num(contrast_floor, norm, apodizer_choice, matrix_dir=matrixdir, rms=rms) - if instrument == 'HiCAT': - c_e2e, c_matrix = consim.contrast_hicat_num(contrast_floor, norm, matrix_dir=matrixdir, rms=rms) - if instrument == 'JWST': - c_e2e, c_matrix = consim.contrast_jwst_num(contrast_floor, norm, matrix_dir=matrixdir, rms=rms) - if instrument == 'RST': - c_e2e, c_matrix = consim.contrast_rst_num(contrast_floor, norm, matrix_dir=matrixdir, rms=rms) - + c_e2e, c_matrix = consim.contrast_general_num(matrix_dir=matrixdir, rms=rms) e2e_rand.append(c_e2e) matrix_rand.append(c_matrix) @@ -208,6 +200,84 @@ def hockeystick_curve(instrument, apodizer_choice=None, matrixdir='', resultdir= log.info(f'\nTotal runtime for pastis_vs_e2e_contrast_calc.py: {runtime} sec = {runtime/60} min') +def hockeystick_curve_class(dir_run=None): + """ + Construct a PASTIS hockeystick contrast curve for validation of the PASTIS matrix, for class telescope instrument. + + The aberration range is a fixed parameter in the function body since it depends on the coronagraph (and telescope) + used. We define how many realizations of a specific WFE rms error we want to run through, and also how many points we + want to fill the aberration range with. At each point we calculate the contrast for all realizations and plot the + mean of this set of results in a figure that shows contrast vs. WFE rms error. + + :param dir_run: string, indicate direction of matrix folder + :return: + """ + + + # d=paths + resultdir = os.path.join(dir_run, 'results') + matrixdir = os.path.join(dir_run, 'matrix_numerical') + config_file = pastis.config.load_config_ini(matrixdir) + instrument = config_file.get('telescope', 'name') + range_points = config_file.getint('hockeystick', 'range_points') + no_realizations = config_file.getint('hockeystick', 'no_realizations') + + + # Keep track of time + start_time = time.time() + + # Create range of WFE RMS values to test + rms_range = np.logspace(config_file.getfloat(instrument, 'valid_range_lower'), + config_file.getfloat(instrument, 'valid_range_upper'), + range_points) + + # Create results directory if it doesn't exist yet + os.makedirs(resultdir, exist_ok=True) + + # Loop over different RMS values and calculate contrast with MATRIX PASTIS and E2E simulation + e2e_contrasts = [] # contrasts from E2E sim + matrix_contrasts = [] # contrasts from matrix PASTIS + + log.info("WFE RMS range: {} nm".format(rms_range, fmt="%e")) + log.info(f"Random realizations: {no_realizations}") + + for i, rms in enumerate(rms_range): + + rms *= u.nm # Making sure this has the correct units + + e2e_rand = [] + matrix_rand = [] + + for j in range(no_realizations): + log.info("CALCULATING CONTRAST FOR {:.4f}".format(rms)) + log.info(f"WFE RMS number {i + 1}/{len(rms_range)}") + log.info(f"Random realization: {j+1}/{no_realizations}") + log.info(f"Total: {(i*no_realizations)+(j+1)}/{len(rms_range)*no_realizations}") + + c_e2e, c_matrix = consim.contrast_general_num(matrix_dir=matrixdir, rms=rms) + e2e_rand.append(c_e2e) + matrix_rand.append(c_matrix) + + e2e_contrasts.append(np.mean(e2e_rand)) + matrix_contrasts.append(np.mean(matrix_rand)) + + # Save contrasts and rms range + np.savetxt(os.path.join(resultdir, 'hockey_rms_range.txt'), rms_range) + np.savetxt(os.path.join(resultdir, 'hockey_e2e_contrasts.txt'), e2e_contrasts) + np.savetxt(os.path.join(resultdir, 'hockey_matrix_contrasts.txt'), matrix_contrasts) + + # Plot + plt.clf() + ppl.plot_hockey_stick_curve(rms_range, matrix_contrasts, e2e_contrasts, + wvln=config_file.getfloat(instrument, 'lambda'), + out_dir=resultdir, + fname_suffix=f'{no_realizations}_realizations_each', + save=True) + + end_time = time.time() + runtime = end_time - start_time + log.info(f'\nTotal runtime for pastis_vs_e2e_contrast_calc.py: {runtime} sec = {runtime/60} min') + if __name__ == '__main__': # Pick one to run diff --git a/pastis/launchers/parameters.py b/pastis/launchers/parameters.py new file mode 100644 index 00000000..03777ba3 --- /dev/null +++ b/pastis/launchers/parameters.py @@ -0,0 +1,95 @@ +''' +Create class to help launcher and keep some parameter in attributs +''' + +import logging + +import pastis.e2e_simulators.telescopes +from pastis.matrix_generation.matrix_building_numerical import MatrixIntensity +from pastis.matrix_generation.matrix_from_efields import MatrixEfield +import pastis.util as util +from pastis.config import CONFIG_PASTIS + +log = logging.getLogger() +init_path = CONFIG_PASTIS.get('local', 'local_data_path') + +def gen_method(dir=init_path): + ''' + Select the right matrix generator + ''' + method = CONFIG_PASTIS.get('generation', 'method') + param = parameters() + param.def_saves() + param.log_off() + + if method == 'intensity': + gen = MatrixIntensity(initial_path=dir, param=param) + elif method == 'Efield': + gen = MatrixEfield(initial_path=dir, param=param) + else: + error_msg = f"{method} inside config.ini file is not a valid generation method!" \ + f"excepted intensity or Efield" + log.error(error_msg) + raise ValueError(error_msg) + log.info(f'Start matrix generation with {method} method') + return gen + +def dir_analysis(dir=init_path): + dir_name = CONFIG_PASTIS.get('local', 'analysis_name') + if dir_name == 'last' : + dir_ana = util.last_folder() + else : + dir_ana = dir + dir_name + return dir_ana + + +class parameters(): + + def __int__(self, initial_path=''): + + super().__init__() + + def def_telescope(self): + self.instrument = CONFIG_PASTIS.get('telescope', 'name') + if self.instrument == 'RST': + self.telescope = pastis.e2e_simulators.telescopes.RST() + elif self.instrument == 'JWST': + self.telescope = pastis.e2e_simulators.telescopes.JWST() + elif self.instrument == 'LUVOIR': + self.telescope = pastis.e2e_simulators.telescopes.LUVOIRA() + else: + error_msg = f"{self.instrument} inside config.ini file is not a valid telescope name!" \ + f"excepted JWST, RST, ATLAST, HiCAT or LUVOIR" + log.error(error_msg) + raise ValueError(error_msg) + + return self.telescope + + def def_saves(self): + self.savepsfs = CONFIG_PASTIS.getboolean('save_data', 'save_psfs') + self.saveopds = CONFIG_PASTIS.getboolean('save_data', 'save_opds') + self.saveefields = CONFIG_PASTIS.getboolean('save_data', 'save_efields') + self.save_coro_floor = CONFIG_PASTIS.getboolean('save_data', 'save_coro_floor') + self.return_coro_simulator = CONFIG_PASTIS.getboolean('save_data', 'coro_simulator') + + def def_analysis(self): + self.calculate_modes = CONFIG_PASTIS.getboolean('analysis','calculate_modes') + self.calculate_sigmas = CONFIG_PASTIS.getboolean('analysis','calculate_sigmas') + self.run_monte_carlo_modes = CONFIG_PASTIS.getboolean('analysis','run_monte_carlo_modes') + self.calc_cumulative_contrast = CONFIG_PASTIS.getboolean('analysis','calc_cumulative_contrast') + self.calculate_mus = CONFIG_PASTIS.getboolean('analysis','calculate_mus') + self.run_monte_carlo_segments = CONFIG_PASTIS.getboolean('analysis','run_monte_carlo_segments') + self.calculate_covariance_matrices = CONFIG_PASTIS.getboolean('analysis','calculate_covariance_matrices') + self.analytical_statistics = CONFIG_PASTIS.getboolean('analysis','analytical_statistics') + self.calculate_segment_based = CONFIG_PASTIS.getboolean('analysis','calculate_segment_based') + + def log_off(self): + mplfm_logger = logging.getLogger('matplotlib.font_manager') + mplcb_logger = logging.getLogger('matplotlib.colorbar') + mplt_logger = logging.getLogger('matplotlib.ticker') + mplbe_logger = logging.getLogger('matplotlib.backends') + + mplfm_logger.setLevel(logging.WARNING) + mplcb_logger.setLevel(logging.WARNING) + mplt_logger.setLevel(logging.WARNING) + mplbe_logger.setLevel(logging.WARNING) diff --git a/pastis/launchers/run.py b/pastis/launchers/run.py new file mode 100644 index 00000000..4955eea4 --- /dev/null +++ b/pastis/launchers/run.py @@ -0,0 +1,37 @@ +""" +Launcher script to start a full RST run: generate matrix and run full PASTIS analysis. +""" +import os + +from pastis.config import CONFIG_PASTIS +from pastis.hockeystick_contrast_curve import hockeystick_curve_class +from pastis.pastis_analysis import run_full_pastis_analysis +import pastis.util as util +import pastis.launchers.parameters as parameters + + +if __name__ == '__main__': + run_gen = CONFIG_PASTIS.getboolean('generation','run') + run_hockey = CONFIG_PASTIS.getboolean('hockeystick','run') + run_analysis = CONFIG_PASTIS.getboolean('analysis','run') + + #Generate PASTIS matrix + if run_gen : + run_matrix = parameters.gen_method() + run_matrix.calc() + + # Set up loggers for data analysis + if run_hockey or run_analysis : + dir_run = parameters.dir_analysis() + util.setup_pastis_logging(dir_run, 'pastis_analysis') + result_dir = os.path.join(dir_run, 'results') + matrix_dir = os.path.join(dir_run, 'matrix_numerical') + + # Then generate hockey stick curve + if run_hockey : + hockeystick_curve_class(dir_run) + + #In development... + # Finally run the analysis + #if run_analysis : + #run_full_pastis_analysis(instrument='RST', run_choice=dir_run, c_target=1e-8) diff --git a/pastis/launchers/run_rst.py b/pastis/launchers/run_rst.py index 823f2487..2062dd75 100644 --- a/pastis/launchers/run_rst.py +++ b/pastis/launchers/run_rst.py @@ -14,10 +14,10 @@ if __name__ == '__main__': # Generate intensity matrix - #run_matrix = MatrixIntensityRST(initial_path=CONFIG_PASTIS.get('local', 'local_data_path')) + run_matrix = MatrixIntensityRST(initial_path=CONFIG_PASTIS.get('local', 'local_data_path')) # Generate E_field matrix - run_matrix = MatrixEfieldRST(initial_path=CONFIG_PASTIS.get('local', 'local_data_path')) + #run_matrix = MatrixEfieldRST(initial_path=CONFIG_PASTIS.get('local', 'local_data_path')) run_matrix.calc() dir_run = run_matrix.overall_dir diff --git a/pastis/matrix_generation/matrix_building_numerical.py b/pastis/matrix_generation/matrix_building_numerical.py index f24cc57b..b02b87f6 100644 --- a/pastis/matrix_generation/matrix_building_numerical.py +++ b/pastis/matrix_generation/matrix_building_numerical.py @@ -4,6 +4,7 @@ Currently supported: JWST + RST LUVOIR HiCAT """ @@ -41,7 +42,7 @@ class PastisMatrix(ABC): instrument = None - def __init__(self, design=None, initial_path=''): + def __init__(self, design=None, initial_path='', param=None): """ Parameters: ---------- @@ -79,7 +80,7 @@ def __init__(self, design=None, initial_path=''): # Record some of the defined parameters log.info(f'Instrument: {tel_suffix}') log.info(f'Wavelength: {self.wvln} m') - log.info(f'Number of segments: {self.nb_seg}') + log.info(f'Number of modes: {self.nb_seg}') log.info(f'Segment list: {self.seglist}') log.info(f'wfe_aber: {self.wfe_aber} m') @@ -98,7 +99,7 @@ class PastisMatrixIntensities(PastisMatrix): """ instrument = None - def __init__(self, design=None, initial_path='', savepsfs=True, saveopds=True): + def __init__(self, design=None, initial_path='', param=None): """ Parameters: ---------- @@ -111,10 +112,9 @@ def __init__(self, design=None, initial_path='', savepsfs=True, saveopds=True): saveopds: bool Whether to save images of pair-wise aberrated pupils to disk or not """ - super().__init__(design=design, initial_path=initial_path) + self.param = param + super().__init__(design=design, initial_path=initial_path, param=param) - self.savepsfs = savepsfs - self.saveopds = saveopds self.calculate_matrix_pair = None os.makedirs(os.path.join(self.resDir, 'psfs'), exist_ok=True) @@ -126,8 +126,13 @@ def calc(self): """ Main method that calculates the PASTIS matrix """ start_time = time.time() + self.telescope = self.param.def_telescope() + # Calculate coronagraph floor, and normalization factor from direct image self.calculate_ref_image() + self.setup_deformable_mirror() + if self.general: + self.nb_seg = self.telescope.number_all_modes self.setup_one_pair_function() self.calculate_contrast_matrix() self.calculate_pastis_from_contrast_matrix() @@ -195,7 +200,7 @@ def calculate_pastis_from_contrast_matrix(self): """ Take the contrast matrix and calculate the PASTIS matrix from it. """ # Calculate the PASTIS matrix from the contrast matrix: analytical matrix element calculation and normalization - self.matrix_pastis = pastis_from_contrast_matrix(self.contrast_matrix, self.seglist, self.wfe_aber, float(self.contrast_floor)) + self.matrix_pastis = pastis_from_contrast_matrix(self.contrast_matrix, self.seglist, self.wfe_aber, float(self.telescope.contrast_floor)) # Save matrix to file filename_matrix = f'pastis_matrix' @@ -207,6 +212,11 @@ def calculate_pastis_from_contrast_matrix(self): def calculate_ref_image(self): """ Create the attributes self.norm, self.contrast_floor and self.coro_simulator. """ + @abstractmethod + def setup_deformable_mirror(self): + """ Set up the deformable mirror for the modes you're using, if necessary, and define the total number of mode actuators. """ + pass + @abstractmethod def setup_one_pair_function(self): """ Create an attribute that is the partial function that can calculate the contrast from one aberrated @@ -297,7 +307,7 @@ def calculate_unaberrated_contrast_and_normalization(instrument, design=None, re rst_sim = webbpsf_imaging.set_up_cgi() # Calculate direct reference images for contrast normalization - rst_direct = rst_sim.raw_PSF() + rst_direct = rst_sim.raw_coronagraph() direct = rst_direct.calc_psf(nlambda=1, fov_arcsec=1.6) direct_psf = direct[0].data norm = direct_psf.max() @@ -538,7 +548,7 @@ def _rst_matrix_one_pair(norm, wfe_aber, resDir, savepsfs, saveopds, actuator_pa if saveopds: opd_name = f'opd_actuator_{actuator_pair[0]}-{actuator_pair[1]}' plt.clf() - plt.figure(figsize=(8, 8)) + plt.figure(figsize=(nb_actu, nb_actu)) rst_cgi.dm1.display(what='opd', opd_vmax=wfe_aber, colorbar_orientation='horizontal', title='Aberrated actuator pair') plt.savefig(os.path.join(resDir, 'OTE_images', opd_name + '.pdf')) @@ -552,6 +562,47 @@ def _rst_matrix_one_pair(norm, wfe_aber, resDir, savepsfs, saveopds, actuator_pa return contrast, actuator_pair +def general_matrix_one_pair(telescope, norm, wfe_aber, resDir, savepsfs, saveopds, mode_pair): + """ + Function to calculate general mean contrast of pair in telescope. + :param norm: float, direct PSF normalization factor (peak pixel of direct PSF) + :param wfe_aber: calibration aberration per segment in m + :param resDir: str, directory for matrix calculations + :param savepsfs: bool, if True, all PSFs will be saved to disk individually, as fits files + :param saveopds: bool, if True, all pupil surface maps of aberrated segment pairs will be saved to disk as PDF + :param mode_pair: tuple, pair of actuators to aberrate. If same segment gets passed in both tuple entries, the actuator will be aberrated only once + :return: contrast as float, and segment pair as tuple + """ + # Put aberration on correct segments. If i=j, apply only once! + log.info(f'PAIR: {mode_pair[0]}-{mode_pair[1]}') + + # Put aberration on correct segments. If i=j, apply only once! + telescope.flatten() + telescope.push_mode(mode_pair[0], wfe_aber) + if mode_pair[0] != mode_pair[1]: + telescope.push_mode(mode_pair[1], wfe_aber) + + log.info('Calculating coro image...') + psf = telescope.imaging_psf() + + # Save PSF image to disk + if savepsfs: + filename_psf = f'psf_segment_{mode_pair[0]}-{mode_pair[1]}' + hcipy.write_fits(psf, os.path.join(resDir, 'psfs', filename_psf + '.fits'),psf.shape) + + # Plot deformable mirror WFE and save to disk + if saveopds: + opd_name = f'opd_segnement_{mode_pair[0]}-{mode_pair[1]}' + plt.clf() + telescope.display_opd() + plt.savefig(os.path.join(resDir, 'OTE_images', opd_name + '.pdf')) + + log.info(f'Calculating mean contrast in dark hole {mode_pair[0]}-{mode_pair[1]}') + contrast = telescope.contrast() + + return contrast, mode_pair + + def pastis_from_contrast_matrix(contrast_matrix, seglist, wfe_aber, coro_floor): """ Calculate the final PASTIS matrix from the input contrast matrix (only half filled). @@ -799,6 +850,7 @@ def num_matrix_multiprocess(instrument, design=None, initial_path='', savepsfs=T class MatrixIntensityLuvoirA(PastisMatrixIntensities): instrument = 'LUVOIR' + general = False # temp attribut to not break legacy """ Calculate a PASTIS matrix for LUVOIR-A, using intensity images. """ def __int__(self, design='small', initial_path='', savepsfs=True, saveopds=True): @@ -826,6 +878,7 @@ class MatrixIntensityHicat(PastisMatrixIntensities): """ Calculate a PASTIS matrix for HiCAT, using intensity images. """ def __int__(self, initial_path='', savepsfs=True, saveopds=True): + self.general = False super().__init__(design=None, savepsfs=savepsfs, saveopds=saveopds) def setup_one_pair_function(self): @@ -849,6 +902,7 @@ def calculate_ref_image(self, save_coro_floor=True, save_psfs=True): class MatrixIntensityJWST(PastisMatrixIntensities): instrument = 'JWST' + general = False # temp attribut to not break legacy """ Calculate a PASTIS matrix for JWST, using intensity images. """ def __int__(self, initial_path='', savepsfs=True, saveopds=True): @@ -872,6 +926,7 @@ def calculate_ref_image(self, save_coro_floor=True, save_psfs=True): class MatrixIntensityRST(PastisMatrixIntensities): instrument = 'RST' + general = False # temp attribut to not break legacy """ Calculate a PASTIS matrix for the pupil-plane continuous DM on RST/CGI, using intensity images. """ def __int__(self, initial_path='', savepsfs=True, saveopds=True): @@ -893,7 +948,43 @@ def calculate_ref_image(self, save_coro_floor=False, save_psfs=False): outpath=self.overall_dir) -if __name__ == '__main__': +class MatrixIntensity(PastisMatrixIntensities): + """ Calculate a PASTIS matrix for telescope class, using intensity images. """ + + instrument = CONFIG_PASTIS.get('telescope', 'name') + general = True # temp attribut to not break legacy + + def __int__(self, initial_path='', param=None): + super().__init__(design=None, param=param) + + def setup_one_pair_function(self): + """ Create the partial function that returns the PSF of a single aberrated mode pair. """ + + self.calculate_matrix_pair = functools.partial(general_matrix_one_pair, self.telescope, self.telescope.norm, self.telescope.wfe_aber, self.resDir, + self.param.savepsfs, self.param.saveopds) + + def setup_deformable_mirror(self): + """DM setup if requiered for some telescopes (LUVOIR-A)""" + self.telescope.setup_deformable_mirror() + + def calculate_ref_image(self): + """ Calculate the coronagraph floor, normalization factor from direct image, and get the simulator object. """ + telescope = self.telescope + telescope.normalization_and_dark_hole() + telescope.calculate_unaberrated_contrast() + outpath = self.overall_dir + + log.info(f'contrast floor: {telescope.contrast_floor}') + + if self.param.save_coro_floor: + # Save contrast floor to text file + with open(os.path.join(outpath, 'coronagraph_floor.txt'), 'w') as file: + file.write(f'Coronagraph floor: {telescope.contrast_floor}') + + if self.param.savepsfs: + ppl.plot_direct_coro_dh(telescope.direct_psf, telescope.unaberrated, telescope.dh_mask.astype(bool), outpath) - MatrixIntensityLuvoirA(design='small', initial_path=CONFIG_PASTIS.get('local', 'local_data_path')).calc() - #MatrixIntensityHicat(initial_path=CONFIG_PASTIS.get('local', 'local_data_path')).calc() + if self.param.return_coro_simulator: + return telescope.contrast_floor, telescope.norm, telescope.coro_simulator + else: + return telescope.contrast_floor, telescope.norm diff --git a/pastis/matrix_generation/matrix_from_efields.py b/pastis/matrix_generation/matrix_from_efields.py index 4f42bf62..3c78125a 100644 --- a/pastis/matrix_generation/matrix_from_efields.py +++ b/pastis/matrix_generation/matrix_from_efields.py @@ -24,7 +24,7 @@ class PastisMatrixEfields(PastisMatrix): instrument = None """ Main class for PASTIS matrix calculations from individually 'poked' modes. """ - def __init__(self, design=None, initial_path='', saveefields=True, saveopds=True): + def __init__(self, design=None, initial_path='', param=None): """ Parameters: ---------- @@ -37,10 +37,11 @@ def __init__(self, design=None, initial_path='', saveefields=True, saveopds=True saveopds: bool Whether to save images of pair-wise aberrated pupils to disk or not """ - super().__init__(design=design, initial_path=initial_path) + super().__init__(design=design, initial_path=initial_path , param=param) - self.save_efields = saveefields - self.saveopds = saveopds + self.param = param + self.save_efields = param.saveefields + self.saveopds = self.param.saveopds self.calculate_one_mode = None self.efields_per_mode = [] @@ -65,14 +66,15 @@ def calc(self): def calculate_efields(self): """ Poke each mode individually and calculate the resulting focal plane E-field. """ - for i in range(self.number_all_modes): + for i in range(self.telescope.number_all_modes): self.efields_per_mode.append(self.calculate_one_mode(i)) self.efields_per_mode = np.array(self.efields_per_mode) def calculate_pastis_matrix_from_efields(self): """ Use the individual-mode E-fields to calculate the PASTIS matrix from it. """ - self.matrix_pastis = pastis_matrix_from_efields(self.efields_per_mode, self.efield_ref, self.norm, self.dh_mask, self.wfe_aber) + self.matrix_pastis = pastis_matrix_from_efields(self.efields_per_mode, self.efield_ref, self.telescope.norm, + self.telescope.dh_mask, self.wfe_aber) # Save matrix to file filename_matrix = f'pastis_matrix' @@ -233,7 +235,7 @@ def calculate_ref_efield(self): self.rst_cgi = webbpsf_imaging.set_up_cgi() # Calculate direct reference images for contrast normalization - rst_direct = self.rst_cgi.raw_PSF() + rst_direct = self.rst_cgi.raw_coronagraph() direct = rst_direct.calc_psf(nlambda=1, fov_arcsec=1.6) direct_psf = direct[0].data self.norm = direct_psf.max() @@ -244,7 +246,7 @@ def calculate_ref_efield(self): # Calculate reference E-field in focal plane, without any aberrations applied _trash, inter = self.rst_cgi.calc_psf(nlambda=1, fov_arcsec=1.6, return_intermediates=True) - self.efield_ref = inter[6].wavefront # [6] is the last optic = detector + self.efield_ref = inter[-1].wavefront # [-1] is the last optic = detector def setup_deformable_mirror(self): """DM setup not needed for RST, just define number of total modes""" @@ -255,6 +257,35 @@ def setup_single_mode_function(self): self.rst_cgi, self.resDir, self.save_efields, self.saveopds) +class MatrixEfield(PastisMatrixEfields): + """ + Class to calculate the PASTIS matrix from E-fields. + """ + instrument = CONFIG_PASTIS.get('telescope', 'name') + + def __int__(self, initial_path='', param=None): + super().__init__(design=None, param=param) + + def calculate_ref_efield(self): + # Calculate direct reference images for contrast normalization + self.telescope = self.param.def_telescope() + self.telescope.normalization_and_dark_hole() + self.telescope.calculate_unaberrated_contrast() + + log.info(f'contrast floor: {self.telescope.contrast_floor}') + + # Calculate reference E-field in focal plane, without any aberrations applied + self.efield_ref = self.telescope.imaging_efield() + + def setup_deformable_mirror(self): + """DM setup not needed for RST, just define number of total modes""" + self.telescope.setup_deformable_mirror() + + def setup_single_mode_function(self): + self.calculate_one_mode = functools.partial(general_matrix_single_mode, self.telescope, self.wfe_aber, + self.resDir, self.save_efields, self.saveopds) + + def _luvoir_matrix_single_mode(which_dm, number_all_modes, wfe_aber, luvoir_sim, resDir, saveefields, saveopds, mode_no): """ Calculate the LUVOIR-A mean E-field of one aberrated mode; for PastisMatrixEfields(). @@ -337,9 +368,46 @@ def _rst_matrix_single_mode(wfe_aber, rst_sim, resDir, saveefields, saveopds, mo if saveopds: opd_name = f'opd_actuator_{mode_no}' plt.clf() - plt.figure(figsize=(8, 8)) + plt.figure(figsize=(nb_actu, nb_actu)) rst_sim.dm1.display(what='opd', opd_vmax=wfe_aber, colorbar_orientation='horizontal', title='Aberrated actuator pair') plt.savefig(os.path.join(resDir, 'OTE_images', opd_name + '.pdf')) return efield_focal_plane.wavefront + + +def general_matrix_single_mode(telescope, wfe_aber, resDir, saveefields, saveopds, mode_no): + """ + Function to calculate Electrical field (E_field) of one mode. + :param telescope: + :param wfe_aber: float, calibration aberration per actuator in m + :param resDir: str, directory for matrix calculations + :param saveefields: bool, if True, all E_field will be saved to disk individually, as fits files + :param savepods: bool, if True, all pupil surface maps of aberrated actuators pairs will be saved to disk as PDF + :param mode_no: int, which aberrated actuator to calculate the E-field for + """ + + log.info(f'MODE NUMBER: {mode_no}') + + # Apply calibration aberration to used mode + telescope.flatten() + telescope.push_mode(mode_no, wfe_aber) + + # Calculate coronagraphic E-field + efield_focal_plane = telescope.imaging_efield() + + # Save E field image to disk + if saveefields: + fname_real = f'efield_real_mode{mode_no}' + hcipy.write_fits(efield_focal_plane.real, os.path.join(resDir, 'efields', fname_real + '.fits'), efield_focal_plane.shape) + fname_imag = f'efield_imag_mode{mode_no}' + hcipy.write_fits(efield_focal_plane.imag, os.path.join(resDir, 'efields', fname_imag + '.fits'), efield_focal_plane.shape) + + # Plot deformable mirror WFE and save to disk + if saveopds: + opd_name = f'opd_actuator_{mode_no}' + plt.clf() + telescope.display_opd() + plt.savefig(os.path.join(resDir, 'OTE_images', opd_name + '.pdf')) + + return efield_focal_plane \ No newline at end of file diff --git a/pastis/util.py b/pastis/util.py index 756d2bd8..43fbc5ab 100644 --- a/pastis/util.py +++ b/pastis/util.py @@ -688,3 +688,11 @@ def seg_to_dm_xy(actuator_total, segment): actuator_pair_y = (segment-actuator_pair_x)/actuator_total return actuator_pair_x, int(actuator_pair_y) + + +def last_folder(): + local_path = CONFIG_PASTIS.get('local', 'local_data_path') + list_of_folder = glob.glob(os.path.join(local_path, '*/')) + latest_folder = max(list_of_folder, key=os.path.getctime) + + return latest_folder