def main(path, ioutput, pickle_path): import seren3 import pickle, os from seren3.analysis.parallel import mpi snap = seren3.load_snapshot(path, ioutput) halos = snap.halos() halo_ix = None if mpi.host: halo_ix = halos.halo_ix(shuffle=True) dest = {} for i, sto in mpi.piter(halo_ix, storage=dest): h = halos[i] if len(h.g) > 1: mpi.msg("Working on halo %i \t %i" % (i, h.hid)) C = h.clumping_factor() sto.idx = h["id"] sto.result = {"hprops" : h.properties, "C" : C} if mpi.host: if pickle_path is None: pickle_path = "%s/pickle/" % path if os.path.isdir(pickle_path) is False: os.mkdir(pickle_path) fname = "%s/halo_clumping_factor_%05i.p" % (pickle_path, ioutput) pickle.dump( mpi.unpack(dest), open( fname, "wb" ) ) mpi.msg("Done")
def main(path, iout, finder='ctrees', pickle_path=None): import seren3 from seren3.analysis.parallel import mpi mpi.msg("Loading data") snap = seren3.load_snapshot(path, iout) snap.set_nproc(1) # disable multiprocessing halos = snap.halos(finder=finder) halo_ix = None if mpi.host: halo_ix = halos.halo_ix(shuffle=True) mpi.msg("Starting worker loop") dest = {} for i, sto in mpi.piter(halo_ix, storage=dest): h = halos[i] if len(h.s) > 0: mpi.msg("Working on halo %i \t %i" % (i, h.hid)) part_dset = h.p[["mass", "epoch", "id", "age"]].flatten() gas_dset = h.g["mass"].flatten() gas_mass = gas_dset["mass"].in_units(_MASS_UNIT) star_idx = stars(part_dset) part_mass = part_dset["mass"].in_units(_MASS_UNIT) idx_young_stars = np.where( part_dset["age"][star_idx].in_units("Myr") <= 10.) tot_mass = part_mass.sum() + gas_mass.sum() star_mass = part_mass[star_idx][idx_young_stars].sum() if (star_mass > 0): np_star = len(part_dset["mass"][star_idx][idx_young_stars]) np_dm = len(part_dset["mass"]) - np_star mpi.msg("%e \t %e \t %e" % (tot_mass, star_mass, star_mass / tot_mass)) sto.idx = h["id"] sto.result = { "tot_mass": tot_mass, "star_mass": star_mass, "np_dm": np_dm, "np_star": np_star } if mpi.host: import os if pickle_path is None: pickle_path = "%s/pickle/" % path if os.path.isdir(pickle_path) is False: os.mkdir(pickle_path) fname = "%s/abundance_%05i.p" % (pickle_path, iout) with open(fname, 'wb') as f: pickle.dump(mpi.unpack(dest), f)
def find_mmp(self, halo, back_to_iout=None): ''' Locates the most massive progenitor ''' from seren3 import load_snapshot if back_to_iout is None: back_to_iout = self.base.ioutput - 1 hid = halo.hid ioutputs = range(back_to_iout, self.base.ioutput)[::-1] last = self.base.ioutput for iout_prog in ioutputs: # Start with the previous snapshot, find the most massive progenitor and use that prog_snap = load_snapshot(self.base.path, iout_prog) prog_halos = prog_snap.halos(finder='ctrees') mmp_id = self._find_mmp(hid, prog_halos) if mmp_id is None: print 'Unable to fing progenitor in output %i.\nReturning last know progenitor (output %i)' % ( iout_prog, last) return hid, prog_halos else: hid = mmp_id last = iout_prog return hid, prog_halos
def main(path, iout, pickle_path): import seren3 from seren3.analysis.parallel import mpi import pickle, os mpi.msg("Loading snapshot: %05i" % iout) snapshot = seren3.load_snapshot(path, iout) snapshot.set_nproc(1) halos = snapshot.halos(finder="ctrees") nhalos = len(halos) dest = {} for i, sto in mpi.piter(range(nhalos), storage=dest): mpi.msg("%i / %i" % (i + 1, nhalos)) h = halos[i] sto.idx = h.hid sto.result = {"mw": mass_weighted_average(h.g), "Mvir": h["Mvir"]} mpi.msg("Done. Waiting.") if mpi.host: if pickle_path is None: pickle_path = "%s/pickle/" % snapshot.path if os.path.isdir(pickle_path) is False: os.mkdir(pickle_path) fname = "%s/trec_mw_halo_av_%05i.p" % (pickle_path, snapshot.ioutput) mpi.msg("Saving data to pickle file: %s" % fname) with open(fname, "wb") as f: pickle.dump(mpi.unpack(dest), f)
def main(path, iout, nbins=50): import seren3, pickle from seren3.analysis.parallel import mpi from seren3.exceptions import NoParticlesException mpi.msg("Loading data...") snap = seren3.load_snapshot(path, iout) snap.set_nproc(1) # disable multiprocessing halos = snap.halos() t = current_age(snap) agerange = [0, t] # halo_spheres = np.array( [ {'id' : h.hid, 'reg' : h.sphere, 'mvir' : h['mvir'].v} for h in halos ] ) halo_spheres = halos.mpi_spheres() dest = {} for h, sto in mpi.piter(halo_spheres, storage=dest): sphere = h["reg"] subsnap = snap[sphere] try: dset = {} dset["sSFR"] = subsnap.s["sSFR"].flatten(nbins=nbins, agerange=agerange) dset["mvir"] = h["mvir"] sto.idx = h["id"] sto.result = dset except NoParticlesException as e: mpi.msg(e.message) pickle.dump(dest, open("./integrate_sSFR_%05i.p" % iout, "wb"))
def add_mass_to_dbs(path, iout, pickle_path): import seren3 from seren3.utils import derived_utils from seren3.analysis.parallel import mpi from seren3.core.serensource import DerivedDataset import pickle, os mpi.msg("Loading snapshot") snap = seren3.load_snapshot(path, iout) snap.set_nproc(1) halos = snap.halos(finder="ctrees") db = load_db(snap.path, snap.ioutput) hids = np.array(db.keys()) star_Nion_d_fn = derived_utils.get_derived_field(snap.s, "Nion_d") nIons = snap.info_rt["nIons"] dest = {} for i, sto in mpi.piter(hids, storage=dest, print_stats=True): hid = int(i) h = halos.with_id(hid) star_dset = h.s[["age", "metal", "mass"]].flatten() star_mass = star_dset["mass"] star_age = star_dset["age"] star_metal = star_dset["metal"] dict_stars = {"age": star_age, "metal": star_metal, "mass": star_mass} dset_stars = DerivedDataset(snap.s, dict_stars) Nion_d_all_groups = snap.array(np.zeros(len(dset_stars["age"])), "s**-1 Msol**-1") for ii in range(nIons): Nion_d_now = star_Nion_d_fn(snap, dset_stars, group=ii + 1, dt=0.) Nion_d_all_groups += Nion_d_now sto.idx = hid sto.result = db[hid] sto.result["star_mass"] = star_mass sto.result["Nion_d_now"] = Nion_d_all_groups if mpi.host: if pickle_path is None: pickle_path = "%s/pickle/%s/" % (path, halos.finder.lower()) # pickle_path = "%s/" % path if os.path.isdir(pickle_path) is False: os.mkdir(pickle_path) unpacked_dest = mpi.unpack(dest) fesc_dict = {} for i in range(len(unpacked_dest)): fesc_dict[int(unpacked_dest[i].idx)] = unpacked_dest[i].result pickle.dump( fesc_dict, open( "%s/fesc_database_no_filt_denoise_%05i.p" % (pickle_path, iout), "wb"))
def get_rho_proj_list(halo_catalogue, h, back_to_aexp=0., **kwargs): ''' Iterate through progenitors and make projections ''' import seren3 import numpy as np from seren3.analysis.visualization import engines, operators from seren3.utils import camera_utils field = "rho" proj_list = [] C = h.base.C info = h.info length_fac = 2. camera = h.camera() camera.region_size *= length_fac camera.distance *= length_fac camera.far_cut_depth *= length_fac extent = camera_utils.extent(h.base, camera) unit_sd = (info["unit_density"] * info["unit_length"]).express(C.Msun * C.kpc**-2) # surface density unit kwargs["surf_qty"] = True for iout in [116, 115, 114, 113, 112, 111, 110, 109, 108, 107]: print iout tmp = seren3.load_snapshot(h.base.path, iout) eng = engines.SurfaceDensitySplatterEngine(tmp.g) # op = operators.MassWeightedOperator(field, info["unit_density"]) # eng = engines.CustomSplatterEngine(h.g, field, op, extra_fields=['rho']) proj = eng.process(camera, **kwargs) proj_list.append(proj) eng = engines.SurfaceDensitySplatterEngine(h.g) # op = operators.MassWeightedOperator(field, info["unit_density"]) # eng = engines.CustomSplatterEngine(h.g, field, op, extra_fields=['rho']) proj = eng.process(camera, **kwargs) proj_map = proj.map.T * unit_sd vmin, vmax = (np.log10(proj_map.min()), np.log10(proj_map.max())) proj_list.append(proj) for prog in halo_catalogue.iterate_progenitors(h, back_to_aexp=back_to_aexp): # camera.center = prog.subsnap.region.center info = prog.info # op = operators.MassWeightedOperator(field, info["unit_density"]) eng = engines.SurfaceDensitySplatterEngine(prog.g) # eng = engines.CustomSplatterEngine(prog.g, field, op, extra_fields=['rho']) proj_list.append(eng.process(camera, **kwargs)) return proj_list[::-1], extent, (vmin, vmax)
def main(path, iout, NSLOTS): import seren3 print path, iout snap = seren3.load_snapshot(path, iout) halos = snap.halos(finder='ahf') if not halos.can_load(NSLOTS=NSLOTS): halos.run() else: print "Halo catalogue already exists"
def main(path, iout): import seren3 snap = seren3.load_snapshot(path, iout) halos = snap.h h = halos.sort('mvir')[0] subsnap = h.subsnap center = subsnap.region.center pos = subsnap.g["pos"].flatten(center=center)
def main(path, iout, pickle_path): import seren3 from seren3.core.serensource import DerivedDataset from seren3.utils import derived_utils from seren3.analysis import outflows from seren3.analysis.parallel import mpi from seren3.exceptions import NoParticlesException import pickle, os mpi.msg("Loading snapshot") snap = seren3.load_snapshot(path, iout) snap.set_nproc(1) halos = None halos = snap.halos(finder="ctrees") halo_ix = None if mpi.host: halo_ix = halos.halo_ix(shuffle=True) dest = {} for i, sto in mpi.piter(halo_ix, storage=dest, print_stats=True): h = halos[i] if (len(h.s) > 0): sto.idx = int(h["id"]) dm_dset = h.d["mass"].flatten() gas_dset = h.g["mass"].flatten() star_dset = h.s["mass"].flatten() tot_mass = dm_dset["mass"].in_units("Msol h**-1").sum() + star_dset["mass"].in_units("Msol h**-1").sum()\ + gas_dset["mass"].in_units("Msol h**-1").sum() F, h_im = outflows.dm_by_dt(h.subsnap, filt=False, nside=2**3, denoise=True) sto.result = {"F" : F, "h_im" : h_im, "tot_mass" : tot_mass, \ "hprops" : h.properties} if mpi.host: if pickle_path is None: pickle_path = "%s/pickle/%s/" % (path, halos.finder.lower()) # pickle_path = "%s/" % path if os.path.isdir(pickle_path) is False: os.mkdir(pickle_path) unpacked_dest = mpi.unpack(dest) fesc_dict = {} for i in range(len(unpacked_dest)): fesc_dict[int(unpacked_dest[i].idx)] = unpacked_dest[i].result pickle.dump( fesc_dict, open( "%s/mass_flux_database_no_filt_quarter_rvir_denoise_%05i.p" % (pickle_path, iout), "wb"))
def main(path, iout, pickle_path): import seren3 from seren3.analysis.parallel import mpi from seren3.analysis.plots import histograms import pickle, os mpi.msg("Loading snapshot %05i" % iout) snapshot = seren3.load_snapshot(path, iout) halos = snapshot.halos(finder="ctrees") nhalos = len(halos) # Bracket min/max nH if mpi.host: mpi.msg("Bracketing nH...") min_nH = np.inf; max_nH = -np.inf for nH in snapshot.g["nH"]: if nH.min() < min_nH: min_nH = nH.min() if nH.max() > max_nH: max_nH = nH.max() data = {"min" : np.log10(min_nH), "max" : np.log10(max_nH)} mpi.msg("Done") else: data = None mpi.msg("bcast nH") data = mpi.comm.bcast(data, root=0) mpi.msg("Done. Processing halos...") min_nH = data["min"]; max_nH = data["max"] x_range = (min_nH, max_nH) mpi.msg("x_range (log): (%f, %f)" % (min_nH, max_nH)) snapshot.set_nproc(1) dest = {} for i, sto in mpi.piter(range(nhalos), storage=dest): h = halos[i] if len(h.g["nH"].f) > 0: P, C, bincenters, dx = histograms.pdf_cdf(h.g, "nH", density=True, \ cumulative=True, x_range=x_range) sto.idx = h.hid sto.result = {"P" : P, "C" : C, "bincenters" : bincenters, "dx" : dx} if mpi.host: if pickle_path is None: pickle_path = "%s/pickle/" % snapshot.path if os.path.isdir(pickle_path) is False: os.mkdir(pickle_path) fname = "%s/pdf_cdf_halo_%05i.p" % (pickle_path, snapshot.ioutput) mpi.msg("Saving data to pickle file: %s" % fname) with open(fname, "wb") as f: pickle.dump(mpi.unpack(dest), f)
def main(path, ioutput, txt_fname, out_fname=None, REGION_MODE=_REGION_MODE): ''' Reads the snapshot (to compute domain decomposition at this time) and outputs a set of idomains required. You can either specify ''' import csv, seren3 # Load the snapshot print "Loading snapshot" snapshot = seren3.load_snapshot(path, ioutput) domains = [] # Open the csv file and collect the desired regions with open(txt_fname, "r") as csvfile: reader = csv.reader(csvfile) print "Reading csv file" for row in reader: if (hasattr(row, "__iter__") and row[0].startswith("#")): continue else: x,y,z,l = [float(ri) for ri in row] pos = [x,y,z] reg = None if REGION_MODE == RegionMode.CUBE: reg = snapshot.get_cube(pos, l) elif REGION_MODE == RegionMode.SPHERE: reg = snapshot.get_sphere(pos, l) else: raise Exception("Unknown region mode: %s" % REGION_MODE) # Compute minimal number of domains to bound this region bbox = reg.get_bounding_box() idomains = snapshot.cpu_list(bbox) domains.extend(idomains) # Use a set to remove duplicates domains = sorted(list(set(domains))) # Converts to set then back to a list if (out_fname is not None): print "Writing output" with open(out_fname, "w") as f: for i in range(len(domains)): f.write("%i" % domains[i]) if (i < len(domains) - 1): f.write("\n") print "Done" # Return the list return domains
def main(path, ioutput, pickle_path, sed_type): import pickle, os import seren3 from seren3.analysis.parallel import mpi from seren3.utils.sed import io mpi.msg("Loading snapshot") snap = seren3.load_snapshot(path, ioutput) snap.set_nproc(1) halos = None halos = snap.halos(finder="ctrees") halo_ix = None if mpi.host: halo_ix = halos.halo_ix(shuffle=True) nml = snap.nml RT_PARAMS_KEY = nml.NML.RT_PARAMS # SED = io.read_seds() SED = io.interp_bpass_to_bc03(sed_type=sed_type) if mpi.host: print SED, sed_type dest = {} kwargs = {"lambda_A": _LAMBDA_A, "sed": SED} for i, sto in mpi.piter(halo_ix, storage=dest, print_stats=True): h = halos[i] if (len(h.s) > 0): dset = h.s[["luminosity", "age"]].flatten(**kwargs) sto.idx = int(h["id"]) sto.result = {"age" : dset["age"], "L" : dset["luminosity"], "lambda_A" : _LAMBDA_A, \ "hprops" : h.properties} if mpi.host: if pickle_path is None: pickle_path = "%s/pickle/%s/" % (path, halos.finder.lower()) # pickle_path = "%s/" % path if os.path.isdir(pickle_path) is False: os.mkdir(pickle_path) unpacked_dest = mpi.unpack(dest) lum_dict = {} for i in range(len(unpacked_dest)): lum_dict[int(unpacked_dest[i].idx)] = unpacked_dest[i].result pickle.dump( lum_dict, open( "%s/luminosity_lambdaA_%s_database_%05i.p" % (pickle_path, int(_LAMBDA_A), ioutput), "wb"))
def main(path, ioutput, pickle_path): import seren3 from seren3.analysis import stars from seren3.analysis.parallel import mpi snap = seren3.load_snapshot(path, ioutput) age_max = snap.array(0.9279320933559091, "Gyr") age_min = snap.array(0., "Gyr") agerange = [age_min, age_max] snap.set_nproc(1) halos = snap.halos() nhalos = len(halos) halo_ix = None if mpi.host: halo_ix = halos.halo_ix(shuffle=True) dest = {} for i, sto in mpi.piter(halo_ix, storage=dest, print_stats=False): h = halos[i] dset = h.s[["age", "mass"]].flatten() if (len(dset["age"]) > 0): star_mass = h.s["mass"].flatten()["mass"].in_units("Msol") sSFR, SFR, lookback_time, binsize = stars.sfr(h, dset=dset, ret_sSFR=True, agerange=agerange) # sSFR = SFR / star_mass.sum() rho_sfr = stars.gas_SFR_density(h) mpi.msg("%i %e %e" % (h["id"], SFR.max(), rho_sfr.max())) sto.idx = h["id"] sto.result = {"SFR" : SFR, "sSFR" : sSFR, "lookback-time" : lookback_time,\ "binsize" : binsize, "rho_sfr" : rho_sfr, "Mvir" : h["Mvir"]} if mpi.host: import os, pickle if pickle_path is None: pickle_path = "%s/pickle/" % path if os.path.isdir(pickle_path) is False: os.mkdir(pickle_path) fname = "%s/sfr_halos_%05i.p" % (pickle_path, ioutput) pickle.dump( mpi.unpack(dest), open( fname, "wb" ) )
def main(path, iout, finder, pickle_path): import seren3 from seren3.analysis.escape_fraction import fesc from seren3.analysis.parallel import mpi from seren3.exceptions import NoParticlesException import pickle, os, random mpi.msg("Loading snapshot") snap = seren3.load_snapshot(path, iout) snap.set_nproc(1) halos = None mpi.msg("Using halo finder: %s" % finder) halos = snap.halos(finder=finder) halo_ix = None if mpi.host: halo_ix = halos.halo_ix(shuffle=True) dest = {} for i, sto in mpi.piter(halo_ix, storage=dest, print_stats=True): h = halos[i] if (h["pid"] == -1 and len(h.s) > 0): sto.idx = h.hid mpi.msg("%i" % h.hid) try: pdset = h.p["mass"].flatten() gdset = h.g["mass"].flatten() tot_mass = pdset["mass"].in_units("Msol h**-1").sum()\ + gdset["mass"].in_units("Msol h**-1").sum() h_fesc = fesc(h.subsnap, nside=2**3, filt=True, do_multigroup=True) mpi.msg("%1.2e \t %1.2e" % (h["Mvir"], h_fesc)) sto.result = {"fesc" : h_fesc, "tot_mass" : tot_mass, "hprops" : h.properties} except NoParticlesException as e: mpi.msg(e.message) except IndexError: mpi.msg("Index error - skipping") if mpi.host: if pickle_path is None: pickle_path = "%s/pickle/%s/" % (path, halos.finder.lower()) # pickle_path = "%s/" % path if os.path.isdir(pickle_path) is False: os.mkdir(pickle_path) # pickle.dump( mpi.unpack(dest), open( "%s/fesc_multi_group_filt_%05i.p" % (pickle_path, iout), "wb" ) ) pickle.dump( mpi.unpack(dest), open( "%s/fesc_do_multigroup_filt_no_age_limit%05i.p" % (pickle_path, iout), "wb" ) )
def main(path, ioutput): import random import numpy as np import seren3 from seren3.analysis.parallel import mpi from seren3.exceptions import NoParticlesException from seren3.analysis.escape_fraction import time_integrated_fesc, integrate_fesc mpi.msg("Loading snapshot...") snap = seren3.load_snapshot(path, ioutput) snap.set_nproc(1) halos = snap.halos(finder="ctrees") halo_ix = np.array(range(len(halos))) random.shuffle(halo_ix) dest = {} for i, sto in mpi.piter(halo_ix, storage=dest): h = halos[i] mpi.msg(h.hid) try: tint_fesc_hist, I1, I2, lbtime = time_integrated_fesc( h, 0., return_data=True) fesc = I1 / I2 sto.idx = h.hid sto.result = {'tint_fesc_hist' : tint_fesc_hist, 'fesc' : fesc, 'I1' : I1, \ 'I2' : I2, 'lbtime' : lbtime, 'Mvir' : h["Mvir"]} except NoParticlesException as e: # mpi.msg(e.message) continue if mpi.host: import pickle, os pickle_path = "%s/pickle/%s/" % (snap.path, halos.finder) if not os.path.isdir(pickle_path): os.mkdir(pickle_path) pickle.dump( mpi.unpack(dest), open("%s/time_int_fesc_%05i.p" % (pickle_path, snap.ioutput), "wb"))
def main(path, iout, field, pickle_path=None): import seren3 import pickle, os from seren3.analysis.parallel import mpi mpi.msg("Loading data") snap = seren3.load_snapshot(path, iout) # snap.set_nproc(1) # disbale multiprocessing/threading snap.set_nproc(8) halos = snap.halos() halo_ix = None if mpi.host: halo_ix = halos.halo_ix(shuffle=True) dest = {} for i, sto in mpi.piter(halo_ix, storage=dest): h = halos[i] mpi.msg("Working on halo %i \t %i" % (i, h.hid)) # vw = _volume_weighted_average(field, h) mw = _mass_weighted_average(field, h) if (np.isinf(mw) or np.isnan(mw)): continue # vw = _volume_weighted_average_cube(snap, field, h) mpi.msg("%i \t %1.2e" % (h.hid, mw)) sto.idx = h["id"] # sto.result = {"vw" : vw, "mw" : mw} sto.result = {"mw": mw} if mpi.host: if pickle_path is None: pickle_path = "%s/pickle/" % path if os.path.isdir(pickle_path) is False: os.mkdir(pickle_path) fname = "%s/%s_halo_av_%05i.p" % (pickle_path, field, iout) pickle.dump(mpi.unpack(dest), open(fname, "wb")) mpi.msg("Done")
def main(path, iout, pickle_path): import pickle, os import seren3 from seren3.analysis import baryon_fraction from seren3.analysis.parallel import mpi back_to_z = 20. back_to_aexp = 1. / (1. + back_to_z) mpi.msg("Loading data") snap = seren3.load_snapshot(path, iout) snap.set_nproc(1) # disbale multiprocessing/threading halos = snap.halos(finder="ctrees") nhalos = len(halos) halo_ix = None if mpi.host: halo_ix = halos.halo_ix(shuffle=True) dest = {} for i, sto in mpi.piter(halo_ix, storage=dest, print_stats=True): # mpi.msg("%i / %i" % (i, nhalos)) h = halos[i] tint_fbaryon_hist, I1, I2, lbtime, fbaryon_dict, tidal_force_tdyn_dict, age_dict = baryon_fraction.time_integrated_fbaryon(h, back_to_aexp) sto.idx = h["id"] sto.result = {"tint_fbaryon" : tint_fbaryon_hist, "I1" : I1, "I2" : I2, "lbtime" : lbtime, "fbaryon_dict" : fbaryon_dict, \ "tidal_force_tdyn_dict" : tidal_force_tdyn_dict, "age_dict" : age_dict} if mpi.host: if pickle_path is None: pickle_path = "%s/pickle/" % path if os.path.isdir(pickle_path) is False: os.mkdir(pickle_path) fname = "%s/tint_fbaryon_tdyn_%05i_fixed.p" % (pickle_path, iout) pickle.dump( mpi.unpack(dest), open( fname, "wb" ) )
def plot(path, iout, pickle_path, field, the_mass_bins=[7., 8., 9., 10.], ax=None, **kwargs): import pickle from seren3.analysis.plots import fit_scatter from seren3.utils import flatten_nested_array import matplotlib.pylab as plt snap = seren3.load_snapshot(path, iout) fname = "%s/cdf_%s_halos_%05i.p" % (pickle_path, field, iout) data = pickle.load(open(fname, 'rb')) cdf_halos = [] bin_centre_halos = [] Mvir = [] for i in range(len(data)): res = data[i].result cdf_halos.append(res["C"]) bin_centre_halos.append(res["bc"]) Mvir.append(res["Mvir"]) cdf_halos = np.array(cdf_halos) bin_centre_halos = np.array(bin_centre_halos) Mvir = np.array(Mvir) idx = np.where(np.log10(Mvir) >= 6.5) cdf_halos = cdf_halos[idx] bin_centre_halos = bin_centre_halos[idx] Mvir = Mvir[idx] # Bin idx mass_bins = np.digitize(np.log10(Mvir), the_mass_bins, right=True) binned_cdf = {} nbins = kwargs.pop("nbins", 5) for i in range(len(the_mass_bins) + 1): if (i == len(the_mass_bins)): x, y = (flatten_nested_array(bin_centre_halos), flatten_nested_array(cdf_halos)) idx = np.where(~np.isnan(y)) binned_cdf["all"] = fit_scatter(x[idx], y[idx], ret_sterr=True, ret_n=True, nbins=nbins) else: idx = np.where(mass_bins == i) x, y = (flatten_nested_array(bin_centre_halos[idx]), flatten_nested_array(cdf_halos[idx])) idx = np.where(~np.isnan(y)) binned_cdf[i] = fit_scatter(x[idx], y[idx], ret_sterr=True, ret_n=True, nbins=nbins) binned_cdf['the_mass_bins'] = the_mass_bins if ax is None: return binned_cdf # ax = plt.gca() cols = None if "cols" not in kwargs: from seren3.utils import plot_utils # cols = plot_utils.ncols(len(the_mass_bins)+1, cmap=kwargs.pop("cmap", "jet"))[::-1] cols = ["r", "b", "darkorange", "k"] else: cols = kwargs.pop("cols") ls = kwargs.pop("linestyle", "-") for i, c in zip(range(len(the_mass_bins)), cols): x, y, std, stderr, n = binned_cdf[i] if i == len(the_mass_bins) - 1: upper = r"$\infty$" else: upper = "%1.1f" % the_mass_bins[i + 1] lower = "%1.1f" % the_mass_bins[i] label = "[%s, %s)" % (lower, upper) print label, n ax.errorbar(x, y, yerr=std, label=label, linewidth=3., \ color=c, linestyle=ls, capsize=5) ax.set_xlabel(r"log$_{10}$ n$_{\mathrm{HI}}$ [m$^{-3}$]") ax.set_ylabel("CDF") if kwargs.pop("legend", True): # Shrink current axis by 20% box = ax.get_position() ax.set_position( [box.x0, box.y0 + box.height * 0.15, box.width, box.height * 0.9]) # Put a legend to the right of the current axis leg = ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.15), fancybox=True, shadow=False, ncol=6, prop={"size": 18}) leg.set_title(r"log$_{10}$(Mvir) [M$_{\odot}$/h]", prop={'size': 'x-large'}) if kwargs.pop("label_nstar", False): fs = 15 n_star = snap.array(snap.info_rt["n_star"].express(snap.C.H_cc), "cm**-3").in_units("m**-3") ax.vlines(np.log10(n_star), -1, 2, color='k', linestyle='-.') # ax1.text(np.log10(snapshot.info_rt['n_star'].express(snapshot.C.H_cc)) + xr*0.01, ymax/2. + .01, r"n$_{*}$", color='r', fontsize=fs) # ax1.text(np.log10(n_star) + xr*0.01, 0.076, r"n$_{*}$", color='r', fontsize=fs) ax.set_ylim(-0.05, 1.05) # ax.set_yscale("log") return binned_cdf
def main(path, iout, pickle_path): import seren3 from seren3.core.serensource import DerivedDataset from seren3.utils import derived_utils from seren3.analysis.escape_fraction import fesc from seren3.analysis.parallel import mpi from seren3.exceptions import NoParticlesException import pickle, os mpi.msg("Loading snapshot") snap = seren3.load_snapshot(path, iout) snap.set_nproc(1) halos = None star_Nion_d_fn = derived_utils.get_derived_field(snap.s, "Nion_d") nIons = snap.info_rt["nIons"] halos = snap.halos(finder="ctrees") halo_ix = None if mpi.host: halo_ix = halos.halo_ix(shuffle=True) dest = {} for i, sto in mpi.piter(halo_ix, storage=dest, print_stats=True): h = halos[i] star_dset = h.s[["age", "metal", "mass"]].flatten() if (len(star_dset["mass"]) > 0): dt = h.dt ix_young = np.where( (star_dset["age"].in_units("Gyr") - dt.in_units("Gyr")) >= 0.) if len(ix_young[0] > 0): # Work on this halo sto.idx = int(h["id"]) try: dm_dset = h.d["mass"].flatten() gas_dset = h.g["mass"].flatten() star_mass = star_dset["mass"] star_age = star_dset["age"] star_metal = star_dset["metal"] dict_stars = { "age": star_age, "metal": star_metal, "mass": star_mass } dset_stars = DerivedDataset(snap.s, dict_stars) Nion_d_all_groups = snap.array( np.zeros(len(dset_stars["age"])), "s**-1 Msol**-1") for ii in range(nIons): Nion_d_now = star_Nion_d_fn(snap, dset_stars, group=ii + 1, dt=0.) Nion_d_all_groups += Nion_d_now tot_mass = dm_dset["mass"].in_units("Msol h**-1").sum() + star_dset["mass"].in_units("Msol h**-1").sum()\ + gas_dset["mass"].in_units("Msol h**-1").sum() h_fesc = fesc(h.subsnap, nside=2**3, filt=False, do_multigroup=True, denoise=True) mpi.msg("%1.2e \t %1.2e" % (h["Mvir"], h_fesc)) sto.result = {"fesc" : h_fesc, "tot_mass" : tot_mass, \ "Nion_d_now" : Nion_d_all_groups, "star_mass" : star_mass,\ "hprops" : h.properties} except NoParticlesException as e: mpi.msg(e.message) if mpi.host: if pickle_path is None: pickle_path = "%s/pickle/%s/" % (path, halos.finder.lower()) # pickle_path = "%s/" % path if os.path.isdir(pickle_path) is False: os.mkdir(pickle_path) unpacked_dest = mpi.unpack(dest) fesc_dict = {} for i in range(len(unpacked_dest)): fesc_dict[int(unpacked_dest[i].idx)] = unpacked_dest[i].result #pickle.dump( fesc_dict, open( "%s/fesc_database_no_filt_denoise_%05i.p" % (pickle_path, iout), "wb" ) ) pickle.dump( fesc_dict, open("%s/time_int_fesc_all_halos_%05i.p" % (pickle_path, iout), "wb"))
def main(path, iout, pickle_path): ''' Compute nH CDF of all halos and bin by Mvir ''' from seren3.analysis.parallel import mpi from seren3.analysis.plots import histograms snap = seren3.load_snapshot(path, iout) nH_range = None # nH_range = np.array([ 1.59355764e+00, 7.93249184e+09]) # nH_range = np.array([ 1.0e-01, 1.0e+10]) if nH_range is None: if mpi.host: print "Bracketing %s" % _FIELD min_nH = np.inf max_nH = -np.inf dset = None if _MEM_OPT: count = 1 ncpu = snap.info["ncpu"] for dset in snap.g[_FIELD]: print "%i/%i cpu" % (count, ncpu) count += 1 nH = dset[_FIELD] if nH.min() < min_nH: min_nH = nH.min() if nH.max() > max_nH: max_nH = nH.max() else: print "Loading full dataset" dset = snap.g[_FIELD].flatten() print "Loaded full dataset" min_nH = dset[_FIELD].min() max_nH = dset[_FIELD].max() nH_range = np.array([min_nH, max_nH]) print nH_range del dset nH_range = mpi.comm.bcast(nH_range, root=mpi.HOST_RANK) mpi.msg(nH_range) snap.set_nproc(1) halos = snap.halos() nhalos = len(halos) halo_ix = None if mpi.host: halo_ix = halos.halo_ix(shuffle=True) dest = {} for i, sto in mpi.piter(halo_ix, storage=dest): mpi.msg("%i / %i" % (i, nhalos)) h = halos[i] sto.idx = h["id"] if len(h.g) > 0: P, C, bc, dx = histograms.pdf_cdf(h.g, _FIELD, cumulative=True, plot=False, x_range=nH_range) if (P.max() > 0.): sto.result = {"C": C, "bc": bc, "Mvir": h["Mvir"]} if mpi.host: import os, pickle if pickle_path is None: pickle_path = "%s/pickle/" % path if os.path.isdir(pickle_path) is False: os.mkdir(pickle_path) fname = "%s/cdf_%s_halos_%05i.p" % (pickle_path, _FIELD, iout) pickle.dump(mpi.unpack(dest), open(fname, "wb"))
import seren3, sys from seren3.analysis.visualization import EngineMode path = sys.argv[1] iout = int(sys.argv[2]) snap = seren3.load_snapshot(path, iout) halos = snap.halos(finder='ctrees') h = halos.sort("Mvir")[0] proj = h.d.projection("mass", mode=EngineMode.SPLATTER)