# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

#from autograd.scipy.stats import norm 
#from autograd.scipy.stats import gamma
#import autograd.numpy as np
from scipy.stats import norm
from scipy.stats import gamma
import numpy as np
import matplotlib.pyplot as plt
from astropy.table import Table, unique, Column, hstack, vstack
#from autograd import grad
from xdgmm import XDGMM
#from autograd import hessian
from scipy.optimize import minimize
import matplotlib as mpl

# +
#import autograd.numpy as np
#import autograd.numpy.linalg as npla
#import autograd.scipy.misc as scpm
import scipy.misc as scpm
#import scipy.special.logsumexp as logsumexp

def mog_logprob(x, means, icovs, lndets, pis):
    """ compute the log likelihood according to a mixture of gaussians
        with means  = [mu0, mu1, ... muk]
             icovs  = [C0^-1, ..., CK^-1]
             lndets = ln [|C0|, ..., |CK|]
             pis    = [pi1, ..., piK] (sum to 1)
        at locations given by x = [x1, ..., xN]
    """
    xx = np.atleast_2d(x)
    D  = xx.shape[1]
    centered = xx[:,:,np.newaxis] - means.T[np.newaxis,:,:]
    solved   = np.einsum('ijk,lji->lki', icovs, centered)
    logprobs = - 0.5*np.sum(solved * centered, axis=1) - (D/2.)*np.log(2*np.pi) \
               - 0.5*lndets + np.log(pis)
    logprob  = scpm.logsumexp(logprobs, axis=1)
    if np.isscalar(x) or len(x.shape) == 1:
        return logprob[0]
    else:
        return logprob



# +
def predicted_color(A, c, alpha=0.3):
    return c + A*alpha

def predicted_magnitude(A, M, d, beta=0.16):
    return M + A*beta + 5.*np.log10(d*1e3/10.)

def logjoint(theta, chat, mhat, varpihat, sigmac, sigmam, sigmavarpi, pis, means, icovs, lndets):
    lnA, M, c, lnd = theta
    A, d      = np.exp(lnA), np.exp(lnd)
    lnp_c     = norm.logpdf(chat, predicted_color(A, c), sigmac)
    lnp_m     = norm.logpdf(mhat, predicted_magnitude(A, M, d), sigmam)
    lnp_varpi = norm.logpdf(varpihat, 1./d, sigmavarpi)
    lnp_Mc    = mog_logprob(np.array([c,M]), means, icovs, lndets, pis)

    if (A >= 0) & (A < 5):
        lnp_A = 0
    else:
        lnp_A = -np.inf
    
    if (d > 0) & (d < 10): 
        lnp_d = 0
    else: 
        lnp_d = -np.inf

    string = 'p(c):{0} p(m):{1} p(varpi):{2} p(Mc):{3} p(A):{4} p(d):{5}'.format(lnp_c, lnp_m, lnp_varpi, lnp_Mc, lnp_A, lnp_d)
    #print(string)
    return np.sum([lnp_c, lnp_m, lnp_varpi, lnp_Mc, lnp_A, lnp_d])


# +
def plotXdgmm(xdgmm, ax, c='k', lw=1, label='prior', step=0.001):
    ts = np.arange(0, 2. * np.pi, step) #magic
    amps = xdgmm.weights
    mus = xdgmm.mu
    Vs = xdgmm.V
    for gg in range(xdgmm.n_components):
        if amps[gg] == np.max(amps):
            label=label
        else:
            label=None
        w, v = np.linalg.eigh(Vs[gg])
        points = np.sqrt(w[0]) * (v[:, 0])[:,None] * (np.cos(ts))[None, :] + \
                 np.sqrt(w[1]) * (v[:, 1])[:,None] * (np.sin(ts))[None, :] + \
                 mus[gg][:, None]
        ax.plot(points[0,:], points[1,:], c, lw=lw, alpha=amps[gg]/np.max(amps), rasterized=True, label=label)

def plotgaussian1sigma(ax, mu, V, c='k', lw=1, label='prior', step=0.001, alpha=1.0):
    ts = np.arange(0, 2. * np.pi, step) #magic
    w, v = np.linalg.eigh(V)
    points = np.sqrt(w[0]) * (v[:, 0])[:,None] * (np.cos(ts))[None, :] + \
             np.sqrt(w[1]) * (v[:, 1])[:,None] * (np.sin(ts))[None, :] + \
             mu[:, None]
    ax.plot(points[0,:], points[1,:], c, lw=lw, alpha=alpha, rasterized=True, label=label)
    return points


# -

ncomp = 256
filename = 'rjce_fullcmd_sfd_0.01_{0}G.fits'.format(ncomp)
xdgmm = XDGMM(filename=filename, method='Bovy')
lndets = np.array([np.linalg.slogdet(s)[1] for s in xdgmm.V])
icovs = np.array([np.linalg.inv(s) for s in xdgmm.V])
means = xdgmm.mu
pis = xdgmm.weights


