Python save_figure.save_figure函数代码示例

def plot_vel_width_sims(sims, snap, log=False):
    """Plot velocity widths for a series of simulations"""
    vel_data.plot_prochaska_2008_data()
    for sss in sims:
        #Make abs. plot
        hspec = get_hspec(sss, snap)
        hspec.plot_vel_width("Si", 2, color=colors[sss], ls=lss[sss])
    outstr = "cosmo_vel_width_z"+str(snap)
    if log:
        ax = plt.gca()
        ax.set_yscale('log')
        plt.ylim(5e-2,10)
        outstr+="_log"
    else:
        plt.ylim(1e-2,2)
    hspec = get_hspec(5, snap, box=10)
    hspec.label=labels["S"]
    hspec.plot_vel_width("Si", 2, color=colors["S"], ls="--")
    hspec.plot_vw_errors("Si", 2, samples=100,cumulative=False, color=colors2["S"])
    plt.xlabel(r"$v_\mathrm{90}$ (km s$^{-1}$)")
#     plt.ylabel("Spectra per log interval")
    plt.xlim(10,1000)
    plt.xticks([10, 100, 1000],["10","100","1000"])
    plt.legend(loc=2,ncol=3)
    save_figure(path.join(outdir,outstr))
    plt.clf()

def save_plot(elem, ion, suffix, ss_corr=True):
    """Helper function to save a nicely named file"""
    filename = elem+str(ion)+"_"+suffix
    if ss_corr:
        save_figure(path.join(outdir,filename+"_photo_atten"))
    else:
        save_figure(path.join(outdir,filename+"_no_atten"))

def plot_SivsHI(mets=0.05):
    """
        Plot the SiII fraction as a function of density, for some metallicity.
        Mets is an array, metallicity as a fraction of solar.
    """
    if np.size(mets) == 1:
        mets = np.array([mets,])
    tab = cc.CloudyTable(3)

    #The hydrogen density in atoms/cm^3
    dens = np.logspace(-5,2,100)

    #Roughly mean DLA metallicity

    tabHI = cg.RahmatiRT(3, 0.71)

    tempHI = 1e4*np.ones_like(dens)
    fracHI = tabHI.neutral_fraction(dens,tempHI)
    plt.semilogx(dens, fracHI, color="red",ls="--")

    ls = [":","-","-."]
    for met in mets:
        metSi = tab.get_solar("Si")*met*np.ones_like(dens)
        fracSi = tab.ion("Si",2,metSi,dens)
        plt.semilogx(dens, fracSi, color="green",ls=ls.pop())

    plt.xlabel(r"$\rho_\mathrm{H}\; (\mathrm{amu}/\mathrm{cm}^3$)")
    plt.ylabel(r"$\mathrm{m}_\mathrm{SiII} / \mathrm{m}_\mathrm{Si}$")
    plt.show()
    save_figure(path.join(outdir,"Si_fracs"))

def plot_halos(sim,hh):
    """Plot the halo closest in mass and position in the given sim to the halo of the given number in sim 7."""
    ahalo = dp.PrettyHalo(myname.get_name(sim),3,20000)
    ahalo7 = dp.PrettyHalo(myname.get_name(7),3,20000)
    (mass, cofm, radii) = _load_halo(ahalo, 30)
    (mass7, cofm7, _) = _load_halo(ahalo7, 30)
    dist = np.sum((cofm7[hh] - cofm)**2, axis=1)
    nn = np.where( dist == np.min(dist))
    rhh = hh
    hh = np.ravel(nn)[0]
    print "Requested: ",rhh,"Got: ",hh," dist:",np.sqrt(dist[nn])," r-mass:",mass[hh]/mass7[rhh], "ormass: ",mass7[rhh]
    plt.title(r"Central Halo: $"+dp.pr_num(ahalo.sub_mass[hh]/0.72/1e11)+r"\times 10^{11} M_\odot$")
    ahalo.plot_pretty_halo(hh)
    #Backwards because someone is a fortran programmer
    circle=plt.Circle((0,0),radii[hh],color="black",fill=False)
    ax = plt.gca()
    ax.add_artist(circle)
