def main(path, pickle_path):
import pickle
from seren3.core.simulation import Simulation
from seren3.analysis.parallel import mpi
mpi.msg("Loading simulation")
sim = Simulation(path)
ioutputs = sim.numbered_outputs
dest = {}
for i, sto in mpi.piter(ioutputs, storage=dest):
mpi.msg("Analysing snapshot %05i" % i)
snap = sim[i]
snap.set_nproc(1) # disable multiprocessing on dset read
z = snap.z
vw = snap.quantities.volume_weighted_average("xHII", mem_opt=MEM_OPT)
mw = snap.quantities.mass_weighted_average("xHII", mem_opt=MEM_OPT)
sto.idx = i
sto.result = {"z": z, "volume_weighted": vw, "mass_weighted": mw}
if mpi.host:
if pickle_path is None:
pickle_path = "%s/pickle/" % sim.path
unpacked_dest = unpack(dest)
pickle.dump(unpacked_dest,
open("%s/xHII_reion_history.p" % pickle_path, "wb"))
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, 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 apply_bias(ic, species, cube, dx, pad, field, ax=None):
origin = np.array(cube - float(dx) / 2. - pad).astype(np.int32)
dx_eps = float(dx) + float(2 * pad)
if debug:
mpi.msg("Loading fields")
delta_cube = None
if (field == 'vel'):
# Velocity field
delta_cube = ic.lazy_load_periodic('vel%s%s' % (species, ax), origin,
int(dx_eps))
elif (field == 'rho'):
delta_cube = ic.lazy_load_periodic('delta%s' % species, origin,
int(dx_eps))
else:
mpi.msg("Unknown species %s" % species)
mpi.terminate(500)
vbc_cube = ic.lazy_load_periodic('vbc', origin, int(dx_eps))
if debug:
mpi.msg("Done")
# Compute the PS bias
if debug:
mpi.msg("Computing bias")
k = bc = bb = None
if (field == 'vel'):
k, bc, bb = drift_velocity.compute_velocity_bias(ic, vbc_cube)
else:
k, bc, bb = drift_velocity.compute_bias(ic, vbc_cube)
# Apply bias to this patch
if debug:
mpi.msg("Applying bias")
if 'b' == species:
modified_delta = drift_velocity.apply_density_bias(ic,
k,
bb,
delta_cube.shape[0],
delta_x=delta_cube)
elif 'c' == species:
modified_delta = drift_velocity.apply_density_bias(ic,
k,
bc,
delta_cube.shape[0],
delta_x=delta_cube)
else:
mpi.terminate("Unknown species %s" % species)
x_shape, y_shape, z_shape = modified_delta.shape
# Remove padding
modified_delta = modified_delta[0 + pad:x_shape - pad,
0 + pad:y_shape - pad,
0 + pad:z_shape - pad]
return modified_delta
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, 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 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 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, 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, 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 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, 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):
import os
from seren3.core.simulation import Simulation
from seren3.analysis.parallel import mpi
from seren3.scripts.mpi import fbaryon
mpi.msg("Loading simulation")
sim = Simulation(path)
ioutputs = sim.numbered_outputs
dest = {}
for i, sto in mpi.piter(ioutputs, storage=dest):
if os.path.isfile("%s/pickle/fbaryon_%05i.p" % (sim.path, i)):
mpi.msg("Analysing snapshot %05i" % i)
snap = sim[i]
snap.set_nproc(1)
Mc, alpha, corr = fbaryon.fit(snap)
sto.idx = i
sto.result = {"Mc": Mc, "alpha": alpha, "corr": corr, "z": snap.z}
else:
mpi.msg("Pickle file not found, skipping output %05i" % i)
sto.result = None
if mpi.host:
import pickle
pickle.dump(mpi.unpack(dest), open("%s/Mc_alpha_fit.p" % path, "wb"))
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 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, 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") )
z = snap.z
vw = snap.quantities.volume_weighted_average("xHII", mem_opt=MEM_OPT)
mw = snap.quantities.mass_weighted_average("xHII", mem_opt=MEM_OPT)
sto.idx = i
sto.result = {"z": z, "volume_weighted": vw, "mass_weighted": mw}
if mpi.host:
if pickle_path is None:
pickle_path = "%s/pickle/" % sim.path
unpacked_dest = unpack(dest)
pickle.dump(unpacked_dest,
open("%s/xHII_reion_history.p" % pickle_path, "wb"))
if __name__ == "__main__":
import sys
path = sys.argv[1]
pickle_path = None
if len(sys.argv) > 2:
from seren3.analysis.parallel import mpi
pickle_path = sys.argv[2]
if mpi.host:
mpi.msg("pickle path: %s" % pickle_path)
try:
main(path, pickle_path)
except Exception as e:
from seren3.analysis.parallel import mpi
mpi.terminate(500, e=e)
def make_movie(families,
out_fname,
field="rho",
camera_func=None,
mpi=True,
**kwargs):
'''
Parameters
----------
families : :iterable:class:`~seren3.core.snapshot.Family`
iterable of families to make images from
field : ``string`` (default rho)
field to make projections of.
camera : :class:`~pymses.analysis.camera.Camera`
camera containing all the view params
Returns
-------
TODO
'''
import os, sys
from seren3.analysis import visualization
if camera_func is None:
camera_func = lambda family: family.camera()
fraction = kwargs.pop("fraction", 0.01)
cmap = kwargs.pop("cmap", "YlOrRd")
verbose = kwargs.pop("verbose", config.get("general", "verbose"))
fps = kwargs.pop("fps", 25)
try:
if mpi:
import pymses
from seren3.analysis.parallel import mpi
for i in mpi.piter(range(len(families))):
if verbose:
mpi.msg("Image %i/%i" % (i, len(families)))
family = families[i]
cam = camera_func(family)
proj = visualization.Projection(family,
field,
camera=cam,
multi_processing=False,
fraction=True,
vol_weighted=True,
**kwargs)
fname = _get_fname(family, field)
proj.save_PNG(img_fname=fname, fraction=fraction, cmap=cmap)
mpi.comm.barrier()
if mpi.host:
_run_mencoder(out_fname, fps)
else:
fnames = []
for i in range(len(families)):
if verbose:
print "Image %i/%i" % (i, len(families))
family = families[i]
cam = camera_func(family)
proj = visualization.Projection(family,
field,
camera=cam,
multi_processing=True,
**kwargs)
fname = _get_fname(family, field)
proj.save_PNG(img_fname=fname, fraction=fraction, cmap=cmap)
fnames.append(fname)
_run_mencoder(out_fname, fps)
except Exception as e:
_cleanup()
return e
except KeyboardInterrupt:
_cleanup()
sys.exit()
return 0
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)
'''
from seren3.analysis.parallel import mpi
import numpy as np
# If we are host (rank 0), then initialise the iterable
if mpi.host:
data = np.linspace(1, 10, num=10, dtype=np.int32)
print data
# Set the variable to None, host will scatter data to us
else:
data = None
# Dictionaries can be passed to piter to store the results. This is automatically gathered to host
# upon for loop exit
dest = {}
for i, res in mpi.piter(data, storage=dest):
# Set the idx variable of the yielded Result object to the value we were given
res.idx = i
# Set the result variable to whatever we want
res.result = (mpi.rank+1)*(i**2)
# Print the results along with our rank. Rank 1 -> size should show empty dictionaries, host has
# all the results.
mpi.msg(dest)
# We can unpack the dictionary easily
if mpi.host:
unpacked_dest = mpi.unpack(dest)
mpi.msg(unpacked_dest)
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"))
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, 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, level, patch_size):
'''
Writes a new set of grafIC initial conditions with a drift velocity dependent
bias in the power spectrum
'''
from seren3.core import grafic_snapshot
from seren3.analysis import drift_velocity as vbc_utils
from seren3.analysis.parallel import mpi
from seren3.utils import divisors
mpi.msg("Loading initial conditions")
ics = grafic_snapshot.load_snapshot(path, level, sample_fft_spacing=False)
if mpi.host:
# Make sure vbc field exists on disk
if not ics.field_exists_on_disk("vbc"):
ics.write_field(ics["vbc"], "vbc")
div = np.array([float(i) for i in divisors(ics.header.N, mode='yield')])
idx = np.abs((ics.header.N / div) * ics.dx - patch_size).argmin()
ncubes = int(div[idx])
# Compute cube positions in cell units
cubes = dx = None
if mpi.host:
mpi.msg("Using %i cubes per dimension." % ncubes)
cubes, dx = vbc_utils.cube_positions(ics, ncubes)
cubes = np.array(cubes)
dx = mpi.comm.bcast(dx, root=0)
pad = 8
############################## WORK LOOP ######################################
# Iterate over patch positions in parallel
dest = {}
for patch, sto in mpi.piter(cubes, storage=dest, print_stats=True):
origin = np.array(patch - float(dx) / 2. - pad, dtype=np.int64)
dx_eps = float(dx) + float(2 * pad)
delta = vbc = None
if (VERBOSE): mpi.msg("Loading patch: %s" % patch)
delta = ics.lazy_load_periodic("deltab", origin, int(dx_eps))
vbc = ics.lazy_load_periodic("vbc", origin, int(dx_eps))
# Compute the bias
if (VERBOSE): mpi.msg("Computing bias")
k_bias, b_cdm, b_b = vbc_utils.compute_bias(ics, vbc)
# Convolve with field
if (VERBOSE): mpi.msg("Performing convolution")
delta_biased = vbc_utils.apply_density_bias(ics,
k_bias,
b_b,
delta.shape[0],
delta_x=delta)
# Remove the padded region
x_shape, y_shape, z_shape = delta_biased.shape
delta_biased = delta_biased[0 + pad:x_shape - pad,
0 + pad:y_shape - pad,
0 + pad:z_shape - pad]
# Store
biased_patch = Patch(patch, dx, delta_biased)
sto.result = biased_patch.__dict__
############################## END OF WORK LOOP ###############################
if mpi.host:
import os
# Write new ICs
dest = mpi.unpack(dest)
output_field = np.zeros(ics.header.nn)
for item in dest:
result = item.result
patch = result["patch"]
dx = result["dx"]
delta_biased = result["field"]
# Bounds of this patch
x_min, x_max = (int((patch[0]) - (dx / 2.)),
int((patch[0]) + (dx / 2.)))
y_min, y_max = (int((patch[1]) - (dx / 2.)),
int((patch[1]) + (dx / 2.)))
z_min, z_max = (int((patch[2]) - (dx / 2.)),
int((patch[2]) + (dx / 2.)))
# Place into output
output_field[x_min:x_max, y_min:y_max, z_min:z_max] = delta_biased
# Write the initial conditions
ics_dir = "%s/ics_ramses_vbc/" % ics.level_dir
if not os.path.isdir(ics_dir):
os.mkdir(ics_dir)
out_dir = "%s/level_%03i/" % (ics_dir, level)
if not os.path.isdir(out_dir):
os.mkdir(out_dir)
ics.write_field(output_field, "deltab", out_dir=out_dir)
int((patch[0]) + (dx / 2.)))
y_min, y_max = (int((patch[1]) - (dx / 2.)),
int((patch[1]) + (dx / 2.)))
z_min, z_max = (int((patch[2]) - (dx / 2.)),
int((patch[2]) + (dx / 2.)))
# Place into output
output_field[x_min:x_max, y_min:y_max, z_min:z_max] = delta_biased
# Write the initial conditions
ics_dir = "%s/ics_ramses_vbc/" % ics.level_dir
if not os.path.isdir(ics_dir):
os.mkdir(ics_dir)
out_dir = "%s/level_%03i/" % (ics_dir, level)
if not os.path.isdir(out_dir):
os.mkdir(out_dir)
ics.write_field(output_field, "deltab", out_dir=out_dir)
if __name__ == "__main__":
import sys
path = sys.argv[1]
level = int(sys.argv[2])
patch_size = float(sys.argv[3])
try:
main(path, level, patch_size)
except Exception as e:
from seren3.analysis.parallel import mpi
mpi.msg("Caught exception: %s" % e.message)
mpi.terminate(500, e=e)
# Place into output
output_field[x_min:x_max, y_min:y_max, z_min:z_max] = delta_biased
# Write the initial conditions
ics_dir = "%s/ics_ramses_vbc/" % ics.level_dir
if not os.path.isdir(ics_dir):
os.mkdir(ics_dir)
out_dir = "%s/level_%03i/" % (ics_dir, level)
if not os.path.isdir(out_dir):
os.mkdir(out_dir)
ics.write_field(output_field, "deltab", out_dir=out_dir)
if __name__ == "__main__":
import sys
import traceback
path = sys.argv[1]
level = int(sys.argv[2])
patch_size = float(sys.argv[3])
try:
main(path, level, patch_size)
except Exception as e:
from seren3.analysis.parallel import mpi
mpi.msg("Caught exception (message): %s" % e.message)
mpi.msg(traceback.format_exc())
mpi.terminate(500, e=e)
print("Done!")
# Write the initial conditions
ics_dir = "%s/ics_ramses_vbc/" % ics.level_dir
if not os.path.isdir(ics_dir):
os.mkdir(ics_dir)
out_dir = "%s/level_%03i/" % (ics_dir, level)
if not os.path.isdir(out_dir):
os.mkdir(out_dir)
ics.write_field(output_field, "deltab", out_dir=out_dir)
if __name__ == "__main__":
import sys
path = sys.argv[1]
level = int(sys.argv[2])
patch_size = float(sys.argv[3])
# Added by LC to stop empty file loading crashing the script
import warnings
warnings.simplefilter("ignore", UserWarning)
print("UserWarning set to ignore")
# main(path, level, patch_size)
try:
main(path, level, patch_size)
except Exception as e:
from seren3.analysis.parallel import mpi
mpi.msg("Caught exception (message): %s" % e.message)
mpi.msg("Caught exception (args): %s" % e.args)
mpi.terminate(500, e=e)
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" ) )
