def create_density_maps():
# Get command line arguments
args = {}
if comm_rank == 0:
print(':Registered %d processes' % comm_size)
args["simdir"] = sys.argv[1]
args["hfdir"] = sys.argv[2]
args["snapnum"] = int(sys.argv[3])
args["zs"] = float(sys.argv[4]) / 10
args = comm.bcast(args)
label = args["simdir"].split('/')[-2].split('_')[2]
# Organize devision of Sub-&Halos over Processes on Proc. 0
if comm_rank == 0:
# Load simulation
s = read_hdf5.snapshot(args["snapnum"], args["simdir"])
s.read(["Coordinates", "Masses", "GFM_StellarFormationTime"],
parttype=[0, 1, 4]) #[0,1,4]
scale = 1e-3 * s.header.hubble
# Define Cosmology
cosmo = LambdaCDM(H0=s.header.hubble * 100,
Om0=s.header.omega_m,
Ode0=s.header.omega_l)
cosmosim = {
'omega_M_0': s.header.omega_m,
'omega_lambda_0': s.header.omega_l,
'omega_k_0': 0.0,
'h': s.header.hubble
}
redshift = s.header.redshift
print(': Redshift: %f' % redshift)
# Sort Sub-&Halos over Processes
df = pd.read_csv(args["hfdir"] + 'halos_%d.dat' % args["snapnum"],
sep='\s+',
skiprows=16,
usecols=[0, 2, 4, 9, 10, 11],
names=['ID', 'Mvir', 'Vrms', 'X', 'Y', 'Z'])
df = df[df['Mvir'] > 5e11]
sh_id = df['ID'].values.astype('float64')
sh_vrms = df['Vrms'].values.astype('float64')
sh_x = df['X'].values.astype('float64')
sh_y = df['Y'].values.astype('float64')
sh_z = df['Z'].values.astype('float64')
del df
hist_edges = procdiv.histedges_equalN(sh_x, comm_size)
SH = procdiv.cluster_subhalos(sh_id, sh_vrms, sh_x, sh_y, sh_z,
hist_edges, comm_size)
# Calculate overlap for particle cuboids
c = (const.c).to_value('km/s')
fov_rad = 4 * np.pi * (np.percentile(SH['Vrms'], 90) / c)**2
sh_dist = (cosmo.comoving_distance(redshift)).to_value('Mpc')
alpha = 6 # multiplied by 4 because of Oguri&Marshall
overlap = 0.5 * alpha * fov_rad * sh_dist #[Mpc] half of field-of-view
print('Cuboids overlap is: %f [Mpc]' % overlap)
# Sort Particles over Processes
## Dark Matter
DM = {
'Mass': (s.data['Masses']['dm']).astype('float64'),
'Pos': (s.data['Coordinates']['dm'] * scale).astype('float64')
}
DM = procdiv.cluster_particles(DM, hist_edges, comm_size)
## Gas
Gas = {
'Mass': (s.data['Masses']['gas']).astype('float64'),
'Pos': (s.data['Coordinates']['gas'] * scale).astype('float64')
}
Gas = procdiv.cluster_particles(Gas, hist_edges, comm_size)
## Stars
age = (s.data['GFM_StellarFormationTime']['stars']).astype('float64')
Star = {
'Mass': (s.data['Masses']['stars'][age >= 0]).astype('float64'),
'Pos': (s.data['Coordinates']['stars'][age >= 0, :] *
scale).astype('float64')
}
del age
Star = procdiv.cluster_particles(Star, hist_edges, comm_size)
else:
c = None
alpha = None
overlap = None
cosmosim = None
cosmo = None
redshift = None
hist_edges = None
SH = {
'ID': None,
'Vrms': None,
'X': None,
'Y': None,
'Z': None,
'split_size_1d': None,
'split_disp_1d': None
}
DM = {
'Mass': None,
'X': None,
'Y': None,
'Z': None,
'split_size_1d': None,
'split_disp_1d': None
}
Gas = {
'Mass': None,
'X': None,
'Y': None,
'Z': None,
'split_size_1d': None,
'split_disp_1d': None
}
Star = {
'Mass': None,
'X': None,
'Y': None,
'Z': None,
'split_size_1d': None,
'split_disp_1d': None
}
# Broadcast variables over all processors
sh_split_size_1d = comm.bcast(SH['split_size_1d'], root=0)
dm_split_size_1d = comm.bcast(DM['split_size_1d'], root=0)
gas_split_size_1d = comm.bcast(Gas['split_size_1d'], root=0)
star_split_size_1d = comm.bcast(Star['split_size_1d'], root=0)
c = comm.bcast(c, root=0)
alpha = comm.bcast(alpha, root=0)
overlap = comm.bcast(overlap, root=0)
cosmo = comm.bcast(cosmo, root=0)
redshift = comm.bcast(redshift, root=0)
hist_edges = comm.bcast(hist_edges, root=0)
SH = procdiv.scatter_subhalos(SH,
sh_split_size_1d,
comm_rank,
comm,
root_proc=0)
DM = procdiv.scatter_particles(DM,
dm_split_size_1d,
comm_rank,
comm,
root_proc=0)
Gas = procdiv.scatter_particles(Gas,
gas_split_size_1d,
comm_rank,
comm,
root_proc=0)
Star = procdiv.scatter_particles(Star,
star_split_size_1d,
comm_rank,
comm,
root_proc=0)
print(
': Proc. %d got: \n\t %d Sub-&Halos \n\t %d dark matter \n\t %d gas \n\t %d stars \n'
% (comm_rank, int(sh_split_size_1d[comm_rank]),
int(dm_split_size_1d[comm_rank]), int(gas_split_size_1d[comm_rank]),
int(star_split_size_1d[comm_rank])))
## Run over Sub-&Halos
zl = redshift
zs = args["zs"]
ncells = [512, 256, 128]
nparts = [1, 2, 4, 8]
M200 = np.ones(len(SH['ID']))
ID = np.ones(len(SH['ID']))
Rein = np.ones((len(SH['ID']), len(ncells), len(nparts)))
for ll in range(len(SH['ID'])):
# Define field-of-view
fov_rad = 4 * np.pi * (SH['Vrms'][ll] / c)**2
sh_dist = (cosmo.comoving_distance(redshift)).to_value('Mpc')
fov_Mpc = alpha * fov_rad * sh_dist #[Mpc] is it the diameter?
fov_arc = (fov_Mpc / cf.Da(zl, cosmo) * u.rad).to_value('arcsec')
sigma_cr = sigma_crit(zl, zs, cosmo).to_value('Msun Mpc-2')
# Check cuboid boundary condition,
# that all surface densities are filled with particles
if ((SH['Pos'][ll,0]-hist_edges[comm_rank] < overlap) or
(hist_edges[comm_rank+1]-overlap < \
SH['Pos'][ll,0]-hist_edges[comm_rank])):
if fov_Mpc * 0.45 > overlap:
print("FOV is bigger than cuboids overlap: %f > %f" % \
(fov_Mpc*0.45, overlap))
continue
## Run over different Ncells
for cc in range(len(ncells)):
dsx_arc = fov_arc / ncells[cc] #[arcsec] pixel size
## Run over particle reductions
for mm in range(len(nparts)):
smlpixel = 20 # maximum smoothing pixel length
pos, indx = dmap.select_particles(
Gas['Pos'],
SH['Pos'][ll], #*a/h,
fov_Mpc,
'box')
gas_sigma = dmap.projected_density_pmesh_adaptive(
pos[::nparts[mm], :],
Gas['Mass'][indx][::nparts[mm]],
fov_Mpc,
ncells[cc],
hmax=smlpixel)
pos, indx = dmap.select_particles(
Star['Pos'],
SH['Pos'][ll], #*a/h,
fov_Mpc,
'box')
star_sigma = dmap.projected_density_pmesh_adaptive(
pos[::nparts[mm], :],
Star['Mass'][indx][::nparts[mm]],
fov_Mpc,
ncells[cc],
hmax=smlpixel)
pos, indx = dmap.select_particles(
DM['Pos'],
SH['Pos'][ll], #*a/h,
fov_Mpc,
'box')
dm_sigma = dmap.projected_density_pmesh_adaptive(
pos[::nparts[mm], :],
DM['Mass'][indx][::nparts[mm]],
fov_Mpc, #[Mpc]
ncells[cc],
hmax=smlpixel)
tot_sigma = dm_sigma + gas_sigma + star_sigma
# Make sure that density-map is filled
while 0.0 in tot_sigma:
smlpixel += 5
dm_sigma = dmap.projected_density_pmesh_adaptive(
pos[::nparts[mm], :],
DM['Mass'][indx][::nparts[mm]],
fov_Mpc, #[Mpc]
ncells[cc],
hmax=smlpixel)
tot_sigma = dm_sigma + gas_sigma + star_sigma
#tmap.plotting(tot_sigma, ncells[cc], fov_Mpc, zl)
# initialize the coordinates of grids (light rays on lens plan)
lpv = np.linspace(-(fov_arc - dsx_arc) / 2,
(fov_arc - dsx_arc) / 2, ncells[cc])
lp1, lp2 = np.meshgrid(lpv, lpv) #[arcsec]
fig = plt.figure()
ax = fig.add_subplot(111)
# Calculate convergence map
kappa = tot_sigma / sigma_cr
# Calculate Deflection Maps
alpha1, alpha2, mu_map, phi, detA, lambda_t = cal_lensing_signals(
kappa, fov_arc, ncells[cc])
# Calculate Einstein Radii
Rein[ll, cc, mm] = einstein_radii(lp1, lp2, detA, lambda_t, zl,
cosmo, ax, 'med', ll)
#print('Rein = %f' % Rein[ll, cc, mm])
ID[ll] = SH['ID'][ll]
#plt.close(fig)
output = {}
for cc in range(len(ncells)):
for mm in range(len(nparts)):
output[(str(ncells[cc]), str(nparts[mm]))] = Rein[:, cc, mm]
df = pd.DataFrame.from_dict(output)
df['ID'] = ID
#self.df = pd.concat([self.df, dfp], axis=1)
fname = 'DMConvTest_' + label + '_' + str(comm_rank) + '_zs150.h5'
df.to_hdf(fname, key='Rein', mode='w')
plt.close(fig)
def lensing_signal():
# Get command line arguments
args = {}
args["simdir"] = sys.argv[1]
args["dmdir"] = sys.argv[2]
args["lcdir"] = sys.argv[3]
args["outbase"] = sys.argv[4]
args["ncells"] = int(sys.argv[5])
#args["simdir"] = '/cosma6/data/dp004/dc-arno1/SZ_project/full_physics/L62_N512_F5_kpc/'
#args["dmdir"] = '/cosma5/data/dp004/dc-beck3/StrongLensing/DensityMap/full_physics/L62_N512_F5_kpc/Lightcone/'
#args["lcdir"] = '/cosma5/data/dp004/dc-beck3/StrongLensing/LightCone/full_physics/Rockstar/LC_SN_L62_N512_F5_kpc'
#args["outbase"] = '/cosma5/data/dp004/dc-beck3/StrongLensing/LensingMap/full_physics/Rockstar/L62_N512_F5_kpc/Lightcone/'
#args["ncells"] = 512
# Names of all available Density maps
dmfile = glob.glob(args["dmdir"]+'*.h5')
dmfile.sort(key = lambda x: x[-4])
# Names of all available Lightcones
lcfile = glob.glob(args["lcdir"]+'*.h5')
lcfile.sort(key = lambda x: x[-4])
lenslistinit(); srclistinit()
# Run through files
for ff in range(len(dmfile)):
print('\n')
print('------------- \n Reading Files: \n%s\n%s' % \
(dmfile[ff].split('/')[-2:], lcfile[ff].split('/')[-2:]))
# Load density maps
dmf = h5py.File(dmfile[ff], 'r')
dmdf = pd.DataFrame({'HF_ID' : dmf['HF_ID'],
'LC_ID' : dmf['LC_ID'],
'fov_Mpc' : dmf['fov_Mpc']})
s1 = pd.Series(dict(list(enumerate(dmf['density_map']))), index=dmdf.index)
dmdf['density_map'] = s1
dmdf = dmdf.set_index('LC_ID')
# Load Lightcones
lcf = h5py.File(lcfile[ff], 'r')
lcdf = pd.DataFrame({'HF_ID' : lcf['HF_ID'].value,
'LC_ID' : lcf['LC_ID'].value,
'zl' : lcf['Halo_z'].value,
'vrms' : lcf['VelDisp'].value,
'snapnum' : lcf['snapnum'].value,
'fov_Mpc' : lcf['FOV'][:][1]})
lcdf = lcdf.set_index('LC_ID')
srcdf = {'Src_ID' : lcf['Src_ID'].value,
'zs' : lcf['Src_z'].value,
'SrcPosSky' : lcf['SrcPosSky'].value,
'SrcAbsMag' : lcf['SrcAbsMag'].value}
print('The minimum Vrms is %f' % (np.min(lcdf['vrms'].values)))
dmdf = dmdf.sort_values(by=['LC_ID'])
lcdf = lcdf.sort_values(by=['LC_ID'])
# sanity check
assert len(lcdf.index.intersection(dmdf.index)) == len(dmdf.index.values)
lcdf['density_map'] = dmdf['density_map']
#lcdf = lcdf.sort_values(by=['snapnum'])
s = read_hdf5.snapshot(45, args["simdir"])
# Cosmological Parameters
cosmo = LambdaCDM(H0=s.header.hubble*100,
Om0=s.header.omega_m,
Ode0=s.header.omega_l)
# Run through lenses
print('There are %d lenses in file' % (len(lcdf.index.values)))
for ll in range(len(lcdf.index.values)):
lens = lcdf.iloc[ll]
#print('working on lens %d' % lens['HF_ID'])
# convert. box size and pixels size from ang. diam. dist. to arcsec
FOV_arc = (lens['fov_Mpc']/cf.Da(lens['zl'], cosmo)*u.rad).to_value('arcsec')
dsx_arc = FOV_arc/args["ncells"] #[arcsec] pixel size
# initialize the coordinates of grids (light rays on lens plan)
lpv = np.linspace(-(FOV_arc-dsx_arc)/2, (FOV_arc-dsx_arc)/2, args["ncells"])
lp1, lp2 = np.meshgrid(lpv, lpv) #[arcsec]
zs, Src_ID, SrcPosSky = lt.source_selection(
srcdf['Src_ID'], srcdf['zs'], srcdf['SrcPosSky'],
lcdf.index.values[ll])
# Run through sources
check_for_sources = 0
fig = plt.figure()
ax = fig.add_subplot(111)
for ss in range(len(Src_ID)):
# Calculate critical surface density
sigma_cr = lt.sigma_crit(lens['zl'],
zs[ss],
cosmo).to_value('Msun Mpc-2')
# convert source position from Mpc to arcsec
beta = lt.mpc2arc(SrcPosSky[ss])
beta = [bb*1e-3 for bb in beta]
# Calculate convergence map
kappa = lens['density_map']/sigma_cr
#fig = plt.figure()
#ax = fig.add_subplot(111)
# Calculate Deflection Maps
alpha1, alpha2, mu_map, phi, detA, lambda_t = lt.cal_lensing_signals(
kappa, FOV_arc, args["ncells"])
# Calculate Einstein Radii in [arcsec]
Ncrit, curve_crit, curve_crit_tan, Rein = lt.einstein_radii(
lp1, lp2, detA, lambda_t, lens['zl'], cosmo, ax, 'med')
#if Rein == 0. or math.isnan(Rein):
# print('!!! Rein is 0. or NaN')
# continue
# Calculate Time-Delay and Magnification
n_imgs, delta_t, mu, theta = lt.timedelay_magnification(
mu_map, phi, dsx_arc, args["ncells"],
lp1, lp2, alpha1, alpha2, beta,
zs[ss], lens['zl'], cosmo)
if n_imgs > 1:
# Tree Branch 2
s_srcID.append(Src_ID[ss])
s_zs.append(zs[ss])
s_beta.append(beta)
#s_lensplane.append([lp1, lp2])
s_detA.append(detA)
s_tancritcurves.append(curve_crit_tan)
s_einsteinradius.append(Rein) #[arcsec]
# Tree Branch 3
s_theta.append(theta)
s_deltat.append(delta_t)
s_mu.append(mu)
check_for_sources = 1
#print(' -> %d multiple lensed images' % (n_imgs))
if check_for_sources == 1:
# Tree Branch 1
l_HFID.append(int(lens['HF_ID']))
l_haloID.append(int(lcdf.index.values[ll]))
l_snapnum.append(int(lens['snapnum']))
l_zl.append(lens['zl'])
#l_haloposbox.append(HaloPosBox[ll])
# Tree Branch 2
l_srcID.append(s_srcID)
l_zs.append(s_zs)
l_srcbeta.append(s_beta)
#l_lensplane.append(s_lensplane)
l_detA.append(s_detA)
l_tancritcurves.append(s_tancritcurves)
l_einsteinradius.append(s_einsteinradius)
# Tree Branch 3
l_srctheta.append(s_theta)
l_deltat.append(s_deltat)
l_mu.append(s_mu)
srclistinit()
check_for_sources = 0
print('Save data of lens %d' % ll)
########## Save to File ########
tree = plant_Tree()
## Tree Branches of Node 1 : Lenses
#tree['HF_ID'] = l_HFID
#tree['snapnum'] = args["snapnum"]
#tree['zl'] = redshift
#tree['zs'] = zs
## Tree Branches of Node 1 : Sources
#tree['Sources']['beta'] = l_srcbeta
#tree['Sources']['TCC'] = l_tancritcurves
#tree['Sources']['Rein'] = l_einsteinradius
#for imgs in range(len(l_mu)):
# # Tree Branches of Node 2 : Multiple Images
# tree['Sources']['theta'][imgs] = l_srctheta[imgs]
# tree['Sources']['delta_t'][imgs] = l_deltat[imgs]
# tree['Sources']['mu'][imgs] = l_mu[imgs]
# Tree Branches of Node 1 : Lenses
tree['LC_ID'] = l_haloID
tree['HF_ID'] = l_HFID
tree['snapnum'] = l_snapnum
tree['zl'] = l_zl
#tree['HaloPosBox'] = l_haloposbox
for sid in range(len(l_haloID)):
# Tree Branches of Node 2 : Sources
tree['Sources']['Src_ID'][sid] = l_srcID[sid]
tree['Sources']['zs'][sid] = l_zs[sid]
tree['Sources']['beta'][sid] = l_srcbeta[sid]
#tree['Sources']['LP'][sid] = l_lensplane[sid]
tree['Sources']['detA'][sid] = l_detA[sid]
tree['Sources']['TCC'][sid] = l_tancritcurves[sid]
tree['Sources']['Rein'][sid] = l_einsteinradius[sid]
for imgs in range(len(l_srcID[sid])):
# Tree Branches of Node 3 : Multiple Images
tree['Sources']['theta'][sid][imgs] = l_srctheta[sid][imgs]
tree['Sources']['delta_t'][sid][imgs] = l_deltat[sid][imgs]
tree['Sources']['mu'][sid][imgs] = l_mu[sid][imgs]
label = args["simdir"].split('/')[-2].split('_')[2]
filename = args["outbase"]+'LM_%s.pickle' % (label)
filed = open(filename, 'wb')
pickle.dump(tree, filed)
filed.close()
plt.close(fig)
def generate_lens_map(lenses, cpunum, LC, Halo_HF_ID, Halo_ID, Halo_z, Rvir,
snapnum, snapfile, h, scale, Ncells, HQ_dir, sim, sim_phy,
sim_name, HaloPosBox, HaloVel, cosmo, results_per_cpu):
"""
Input:
ll: halo array indexing
LC: Light-cone dictionary
Halo_ID: ID of Halo
Halo_z: redshift of Halo
Rvir: virial radius in [Mpc]
previous_snapnum:
snapnum
Output:
"""
print('Process %s started' % mp.current_process().name)
first_lens = lenses[0]
previous_snapnum = snapnum[first_lens]
memtrack = 0
file = open('F6_test_'+str(mp.current_process().name)+'.txt','w')
lenslistinit()
# Run through lenses
for ll in range(first_lens, lenses[-1]):
zs, Src_ID, SrcPosSky = source_selection(LC['Src_ID'], LC['Src_z'],
LC['SrcPosSky'], Halo_ID[ll])
zl = Halo_z[ll]
Lbox = Rvir[ll]*0.3*u.Mpc
FOV = Lbox.to_value('Mpc')
# converting box size and pixels size from ang. diam. dist. to arcsec
FOV_arc = (FOV/cf.Da(zl, cosmo)*u.rad).to_value('arcsec') #[arcsec] box size
dsx_arc = FOV_arc/Ncells #[arcsec] pixel size
# initialize the coordinates of grids (light rays on lens plan)
lp1, lp2 = cf.make_r_coor(FOV_arc, Ncells) #[arcsec]
# Only load new particle data if lens is at another snapshot
if (previous_snapnum != snapnum[ll]) or (ll == first_lens):
snap = snapfile % (snapnum[ll], snapnum[ll])
# 0 Gas, 1 DM, 4 Star[Star=+time & Wind=-time], 5 BH
DM_pos = readsnap.read_block(snap, 'POS ', parttype=1)*scale #[Mpc]
DM_mass = readsnap.read_block(snap, 'MASS', parttype=1)*1e10/h
Gas_pos = readsnap.read_block(snap, 'POS ', parttype=0)*scale #[Mpc]
Gas_mass = readsnap.read_block(snap, 'MASS', parttype=0)*1e10/h
Star_pos = readsnap.read_block(snap, 'POS ', parttype=4)*scale #[Mpc]
Star_age = readsnap.read_block(snap, 'AGE ', parttype=4)
Star_mass = readsnap.read_block(snap, 'MASS', parttype=4)
Star_pos = Star_pos[Star_age >= 0]
Star_mass = Star_mass[Star_age >= 0]*1e10/h
del Star_age
BH_pos = readsnap.read_block(snap, 'POS ', parttype=5)*scale
BH_mass = readsnap.read_block(snap, 'MASS', parttype=5)*1e10/h
file.write(str(mp.current_process().name) + 'read particles \n')
previous_snapnum = snapnum[ll]
DM_sigma, xs, ys = projected_surface_density(ll,
DM_pos, #[Mpc]
DM_mass,
HaloPosBox[ll],
fov=FOV, #[Mpc]
bins=Ncells,
smooth=False,
smooth_fac=0.5,
neighbour_no=32)
Gas_sigma, xs, ys = projected_surface_density(ll, Gas_pos, #*a/h,
Gas_mass,
HaloPosBox[ll], #*a/h,
fov=FOV,
bins=Ncells,
smooth=False,
smooth_fac=0.5,
neighbour_no=32)
Star_sigma, xs, ys = projected_surface_density(ll, Star_pos, #*a/h,
Star_mass,
HaloPosBox[ll], #*a/h,
fov=FOV,
bins=Ncells,
smooth=False,
smooth_fac=0.5,
neighbour_no=8)
file.write(str(mp.current_process().name) + 'created surfce densities \n')
# point sources need to be smoothed by > 1 pixel to avoid artefacts
tot_sigma = DM_sigma + Gas_sigma + Star_sigma
srclistinit()
# Run through Sources
check_for_sources = 0
for ss in range(len(Src_ID)):
# Calculate critical surface density
sigma_cr = sigma_crit(zl, zs[ss], cosmo).to_value('Msun Mpc-2')
kappa = tot_sigma/sigma_cr
fig = plt.figure()
ax = fig.add_subplot(111)
kappa = gaussian_filter(kappa, sigma=3)
# Calculate Deflection Maps
alpha1, alpha2, mu_map, phi, detA, lambda_t = cal_lensing_signals(kappa,
FOV_arc,
Ncells)
file.write(str(ll)+'; '+str(ss)+'; '+str(mp.current_process().name) + ' lensing signals \n')
# Calculate Einstein Radii
Ncrit, curve_crit, curve_crit_tan, Rein = einstein_radii(lp1, lp2,
detA,
lambda_t,
zl, cosmo,
ax, 'med')
file.write(str(mp.current_process().name)+' Rein calc \n')
# Calculate Time-Delay and Magnification
n_imgs, delta_t, mu, theta, beta = timedelay_magnification(mu_map, phi,
dsx_arc,
Ncells, lp1, lp2,
alpha1, alpha2,
SrcPosSky[ss],
zs[ss],
zl, cosmo)
file.write(str(mp.current_process().name)+' time delay \n')
if n_imgs > 1:
file.write(str(mp.current_process().name)+' adding multi source --------------- \n')
print('%s, %s, %s' % (str(mp.current_process().name),
str(ll),
str(ss)))
# Tree Branch 2
s_srcID.append(Src_ID[ss])
print('-- zs[ss]', zs)
s_zs.append(zs[ss])
s_beta.append(beta)
#s_lensplane.append([lp1, lp2])
s_detA.append(detA)
s_tancritcurves.append(curve_crit_tan)
s_einsteinradius.append(Rein)
# Tree Branch 3
s_theta.append(theta)
s_deltat.append(delta_t)
s_mu.append(mu)
check_for_sources = 1
if check_for_sources == 1:
# Tree Branch 1
l_HFID.append(Halo_HF_ID[ll])
l_haloID.append(Halo_ID[ll])
l_snapnum.append(int(snapnum[ll]))
l_zl.append(Halo_z[ll])
l_haloposbox.append(HaloPosBox[ll])
l_halovel.append(HaloVel[ll])
# Tree Branch 2
l_srcID.append(s_srcID)
l_zs.append(s_zs)
l_srcbeta.append(s_beta)
#l_lensplane.append(s_lensplane)
l_detA.append(s_detA)
l_tancritcurves.append(s_tancritcurves)
l_einsteinradius.append(s_einsteinradius)
# Tree Branch 3
l_srctheta.append(s_theta)
l_deltat.append(s_deltat)
l_mu.append(s_mu)
#memuseout = (sys.getsizeof(l_HFID) + sys.getsizeof(l_haloID) + \
# sys.getsizeof(l_snapnum) + sys.getsizeof(l_zl) + \
# sys.getsizeof(l_haloposbox) + sys.getsizeof(l_halovel) + \
# sys.getsizeof(l_srcID) + sys.getsizeof(l_zs) + \
# sys.getsizeof(l_srcbeta) + sys.getsizeof(l_detA) + \
# sys.getsizeof(l_tancritcurves) + sys.getsizeof(l_einsteinradius) + \
# sys.getsizeof(l_srctheta) + sys.getsizeof(l_deltat) + \
# sys.getsizeof(l_mu))/1024**3 #[GB]
memusetot = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024**3 #[GB]
print('::::::::::::: Tot. Memory Size [GB]: ', memusetot)
if memusetot > 5:
########## Save to File ########
print('save file because it is too big')
tree = plant_Tree()
#tree = grow_Tree()
# Tree Branches of Node 1 : Lenses
tree['Halo_ID'] = l_haloID
tree['HF_ID'] = l_HFID
tree['snapnum'] = l_snapnum
tree['zl'] = l_zl
tree['HaloPosBox'] = l_haloposbox
tree['HaloVel'] = l_halovel
for sid in range(len(l_haloID)):
# Tree Branches of Node 2 : Sources
tree['Sources']['Src_ID'][sid] = l_srcID[sid]
tree['Sources']['zs'][sid] = l_zs[sid]
tree['Sources']['beta'][sid] = l_srcbeta[sid]
#tree['Sources']['LP'][sid] = l_lensplane[sid]
tree['Sources']['detA'][sid] = l_detA[sid]
tree['Sources']['TCC'][sid] = l_tancritcurves[sid]
tree['Sources']['Rein'][sid] = l_einsteinradius[sid]
for imgs in range(len(l_srcID[sid])):
# Tree Branches of Node 3 : Multiple Images
tree['Sources']['theta'][sid][imgs] = l_srctheta[sid][imgs]
tree['Sources']['delta_t'][sid][imgs] = l_deltat[sid][imgs]
tree['Sources']['mu'][sid][imgs] = l_mu[sid][imgs]
lm_dir = HQ_dir+'LensingMap/'+sim_phy[sim]+'/'+sim_name[sim]+'/'
ensure_dir(lm_dir)
filename = lm_dir+'LM_'+mp.current_process().name+'_' + \
str(memtrack)+'.pickle'
filed = open(filename, 'wb')
pickle.dump(tree, filed)
filed.close()
plt.close(fig)
memtrack += 1
lenslistinit()
srclistinit()
########## Save to File ########
tree = plant_Tree()
#tree = grow_Tree()
# Tree Branches of Node 1 : Lenses
tree['Halo_ID'] = l_haloID
tree['HF_ID'] = l_HFID
tree['snapnum'] = l_snapnum
tree['zl'] = l_zl
tree['HaloPosBox'] = l_haloposbox
tree['HaloVel'] = l_halovel
for sid in range(len(l_haloID)):
# Tree Branches of Node 2 : Sources
tree['Sources']['Src_ID'][sid] = l_srcID[sid]
tree['Sources']['zs'][sid] = l_zs[sid]
tree['Sources']['beta'][sid] = l_srcbeta[sid]
#tree['Sources']['LP'][sid] = l_lensplane[sid]
tree['Sources']['detA'][sid] = l_detA[sid]
tree['Sources']['TCC'][sid] = l_tancritcurves[sid]
tree['Sources']['Rein'][sid] = l_einsteinradius[sid]
for imgs in range(len(l_srcID[sid])):
# Tree Branches of Node 3 : Multiple Images
tree['Sources']['theta'][sid][imgs] = l_srctheta[sid][imgs]
tree['Sources']['delta_t'][sid][imgs] = l_deltat[sid][imgs]
tree['Sources']['mu'][sid][imgs] = l_mu[sid][imgs]
file.close()
lm_dir = HQ_dir+'LensingMap/'+sim_phy[sim]+'/'+sim_name[sim]+'/'
ensure_dir(lm_dir)
filename = lm_dir+'LM_'+mp.current_process().name+'_' + \
str(memtrack)+'.pickle'
filed = open(filename, 'wb')
pickle.dump(tree, filed)
filed.close()
plt.close(fig)
def lensing_signal():
# Get command line arguments
args = {}
#args["snapnum"] = int(sys.argv[1])
#args["ncells"] = int(sys.argv[2])
#args["simdir"] = sys.argv[3]
#args["dmdir"] = sys.argv[4]
#args["outbase"] = sys.argv[5]
args["snapnum"] = 40
args["ncells"] = 1024
args["simdir"] = "/cosma6/data/dp004/dc-arno1/SZ_project/full_physics/L62_N512_GR_kpc/"
args["dmdir"] = "/cosma5/data/dp004/dc-beck3/StrongLensing/DensityMap/full_physics/L62_N512_GR_kpc/z_40/"
args["outbase"] = "/cosma5/data/dp004/dc-beck3/StrongLensing/LensingMap/full_physics/Rockstar/L62_N512_GR_kpc/Box/"
# Organize devision of Sub-&Halos over Processes on Proc. 0
s = read_hdf5.snapshot(args["snapnum"], args["simdir"])
fname = glob.glob(args["dmdir"]+'*.h5')
ffname = []
for ff in range(len(fname)):
if (os.path.getsize(fname[ff])/(1024*1024.0)) < 1:
fname[ff] = 0
else:
ffname.append(fname[ff])
ffname = np.asarray(ffname)
print(ffname)
# Cosmological Parameters
cosmo = LambdaCDM(H0=s.header.hubble*100,
Om0=s.header.omega_m,
Ode0=s.header.omega_l)
redshift = s.header.redshift
# Calculate critical surface density
zl = redshift
zs = 2.
sigma_cr = sigma_crit(zl, zs, cosmo).to_value('Msun Mpc-2')
lenslistinit(); srclistinit()
# Run through files
for ff in range(len(ffname))[:]:
print('Reading File: %s' % (fname[ff].split('/')[-2:]))
dmf = h5py.File(ffname[ff], 'r')
print(len(dmf['subhalo_id'][:]))
# Run through lenses
for ll in range(len(dmf['subhalo_id']))[1:]:
print('Works on lens %d with ID %d' % \
(ll, dmf['subhalo_id'][ll]))
# convert. box size and pixels size from ang. diam. dist. to arcsec
FOV_arc = (dmf['fov_width'][ll]/cf.Da(zl, cosmo)*u.rad).to_value('arcsec') #[arcsec] box size
dsx_arc = FOV_arc/args["ncells"] #[arcsec] pixel size
# initialize the coordinates of grids (light rays on lens plan)
#lp1, lp2 = cf.make_r_coor(FOV_arc, args["ncells"]) #[arcsec]
lpv = np.linspace(-(FOV_arc-dsx_arc)/2, (FOV_arc-dsx_arc)/2, args["ncells"])
lp1, lp2 = np.meshgrid(lpv, lpv) #[arcsec]
# Calculate convergence map
kappa = dmf['density_map'][ll]/sigma_cr
#print('The Kappa place has a max of %f and min of %f' %
# (np.max(kappa), np.min(kappa)))
fig = plt.figure()
ax = fig.add_subplot(111)
# Calculate Deflection Maps
alpha1, alpha2, mu_map, phi, detA, lambda_t = cal_lensing_signals(kappa,
FOV_arc,
args["ncells"])
# Calculate Einstein Radii
print('finish cal_lensing_signals')
Ncrit, curve_crit, curve_crit_tan, Rein = einstein_radii(lp1, lp2,
detA,
lambda_t,
zl, cosmo,
ax, 'med')
print('finish einstein_radii')
# Calculate Time-Delay and Magnification
snia_pos = np.array([0, 0, 0])
n_imgs, delta_t, mu, theta, beta = timedelay_magnification(
mu_map, phi, dsx_arc, args["ncells"],
lp1, lp2, alpha1, alpha2, snia_pos, zs, zl, cosmo)
print('finish timedelay_magnification')
if n_imgs > 1:
print(dmf['subhalo_id'][ll])
print('11111')
l_HFID.append(int(dmf['subhalo_id'][ll]))
print('22222')
print('timedelay_magnification', n_imgs)
# Tree Branch 1
#l_HFID.append(int(dmf['subhalo_id'][ll]))
# Tree Branch 2
l_srcbeta.append(beta)
l_tancritcurves.append(curve_crit_tan)
l_einsteinradius.append(Rein)
# Tree Branch 3
l_srctheta.append(theta)
l_deltat.append(delta_t)
l_mu.append(mu)
########## Save to File ########
print('Plant tree of %d lenses' % (len(l_HFID)))
tree = plant_Tree()
# Tree Branches of Node 1 : Lenses
tree['HF_ID'] = l_HFID
tree['snapnum'] = args["snapnum"]
tree['zl'] = redshift
tree['zs'] = 2.
# Tree Branches of Node 1 : Sources
tree['Sources']['beta'] = l_srcbeta
tree['Sources']['TCC'] = l_tancritcurves
tree['Sources']['Rein'] = l_einsteinradius
for imgs in range(len(l_mu)):
# Tree Branches of Node 2 : Multiple Images
tree['Sources']['theta'][imgs] = l_srctheta[imgs]
tree['Sources']['delta_t'][imgs] = l_deltat[imgs]
tree['Sources']['mu'][imgs] = l_mu[imgs]
print('tree created')
label = args["simdir"].split('/')[-2].split('_')[2]
filename = args["outbase"]+'LM_%s_z%d.pickle' % (label, args["snapnum"])
filed = open(filename, 'wb')
pickle.dump(tree, filed)
filed.close()
plt.close(fig)
print('plt close')
def raytrace_grid_maps_for_zs(self, ZS=None):
"""
Performs ray tracing at the source planes at redshift ZS. Note: the output files will be closed after
this function executes, so if needed to be called again, this object will need to be re-initialized.
Parameters
----------
ZS : float array
redshifts at which to perform ray tracing on the grid points. If None, then use the
higher redshift edge of each lens plane in the grid maps as a source plane. Defaults
to None.
max_planes : int
maximum number of planes about the halo plane to include in the ray tracing (e.g. if
max_planes = 4, then include at most four foreground planes and four background planes
in the calculation. Sources will still be placed to the full depth of the cutout, but
will be separated by empty space beyond the nine total included selected planes.
"""
self.print(
'\n ---------- creating ray trace maps for halo {} ---------- '.
format(self.inp.halo_id))
assert self.z_lens_planes is not None, "read_grid_maps_zs0() must be called before ray-tracing"
ZS0 = self.inp.zs0
ncc = self.inp.nnn
if (ZS is None):
ZS = self.z_lens_planes
# loop over source redshifts
for i in range(len(ZS)):
# Rescale Lens Data (zs0->zs)
zs = ZS[i]
zl_array = self.zmedian_lens_planes[self.zmedian_lens_planes <= (
zs)]
nzlp = len(zl_array)
alpha1_array = np.zeros((nzlp, ncc, ncc))
alpha2_array = np.zeros((nzlp, ncc, ncc))
kappa0_array = np.zeros((nzlp, ncc, ncc))
shear1_array = np.zeros((nzlp, ncc, ncc))
shear2_array = np.zeros((nzlp, ncc, ncc))
for j in range(nzlp):
rescale = cf.Da(ZS0) / cf.Da2(zl_array[j], ZS0) * cf.Da2(
zl_array[j], zs) / cf.Da(zs)
kappa0_array[j] = self.kappa_zs0[j] * rescale
alpha1_array[j] = self.alpha1_zs0[j] * rescale
alpha2_array[j] = self.alpha2_zs0[j] * rescale
shear1_array[j] = self.shear1_zs0[j] * rescale
shear2_array[j] = self.shear2_zs0[j] * rescale
# Ray-tracing
self.print(
'-------- ray tracing at source plane {:.3f} --------'.format(
zs))
af1, af2, kf0, sf1, sf2 = self._ray_tracing_all(
alpha1_array, alpha2_array, kappa0_array, shear1_array,
shear2_array, zl_array, zs)
self.print('max values:')
self.print("kf0 = {}".format(np.max(kf0)))
self.print("af1 = {}".format(np.max(af1)))
self.print("af2 = {}".format(np.max(af2)))
self.print("sf1 = {}".format(np.max(sf1)))
self.print("sf2 = {}".format(np.max(sf2)))
# Save Outputs
zs_group = 'plane{}'.format(i)
self.out_file.create_group(zs_group)
self.out_file[zs_group]['zs'] = np.atleast_1d(zs).astype('float32')
self.out_file[zs_group]['kappa0'] = kf0.astype('float32')
self.out_file[zs_group]['alpha1'] = af1.astype('float32')
self.out_file[zs_group]['alpha2'] = af2.astype('float32')
self.out_file[zs_group]['shear1'] = sf1.astype('float32')
self.out_file[zs_group]['shear2'] = sf2.astype('float32')
# debugging; remove this later
self.out_file[zs_group]['kappa0_array'] = kappa0_array
self.out_file[zs_group]['alpha1_array'] = alpha1_array
self.out_file[zs_group]['alpha2_array'] = alpha2_array
self.out_file[zs_group]['shear1_array'] = shear1_array
self.out_file[zs_group]['shear2_array'] = shear2_array
self.out_file.close()
self.print('done')
def _rec_read_xj(self, alpha1_array, alpha2_array, zln, zs, n):
xx1 = self.inp.xi1
xx2 = self.inp.xi2
dsi = self.inp.dsx_arc
nx1 = self.inp.nnn
nx2 = self.inp.nnn
if n == 0:
return xx1 * 0.0, xx2 * 0.0
if n == 1:
xx1.astype('double').tofile(self.inp.xj_path + str(n - 1) +
"_xj1.bin")
xx2.astype('double').tofile(self.inp.xj_path + str(n - 1) +
"_xj2.bin")
return xx1, xx2
if n == 2:
z2 = zln[1]
z1 = zln[0]
z0 = 0
x01 = xx1 * 0.0
x02 = xx2 * 0.0
try:
x11 = np.fromfile(self.inp.xj_path + str(n - 2) + "_xj1.bin",
dtype='double').reshape((nx1, nx2))
x12 = np.fromfile(self.inp.xj_path + str(n - 2) + "_xj2.bin",
dtype='double').reshape((nx1, nx2))
except:
self.print(
"No {} files, recalculate XJ...".format(self.inp.xj_path +
str(n - 2) +
"_xj1.bin"))
x11 = xx1
x12 = xx2
aim11 = alpha1_array[n - 2]
aim12 = alpha2_array[n - 2]
ahm11 = cf.ai_to_ah(aim11, z1, zs)
ahm12 = cf.ai_to_ah(aim12, z1, zs)
bij = cf.Da(z1) * cf.Da2(z0, z2) / (cf.Da(z2) * cf.Da2(z0, z1))
x21 = x11 * bij - (bij - 1) * x01 - ahm11 * cf.Da2(z1,
z2) / cf.Da(z2)
x22 = x12 * bij - (bij - 1) * x02 - ahm12 * cf.Da2(z1,
z2) / cf.Da(z2)
x21.astype('double').tofile(self.inp.xj_path + str(n - 1) +
"_xj1.bin")
x22.astype('double').tofile(self.inp.xj_path + str(n - 1) +
"_xj2.bin")
return x21, x22
if n > 2:
zi = zln[n - 1]
zim1 = zln[n - 2]
zim2 = zln[n - 3]
try:
xjm21 = np.fromfile(self.inp.xj_path + str(n - 3) + "_xj1.bin",
dtype='double').reshape((nx1, nx2))
xjm22 = np.fromfile(self.inp.xj_path + str(n - 3) + "_xj2.bin",
dtype='double').reshape((nx1, nx2))
except:
self.print(
"No {} files, recalculate XJ...".format(self.inp.xj_path +
str(n - 3) +
"_xj1.bin"))
xjm21, xjm22 = self._rec_read_xj(alpha1_array, alpha2_array,
zln, zs, n - 2)
try:
xjm11 = np.fromfile(self.inp.xj_path + str(n - 2) + "_xj1.bin",
dtype='double').reshape((nx1, nx2))
xjm12 = np.fromfile(self.inp.xj_path + str(n - 2) + "_xj2.bin",
dtype='double').reshape((nx1, nx2))
except:
self.print(
"No {} files, recalculate XJ...".format(self.inp.xj_path +
str(n - 2) +
"_xj1.bin"))
xjm11, xjm12 = self._rec_read_xj(alpha1_array, alpha2_array,
zln, zs, n - 1)
aijm11 = cf.call_inverse_cic(alpha1_array[n - 2], 0.0, 0.0, xjm11,
xjm12, dsi)
aijm12 = cf.call_inverse_cic(alpha2_array[n - 2], 0.0, 0.0, xjm11,
xjm12, dsi)
ahjm11 = cf.ai_to_ah(aijm11, zim1, zs)
ahjm12 = cf.ai_to_ah(aijm12, zim1, zs)
bij = cf.Da(zim1) * cf.Da2(zim2, zi) / cf.Da(zi) / cf.Da2(
zim2, zim1)
xj1 = xjm11 * bij - (bij - 1) * xjm21 - ahjm11 * cf.Da2(
zim1, zi) / cf.Da(zi)
xj2 = xjm12 * bij - (bij - 1) * xjm22 - ahjm12 * cf.Da2(
zim1, zi) / cf.Da(zi)
xj1.astype('double').tofile(self.inp.xj_path + str(n - 1) +
"_xj1.bin")
xj2.astype('double').tofile(self.inp.xj_path + str(n - 1) +
"_xj2.bin")
return xj1, xj2