def jeans_scale_interp_fn(aexp, jeans_scale): from scipy import interpolate from seren3.analysis.plots import fit_scatter bc, mean, std = fit_scatter(np.log10(aexp), np.log10(jeans_scale), nbins=10) fn = interpolate.interp1d(bc, mean, fill_value="extrapolate") return fn
def binned_by_density(density, field, xlogrange=True, ylogrange=True, thresh=None, den_field='nH', nbins=100): ''' Bin a quantity by density thresh: lambda function to return indicies of cells to keep using np.where ''' import numpy as np from seren3.analysis.plots import fit_scatter if thresh: idx = thresh(density, field) field = field[idx] density = density[idx] if xlogrange: density = np.log10(density) if ylogrange: field = np.log10(field) bin_centres, mean, std = fit_scatter(density, field, nbins=nbins) return bin_centres, mean, std
def plot(sims, ioutputs, labels, colours, pickle_paths=None): import pickle import numpy as np import matplotlib.pylab as plt from seren3.analysis import plots if (pickle_paths is None): pickle_paths = ["%s/pickle/" % sim.path for sim in sims] fig, ax = plt.subplots() for sim, ioutput, ppath, label, c in zip(sims, ioutputs, pickle_paths, labels, colours): fname = "%s/ConsistentTrees/halo_clumping_factor_%05i.p" % (ppath, ioutput) data = pickle.load( open(fname, "rb") ) mvir = np.zeros(len(data)) C = np.zeros(len(data)) for i in range(len(data)): res = data[i].result mvir[i] = res["hprops"]["mvir"] C[i] = res["C"] bc, mean, std = plots.fit_scatter(np.log10(mvir), np.log10(C)) ax.scatter(mvir, np.log10(C), color=c, s=15, marker="+", alpha=0.25) e = ax.errorbar(10**bc, mean, yerr=std, color=c, label=label,\ fmt="o", markerfacecolor=c, mec='k', capsize=2, capthick=2, elinewidth=2, linestyle='-') ax.set_xlabel(r"M$_{\mathrm{vir}}$ [M$_{\odot}$/h]") ax.set_ylabel(r"log$_{10}$ $C$") ax.set_xscale("log") ax.legend()
def nH_halo_cdf_fit(snapshot, the_mass_bins=[8., 9., 10.], nbins=10): import numpy as np from seren3.analysis.plots import fit_scatter from seren3.utils import flatten_nested_array # First, compute min/max nH in the snapshot print "Bracketing nH" min_nH = np.inf max_nH = -np.inf for dset in snapshot.g["nH"]: nH = dset["nH"] if nH.min() < min_nH: min_nH = nH.min() if nH.max() > max_nH: max_nH = nH.max() nH_range = (min_nH, max_nH) halos = snapshot.halos(finder='ctrees') nhalos = len(halos) cdf_halos = [] bin_centre_halos = [] Mvir = [] for i in range(nhalos): print '(%i / %i)' % (i + 1, nhalos) h = halos[i] if len(h.g) > 0: P, C, bc, dx = pdf_cdf(h.g, "nH", cumulative=True, plot=False, bins=50, x_range=nH_range, density=True) cdf_halos.append(C) bin_centre_halos.append(bc) Mvir.append(h["Mvir"]) cdf_halos = np.array(cdf_halos) bin_centre_halos = np.array(bin_centre_halos) Mvir = np.array(Mvir) # Bin idx mass_bins = np.digitize(np.log10(Mvir), the_mass_bins, right=True) binned_cdf = {} for i in range(len(the_mass_bins)): idx = np.where(mass_bins == i) x, y = (flatten_nested_array(bin_centre_halos[idx]), flatten_nested_array(cdf_halos[idx])) binned_cdf[i] = fit_scatter(x, y, ret_sterr=True) binned_cdf['the_mass_bins'] = the_mass_bins return binned_cdf
def _plot(fesc, tot_mass, label, alpha, ax, c, nbins): import numpy as np import random # from seren3.analysis.plots import fit_scatter from seren3.analysis import plots reload(plots) print len(fesc) # Deal with fesc>1. following Kimm & Cen 2014 bad = np.where( fesc > 1. ) nbad = float(len(bad[0])) print "%f %% of points above 1" % ((nbad/float(len(fesc)))*100.) for ix in bad: fesc[ix] = random.uniform(0.9, 1.0) keep = np.where( np.logical_and(fesc >= 0., fesc <=1.) ) fesc = fesc[keep] tot_mass = tot_mass[keep] log_tot_mass = np.log10(tot_mass) log_fesc = np.log10(100. * fesc) fesc_percent = fesc * 100. remove = np.where( np.logical_or( np.isinf(log_fesc), np.isnan(log_fesc) ) ) log_fesc = np.delete( log_fesc, remove ) fesc_percent = np.delete( fesc_percent, remove ) log_tot_mass = np.delete( log_tot_mass, remove ) remove = np.where( log_fesc < -1. ) log_fesc = np.delete( log_fesc, remove ) fesc_percent = np.delete( fesc_percent, remove ) log_tot_mass = np.delete( log_tot_mass, remove ) print len(log_fesc) # bin_centres, mean, std_dev, std_err = plots.fit_scatter(log_tot_mass, log_fesc, nbins=nbins, ret_sterr=True) bin_centres, mean, std_dev, std_err = plots.fit_scatter(log_tot_mass, fesc_percent, nbins=nbins, ret_sterr=True) # bin_centres, median = plots.fit_median(log_tot_mass, fesc_percent, nbins=nbins) # print bin_centres, mean # Plot ax.scatter(log_tot_mass, fesc_percent, marker='+', color=c, alpha=alpha) # ax.errorbar(bin_centres, mean, yerr=std_dev, color=c, label=label, linewidth=2.) e = ax.errorbar(bin_centres, mean, yerr=std_dev, color=c, label=label,\ fmt="o", markerfacecolor=c, mec='k', capsize=2, capthick=2, elinewidth=2, linestyle='-') # e = ax.plot(bin_centres, median, color=c, label=label, linestyle='-', linewidth=2.) ax.set_yscale("log")
def plot(sims, iout, labels, cols, ax=None, **kwargs): import numpy as np import matplotlib.pylab as plt from seren3.analysis import plots if (ax is None): ax = plt.gca() ls = ["-", "--"] lw = [3., 1.5] for sim, label, col, lsi, lwi in zip(sims, labels, cols, ls, lw): snap = sim[iout] nphotons_escaped, tot_mass_outflowed, mvir = fesc_tot_outflow(snap) print "%e" % nphotons_escaped.sum() log_mvir = np.log10(mvir) x = np.log10(tot_mass_outflowed) y = np.log10(nphotons_escaped) ix = np.where(np.logical_and(log_mvir >= 7.5, x >= 5.5)) x = x[ix] y = y[ix] ix = np.where(np.logical_and(np.isfinite(x), np.isfinite(y))) x = x[ix] y = y[ix] bc, mean, std, sterr = plots.fit_scatter(x, y, ret_sterr=True, **kwargs) ax.scatter(x, y, alpha=0.10, s=5, color=col) e = ax.errorbar(bc, mean, yerr=std, color=col, label=label,\ fmt="o", markerfacecolor=col, mec='k',\ capsize=2, capthick=2, elinewidth=2, linewidth=lwi, linestyle=lsi) # ax.plot(bc, mean, color=col, label=None, linewidth=3., linestyle="-") # ax.fill_between(bc, mean-std, mean+std, facecolor=col, alpha=0.35, interpolate=True, label=label) ax.set_xlabel( r"log$_{10}$ $\int_{0}^{t_{\mathrm{H}}}$ $\vec{F}_{+}(t)$ $dt$ [M$_{\odot}$]", fontsize=20) ax.set_ylabel( r'log$_{10}$ $\int_{0}^{t_{\mathrm{H}}}$ $\dot{\mathrm{N}}_{\mathrm{ion}}(t)$ f$_{\mathrm{esc}}$ ($t$) $dt$ [#]', fontsize=20) ax.legend(loc='lower right', frameon=False, prop={"size": 16})
def read_and_smooth_bpass_seds(sed_type="bin", nbins=6900): from seren3.analysis import plots agebins, zbins, Ls, SEDs = read_bpass_seds(sed_type=sed_type) Ls_smoothed = np.zeros(nbins) SEDs_smoothed = np.zeros((nbins, len(agebins), len(zbins))) for i in range(len(agebins)): for j in range(len(zbins)): bc, mean, std = plots.fit_scatter(Ls, np.log10(SEDs[:,i,j]), nbins=nbins) SEDs_smoothed[:,i,j] = 10**mean Ls_smoothed = bc return agebins, zbins, Ls_smoothed, SEDs_smoothed
def plot_integrated_nion_esc(snapshot, ax=None, **kwargs): ''' Plot cumulative photons escaped as a function of halo mass ''' import matplotlib.pylab as plt from seren3.analysis import plots if (ax is None): ax = plt.gca() color = kwargs.pop("color", "k") nbins = kwargs.pop("nbins", 5) label = kwargs.pop("label", None) ls = kwargs.pop("ls", "-") lw = kwargs.pop("lw", 1.) mvir, nphotons_escaped = compute_integrated_nion_esc(snapshot) print "%e" % nphotons_escaped.sum() log_mvir = np.log10(mvir) log_nphotons_escaped = np.log10(nphotons_escaped) ix = np.where( np.logical_and(np.isfinite(log_nphotons_escaped), log_mvir >= 7.5) ) log_nphotons_escaped = log_nphotons_escaped[ix] log_mvir = log_mvir[ix] bc, mean, std, sterr = plots.fit_scatter(log_mvir, log_nphotons_escaped, nbins=nbins, ret_sterr=True) ax.scatter(log_mvir, log_nphotons_escaped, s=10, color=color, alpha=0.1) e = ax.errorbar(bc, mean, yerr=std, color=color, label=label,\ fmt="o", markerfacecolor=color, mec='k', capsize=2, capthick=2, elinewidth=2, linestyle=ls, linewidth=lw) if (kwargs.pop("legend", False)): ax.legend(loc="lower right", frameon=False, prop={"size":16}) ax.set_xlabel(r'log$_{10}$ M$_{\mathrm{vir}}$ [M$_{\odot}$/h]', fontsize=20) ax.set_ylabel(r'log$_{10}$ $\int_{0}^{t_{\mathrm{H}}}$ $\dot{\mathrm{N}}_{\mathrm{ion}}(t)$ f$_{\mathrm{esc}}$ ($t$) $dt$ [#]', fontsize=20)
def plot(simulations, ioutputs, labels, colours, nbins=10, plot_baryon_fraction=False, compare=True, dm_particle_cutoff=100, star_particle_cutoff=1, pickle_paths=None): import pickle from pynbody.plot.stars import moster, behroozi import matplotlib.pylab as plt from seren3.analysis.plots import fit_scatter import matplotlib matplotlib.rcParams['axes.linewidth'] = 1.5 matplotlib.rcParams['xtick.labelsize'] = 18 matplotlib.rcParams['ytick.labelsize'] = 18 matplotlib.rcParams['axes.labelsize'] = 22 if (pickle_paths is None): pickle_paths = [ "%s/pickle/ConsistentTrees/" % sim.path for sim in simulations ] fig, ax = plt.subplots(figsize=(8, 8)) # fig, ax = plt.subplots() for sim, ioutput, label, c, ppath in zip(simulations, ioutputs, labels, colours, pickle_paths): snap = sim[ioutput] cosmo = snap.cosmo cosmic_mean = cosmo["omega_b_0"] / cosmo["omega_M_0"] data = None fname = "%s/abundance_%05i.p" % (ppath, snap.ioutput) with open(fname, 'rb') as f: data = pickle.load(f) nrecords = len(data) mass = np.zeros(nrecords) stellar_mass = np.zeros(nrecords) np_dm = np.zeros(nrecords) np_star = np.zeros(nrecords) for i in range(nrecords): res = data[i].result mass[i] = res["tot_mass"] stellar_mass[i] = res["star_mass"] np_dm[i] = res["np_dm"] np_star[i] = res["np_star"] del res # print np_star.min(), np_star.max() # idx = np.where(np.logical_and( np_dm >= dm_particle_cutoff, np_star >= star_particle_cutoff )) idx = np.where( np.logical_and( np.log10(mass) >= 7.5, np_star >= star_particle_cutoff)) mass = mass[idx] stellar_mass = stellar_mass[idx] print np.log10(mass.min()), np.log10(mass.max()) stellar_mass = (stellar_mass / mass) / cosmic_mean stellar_mass *= 100 # % bc, mean, std = fit_scatter(mass, stellar_mass, nbins=10) bc, mean, log_std, log_sterr = fit_scatter(np.log10(mass), np.log10(stellar_mass), ret_sterr=True, nbins=nbins) # std = (log_std/0.434) * 10**mean # sterr = (log_sterr/0.434) * 10**mean if compare: x = snap.array(np.logspace(np.log10(2.5e7), np.log10(1.5e10), 100), _MASS_UNIT) ystarmasses, errors = moster(x.in_units("Msol"), snap.z) ystarmasses = snap.array(ystarmasses, "Msol").in_units(_MASS_UNIT) ystarmasses = (ystarmasses / x) / cosmic_mean ystarmasses *= 100 # print errors.min(), errors.max() moster_c = "#BBBBBB" ax.fill_between(x, np.log10(np.array(ystarmasses) / np.array(errors)), y2=np.log10( np.array(ystarmasses) * np.array(errors)), facecolor=moster_c, color=moster_c, label='Moster et al (2013)', alpha=0.5) idx_extrap = np.where(np.log10(x) < np.log10(1e10 * 0.702)) idx_no_extrap = np.where(np.log10(x) >= np.log10(1e10 * 0.702)) ax.plot(x[idx_no_extrap], np.log10(ystarmasses)[idx_no_extrap], color="#909497", linewidth=5.) ax.plot(x[idx_extrap], np.log10(ystarmasses)[idx_extrap], color="#909497", linestyle="--", linewidth=5.) # behroozi_c = "#F08080" behroozi_c = "lightskyblue" ystarmasses, errors = behroozi(x.in_units("Msol"), snap.z) ystarmasses = snap.array(ystarmasses, "Msol").in_units(_MASS_UNIT) ystarmasses = (ystarmasses / x) / cosmic_mean ystarmasses *= 100 ax.fill_between(x, np.log10(np.array(ystarmasses) / np.array(errors)), y2=np.log10( np.array(ystarmasses) * np.array(errors)), facecolor=behroozi_c, color=behroozi_c, label='Behroozi et al (2013)', alpha=0.5) idx_extrap = np.where(np.log10(x) < np.log10(1e10 * 0.702)) idx_no_extrap = np.where(np.log10(x) >= np.log10(1e10 * 0.702)) ax.plot(x[idx_no_extrap], np.log10(ystarmasses)[idx_no_extrap], color="#3498DB", linewidth=5.) ax.plot(x[idx_extrap], np.log10(ystarmasses)[idx_extrap], color="#3498DB", linestyle="--", linewidth=5.) compare = False if plot_baryon_fraction: baryon_fraction = x * cosmic_mean y = ((0.1 * baryon_fraction) / x) / cosmic_mean y *= 100. # ax.loglog(mass, baryon_fraction, linestyle='dotted', label=r'$\Omega_{\mathrm{b}} / \Omega_{\mathrm{M}}$') ax.plot(x, y, linestyle='dotted', linewidth=3., label=r'0.1 $\Omega_{\mathrm{b}} / \Omega_{\mathrm{M}}$') plot_baryon_fraction = False # ax.scatter(mass, stellar_mass, color=c, alpha=0.4, s=15) # ax.errorbar(10**bc, 10**mean, yerr=std, color=c, linewidth=3., linestyle="-", zorder=10, label=label) # ax.plot(10**bc, 10**mean, color=c, linewidth=3., linestyle="-", zorder=10, label=label) # ax.plot(10**bc, 10**(mean+std), color=c, linewidth=3., linestyle="--", zorder=10) # ax.plot(10**bc, 10**(mean-std), color=c, linewidth=2., linestyle="--", zorder=10) # ax.plot(10**bc, mean, color=c, linewidth=3., linestyle="-", zorder=10, label=label) # ax.plot(10**bc, mean + log_std, color=c, linewidth=1.5, linestyle="--", zorder=10) # ax.plot(10**bc, mean - log_std, color=c, linewidth=1.5, linestyle="--", zorder=10) # ax.errorbar(10**bc, mean, yerr=log_std, color=c, linewidth=3., linestyle="none", zorder=10, label=label) e = ax.errorbar(10**bc[1:], mean[1:], yerr=log_std[1:], color=c, label=label,\ fmt="o", markerfacecolor=c, mec='k', capsize=2, capthick=2, elinewidth=2, linestyle='-', linewidth=2.) text_pos = (1.e8, np.log10(1)) text = 'z = %1.2f' % snap.z ax.set_xlabel(r'M$_{\mathrm{vir}}$ [M$_{\odot}$/h]') # ax.set_ylabel(r'M$_{*}$ [M$_{\odot}$/h]') ax.set_ylabel( r'log$_{10}$(M$_{*}$/M$_{\mathrm{h}}$)/($\Omega_{\mathrm{b}}$/$\Omega_{\mathrm{M}}$) [%]' ) # box = ax.get_position() # ax.set_position([box.x0, box.y0 + box.height * 0.1, # box.width, box.height * 0.9]) # # Put a legend below current axis # ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.13), # fancybox=True, shadow=False, ncol=5, prop={'size':10.5}) ax.legend(prop={'size': 15.}, loc="lower right") ax.set_xscale("log") #; ax.set_yscale("log") ax.text(text_pos[0], text_pos[1], text, color="k", size=18) ax.set_xlim(6.5e7, x.max()) plt.tight_layout() plt.savefig("/home/ds381/RTX_HD_stellar_abundance.pdf", format="pdf") plt.show()
def plot_Mc_z_xHII(simulations, simulation_ioutputs, pickle_paths, ann_out_paths, labels, \ colours, NN, fix_alpha=False, use_lmfit=True, **kwargs): ''' Plot the evolution of Mc against z and xHII ''' import matplotlib.pylab as plt from seren3.cosmology import hoeft_Mc from seren3.analysis import plots from seren3.analysis.plots import reion from seren3.analysis.baryon_fraction import tidal_force from scipy.interpolate import interp1d import matplotlib.patheffects as path_effects reload(baryon_fraction) reload(neural_net2) reload(tidal_force) _Y_LIM = (5e6, 2e8) ANN_iouts = [106, 100, 90, 80, 70, 60, 48, 42] weight = kwargs.pop("weight", "mw") # pdf_sampling = kwargs.pop("pdf_sampling", False) # filter_ftidal = kwargs.pop("filter_ftidal", False) filter_ftidal = False pdf_sampling = filter_ftidal print "filter tidal force? ", filter_ftidal fig, axes = plt.subplots(2, 2, figsize=(10, 10)) axs = axes[0, :] err_str = "stderr" if use_lmfit else "sigma" for simulation, ioutputs, ppath, ann_out_path, label, color in zip(simulations, \ simulation_ioutputs, pickle_paths, ann_out_paths, labels, colours): print simulation z_xHII, xHII_vw, xHII_mw = reion.load_reionization_history( simulation, pickle_path=ppath) # Keep only z<18 idx = np.where(z_xHII <= 18.) z_xHII = z_xHII[idx] xHII_vw = xHII_vw[idx] xHII_mw = xHII_mw[idx] # Fit and interpolate reion. history bc, mean, std, stderr = plots.fit_scatter(z_xHII, xHII_mw, nbins=25, ret_sterr=True) fn_xHII = interp1d(z_xHII, xHII_mw, fill_value="extrapolate") # fn_xHII = interp1d(bc, mean, fill_value="extrapolate") fn_xHII_std = interp1d(bc, std, fill_value="extrapolate") # Compute Mc at each output z = np.zeros(len(ioutputs)) Mc = np.zeros(len(ioutputs)) Mc_err = np.zeros(len(ioutputs)) # Mc_ann = np.zeros(len(ioutputs)) # Mc_ann_err = np.zeros(len(ioutputs)) z_ann = np.zeros(len(ANN_iouts)) Mc_ann = np.zeros(len(ANN_iouts)) Mc_ann_err = np.zeros(len(ANN_iouts)) z_ann_ftidal_pdf = np.zeros(len(ANN_iouts)) Mc_ann_ftidal_pdf = np.zeros(len(ANN_iouts)) Mc_ann_err_ftidal_pdf = np.zeros(len(ANN_iouts)) ANN_count = 0 for i in range(len(ioutputs)): ioutput = ioutputs[i] snapshot = simulation[ioutput] cosmo = snapshot.cosmo cosmic_mean_b = cosmo["omega_b_0"] / cosmo["omega_M_0"] # hids, mvir, fb, tidal_force_tdyn, pid, np_dm, ncell = filter_and_load_data(snapshot, pickle_path=ppath) log_mvir, fb, ftidal, xHII, T, T_U, pid = neural_net2.load_training_arrays( snapshot, pickle_path=ppath, weight="mw") # Filter ftidal? if filter_ftidal: P, C, bincenters, dx, x_pdf, y_2_pdf, ( ftidal_pdf, indexes, peaks_x, params, sigma) = tidal_force.tidal_force_pdf(snapshot) idx = np.where(ftidal <= 10**peaks_x[-1]) # idx = np.where(np.logical_and(ftidal < 0.25, pid == -1)) log_mvir = log_mvir[idx] fb = fb[idx] ftidal = ftidal[idx] xHII = xHII[idx] T = T[idx] T_U = T_U[idx] pid = pid[idx] mvir = 10**log_mvir Mc_fit_dict = baryon_fraction.fit(mvir, fb, fix_alpha, use_lmfit=use_lmfit, **cosmo) z[i] = snapshot.z Mc[i] = Mc_fit_dict["Mc"]["fit"] Mc_err[i] = Mc_fit_dict["Mc"][err_str] * 5 # 3 sigma # alpha = Mc_fit_dict["alpha"]["fit"] # Neural net if ioutput in ANN_iouts: print simulation, ioutput, ann_out_path, NN ann_Mc_fit_dict = neural_net2.compute_ANN_Mc( simulation, ioutput, ann_out_path, NN, use_lmfit=use_lmfit, fix_alpha=fix_alpha, pdf_sampling=False) # ann_Mc_fit_dict = neural_net2.compute_ANN_Mc(simulation, ioutput, ann_out_path, NN, alpha=alpha, use_lmfit=use_lmfit, fix_alpha=True) z_ann[ANN_count] = snapshot.z Mc_ann[ANN_count] = ann_Mc_fit_dict["Mc"]["fit"] Mc_ann_err[ ANN_count] = ann_Mc_fit_dict["Mc"][err_str] * 5 # 3 sigma ann_Mc_fit_dict = neural_net2.compute_ANN_Mc( simulation, ioutput, ann_out_path, NN, use_lmfit=use_lmfit, fix_alpha=fix_alpha, pdf_sampling=True) z_ann_ftidal_pdf[ANN_count] = snapshot.z Mc_ann_ftidal_pdf[ANN_count] = ann_Mc_fit_dict["Mc"]["fit"] Mc_ann_err_ftidal_pdf[ ANN_count] = ann_Mc_fit_dict["Mc"][err_str] * 5 # 3 sigma ANN_count += 1 # Plot # axs[0].errorbar(z, Mc, yerr=Mc_err, linewidth=2., color=color) axs[0].fill_between(z, Mc - Mc_err, Mc + Mc_err, facecolor=color, alpha=0.35, interpolate=True, label=label) #, transform=trans) if not pdf_sampling: e = axs[0].errorbar(z_ann, Mc_ann, yerr=Mc_ann_err, color=color, label="%s ANN" % label,\ fmt="o", markerfacecolor=color, mec='k', capsize=2, capthick=2, elinewidth=2, linestyle='None') else: e = axs[0].errorbar(z_ann, Mc_ann_ftidal_pdf, yerr=Mc_ann_err_ftidal_pdf, color=color, label="%s ANN" % label,\ fmt="o", markerfacecolor=color, mec='k', capsize=2, capthick=2, elinewidth=2, linestyle='None') # e[1][0].set_path_effects([path_effects.Stroke(linewidth=4, foreground='black'), # path_effects.Normal()]) # e[1][1].set_path_effects([path_effects.Stroke(linewidth=4, foreground='black'), # path_effects.Normal()]) # e[2][0].set_path_effects([path_effects.Stroke(linewidth=4, foreground='black'), # path_effects.Normal()]) axs[1].fill_between(fn_xHII(z), Mc - Mc_err, Mc + Mc_err, facecolor=color, alpha=0.35, interpolate=True, label=label) #, transform=trans) # axs[1].errorbar(fn_xHII(z), Mc, xerr=fn_xHII_std(z), yerr=Mc_err, label=label, linewidth=2., color=color) if not pdf_sampling: e = axs[1].errorbar(fn_xHII(z_ann), Mc_ann, xerr=fn_xHII_std(z_ann), yerr=Mc_ann_err, color=color, label="%s ANN" % label,\ fmt="o", markerfacecolor=color, mec='k', capsize=2, capthick=2, elinewidth=2, linestyle='-') else: e = axs[1].errorbar(fn_xHII(z_ann), Mc_ann_ftidal_pdf, xerr=fn_xHII_std(z_ann), yerr=Mc_ann_err_ftidal_pdf, color=color, label="%s ANN" % label,\ fmt="o", markerfacecolor=color, mec='k', capsize=2, capthick=2, elinewidth=2, linestyle='-') # c2 = None # if color == "b": # c2 = "c" # elif color == "g": # c2 = "k" # else: # c2 = "m" # e = axs[1].errorbar(fn_xHII(z_ann_ftidal_pdf), Mc_ann_ftidal_pdf, xerr=fn_xHII_std(z_ann), yerr=Mc_ann_err_ftidal_pdf, color=c2, label="%s ANN PDF SAMPLED" % label,\ # fmt="o", markerfacecolor=c2, mec='k', capsize=2, capthick=2, elinewidth=2, linestyle='-.', linewidth=3.) # Plot the Ok08 fit Ok_z = np.linspace(6, 10, 100) Ok_fn = baryon_fraction.Okamoto_Mc_fn() Ok_Mc = np.array([Ok_fn(i) for i in Ok_z]) axs[0].plot(Ok_z, Ok_Mc, color='k', label="Okamoto et al. 08", linestyle='-.') Hoeft_Mc = [hoeft_Mc(zi, omega_m=cosmo["omega_M_0"]) for zi in Ok_z] axs[0].plot(Ok_z, Hoeft_Mc, color='m', label="Hoeft et al. 06", linestyle='-.') # Axis labels and limits axs[0].set_xlabel(r"$z$") axs[0].set_ylabel(r"$M_{\mathrm{c}}$ [M$_{\odot}$/h]") axs[0].set_xlim(5.5, 12.5) axs[0].set_ylim(_Y_LIM[0], _Y_LIM[1]) axs[1].set_xlabel(r"$\langle x_{\mathrm{HII}} \rangle_{\mathrm{M}}$") axs[1].set_ylabel(r"$M_{\mathrm{c}}$ [M$_{\odot}$/h]") # axs[1].set_xlim(0.0, 1.05) axs[0].set_ylim(_Y_LIM[0], _Y_LIM[1]) for ax in axs.flatten(): ax.set_yscale("log") ax.legend(prop={'size': 10}) # Tidal cutoff axs = axes[1, :] filter_ftidal = True pdf_sampling = filter_ftidal for simulation, ioutputs, ppath, ann_out_path, label, color in zip(simulations, \ simulation_ioutputs, pickle_paths, ann_out_paths, labels, colours): print simulation z_xHII, xHII_vw, xHII_mw = reion.load_reionization_history( simulation, pickle_path=ppath) # Keep only z<18 idx = np.where(z_xHII <= 18.) z_xHII = z_xHII[idx] xHII_vw = xHII_vw[idx] xHII_mw = xHII_mw[idx] # Fit and interpolate reion. history bc, mean, std, stderr = plots.fit_scatter(z_xHII, xHII_mw, nbins=25, ret_sterr=True) fn_xHII = interp1d(z_xHII, xHII_mw, fill_value="extrapolate") # fn_xHII = interp1d(bc, mean, fill_value="extrapolate") fn_xHII_std = interp1d(bc, std, fill_value="extrapolate") # Compute Mc at each output z = np.zeros(len(ioutputs)) Mc = np.zeros(len(ioutputs)) Mc_err = np.zeros(len(ioutputs)) # Mc_ann = np.zeros(len(ioutputs)) # Mc_ann_err = np.zeros(len(ioutputs)) z_ann = np.zeros(len(ANN_iouts)) Mc_ann = np.zeros(len(ANN_iouts)) Mc_ann_err = np.zeros(len(ANN_iouts)) z_ann_ftidal_pdf = np.zeros(len(ANN_iouts)) Mc_ann_ftidal_pdf = np.zeros(len(ANN_iouts)) Mc_ann_err_ftidal_pdf = np.zeros(len(ANN_iouts)) ANN_count = 0 for i in range(len(ioutputs)): ioutput = ioutputs[i] snapshot = simulation[ioutput] cosmo = snapshot.cosmo cosmic_mean_b = cosmo["omega_b_0"] / cosmo["omega_M_0"] # hids, mvir, fb, tidal_force_tdyn, pid, np_dm, ncell = filter_and_load_data(snapshot, pickle_path=ppath) log_mvir, fb, ftidal, xHII, T, T_U, pid = neural_net2.load_training_arrays( snapshot, pickle_path=ppath, weight="mw") # Filter ftidal? if filter_ftidal: P, C, bincenters, dx, x_pdf, y_2_pdf, ( ftidal_pdf, indexes, peaks_x, params, sigma) = tidal_force.tidal_force_pdf(snapshot) idx = np.where(ftidal <= 10**peaks_x[-1]) # idx = np.where(np.logical_and(ftidal < 0.25, pid == -1)) log_mvir = log_mvir[idx] fb = fb[idx] ftidal = ftidal[idx] xHII = xHII[idx] T = T[idx] T_U = T_U[idx] pid = pid[idx] mvir = 10**log_mvir Mc_fit_dict = baryon_fraction.fit(mvir, fb, fix_alpha, use_lmfit=use_lmfit, **cosmo) z[i] = snapshot.z Mc[i] = Mc_fit_dict["Mc"]["fit"] Mc_err[i] = Mc_fit_dict["Mc"][err_str] * 5 # 3 sigma # alpha = Mc_fit_dict["alpha"]["fit"] # Neural net if ioutput in ANN_iouts: print simulation, ioutput, ann_out_path, NN ann_Mc_fit_dict = neural_net2.compute_ANN_Mc( simulation, ioutput, ann_out_path, NN, use_lmfit=use_lmfit, fix_alpha=fix_alpha, pdf_sampling=False) # ann_Mc_fit_dict = neural_net2.compute_ANN_Mc(simulation, ioutput, ann_out_path, NN, alpha=alpha, use_lmfit=use_lmfit, fix_alpha=True) z_ann[ANN_count] = snapshot.z Mc_ann[ANN_count] = ann_Mc_fit_dict["Mc"]["fit"] Mc_ann_err[ ANN_count] = ann_Mc_fit_dict["Mc"][err_str] * 5 # 3 sigma ann_Mc_fit_dict = neural_net2.compute_ANN_Mc( simulation, ioutput, ann_out_path, NN, use_lmfit=use_lmfit, fix_alpha=fix_alpha, pdf_sampling=True) z_ann_ftidal_pdf[ANN_count] = snapshot.z Mc_ann_ftidal_pdf[ANN_count] = ann_Mc_fit_dict["Mc"]["fit"] Mc_ann_err_ftidal_pdf[ ANN_count] = ann_Mc_fit_dict["Mc"][err_str] * 5 # 3 sigma ANN_count += 1 # Plot # axs[0].errorbar(z, Mc, yerr=Mc_err, linewidth=2., color=color) axs[0].fill_between(z, Mc - Mc_err, Mc + Mc_err, facecolor=color, alpha=0.35, interpolate=True, label=label) #, transform=trans) if not pdf_sampling: e = axs[0].errorbar(z_ann, Mc_ann, yerr=Mc_ann_err, color=color, label="%s ANN" % label,\ fmt="o", markerfacecolor=color, mec='k', capsize=2, capthick=2, elinewidth=2, linestyle='None') else: e = axs[0].errorbar(z_ann, Mc_ann_ftidal_pdf, yerr=Mc_ann_err_ftidal_pdf, color=color, label="%s ANN" % label,\ fmt="o", markerfacecolor=color, mec='k', capsize=2, capthick=2, elinewidth=2, linestyle='None') # e[1][0].set_path_effects([path_effects.Stroke(linewidth=4, foreground='black'), # path_effects.Normal()]) # e[1][1].set_path_effects([path_effects.Stroke(linewidth=4, foreground='black'), # path_effects.Normal()]) # e[2][0].set_path_effects([path_effects.Stroke(linewidth=4, foreground='black'), # path_effects.Normal()]) axs[1].fill_between(fn_xHII(z), Mc - Mc_err, Mc + Mc_err, facecolor=color, alpha=0.35, interpolate=True, label=label) #, transform=trans) # axs[1].errorbar(fn_xHII(z), Mc, xerr=fn_xHII_std(z), yerr=Mc_err, label=label, linewidth=2., color=color) if not pdf_sampling: e = axs[1].errorbar(fn_xHII(z_ann), Mc_ann, xerr=fn_xHII_std(z_ann), yerr=Mc_ann_err, color=color, label="%s ANN" % label,\ fmt="o", markerfacecolor=color, mec='k', capsize=2, capthick=2, elinewidth=2, linestyle='-') else: e = axs[1].errorbar(fn_xHII(z_ann), Mc_ann_ftidal_pdf, xerr=fn_xHII_std(z_ann), yerr=Mc_ann_err_ftidal_pdf, color=color, label="%s ANN" % label,\ fmt="o", markerfacecolor=color, mec='k', capsize=2, capthick=2, elinewidth=2, linestyle='-') # c2 = None # if color == "b": # c2 = "c" # elif color == "g": # c2 = "k" # else: # c2 = "m" # e = axs[1].errorbar(fn_xHII(z_ann_ftidal_pdf), Mc_ann_ftidal_pdf, xerr=fn_xHII_std(z_ann), yerr=Mc_ann_err_ftidal_pdf, color=c2, label="%s ANN PDF SAMPLED" % label,\ # fmt="o", markerfacecolor=c2, mec='k', capsize=2, capthick=2, elinewidth=2, linestyle='-.', linewidth=3.) # Plot the Ok08 fit Ok_z = np.linspace(6, 10, 100) Ok_fn = baryon_fraction.Okamoto_Mc_fn() Ok_Mc = np.array([Ok_fn(i) for i in Ok_z]) axs[0].plot(Ok_z, Ok_Mc, color='k', label="Okamoto et al. 08", linestyle='-.') Hoeft_Mc = [hoeft_Mc(zi, omega_m=cosmo["omega_M_0"]) for zi in Ok_z] axs[0].plot(Ok_z, Hoeft_Mc, color='m', label="Hoeft et al. 06", linestyle='-.') # Axis labels and limits axs[0].set_xlabel(r"$z$") axs[0].set_ylabel(r"$M_{\mathrm{c}}$ [M$_{\odot}$/h]") axs[0].set_xlim(5.5, 12.5) axs[0].set_ylim(_Y_LIM[0], _Y_LIM[1]) axs[1].set_xlabel(r"$\langle x_{\mathrm{HII}} \rangle_{\mathrm{M}}$") axs[1].set_ylabel(r"$M_{\mathrm{c}}$ [M$_{\odot}$/h]") # axs[1].set_xlim(0.0, 1.05) axs[0].set_ylim(_Y_LIM[0], _Y_LIM[1]) for ax in axs.flatten(): ax.set_yscale("log") ax.legend(prop={'size': 10}) fig.tight_layout()
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
# cols = ["r", "b"] data = {"BIN": load_fesc(bin_sim[iout]), "SIN": load_fesc(sin_sim[iout])} binned_data = {} cosmo = bin_sim[iout].cosmo c = cols for key, ls in zip(data.keys(), ["-", "--", ":"]): log_mvir, fesc, tint_fesc, hids = data[key] if (iout == bin_sim.redshift(9.)): cosmo["z"] = 9. if (iout == bin_sim.redshift(11)): cosmo["z"] = 11. # log_mvir, fesc, tint_fesc, nphotons = data[key] # fesc_percent = tint_fesc * 100. fesc_percent = fesc * 100. bin_centres, mean, std, sterr = fit_scatter(log_mvir, fesc_percent, nbins=nbins, ret_sterr=True) binned_data[key] = (bin_centres, mean, std) # bin_centres, median = fit_median(log_mvir, fesc_percent, nbins=nbins) ax.scatter(log_mvir, fesc_percent, alpha=0.25, color=c, s=5, marker=".") # e = plt.errorbar(bin_centres, median, yerr=std, color=c, label='%s z=%1.2f' % (key, cosmo['z']),\ # fmt="o", markerfacecolor=c, mec='k', capsize=2, capthick=2, elinewidth=2, linestyle='-', linewidth=2.) e = ax.errorbar(bin_centres, mean,
return log_mvir, fesc, tint_fesc, nphotons from scipy import interpolate nbins = 5 ax = plt.gca() cosmo["z"] = 6 data = {"BIN": load_tint_fesc(bin_snap), "SIN": load_tint_fesc(sin_snap)} cols = ["r", "b"] binned_data = {} for key, ls, lw, c in zip(data.keys(), ["-", "--"], [3., 1.5], cols): # log_mvir, fesc, tint_fesc, hids = data[key] log_mvir, fesc, tint_fesc, nphotons = data[key] fesc_percent = tint_fesc * 100. bin_centres, mean, std = fit_scatter(log_mvir, fesc_percent, nbins=nbins) binned_data[key] = (bin_centres, mean, std) # bin_centres, median = fit_median(log_mvir, log_fesc, nbins=nbins) ax.scatter(log_mvir, fesc_percent, alpha=0.1, color=c, s=5) fn = interpolate.interp1d(bin_centres, np.log10(mean), fill_value="extrapolate") mass_milky_way = 1e12 # approx halo mass in solar masses fesc_milky_way = 10**fn(np.log10(mass_milky_way)) print key, "fesc Milky Way = %1.2f" % fesc_milky_way x = np.linspace(8, 12, 100) y = 10**fn(x) # ax.plot(x, y, color=c, linestyle=":", linewidth=5)
def plot_Mc_var(sims, sim_iouts, pickle_paths, labels, cols, fix_alpha=True, tidal_force_cutoff=None): import pickle import matplotlib.pylab as plt # from seren3.analysis.plots import fit_scatter, obs_errors from seren3.analysis import plots from seren3.utils import tau as tau_mod from scipy.interpolate import interp1d reload(plots) # fig, axs = plt.subplots(2, 2, figsize=(11,10)) fig1, axs1 = plt.subplots(1, 2, figsize=(10,5)) fig2, axs2 = plt.subplots(2, 1, figsize=(6,10)) plot_PLANCK=True plot_obs=True count = 0 for sim, iouts, ppath, label, c in zip(sims, sim_iouts, pickle_paths, labels, cols): print ppath cosmo = sim[iouts[0]].cosmo # data = pickle.load( open("%s/T_time_averaged.p" % ppath, "rb") ) data = pickle.load( open("%s/xHII_reion_history.p" % ppath, "rb") ) z = np.zeros(len(data)) var = np.zeros(len(data)) for i in range(len(data)): res = data[i].result z[i] = res['z'] # var[i] = res['mw'] var[i] = res["volume_weighted"] # var[i] = res["mass_weighted"] tau, redshifts = tau_mod.interp_xHe(var, z, sim) # tau_mod.plot(tau, redshifts, ax=axs[1,1], plot_PLANCK=plot_PLANCK, label=label, color=c) tau_mod.plot(tau, redshifts, ax=axs2[1], plot_PLANCK=plot_PLANCK, label=label, color=c) plot_PLANCK=False idx = np.where(z <= 18.) z = z[idx] var = var[idx] bc, mean, std, stderr = plots.fit_scatter(z, var, ret_sterr=True) # fn = interp1d(bc, mean, fill_value="extrapolate") fn = interp1d(z, var, fill_value="extrapolate") # Compute Mc output = fit_sim_iouts(iouts, ppath, fix_alpha, True,\ tidal_force_cutoff, 20., 1., **cosmo) # 0.15, 50., 1., **cosmo) # 0.3, 20., 1., **cosmo) z_sim = np.array([sim[i].z for i in iouts]) Mc = np.zeros(len(output)) Mc_stderr = np.zeros(len(output)) for i in range(len(output)): iout = iouts[i] Mc[i] = output[iout]["Mc"]["fit"] Mc_stderr[i] = output[iout]["Mc"]["stderr"] # Plot Mc # axs[0,0].errorbar(z_sim, Mc, yerr=Mc_stderr, label=label, linewidth=2., color=c) # axs[0,1].errorbar(fn(z_sim), Mc, yerr=Mc_stderr, label=label, linewidth=2., color=c) # axs[1,0].plot(z, 1.-var, label=label, color=c) axs1[0].errorbar(z_sim, Mc, yerr=Mc_stderr, label=label, linewidth=2., color=c) axs1[1].errorbar(fn(z_sim), Mc, yerr=Mc_stderr, label=label, linewidth=2., color=c) # Plot neutral fraction axs2[0].plot(z, 1.-var, color=c) # # Gamma # data = pickle.load( open("%s/Gamma_time_averaged.p" % ppath, "rb") ) # z = np.zeros(len(data)) # var = np.zeros(len(data)) # for i in range(len(data)): # res = data[i].result # z[i] = res['z'] # # var[i] = res['mw'] # var[i] = res["vw"] # # var[i] = res["mass_weighted"] # axs2[2].plot(z, np.log10(var), color=c) if (plot_obs): plots.obs_errors("xv", ax=axs2[0]) # obs_errors("Gamma", ax=axs2[0]) plot_obs = False axs2[0].legend(loc="lower right") count += 1 #p = axs[1,0].errorbar(bc, mean, yerr=stderr, linewidth=1, label=label) #children = p.get_children() #l = children[0] #axs[1,0].plot(bc, fn(bc), linewidth=4., linestyle='--', color=l.get_color(), zorder=10) Ok_z = np.linspace(6, 10, 100) Ok_fn = Okamoto_Mc_fn() Ok_Mc = np.array([Ok_fn(i) for i in Ok_z]) # axs[0,0].plot(Ok_z, Ok_Mc, color='g', label="Okamoto et al. 08", linestyle='-.') # axs[0,0].set_xlabel(r"$z$") # axs[0,0].set_ylabel(r"$M_{\mathrm{c}}$ [M$_{\odot}$/h]") # axs[0,0].set_xlim(6, 14) # axs[0,0].set_ylim(1e7, 2e8) # axs[0,1].set_xlabel(r"$\langle x_{\mathrm{HII}} \rangle_{V}$") # axs[0,1].set_ylabel(r"$M_{\mathrm{c}}$ [M$_{\odot}$/h]") # # axs[0,1].set_xscale("log") # axs[0,1].set_xlim(0.2, 1.05) # axs[0,1].set_ylim(1e7, 2e8) # axs[1,0].set_xlabel(r"$z$") # axs[1,0].set_ylabel(r"$\langle x_{\mathrm{HI}} \rangle_{V}$") # # axs[1,0].set_xlim(6, 18) # axs[1,0].set_xlim(6, 14) axs1[0].plot(Ok_z, Ok_Mc, color='g', label="Okamoto et al. 08", linestyle='-.') axs1[0].set_xlabel(r"$z$") axs1[0].set_ylabel(r"$M_{\mathrm{c}}$ [M$_{\odot}$/h]") # axs1[0].set_xlim(6, 14) axs1[0].set_ylim(1e7, 2e8) axs1[1].set_xlabel(r"$\langle x_{\mathrm{HII}} \rangle_{V}$") axs1[1].set_ylabel(r"$M_{\mathrm{c}}$ [M$_{\odot}$/h]") # axs[0,1].set_xscale("log") axs1[1].set_xlim(0.2, 1.05) axs1[1].set_ylim(1e7, 2e8) axs2[0].set_xlabel(r"$z$") axs2[0].set_ylabel(r"$\langle x_{\mathrm{HI}} \rangle_{V}$") # axs2[2].set_xlabel(r"$z$") # axs2[2].set_ylabel(r"$\langle \Gamma \rangle_{V}$") # axs[1,0].set_xlim(6, 18) axs2[0].set_xlim(5.5, 16) axs2[1].set_xlim(5.5, 16) # axs2[2].set_xlim(5.5, 16) axs1[0].set_yscale("log") axs1[1].set_yscale("log") # axs1[2].set_yscale("log") # axs2[0].set_yscale("log") axs1[0].legend() # axs1[1].legend() # axs2[0].legend() axs2[1].legend() # for ax in axs.flatten()[:-1]: # ax.set_yscale("log") # for ax in axs.flatten(): # ax.legend() fig1.tight_layout() fig2.tight_layout()
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 plot_fb_panels_ANN(snapshot, sim_name, pickle_path, out_dir, NN, weight="mw", weight_label="M", **kwargs): import numpy as np import matplotlib.pylab as plt from seren3.analysis.plots import fit_scatter from seren3.analysis.baryon_fraction.neural_net2 import _MVIR_MIN, _MVIR_MAX, _FB_MIN, _FB_MAX reload(neural_net2) log_mvir, fb, ftidal, xHII, T, T_U, pid_scaled = neural_net2.load_training_arrays( snapshot, pickle_path=pickle_path, weight=weight) cosmo = snapshot.cosmo cosmic_mean_b = cosmo["omega_b_0"] / cosmo["omega_M_0"] mvir = 10**log_mvir fb_cosmic_mean = fb / cosmic_mean_b bc, mean, std, sterr = fit_scatter(log_mvir, fb_cosmic_mean, nbins=10, ret_sterr=True) y_min = 0. y_max = (fb / cosmic_mean_b).max() pdf_sampling = kwargs.pop("pdf_sampling", False) def _reverse_scaling_xHII(arr): unscaled = arr / 2. unscaled += 0.5 return unscaled def reverse_scaling(data, y_idx): x, y = (data[_X_IDX], data[y_idx]) def _reverse_scaling_mvir(arr): unscaled = arr / 2. unscaled += 0.5 unscaled *= (_MVIR_MAX - _MVIR_MIN) unscaled += _MVIR_MIN return unscaled def _reverse_scaling_fb(arr): unscaled = arr / 2. unscaled += 0.5 unscaled *= (_FB_MAX - _FB_MIN) unscaled += _FB_MIN return unscaled # return _reverse_scaling(x, x_orig), _reverse_scaling(y, y_orig) # return x, _reverse_scaling2(y, y_orig) return _reverse_scaling_mvir(x), _reverse_scaling_fb(y) fig, axes = plt.subplots(nrows=4, ncols=2, sharex=True, figsize=(16, 12)) fig.subplots_adjust(hspace=0.1) fig.subplots_adjust(wspace=0.25) def _plot(mvir, fb, carr, ax, **kwargs): ax.errorbar(10**bc, mean, yerr=std, linewidth=2., color="k", linestyle="--") return ax.scatter(mvir, fb, c=carr, **kwargs) labels = [r"log$_{10} \langle F_{\mathrm{tidal}} \rangle_{t_{\mathrm{dyn}}}$",\ r"$\langle x_{\mathrm{HII}} \rangle_{\mathrm{%s}}$" % weight_label,\ # r"log$_{10} \langle x_{\mathrm{HII}} \rangle_{\mathrm{%s}}$" % weight_label,\ r"log$_{10} \langle T \rangle_{\mathrm{%s}}$/$T_{\mathrm{vir}}$" % weight_label,\ r"log$_{10}$ T/|U|", r"pid"] axs = axes[:, 0] axs[0].set_title("%s z = %1.2f" % (sim_name, snapshot.z)) count = 0 cbar = None for ax, carr, lab in zip(axs.flatten(), [np.log10(ftidal), xHII, T, T_U], labels): # for ax, carr, lab in zip(axs.flatten(), [np.log10(ftidal), np.log10(xHII), T], labels): sp = None if (count == 0): sp = _plot(mvir, fb_cosmic_mean, carr, ax, vmin=-2, vmax=0.5, **kwargs) if (count == 2): sp = _plot(mvir, fb_cosmic_mean, carr, ax, vmin=-0.5, vmax=1.5, **kwargs) else: sp = _plot(mvir, fb_cosmic_mean, carr, ax, **kwargs) if (count == 1): cbar = fig.colorbar(sp, ax=ax, ticks=[1., 0.75, 0.5, 0.25, 0.]) else: cbar = fig.colorbar(sp, ax=ax) cbar.set_label(lab) if (count == 0): cbar.set_clim(-2, 0.5) ax.set_xlim(5e6, 1.5e10) # ax.set_xlabel(r"log$_{10}$ M$_{\mathrm{vir}}$ [M$_{\odot}$/h]") ax.set_ylabel( r"f$_{\mathrm{b}}$[$\Omega_{\mathrm{b}}$/$\Omega_{\mathrm{M}}$]") ax.set_xscale("log") count += 1 ax.set_ylim(y_min, y_max) axs[-1].set_xlabel(r"M$_{\mathrm{vir}}$ [M$_{\odot}$/h]") axs = axes[:, 1] axs[0].set_title("ANN z = %1.2f" % (snapshot.z)) ioutput = snapshot.ioutput dir_name = _get_dir_name(out_dir, ioutput) input_fname = _get_input_fname(dir_name, ioutput, weight, pdf_sampling) input_data = np.loadtxt(input_fname, unpack=True) log_mvir_unscaled, fb_unscaled = reverse_scaling(input_data, _Y_IDX) mvir_unscaled = 10**log_mvir_unscaled # log_ftidal = np.log10(input_data[1]) # log_xHII = np.log10(input_data[2]) log_ftidal = input_data[1] log_xHII = input_data[2] # xHII = _reverse_scaling_xHII(input_data[2]) T = input_data[3] T_U = input_data[4] results_fname = _get_results_fname(dir_name, ioutput, weight, NN, pdf_sampling) results_data = np.loadtxt(results_fname, unpack=True) tmp, fb_unscaled = reverse_scaling(results_data, _Y_IDX_RES) fb_cosmic_mean_unscaled = fb_unscaled / cosmic_mean_b # fig.suptitle('ANN z = %1.2f' % snapshot.z, fontsize=16) bc, mean, std, sterr = fit_scatter(log_mvir_unscaled, fb_cosmic_mean_unscaled, nbins=15, ret_sterr=True) count = 0 for ax, carr, lab in zip(axs.flatten(), [log_ftidal, 10**log_xHII, T, T_U], labels): # for ax, carr, lab in zip(axs.flatten(), [log_ftidal, log_xHII, T], labels): sp = None if (count == 0): sp = _plot(mvir_unscaled, fb_cosmic_mean_unscaled, carr, ax, vmin=-2, vmax=0.5, **kwargs) if (count == 2): sp = _plot(mvir_unscaled, fb_cosmic_mean_unscaled, carr, ax, vmin=-0.5, vmax=1.5, **kwargs) else: sp = _plot(mvir_unscaled, fb_cosmic_mean_unscaled, carr, ax, **kwargs) if (count == 1): cbar = fig.colorbar(sp, ax=ax, ticks=[1., 0.75, 0.5, 0.25, 0.]) else: cbar = fig.colorbar(sp, ax=ax) cbar.set_label(lab) if (count == 0): cbar.set_clim(-2, 0.5) ax.set_xlim(5e6, 1.5e10) # ax.set_xlabel(r"log$_{10}$ M$_{\mathrm{vir}}$ [M$_{\odot}$/h]") ax.set_ylabel( r"f$_{\mathrm{b}}$[$\Omega_{\mathrm{b}}$/$\Omega_{\mathrm{M}}$]") ax.set_xscale("log") count += 1 ax.set_ylim(y_min, y_max) axs[-1].set_xlabel(r"M$_{\mathrm{vir}}$ [M$_{\odot}$/h]")
def plot_dSFR_mvir(sims, iouts, labels, cols, pickle_paths=None, **kwargs): ''' Plot star formation rate density against halo virial mass ''' import pickle import matplotlib.pylab as plt from seren3.analysis.plots import fit_scatter dSFR_units = kwargs.pop("units", "Msol yr**-1 kpc**-3") nbins = kwargs.pop("nbins", 25) legend_size = kwargs.pop("lsize", 18) fig, axs = plt.subplots(nrows=1, ncols=2, figsize=(14, 6)) def get_fname_halos(ppath, iout): return "%s/dSFR_halo_av_%05i.p" % (ppath, iout) def get_fname(ppath): return "%s/dSFR_time_averaged.p" % (ppath) ax = axs[0] if (pickle_paths is None): pickle_paths = ["%s/pickle/" % sim.path for sim in sims] for sim, iout, ppath, label, c in zip(sims, iouts, pickle_paths, labels, cols): fname = get_fname_halos(ppath, iout) data = pickle.load(open(fname, "rb")) snap = sim[iout] halos = snap.halos() part_mass = snap.quantities.particle_mass() mvir_dict = {} for h in halos: npdm = h["Mvir"] / part_mass mvir_dict[int(h["id"])] = {"Mvir": h["Mvir"], "np_dm": npdm} nrecrods = len(data) dSFR_mw = np.zeros(nrecrods) mvir = np.zeros(nrecrods) np_dm = np.zeros(nrecrods) for i in range(nrecrods): res = data[i].result dSFR_mw[i] = res["mw"].in_units(dSFR_units) mvir[i] = mvir_dict[int(data[i].idx)]["Mvir"] np_dm[i] = mvir_dict[int(data[i].idx)]["np_dm"] idx = np.where(np.logical_and(np_dm >= 50, ~np.isnan(dSFR_mw))) mvir = mvir[idx] dSFR_mw = dSFR_mw[idx] x = np.log10(mvir) y = np.log10(dSFR_mw) # Plot # plt.scatter(10**x, y, color=c, **kwargs) bc, mean, std, sterr = fit_scatter(x, y, ret_sterr=True, nbins=nbins) ax.errorbar(10**bc, mean, yerr=sterr, color=c, linewidth=2., linestyle="-", label=label) ax.legend(loc="lower right", prop={'size': legend_size}) # plt.yscale("log") ax.set_xscale("log") ax.set_xlabel(r"M$_{\mathrm{vir}}$ [M$_{\odot}$/h]") ax.set_ylabel( r"log$_{10}$ $\langle \rho_{\mathrm{SFR}} \rangle_{M}$ [M$_{\odot}$ yr$^{-1}$ kpc$^{-3}$]" ) # plt.xlim(1e7, 2e10) # plt.ylim(5e-9, 5e-1) # ax.set_ylim(-8, -2.5) ax.set_ylim(-8, y.max()) # ax.set_title("z = %1.2f" % snap.z) text_pos = (2e7, 1) text = "z = %1.2f" % snap.z ax.text(text_pos[0], text_pos[1], text, color="k", size="x-large") # Time evol. ax = axs[1] min_v = np.inf max_v = -np.inf for sim, ppath, label, c in zip(sims, pickle_paths, labels, cols): fname = get_fname(ppath) data = pickle.load(open(fname, "rb")) nrecords = len(data) z = np.zeros(nrecords) vw = np.zeros(nrecords) mw = np.zeros(nrecords) for i in range(nrecords): res = data[i].result if (res["z"] < 6): break z[i] = res["z"] vw[i] = res["vw"].in_units(dSFR_units) mw[i] = res["mw"].in_units(dSFR_units) x = z y = np.log10(mw) # ax.plot(z, log_vw, linewidth=2., color=c, label=label, linestyle="-") # ax.plot(z, log_mw, linewidth=2., color=c, label=label, linestyle="-") bc, mean, std, sterr = fit_scatter(x, y, ret_sterr=True, nbins=nbins + 2) ax.errorbar(bc, mean, yerr=std, color=c, linewidth=2., linestyle="-", label=label) ax.legend(loc="lower left", prop={'size': legend_size}) ax.set_xlabel(r"$z$") ax.set_ylabel( r"log$_{10}$ $\langle \rho_{\mathrm{SFR}} \rangle_{M}$ [M$_{\odot}$ yr$^{-1}$ kpc$^{-3}$]" ) ax.set_ylim(-7., 0.) fig.tight_layout(w_pad=1.) plt.show()
def plot_nion_esc(snapshot, ax=None, **kwargs): # Plot instantaneous photons escaped as a function of halo mass; adapated from David code which calculated integrated rate import matplotlib.pylab as plt from seren3.analysis import plots from seren3.scripts.mpi import time_int_fesc_all_halos from seren3.array import SimArray if (ax is None): ax = plt.gca() color = kwargs.pop("color", "k") nbins = kwargs.pop("nbins", 5) label = kwargs.pop("label", None) ls = kwargs.pop("ls", "-") lw = kwargs.pop("lw", 1.) #edited from here... fesc_db = time_int_fesc_all_halos.load(snapshot) nphotons_escaped = np.zeros(len(fesc_db)) mvir = np.zeros(len(fesc_db)) for i in range( len(fesc_db) ): #collation of instantaneous fesc data vs halo mass data for each halo hid = int(fesc_db[i].idx) fesc_res = fesc_db[i].result nphotons_escaped[i] = fesc_res["I1"][0] mvir[i] = fesc_res["Mvir"] # to here. The above collects the instantaneous fesc data attributed # to each halo. The original function used a cumulative calculation (integral) log_mvir = np.log10(mvir) log_nphotons = np.log10(nphotons_escaped) ix = np.where(np.logical_and( np.isfinite(log_nphotons), log_mvir >= 7.5)) #filters out low mass halos below significant resolution log_nphotons = log_nphotons[ix] log_mvir = log_mvir[ix] bc, mean, std, sterr = plots.fit_scatter(log_mvir, log_nphotons, nbins=nbins, ret_sterr=True) ax.scatter(log_mvir, log_nphotons, s=10, color='blue', alpha=0.233) #alpha determines colour intensity e = ax.errorbar(bc, mean, yerr=std, color=color, label=label,\ fmt="o", markerfacecolor=color, mec='k', capsize=2, capthick=2, elinewidth=2, linestyle=ls, linewidth=lw) if (kwargs.pop("legend", False)): ax.legend(loc="lower right", frameon=False, prop={"size": 16}) ax.set_xlabel(r'log$_{10}$ M$_{\mathrm{vir}}$ [M$_{\odot}$/h]', fontsize=20) ax.set_ylabel( r'log$_{10}$ $\dot{\mathrm{N}}_{\mathrm{ion}}(t)$ f$_{\mathrm{esc}}$ ($t$) [#$s^{-1}$]', fontsize=20) plt.show()
def plot_mass_function(self, units='Msol h**-1', kern='ST', ax=None,\ plot_Tvir=False, label_z=False, nbins=100, label=None, show=False, **kwargs): ''' Plot the Halo mass function and (optionally) Tvir on twinx axis Params: kern: The analytical kernal to use plot_Tvir: Calculates the Virial temperature in Kelvin for a halo of a given mass using equation 26 of Barkana & Loeb. ''' import matplotlib.pylab as plt from seren3.analysis.plots import fit_scatter snapshot = self.base if ax is None: fig = plt.figure() ax = fig.add_subplot(111) if label_z: label = "%s z=%1.3f" % (label, self.base.z) c = kwargs.pop("color", "b") mbinmps, y, mbinsize = self.mass_function(units=units, nbins=nbins, **kwargs) # ax.semilogy(mbinmps, y, 'o', label=label, color=c) bin_centers, mean, std = fit_scatter(mbinmps, y) # ax.errorbar(bin_centers, mean, yerr=std, color=c) e = ax.errorbar(bin_centers, mean, yerr=std, color=c,\ fmt="o", markerfacecolor=c, mec='k', capsize=2, capthick=2, elinewidth=2, linestyle="-", linewidth=2., label=label) if plot_Tvir: import cosmolopy.perturbation as cp cosmo = snapshot.cosmo M_ticks = np.array(ax.get_xticks()) mass = self.base.array(10**M_ticks, units) Tvir = [ float("%1.3f" % np.log10(v)) for v in cp.virial_temp(mass, **cosmo) ] ax2 = ax.twiny() ax2.set_xticks(M_ticks - M_ticks.min()) ax2.set_xticklabels(Tvir) ax2.set_xlabel(r"log$_{10}$(T$_{\mathrm{vir}}$ [K])") if kern is not None: import pynbody # We only need to load one CPU to setup the params dict s = pynbody.snapshot.ramses.RamsesSnap( "%s/output_%05d" % (snapshot.path, snapshot.ioutput), cpus=[1]) M_kern, sigma_kern, N_kern = pynbody.analysis.halo_mass_function( s, kern=kern, log_M_min=mbinmps.min(), log_M_max=mbinmps.max()) # Convert to correct units M_kern.convert_units(mbinmps.units) N_kern.convert_units(y.units) # ax.semilogy(np.log10(M_kern*(snapshot.info['H0']/100)), N_kern, label=kern) ax.semilogy(np.log10(M_kern), N_kern, label=kern, color="grey", linestyle="--") ax.set_xlabel(r'log$_{10}$(M [$%s$])' % mbinmps.units.latex()) ax.set_ylabel('dN / dlog$_{10}$(M [$%s$])' % y.units.latex()) if "title" in kwargs: ax.set_title(kwargs.get("title")) if show: ax.legend() plt.show()
def fit(self, nbins_fit): from seren3.analysis.plots import fit_scatter bc, mean, std, sterr = fit_scatter(np.log10(self.dset[self.den_field]),\ np.log10(self.dset[self.temp_field]), nbins=nbins_fit, ret_sterr=True) return bc, mean, std, sterr
def plot_fb_panels(sim_name, ioutputs, sim_label, weight="mw", weight_label="M", **kwargs): import numpy as np import seren3 import matplotlib.pylab as plt from seren3.analysis.plots import fit_scatter fig, axes = plt.subplots(nrows=4, ncols=len(ioutputs), sharex=True, sharey=True, figsize=(10, 10)) fig.subplots_adjust(hspace=0.05) fig.subplots_adjust(wspace=0.05) sim = seren3.load(sim_name) for i in range(len(ioutputs)): # ioutput = sim.redshift(zi) ioutput = ioutputs[i] snapshot = sim[ioutput] pickle_path = "%s/pickle/" % snapshot.path axs = axes[:, i] axs[0].set_title("z = %1.2f" % (snapshot.z)) log_mvir, fb, ftidal, xHII, T, T_U, pid = neural_net2.load_training_arrays( snapshot, pickle_path=pickle_path, weight=weight) cosmo = snapshot.cosmo cosmic_mean_b = cosmo["omega_b_0"] / cosmo["omega_M_0"] y_min = 0. y_max = (fb / cosmic_mean_b).max() fb_cosmic_mean = fb / cosmic_mean_b mvir = 10**log_mvir log_xHI = np.log10(1. - xHII) log_ftidal = np.log10(ftidal) bc, mean, std, sterr = fit_scatter(log_mvir, fb_cosmic_mean, nbins=10, ret_sterr=True) def _plot(mvir, fb, carr, ax, **kwargs): ax.errorbar(10**bc, mean, yerr=std, linewidth=3., color="k", linestyle="--") return ax.scatter(mvir, fb, c=carr, **kwargs) labels = [r"log$_{10} \langle F_{\mathrm{tidal}} \rangle_{t_{\mathrm{dyn}}}$",\ r"$\langle x_{\mathrm{HII}} \rangle_{\mathrm{%s}}$" % weight_label,\ # r"log$_{10} \langle x_{\mathrm{HII}} \rangle_{\mathrm{%s}}$" % weight_label,\ r"log$_{10} \langle T \rangle_{\mathrm{%s}}$/$T_{\mathrm{vir}}$" % weight_label ,\ r"log$_{10}$ T/|U|"] count = 0 cbar = None for ax, carr, lab in zip(axs.flatten(), [np.log10(ftidal), xHII, T, T_U], labels): ax.errorbar(10**bc, mean, yerr=std, linewidth=1.5, color="k", linestyle="--") plt.xlim(5e6, 1.5e10) sp = None if (count == 0): sp = ax.scatter(mvir, fb_cosmic_mean, c=carr, vmin=-1, vmax=4, **kwargs) if (count == 2): sp = ax.scatter(mvir, fb_cosmic_mean, c=carr, vmin=-0.5, vmax=1.5, **kwargs) else: sp = ax.scatter(mvir, fb_cosmic_mean, c=carr, **kwargs) if (count == 1): cbar = fig.colorbar(sp, ax=ax, ticks=[1., 0.75, 0.5, 0.25, 0.]) else: cbar = fig.colorbar(sp, ax=ax) if (i == len(ioutputs) - 1): cbar.set_label(lab) # ax.set_xlabel(r"log$_{10}$ M$_{\mathrm{vir}}$ [M$_{\odot}$/h]") if (i == 0): ax.set_ylabel( r"f$_{\mathrm{b}}$[$\Omega_{\mathrm{b}}$/$\Omega_{\mathrm{M}}$]" ) ax.set_xscale("log") count += 1 ax.set_ylim(y_min, y_max) axs[-1].set_xlabel(r"M$_{\mathrm{vir}}$ [M$_{\odot}$/h]") for ax in axes.flatten(): ax.set_xlim(5e6, 1.5e10) fig.tight_layout()