Skip to content

GitLab

Commit dfffbc3a authored by Arthur Vigan's avatar Arthur Vigan
Browse files

Parallelized version of the code

parent bb6415e7
Hide whitespace changes
Inline Side-by-side
hrm-hc/MAIN_multiproc.py
View file @ dfffbc3a
......@@ -6,6 +6,7 @@ from scipy.interpolate import interp2d
from numpy import pi
import os
import sys
import multiprocessing as mp
from time import localtime, strftime
import time
import matplotlib.pyplot as plt
......@@ -13,6 +14,105 @@ from astropy.io import fits
from utils_basic import minterp2d, DISPER, VECT, VECT1, BEAMSHIFT, ELEVATION, PARANG, ft, MAKE_ELT, PSF, LD2MAS, MAS2LD, M4Magic, DSP, EELTNLS, APPLY_CROSSTALK, DISPER_ADC
def compute_image_at_wavelength(i, k):
#st = time.time()
WF_CP_0, WF_NCPA_0 = BEAMSHIFT(N, SCREEN, Pupil_shift_vect_0, D_opt_vect, Z_opt_vect, lambda_wfs, lambda_vect[k], 90-ROT_ANG_vect[i], ROT_ANG_vect[0], rot_flag*rot_vect, com_vect, gamma)
#log.info(time.time()-st)
WF_CP_0 = np.squeeze(WF_CP_0)
WF_NCPA_0 = np.squeeze(WF_NCPA_0)
WF_0=(WF_CP_0-WF_corr_0+EELT_phase_0+WF_NCPA_0)
DSP_V, DSP_I = DSP(WF_0, BigD/D, np.transpose(ORIGINAL_PUPIL), 0, 1000)
log.info('HA={0}: Aberration level (0-1000 lambda/D) at {1:.1f} nm: {2:.1f} nm'.format(HA_0, lambda_vect[k]*1e9, 1e9*DSP_V))
#DSP_V, DSP_I = DSP(WF_CP_0+EELT_phase_0+WF_NCPA_0, 1132/1024, PUPILLE_0, 2, 30)
#log.info('HA={0}: Aberration level pre correction (2-30 lambda/D) at {1} nm: {2} nm'.format(HA_0, lambda_vect[k]*1e9, 1e9*DSP_V))
#DSP_V, DSP_I = DSP(WF_0, 1132/1024, PUPILLE_0, 2, 30)
#log.info('HA={0}: Aberration level (2-30 lambda/D) at {1} nm: {2} nm'.format(HA_0, lambda_vect[k]*1e9, 1e9*DSP_V))
OWA = MAS2LD(FOV, 38.542, lambda_vect[k])/2
if ADC_flag == 1:
#alpha_disp_0 = MAS2LD(1000*DISPER_ADC(1.45, lambda_vect[k]*1e6, 90-ELEVATION(DEC, HA_0), z_min, z_max), 38.542, lambda_vect[k])-MAS2LD(1000*DISPER_ADC(1.45, lambda_vect[-1]*1e6, 90-ELEVATION(DEC, HA_0), z_min, z_max), 38.542, lambda_vect[k])
alpha_disp_0 = MAS2LD(1000*DISPER_ADC(lambda_vect[0]*1e6, lambda_vect[k]*1e6, 90-ELEVATION(DEC, HA_0), z_min, z_max), 38.542, lambda_vect[k])
else:
#alpha_disp_0 = MAS2LD(1000*DISPER(1.45, lambda_vect[k]*1e6, 90-ELEVATION(DEC, HA_0)), 38.542, lambda_vect[k])-MAS2LD(1000*DISPER(1.45, lambda_vect[-1]*1e6, 90-ELEVATION(DEC, HA_0)), 38.542, lambda_vect[k])
alpha_disp_0 = MAS2LD(1000*DISPER(lambda_vect[0]*1e6, lambda_vect[k]*1e6, 90-ELEVATION(DEC, HA_0)), 38.542, lambda_vect[k])
#log.info('Differential dispersion angle with 2.45um: {0}'.format(alpha_disp_0))
x, X, Y, R, T = VECT(len(SP), 1132/1024)
E_0 = ft(np.transpose(SP)*(1-MISSSEG)*np.exp(2*1j*pi*WF_0/lambda_vect[k])*np.exp(-2*1j*pi*alpha_disp_0*Y), 1, BigD/D, BigD/D, 2*OWA, 2*OWA, 1, N_im, N_im)
#Commented to save time ; halo only
I_0 = np.abs(E_0)**2
#HALO
FT2DSP_norm = 1e-18*(2*pi/lambda_vect[k])**2
# Normalisation by maximum of phase autocorrelation (in 0) = phase variance
# then PHASE STRUCTURE FUNCTION (Roddier) = 2*(autocol(0, 0) - autocol(u, v))
# then Atmospheric Optical Transfert function
FT_PST = RENORM_FT_DSP*FT2DSP_norm
#Rotation of FT_PST to simulate changing wind direction
FT_PST = rotate(FT_PST, Wind_Angle, order=1, reshape=False)
PH_STR = 2*(FT_PST[int(size_dsp//2+1), int(size_dsp//2+1)]-FT_PST)
# Optical transfer function of the atmosphere (long exposure model)
FTO_atm = np.exp(-PH_STR/2.)
a = (time.time()-start_time)/60.
I_s_0 = np.abs(ft(np.transpose(SP)*(1-MISSSEG)*np.exp(2*1j*pi*WF_0/lambda_vect[k])*np.exp(2*1j*pi*alpha_disp_0*Y), 1, BigD/D, BigD/D, 2*OWA, 2*OWA, 1, len(dsp), len(dsp)))**2
FTO_tel_0 = ft(I_s_0, 1, 2*OWA, 2*OWA, DiamFTO, DiamFTO, 1, len(dsp), len(dsp))
I_s_post_0 = np.abs(ft(FTO_tel_0*FTO_atm, 1, DiamFTO, DiamFTO, 2*OWA, 2*OWA, 1, N_im, N_im))
I_s_0_ForSR = np.abs(ft(np.transpose(ORIGINAL_PUPIL), 1, BigD/D, BigD/D, 2*OWA, 2*OWA, 1, len(dsp), len(dsp)))**2
FTO_tel_0_ForSR = ft(I_s_0_ForSR, 1, 2*OWA, 2*OWA, DiamFTO, DiamFTO, 1, len(dsp), len(dsp))
SR = np.sum(np.abs(FTO_tel_0_ForSR*FTO_atm))/np.sum(np.abs(FTO_tel_0_ForSR))
log.info('Wavelength {:4d}/{:4d}'.format(k+1, N_LD))
log.info(' => Strehl at {:4.2f}nm: {:4.2f}'.format(lambda_vect[k]*1e9, SR))
#Normalization of I_s_post_0 to reflect a unit energy arriving on
#the pupil, i.e., we just have to multiply the PSF by the number of
#photons collected by the aperture. Careful: this normalization
#takes into account the central obscuration, and the apodizer
#throughput.
I_s_post_0 = I_s_post_0/PUP_NORM_INTENSITY*(2*OWA/N_im)**2
#[u, U, V, Ru, Tu]=VECT(N_im, 2*OWA);
#REG_100=(Ru<MAS2LD(125, 38.542, lambda_vect(k)*1e9)).*(Ru>MAS2LD(75, 38.542, lambda_vect(k)*1e9));
#REG_200=(Ru<MAS2LD(225, 38.542, lambda_vect(k)*1e9)).*(Ru>MAS2LD(175, 38.542, lambda_vect(k)*1e9));
#I_stat_nohalo=I_0/max(I_0(:));
#I_stat=I_s_post_0/max(I_s_post_0(:));
#disp(['5 sigma static Contrast at 100 mas - NO HALO: ' num2str(5*std(I_stat_nohalo(REG_100==1)))])
#disp(['5 sigma static Contrast at 200 mas - NO HALO: ' num2str(5*std(I_stat_nohalo(REG_200==1)))])
#disp(['Static Contrast -- std -- w/ E2E model: ' num2str(5*std(I_stat(REG==1)))])
#fits.writeto(path_directory + 'PSF_HALO_ON_Nexp{0:03d}_Nlbd{1:04d}.fits'.format(i, k), I_s_post_0)
# I1D = I_s_post_0[:, 107]
# I1D = I1D/sum(I1D)
# V1D = Y_mask[:, 107]
# alpha_shift = -sum(I1D*V1D)
# shift_value[k] = alpha_shift/(x_mask[1]-x_mask[0])
I_temp[k] = I_s_post_0
I_temp_nomask[k] = I_s_post_0
# if exhaustive_writing_flag==1:
# fits.writeto(path_directory + 'PSF_HALO_OFF_Nexp{0}_Nlbd{1}.fits'.format(i, k), I_0)
# WF_small_0_f = interp2d(x, x, WF_0, kind='linear')
# WF_small_0 = WF_small_0_f(x2, x2)
# WF_small_0 = WF_small_0*P_small_0
# fits.writeto(path_directory + 'WF_Nexp{0}_Nlbd{1}.fits'.format(i, k), WF_small_0)
return k, I_s_post_0, SR
if __name__ == '__main__':
# timing
start_time = time.time()
......@@ -133,7 +233,7 @@ if __name__ == '__main__':
K2_high = [2.197, 2.294, 2.390, 17000]
if BAND=='TEST_H':
N_LD = 3
N_LD = 6
LBD_min = H_band[0]*1e-6
LBD_max = H_band[2]*1e-6
if APODIZER == 'SP_1':
......@@ -362,7 +462,7 @@ if __name__ == '__main__':
os.mkdir(path_directory)
else:
print('Old directory exists!')
stop
logging.basicConfig(filename=path_directory + 'simulation_log.txt', filemode='w',
format='%(message)s', level=logging.DEBUG)
log = logging.getLogger('harmoni-hc')
......@@ -659,107 +759,24 @@ if __name__ == '__main__':
MAXFTDSP = np.max(FT_dsp)
RENORM_FT_DSP = FT_dsp/MAXFTDSP*SUMDSP
ncpu = 4 # int(os.environ.get('PBS_NUM_PPN', mp.cpu_count()))
log.info('Using {} CPU to compute images'.format(ncpu))
pool = mp.Pool(processes=ncpu)
tasks = []
for k in range(N_LD):
#st = time.time()
WF_CP_0, WF_NCPA_0 = BEAMSHIFT(N, SCREEN, Pupil_shift_vect_0, D_opt_vect, Z_opt_vect, lambda_wfs, lambda_vect[k], 90-ROT_ANG_vect[i], ROT_ANG_vect[0], rot_flag*rot_vect, com_vect, gamma)
#log.info(time.time()-st)
WF_CP_0 = np.squeeze(WF_CP_0)
WF_NCPA_0 = np.squeeze(WF_NCPA_0)
WF_0=(WF_CP_0-WF_corr_0+EELT_phase_0+WF_NCPA_0)
DSP_V, DSP_I = DSP(WF_0, BigD/D, np.transpose(ORIGINAL_PUPIL), 0, 1000)
log.info('HA={0}: Aberration level (0-1000 lambda/D) at {1} nm: {2} nm'.format(HA_0, lambda_vect[k]*1e9, 1e9*DSP_V))
#DSP_V, DSP_I = DSP(WF_CP_0+EELT_phase_0+WF_NCPA_0, 1132/1024, PUPILLE_0, 2, 30)
#log.info('HA={0}: Aberration level pre correction (2-30 lambda/D) at {1} nm: {2} nm'.format(HA_0, lambda_vect[k]*1e9, 1e9*DSP_V))
#DSP_V, DSP_I = DSP(WF_0, 1132/1024, PUPILLE_0, 2, 30)
#log.info('HA={0}: Aberration level (2-30 lambda/D) at {1} nm: {2} nm'.format(HA_0, lambda_vect[k]*1e9, 1e9*DSP_V))
OWA = MAS2LD(FOV, 38.542, lambda_vect[k])/2
if ADC_flag == 1:
#alpha_disp_0 = MAS2LD(1000*DISPER_ADC(1.45, lambda_vect[k]*1e6, 90-ELEVATION(DEC, HA_0), z_min, z_max), 38.542, lambda_vect[k])-MAS2LD(1000*DISPER_ADC(1.45, lambda_vect[-1]*1e6, 90-ELEVATION(DEC, HA_0), z_min, z_max), 38.542, lambda_vect[k])
alpha_disp_0 = MAS2LD(1000*DISPER_ADC(lambda_vect[0]*1e6, lambda_vect[k]*1e6, 90-ELEVATION(DEC, HA_0), z_min, z_max), 38.542, lambda_vect[k])
else:
#alpha_disp_0 = MAS2LD(1000*DISPER(1.45, lambda_vect[k]*1e6, 90-ELEVATION(DEC, HA_0)), 38.542, lambda_vect[k])-MAS2LD(1000*DISPER(1.45, lambda_vect[-1]*1e6, 90-ELEVATION(DEC, HA_0)), 38.542, lambda_vect[k])
alpha_disp_0 = MAS2LD(1000*DISPER(lambda_vect[0]*1e6, lambda_vect[k]*1e6, 90-ELEVATION(DEC, HA_0)), 38.542, lambda_vect[k])
#log.info('Differential dispersion angle with 2.45um: {0}'.format(alpha_disp_0))
x, X, Y, R, T = VECT(len(SP), 1132/1024)
E_0 = ft(np.transpose(SP)*(1-MISSSEG)*np.exp(2*1j*pi*WF_0/lambda_vect[k])*np.exp(-2*1j*pi*alpha_disp_0*Y), 1, BigD/D, BigD/D, 2*OWA, 2*OWA, 1, N_im, N_im)
#Commented to save time ; halo only
I_0 = np.abs(E_0)**2
#HALO
FT2DSP_norm = 1e-18*(2*pi/lambda_vect[k])**2
# Normalisation by maximum of phase autocorrelation (in 0) = phase variance
# then PHASE STRUCTURE FUNCTION (Roddier) = 2*(autocol(0, 0) - autocol(u, v))
# then Atmospheric Optical Transfert function
FT_PST = RENORM_FT_DSP*FT2DSP_norm
#Rotation of FT_PST to simulate changing wind direction
FT_PST = rotate(FT_PST, Wind_Angle, order=1, reshape=False)
PH_STR = 2*(FT_PST[int(size_dsp//2+1), int(size_dsp//2+1)]-FT_PST)
# Optical transfer function of the atmosphere (long exposure model)
FTO_atm = np.exp(-PH_STR/2.)
a = (time.time()-start_time)/60.
try:
t_iter = a / (i*N_LD + k)
except ZeroDivisionError:
t_iter = 0
t_left = (N_LD*len(HA_vect)-(i*N_LD + k))*t_iter
I_s_0 = np.abs(ft(np.transpose(SP)*(1-MISSSEG)*np.exp(2*1j*pi*WF_0/lambda_vect[k])*np.exp(2*1j*pi*alpha_disp_0*Y), 1, BigD/D, BigD/D, 2*OWA, 2*OWA, 1, len(dsp), len(dsp)))**2
FTO_tel_0 = ft(I_s_0, 1, 2*OWA, 2*OWA, DiamFTO, DiamFTO, 1, len(dsp), len(dsp))
I_s_post_0 = np.abs(ft(FTO_tel_0*FTO_atm, 1, DiamFTO, DiamFTO, 2*OWA, 2*OWA, 1, N_im, N_im))
I_s_0_ForSR = np.abs(ft(np.transpose(ORIGINAL_PUPIL), 1, BigD/D, BigD/D, 2*OWA, 2*OWA, 1, len(dsp), len(dsp)))**2
FTO_tel_0_ForSR = ft(I_s_0_ForSR, 1, 2*OWA, 2*OWA, DiamFTO, DiamFTO, 1, len(dsp), len(dsp))
SR[i, k] = np.sum(np.abs(FTO_tel_0_ForSR*FTO_atm))/np.sum(np.abs(FTO_tel_0_ForSR))
log.info('Wavelength {:4d}/{:4d} - {:.2f} min left'.format(k+1, N_LD, t_left))
log.info(' => Strehl at {:4.2f}nm: {:4.2f}'.format(lambda_vect[k]*1e9, SR[i, k]))
#Normalization of I_s_post_0 to reflect a unit energy arriving on
#the pupil, i.e., we just have to multiply the PSF by the number of
#photons collected by the aperture. Careful: this normalization
#takes into account the central obscuration, and the apodizer
#throughput.
I_s_post_0 = I_s_post_0/PUP_NORM_INTENSITY*(2*OWA/N_im)**2
#[u, U, V, Ru, Tu]=VECT(N_im, 2*OWA);
#REG_100=(Ru<MAS2LD(125, 38.542, lambda_vect(k)*1e9)).*(Ru>MAS2LD(75, 38.542, lambda_vect(k)*1e9));
#REG_200=(Ru<MAS2LD(225, 38.542, lambda_vect(k)*1e9)).*(Ru>MAS2LD(175, 38.542, lambda_vect(k)*1e9));
#I_stat_nohalo=I_0/max(I_0(:));
#I_stat=I_s_post_0/max(I_s_post_0(:));
#disp(['5 sigma static Contrast at 100 mas - NO HALO: ' num2str(5*std(I_stat_nohalo(REG_100==1)))])
#disp(['5 sigma static Contrast at 200 mas - NO HALO: ' num2str(5*std(I_stat_nohalo(REG_200==1)))])
#disp(['Static Contrast -- std -- w/ E2E model: ' num2str(5*std(I_stat(REG==1)))])
#fits.writeto(path_directory + 'PSF_HALO_ON_Nexp{0:03d}_Nlbd{1:04d}.fits'.format(i, k), I_s_post_0)
# I1D = I_s_post_0[:, 107]
# I1D = I1D/sum(I1D)
# V1D = Y_mask[:, 107]
# alpha_shift = -sum(I1D*V1D)
# shift_value[k] = alpha_shift/(x_mask[1]-x_mask[0])
I_temp[k, :, :] = I_s_post_0
I_temp_nomask[k, :, :] = I_s_post_0
# if exhaustive_writing_flag==1:
# fits.writeto(path_directory + 'PSF_HALO_OFF_Nexp{0}_Nlbd{1}.fits'.format(i, k), I_0)
# WF_small_0_f = interp2d(x, x, WF_0, kind='linear')
# WF_small_0 = WF_small_0_f(x2, x2)
# WF_small_0 = WF_small_0*P_small_0
# fits.writeto(path_directory + 'WF_Nexp{0}_Nlbd{1}.fits'.format(i, k), WF_small_0)
# compute_image_at_wavelength(i, k)
tasks.append(pool.apply_async(compute_image_at_wavelength, args=(i, k)))
pool.close()
pool.join()
# get and reorder results
for task in tasks:
k, img, strehl = task.get()
I_temp[k] = img
I_temp_nomask[k] = img
SR[i, k] = strehl
if crosstalk_flag == 1:
I_temp = APPLY_CROSSTALK(I_temp)
......@@ -787,7 +804,6 @@ if __name__ == '__main__':
MASK_temp = shift(MASK, [-shift_pix_FPM, 0], output=None, order=1, mode='constant', cval=0.0, prefilter=True)
for k in range(N_LD):
if ADC_flag == 1:
shift_value = 1000*DISPER_ADC(lambda_vect[0]*1e6, lambda_vect[k]*1e6, 90-ELEVATION(DEC, HA_0), z_min, z_max)
shift_value /= MASperPIXEL
......@@ -797,9 +813,9 @@ if __name__ == '__main__':
#print(shift_value*MASperPIXEL)
I_temp[k, :, :] = I_temp[k, :, :]*np.flipud(MASK_temp)
I_temp[k] = I_temp[k]*np.flipud(MASK_temp)
I_temp[k, :, :] = shift(I_temp[k, :, :], [shift_value, 0], output=None, order=1, mode='constant', cval=0.0, prefilter=True)
I_temp[k] = shift(I_temp[k], [shift_value, 0], output=None, order=1, mode='constant', cval=0.0, prefilter=True)
image_filename = path_directory + 'PSF_HALO_ON_masked_centered_Nexp{0:04d}.fits'.format(i)
image_filename_nomask = path_directory + 'PSF_HALO_ON_notmasked_centered_Nexp{0:04d}.fits'.format(i)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment