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