예제 #1
0
def compute_model(options):

    import numpy
    import astropy.io.fits as fits
    import JLA_library as JLA
    from astropy.table import Table
    from astropy.cosmology import FlatwCDM

    # -----------  Read in the configuration file ------------

    params = JLA.build_dictionary(options.config)

    # -----------  Read in the SN ordering ------------------------
    SNeList = numpy.genfromtxt(options.SNlist,
                               usecols=(0, 2),
                               dtype='S30,S200',
                               names=['id', 'lc'])
    nSNe = len(SNeList)

    for i, SN in enumerate(SNeList):
        SNeList['id'][i] = SNeList['id'][i].replace('lc-',
                                                    '').replace('.list', '')

    lcfile = JLA.get_full_path(params[options.lcfits])
    SNe = Table.read(lcfile, format='fits')

    print 'There are %d SNe' % (nSNe)

    #z=numpy.array([])
    #offset=numpy.array([])
    Om_0 = 0.303  # JLA value in the wCDM model

    cosmo1 = FlatwCDM(name='SNLS3+WMAP7', H0=70.0, Om0=Om_0, w0=-1.0)
    cosmo2 = FlatwCDM(name='SNLS3+WMAP7', H0=70.0, Om0=Om_0, w0=-1.024)

    # For the JLA SNe
    redshift = SNe['zcmb']
    replace = (redshift < 0)
    # For the non JLA SNe
    redshift[replace] = SNe[replace]['zhel']

    Delta_M = 5 * numpy.log10(
        cosmo1.luminosity_distance(redshift) /
        cosmo2.luminosity_distance(redshift))

    # Build the covariance matrix. Note that only magnitudes are affected
    Zero = numpy.zeros(nSNe)
    H = numpy.concatenate((Delta_M, Zero, Zero)).reshape(3,
                                                         nSNe).ravel(order='F')
    C_model = numpy.matrix(H).T * numpy.matrix(H)

    date = JLA.get_date()
    fits.writeto('C_model_%s.fits' % (date),
                 numpy.array(C_model),
                 clobber=True)

    return None
예제 #2
0
def ln_likelihood_2d(mu_obs, inv_covmat, z_vector, theta, choice):
    if choice=='FL2':
        cosmo=FlatLambdaCDM(Om0=theta[0], H0=100*theta[1])
    if choice=='OL2':
        cosmo = LambdaCDM(H0=70, Om0=theta[0], Ode0=theta[1])
    if choice=='FW2':
        cosmo=FlatwCDM(H0=73.8, Om0=theta[0], w0=theta[1])
    if choice=='FW3':
        cosmo=FlatwCDM(Om0=theta[0], w0=theta[1], H0=100*theta[2])
    if choice=='OL3':
        cosmo=LambdaCDM(Om0=theta[0], Ode0=theta[1], H0=100*theta[2])
    mu_th=cosmo.distmod(z_vector).value
    r=(mu_obs-mu_th).reshape(-1, 1) #Check bracketing
    chi_2=np.sum(np.matmul(np.matmul(r.T, inv_covmat), r))
    return (-chi_2/2.0)
예제 #3
0
def lnlike(theta, x, y, z, xerr, yerr, zerr):
    alpha, beta, h0 = theta

    Or = 4.153e-5 * h0**(-2)
    Om = 0.3
    w0 = -1.0

    cosmo = FlatwCDM(H0=h0 * 100, Om0=Om, w0=w0)
    #---------------------------------------------------------------------------
    ixG = np.where(z > 10)
    ixH = np.where(z < 10)

    Mum = z * 0.0
    MumErr = z * 0.0

    Mum[ixG] = z[ixG]
    MumErr[ixG] = zerr[ixG]

    Mum[ixH] = 5.0 * np.log10(cosmo.luminosity_distance(z[ixH]).value) + 25.0
    MumErr[ixH] = (5.0 / np.log(10.0)) * (zerr[ixH] / z[ixH])

    Mu = 2.5 * (beta * x + alpha) - 2.5 * y - 100.195
    MuErr = 2.5 * np.sqrt((yerr)**2 + beta**2 * (xerr)**2)

    R = (Mu - Mum)
    W = 1.0 / (MuErr**2 + MumErr**2)

    xsq = np.sum(R**2 * W)
    llq = -0.5 * xsq
    return (llq, xsq, R, Mum)
