def __init__(self, pymses_snapshot, finder, filename=None, **kwargs):
# import weakref
# self._base = weakref.ref(pymses_snapshot)
self._base = weakref.ref(pymses_snapshot)
self.finder = finder
self.finder_base_dir = "%s/%s" % (
self.base.path, config.get("halo",
"%s_base" % self.finder.lower()))
can_load, message = self.can_load(**kwargs)
if can_load is False:
print "Unable to load catalogue: %s" % message
logger.error("Unable to load catalogue: %s" % message)
return
self.filename = self.get_filename(
**kwargs) if filename is None else filename
self.boxsize = self.get_boxsize(**kwargs) # Mpccm/h
if (verbose):
print "%sCatalogue: loading halos..." % self.finder,
sys.stdout.flush()
self.load(**kwargs)
if (verbose):
print 'Loaded %d halos' % len(self)
def star_Nion_d(context, dset, dt=0., group=1):
'''
Computes the number of ionisiing photons produced by a stellar population per solar mass per second
'''
from seren3.array import SimArray
from seren3.utils.sed import io
from seren3.exceptions import NoParticlesException
from seren3 import config
# from seren3.analysis import interpolate
from scipy.interpolate import interp2d
verbose = config.get("general", "verbose")
Z_sun = 0.02 # metallicity of the sun
nGroups = context.info["nGroups"]
if (verbose):
print 'Computing Nion_d for photon group %i/%i' % (group, nGroups)
nIons = context.info["nIons"]
nPhotons_idx = 0 # index of photon number in SED
# Load the SED table
agebins, zbins, SEDs = io.read_seds_from_lists(context.path, nGroups,
nIons)
igroup = group - 1
fn = interp2d(zbins / Z_sun, agebins, SEDs[:, :, igroup, nPhotons_idx])
age = dset["age"].in_units("Gyr")
Z = dset["metal"] / Z_sun # in units of solar metalicity
# Which star particles should we keep
if dt != 0.:
age -= dt.in_units("Gyr")
# keep = np.where( np.logical_and(age >= 0., age.in_units("Myr") <= 10.) )
keep = np.where(age >= 0.)
age = age[keep]
Z = Z[keep]
if len(age) == 0:
raise NoParticlesException("No particles with (age - dt) > 0",
"star_Nion_d")
# interpolate photon production rate from SED
nStars = len(age)
nPhotons = np.zeros(nStars)
for i in xrange(nStars):
nPhotons[i] = fn(Z[i], age[i])
# nPhotons = interpolate.interpolate2d(age, Z, agebins, zbins, SEDs[:,:,igroup,nPhotons_idx])
# Multiply by (SSP) escape fraction and return
nml = context.nml
NML_KEYS = nml.NML
rt_esc_frac = float(nml[NML_KEYS.RT_PARAMS]['rt_esc_frac'].replace(
'd', 'e'))
Nion_d = SimArray(rt_esc_frac * nPhotons, "s**-1 Msol**-1")
Nion_d.set_field_latex("$\\dot{N_{\\mathrm{ion}}}$")
return Nion_d
def Okamoto_Mc_fn():
from seren3 import config
from scipy import interpolate
# from seren3.analysis.interpolate import extrap1d
fname = "%s/Mc_Okamoto08.txt" % config.get('data', 'data_dir')
data = np.loadtxt(fname)
ok_a, ok_z, ok_Mc = data.T
fn = interpolate.interp1d(ok_z, np.log10(ok_Mc), fill_value="extrapolate")
return lambda z: 10**fn(z)
def halos(self, finder=config.get("halo", "default_finder"), **kwargs):
if finder.lower() == 'ahf':
from seren3.halos.halos import AHFCatalogue
return AHFCatalogue(self, **kwargs)
elif finder.lower() == 'rockstar':
from seren3.halos.halos import RockstarCatalogue
return RockstarCatalogue(self, **kwargs)
elif finder.lower() == 'ctrees':
from seren3.halos.halos import ConsistentTreesCatalogue
return ConsistentTreesCatalogue(self, **kwargs)
else:
raise Exception("Unknown halo finder: %s" % finder)
def list_custom_cmaps():
import glob, re
from seren3 import config
_REGEX = ".+?(?=_256)"
cmap_dir = "%s/cmaps/" % config.get("data", "data_dir")
files = glob.glob("%s/*.txt" % cmap_dir)
cmap_names = []
for f in files:
m = re.search(_REGEX, f)
cmap_name = m.group(0).split("/")[-1]
cmap_names.append(cmap_name)
return cmap_names
def read_fan06():
'''
Reads the Fan et al. 2006 quasar observations for optical depth
'''
import csv
from seren3 import config
data_dir = config.get("data", "data_dir")
lyalpha_fname = "%s/obs/fan06_lyalpha.csv" % data_dir
#lybeta_fname = "%s/fan06_lybeta.csv" % data_dir
#lygamma_fname = "%s/fan06_lygamma.csv" % data_dir
# TODO - Make sure to apply Fan et al. 06 beta/gamma conversion factors
# fnames = [lyalpha_fname, lybeta_fname, lygamma_fname]
with open(lyalpha_fname, "rb") as csvfile:
reader = csv.reader(csvfile)
rows = []
for row in reader:
rows.append(row)
quasars = {}
key = None
for row in rows:
if row[0].startswith("J"):
key = row[0]
if key in quasars:
# quasars[key]["zem"].append(float(row[1]))
quasars[key]["zabs"].append(float(row[2]))
quasars[key]["T"].append(float(row[3]))
quasars[key]["sigma"].append(float(row[4]))
elif key:
quasars[key] = {
"zem": float(row[1]),
"zabs": [],
"T": [],
"sigma": []
}
return quasars
def load_custom_cmaps(cmap_name):
import numpy as np
import glob, re
from seren3 import config
if (cmap_name == "parula"):
from seren3.utils import parula
return parula.parula_map
cmap_dir = "%s/cmaps/" % config.get("data", "data_dir")
_REGEX = ".+?(?=_256)"
files = glob.glob("%s/*.txt" % cmap_dir)
for f in files:
m = re.search(_REGEX, f)
name = m.group(0).split("/")[-1]
if cmap_name == name:
C = np.loadtxt(f)
cm = matplotlib.colors.ListedColormap(C / 255.0)
return cm
raise IOError("No cmap found with name: %s" % cmap_name)
'''
Module for interpolating/calculating magnitues / luminosities.
Adapted from https://pynbody.github.io/pynbody/_modules/pynbody/analysis/luminosity.html to work with seren
'''
import numpy as np
import os
from seren3.analysis.interpolate import interpolate2d
from seren3 import config
_data_dir = config.get("data", "data_dir")
def plot_many(sims, iouts, labels, nbins):
import matplotlib.pylab as plt
bouwens_cols = ["k", "darkorange", "m"]
ax = plt.gca()
ax.plot([-99, -99], [-99, -99], color="r", linewidth=2., label=labels[0])
# ax.plot([-99, -99], [-99, -99], color="b", linewidth=2., label=labels[1])
for iout, bc in zip(iouts, bouwens_cols):
snap1 = sims[0][iout]
snap2 = sims[1][iout]
lums1 = calc_luminosities_halos(snap1, lambda_A=1600)
lums2 = calc_luminosities_halos(snap2, lambda_A=1600)
luminosity_function(lums1,
snap1.z,
def read_bouwens_2015():
import os
from seren3 import config
fname = "bouwens_2015.csv"
data_dir = os.path.join(config.get("data", "data_dir"), "obs/")
data = {}
with open("%s/%s" % (data_dir, fname), "Urb") as f:
import csv
reader = csv.reader(f)
z = np.inf
for line in reader:
if line[0].startswith('z'):
z = int(line[0][4])
data[z] = {'M1600': [], 'phi': [], 'err': []}
elif z is not np.inf:
M, phi, err = [l.strip() for l in line]
data[z]['M1600'].append(float(M))
if phi.endswith('b'):
# Upper limit
data[z]['phi'].append(phi)
data[z]['err'].append(np.nan)
else:
data[z]['phi'].append(float(phi))
data[z]['err'].append(float(err))
return data
# def plot_seds(ages=[0., 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 16., 25., 50., 100., 252., 502., 1005., 10005.]):
# '''
# Test function to ensure SED tables are read correctly
# '''
# import matplotlib.cm as cm
# import matplotlib.pylab as plt
# import os
# HI = 912 # A
# HeI = 505
# HeII = 228
# agebins, zbins, Ls, SEDs = read_seds(os.getenv("RAMSES_SED_DIR"))
# nLs = len(Ls)
# Z_sun = 0.02 # Solar metallicity
# z_idx = (np.abs(Z_sun - zbins)).argmin()
# lambda_range = np.linspace(100., 1000., 250)
# # ages = [0., 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 16.,
# # 25., 50., 100., 252., 502., 1005., 10005.] # Myr
# ages = [1., 10.] # , 100., 1000., 10000.] # Myr
# color = cm.rainbow(np.linspace(0, 1, len(ages)))
# ax = plt.subplot(111)
# for age, c in zip(ages, color):
# y = []
# age_idx = (np.abs(age * 1e6 - agebins)).argmin()
# for lambda_A in lambda_range:
# ii = 0
# while (ii < nLs) and (Ls[ii] < lambda_A):
# ii += 1
# y.append(SEDs[ii, age_idx, z_idx])
# ax.plot(lambda_range, y, color=c, label='BC03 %1.0f Myrs' %
# (age), linestyle='--', linewidth=3.)
# print 'PLOTTING BPASS'
# data = np.loadtxt(
# '/lustre/scratch/astro/ds381/SEDs/BPASS/SEDS/sed.bpass.instant.nocont.bin.z020').T
# idx = np.where(np.logical_and(data[0] >= 100., data[0] <= 1000.))
# ages = [1, 11] # , 21, 31, 41]
# labels = [6, 7, 8, 9, 10]
# c = plt.cm.rainbow(np.linspace(0, 1, len(ages)))
# for i, col in zip(range(len(ages)), c):
# plt.semilogy(data[0][idx], data[ages[i]][idx] / 1.e6, color=col, label=r'BPASS %1.0f Myr' %
# ((10.**labels[i]) / 10.**6), linewidth=1., linestyle='-')
# plt.legend()
# # Shrink current axis by 20%
# box = ax.get_position()
# ax.set_position([box.x0, box.y0, box.width * 0.9, box.height])
# # Put a legend to the right of the current axis
# ax.legend(loc='center left', bbox_to_anchor=(1, 0.5), prop={'size': 12.})
# ax.set_yscale('log')
# plt.xlabel(r'$\lambda$ [$\AA$]')
# plt.ylabel(r'J$_{\lambda}$ [L$_{\odot}$/M$_{\odot}$/$\AA$]')
# ax.axvline(x=HI, linestyle='--', color='k')
# ax.axvline(x=HeI, linestyle='--', color='k')
# ax.axvline(x=HeII, linestyle='--', color='k')
# plt.show()
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
from seren3 import config
_IMAGE_DIR = config.get("data", "movie_img_dir")
def _get_fname(family, field):
return "%s/_tmp_%s_%s_%05i.png" % (_IMAGE_DIR, family.family, field,
family.ioutput)
def _cleanup():
import os
#os.system("rm %s/_tmp*.png" % _IMAGE_DIR)
def _run_mencoder(out_fname, fps, remove_files=True):
from seren3.utils import which
import subprocess, os
mencoder = which("mencoder")
args = "-mf w=800:h=600:fps=%i:type=png -ovc lavc -lavcopts vcodec=mpeg4:mbd=2:trell -oac copy" % fps
# Make the movie
os.chdir(_IMAGE_DIR)
exe = "{mencoder} mf://_tmp*.png {args} -o {out_dir}/{out_fname}".format(mencoder=mencoder, \
out_dir=_IMAGE_DIR, args=args, out_fname=out_fname)
p = subprocess.check_output(exe, shell=True)
# Remove _tmp files
if remove_files:
_cleanup()
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" ) )
if __name__ == "__main__":
import sys
from seren3 import config
path = sys.argv[1]
iout = int(sys.argv[2])
finder = config.get("halo", "default_finder")
pickle_path = None
if len(sys.argv) > 3:
finder = sys.argv[3]
if len(sys.argv) > 4:
pickle_path = sys.argv[4]
main(path, iout, finder, pickle_path)
# try:
# main(path, iout, finder, pickle_path)
# except Exception as e:
# from seren3.analysis.parallel import mpi
# print 'Caught exception: ', e
# mpi.terminate(500, e=e)
Heavily based on halos.py from pynbody, so credit to those guys
Rewritten to allow loading of Rockstar catalogues when using any module
@author dsullivan, bthompson
'''
import seren3
from seren3 import config
from seren3.core.snapshot import Family
from seren3.array import SimArray
import numpy as np
import logging
import abc
import sys
import weakref
verbose = config.get("general", "verbose")
logger = logging.getLogger('seren3.halos.halos')
class Halo(object):
"""
Object to represent a halo and allow filtered data access
"""
def __init__(self, properties, snapshot, units, boxsize):
self.hid = properties["id"]
self.properties = properties
self.units = units
self.boxsize = boxsize
self._base = weakref.ref(snapshot)
self._subsnap = None
