Python cosmology.WMAP9类代码示例

def make_cosmic_sfr(CSP_dir):
    SFHs_fnames = glob.glob(os.path.join(CSP_dir, 'SFHs_*.fits'))
    nsubpersfh = fits.getval(SFHs_fnames[0], ext=0, keyword='NSUBPER')
    SFHs = np.row_stack(
        [fits.getdata(fn_, 'allsfhs') / fits.getdata(fn_, 'mformed')[::nsubpersfh, None]
         for fn_ in SFHs_fnames])
    ts = fits.getdata(SFHs_fnames[0], 'allts')
    zs, zmasks = zip(*[masked_z_at_value(WMAP9.age, t_ * u.Gyr) for t_ in ts])
    zs = np.ma.array(zs, mask=zmasks)

    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)

    SFRDmean = SFHs.mean(axis=0) / WMAP9.scale_factor(zs)**3.

    ax.plot(1. / WMAP9.scale_factor(zs), SFRDmean / SFRDmean.max())
    ax.set_xlim([1., 1. / WMAP9.scale_factor(10.)])
    ax.set_yscale('log')
    ax.set_ylim([.009, 1.05])
    ax.set_xlabel(r'$\frac{1}{a}$', size='x-small')
    ax.set_ylabel(r'$\log{\psi}$', size='x-small')
    ax.tick_params(labelsize='x-small', which='both')

    ax_ = ax.twiny()
    ax_.set_xlim(ax.get_xlim())
    zticks = np.linspace(0., 10., 11)
    inv_sf_ticks = 1. / WMAP9.scale_factor(zticks)
    ax_.set_xticks(inv_sf_ticks, minor=False)
    ax_.set_xticklabels(zticks)
    ax_.tick_params(labelsize='x-small')
    ax_.set_xlabel(r'$z$', size='x-small')

    fig.suptitle('``Cosmic" SFR', size='small')

    savefig(fig, 'CosmicSFR.png', CSP_dir, close=True)

def plot_cosmology(zs, mu_mcmc, mu_minuit, std_mcmc, std_minuit, n):
    import matplotlib.pyplot as plt
    from matplotlib import gridspec
    from matplotlib.ticker import MaxNLocator
    fig = plt.figure(figsize=(4.5, 5.5))
    gs = gridspec.GridSpec(2, 1, height_ratios=[2, 1], hspace=0.0, wspace=0.0)
    ax0 = fig.add_subplot(gs[0])
    ax1 = fig.add_subplot(gs[1], sharex=ax0)
    axes = [ax0, ax1]
    zsort = sorted(zs)
    distmod = WMAP9.distmod(zsort).value - 19.3
    distmod2 = WMAP9.distmod(zs).value - 19.3
    ms = 2
    alpha = 0.4
    axes[0].errorbar(zs, mu_minuit, yerr=np.sqrt(std_minuit*std_minuit + 0.1*0.1),
                     ms=ms, fmt='o', label=r"Max. Likelihood", color="r", alpha=alpha)
    axes[0].errorbar(zs, mu_mcmc, yerr=np.sqrt(std_mcmc*std_mcmc + 0.1*0.1), fmt='o',
                     ms=ms, label=r"MCMC", color="b", alpha=alpha)
    axes[1].errorbar(zs, mu_minuit - distmod2, yerr=np.sqrt(std_minuit*std_minuit + 0.1*0.1),
                     ms=ms, fmt='o', label=r"Max. Likelihood", color="r", alpha=alpha)
    axes[1].errorbar(zs, mu_mcmc - distmod2, yerr=np.sqrt(std_mcmc*std_mcmc + 0.1*0.1), fmt='o',
                     ms=ms, label=r"MCMC", color="b", alpha=alpha)
    axes[0].plot(zsort, distmod, 'k')
    axes[1].axhline(0, color='k')
    axes[1].set_xlabel("$z$")
    axes[0].set_ylabel(r"$\mu$")
    axes[1].set_ylabel(r"$\mu_{\rm obs} - \mu(\mathcal{C})$")
    axes[0].legend(loc=2, frameon=False)
    plt.setp(ax0.get_xticklabels(), visible=False)
    ax0.yaxis.set_major_locator(MaxNLocator(7, prune="lower"))
    ax1.yaxis.set_major_locator(MaxNLocator(3))

    fig.savefig("output/obs_cosmology_%s.png" % n, bbox_inches="tight", dpi=300, transparent=True)
    fig.savefig("output/obs_cosmology_%s.pdf" % n, bbox_inches="tight", dpi=300, transparent=True)