def get_dist(H0, Om0, z, w=None):
    h = H0 / 100
    if w:
        cosmos = FlatwCDM(H0=H0, Om0=Om0, w0=w)
    else:
        cosmos = FlatLambdaCDM(H0=H0, Om0=Om0)
    return cosmos.comoving_distance(z).value * h
예제 #5
0
    def __init__(self):
        self.M0 = -19.3
        self.musd = 0.1
        self.Rc = 0.1
        self.Rx = 1.
        self.x1Star = 0.0
        self.cStar = 0.0
        self._alpha = 0.13
        self._beta = 2.56
        self.cosmo = FlatwCDM(H0=72., Om0=.3, w0=-1.)
        self._SetSurveyParams()

        self._Generate()
예제 #6
0
파일: plotting.py 프로젝트: LSSTDESC/SN-PWV
def plot_cosmology_fit(
        data: pd.DataFrame, abs_mag: Numeric, H0: Numeric, Om0: Numeric,
        w0: Numeric, alpha: Numeric,
        beta: Numeric) -> Tuple[plt.figure, np.ndarray, np.ndarray]:
    """Plot a cosmological fit to a set of supernova data.

    Args:
        data: Results from the snat_sim fitting pipeline
        abs_mag: Intrinsic absolute magnitude of SNe Ia
        H0: Fitted Hubble constant at z = 0 in [km/sec/Mpc]
        Om0: Omega matter density in units of the critical density at z=0
        w0: Dark energy equation of state
        alpha: Fitted nuisance parameter for supernova stretch correction
        beta: Fitted nuisance parameter for supernova color correction

    Returns:
        The matplotlib figure, fitted distance modulus, and tabulated residuals
    """

    data = data.sort_values('z')

    fitted_mu = FlatwCDM(H0=H0, Om0=Om0, w0=w0).distmod(data.z).value
    measured_mu = data.snat_sim.calc_distmod(
        abs_mag) + alpha * data.x1 - beta * data.c
    residuals = measured_mu - fitted_mu

    fig, (top_ax,
          bottom_ax) = plt.subplots(2,
                                    sharex='col',
                                    gridspec_kw={'height_ratios': [2, 1]})

    top_ax.errorbar(data.z, measured_mu, yerr=data.mb_err, linestyle='')
    top_ax.scatter(data.z, measured_mu, s=1)
    top_ax.plot(data.z, fitted_mu, color='k', alpha=.75)

    bottom_ax.axhline(0, color='k', alpha=.75, linestyle='--')
    bottom_ax.errorbar(data.z, residuals, yerr=data.mb_err, linestyle='')
    bottom_ax.scatter(data.z, residuals, s=1)

    # Style the plot
    top_ax.set_ylabel(r'$\mu = m^*_B - M_B + \alpha x_1 - \beta c$')
    bottom_ax.set_ylabel('Residuals')
    bottom_ax_lim = max(np.abs(bottom_ax.get_ylim()))
    bottom_ax.set_ylim(-bottom_ax_lim, bottom_ax_lim)
    bottom_ax.set_xlim(xmin=0)
    fig.subplots_adjust(hspace=0.1)
    return fig, fitted_mu, residuals
예제 #7
0
    def cosmo(self, kwargs):
        """

        :param kwargs: keyword arguments of parameters (can include others not used for the cosmology)
        :return: astropy.cosmology instance
        """
        if self._cosmology == "FLCDM":
            cosmo = FlatLambdaCDM(H0=kwargs['h0'], Om0=kwargs['om'])
        elif self._cosmology == "FwCDM":
            cosmo = FlatwCDM(H0=kwargs['h0'], Om0=kwargs['om'], w0=kwargs['w'])
        elif self._cosmology == "w0waCDM":
            cosmo = w0waCDM(H0=kwargs['h0'], Om0=kwargs['om'], Ode0=1.0 - kwargs['om'], w0=kwargs['w0'], wa=kwargs['wa'])
        elif self._cosmology == "oLCDM":
            cosmo = LambdaCDM(H0=kwargs['h0'], Om0=kwargs['om'], Ode0=1.0 - kwargs['om'] - kwargs['ok'])
        else:
            raise ValueError("Cosmology %s is not supported" % self._cosmology)
        return cosmo
