def ana_matrix_jwst():
# Keep track of time
start_time = time.time() # runtime is currently around 11 minutes
log.info('Building analytical matrix for JWST\n')
# Parameters
datadir = os.path.join(CONFIG_PASTIS.get('local', 'local_data_path'),
'active')
which_tel = CONFIG_PASTIS.get('telescope', 'name')
resDir = os.path.join(datadir, 'matrix_analytical')
nb_seg = CONFIG_PASTIS.getint(which_tel, 'nb_subapertures')
nm_aber = CONFIG_PASTIS.getfloat(which_tel,
'calibration_aberration') * u.nm
zern_number = CONFIG_PASTIS.getint('calibration',
'local_zernike') # Noll convention!
zern_mode = util.ZernikeMode(
zern_number) # Create Zernike mode object for easier handling
# If subfolder "matrix_analytical" doesn't exist yet, create it.
if not os.path.isdir(resDir):
os.mkdir(resDir)
#-# Generating the PASTIS matrix
matrix_direct = np.zeros(
[nb_seg,
nb_seg]) # Generate empty matrix for contrast values from loop.
all_ims = []
all_dhs = []
all_contrasts = []
for i in range(nb_seg):
for j in range(nb_seg):
log.info('STEP: {}-{} / {}-{}'.format(i + 1, j + 1, nb_seg,
nb_seg))
# Putting aberration only on segments i and j
tempA = np.zeros([nb_seg])
tempA[i] = nm_aber.value
tempA[j] = nm_aber.value
tempA *= u.nm # making sure this array has the right units
# Create PASTIS image and save full image as well as DH image
temp_im_am, full_psf = impastis.analytical_model(zern_number,
tempA,
cali=True)
filename_psf = 'psf_' + zern_mode.name + '_' + zern_mode.convention + str(
zern_mode.index) + '_segs_' + str(i + 1) + '-' + str(j + 1)
util.write_fits(full_psf,
os.path.join(resDir, 'psfs',
filename_psf + '.fits'),
header=None,
metadata=None)
all_ims.append(full_psf)
filename_dh = 'dh_' + zern_mode.name + '_' + zern_mode.convention + str(
zern_mode.index) + '_segs_' + str(i + 1) + '-' + str(j + 1)
util.write_fits(temp_im_am,
os.path.join(resDir, 'darkholes',
filename_dh + '.fits'),
header=None,
metadata=None)
all_dhs.append(temp_im_am)
contrast = np.mean(temp_im_am[np.where(temp_im_am != 0)])
matrix_direct[i, j] = contrast
log.info(f'contrast = {contrast}')
all_contrasts.append(contrast)
all_ims = np.array(all_ims)
all_dhs = np.array(all_dhs)
all_contrasts = np.array(all_contrasts)
# Filling the off-axis elements
matrix_two_N = np.copy(
matrix_direct
) # This is just an intermediary copy so that I don't mix things up.
matrix_pastis = np.copy(
matrix_direct) # This will be the final PASTIS matrix.
for i in range(nb_seg):
for j in range(nb_seg):
if i != j:
matrix_off_val = (matrix_two_N[i, j] - matrix_two_N[i, i] -
matrix_two_N[j, j]) / 2.
matrix_pastis[i, j] = matrix_off_val
log.info('Off-axis for i{}-j{}: {}'.format(
i + 1, j + 1, matrix_off_val))
# Normalize matrix for the input aberration
matrix_pastis /= np.square(nm_aber.value)
# Save matrix to file
filename = 'PASTISmatrix_' + zern_mode.name + '_' + zern_mode.convention + str(
zern_mode.index)
util.write_fits(matrix_pastis,
os.path.join(resDir, filename + '.fits'),
header=None,
metadata=None)
log.info(f'Matrix saved to: {os.path.join(resDir, filename + ".fits")}')
# Save the PSF and DH image *cubes* as well (as opposed to each one individually)
util.write_fits(all_ims,
os.path.join(resDir, 'psfs', 'psf_cube' + '.fits'),
header=None,
metadata=None)
util.write_fits(all_dhs,
os.path.join(resDir, 'darkholes', 'dh_cube' + '.fits'),
header=None,
metadata=None)
np.savetxt(os.path.join(resDir, 'pair-wise_contrasts.txt'),
all_contrasts,
fmt='%e')
# Tell us how long it took to finish.
end_time = time.time()
log.info(
f'Runtime for matrix_building.py: {end_time - start_time}sec = {(end_time - start_time) / 60}min'
)
log.info('Data saved to {}'.format(resDir))
def num_matrix_multiprocess(instrument, design=None, savepsfs=True, saveopds=True):
"""
Generate a numerical/semi-analytical PASTIS matrix.
Multiprocessed script to calculate PASTIS matrix. Implementation adapted from
hicat.scripts.stroke_minimization.calculate_jacobian
:param instrument: str, what instrument (LUVOIR, HiCAT, JWST) to generate the PASTIS matrix for
:param design: str, optional, default=None, which means we read from the configfile: what coronagraph design
to use - 'small', 'medium' or 'large'
: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
:return: overall_dir: string, experiment directory
"""
# Keep track of time
start_time = time.time() # runtime is currently around 150 minutes
### Parameters
# Create directory names
tel_suffix = f'{instrument.lower()}'
if instrument == 'LUVOIR':
if design is None:
design = CONFIG_PASTIS.get('LUVOIR', 'coronagraph_design')
tel_suffix += f'-{design}'
overall_dir = util.create_data_path(CONFIG_PASTIS.get('local', 'local_data_path'), telescope=tel_suffix)
os.makedirs(overall_dir, exist_ok=True)
resDir = os.path.join(overall_dir, 'matrix_numerical')
# Create necessary directories if they don't exist yet
os.makedirs(resDir, exist_ok=True)
os.makedirs(os.path.join(resDir, 'OTE_images'), exist_ok=True)
os.makedirs(os.path.join(resDir, 'psfs'), exist_ok=True)
# Set up logger
util.setup_pastis_logging(resDir, f'pastis_matrix_{tel_suffix}')
log.info(f'Building numerical matrix for {tel_suffix}\n')
# Read calibration aberration
zern_number = CONFIG_PASTIS.getint('calibration', 'local_zernike')
zern_mode = util.ZernikeMode(zern_number) # Create Zernike mode object for easier handling
# General telescope parameters
nb_seg = CONFIG_PASTIS.getint(instrument, 'nb_subapertures')
seglist = util.get_segment_list(instrument)
wvln = CONFIG_PASTIS.getfloat(instrument, 'lambda') * 1e-9 # m
wfe_aber = CONFIG_PASTIS.getfloat(instrument, 'calibration_aberration') * 1e-9 # m
# Record some of the defined parameters
log.info(f'Instrument: {tel_suffix}')
log.info(f'Wavelength: {wvln} m')
log.info(f'Number of segments: {nb_seg}')
log.info(f'Segment list: {seglist}')
log.info(f'wfe_aber: {wfe_aber} m')
log.info(f'Total number of segment pairs in {instrument} pupil: {len(list(util.segment_pairs_all(nb_seg)))}')
log.info(f'Non-repeating pairs in {instrument} pupil calculated here: {len(list(util.segment_pairs_non_repeating(nb_seg)))}')
# Copy configfile to resulting matrix directory
util.copy_config(resDir)
# Calculate coronagraph floor, and normalization factor from direct image
contrast_floor, norm = calculate_unaberrated_contrast_and_normalization(instrument, design, return_coro_simulator=False,
save_coro_floor=True, save_psfs=False, outpath=overall_dir)
# Figure out how many processes is optimal and create a Pool.
# Assume we're the only one on the machine so we can hog all the resources.
# We expect numpy to use multithreaded math via the Intel MKL library, so
# we check how many threads MKL will use, and create enough processes so
# as to use 100% of the CPU cores.
# You might think we should divide number of cores by 2 to get physical cores
# to account for hyperthreading, however empirical testing on telserv3 shows that
# it is slightly more performant on telserv3 to use all logical cores
num_cpu = multiprocessing.cpu_count()
# try:
# import mkl
# num_core_per_process = mkl.get_max_threads()
# except ImportError:
# # typically this is 4, so use that as default
# log.info("Couldn't import MKL; guessing default value of 4 cores per process")
# num_core_per_process = 4
num_core_per_process = 1 # NOTE: this was changed by Scott Will in HiCAT and makes more sense, somehow
num_processes = int(num_cpu // num_core_per_process)
log.info(f"Multiprocess PASTIS matrix for {instrument} will use {num_processes} processes (with {num_core_per_process} threads per process)")
# Set up a function with all arguments fixed except for the last one, which is the segment pair tuple
if instrument == 'LUVOIR':
calculate_matrix_pair = functools.partial(_luvoir_matrix_one_pair, design, norm, wfe_aber, zern_mode, resDir,
savepsfs, saveopds)
if instrument == 'HiCAT':
# Copy used BostonDM maps to matrix folder
shutil.copytree(CONFIG_PASTIS.get('HiCAT', 'dm_maps_path'), os.path.join(resDir, 'hicat_boston_dm_commands'))
calculate_matrix_pair = functools.partial(_hicat_matrix_one_pair, norm, wfe_aber, resDir, savepsfs, saveopds)
if instrument == 'JWST':
calculate_matrix_pair = functools.partial(_jwst_matrix_one_pair, norm, wfe_aber, resDir, savepsfs, saveopds)
# Iterate over all segment pairs via a multiprocess pool
mypool = multiprocessing.Pool(num_processes)
t_start = time.time()
results = mypool.map(calculate_matrix_pair, util.segment_pairs_non_repeating(nb_seg)) # this util function returns a generator
t_stop = time.time()
log.info(f"Multiprocess calculation complete in {t_stop-t_start}sec = {(t_stop-t_start)/60}min")
# Unscramble results
# results is a list of tuples that contain the return from the partial function, in this case: result[i] = (c, (seg1, seg2))
contrast_matrix = np.zeros([nb_seg, nb_seg]) # Generate empty matrix
for i in range(len(results)):
# Fill according entry in the matrix and subtract baseline contrast
contrast_matrix[results[i][1][0], results[i][1][1]] = results[i][0] - contrast_floor
mypool.close()
# Save all contrasts to disk, WITH subtraction of coronagraph floor
hcipy.write_fits(contrast_matrix, os.path.join(resDir, 'pair-wise_contrasts.fits'))
plt.figure(figsize=(10, 10))
plt.imshow(contrast_matrix)
plt.colorbar()
plt.savefig(os.path.join(resDir, 'contrast_matrix.pdf'))
# Calculate the PASTIS matrix from the contrast matrix: off-axis elements and normalization
matrix_pastis = pastis_from_contrast_matrix(contrast_matrix, seglist, wfe_aber)
# Save matrix to file
filename_matrix = f'PASTISmatrix_num_{zern_mode.name}_{zern_mode.convention + str(zern_mode.index)}'
hcipy.write_fits(matrix_pastis, os.path.join(resDir, filename_matrix + '.fits'))
ppl.plot_pastis_matrix(matrix_pastis, wvln*1e9, out_dir=resDir, save=True) # convert wavelength to nm
log.info(f'Matrix saved to: {os.path.join(resDir, filename_matrix + ".fits")}')
# Tell us how long it took to finish.
end_time = time.time()
log.info(f'Runtime for matrix_building_numerical.py/multiprocess: {end_time - start_time}sec = {(end_time - start_time)/60}min')
log.info(f'Data saved to {resDir}')
return overall_dir
def num_matrix_jwst():
"""
Generate a numerical PASTIS matrix for a JWST coronagraph.
-- Depracated function, the LUVOIR PASTIS matrix is better calculated with num_matrix_multiprocess(), which can
do this for your choice of one of the implemented instruments (LUVOIR, HiCAT, JWST). --
All inputs are read from the (local) configfile and saved to the specified output directory.
"""
import webbpsf
from e2e_simulators import webbpsf_imaging as webbim
# Set WebbPSF environment variable
os.environ['WEBBPSF_PATH'] = CONFIG_PASTIS.get('local', 'webbpsf_data_path')
# Keep track of time
start_time = time.time() # runtime is currently around 21 minutes
log.info('Building numerical matrix for JWST\n')
# Parameters
overall_dir = util.create_data_path(CONFIG_PASTIS.get('local', 'local_data_path'), telescope='jwst')
resDir = os.path.join(overall_dir, 'matrix_numerical')
which_tel = CONFIG_PASTIS.get('telescope', 'name')
nb_seg = CONFIG_PASTIS.getint(which_tel, 'nb_subapertures')
im_size_e2e = CONFIG_PASTIS.getint('numerical', 'im_size_px_webbpsf')
inner_wa = CONFIG_PASTIS.getint(which_tel, 'IWA')
outer_wa = CONFIG_PASTIS.getint(which_tel, 'OWA')
sampling = CONFIG_PASTIS.getfloat(which_tel, 'sampling')
fpm = CONFIG_PASTIS.get(which_tel, 'focal_plane_mask') # focal plane mask
lyot_stop = CONFIG_PASTIS.get(which_tel, 'pupil_plane_stop') # Lyot stop
filter = CONFIG_PASTIS.get(which_tel, 'filter_name')
wfe_aber = CONFIG_PASTIS.getfloat(which_tel, 'calibration_aberration') * u.nm
wss_segs = webbpsf.constants.SEGNAMES_WSS_ORDER
zern_max = CONFIG_PASTIS.getint('zernikes', 'max_zern')
zern_number = CONFIG_PASTIS.getint('calibration', 'local_zernike')
zern_mode = util.ZernikeMode(zern_number) # Create Zernike mode object for easier handling
wss_zern_nb = util.noll_to_wss(zern_number) # Convert from Noll to WSS framework
# Create necessary directories if they don't exist yet
os.makedirs(overall_dir, exist_ok=True)
os.makedirs(resDir, exist_ok=True)
os.makedirs(os.path.join(resDir, 'OTE_images'), exist_ok=True)
os.makedirs(os.path.join(resDir, 'psfs'), exist_ok=True)
os.makedirs(os.path.join(resDir, 'darkholes'), exist_ok=True)
# Create the dark hole mask.
pup_im = np.zeros([im_size_e2e, im_size_e2e]) # this is just used for DH mask generation
dh_area = util.create_dark_hole(pup_im, inner_wa, outer_wa, sampling)
# Create a direct WebbPSF image for normalization factor
fake_aber = np.zeros([nb_seg, zern_max])
psf_perfect = webbim.nircam_nocoro(filter, fake_aber)
normp = np.max(psf_perfect)
psf_perfect = psf_perfect / normp
# Set up NIRCam coro object from WebbPSF
nc_coro = webbpsf.NIRCam()
nc_coro.filter = filter
nc_coro.image_mask = fpm
nc_coro.pupil_mask = lyot_stop
# Null the OTE OPDs for the PSFs, maybe we will add internal WFE later.
nc_coro, ote_coro = webbpsf.enable_adjustable_ote(nc_coro) # create OTE for coronagraph
nc_coro.include_si_wfe = False # set SI internal WFE to zero
#-# Generating the PASTIS matrix and a list for all contrasts
contrast_matrix = np.zeros([nb_seg, nb_seg]) # Generate empty matrix
all_psfs = []
all_dhs = []
all_contrasts = []
log.info(f'wfe_aber: {wfe_aber}')
for i in range(nb_seg):
for j in range(nb_seg):
log.info(f'\nSTEP: {i+1}-{j+1} / {nb_seg}-{nb_seg}')
# Get names of segments, they're being addressed by their names in the ote functions.
seg_i = wss_segs[i].split('-')[0]
seg_j = wss_segs[j].split('-')[0]
# Put the aberration on the correct segments
Aber_WSS = np.zeros([nb_seg, zern_max]) # The Zernikes here will be filled in the WSS order!!!
# Because it goes into _apply_hexikes_to_seg().
Aber_WSS[i, wss_zern_nb - 1] = wfe_aber.to(u.m).value # Aberration on the segment we're currently working on;
# convert to meters; -1 on the Zernike because Python starts
# numbering at 0.
Aber_WSS[j, wss_zern_nb - 1] = wfe_aber.to(u.m).value # same for other segment
# Putting aberrations on segments i and j
ote_coro.reset() # Making sure there are no previous movements on the segments.
ote_coro.zero() # set OTE for coronagraph to zero
# Apply both aberrations to OTE. If i=j, apply only once!
ote_coro._apply_hexikes_to_seg(seg_i, Aber_WSS[i, :]) # set segment i (segment numbering starts at 1)
if i != j:
ote_coro._apply_hexikes_to_seg(seg_j, Aber_WSS[j, :]) # set segment j
# If you want to display it:
# ote_coro.display_opd()
# plt.show()
# Save OPD images for testing
opd_name = f'opd_{zern_mode.name}_{zern_mode.convention + str(zern_mode.index)}_segs_{i+1}-{j+1}'
plt.clf()
ote_coro.display_opd()
plt.savefig(os.path.join(resDir, 'OTE_images', opd_name + '.pdf'))
log.info('Calculating WebbPSF image')
image = nc_coro.calc_psf(fov_pixels=int(im_size_e2e), oversample=1, nlambda=1)
psf = image[0].data / normp
# Save WebbPSF image to disk
filename_psf = f'psf_{zern_mode.name}_{zern_mode.convention + str(zern_mode.index)}_segs_{i+1}-{j+1}'
util.write_fits(psf, os.path.join(resDir, 'psfs', filename_psf + '.fits'), header=None, metadata=None)
all_psfs.append(psf)
log.info('Calculating mean contrast in dark hole')
dh_intensity = psf * dh_area
contrast = np.mean(dh_intensity[np.where(dh_intensity != 0)])
log.info(f'contrast: {contrast}')
# Save DH image to disk and put current contrast in list
filename_dh = f'dh_{zern_mode.name}_{zern_mode.convention + str(zern_mode.index)}_segs_{i+1}-{j+1}'
util.write_fits(dh_intensity, os.path.join(resDir, 'darkholes', filename_dh + '.fits'), header=None, metadata=None)
all_dhs.append(dh_intensity)
all_contrasts.append(contrast)
# Fill according entry in the matrix
contrast_matrix[i,j] = contrast
# Transform saved lists to arrays
all_psfs = np.array(all_psfs)
all_dhs = np.array(all_dhs)
all_contrasts = np.array(all_contrasts)
# Filling the off-axis elements
matrix_two_N = np.copy(contrast_matrix) # This is just an intermediary copy so that I don't mix things up.
matrix_pastis = np.copy(contrast_matrix) # This will be the final PASTIS matrix.
for i in range(nb_seg):
for j in range(nb_seg):
if i != j:
matrix_off_val = (matrix_two_N[i,j] - matrix_two_N[i,i] - matrix_two_N[j,j]) / 2.
matrix_pastis[i,j] = matrix_off_val
log.info(f'Off-axis for i{i+1}-j{j+1}: {matrix_off_val}')
# Normalize matrix for the input aberration
matrix_pastis /= np.square(wfe_aber.value)
# Save matrix to file
filename_matrix = f'PASTISmatrix_num_{zern_mode.name}_{zern_mode.convention + str(zern_mode.index)}'
util.write_fits(matrix_pastis, os.path.join(resDir, filename_matrix + '.fits'), header=None, metadata=None)
log.info(f'Matrix saved to: {os.path.join(resDir, filename_matrix + ".fits")}')
# Save the PSF and DH image *cubes* as well (as opposed to each one individually)
util.write_fits(all_psfs, os.path.join(resDir, 'psfs', 'psf_cube.fits'), header=None, metadata=None)
util.write_fits(all_dhs, os.path.join(resDir, 'darkholes', 'dh_cube.fits'), header=None, metadata=None)
np.savetxt(os.path.join(resDir, 'pair-wise_contrasts.txt'), all_contrasts, fmt='%e')
# Tell us how long it took to finish.
end_time = time.time()
log.info(f'Runtime for matrix_building.py: {end_time - start_time}sec = {(end_time - start_time) / 60}min')
log.info(f'Data saved to {resDir}')
def num_matrix_luvoir(design, savepsfs=False, saveopds=True):
"""
Generate a numerical PASTIS matrix for a LUVOIR A coronagraph.
-- Depracated function, the LUVOIR PASTIS matrix is better calculated with num_matrix_multiprocess(), which can
do this for your choice of one of the implemented instruments (LUVOIR, HiCAT, JWST). --
All inputs are read from the (local) configfile and saved to the specified output directory.
The LUVOIR STDT delivery in May 2018 included three different apodizers
we can work with, you pick which of the three you want with the 'design' parameter.
:param design: string, what coronagraph design to use - 'small', 'medium' or 'large'
:param savepsfs: bool, if True, all PSFs will be saved to disk individually, as fits files, additionally to the
total PSF cube. If False, the total cube will still get saved at the very end of the script.
:param saveopds: bool, if True, all pupil surface maps of aberrated segment pairs will be saved to disk as PDF
:return overall_dir: string, experiment directory
"""
# Keep track of time
start_time = time.time()
### Parameters
# System parameters
overall_dir = util.create_data_path(CONFIG_PASTIS.get('local', 'local_data_path'), telescope='luvoir-'+design)
os.makedirs(overall_dir, exist_ok=True)
resDir = os.path.join(overall_dir, 'matrix_numerical')
# Create necessary directories if they don't exist yet
os.makedirs(resDir, exist_ok=True)
os.makedirs(os.path.join(resDir, 'OTE_images'), exist_ok=True)
os.makedirs(os.path.join(resDir, 'psfs'), exist_ok=True)
# Set up logger
util.setup_pastis_logging(resDir, f'pastis_matrix_{design}')
log.info('Building numerical matrix for LUVOIR\n')
# Read calibration aberration
zern_number = CONFIG_PASTIS.getint('calibration', 'local_zernike')
zern_mode = util.ZernikeMode(zern_number) # Create Zernike mode object for easier handling
# General telescope parameters
nb_seg = CONFIG_PASTIS.getint('LUVOIR', 'nb_subapertures')
wvln = CONFIG_PASTIS.getfloat('LUVOIR', 'lambda') * 1e-9 # m
diam = CONFIG_PASTIS.getfloat('LUVOIR', 'diameter') # m
wfe_aber = CONFIG_PASTIS.getfloat('LUVOIR', 'calibration_aberration') * 1e-9 # m
# Image system parameters
sampling = CONFIG_PASTIS.getfloat('LUVOIR', 'sampling')
# Record some of the defined parameters
log.info(f'LUVOIR apodizer design: {design}')
log.info(f'Wavelength: {wvln} m')
log.info(f'Telescope diameter: {diam} m')
log.info(f'Number of segments: {nb_seg}')
log.info(f'Sampling: {sampling} px per lambda/D')
log.info(f'wfe_aber: {wfe_aber} m')
# Copy configfile to resulting matrix directory
util.copy_config(resDir)
### Instantiate Luvoir telescope with chosen apodizer design
optics_input = CONFIG_PASTIS.get('LUVOIR', 'optics_path')
luvoir = LuvoirAPLC(optics_input, design, sampling)
### Reference images for contrast normalization and coronagraph floor
unaberrated_coro_psf, ref = luvoir.calc_psf(ref=True, display_intermediate=False, return_intermediate=False)
norm = np.max(ref)
dh_intensity = (unaberrated_coro_psf / norm) * luvoir.dh_mask
contrast_floor = np.mean(dh_intensity[np.where(luvoir.dh_mask != 0)])
log.info(f'contrast floor: {contrast_floor}')
### Generating the PASTIS matrix and a list for all contrasts
contrast_matrix = np.zeros([nb_seg, nb_seg]) # Generate empty matrix
all_psfs = []
all_contrasts = []
for i in range(nb_seg):
for j in range(nb_seg):
log.info(f'\nSTEP: {i+1}-{j+1} / {nb_seg}-{nb_seg}')
# Put aberration on correct segments. If i=j, apply only once!
luvoir.flatten()
luvoir.set_segment(i+1, wfe_aber/2, 0, 0)
if i != j:
luvoir.set_segment(j+1, wfe_aber/2, 0, 0)
log.info('Calculating coro image...')
image, inter = luvoir.calc_psf(ref=False, display_intermediate=False, return_intermediate='intensity')
# Normalize PSF by reference image
psf = image / norm
all_psfs.append(psf.shaped)
# Save image to disk
if savepsfs: # TODO: I might want to change this to matplotlib images since I save the PSF cube anyway.
filename_psf = f'psf_{zern_mode.name}_{zern_mode.convention + str(zern_mode.index)}_segs_{i+1}-{j+1}'
hcipy.write_fits(psf, os.path.join(resDir, 'psfs', filename_psf + '.fits'))
# Save OPD images for testing
if saveopds:
opd_name = f'opd_{zern_mode.name}_{zern_mode.convention + str(zern_mode.index)}_segs_{i+1}-{j+1}'
plt.clf()
hcipy.imshow_field(inter['seg_mirror'], mask=luvoir.aperture, cmap='RdBu')
plt.savefig(os.path.join(resDir, 'OTE_images', opd_name + '.pdf'))
log.info('Calculating mean contrast in dark hole')
dh_intensity = psf * luvoir.dh_mask
contrast = np.mean(dh_intensity[np.where(luvoir.dh_mask != 0)])
log.info(f'contrast: {float(contrast)}') # contrast is a Field, here casting to normal float
all_contrasts.append(contrast)
# Fill according entry in the matrix and subtract baseline contrast
contrast_matrix[i,j] = contrast - contrast_floor
# Transform saved lists to arrays
all_psfs = np.array(all_psfs)
all_contrasts = np.array(all_contrasts)
# Save the PSF image *cube* as well (as opposed to each one individually)
hcipy.write_fits(all_psfs, os.path.join(resDir, 'psfs', 'psf_cube.fits'),)
np.savetxt(os.path.join(resDir, 'pair-wise_contrasts.txt'), all_contrasts, fmt='%e')
# Filling the off-axis elements
log.info('\nCalculating off-axis matrix elements...')
matrix_two_N = np.copy(contrast_matrix) # This is just an intermediary copy so that I don't mix things up.
matrix_pastis = np.copy(contrast_matrix) # This will be the final PASTIS matrix.
for i in range(nb_seg):
for j in range(nb_seg):
if i != j:
matrix_off_val = (matrix_two_N[i,j] - matrix_two_N[i,i] - matrix_two_N[j,j]) / 2.
matrix_pastis[i,j] = matrix_off_val
log.info(f'Off-axis for i{i+1}-j{j+1}: {matrix_off_val}')
# Normalize matrix for the input aberration - this defines what units the PASTIS matrix will be in. The PASTIS
# matrix propagation function (util.pastis_contrast()) then needs to take in the aberration vector in these same
# units. I have chosen to keep this to 1nm, so, we normalize the PASTIS matrix to units of nanometers.
matrix_pastis /= np.square(wfe_aber * 1e9) # 1e9 converts the calibration aberration back to nanometers
# Save matrix to file
filename_matrix = f'PASTISmatrix_num_{zern_mode.name}_{zern_mode.convention + str(zern_mode.index)}'
hcipy.write_fits(matrix_pastis, os.path.join(resDir, filename_matrix + '.fits'))
log.info(f'Matrix saved to: {os.path.join(resDir, filename_matrix + ".fits")}')
# Tell us how long it took to finish.
end_time = time.time()
log.info(f'Runtime for matrix_building.py: {end_time - start_time}sec = {(end_time - start_time) / 60}min')
log.info(f'Data saved to {resDir}')
return overall_dir
zern_number = CONFIG_PASTIS.getint(
'calibration',
'local_zernike') # Which (Noll) Zernike we are calibrating for
wss_zern_nb = util.noll_to_wss(
zern_number) # Convert from Noll to WSS framework
# If subfolder "calibration" doesn't exist yet, create it.
if not os.path.isdir(outDir):
os.mkdir(outDir)
# If subfolder "images" in "calibration" doesn't exist yet, create it.
if not os.path.isdir(os.path.join(outDir, 'images')):
os.mkdir(os.path.join(outDir, 'images'))
# Create Zernike mode object for easier handling
zern_mode = util.ZernikeMode(zern_number)
# Create NIRCam objects, one for perfect PSF and one with coronagraph
log.info('Setting up the E2E simulation.')
nc = webbpsf.NIRCam()
# Set filter
nc.filter = filter
# Same for coronagraphic case
nc_coro = webbpsf.NIRCam()
nc_coro.filter = filter
# Add coronagraphic elements to nc_coro
nc_coro.image_mask = fpm
nc_coro.pupil_mask = lyot_stop
def analytical_model(zernike_pol, coef, cali=False):
"""
:param zernike_pol:
:param coef:
:param cali: bool; True if we already have calibration coefficients to use. False if we still need to create them.
:return:
"""
#-# Parameters
dataDir = os.path.join(CONFIG_PASTIS.get('local', 'local_data_path'),
'active')
telescope = CONFIG_PASTIS.get('telescope', 'name')
nb_seg = CONFIG_PASTIS.getint(telescope, 'nb_subapertures')
tel_size_m = CONFIG_PASTIS.getfloat(telescope, 'diameter') * u.m
real_size_seg = CONFIG_PASTIS.getfloat(
telescope, 'flat_to_flat'
) # in m, size in meters of an individual segment flatl to flat
size_seg = CONFIG_PASTIS.getint(
'numerical',
'size_seg') # pixel size of an individual segment tip to tip
wvln = CONFIG_PASTIS.getint(telescope, 'lambda') * u.nm
inner_wa = CONFIG_PASTIS.getint(telescope, 'IWA')
outer_wa = CONFIG_PASTIS.getint(telescope, 'OWA')
tel_size_px = CONFIG_PASTIS.getint(
'numerical', 'tel_size_px') # pupil diameter of telescope in pixels
im_size_pastis = CONFIG_PASTIS.getint(
'numerical', 'im_size_px_pastis') # image array size in px
sampling = CONFIG_PASTIS.getfloat(telescope, 'sampling') # sampling
size_px_tel = tel_size_m / tel_size_px # size of one pixel in pupil plane in m
px_sq_to_rad = (size_px_tel * np.pi / tel_size_m) * u.rad
zern_max = CONFIG_PASTIS.getint('zernikes', 'max_zern')
sz = CONFIG_PASTIS.getint(
'ATLAST',
'im_size_lamD_hcipy') # image size in lam/D, only used in ATLAST case
# Create Zernike mode object for easier handling
zern_mode = util.ZernikeMode(zernike_pol)
#-# Mean subtraction for piston
if zernike_pol == 1:
coef -= np.mean(coef)
#-# Generic segment shapes
if telescope == 'JWST':
# Load pupil from file
pupil = fits.getdata(
os.path.join(dataDir, 'segmentation', 'pupil.fits'))
# Put pupil in randomly picked, slightly larger image array
pup_im = np.copy(pupil) # remove if lines below this are active
#pup_im = np.zeros([tel_size_px, tel_size_px])
#lim = int((pup_im.shape[1] - pupil.shape[1])/2.)
#pup_im[lim:-lim, lim:-lim] = pupil
# test_seg = pupil[394:,197:315] # this is just so that I can display an individual segment when the pupil is 512
# test_seg = pupil[:203,392:631] # ... when the pupil is 1024
# one_seg = np.zeros_like(test_seg)
# one_seg[:110, :] = test_seg[8:, :] # this is the centered version of the individual segment for 512 px pupil
# Creat a mini-segment (one individual segment from the segmented aperture)
mini_seg_real = poppy.NgonAperture(
name='mini', radius=real_size_seg
) # creating real mini segment shape with poppy
#test = mini_seg_real.sample(wavelength=wvln, grid_size=flat_diam, return_scale=True) # fix its sampling with wavelength
mini_hdu = mini_seg_real.to_fits(wavelength=wvln,
npix=size_seg) # make it a fits file
mini_seg = mini_hdu[
0].data # extract the image data from the fits file
elif telescope == 'ATLAST':
# Create mini-segment
pupil_grid = hcipy.make_pupil_grid(dims=tel_size_px,
diameter=real_size_seg)
focal_grid = hcipy.make_focal_grid(
pupil_grid, sampling, sz, wavelength=wvln.to(
u.m).value) # fov = lambda/D radius of total image
prop = hcipy.FraunhoferPropagator(pupil_grid, focal_grid)
mini_seg_real = hcipy.hexagonal_aperture(circum_diameter=real_size_seg,
angle=np.pi / 2)
mini_seg_hc = hcipy.evaluate_supersampled(
mini_seg_real, pupil_grid, 4
) # the supersampling number doesn't really matter in context with the other numbers
mini_seg = mini_seg_hc.shaped # make it a 2D array
# Redefine size_seg if using HCIPy
size_seg = mini_seg.shape[0]
# Make stand-in pupil for DH array
pupil = fits.getdata(
os.path.join(dataDir, 'segmentation', 'pupil.fits'))
pup_im = np.copy(pupil)
#-# Generate a dark hole mask
#TODO: simplify DH generation and usage
dh_area = util.create_dark_hole(
pup_im, inner_wa, outer_wa, sampling
) # this might become a problem if pupil size is not same like pastis image size. fine for now though.
if telescope == 'ATLAST':
dh_sz = util.zoom_cen(dh_area, sz * sampling)
#-# Import information form segmentation script
Projection_Matrix = fits.getdata(
os.path.join(dataDir, 'segmentation', 'Projection_Matrix.fits'))
vec_list = fits.getdata(
os.path.join(dataDir, 'segmentation', 'vec_list.fits')) # in pixels
NR_pairs_list = fits.getdata(
os.path.join(dataDir, 'segmentation', 'NR_pairs_list_int.fits'))
# Figure out how many NRPs we're dealing with
NR_pairs_nb = NR_pairs_list.shape[0]
#-# Chose whether calibration factors to do the calibraiton with
if cali:
filename = 'calibration_' + zern_mode.name + '_' + zern_mode.convention + str(
zern_mode.index)
ck = fits.getdata(
os.path.join(dataDir, 'calibration', filename + '.fits'))
else:
ck = np.ones(nb_seg)
coef = coef * ck
#-# Generic coefficients
# the coefficients in front of the non redundant pairs, the A_q in eq. 13 in Leboulleux et al. 2018
generic_coef = np.zeros(
NR_pairs_nb
) * u.nm * u.nm # setting it up with the correct units this will have
for q in range(NR_pairs_nb):
for i in range(nb_seg):
for j in range(i + 1, nb_seg):
if Projection_Matrix[i, j, 0] == q + 1:
generic_coef[q] += coef[i] * coef[j]
#-# Constant sum and cosine sum - calculating eq. 13 from Leboulleux et al. 2018
if telescope == 'JWST':
i_line = np.linspace(-im_size_pastis / 2., im_size_pastis / 2.,
im_size_pastis)
tab_i, tab_j = np.meshgrid(i_line, i_line)
cos_u_mat = np.zeros(
(int(im_size_pastis), int(im_size_pastis), NR_pairs_nb))
elif telescope == 'ATLAST':
i_line = np.linspace(-(2 * sz * sampling) / 2.,
(2 * sz * sampling) / 2., (2 * sz * sampling))
tab_i, tab_j = np.meshgrid(i_line, i_line)
cos_u_mat = np.zeros((int((2 * sz * sampling)), int(
(2 * sz * sampling)), NR_pairs_nb))
# Calculating the cosine terms from eq. 13.
# The -1 with each NR_pairs_list is because the segment names are saved starting from 1, but Python starts
# its indexing at zero, so we have to make it start at zero here too.
for q in range(NR_pairs_nb):
# cos(b_q <dot> u): b_q with 1 <= q <= NR_pairs_nb is the basis of NRPS, meaning the distance vectors between
# two segments of one NRP. We can read these out from vec_list.
# u is the position (vector) in the detector plane. Here, those are the grids tab_i and tab_j.
# We need to calculate the dot product between all b_q and u, so in each iteration (for q), we simply add the
# x and y component.
cos_u_mat[:, :, q] = np.cos(
px_sq_to_rad *
(vec_list[NR_pairs_list[q, 0] - 1, NR_pairs_list[q, 1] - 1, 0] *
tab_i) + px_sq_to_rad *
(vec_list[NR_pairs_list[q, 0] - 1, NR_pairs_list[q, 1] - 1, 1] *
tab_j)) * u.dimensionless_unscaled
sum1 = np.sum(
coef**2
) # sum of all a_{k,l} in eq. 13 - this works only for single Zernikes (l fixed), because np.sum would sum over l too, which would be wrong.
if telescope == 'JWST':
sum2 = np.zeros(
(int(im_size_pastis), int(im_size_pastis))
) * u.nm * u.nm # setting it up with the correct units this will have
elif telescope == 'ATLAST':
sum2 = np.zeros(
(int(2 * sz * sampling), int(2 * sz * sampling))) * u.nm * u.nm
for q in range(NR_pairs_nb):
sum2 = sum2 + generic_coef[q] * cos_u_mat[:, :, q]
#-# Local Zernike
if telescope == 'JWST':
# Generate a basis of Zernikes with the mini segment being the support
isolated_zerns = zern.hexike_basis(nterms=zern_max,
npix=size_seg,
rho=None,
theta=None,
vertical=False,
outside=0.0)
# Calculate the Zernike that is currently being used and put it on one single subaperture, the result is Zer
# Apply the currently used Zernike to the mini-segment.
if zernike_pol == 1:
Zer = np.copy(mini_seg)
elif zernike_pol in range(2, zern_max - 2):
Zer = np.copy(mini_seg)
Zer = Zer * isolated_zerns[zernike_pol - 1]
# Fourier Transform of the Zernike - the global envelope
mf = mft.MatrixFourierTransform()
ft_zern = mf.perform(Zer, im_size_pastis / sampling, im_size_pastis)
elif telescope == 'ATLAST':
isolated_zerns = hcipy.make_zernike_basis(num_modes=zern_max,
D=real_size_seg,
grid=pupil_grid,
radial_cutoff=False)
Zer = hcipy.Wavefront(mini_seg_hc * isolated_zerns[zernike_pol - 1],
wavelength=wvln.to(u.m).value)
# Fourier transform the Zernike
ft_zern = prop(Zer)
#-# Final image
if telescope == 'JWST':
# Generating the final image that will get passed on to the outer scope, I(u) in eq. 13
intensity = np.abs(ft_zern)**2 * (sum1.value + 2. * sum2.value)
elif telescope == 'ATLAST':
intensity = ft_zern.intensity.shaped * (sum1.value + 2. * sum2.value)
# PASTIS is only valid inside the dark hole, so we cut out only that part
if telescope == 'JWST':
tot_dh_im_size = sampling * (outer_wa + 3)
intensity_zoom = util.zoom_cen(
intensity, tot_dh_im_size
) # zoom box is (owa + 3*lambda/D) wide, in terms of lambda/D
dh_area_zoom = util.zoom_cen(dh_area, tot_dh_im_size)
dh_psf = dh_area_zoom * intensity_zoom
elif telescope == 'ATLAST':
dh_psf = dh_sz * intensity
"""
# Create plots.
plt.subplot(1, 3, 1)
plt.imshow(pupil, origin='lower')
plt.title('JWST pupil and diameter definition')
plt.plot([46.5, 464.5], [101.5, 409.5], 'r-') # show how the diagonal of the pupil is defined
plt.subplot(1, 3, 2)
plt.imshow(mini_seg, origin='lower')
plt.title('JWST individual mini-segment')
plt.subplot(1, 3, 3)
plt.imshow(dh_psf, origin='lower')
plt.title('JWST dark hole')
plt.show()
"""
# dh_psf is the image of the dark hole only, the pixels outside of it are zero
# intensity is the entire final image
return dh_psf, intensity