def load_model(nbins_sfh=6,sigma=0.3,datfile=None,objname=None, **extras):

    # we'll need this to access specific model parameters
    n = [p['name'] for p in model_params]

    # create SFH bins
    with open(datfile,'r') as f:
        data = json.load(f)
    zred = float(data[objname]['redshift'])
    model_params[n.index('zred')]['init'] = zred
    tuniv = WMAP9.age(zred).value*1e9

    # now construct the nonparametric SFH
    # set number of components
    # set logsfr_ratio prior
    # propagate to agebins
    model_params[n.index('agebins')]['N'] = nbins_sfh
    model_params[n.index('mass')]['N'] = nbins_sfh
    model_params[n.index('logsfr_ratios')]['N'] = nbins_sfh-1
    model_params[n.index('logsfr_ratios')]['init'] = np.full(nbins_sfh-1,0.0) # constant SFH
    model_params[n.index('logsfr_ratios')]['prior'] = SFR_Ratio(mean=np.full(nbins_sfh-1,0.0),sigma=np.full(nbins_sfh-1,sigma))

    model_params.append({'name': 'tuniv', 'N': 1,
                            'isfree': False,
                            'init': tuniv})

    # set mass-metallicity prior
    model_params[n.index('massmet')]['prior'] = MassMet(z_mini=-1.98, z_maxi=0.19, mass_mini=7, mass_maxi=12.5)

    return sedmodel.SedModel(model_params)

def load_model(agelims=[], nbins_sfh=7, sigma=0.3, df=2, **extras):

    # we'll need this to access specific model parameters
    n = [p['name'] for p in model_params]

    # first calculate t_universe at z=1
    tuniv = WMAP9.age(1.0).value*1e9

    # now construct the nonparametric SFH
    # current scheme:  last bin is 15% age of the Universe, first two are 0-30, 30-100
    # remaining N-3 bins spaced equally in logarithmic space
    tbinmax = (tuniv*0.85)
    agelims = agelims[:2] + np.linspace(agelims[2],np.log10(tbinmax),nbins_sfh-2).tolist() + [np.log10(tuniv)]
    agebins = np.array([agelims[:-1], agelims[1:]])

    # load nvariables and agebins
    model_params[n.index('agebins')]['N'] = nbins_sfh
    model_params[n.index('agebins')]['init'] = agebins.T
    model_params[n.index('mass')]['N'] = nbins_sfh
    model_params[n.index('logsfr_ratios')]['N'] = nbins_sfh-1
    model_params[n.index('logsfr_ratios')]['init'] = np.full(nbins_sfh-1,0.0) # constant SFH
    model_params[n.index('logsfr_ratios')]['prior'] = priors.StudentT(mean=np.full(nbins_sfh-1,0.0),
                                                                      scale=np.full(nbins_sfh-1,sigma),
                                                                      df=np.full(nbins_sfh-1,df))

    # insert redshift into model dictionary
    model_params[n.index('zred')]['init'] = 0.0

    return sedmodel.SedModel(model_params)

def load_model(datname='', objname='', **extras):

    ###### REDSHIFT ######
    hdulist = fits.open(datname)
    idx = hdulist[1].data['Name'] == objname
    zred =  hdulist[1].data['cz'][idx][0] / 3e5
    hdulist.close()

    #### TUNIV #####
    tuniv = WMAP9.age(zred).value

    #### TAGE #####
    tage_init = 1.1
    tage_mini  = 0.11      # FSPS standard
    tage_maxi = tuniv

    #### INSERT MAXIMUM AGE AND REDSHIFT INTO MODEL PARAMETER DICTIONARY ####
    pnames = [m['name'] for m in model_params]
    zind = pnames.index('zred')
    model_params[zind]['init'] = zred
    tind = pnames.index('tage')
    model_params[tind]['prior_args']['maxi'] = tuniv
    model = BurstyModel(model_params)

    return model

def load_model(nbins_sfh=7,sigma=0.3,df=2.,agelims=None,objname=None, **extras):

    # we'll need this to access specific model parameters
    n = [p['name'] for p in model_params]

    # replace nbins_sfh
    nbins_sfh = 4 + (int(objname)-1) / 9

    # create SFH bins
    zred = model_params[n.index('zred')]['init']
    tuniv = WMAP9.age(zred).value

    # now construct the nonparametric SFH
    # current scheme: six bins, four spaced equally in logarithmic 
    # last bin is 15% age of the Universe, first two are 0-30, 30-100
    tbinmax = (tuniv*0.85)*1e9
    agelims = agelims[:2] + np.linspace(agelims[2],np.log10(tbinmax),nbins_sfh-2).tolist() + [np.log10(tuniv*1e9)]
    agebins = np.array([agelims[:-1], agelims[1:]])

    # load nvariables and agebins
    model_params[n.index('agebins')]['N'] = nbins_sfh
    model_params[n.index('agebins')]['init'] = agebins.T
    model_params[n.index('mass')]['N'] = nbins_sfh
    model_params[n.index('logsfr_ratios')]['N'] = nbins_sfh-1
    model_params[n.index('logsfr_ratios')]['init'] = np.full(nbins_sfh-1,0.0) # constant SFH
    model_params[n.index('logsfr_ratios')]['prior'] = priors.StudentT(mean=np.full(nbins_sfh-1,0.0),
                                                                      scale=np.full(nbins_sfh-1,sigma),
                                                                      df=np.full(nbins_sfh-1,df))

    return sedmodel.SedModel(model_params)

def load_model(objname=None, datdir=None, nbins_sfh=7, sigma=0.3, df=2, agelims=[], zred=None, runname=None, **extras):

    # we'll need this to access specific model parameters
    n = [p['name'] for p in model_params]

    # first calculate redshift and corresponding t_universe
    # if no redshift is specified, read from file
    if zred is None:
        datname = datdir + objname.split('_')[0] + '_' + runname + '.dat'
        dat = ascii.read(datname)
        idx = dat['phot_id'] == int(objname.split('_')[-1])
        zred = float(dat['z_best'][idx])
    tuniv = WMAP9.age(zred).value*1e9
    model_params[n.index('zred')]['init'] = zred

    # now construct the nonparametric SFH
    # current scheme: last bin is 15% age of the Universe, first two are 0-30, 30-100
    # remaining N-3 bins spaced equally in logarithmic space
    tbinmax = (tuniv*0.85)
    agelims = agelims[:2] + np.linspace(agelims[2],np.log10(tbinmax),nbins_sfh-2).tolist() + [np.log10(tuniv)]
    agebins = np.array([agelims[:-1], agelims[1:]])

    # load nvariables and agebins
    model_params[n.index('agebins')]['N'] = nbins_sfh
    model_params[n.index('agebins')]['init'] = agebins.T
    model_params[n.index('mass')]['N'] = nbins_sfh
    model_params[n.index('logsfr_ratios')]['N'] = nbins_sfh-1
    model_params[n.index('logsfr_ratios')]['init'] = np.full(nbins_sfh-1,0.0) # constant SFH
    model_params[n.index('logsfr_ratios')]['prior'] = priors.StudentT(mean=np.full(nbins_sfh-1,0.0),
                                                                      scale=np.full(nbins_sfh-1,sigma),
                                                                      df=np.full(nbins_sfh-1,df))

    return sedmodel.SedModel(model_params)

def cosmoComVol(redshift,
                WMAP9=False,
                H0=70.0,
                Om0=0.30,
                Planck15=False,
                Gpc=False):
    """
    Get the Comoving Volume at redshift=z.

    This is simply a wrapper of astropy.cosmology
    The input redsfhit can be an array
    """
    if WMAP9:
        from astropy.cosmology import WMAP9 as cosmo
    elif Planck15:
        from astropy.cosmology import Planck15 as cosmo
    else:
        from astropy.cosmology import FlatLambdaCDM
        cosmo = FlatLambdaCDM(H0=H0, Om0=Om0)

    v = cosmo.comoving_volume(redshift)

    if not Gpc:
        return v.value
    else:
        return v.to(u.Gpc).value