예제 #8
0
def log_prob_ddt(theta, lenses, cosmology):
    """
	Compute the likelihood of the given cosmological parameters against the
	modeled angular diameter distances of the lenses.

    Parameters
    ----------
	theta: list
        loat folded cosmological parameters.
	lenses: list
        lens objects (currently either GLEELens or LenstronomyLens).
	cosmology: string
        keyword indicating the choice of cosmology to work with.
	"""

    lp = log_prior(theta, cosmology)
    if not np.isfinite(lp):
        return -np.inf
    else:
        logprob = lp
        if cosmology == "FLCDM":
            h0, om = theta
            cosmo = FlatLambdaCDM(H0=h0, Om0=om)
        elif cosmology == "FwCDM":
            h0, om, w = theta
            cosmo = FlatwCDM(H0=h0, Om0=om, w0=w)
        elif cosmology == "oLCDM":
            h0, om, ok = theta
            # assert we are not in a crazy cosmological situation that prevents
            # computing the angular distance integral
            if np.any([
                    ok * (1.0 + lens.zsource)**2 + om *
                (1.0 + lens.zsource)**3 + (1.0 - om - ok) <= 0
                    for lens in lenses
            ]):
                return -np.inf
            else:
                cosmo = LambdaCDM(H0=h0, Om0=om, Ode0=1.0 - om - ok)
        else:
            raise ValueError("I don't know the cosmology %s" % cosmology)

        for lens in lenses:
            logprob += log_like_add(lens=lens, cosmo=cosmo)

        return logprob
예제 #9
0
def compute_rel_size(options):
    import numpy
    import astropy.io.fits as fits
    from astropy.table import Table
    import JLA_library as JLA
    from astropy.cosmology import FlatwCDM
    import os

    # -----------  Read in the configuration file ------------

    params = JLA.build_dictionary(options.config)

    # ---------- Read in the SNe list -------------------------

    SNeList = numpy.genfromtxt(options.SNlist,
                               usecols=(0, 2),
                               dtype='S30,S200',
                               names=['id', 'lc'])

    for i, SN in enumerate(SNeList):
        SNeList['id'][i] = SNeList['id'][i].replace('lc-',
                                                    '').replace('.list', '')

    # -----------  Read in the data JLA --------------------------

    lcfile = JLA.get_full_path(params[options.lcfits])
    SNe = Table.read(lcfile, format='fits')

    nSNe = len(SNe)
    print 'There are %d SNe in this sample' % (nSNe)

    # sort it to match the listing in options.SNlist
    indices = JLA.reindex_SNe(SNeList['id'], SNe)
    SNe = SNe[indices]

    # ---------- Compute the Jacobian ----------------------
    # The Jacobian is an m by 4 matrix, where m is the number of SNe
    # The columns are ordered in terms of Om, w, alpha and beta

    J = []
    JLA_result = {
        'Om': 0.303,
        'w': -1.00,
        'alpha': 0.141,
        'beta': 3.102,
        'M_B': -19.05
    }
    offset = {'Om': 0.01, 'w': 0.01, 'alpha': 0.01, 'beta': 0.01, 'M_B': 0.01}
    nFit = 4

    cosmo1 = FlatwCDM(name='SNLS3+WMAP7',
                      H0=70.0,
                      Om0=JLA_result['Om'],
                      w0=JLA_result['w'])

    # Varying Om
    cosmo2 = FlatwCDM(name='SNLS3+WMAP7',
                      H0=70.0,
                      Om0=JLA_result['Om'] + offset['Om'],
                      w0=JLA_result['w'])
    J.append(5 * numpy.log10((cosmo1.luminosity_distance(SNe['zcmb']) /
                              cosmo2.luminosity_distance(SNe['zcmb']))[:, 0]))

    # varying alpha
    J.append(1.0 * offset['alpha'] * SNe['x1'][:, 0])

    # varying beta
    J.append(-1.0 * offset['beta'] * SNe['color'][:, 0])

    # varying M_B

    J.append(offset['M_B'] * numpy.ones(nSNe))

    J = numpy.matrix(
        numpy.concatenate((J)).reshape(nSNe, nFit, order='F') * 100.)

    # Set up the covariance matrices

    systematic_terms = [
        'bias', 'cal', 'host', 'dust', 'model', 'nonia', 'pecvel', 'stat'
    ]

    covmatrices = {
        'bias': params['bias'],
        'cal': params['cal'],
        'host': params['host'],
        'dust': params['dust'],
        'model': params['model'],
        'nonia': params['nonia'],
        'pecvel': params['pecvel'],
        'stat': params['stat']
    }

    if options.type in systematic_terms:
        print "Using %s for the %s term" % (options.name, options.type)
        covmatrices[options.type] = options.name

    # Combine the matrices to compute the full covariance matrix, and compute its inverse
    if options.all:
        #read in the user provided matrix, otherwise compute it, and write it out
        C = fits.getdata(JLA.get_full_path(params['all']))
    else:
        C = add_covar_matrices(covmatrices, params['diag'])
        date = JLA.get_date()
        fits.writeto('C_total_%s.fits' % (date), C, clobber=True)

    Cinv = numpy.matrix(C).I

    # Construct eta, a 3n vector

    eta = numpy.zeros(3 * nSNe)
    for i, SN in enumerate(SNe):
        eta[3 * i] = SN['mb']
        eta[3 * i + 1] = SN['x1']
        eta[3 * i + 2] = SN['color']

    # Construct A, a n x 3n matrix
    A = numpy.zeros(nSNe * 3 * nSNe).reshape(nSNe, 3 * nSNe)

    for i in range(nSNe):
        A[i, 3 * i] = 1.0
        A[i, 3 * i + 1] = JLA_result['alpha']
        A[i, 3 * i + 2] = -JLA_result['beta']

    # ---------- Compute W  ----------------------
    # W has shape m * 3n, where m is the number of fit paramaters.

    W = (J.T * Cinv * J).I * J.T * Cinv * numpy.matrix(A)

    # Note that (J.T * Cinv * J) is a m x m matrix, where m is the number of fit parameters

    # ----------- Compute V_x, where x represents the systematic uncertainty

    result = []

    for term in systematic_terms:
        cov = numpy.matrix(fits.getdata(JLA.get_full_path(covmatrices[term])))
        if 'C_stat' in covmatrices[term]:
            # Add diagonal term from Eq. 13 to the magnitude
            sigma = numpy.genfromtxt(
                JLA.get_full_path(params['diag']),
                comments='#',
                usecols=(0, 1, 2),
                dtype='f8,f8,f8',
                names=['sigma_coh', 'sigma_lens', 'sigma_pecvel'])
            for i in range(nSNe):
                cov[3 * i, 3 * i] += sigma['sigma_coh'][i]**2 + sigma[
                    'sigma_lens'][i]**2 + sigma['sigma_pecvel'][i]**2

        V = W * cov * W.T
        result.append(V[0, 0])

    print '%20s\t%5s\t%5s\t%s' % ('Term', 'sigma', 'var', 'Percentage')
    for i, term in enumerate(systematic_terms):
        if options.type != None and term == options.type:
            print '* %18s\t%5.4f\t%5.4f\t%4.1f' % (term, numpy.sqrt(
                result[i]), result[i], result[i] / numpy.sum(result) * 100.)
        else:
            print '%20s\t%5.4f\t%5.4f\t%4.1f' % (term, numpy.sqrt(
                result[i]), result[i], result[i] / numpy.sum(result) * 100.)

    print '%20s\t%5.4f' % ('Total', numpy.sqrt(numpy.sum(result)))

    return
예제 #10
0
## this code plots the offset distribution, and overplots the posterior models.
## it makes a Figure 3 like plot
import numpy as np
import matplotlib.pyplot as plt
from astropy import units as u
from astropy.coordinates import SkyCoord
from astropy.cosmology import FlatwCDM
import astropy.io.fits as pyfits
cosmo = FlatwCDM(H0=70, Om0=0.3)


## some of the models were coded up for testing purposes, but only func4 was used in the end.
def func(xx, rho_0, r0, tau):
    # the gaussian+Rayleigh cent/miscent model
    pr_cor = rho_0 * (1.0 / r0 / np.sqrt(2.0 * np.pi) *
                      np.exp(-(xx / r0)**2 * 0.5))
    pr_mis = (1 - rho_0) * (xx / tau**2) * (np.exp(-0.5 * xx**2 / tau**2))
    pr = pr_cor + pr_mis
    return pr, pr_cor, pr_mis


def func2(xx, rho_0, r0, tau):
    # the two gaussians model with the second width fixed
    pr_cor = rho_0 * (1.0 / r0 / np.sqrt(2.0 * np.pi) *
                      np.exp(-(xx / r0)**2 * 0.5))
    pr_mis = (1 - rho_0) * (1.0 / 0.329 / np.sqrt(2.0 * np.pi) *
                            np.exp(-(xx / 0.329)**2 * 0.5))
    pr = pr_cor + pr_mis
    return pr, pr_cor, pr_mis

예제 #11
0
def distance(z, om=0.3, w=-1):
    #this function takes a redshift and two parameters (om,w)
    cosmo = FlatwCDM(Om0=om, w0=w)
    return cosmom.distmod(z)
예제 #12
0
        return truncated_z


if __name__ == '__main__':

    #Read simlib file before sampling, obs is list of libIDs which is passed to simulation class each time.
    #This will take a while as we also need to sort all times in all bands
    libfile = 'DES_DIFFIMG.SIMLIB'  #NOTE: THE SIMLIB INDEX STARTS AT 1, NOT 0#
    meta, obs = SncosmoSimulation.read_simlib(libfile)
    print type(obs)
    for libid in obs.keys():
        df = obs[libid].to_pandas()
        df = df.sort_values(by='time')
        obs[libid] = Table.from_pandas(df)

    #parameters we might want to vary in sampler
    cosmo = FlatwCDM(name='SNLS3+WMAP7', H0=71.58, Om0=0.262, w0=-1.0)
    alpha = 0.14
    beta = 3.2
    deltaM = 0.0
    zp_offset = [0.0, 0.0, 0.0, 0.0]  #griz

    fm_sim = SncosmoSimulation(simlib_obs_sets=obs,
                               cosmo=cosmo,
                               alpha=alpha,
                               beta=beta,
                               deltaM=deltaM,
                               zp_off=zp_offset,
                               NumSN=50)
예제 #13
0
def wcdm():
    return FlatwCDM(H0=70.0, w0=-0.9, Om0=0.3, Ob0=0.05, Tcmb0=2.7)
예제 #14
0
def make_hubble_plot(fitres_file, m0diff_file, prob_col_name, args):
    logging.info(
        f"Making Hubble plot from FITRES file {fitres_file} and M0DIF file {m0diff_file}"
    )
    # Note that the fitres file has mu and fit 0, m0diff will have to select down to it

    name, sim_num, *_ = fitres_file.split("__")
    sim_num = int(sim_num)

    df = pd.read_csv(fitres_file, delim_whitespace=True, comment="#")
    dfm = pd.read_csv(m0diff_file)
    dfm = dfm[(dfm.name == name) & (dfm.sim_num == sim_num) &
              (dfm.muopt_num == 0) & (dfm.fitopt_num == 0)]

    from astropy.cosmology import FlatwCDM
    import numpy as np
    import matplotlib.pyplot as plt

    df.sort_values(by="zHD", inplace=True)
    dfm.sort_values(by="z", inplace=True)
    dfm = dfm[dfm["MUDIFERR"] < 10]

    ol = dfm.ol_ref.unique()[0]
    w = dfm.w_ref.unique()[0]
    if np.isnan(ol):
        logging.info("Setting ol = 0.689")
        ol = 0.689
    if np.isnan(w):
        logging.info("Setting w = -1")
        w = -1
    alpha = 0
    beta = 0
    sigint = 0
    gamma = r"$\gamma = 0$"
    scalepcc = "NA"
    num_sn_fit = df.shape[0]
    contam_data, contam_true = "", ""

    with gzip.open(fitres_file, "rt") as f:
        for line in f.read().splitlines():
            if "NSNFIT" in line:
                v = int(line.split("=", 1)[1].strip())
                num_sn_fit = v
                num_sn = f"$N_{{SN}} = {v}$"
            if "alpha0" in line and "=" in line and "+-" in line:
                alpha = r"$\alpha = " + line.split("=")[-1].replace(
                    "+-", r"\pm") + "$"
            if "beta0" in line and "=" in line and "+-" in line:
                beta = r"$\beta = " + line.split("=")[-1].replace(
                    "+-", r"\pm") + "$"
            if "sigint" in line and "iteration" in line:
                sigint = r"$\sigma_{\rm int} = " + line.split()[3] + "$"
            if "gamma" in line and "=" in line and "+-" in line:
                gamma = r"$\gamma = " + line.split("=")[-1].replace(
                    "+-", r"\pm") + "$"
            if "CONTAM_TRUE" in line:
                v = max(0.0,
                        float(line.split("=", 1)[1].split("#")[0].strip()))
                n = v * num_sn_fit
                contam_true = f"$R_{{CC, true}} = {v:0.4f} (\\approx {int(n)} SN)$"
            if "CONTAM_DATA" in line:
                v = max(0.0,
                        float(line.split("=", 1)[1].split("#")[0].strip()))
                n = v * num_sn_fit
                contam_data = f"$R_{{CC, data}} = {v:0.4f} (\\approx {int(n)} SN)$"
            if "scalePCC" in line and "+-" in line:
                scalepcc = "scalePCC = $" + line.split(
                    "=")[-1].strip().replace("+-", r"\pm") + "$"
    if prob_col_name is not None:
        prob_label = prob_col_name.replace("PROB_", "").replace("_", " ")
        classifier_text = f"Classifier = {prob_label}"
    else:
        classifier_text = "No Classification"
    label = "\n".join([
        num_sn, alpha, beta, sigint, gamma, scalepcc, contam_true, contam_data,
        classifier_text
    ])
    label = label.replace("\n\n", "\n").replace("\n\n", "\n")
    dfz = df["zHD"]
    zs = np.linspace(dfz.min(), dfz.max(), 500)
    distmod = FlatwCDM(70, 1 - ol, w).distmod(zs).value

    n_trans = 1000
    n_thresh = 0.05
    n_space = 0.3
    subsec = True
    if zs.min() > n_thresh:
        n_space = 0.01
        subsec = False
    z_a = np.logspace(np.log10(min(0.01,
                                   zs.min() * 0.9)), np.log10(n_thresh),
                      int(n_space * n_trans))
    z_b = np.linspace(n_thresh,
                      zs.max() * 1.01, 1 + int((1 - n_space) * n_trans))[1:]
    z_trans = np.concatenate((z_a, z_b))
    z_scale = np.arange(n_trans)

    def tranz(zs):
        return interp1d(z_trans, z_scale)(zs)

    if subsec:
        x_ticks = np.array([0.01, 0.02, 0.05, 0.2, 0.4, 0.6, 0.8, 1.0])
        x_ticks_m = np.array([0.03, 0.04, 0.1, 0.3, 0.5, 0.6, 0.7, 0.9])
    else:
        x_ticks = np.array([0.05, 0.2, 0.4, 0.6, 0.8, 1.0])
        x_ticks_m = np.array([0.1, 0.3, 0.5, 0.6, 0.7, 0.9])
    mask = (x_ticks > z_trans.min()) & (x_ticks < z_trans.max())
    mask_m = (x_ticks_m > z_trans.min()) & (x_ticks_m < z_trans.max())
    x_ticks = x_ticks[mask]
    x_ticks_m = x_ticks_m[mask_m]
    x_tick_t = tranz(x_ticks)
    x_ticks_mt = tranz(x_ticks_m)

    fig, axes = plt.subplots(figsize=(7, 5),
                             nrows=2,
                             sharex=True,
                             gridspec_kw={
                                 "height_ratios": [1.5, 1],
                                 "hspace": 0
                             })
    logging.info(f"Hubble plot prob colour given by column {prob_col_name}")

    if prob_col_name is not None:
        if prob_col_name.upper().startswith("PROB"):
            mask_no_prob = df[prob_col_name] < -1
            df.loc[mask_no_prob, prob_col_name] = 1.0
            df[prob_col_name] = df[prob_col_name].clip(0, 1)

    for resid, ax in enumerate(axes):
        ax.tick_params(which="major", direction="inout", length=4)
        ax.tick_params(which="minor", direction="inout", length=3)
        if resid:
            sub = df["MUMODEL"]
            sub2 = 0
            sub3 = distmod
            ax.set_ylabel(r"$\Delta \mu$")
            ax.tick_params(top=True, which="both")
            alpha = 0.2
            ax.set_ylim(-0.5, 0.5)
        else:
            sub = 0
            sub2 = -dfm["MUREF"]
            sub3 = 0
            ax.set_ylabel(r"$\mu$")
            ax.annotate(label, (0.98, 0.02),
                        xycoords="axes fraction",
                        horizontalalignment="right",
                        verticalalignment="bottom",
                        fontsize=8)
            alpha = 0.7

        ax.set_xlabel("$z$")
        if subsec:
            ax.axvline(tranz(n_thresh),
                       c="#888888",
                       alpha=0.4,
                       zorder=0,
                       lw=0.7,
                       ls="--")

        if prob_col_name is None or df[prob_col_name].min() >= 1.0:
            cc = df["IDSURVEY"]
            vmax = None
            color_prob = False
            cmap = "rainbow"
        else:
            cc = df[prob_col_name]
            vmax = 1.05
            color_prob = True
            cmap = "inferno"

        # Plot each point
        ax.errorbar(tranz(dfz),
                    df["MU"] - sub,
                    yerr=df["MUERR"],
                    fmt="none",
                    elinewidth=0.5,
                    c="#AAAAAA",
                    alpha=0.5 * alpha)
        h = ax.scatter(tranz(dfz),
                       df["MU"] - sub,
                       c=cc,
                       s=1,
                       zorder=2,
                       alpha=alpha,
                       vmax=vmax,
                       cmap=cmap)

        if not args.get("BLIND", []):
            # Plot ref cosmology
            ax.plot(tranz(zs),
                    distmod - sub3,
                    c="k",
                    zorder=-1,
                    lw=0.5,
                    alpha=0.7)

            # Plot m0diff
            ax.errorbar(tranz(dfm["z"]),
                        dfm["MUDIF"] - sub2,
                        yerr=dfm["MUDIFERR"],
                        fmt="o",
                        mew=0.5,
                        capsize=3,
                        elinewidth=0.5,
                        c="k",
                        ms=4)
        ax.set_xticks(x_tick_t)
        ax.set_xticks(x_ticks_mt, minor=True)
        ax.set_xticklabels(x_ticks)
        ax.set_xlim(z_scale.min(), z_scale.max())

        if args.get("BLIND", []):
            ax.set_yticklabels([])
            ax.set_yticks([])
    if color_prob:
        cbar = fig.colorbar(h,
                            ax=axes,
                            orientation="vertical",
                            fraction=0.1,
                            pad=0.01,
                            aspect=40)
        cbar.set_label("Prob Ia")

    fp = fitres_file.replace(".fitres.gz", ".png")
    logging.debug(f"Saving Hubble plot to {fp}")
    fig.savefig(fp, dpi=300, transparent=True, bbox_inches="tight")
    plt.close(fig)