def plot_baryfrac(pp1=None, pp2=None):
from seren3.utils import plot_utils
reload(plot_utils)
cmap = plot_utils.load_custom_cmaps("parula")
path1 = '/lustre/scratch/astro/ds381/simulations/bpass/ramses/'
path2 = '/research/prace/david/bpass/bc03/'
sim1 = seren3.init(path1)
sim2 = seren3.init(path2)
snap1 = sim1[93]
snap2 = sim2[106]
if (pp1 is None):
pp1 = PhasePlot(snap1, nbins=500)
if (pp2 is None):
pp2 = PhasePlot(snap2, nbins=500)
for pp in [pp1, pp2]:
pp.cmap = cmap
pp.annotate_nstar()
pp.annotate_T2_thresh()
# pp.annotate_fit()
pp.vmin = -10
pp.vmax = -2
z = [pp1.snapshot.z, pp1.snapshot.z]
fig, axs = plt.subplots(nrows=1, ncols=2, figsize=(15, 6))
pp1.draw(ax=axs[0], draw_cbar=False)
pp2.draw(ax=axs[1], draw_cbar=False)
# plt.tight_layout()
# cbar1 = fig.colorbar(pp1.im, ax=axs[0])
# cbar2 = fig.colorbar(pp2.im, ax=axs[1])
# cbar1.set_label('f(mass)', labelpad=-5)
# cbar2.set_label('f(mass)', labelpad=-5)
plt.tight_layout()
fig.subplots_adjust(right=0.8)
cbar_ax = fig.add_axes([0.85, 0.15, 0.04, 0.8])
cbar = fig.colorbar(pp2.im, cax=cbar_ax, cmap=cmap)
cbar.set_label('f(mass)', labelpad=-5)
labels = ["HD", "RT2"]
axs[1].set_ylabel('')
# for ax, snap, zi, l in zip(axs, [snap1, snap2], z, labels):
# ax.set_title("%s: z = %1.2f" % (l, zi), fontsize=20)
# ax.set_title(l, fontsize=18)
# plt.show()
plt.savefig("./bc03_phase_diag_parula.pdf", format="pdf")
return pp1, pp2
def main(path, pickle_path):
import seren3
from seren3.analysis.parallel import mpi
import pickle, os
mpi.msg("Loading simulation")
simulation = seren3.init(path)
if mpi.host:
mpi.msg("Averaging field: %s" % field)
iout_start = max(simulation.numbered_outputs[0], 20)
iouts = range(iout_start, max(simulation.numbered_outputs) + 1)
dest = {}
for iout, sto in mpi.piter(iouts, storage=dest):
mpi.msg("%05i" % iout)
snapshot = simulation[iout]
snapshot.set_nproc(1)
sto.result = {"mw" : mass_weighted_average(snapshot.g), \
"z" : snapshot.z}
if mpi.host:
if pickle_path == None:
pickle_path = "%s/pickle/" % path
if os.path.isdir(pickle_path) is False:
os.mkdir(pickle_path)
pickle.dump(
mpi.unpack(dest),
open("%s/%s_mw_time_averaged.p" % (pickle_path, field), "wb"))
def main(path, field, pickle_path):
import seren3
from seren3.analysis import mean_field_rho_mean
from seren3.analysis.parallel import mpi
import pickle, os
mpi.msg("Loading simulation")
simulation = seren3.init(path)
if mpi.host:
mpi.msg("Averaging field: %s" % field)
iout_start = max(simulation.numbered_outputs[0], 1)
iouts = range(iout_start, max(simulation.numbered_outputs)+1)
mpi.msg("Starting with snapshot %05i" % iout_start)
dest = {}
for iout, sto in mpi.piter(iouts, storage=dest, print_stats=True):
mpi.msg("%05i" % iout)
snapshot = simulation.snapshot(iout, verbose=False)
snapshot.set_nproc(1)
mw, rho_mean = mean_field_rho_mean(snapshot, field, ret_rho_mean=True)
sto.idx = iout
sto.result = {"rho_mean" : rho_mean, "mw" : mw, "z" : snapshot.z, "aexp" : snapshot.cosmo["aexp"]}
if mpi.host:
if pickle_path == None:
pickle_path = "%s/pickle/" % path
if os.path.isdir(pickle_path) is False:
os.mkdir(pickle_path)
pickle.dump( mpi.unpack(dest), open("%s/%s_time_averaged_at_rho_mean.p" % (pickle_path, field), "wb") )
def simulation_example(path, zi=6., name=None):
'''
Sometimes, it's useful to start with a simulation object, and then load your snapshot
'''
import seren3
sim = None # init
# Load our simulation
if (name is not None): # load by name
sim = seren3.load(name)
else:
# Just init from the path
sim = seren3.init(path)
print sim
# Now, lets load the snapshot which is closest to the desired redshift
ioutput = sim.redshift(zi) # output number
snapshot = sim[ioutput]
# There are also interpolators for age -> redshift and vise versa, for our chosen cosmology
age_func = sim.age_func(zmax=100., zmin=0.)
z_func = sim.redshift_func(zmax=100., zmin=0.)
age_of_universe = age_func(snapshot.z)
redshift = z_func(age_of_universe)
print snapshot.z, redshift, age_of_universe
def test(path, istart, iend):
import seren3
import numpy as np
sim = seren3.init(path)
families = [sim[i].g for i in range(istart, iend + 1)]
camera_func = lambda family: family.camera(
region_size=np.array([.05, .05]), distance=.05, far_cut_depth=.05)
print camera_func(families[0]).__dict__
return make_movie(families, "output.avi", camera_func=camera_func)
def main(path, pickle_path):
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
from seren3.scripts.mpi import write_fesc_hid_dict
mpi.msg("Loading simulation...")
sim = seren3.init(path)
iout_start = max(sim.numbered_outputs[0], 60)
back_to_aexp = sim[60].info["aexp"]
# iouts = range(iout_start, max(sim.numbered_outputs)+1)
print "IOUT RANGE HARD CODED"
iouts = range(iout_start, 109)
# iouts = [109]
for iout in iouts[::-1]:
snap = sim[iout]
mpi.msg("Working on snapshot %05i" % snap.ioutput)
snap.set_nproc(1)
halos = snap.halos(finder="ctrees")
db = write_fesc_hid_dict.load_db(path, iout)
halo_ids = None
if mpi.host:
halo_ids = db.keys()
random.shuffle(halo_ids)
dest = {}
for i, sto in mpi.piter(halo_ids, storage=dest, print_stats=True):
h = halos.with_id(i)
res = time_integrated_fesc(h, back_to_aexp, db=db, return_data=True)
if (res is not None):
mpi.msg("%05i \t %i \t %i" % (snap.ioutput, h.hid, i))
tint_fesc_hist, I1, I2, lbtime, hids, iouts = res
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"], 'hids' : hids, 'iouts' : iouts}
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_all_halos_%05i.p" % (pickle_path, snap.ioutput), "wb") )
mpi.msg("Waiting...")
mpi.comm.Barrier()
def test_sfr():
'''
Test SFR calculation is working correctly
'''
import seren3
import numpy as np
from scipy import integrate
path = "/research/prace/david/aton/256/"
sim = seren3.init(path)
iout = sim.numbered_outputs[-1] # last snapshot
snap = sim[iout]
snap_sfr, lbtime, bsize = sfr(snap)
dset = snap.s["mass"].flatten()
mstar_tot = dset["mass"].in_units("Msol").sum()
integrated_mstar = integrate.trapz(snap_sfr, lbtime)
assert (np.allclose(
mstar_tot, integrated_mstar,
rtol=1e-2)), "Error: Integrated stellar mass not close to actual."
print "Passed"
def main(path, field, pickle_path):
import seren3
from seren3.analysis.parallel import mpi
import pickle, os
mpi.msg("Loading simulation")
simulation = seren3.init(path)
if mpi.host:
mpi.msg("Averaging field: %s" % field)
iout_start = max(simulation.numbered_outputs[0], 1)
iouts = range(iout_start, max(simulation.numbered_outputs)+1)
mpi.msg("Starting with snapshot %05i" % iout_start)
dest = {}
for iout, sto in mpi.piter(iouts, storage=dest, print_stats=True):
mpi.msg("%05i" % iout)
snapshot = simulation.snapshot(iout, verbose=False)
snapshot.set_nproc(1)
# vw = snapshot.quantities.volume_weighted_average(field, mem_opt=mem_opt)
vw, mw = volume_mass_weighted_average(snapshot, field)
sto.idx = iout
sto.result = {"vw" : vw, "mw" : mw, "z" : snapshot.z}
# if do_mass_weighted:
# mw = snapshot.quantities.mass_weighted_average(field, mem_opt=mem_opt)
# sto.result["mw"] = mw
if mpi.host:
if pickle_path == None:
pickle_path = "%s/pickle/" % path
if os.path.isdir(pickle_path) is False:
os.mkdir(pickle_path)
pickle.dump( mpi.unpack(dest), open("%s/%s_time_averaged.p" % (pickle_path, field), "wb") )
def plot(path, iout, pickle_path, the_mass_bins=[7., 8., 9., 10,], lab='', ax=None, **kwargs):
import pickle
from seren3.analysis.plots import fit_scatter
from seren3.utils import flatten_nested_array
import matplotlib.pylab as plt
sim = seren3.init(path)
# snap = seren3.load_snapshot(path, iout)
snap = sim[iout]
fname = "%s/sfr_halos_%05i.p" % (pickle_path, iout)
data = pickle.load(open(fname, 'rb'))
sfr_halos = []; bin_centre_halos = []; Mvir = []
for i in range(len(data)):
res = data[i].result
sfr_halos.append(res["SFR"].in_units("Msol yr**-1"))
bin_centre_halos.append(res["lookback-time"])
Mvir.append(res["Mvir"])
sfr_halos = np.array(sfr_halos); bin_centre_halos = np.array(bin_centre_halos); Mvir = np.array(Mvir)
idx = np.where(np.log10(Mvir) >= 6.5)
sfr_halos = sfr_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_sfr = {}
nbins = kwargs.pop("nbins", 100)
# nbins = kwargs.pop("nbins", 50)
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(sfr_halos) )
idx = np.where(~np.isnan(y))
binned_sfr["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(sfr_halos[idx]) )
idx = np.where(~np.isnan(y))
binned_sfr[i] = fit_scatter(x[idx], y[idx], ret_sterr=True, ret_n=True, nbins=nbins)
binned_sfr['the_mass_bins'] = the_mass_bins
if ax is None:
# return binned_sfr
fig, ax = plt.subplots(figsize=(8,8))
# ax = plt.gca()
cols = None
if "cols" not in kwargs:
from seren3.utils import plot_utils
cols = plot_utils.ncols(len(the_mass_bins), cmap=kwargs.pop("cmap", "Set1"))[::-1]
# cols = plot_utils.ncols(len(the_mass_bins), cmap=kwargs.pop("cmap", "Set1"))[::-1]
# cols = ["r", "b", "darkorange", "k"]
else:
cols = kwargs.pop("cols")
ls = kwargs.pop("linestyle", "-")
lw = kwargs.get("lw", 2.5)
z_fn = sim.redshift_func(zmax=1000., zmin=0.)
age_fn = sim.age_func()
sim_age = age_fn(snap.z)
for i,c in zip(range(len(the_mass_bins)), cols):
# for i,c in zip([1, len(the_mass_bins)-1], cols):
x,y,std,stderr,n = binned_sfr[i]
age = sim_age - x
age_to_z = z_fn(age)
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 = "BC03 log(M) = [%s, %s)" % (lower, upper)
label = "%s log(M) = [%s, %s)" % (lab, lower, upper)
if kwargs.get("legend", False) is False:
label = None
# print label, n
# ax.errorbar(x, y, yerr=std, label=label, linewidth=3., \
# color=c, linestyle=ls, capsize=5)
from scipy import integrate
from seren3.array import SimArray
integrated_mstar = integrate.trapz(y, SimArray(x, "Gyr").in_units("yr"))
print integrated_mstar
# ax.step(x, y, label=label, linewidth=2.5, color=c, linestyle=ls)
ax.step(age_to_z, y, label=label, linewidth=lw, color=c, linestyle=ls)
# ax.set_xlabel(r"Lookback-time [Gyr]")
ax.set_xlabel(r"$z$")
ax.set_ylabel(r"SFR [M$_{\odot}$ yr$^{-1}$]")
if kwargs.pop("legend", False):
# 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 = ax.legend(loc="upper right", prop={"size":14})
# leg.set_title(r"log$_{10}$(Mvir) [M$_{\odot}$/h]", prop = {'size':'x-large'})
# ax.set_ylim(-0.05, 1.05)
ax.set_yscale("log")
# ax.set_xscale("log")
ax.set_ylim(1e-6, 5e0)
ax.set_xlim(6., 20.)
# ax.set_xlim(0.,)
return binned_sfr
def main(path, iout, pickle_path, allow_estimation=False):
import seren3
import pickle, os
from seren3.analysis.parallel import mpi
import random
mpi.msg("Loading data")
sim = seren3.init(path)
# snap = seren3.load_snapshot(path, iout)
snap = sim.snapshot(iout)
# Age function and age now to compute time since last MM
age_fn = sim.age_func()
# age_now = snap.age
age_now = age_fn(snap.z)
snap.set_nproc(1) # disbale multiprocessing/threading
# Use consistent trees halo catalogues to compute time since last major merger
# finder = "ahf"
finder = "ctrees"
# halos = snap.halos(finder="ctrees")
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):
h = halos[i]
mpi.msg("Working on halo %i \t %i" % (i, h.hid))
part_dset = h.p[["id", "mass", "epoch"]].flatten()
ix_dm = np.where(np.logical_and( part_dset["id"] > 0., part_dset["epoch"] == 0 )) # index of dm particles
ix_stars = np.where( np.logical_and( part_dset["id"] > 0., part_dset["epoch"] != 0 ) ) # index of star particles
np_dm = len(ix_dm[0])
# if np_dm > 50.:
gas_dset = h.g["mass"].flatten()
ncell = len(gas_dset["mass"])
gas_mass_tot = 0.
if ncell == 0 and allow_estimation:
mpi.msg("Estimating gas mass for halo %s" % h["id"])
gas_mass_tot = estimate_unresolved(snap, h)
ncell = 1
else:
gas_mass_tot = gas_dset["mass"].in_units(_MASS_UNIT).sum()
part_mass_tot = part_dset["mass"].in_units(_MASS_UNIT).sum()
star_mass_tot = part_dset["mass"].in_units(_MASS_UNIT)[ix_stars].sum()
# gas_mass_tot = gas_dset["mass"].in_units(_MASS_UNIT).sum()
tot_mass = part_mass_tot + gas_mass_tot
fb = (gas_mass_tot + star_mass_tot)/tot_mass
sto.idx = h["id"]
if finder == "ctrees":
# Compute time since last major merger and store
pid = h["pid"] # -1 if distinct halo
scale_of_last_MM = h["scale_of_last_mm"] # aexp of last major merger
z_of_last_MM = (1./scale_of_last_MM)-1.
age_of_last_MM = age_fn(z_of_last_MM)
time_since_last_MM = age_now - age_of_last_MM
sto.result = {"fb" : fb, "tot_mass" : tot_mass,\
"pid" : pid, "np_dm" : np_dm, "ncell" : ncell,\
"time_since_last_MM" : time_since_last_MM, "hprops" : h.properties}
# else:
# mpi.msg("Skipping halo with %i dm particles" % np_dm)
else:
sto.result = {"fb" : fb, "tot_mass" : tot_mass,\
"np_dm" : np_dm, "ncell" : ncell,\
"hprops" : h.properties}
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/fbaryon_tdyn_%05i.p" % (pickle_path, iout)
if allow_estimation:
fname = "%s/fbaryon_tdyn_gdset_est_%05i.p" % (pickle_path, iout)
pickle.dump( mpi.unpack(dest), open( fname, "wb" ) )
def main(path, iout, pickle_path, allow_estimation=False):
import seren3
import pickle, os
from seren3.analysis.parallel import mpi
import random
def _mass_weighted_average(halo, mass_units="Msol h**-1"):
dset = halo.g[["xHII", "T", "mass"]].flatten()
cell_mass = dset["mass"].in_units(mass_units)
return np.sum(dset["xHII"] * cell_mass) / cell_mass.sum(), np.sum(
dset["T"] * cell_mass) / cell_mass.sum()
mpi.msg("Loading data")
sim = seren3.init(path)
# snap = seren3.load_snapshot(path, iout)
snap = sim.snapshot(iout)
# Age function and age now to compute time since last MM
age_fn = sim.age_func()
# age_now = age_fn(snap.z)
snap.set_nproc(1) # disbale multiprocessing/threading
# Use consistent trees halo catalogues to compute time since last major merger
# finder = "ahf"
finder = "ctrees"
# halos = snap.halos(finder="ctrees")
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):
h = halos[i]
mpi.msg("Working on halo %i \t %i" % (i, h.hid))
fb, tot_mass, ncell, np_dm = baryon_fraction.compute_fb(
h, return_stats=True)
sto.idx = h["id"]
if finder == "ctrees":
# Compute time since last major merger and store
pid = h["pid"] # -1 if distinct halo
scale_now = h["aexp"]
z_now = (1. / scale_now) - 1.
age_now = age_fn(z_now)
scale_of_last_MM = h[
"scale_of_last_mm"] # aexp of last major merger
z_of_last_MM = (1. / scale_of_last_MM) - 1.
age_of_last_MM = age_fn(z_of_last_MM)
xHII_mw, T_mw = _mass_weighted_average(h)
time_since_last_MM = age_now - age_of_last_MM
sto.result = {"fb" : fb, "tot_mass" : tot_mass,\
"pid" : pid, "np_dm" : np_dm, "ncell" : ncell,\
"time_since_last_MM" : time_since_last_MM, "hprops" : h.properties, \
"xHII_mw" : xHII_mw, "T_mw" : T_mw}
# else:
# mpi.msg("Skipping halo with %i dm particles" % np_dm)
else:
sto.result = {"fb" : fb, "tot_mass" : tot_mass,\
"np_dm" : np_dm, "ncell" : ncell,\
"hprops" : h.properties}
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/fbaryon_tdyn_%05i.p" % (pickle_path, iout)
pickle.dump(mpi.unpack(dest), open(fname, "wb"))