· 8 years ago · Jun 09, 2018, 05:12 AM
1import os
2import sys
3import warnings
4import time
5import yt
6import healpy
7from p3lum import *
8from math import sqrt
9import h5py as h5
10import numpy as np
11import yt.utilities.cosmology as co
12from yt.units import mass_hydrogen
13from yt.utilities.parallel_tools.parallel_analysis_interface \
14 import communication_system
15
16########################################################################
17yt.enable_plugins()
18yt.enable_parallelism()
19top_comm = communication_system.communicators[-1]
20########################################################################
21ninner = 4
22nouter = top_comm.size / ninner
23dynamic = nouter > 2
24# -1 dynamic process scheduler
25nouter = (top_comm.size-dynamic) / ninner
26if yt.is_root(): print nouter, ninner, dynamic
27if ninner * nouter + dynamic != top_comm.size and yt.is_root():
28 warnings.warn("Number of MPI processes = %d is indivisible by outer loop parallelism (%d)" % \
29 (top_comm.size, nouter+1))
30########################################################################
31nhalos = -1
32start_halo = 0
33healpix_level = 2
34overwrite_fesc = False
35########################################################################
36
37data_dir = 'fesc-data'
38if yt.is_root():
39 if not os.path.exists('fesc-data'):
40 os.mkdir(data_dir)
41
42fn = sys.argv[-1]
43if not os.path.exists(fn):
44 raise RuntimeError("Dataset %s not found" % (fn))
45ds = yt.load(fn)
46ds.add_particle_filter('p2')
47ds.add_particle_filter('p2_young')
48ds.add_particle_filter('p3_living')
49
50omega = dict(b=0.0449, m=ds.parameters['CosmologyOmegaMatterNow'],
51 l=ds.parameters['CosmologyOmegaLambdaNow'])
52dm_factor = (omega['m'] - omega['b']) / omega['m']
53rvir_factor = dm_factor**(-1./3)
54
55halo_file = "%s/p%s_fesc_halos.h5" % (data_dir, ds)
56ascii_file = "%s/p%s_fesc_halos.txt" % (data_dir, ds)
57if yt.is_root():
58 fpa = open(ascii_file, "w")
59 fpa.write("#%9s %10s %10s %10s %10s %10s %10s %10s\n" % \
60 ("HaloNum", "Time[Myr]", "Mtotal", "Mstar", "M_young", "f_esc", "f_star", "f_gas"))
61 fpa.write("#"*90 + "\n")
62 fpa.close()
63
64
65# Load halo catalog
66hfn = "rockstar_halos/halos_%s.0.bin" % (ds.directory)
67if not os.path.exists(hfn):
68 raise RuntimeError("Halo catalog %s not found" % (hfn))
69hds = yt.load(hfn)
70halos = hds.all_data()
71msort = halos['particle_mass'].argsort().v.astype('int')[::-1]
72if nhalos < 0:
73 nhalos = msort.size
74
75# Get ionizing parameters
76nion_mass = ds.parameters["StarClusterIonizingLuminosity"]
77helium_ion = ds.parameters["StarClusterHeliumIonization"]
78
79top_comm = communication_system.communicators[-1]
80def is_group_root():
81 return communication_system.communicators[-1].rank == 0
82
83top_comm.barrier()
84hfile_exists = os.path.exists(halo_file)
85top_comm.barrier()
86if yt.is_root(): print 'Number of halos = ', nhalos
87
88def open_file(halo_file, first_open=False):
89 hh = None
90 no_overwrite = (hfile_exists and not overwrite_fesc) or (not first_open)
91 if is_group_root():
92 if no_overwrite:
93 hh = h5.File(halo_file, "a")
94 else:
95 hh = h5.File(halo_file, "w")
96 else:
97 if no_overwrite:
98 hh = h5.File(halo_file, "r")
99 return hh
100def close_file(hh):
101 if hh != None:
102 hh.close()
103 return
104
105# Create empty file, if needed
106hf = open_file(halo_file, first_open=True)
107if hf != None:
108 finished_halos = hf.keys()
109else:
110 finished_halos = []
111close_file(hf)
112top_comm.barrier()
113
114first_time = True
115all_halonum = range(start_halo, nhalos)
116files = {}
117for _files, ihalo in yt.parallel_objects(all_halonum, njobs=nouter, dynamic=dynamic,
118 storage=files):
119 ih = msort[ihalo]
120 halo_name = "Halo%8.8d" % (ihalo)
121 this_comm = communication_system.communicators[-1]
122
123 if hfile_exists and not overwrite_fesc:
124 if halo_name in finished_halos:
125 print "%s already exists. Skipping." % (halo_name)
126 continue
127
128 # Create a set of points, corresponding to HEALPix pixel centers
129 # that are r200 away from the center of mass
130 nside = 2**healpix_level
131 npix = 12 * nside**2
132 pixels = healpy.pix2vec(nside, np.arange(npix), nest=True)
133 pixels = np.array(pixels).T
134 rvir = ds.quan(halos['virial_radius'][ih].in_cgs().v, 'cm') * rvir_factor
135 center0 = ds.arr(halos['particle_position'][ih].in_cgs().v, 'cm')
136 sp = ds.sphere(center0, rvir)
137 center = sp.quantities.center_of_mass()
138 pixels = rvir * pixels + center
139
140 # Star particle number and total mass
141 np3 = int(sp['p3_living', 'particle_ones'].sum())
142 np2 = int(sp['p2', 'particle_ones'].sum())
143 np2y = int(sp['p2_young', 'particle_ones'].sum())
144 Mstar2 = sp['p2', 'particle_mass'].sum()
145 Mstar2y = sp['p2_young', 'particle_mass'].sum()
146 Mstar3 = sp['p3_living', 'particle_mass'].sum()
147 Mstar = Mstar2 + Mstar3
148 Mgas, Mdm = sp.quantities.total_mass()
149 Mhalo = Mgas + Mdm
150 M_HI, M_HII = sp.quantities.total_quantity(['H_p0_mass', 'H_p1_mass'])
151
152 # Delete existing data, if requested
153 # if is_group_root():
154 # hf = open_file(halo_file)
155 # if overwrite_fesc and hfile_exists:
156 # for k in hf[halo_name].keys():
157 # if k not in ["center", "rvir", "mass"]:
158 # del hf[halo_name][k]
159 # # if halo_name not in hf.keys():
160 # # gf = hf.create_group(halo_name)
161 # # else:
162 # # gf = hf[halo_name]
163 # #gf.create_group("Pop2_Stars")
164 # #gf.create_group("Pop3_Stars")
165 # close_file(hf)
166
167 Ntot = 0.0
168 Nesc = 0.0
169 Nfesc = 100
170 fesc_all = np.zeros(Nfesc+1)
171 fesc_bins = np.linspace(0, 1, Nfesc+1)
172
173 # Preload HI_Density
174 trash = sp['H_p0_density']
175 if is_group_root():
176 all_sg2 = {}
177 all_sg3 = {}
178
179 ##################
180 # Pop III stars
181 ##################
182 for i in range(np3):
183 star_name = "star%8.8d" % (i)
184 starc = sp['p3_living', 'particle_position'][i]
185 mass = sp['p3_living', 'particle_mass'][i]
186 my_storage = {}
187 for sto, j in yt.parallel_objects(range(npix), njobs=ninner, storage = my_storage):
188 ray = ds.ray(starc, pixels[j,:])
189 ray_length = np.sqrt(np.sum((starc - pixels[j,:]))**2)
190 # Column density in ray segments
191 dN = ray_length * ray['dts'] * ray['H_p0_number_density']
192 sto.result_id = j
193 sto.result = dN.sum()
194 if is_group_root():
195 allN = []
196 for j, val in sorted(my_storage.items()):
197 allN.append(val)
198 allN = ds.arr(allN)
199 E, nion = pop3_lum(mass)
200 # Photo-ionization cross section (Verner 1996)
201 sigma_HI = ds.quan(5.475e-14 * (E[0] / 0.4298 - 1)**2 * \
202 (E[0] / 0.4298)**(-4.0185) * \
203 (1 + sqrt(E[0] / 14.13))**(-2.963), 'cm**2')
204 tau = (allN * sigma_HI).in_units('dimensionless').v
205 fesc = np.exp(-tau)
206 fesc_avg = fesc.mean()
207 ibins = np.searchsorted(fesc_bins, fesc)
208 for j in range(npix):
209 fesc_all[ibins[j]] += nion[0] / npix
210 Nesc += fesc_avg * nion[0]
211 Ntot += nion[0]
212 #hf = open_file(halo_file)
213 #g3 = hf[halo_name]["Pop3_Stars"]
214 #sg = g3.create_group(star_name)
215 sg = {}
216 sg['pos'] = starc.in_units('code_length')
217 sg['mass'] = mass.in_units('Msun')
218 sg['fesc'] = fesc
219 sg['NHI'] = allN.in_units('1/cm**2')
220 sg['fesc_avg'] = fesc_avg
221 sg['Nesc'] = fesc_avg * nion[0]
222 sg['Ntot'] = nion[0]
223 all_sg3[star_name] = sg
224 print "[H%d//PopIII] Ray traced %5d out of %5d" % (ihalo, i, np3)
225 #hf.close()
226
227 ########################
228 # Metal-enriched stars
229 ########################
230 for i in range(np2y):
231 star_name = "star%8.8d" % (i)
232 starc = sp['p2_young', 'particle_position'][i]
233 mass = sp['p2_young', 'particle_mass'][i]
234 my_storage = {}
235 for sto, j in yt.parallel_objects(range(npix), njobs=ninner, storage = my_storage):
236 ray = ds.ray(starc, pixels[j,:])
237 ray_length = np.sqrt(np.sum((starc - pixels[j,:]))**2)
238 # Column density in ray segments
239 dN = ray_length * ray['dts'] * ray['H_p0_number_density']
240 sto.result_id = j
241 sto.result = dN.sum()
242 if is_group_root():
243 allN = []
244 for j, val in sorted(my_storage.items()):
245 allN.append(val)
246 allN = ds.arr(allN)
247 E, nion = pop2_lum(mass, nion_mass, helium_ion=helium_ion)
248 # Photo-ionization cross section (Verner 1996)
249 sigma_HI = ds.quan(5.475e-14 * (E[0] / 0.4298 - 1)**2 * \
250 (E[0] / 0.4298)**(-4.0185) * \
251 (1 + sqrt(E[0] / 14.13))**(-2.963), 'cm**2')
252 tau = (allN * sigma_HI).in_units('dimensionless').v
253 fesc = np.exp(-tau)
254 fesc_avg = fesc.mean()
255 ibins = np.searchsorted(fesc_bins, fesc)
256 for j in range(npix):
257 fesc_all[ibins[j]] += nion[0] / npix
258 Nesc += fesc_avg * nion[0]
259 Ntot += nion[0]
260 #hf = open_file(halo_file)
261 #g2 = hf[halo_name]["Pop2_Stars"]
262 #sg = g2.create_group(star_name)
263 sg = {}
264 sg['pos'] = starc.in_units('code_length')
265 sg['mass'] = mass.in_units('Msun')
266 sg['fesc'] = fesc
267 sg['NHI'] = allN.in_units('1/cm**2')
268 sg['fesc_avg'] = fesc_avg
269 sg['Nesc'] = fesc_avg * nion[0]
270 sg['Ntot'] = nion[0]
271 all_sg2[star_name] = sg
272 #hf.close()
273 print "[H%d//PopII] Ray traced %5d out of %5d" % (ihalo, i, np2y)
274
275 # Write halo properties to file
276 if is_group_root():
277 gfn = halo_file + "_%04d" % (top_comm.rank)
278 hf = open_file(gfn)
279 if halo_name not in hf.keys():
280 gf = hf.create_group(halo_name)
281 else:
282 gf = hf[halo_name]
283 g2 = gf.create_group("Pop2_Stars")
284 for ks, sg in all_sg2.items():
285 s2 = g2.create_group(ks)
286 for k,v in sg.items():
287 s2[k] = v
288 g3 = gf.create_group("Pop3_Stars")
289 for ks, sg in all_sg3.items():
290 s3 = g3.create_group(ks)
291 for k,v in sg.items():
292 s3[k] = v
293 gf["NumberOfPop2Stars"] = np2
294 gf["NumberOfPop3Stars"] = np3
295 gf["Mstar_pop2"] = Mstar2.in_units('Msun')
296 gf["Mstar_pop2_young"] = Mstar2y.in_units('Msun')
297 gf["Mstar_pop3"] = Mstar3.in_units('Msun')
298 gf["Mstar"] = Mstar.in_units('Msun')
299 gf["Nesc"] = Nesc
300 gf["Ntot"] = Ntot
301 if Ntot > 0:
302 fesc = Nesc / Ntot
303 else:
304 fesc = -1
305 gf["fesc"] = fesc
306 gf["fesc_bins"] = fesc_bins
307 gf["fesc_hist"] = fesc_all
308 gf["rvir"] = rvir.in_units('kpc').v
309 gf["mass"] = Mhalo.in_units('Msun').v
310 gf["fgas"] = Mgas / Mhalo
311 gf["fstar"] = Mstar / Mhalo
312 gf["M_HI"] = M_HI.in_units('Msun').v
313 gf["M_HII"] = M_HII.in_units('Msun').v
314 gf["center"] = center
315 close_file(hf)
316 first_time = False
317 _files.result = gfn
318 _files.result_id = top_comm.rank
319
320 # Output to stdout
321 print "Halo %3.3d: Mtot = %s, Mstar = %s, fesc = %10.4f\n" % \
322 (ihalo, Mhalo.in_units('Msun'), Mstar.in_units('Msun'), fesc)
323
324 # Write to ASCII table
325 fpa = open(ascii_file+"_%04d" % (top_comm.rank), "a")
326 fpa.write("%10d %10.4g %10.4g %10.4g %10.4g %10.4f %10.4g %10.4g\n" % \
327 (ihalo, ds.current_time.in_units('Myr'),
328 Mhalo.in_units('Msun'), Mstar.in_units('Msun').v,
329 Mstar2y.in_units('Msun').v, fesc, Mstar/Mhalo, Mgas/Mhalo))
330 fpa.close()
331
332
333if yt.is_root():
334 allfiles = []
335 for j, val in sorted(files.items()):
336 if val not in allfiles:
337 allfiles.append(val)
338 mainfn = halo_file
339 mainfp = h5.File(mainfn, 'w')
340 for fn in allfiles:
341 fp = h5.File(fn, 'r')
342 for k in fp.keys():
343 fp.copy(k, mainfp)
344 fp.close()
345 os.remove(fn)
346 mainfp.close()