#     plot_rvir(ahalo.sub_cofm[hh], cofm, radii,ahalo.sub_radii[hh])
    dp.tight_layout_wrapper()
    save_figure(path.join(outdir,"pretty_"+str(sim)+"_halo_"+str(rhh)))
    plt.clf()
    ahalo.plot_pretty_cut_halo(hh)
    circle=plt.Circle((0,0),radii[hh],color="black",fill=False)
    ax = plt.gca()
    ax.add_artist(circle)
#     plot_rvir(ahalo.sub_cofm[hh], cofm, radii,ahalo.sub_radii[hh]/)
    plt.title(labels[sim]+" at z=3")
    dp.tight_layout_wrapper()
    save_figure(path.join(outdir,"pretty_cut_"+str(sim)+"_halo_"+str(rhh)))
    plt.clf()
    del ahalo

def plot_all_rho(simlist):
    """Make the rho_HI plot with labels etc"""
    #Cosmo0
    for i in simlist:   #xrange(8):
        (zzz,dndx,omegadla) = get_rhohi_dndx(i)
        plt.figure(1)
        plt.plot(zzz,dndx, 'o', color=colors[i], ls=lss[i], label=labels[i])
        plt.figure(2)
        plt.plot(zzz,omegadla, 'o', color=colors[i], ls=lss[i],label=labels[i])
    plt.figure(1)
    plt.xlabel("z")
    plt.ylabel(r"$dN / dX$")
    dla_data.dndx_not()
    dla_data.dndx_pro()
    plt.xlim(2,4)
    plt.ylim(0,0.13)
    plt.legend(loc=4, ncol=3)
    tight_layout_wrapper()
    save_figure(path.join(outdir,"cosmo_dndx"))
    plt.clf()

    plt.figure(2)
    plt.xlabel("z")
    plt.ylabel(r"$10^3 \Omega_\mathrm{DLA}$")
    dla_data.omegahi_not()
#     dla_data.omegahi_pro()
    plt.xlim(2,4)
    plt.ylim(0,2.)
    plt.legend(loc=4, ncol=3)
    tight_layout_wrapper()
    save_figure(path.join(outdir,"cosmo_rhohi"))
    plt.clf()

def plot_halo_stellar_mass():
    """Plot a histogram of nearby halos"""
    for sim in (1,7,9):   #xrange(8):
        halo = myname.get_name(sim, True)
        hms = []
        sms = []
        sfrs = []
        zz = []
        for snap in (1,3,5):
            ahalo = dp.PrettyBox(halo, snap, nslice=10, label=labels[sim])
            (hm, sm, sfr) = ahalo.get_avg_stellar_mass()
            hms.append(hm)
            sms.append(sm)
            sfrs.append(sfr)
            zz.append(redshifts[snap])
        plt.figure(1)
        plt.semilogy(zz, hms, color=colors[sim], ls="-", label=labels[sim])
        plt.semilogy(zz, sms, color=colors[sim], ls="--", label=labels[sim])
        plt.figure(2)
        plt.semilogy(zz, sfrs, color=colors[sim], ls="--", label=labels[sim])
    plt.figure(1)
    plt.legend(loc=1)
    save_figure(path.join(outdir, "halos/avg_smhm_z"+str(snap)))
    plt.figure(2)
    plt.legend(loc=1)
    save_figure(path.join(outdir, "halos/avg_sfr_z"+str(snap)))
    plt.clf()

def plot_dndx_breakdown(sim):
    """Make the rho_HI plot with labels etc"""
    halo = myname.get_name(sim, True, 25)
    snaps = {4:1, 3.5:2, 3:3, 2.5:4, 2:5}
    fractions=[]
    zzz = []
    omegadla = []
    for zzzz in (4, 3.5, 3, 2.5, 2):
        ahalo = dp.PrettyBox(halo, snaps[zzzz], nslice=10)
        (massbins, fracs) = ahalo.get_omega_hi_mass_breakdown(False)
        fractions.append(fracs)
        zzz.append(zzzz)
        omegadla.append(ahalo.line_density())
    fractions = np.array(fractions)
    for i in xrange(np.size(fractions[0,:])-2):
        plt.plot(zzz,fractions[:,i+1], color=colors[i], ls=lss[i], label=dp.pr_num(np.log10(massbins[i]))+" - "+dp.pr_num(np.log10(massbins[i+1])))
    plt.plot(zzz,fractions[:,-1], color=colors[6], ls=lss[6], label="Field")
    plt.plot(zzz,omegadla, color=colors[sim], ls=lss[sim], label="Total")
    plt.xlabel("z")
    plt.ylabel(r"$dN/dX$")
    dla_data.dndx_not()
    dla_data.dndx_pro()
    plt.xlim(2,4)
    plt.ylim(0,0.15)
    plt.legend(loc=1, ncol=2)
    tight_layout_wrapper()
    save_figure(path.join(outdir,"cosmo_dndx_break"+str(sim)))
    plt.clf()

def plot_den(sim, snap, num, subdir="", voff = 0, box=10, elem="Si", ion=2):
    """Plot density"""
    hspec = get_hspec(sim, snap, snr=20., box=box)
    #Adjust the default plot parameters, which do not scale well in a gridspec.
    matplotlib.rc('xtick', labelsize=10)
    matplotlib.rc('ytick', labelsize=10)
    matplotlib.rc('axes', labelsize=10)
    matplotlib.rc('font', size=8)
    matplotlib.rc('lines', linewidth=1.5)
    gs = gridspec.GridSpec(9,2)
    ax3 = plt.subplot(gs[0:4,0])
    plt.sca(ax3)
    xoff = hspec.plot_spectrum(elem,ion,-1,num, flux=False)
    xlim = plt.xlim()
    ax3.xaxis.set_label_position('top')
    ax3.xaxis.tick_top()
    voff += xoff
    ax2 = plt.subplot(gs[5:,0])
    plt.sca(ax2)
    dxlim = hspec.plot_density(elem,ion, num)
    plt.ylabel(r"n$_\mathrm{"+elem+"II}$ (cm$^{-3}$)")
    plt.ylim(ymin=1e-9)
    ax1 = plt.subplot(gs[4,0])
    plt.sca(ax1)
    xscale = dxlim*hspec.velfac/xlim[1]
    hspec.plot_den_to_tau(elem, ion, num, thresh = 1e-9, xlim=200,voff=voff, xscale=xscale)
    ax1.axes.get_xaxis().set_visible(False)
    plt.xlabel("")
    plt.xlim(xlim)
    sdir = path.join(outdir,"spectra/"+subdir)
    if not path.exists(sdir):
        os.mkdir(sdir)
    save_figure(path.join(sdir,str(num)+"_cosmo"+str(sim)+"_"+elem+"_colden"))
    plt.clf()
    matplotlib.rc_file_defaults()

def plot_H2_effect(sim, snap):
    """Load a simulation and plot its cddf"""
    halo = myname.get_name(sim, True)
    savefile = "boxhi_grid_noH2.hdf5"
    ahalo = dp.PrettyBox(halo, snap, nslice=10, savefile=savefile, label=r"No $H_2$")
    ahalo.plot_column_density(color="blue", ls="--", moment=True)
    savefile = "boxhi_grid_H2.hdf5"
    ahalo2 = dp.PrettyBox(halo, snap, nslice=10, savefile=savefile, label=r"$H_2$")
    ahalo2.plot_column_density(color="red",moment=True)
#     savefile = path.join(halo,"snapdir_"+str(snap).rjust(3,'0'),"boxhi_grid_H2-old.hdf5")
#     ahalo2 = dp.PrettyBox(halo, snap, nslice=10, savefile=savefile)
#     ahalo2.plot_column_density(color="green",moment=True)
    dla_data.column_density_data(moment=True)
