Fitting straight line with uncertainties in both variables (quickly:)

quick_xy.py

import numpy as np
import pymc3 as pm
import matplotlib.pyplot as plt
import corner


def quick_xy(x, y, x_err, y_err, plot_corner=True, plot_trace=True,
             plot_posterior=True, plot_fit=True, mu_intercept=0, sd_intercept=20,
             mu_gradient=0, sd_gradient=20, mu_true_x=5, sd_true_x=10,
             n_x_pred=100, alpha_hpd=0.97, robust=False, n_tune=1000):
    """
    Based on @aplavin's adventures on pymc discourse. Currently no censored data and no
    limits.  
    
    :param x:
        Array of x-coordinates.
    :param y:
         Array of x-coordinates.
    :param x_err:
         Array of x-coordinates errors.
    :param y_err:
         Array of y-coordinates errors.
    :param plot_corner: (optional)
        Should we plot posterior using ``corner``? (default: ``True``)
    :param plot_trace: (optional)
        Should we plot trace? (default: ``True``)
    :param plot_posterior: (optional)
        Should we plot posterior using ``pymc``? (default: ``True``)
    :param plot_fit: (optional)
        Should we plot data with fitted line? (default: ``True``)
    :param mu_intercept: (optional)
        Location of the Normal prior on intercept. (default: ``0``)
    :param sd_intercept: (optional)
        Width of the Normal prior on intercept. (default: ``20``)
    :param mu_gradient: (optionl)
        Location of the Normal prior on gradient. (default: ``0``)
    :param sd_gradient: (optional)
        Width of the Normal prior on gradient. (default: ``20``)
    :param mu_true_x: (optional)
        Location of the Normal prior on location of x-points. (default: ``10``)
    :param sd_true_x: (optional)
        Width of the Normal prior on location of x-points. (default: ``10``)
    :param n_x_pred: (optional)
        Number of points on x-axis for making predictions. (default: ``100``)
    :param alpha_hpd: (optional)
        Probability mass of Highest Posterior Density interval. (default: ``0.97``)
    :param robust: (optional)
        Should we use Student t-distribution as likelihood? Set it to ``True`` in case
        outliers in data. (default: ``False``)
    :n_tune: (optional)
        NUTS tuning argument. (default: ``1000``)
        
    :return:
        Dictionary with ``trace`` and dictionary ``figures`` with all figures created.
    """
    figures = dict()
    with pm.Model() as model:
        intercept = pm.Normal('intercept', mu_intercept, sd=sd_intercept)
        gradient = pm.Normal('gradient', mu_gradient, sd=sd_gradient)
        # depends on your prior knowledge of true_x
        true_x = pm.Normal('true_x', mu=mu_true_x, sd=sd_true_x, shape=len(x))

        likelihood_x = pm.Normal('x', mu=true_x, sd=x_err, observed=x)
        true_y = pm.Deterministic('true_y', intercept+gradient*true_x)

        if robust:
            nu = pm.Uniform('nu', lower=1, upper=100)
            likelihood_y = pm.StudentT('y', mu=true_y, sd=y_err, nu=nu,
                                       observed=y)
        else:
            likelihood_y = pm.Normal('y', mu=true_y, sd=y_err, observed=y)

        trace = pm.sample(njobs=4, tune=n_tune)

        if plot_trace:
            if robust:
                varnames = [intercept, gradient, nu]
            else:
                varnames = [intercept, gradient]
            ax1 = pm.traceplot(trace, varnames=varnames)
            fig1 = plt.gcf()
            fig1.savefig("trace.pdf")
            figures.update({"trace": fig1})

        if plot_posterior:
            if robust:
                varnames = [intercept, gradient, nu]
            else:
                varnames = [intercept, gradient]
            ax2 = pm.plot_posterior(trace, varnames=varnames)
            fig2 = plt.gcf()
            fig2.savefig("posterior_pymc.pdf")
            figures.update({"posterior_pymc": fig2})

        if plot_corner:
            if robust:
                names = ["intercept", "gradient", "nu"]
                labels = [r'$\mathrm{gradient}$', r'$\mathrm{intercept}$',
                          r'$\nu$']
            else:
                names = ["intercept", "gradient"]
                labels = [r'$\mathrm{gradient}$', r'$\mathrm{intercept}$']
            samples = np.vstack([trace.get_values(name) for name in names])
            fig_corner = corner.corner(samples.T,
                                       labels=labels,
                                       show_titles=True,
                                       quantiles=[0.16, 0.50, 0.84])
            fig_corner.savefig("posterior_corner.pdf")
            figures.update({"posterior_corner": fig_corner})

    if plot_fit:
        x_pred = np.linspace(np.min(x), np.max(x), n_x_pred)
        fig, axes = plt.subplots(1, 1)
        axes.errorbar(x, y, xerr=x_err, yerr=y_err, fmt='.k')
        mu_pred = np.zeros((len(x_pred), len(trace['gradient'])))
        for iseq, seq in enumerate(x_pred):
            mu_pred[iseq] = trace['gradient']*x_pred[iseq]+trace['intercept']
        mu_mean = mu_pred.mean(axis=1)
        mu_hpd = pm.hpd(mu_pred.T, alpha=1-alpha_hpd)
        l, = axes.plot(x_pred, mu_mean)
        axes.fill_between(x_pred, mu_hpd[:, 0], mu_hpd[:, 1],
                          color=l.get_color(), alpha=0.3)
        fig.savefig("fit.pdf")
        figures.update({"fit": fig})

    return {"trace": trace, "figures": figures}


if __name__ == "__main__":
    x = np.random.uniform(0, 10, 20)
    a = 1
    b = 2
    y = a*x+b
    x_err = np.abs(np.random.normal(0, 0.5, 20))
    y_err = np.abs(np.random.normal(0, 1, 20))
    y += y_err
    x += x_err
    result = fit_xy(x, y, x_err, y_err, robust=False)