def cosmoLookBack(redshift,
                  WMAP9=False,
                  H0=70.0,
                  Om0=0.30,
                  Planck15=False,
                  Myr=False):
    """
    Get the Look-back Time at redshift=z.

    This is simply a wrapper of astropy.cosmology
    The input redsfhit can be an array
    """
    if WMAP9:
        from astropy.cosmology import WMAP9 as cosmo
    elif Planck15:
        from astropy.cosmology import Planck15 as cosmo
    else:
        from astropy.cosmology import FlatLambdaCDM
        cosmo = FlatLambdaCDM(H0=H0, Om0=Om0)

    lbt = cosmo.lookback_time(redshift)

    if not Myr:
        return lbt.value
    else:
        return lbt.to(u.Myr).value

def cosmoAge(redshift,
             WMAP9=False,
             H0=70.0,
             Om0=0.30,
             Planck15=False,
             Myr=False):
    """
    Get the Age of the Universe at redshift=z.

    This is simply a wrapper of astropy.cosmology
    The input redsfhit can be an array
    """
    if WMAP9:
        from astropy.cosmology import WMAP9 as cosmo
    elif Planck15:
        from astropy.cosmology import Planck15 as cosmo
    else:
        from astropy.cosmology import FlatLambdaCDM
        cosmo = FlatLambdaCDM(H0=H0, Om0=Om0)

    age = cosmo.age(redshift)

    if not Myr:
        return age.value
    else:
        return age.to(u.Myr).value

def load_model(nbins_sfh=5,sigma=0.3,df=2, **extras):

    # we'll need this to access specific model parameters
    n = [p['name'] for p in model_params]

    # create SFH bins
    zred = model_params[n.index('zred')]['init']
    tuniv = WMAP9.age(zred).value*1e9

    # now construct the nonparametric SFH
    # set number of components
    # set logsfr_ratio prior
    # propagate to agebins
    model_params[n.index('agebins')]['N'] = nbins_sfh
    model_params[n.index('mass')]['N'] = nbins_sfh
    model_params[n.index('logsfr_ratios')]['N'] = nbins_sfh-1
    model_params[n.index('logsfr_ratios')]['init'] = np.full(nbins_sfh-1,0.0) # constant SFH
    model_params[n.index('logsfr_ratios')]['prior'] = priors.StudentT(mean=np.full(nbins_sfh-1,0.0),
                                                                      scale=np.full(nbins_sfh-1,sigma),
                                                                      df=np.full(nbins_sfh-1,df))
    model_params[n.index('logsfr_ratio30')]['prior'] = priors.StudentT(mean=0.0,
                                                                      scale=sigma,
                                                                      df=df)
    model_params[n.index('logsfr_ratiomax')]['prior'] = priors.StudentT(mean=0.0,
                                                                      scale=sigma,
                                                                      df=df)
    model_params.append({'name': 'tuniv', 'N': 1,
                            'isfree': False,
                            'init': tuniv})

    return sedmodel.SedModel(model_params)

def load_model(alpha_sfh=0.2,agelims=None, **extras):

    # we'll need this to access specific model parameters
    n = [p['name'] for p in model_params]

    # create SFH bins
    zred = model_params[n.index('zred')]['init']
    tuniv = WMAP9.age(zred).value

    # now construct the nonparametric SFH
    # current scheme: six bins, four spaced equally in logarithmic 
    # last bin is 15% age of the Universe, first two are 0-30, 30-100
    tbinmax = (tuniv*0.85)*1e9
    agelims = agelims[:2] + np.linspace(agelims[2],np.log10(tbinmax),len(agelims)-3).tolist() + [np.log10(tuniv*1e9)]
    agebins = np.array([agelims[:-1], agelims[1:]])
    ncomp = len(agelims) - 1

    # load nvariables and agebins
    model_params[n.index('agebins')]['N'] = ncomp
    model_params[n.index('agebins')]['init'] = agebins.T
    model_params[n.index('mass')]['N'] = ncomp
    model_params[n.index('mass')]['init'] = np.full(ncomp,1e6)
    model_params[n.index('mass')]['prior'] = priors.TopHat(mini=np.full(ncomp,1e5), maxi=np.full(ncomp,1e12))

    return sedmodel.SedModel(model_params)