datahigh = Table.read('dustHighLat-result.fits.gz')
datalow = Table.read('dustLowLat-result.fits.gz')
data = vstack((datahigh, datalow))

mhat = data['w2mpro']
w2hat = data['w2mpro']
hhat = data['h_m']
chat = data['h_m'] - data['w2mpro']
varpihat = data['parallax']
sigmac = np.sqrt(data['h_msigcom']**2 + data['w2mpro_error']**2.)
sigmam = data['w2mpro_error']
sigmavarpi = data['parallax_error']
sigmaw2 = data['w2mpro']
sigmah = data['h_msigcom']

sample = xdgmm.sample(int(len(chat)/10))

# +
method= None #'L-BFGS-B'
i = 1
fig, ax = plt.subplots(1, 1, figsize=(10, 10))

xlim = [-0.5, 1.0]
ylim = [-6, 10]
nbins = 100
xbins = np.linspace(xlim[0], xlim[1], nbins)
ybins = np.linspace(ylim[0], ylim[1], nbins)

H, xe, ye = np.histogram2d(sample[:,0], 
                           sample[:,1], 
                           bins=(xbins, ybins))
im = ax.pcolormesh(xe, ye, H.T, norm=mpl.colors.LogNorm(), cmap=plt.get_cmap('Blues'))
im.set_rasterized(True)


for i in range(0, 20):
    args = (chat[i], mhat[i], varpihat[i], sigmac[i], sigmam[i], sigmavarpi[i], pis, means, icovs, lndets)
    theta_0 = np.array([0.1, 1., 0.08, 0.5])
    #                     A, M, c, d = theta

    obj = lambda *args: -1.*logjoint(*args)
    res = minimize(obj, theta_0, args=args, method=method) #, jac=gobj)
    theta_hat = res.x
    print('Dust: {0:.3f}  Absmag: {1:0.3f}  color: {2:0.3f}  distance: {3:0.3f}'.format(theta_hat[0], 
                                                                                        theta_hat[1],
                                                                                        theta_hat[2], 
                                                                                        theta_hat[3]))

    am = mhat[i] - 5.*np.log10(1./(varpihat[i]/1e3*10))
    ax.scatter(chat[i], am, s=10, c='black')
    ax.scatter(theta_hat[2], theta_hat[1], c='red', s=10)
    plt.arrow(chat[i], am, theta_hat[2] - chat[i], theta_hat[1] - am)

    ax.set_xlabel('H-W2', fontsize=20)
    ax.set_ylabel('$M_{W2}$', fontsize=20)
    ax.set_xlim(xlim)
    ax.set_ylim(ylim[::-1])

plt.savefig('posteriors.png', rasterized=True)
# -

xx, yy = np.meshgrid(xe, ye)
values = np.zeros(len(xx.flatten()))
for i, (x, y) in enumerate(zip(xx.flatten(), yy.flatten())):
    values[i] = mog_logprob(np.array([x, y]), means, icovs, lndets, pis)
values[values < -10] = -10

from scipy.integrate import simps
simps(simps(np.exp(values).reshape(xx.shape), ye), xe)

# +
method='L-BFGS-B'
i = 1
fig, ax = plt.subplots(1, 1, figsize=(10, 10))

xlim = [-0.5, 1.0]
ylim = [-6, 10]
nbins = 100
xbins = np.linspace(xlim[0], xlim[1], nbins)
ybins = np.linspace(ylim[0], ylim[1], nbins)

im = ax.contourf(xx, yy, values.reshape(xx.shape), 100, vmin=-10, vmax=10)
fig.colorbar(im,ax=ax)
#im.set_rasterized(True)


