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)