def cosmoDA(redshift,
            WMAP9=False,
            H0=70.0,
            Om0=0.30,
            Planck15=False,
            kpc=False):
    """
    Get the Angular Diameter Distance at redshift=z.

    This is simply a wrapper of astropy.cosmology
    The input redsfhit can be an array
    """
    if WMAP9:
        from astropy.cosmology import WMAP9 as cosmo
    elif Planck15:
        from astropy.cosmology import Planck15 as cosmo
    else:
        from astropy.cosmology import FlatLambdaCDM
        cosmo = FlatLambdaCDM(H0=H0, Om0=Om0)

    da = cosmo.angular_diameter_distance(redshift)

    if not kpc:
        return da.value
    else:
        return da.to(u.kpc).value

 def get_size_flag(self):
     self.DA=np.zeros(len(self.s.SERSIC_TH50))
     #self.DA[self.membflag] = cosmo.angular_diameter_distance(self.s.CLUSTER_REDSHIFT[self.membflag]).value*Mpcrad_kpcarcsec)
     for i in range(len(self.DA)):
         if self.membflag[i]:
             self.DA[i] = cosmo.angular_diameter_distance(self.s.CLUSTER_REDSHIFT[i]).value*Mpcrad_kpcarcsec
         else:
             self.DA[i] = cosmo.angular_diameter_distance(self.s.ZDIST[i]).value*Mpcrad_kpcarcsec
     self.sizeflag=(self.s.SERSIC_TH50*self.DA > minsize_kpc) #& (self.s.SERSIC_TH50 < 20.)

def build_light_cone(fileglob, zs=np.array([6, 6.5, 7]), boxtype='delta_T'):
    zs = np.array(zs)
    files = glob(fileglob)
    zind = {'delta_T': 5, 'Ts_z': 1, 'xH_': 2, 'deltax': 3}
    zsim = []
    dims = []
    lengths = []
    files_keep = []
    for f in files:
        if not os.path.basename(f).startswith(boxtype):
            if not 'updated_smoothed_' +  boxtype in os.path.basename(f):
                continue
        files_keep.append(f)
        parms = os.path.basename(f).split('_')
        zsim.append(np.float(parms[zind[boxtype]][1:]))  # redshifts
        dims.append(np.int(parms[-2]))  # dim of box
        lengths.append(np.float(parms[-1][0:-3]))  # length in Mpc
    files = files_keep
    if (np.max(np.diff(dims)) > 0) or (np.max(np.diff(lengths)) > 0):
        raise(ValueError('Boxes are not all the same size'))
    if (np.max(zs) > np.max(zsim)) or (np.min(zs) < np.min(zsim)):
        raise(ValueError('Requested redshifts outside range of sim.'))
    order = np.argsort(zsim)
    files = [files[i] for i in order]
    zsim = np.array([zsim[i] for i in order])
    zsim0 = zsim[0]
    dim = dims[0]
    length = lengths[0]
    dx = length / dim

    lightcube = np.zeros((dim, dim, len(zs)), dtype=np.float32)
    Ds = cosmo.comoving_distance(zs).value - cosmo.comoving_distance(zsim0).value
    pix1 = map(int, np.floor(Ds / dx) % dim)
    wp1 = Ds / dx - np.floor(Ds / dx)
    pix2 = map(int, np.ceil(Ds / dx) % dim)
    wp2 = 1 - wp1
    box1 = [np.argmax(zsim[zsim <= z]) for z in zs]
    box2 = [i + 1 for i in box1]
    wb2 = (zs - zsim[box1]) / (zsim[box2] - zsim[box1])
    wb1 = (zsim[box2] - zs) / (zsim[box2] - zsim[box1])

    for i, z in enumerate(zs):
        data = np.fromfile(files[box1[i]], dtype=np.float32)
        data = data.reshape((dim, dim, dim))
        slice11 = data[:, :, pix1[i]]
        slice12 = data[:, :, pix2[i]]
        data = np.fromfile(files[box2[i]], dtype=np.float32)
        data = data.reshape((dim, dim, dim))
        slice21 = data[:, :, pix1[i]]
        slice22 = data[:, :, pix2[i]]
        lightcube[:, :, i] = (slice11 * wb1[i] * wp1[i] + slice12 * wb1[i] * wp2[i] +
                              slice21 * wb2[i] * wp1[i] + slice21 * wb2[i] * wp2[i])
    return lightcube