for j in range(0, 20):
    i = np.random.randint(0, len(bright))
    args = (chat[bright[i]], mhat[bright[i]], varpihat[bright[i]], sigmac[bright[i]], sigmam[bright[i]], sigmavarpi[bright[i]], pis, means, icovs, lndets, xdgmm)
    theta_0 = np.array([0.1, 2., 0.08, 1.0])

    obj = lambda *args: -1.*logjoint(*args)
    gobj = grad(obj)
    hobj = hessian(obj)
    res = minimize(obj, theta_0, args=args, method=method, jac=gobj)
    theta_hat = res.x
    sigma_hat = np.linalg.inv(hobj(theta_hat, chat[bright[0]], mhat[bright[0]], 
                                   varpihat[bright[0]], sigmac[bright[0]], sigmam[bright[0]], sigmavarpi[bright[0]], pis, means, icovs, lndets, xdgmm))

    print('Dust: {0:.3f}+/-{1:0.3f}  Absmag: {2:0.3f}+-{3:0.3f}  color: {4:0.3f}+-{5:0.3f}  distance: {6:0.3f}+-{7:0.3f}'.format(theta_hat[0], np.sqrt(sigma_hat[0,0]), 
                                                                                  theta_hat[1], np.sqrt(sigma_hat[1,1]),
                                                                                  theta_hat[2], np.sqrt(sigma_hat[2,2]),
                                                                                  theta_hat[3], np.sqrt(sigma_hat[3,3])))


    am = mhat[bright[i]] - 5.*np.log10(1./(varpihat[bright[i]]/1e3*10))
    #plotXdgmm(xdgmm, ax[1], lw=2)
    V = np.zeros((2,2))
    V[0,0] = sigmac[0]**2.
    V[1,1] = sigmam[0]**2 + sigmavarpi[0]**2
    p = plotgaussian1sigma(ax, np.array([chat[bright[i]], am]), V, alpha=0.5) 
    V = np.zeros((2,2))
    V[0,0] = sigma_hat[2,2]
    V[1,0] = sigma_hat[2,1]
    V[0,1] = sigma_hat[1,2]
    V[1,1] = sigma_hat[1,1]
    p = plotgaussian1sigma(ax, np.array([theta_hat[2], theta_hat[1]]), V, c='red')
    #A, M, c, d = theta

    plt.arrow(chat[bright[i]], am, theta_hat[2] - chat[bright[i]], theta_hat[1] - am)

    ax.set_xlabel('H-W2', fontsize=20)
    ax.set_ylabel('$M_{W2}$', fontsize=20)
    ax.set_xlim(xlim)
    ax.set_ylim(ylim[::-1])
    #ax.set_title(l, fontsize=20)
plt.savefig('posteriors.png', rasterized=True)

# -

p(c):-6.217428764152867 p(m):-49.720487084570664 p(varpi):-2857.7659418150315 p(Mc):-4.268768121303194 p(A):-0.1 p(d):-1.0
p(c):-21.72984608241658 p(m):-533.6503819365388 p(varpi):-1937.6469227094653 p(Mc):-7.56569717726221 p(A):-0.053345126719801426 p(d):-0.7029512825086909                        

lnp_c,             lnp_m,                lnp_varpi,           lnp_Mc,             lnp_A,             lnp_d
2.1611949728028126 -90.33760374614414 -21.50742721543208 -12.24564139236691 -0.9239385332046727 -3.2265236261987185
-8.424173677612075 -14.177497800096212 -37.45983009577268 4.6380018062232455 -0.9675665739445588 -3.2262035033735943


# +



plt.contourf(xx, yy, values.reshape(xx.shape), 100, vmin=-10, vmax=10, cmap=plt.get_cmap('Blues'))
plt.colorbar()
plt.gca().invert_yaxis()

# -

np.exp(values)

# +
fig, ax = plt.subplots(1, 3, figsize=(30, 10))

xlim = [-0.5, 1.0]
ylim = [-6, 10]
nbins = 100
xbins = np.linspace(xlim[0], xlim[1], nbins)
ybins = np.linspace(ylim[0], ylim[1], nbins)

for c, ab, a in zip([chat, sample[:,0]], [absmag, sample[:,1]], [ax[0], ax[2]]):
    H, xe, ye = np.histogram2d(c, 
                           ab, 
                           bins=(xbins, ybins))
    im = a.pcolormesh(xe, ye, H.T,
      norm=mpl.colors.LogNorm(), cmap=plt.get_cmap('Blues'))#vmax=75))#, vmax=20) #, vmin=-100, vmax=100)
    im.set_rasterized(True)

plotXdgmm(xdgmm, ax[1], lw=2)
V = np.zeros((2,2))
V[0,0] = sigmac[0]**2.
V[1,1] = sigmam[0]**2 + sigmavarpi[0]**2
p = plotgaussian1sigma(ax[2], np.array([chat[bright[i]], mhat[bright[i]] - 5.*np.log10(1./(varpihat[bright[i]]/1e3*10))]), V) 
V = np.zeros((2,2))
V[0,0] = sigma_hat[2,2]
V[1,0] = sigma_hat[2,1]
V[0,1] = sigma_hat[1,2]
V[1,1] = sigma_hat[1,1]
p = plotgaussian1sigma(ax[2], np.array([theta_hat[2], theta_hat[1]]), V)
#A, M, c, d = theta
labels = ['data', 'gaussian mixture', 'samples']

for a , l in zip(ax, labels):
    a.set_xlabel('H-W2', fontsize=20)
    a.set_ylabel('$M_{W2}$', fontsize=20)
    a.set_xlim(xlim)
    a.set_ylim(ylim[::-1])
    a.set_title(l, fontsize=20)
# -

theta_hat

print(p)

plot gmm
plot chat, mhat-> Mhat
plot c, M
plot arrow




StarcoderdataJupyterScriptsDedupFiltered