#     dla_data.noterdaeme_12_data(path.join(path.dirname(__file__),"../dla_data"), moment=True)
    plt.legend(loc=3)
    plt.xlim(1e20,2e22)
    plt.ylim(1e-5,0.1)
#     plt.title("CDDF for "+labels[sim]+" at z="+str(redshifts[snap]))
    save_figure(path.join(outdir, "cosmo"+str(sim)+"_H2_"+str(snap)))
    plt.clf()
    cddf_base = ahalo.column_density_function()
    cddf = ahalo2.column_density_function()
    plt.semilogx(cddf_base[0], np.log10(cddf[1]/cddf_base[1]), color=colors[sim], ls=lss[sim])
    plt.ylim(-0.5,0.5)
    tight_layout_wrapper()
    ax = plt.gca()
    ylab = ax.set_ylabel(r"$N_\mathrm{HI} f(N)$")
    save_figure(path.join(outdir, "cosmo_rel"+str(sim)+"_H2_"+str(snap)))
    plt.clf()

def plot_covering_frac(sim, snap, ff=True):
    """Load a simulation and plot its cddf"""
    halo = myname.get_name(sim, ff)
    ahalo = dp.PrettyBox(halo, snap, nslice=10)
    ahalo.plot_sigma_DLA()
    del ahalo
    save_figure(path.join(outdir, "cosmo"+str(sim)+"_covering_z"+str(snap)))
    plt.clf()

def test_tescari_halos(sim, snap):
    """Plot velocity width for spectra through the center of halos, like in Tescari 2009"""
    halo = myname.get_name(sim, box=10)
    hspec = ps.VWPlotSpectra(snap, halo, savefile="halo_spectra_2.hdf5",cdir=path.expanduser("~/codes/cloudy_tables/ion_out_no_atten/"))
    hspec.plot_vel_width("Si", 2, color="red")
    vel_data.plot_prochaska_2008_data()
    save_figure(path.join(outdir,"cosmo_tescari_halos"))
    plt.clf()

def plot_UVB_effect():
    """Load a simulation and plot its cddf"""
    for i in (0,5,7):
        halo = myname.get_name(i, True)
        ahalo = dp.PrettyBox(halo, 3, nslice=10)
        ahalo.plot_column_density(color=colors[i], ls=lss[i],moment=True)
    dla_data.column_density_data(moment=True)
    save_figure(path.join(outdir, "cosmo_UVB_3"))
    plt.clf()

def collis_dist(name, ahalos):
    """Plot the collisional ionisation fractions"""
    for ahalo in ahalos:
        ahalo.plot_collisional_fraction()
    plt.xlabel("r perp (kpc)")
    plt.ylabel(r"Collisional ionisation fraction")
    plt.ylim(0, 1.5)
    plt.legend()
    save_figure(path.join(outdir,name+"_CIV_collisional"))
    plt.clf()

def mass_hist(name, ahalos):
    """Plot a histogram of halo masses"""
    for ahalo in ahalos:
        ahalo.plot_mass_hist()
    plt.xlabel(r"Halo mass ($M_\odot$)")
    plt.ylabel(r"PDF")
    plt.ylim(0, 1.5)
    plt.legend()
    save_figure(path.join(outdir,name+"_CIV_mass_hist"))
    plt.clf()

def CIV_eq_ratio(name, ahalos):
    """Carbon IV equivalent width ratio"""
    for ahalo in ahalos:
        ahalo.plot_eq_width_ratio()
    plt.xlabel("r perp (kpc)")
    plt.ylabel(r"$\eta_\mathrm{pair} / \eta_\mathrm{DLA}$")
    plt.ylim(0, 1.5)
    plt.legend()
    save_figure(path.join(outdir,name+"_CIV_eq_ratio"))
    plt.clf()