import numpy as np
from matplotlib import pyplot as plt
from sklearn.linear_model import LinearRegression as LinearRegression_sk
from astropy.cosmology import LambdaCDM
from astroML.datasets import generate_mu_z

linear_sk = LinearRegression_sk()
linear_sk.fit(z_sample[:,None], mu_sample)
mu_fit_sk = linear_sk.predict(z[:, None])

#------------------------------------------------------------
# Plot the results
fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111)
ax.plot(z, mu_fit_sk, '-k')
ax.plot(z, mu_true, '--', c='gray')
ax.errorbar(z_sample, mu_sample, dmu, fmt='.k', ecolor='gray', lw=1)
ax.set_xlim(0.01, 1.8)
ax.set_ylim(36.01, 48)
ax.set_ylabel(r'$\mu$')
ax.set_xlabel(r'$z$')
plt.show()