In [1]:
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
import scipy.optimize
%cd ..
from starlight.models import SimpleHRDModel
cmap = plt.cm.gray_r
In [2]:
nbins_perdim = 20
ncols = 1
nbins = nbins_perdim**(ncols+1)
lines = [np.linspace(0, 2, nbins_perdim) for i in range(ncols+1)]
lines[0] = np.linspace(0, 4, nbins_perdim)
mugrids = np.meshgrid(*lines)
siggrids = np.meshgrid(*[np.repeat((l[1]-l[0])/2., nbins_perdim) for l in lines])
binmus = np.vstack([g.ravel() for g in mugrids]).T
binmus[:, 0] += 1.0
binsigs = np.vstack([g.ravel() for g in siggrids]).T
numbumps = 5
alphas = np.repeat(0., nbins)
alphas[np.random.choice(nbins, numbumps, replace=False)] = 1/numbumps
binamps = np.random.dirichlet(alphas)
In [3]:
def hrd(x, amps, mus, sigs):
nbins = amps.size
npoints, ndim = x.shape
y = np.zeros((npoints, ))
for b in range(nbins):
y += amps[b] * np.prod(
np.exp(-0.5*((x-mus[b, None, :])/sigs[b, None, :])**2) /
sigs[b, None, :]/np.sqrt(2*np.pi), axis=1)
return y
fbins = 100
flines = [np.linspace(0, 2, fbins) for i in range(ncols+1)]
flines[0] = np.linspace(0, 4, fbins)
fgrids = np.meshgrid(*flines)
xfine = np.vstack([g.ravel() for g in fgrids]).T
true_model = hrd(xfine, binamps, binmus, binsigs).reshape(fgrids[0].shape)
In [4]:
nobj = 10000
varpi_fracerror = np.random.uniform(low=0.001, high=0.05, size=nobj)
mags_fracerror = np.random.uniform(low=0.001, high=0.05, size=nobj)
model = SimpleHRDModel()
#nbins, binamps, binmus, binsigs = model.draw_bins(nbins_perdim, ncols)
absmags, colors, distances, bins =\
model.draw_properties(binamps, binmus, binsigs, nobj)
dist_min, dist_max = 0.8*np.min(distances), 1.2*np.max(distances)
varpi, varpi_err, obsmags, obsmags_err, obscolors, obscolors_err =\
model.draw_data(absmags, colors, distances,
varpi_fracerror, mags_fracerror)
model.set_data(binmus, binsigs, varpi, varpi_err,
obsmags, obsmags_err, obscolors, obscolors_err, dist_min, dist_max)
x_true = model.combine_params(distances, binamps)
obsabsmag = obsmags + 5*np.log10(varpi) - 10
noisy_absmag_colors = np.vstack((obsabsmag, obscolors.T)).T
In [5]:
num_samples = 1000
dist_err = np.vstack([1/(varpi + varpi_err*np.random.randn(nobj)) for i in range(num_samples)]).std(axis=0)
In [6]:
cmap = plt.cm.gray_r
cmap.set_gamma(1)
s = 1
fig, axs = plt.subplots(2, 2, figsize=(8, 8), sharex=True, sharey=True)
axs = axs.ravel()
axs[0].pcolormesh(fgrids[0], fgrids[1], true_model, cmap=cmap)
axs[1].scatter(absmags, colors, marker='.', s=s)
axs[2].scatter(obsabsmag, obscolors, marker='.', s=s)
#axs[3].pcolormesh(fgrids[0], fgrids[1], true_model, cmap=cmap)
axs[3].errorbar(obsabsmag, obscolors, xerr=np.sqrt(obsmags_err**2 + dist_err**2), yerr=obscolors_err,
fmt='o', markersize=1, color='k', lw=1, capsize=0, alpha=0.1)
fig.tight_layout()
In [7]:
num_samples = 1000
bins_samples, binamps_samples = model.gibbs_sampler(num_samples)
In [8]:
#distances2 = distances_samples.mean(axis=0)
bins2 = bins_samples.mean(axis=0).astype(int)
binamps2 = binamps_samples.mean(axis=0)
#np.isfinite(distances_samples).sum() / num_samples
In [9]:
# TODO: better representation of bin locations (in 2D plane with sample average)
distances_samples2 = np.vstack([1/(varpi + varpi_err*np.random.randn(nobj)) for i in range(num_samples)])
fig, axs = plt.subplots(4, 2, figsize=(8, 8))
axs = axs.ravel()
js = np.where(varpi_fracerror > 0.01)[0]
np.random.shuffle(js)
js = js[:axs.shape[0]]
tgrid = model.probgrid
for i in range(axs.shape[0]):
axs[i].axvline(bins[js[i]], c='k', lw=1, zorder=0)
axs[i].hist(bins_samples[:, js[i]], nbins, histtype='step', zorder=1, normed=True, range=[0, nbins])
axs[i].plot(tgrid[js[i], :], ls='steps')
#axs[i, 1].hist(distances_samples[:, js[i]], 30, normed=True, color='b', histtype='step')#, range=[0., 0.4]
#axs[i, 1].hist(distances_samples2[:, js[i]], 30, normed=True, color='r', histtype='step')#, range=[0., 0.4]
#axs[i, 1].axvline(distances[js[i]], c='k', lw=1)
#axs[i, 1].axvline(distances2[js[i]], c='b', lw=1, ls='--')
#axs[i, 1].axvline(1/varpi[js[i]], c='r', lw=1, ls='--')
fig.tight_layout()
In [10]:
inf_model = hrd(xfine, binamps_samples.mean(axis=0), binmus, binsigs).reshape(fgrids[0].shape)
inf_model_std = hrd(xfine, binamps_samples.std(axis=0), binmus, binsigs).reshape(fgrids[0].shape)
fig, axs = plt.subplots(1, 3, figsize=(12, 3))
axs[0].pcolormesh(fgrids[0], fgrids[1], true_model, cmap=cmap)
axs[1].pcolormesh(fgrids[0], fgrids[1], inf_model, cmap=cmap)
axs[2].pcolormesh(fgrids[0], fgrids[1], inf_model_std, cmap=cmap)
Out[10]:
In [11]:
obsabsmag2 = obsmags - 5*np.log10(distances2) - 10
fig, axs = plt.subplots(2, 2, figsize=(8, 8))
axs = axs.ravel()
ind = np.abs(distances - 1/varpi) > 0.01
axs[1].hist2d(distances[ind], distances2[ind], cmap=cmap, bins=50, range=[[0.1, 0.4], [0.1, 0.4]])
axs[0].hist2d(distances[ind], 1/varpi[ind], cmap=cmap, bins=50, range=[[0.1, 0.4], [0.1, 0.4]])
for i in range(2):
#axs[i].plot([0.1, 0.3], [0.1, 0.3])
axs[i].set_xlim([0.1, 0.4])
axs[i].set_ylim([0.1, 0.4])
axs[3].errorbar(absmags, obsabsmag2, fmt='o', markersize=1)
axs[2].errorbar(absmags, obsabsmag, fmt='o', markersize=1)
omn, omx = np.min(obsabsmag), np.max(obsabsmag)
for i in range(2, 4):
axs[i].plot([omn, omx], [omn, omx])
axs[i].set_xlim([omn, omx])
axs[i].set_ylim([omn, omx])
fig, axs = plt.subplots(1, 2, figsize=(8, 3.5), sharex=True, sharey=True)
axs = axs.ravel()
axs[0].pcolormesh(fgrids[0], fgrids[1], true_model, cmap=cmap)
axs[0].scatter(obsabsmag, obscolors, marker='.', s=s, c='y')
axs[1].pcolormesh(fgrids[0], fgrids[1], true_model, cmap=cmap)
axs[1].scatter(obsabsmag2, obscolors, marker='.', s=s, c='y')
In [ ]:
In [ ]:
In [ ]: