#!/usr/bin/env python
# coding: utf-8

# In[1]:


from IPython.display import HTML

HTML('''<script>
code_show=true; 
function code_toggle() {
 if (code_show){
 $('div.input').hide();
 } else {
 $('div.input').show();
 }
 code_show = !code_show
} 
$( document ).ready(code_toggle);
</script>
<form action="javascript:code_toggle()"><input type="submit" value="Click here to toggle on/off the raw code."></form>''')


# In[2]:


import os
import numpy
from scipy import optimize
import pandas
import fitsio
import corner
import bovy_mcmc
import gaia_tools.load
from galpy.util import bovy_plot, bovy_coords
get_ipython().run_line_magic('pylab', 'inline')
import seaborn as sns
from matplotlib import gridspec
import matplotlib.lines as mlines
from matplotlib.ticker import NullFormatter
nullfmt= NullFormatter()
save_figures= False
numpy.random.seed(1)


# # Measuring the Oort constants with TGAS

# Load the data:

# In[3]:


tgas= gaia_tools.load.tgas()
# apass data downloaded from Dustin Lang's http://portal.nersc.gov/project/cosmo/temp/dstn/gaia/
apass= fitsio.read('tgas-matched-apass-dr9.fits')
bv= (apass['bmag']-apass['vmag'])


# Compute the proper-motion covariance matrix in Galactic coordinates:

# In[4]:


pmllbb= bovy_coords.pmrapmdec_to_pmllpmbb(tgas['pmra'],tgas['pmdec'],
                                         tgas['ra'],tgas['dec'],degree=True)
cov_pmradec= numpy.array([[tgas['pmra_error']**2.,
                                   tgas['pmra_pmdec_corr']*tgas['pmra_error']*tgas['pmdec_error']],
                                 [tgas['pmra_pmdec_corr']*tgas['pmra_error']*tgas['pmdec_error'],
                                 tgas['pmdec_error']**2.]]).T
cov_pmllbb= bovy_coords.cov_pmrapmdec_to_pmllpmbb(cov_pmradec,
                                                  tgas['ra'],tgas['dec'],degree=True)


# Cuts to define the main sequence and plot the color-magnitude diagram:

# In[5]:


def main_sequence_cut(bv,low=False):
    if low:
        #(-0.1,1.75),(0.5,5.5)
        #(0.5,5.5),(1.3,9)
        #(1.3,9),(1.5,12)
        out= numpy.zeros_like(bv)
        indx= bv < 0.5
        out[indx]= 3.75/0.6*(bv[indx]+0.1)+1.75
        indx= (bv >= 0.5)*(bv < 1.3)
        out[indx]= 3.5/0.8*(bv[indx]-0.5)+5.5
        indx= bv >= 1.3
        out[indx]= 3./0.2*(bv[indx]-1.3)+9.
    else:
        """(0.8,3.5),(0.2,0)
        (0.8,3.5),(1,6)
        (1,6),(8,1.5)
        (8,1.5),(1.7,10)"""
        out= numpy.zeros_like(bv)
        indx= bv < 0.8
        out[indx]= 3.5/0.6*(bv[indx]-0.2)+0.
        indx= (bv >= 0.8)*(bv < 1.)
        out[indx]= 2.5/0.2*(bv[indx]-0.8)+3.5
        indx= (bv >= 1.)*(bv < 1.5)
        out[indx]= 2./0.5*(bv[indx]-1.)+6.
        indx= (bv >= 1.5)
        out[indx]= 2./0.2*(bv[indx]-1.5)+8.
    return out


# In[6]:


print("Number of TGAS sources with good parallaxes:",
      numpy.sum((tgas['parallax']/tgas['parallax_error'] > 10.)*(1./tgas['parallax'] < 0.5)))


# In[7]:


bovy_plot.bovy_print(axes_labelsize=17.,text_fontsize=12.,xtick_labelsize=15.,ytick_labelsize=15.)
figsize(5,7)
indx= (tgas['parallax']/tgas['parallax_error'] > 10.)\
        *(True-numpy.isnan(apass['vmag']))*(True-numpy.isnan(apass['bmag']))\
    *(apass['vmag'] != 0.)*(apass['bmag'] != 0.)\
    *(1./tgas['parallax'] < 0.5)
print("Number of TGAS sources with good parallaxes and APASS photometry:",numpy.sum(indx))
dm= -5.*numpy.log10(tgas['parallax'])+10.
bovy_plot.bovy_plot((apass['bmag']-apass['vmag'])[indx],
                    (apass['vmag']-dm)[indx],color='k',
                    scatter=True,alpha=0.2,rasterized=True,
                    s=(1.-0.9*numpy.exp(-0.5*((apass['bmag']-apass['vmag'])[indx]-0.6)**2./0.3**2.))*.1,
                    xrange=[-0.1,2.],yrange=[10.,-2.],
                    xlabel=r'$B-V$',ylabel=r'$M_V$')
bvs= numpy.linspace(-1.,2.,201)
plot(bvs,main_sequence_cut(bvs),'-',lw=1.5,color='0.4')
plot(bvs,main_sequence_cut(bvs,low=True),'-',lw=1.5,color='0.4')
if save_figures:
    plt.savefig(os.path.join(os.getenv('PAPERSDIR'),'2016-oort','cmd.pdf'),bbox_inches='tight')


# The sample after cutting to the main sequence:

# In[8]:


figsize(5,7)
indx= (tgas['parallax']/tgas['parallax_error'] > 10.)\
        *(True-numpy.isnan(apass['vmag']))*(True-numpy.isnan(apass['bmag']))\
    *(apass['vmag'] != 0.)*(apass['bmag'] != 0.)\
    *((apass['vmag']-dm) > main_sequence_cut(bv))\
    *((apass['vmag']-dm) < main_sequence_cut(bv,low=True))
print("Number of sources on main-sequence:",numpy.sum(indx))
print("Number of sources on main-sequence with good colors:",numpy.sum(indx*(bv > 0.2)*(bv < 1.)))
print("Typical distance in kpc:",numpy.median(1./tgas['parallax'][indx]))
dm= -5.*numpy.log10(tgas['parallax'])+10.
bovy_plot.bovy_plot((apass['bmag']-apass['vmag'])[indx],
                    (apass['vmag']-dm)[indx],color='k',
                    scatter=True,alpha=0.2,rasterized=True,
                    s=(1.-0.9*numpy.exp(-0.5*((apass['bmag']-apass['vmag'])[indx]-0.6)**2./0.3**2.))*.1,
                    xrange=[-0.1,2.],yrange=[10.,-2.],
                    xlabel=r'$B-V$',ylabel=r'$M_V$')
bvs= numpy.linspace(-1.,2.,201)
plot(bvs,main_sequence_cut(bvs),'-',lw=1.5,color='0.6')
plot(bvs,main_sequence_cut(bvs,low=True),'-',lw=1.5,color='0.6')


# The distance distribution:

# In[9]:


figsize(6,4)
_= bovy_plot.bovy_hist(1000./tgas['parallax'][indx],bins=301,histtype='step',
                      range=[0.,500.],lw=2.,
                      xlabel=r'$\mathrm{Distance}\,(\mathrm{kpc})$')
axvline(numpy.median(1000./tgas['parallax'][indx]),color='k')


# The overall trend in $\mu_l$ with $l$:

# In[10]:


figsize(12,4)
tindx= indx*(bv > 0.4)*(bv < 0.7)
bovy_plot.bovy_plot(tgas['l'][tindx],pmllbb[tindx,0]*4.74047,'k,',
                    xrange=[-20.,380.],yrange=[-400.,400.],
                   xlabel=r'$\mathrm{Galactic\ longitude}\,(\mathrm{deg})$',
                   ylabel=r'$\mu_l\,(\mathrm{km\,s}^{-1}\,\mathrm{kpc}^{-1})$')


# Let's look at the correlation between parallax and proper motion and whether it is dangerous to ignore it (by copying some code in $\texttt{galpy.util.bovy_coords}$):

# In[11]:


ra, dec= tgas['ra']/180.*numpy.pi, tgas['dec']/180.*numpy.pi
theta,dec_ngp,ra_ngp= bovy_coords.get_epoch_angles(2000.)
#Whether to use degrees and scalar input is handled by decorators
dec[dec == dec_ngp]+= 10.**-16 #deal w/ pole.
sindec_ngp= numpy.sin(dec_ngp)
cosdec_ngp= numpy.cos(dec_ngp)
sindec= numpy.sin(dec)
cosdec= numpy.cos(dec)
sinrarangp= numpy.sin(ra-ra_ngp)
cosrarangp= numpy.cos(ra-ra_ngp)
cosphi= sindec_ngp*cosdec-cosdec_ngp*sindec*cosrarangp
sinphi= sinrarangp*cosdec_ngp
norm= numpy.sqrt(cosphi**2.+sinphi**2.)
cosphi/= norm
sinphi/= norm
cov_plxpmradec= numpy.array([[tgas['parallax_error']**2.,
                              tgas['parallax_pmra_corr']*tgas['parallax_error']*tgas['pmra_error'],
                              tgas['parallax_pmdec_corr']*tgas['parallax_error']*tgas['pmdec_error']],
                             [tgas['parallax_pmra_corr']*tgas['parallax_error']*tgas['pmra_error'],
                              tgas['pmra_error']**2.,
                              tgas['pmra_pmdec_corr']*tgas['pmra_error']*tgas['pmdec_error']],
                             [tgas['parallax_pmdec_corr']*tgas['parallax_error']*tgas['pmdec_error'],
                              tgas['pmra_pmdec_corr']*tgas['pmra_error']*tgas['pmdec_error'],
                              tgas['pmdec_error']**2.]]).T
cov_plxpmllbb= numpy.zeros_like(cov_plxpmradec)
for ii in range(len(tgas)):
    P= numpy.array([[1.,0.,0.],[0.,cosphi[ii],-sinphi[ii]],[0.,sinphi[ii],cosphi[ii]]])
    cov_plxpmllbb[ii]= numpy.dot(P,numpy.dot(cov_plxpmradec[ii],P.T))


# In[12]:


figsize(14,4)
plxpmll_corr= cov_plxpmllbb[:,0,1]/numpy.sqrt(cov_plxpmllbb[:,0,0]*cov_plxpmllbb[:,1,1])
bovy_plot.bovy_plot(tgas['l'][indx],
                    plxpmll_corr[indx],'k,',alpha=0.1,rasterized=True,
                   xrange=[-20.,380.],
                   yrange=[-1.3,1.3],
                   xlabel=r'$\mathrm{Galactic\ longitude}\,(\mathrm{deg})$',
                   ylabel=r'$\mathrm{Parallax}-\mu_l\ \mathrm{correlation}$')
sindx= numpy.argsort(tgas['l'][indx])
wind= numpy.sum(indx)//360*2
plot(tgas['l'][indx][sindx][::wind],
     pandas.rolling_apply(plxpmll_corr[indx][sindx],wind,lambda x: numpy.mean(x),center=True)[::wind],
     'o')


# In[13]:


figsize(14,4)
plxpmbb_corr= cov_plxpmllbb[:,0,2]/numpy.sqrt(cov_plxpmllbb[:,0,0]*cov_plxpmllbb[:,2,2])
bovy_plot.bovy_plot(tgas['l'][indx],
                    plxpmbb_corr[indx],'k,',alpha=0.1,rasterized=True,
                   xrange=[-20.,380.],
                   yrange=[-1.3,1.3],
                   xlabel=r'$\mathrm{Galactic\ longitude}\,(\mathrm{deg})$',
                   ylabel=r'$\mathrm{Parallax}-\mu_b\ \mathrm{correlation}$')
sindx= numpy.argsort(tgas['l'][indx])
plot(tgas['l'][indx][sindx][::wind],
     pandas.rolling_apply(plxpmbb_corr[indx][sindx],wind,lambda x: numpy.mean(x),center=True)[::wind],
     'o')


# There are substantial trends, but by decorrelating the uncertainties (see below), we can show that these give much smaller contributions to the measured Oort constants than the measurement uncertainties.

# ## Fit for the Oort constants
# 
# Model is that
# \begin{equation}
# \mu_l (l) = (A\,\cos 2l - C\,\sin 2l + B)\,\cos b + \varpi\,(u_0\,\sin l - v_0\,\cos l) + \mathrm{scatter}\,,
# \end{equation}
# and
# \begin{equation}
# \mu_b (l) = -(A\,\sin 2l + C\,\cos 2l + K)\,\sin b\,\cos b + \varpi\,(u_0\,\cos l + v_0\,\sin l)\,\sin b - \varpi\,w_0\,\cos b+ \mathrm{scatter}\,.
# \end{equation}

# In[14]:


def mloglike(*args):
    return -loglike(*args)
def loglike(p,cos_twol,sin_twol,sin_l,cos_l,sin_b,cos_b,data_pml,data_pmb,data_pm_cov,data_plx):
    #p= [A,B,C,K,u0,v0,w0,scatter_mul,scatter_mub]
    A= p[0]
    B= p[1]
    C= p[2]
    K= p[3]
    u0= p[4]
    v0= p[5]
    w0= p[6]
    scatter_pml= numpy.exp(p[7])
    scatter_pmb= numpy.exp(p[8])
    pred_pml_l= (A*cos_twol-C*sin_twol+B)*cos_b+u0*sin_l*data_plx-v0*cos_l*data_plx
    pred_pmb_l= -(A*sin_twol+C*cos_twol+K)*sin_b*cos_b+sin_b*(u0*cos_l*data_plx+v0*sin_l*data_plx)-w0*cos_b*data_plx
    return -0.5*numpy.sum((pred_pml_l-data_pml)**2./(data_pm_cov[:,0,0]+scatter_pml**2.)\
                          +numpy.log(data_pm_cov[:,0,0]+scatter_pml**2.))\
           -0.5*numpy.sum((pred_pmb_l-data_pmb)**2./(data_pm_cov[:,1,1]+scatter_pmb**2.)\
                          +numpy.log(data_pm_cov[:,1,1]+scatter_pmb**2.))
def fit_sample(tindx,plots=False,disp=False,bvmin=None,bvmax=None,subaxi=False,noxlabel=False,showmub=False,
              addlegend=False,decorrelate_errors=False):
    """
    NAME:
        fit_sample
    PURPOSE:
        fit the Oort constants to the mu_l(l) and mu_b(l) trends in TGAS
    INPUT:
        tindx - index of the subsample to fit
        plots= (False) Plot the proper motion corrected for solar motion in l (or b if showmub) and the best-fit model?
        disp= (False) print the convergence message from the optimizer?
        bvmin= (None) Minimum B-V of the sample (for label on plot)
        bvmax= (None) Maximum B-V of the sample (for label on plot)        
        subaxi= (False) For the plot, subtract the axisymmetric signal (A,B)
        noxlabel= (False) Don't include the xlabel
        showmub= (False) Plot mu_b(l) instead of mu_l(l)
        addlegend= (False) Add a legend for the line
        decorrelate_errors= (False) if True, de-correlate the parallax and proper motion errors by adding 
                            random, uncorrelated noise
    OUTPUT:
        best-fit parameters
    HISTORY:
       2016-10-18 - Started - Bovy (UofT/CCA)
    """
    p= [20.,-10.,-10.,-10.,0.,0.,0.,numpy.log(30.),numpy.log(30.)]
    tpml= pmllbb[tindx,0]*4.74047
    tpmb= pmllbb[tindx,1]*4.74047
    tpmcov= cov_pmllbb[tindx]*4.74047**2.
    tl= tgas['l'][tindx]/180.*numpy.pi
    tb= tgas['b'][tindx]/180.*numpy.pi
    if decorrelate_errors:
        tplx= tgas['parallax'][tindx]+numpy.random.normal(size=numpy.sum(tindx))\
            *4.*tgas['parallax_error'][tindx]
        tpml+= 4.*numpy.sqrt(tpmcov[:,0,0])*numpy.random.normal(size=numpy.sum(tindx))
        tpmb+= 4.*numpy.sqrt(tpmcov[:,1,1])*numpy.random.normal(size=numpy.sum(tindx))
        tpmcov[:,0,0]*= 16.
        tpmcov[:,1,1]*= 16.
    else:
        tplx= tgas['parallax'][tindx]
    pout= optimize.fmin_powell(mloglike,p,
                               args=(numpy.cos(2*tl),
                                     numpy.sin(2*tl),
                                     numpy.sin(tl),
                                     numpy.cos(tl),
                                     numpy.sin(tb),
                                     numpy.cos(tb),
                                     tpml,tpmb,tpmcov,
                                     tplx),
                              disp=disp)
    if plots:
        A= pout[0]
        B= pout[1]
        C= pout[2]
        K= pout[3]
        u0= pout[4]
        v0= pout[5]
        w0= pout[6]
        if showmub:
            cindx= tindx*(numpy.fabs(tgas['b']) > 40.)*(numpy.fabs(tgas['b']) < 50.)
            corr= -(pmllbb[cindx,1]*4.74047\
                   +w0*numpy.cos(tgas['b'][cindx]/180.*numpy.pi)*tgas['parallax'][cindx]\
                   -numpy.sin(tgas['b'][cindx]/180.*numpy.pi)*tgas['parallax'][cindx]\
                   *(u0*numpy.cos(tgas['l'][cindx]/180.*numpy.pi)+v0*numpy.sin(tgas['l'][cindx]/180.*numpy.pi)))\
                   /numpy.sin(tgas['b'][cindx]/180.*numpy.pi)/numpy.cos(tgas['b'][cindx]/180.*numpy.pi)
        else:
            cindx= tindx*(numpy.fabs(tgas['b']) < 20.)
            corr= (pmllbb[cindx,0]*4.74047\
                                -u0*numpy.sin(tgas['l'][cindx]/180.*numpy.pi)\
                                     *tgas['parallax'][cindx]\
                                 +v0*numpy.cos(tgas['l'][cindx]/180.*numpy.pi)\
                                     *tgas['parallax'][cindx])\
                                    /numpy.cos(tgas['b'][cindx]/180.*numpy.pi)
            if subaxi:
                corr-= A*numpy.cos(2.*tgas['l'][cindx]/180.*numpy.pi)+B
        sindx= numpy.argsort(tgas['l'][cindx])
        if subaxi:
            wind= numpy.sum(cindx)//360*8
        else:
            wind= numpy.sum(cindx)//360*(8+4*showmub)
        if noxlabel:
            xlabel= None
        else:
            xlabel=r'$\mathrm{Galactic\ longitude}\,(\mathrm{deg})$'
        if showmub:
            ylabel=r'$\Delta\mu_b\,(\mathrm{km\,s}^{-1}\,\mathrm{kpc}^{-1})$'
        else:
            ylabel=r'$\Delta\mu_l\,(\mathrm{km\,s}^{-1}\,\mathrm{kpc}^{-1})$'
        bovy_plot.bovy_plot(tgas['l'][cindx][sindx][::wind],
                            pandas.rolling_apply(corr[sindx],wind,lambda x: numpy.mean(x),center=True)[::wind],'ko',
                            xrange=[-20.,380.],
                            yrange=[-76.+30.*subaxi,76.-30.*subaxi],
                            gcf=True,
                            xlabel=xlabel,
                            ylabel=ylabel)
        errorbar(tgas['l'][cindx][sindx][::wind],
                 pandas.rolling_apply(corr[sindx],wind,lambda x: numpy.mean(x),center=True)[::wind],
                 yerr=pandas.rolling_apply(corr,wind,lambda x: numpy.std(x)/numpy.sqrt(len(x)),center=True)[::wind],
                 marker='None',ls='none',capsize=0,color='k')
        ls= numpy.linspace(0.,360.,201)
        if showmub:
            pred_pm_l= A*numpy.sin(2.*ls/180.*numpy.pi)+C*numpy.cos(2.*ls/180.*numpy.pi)+K
        else:
            if subaxi:
                pred_pm_l= -C*numpy.sin(2.*ls/180.*numpy.pi)
            else:
                pred_pm_l= A*numpy.cos(2.*ls/180.*numpy.pi)-C*numpy.sin(2.*ls/180.*numpy.pi)+B
        plot(ls,pred_pm_l,lw=4.,color=sns.color_palette()[0])
        if noxlabel:
            gca().xaxis.set_major_formatter(nullfmt)
        bovy_plot.bovy_text(r'$%.1f < B-V < %.1f$' % (bvmin,bvmax),
                           size=17.,top_left=True)
        if showmub:
            bovy_plot.bovy_text(r'$40^\circ < |b| < 50^\circ$',
                               size=17.,top_right=True)
        else:
            bovy_plot.bovy_text(r'$|b| < 20^\circ$',
                               size=17.,top_right=True)
        if addlegend:
            if showmub:
                legendlabel= r'$\Delta\mu_b = A\,\sin 2l + C\,\cos 2l+K$'
            else:
                legendlabel= r'$\Delta \mu_l = A\,\cos 2l - C\,\sin 2l + B$'
            pyplot.legend((mlines.Line2D([], [],lw=4.,color=sns.color_palette()[0],ls='-',marker='None'),),
                          (legendlabel,),
                           loc='lower left',#bbox_to_anchor=(.02,.02),
                           numpoints=1,
                           prop={'size':20},
                           frameon=False)
    return pout
def mcmc_sample(tindx,nsamples=1000,p=None,plots=False):
    """Similar to fit_sample, but run MCMC"""
    if p is None:
        p= fit_sample(tindx,plots=False)
    samples=\
        bovy_mcmc.markovpy(p,0.2,
                           loglike,(numpy.cos(2*tgas['l'][tindx]/180.*numpy.pi),
                                     numpy.sin(2*tgas['l'][tindx]/180.*numpy.pi),
                                     numpy.sin(tgas['l'][tindx]/180.*numpy.pi),
                                     numpy.cos(tgas['l'][tindx]/180.*numpy.pi),
                                     numpy.sin(tgas['b'][tindx]/180.*numpy.pi),
                                     numpy.cos(tgas['b'][tindx]/180.*numpy.pi),
                                     pmllbb[tindx,0]*4.74047,pmllbb[tindx,1]*4.74047,
                                     cov_pmllbb[tindx]*4.74047**2.,
                                     tgas['parallax'][tindx]),
                           isDomainFinite=[[False,False] for ii in range(len(p))],
                           domain= [[0.,0.] for ii in range(len(p))],
                           nsamples=nsamples,
                           nwalkers=2*len(p))
    samples= numpy.array(samples).T
    if plots: plot_samples(samples)
    return samples
def plot_samples(samples,fitmub=False):
    labels= [r'$A\,(\mathrm{km\,s}^{-1}\,\mathrm{kpc}^{-1})$',
             r'$B\,(\mathrm{km\,s}^{-1}\,\mathrm{kpc}^{-1})$',
             r'$C\,(\mathrm{km\,s}^{-1}\,\mathrm{kpc}^{-1})$',
             r'$K\,(\mathrm{km\,s}^{-1}\,\mathrm{kpc}^{-1})$',
             r'$u_0\,(\mathrm{km\,s}^{-1})$',
             r'$v_0\,(\mathrm{km\,s}^{-1})$',
             r'$w_0\,(\mathrm{km\,s}^{-1})$']
    corner.corner(samples[:7].T,quantiles=[0.16, 0.5, 0.84],labels=labels,
                  show_titles=True, title_args={"fontsize": 12})
    return None


# ### Some examples of fitting

# In[15]:


figsize(6,4)
tindx= indx*(bv > .2)*(bv < 0.4)
print fit_sample(tindx,plots=True,disp=True,bvmin=0.2,bvmax=0.4,subaxi=False)


# In[16]:


figsize(6,4)
tindx= indx*(bv > .4)*(bv < 0.6)
print fit_sample(tindx,plots=True,disp=True,bvmin=0.4,bvmax=0.6,subaxi=False)


# In[17]:


figsize(6,4)
tindx= indx*(bv > .6)*(bv < 0.8)
print fit_sample(tindx,plots=True,disp=True,bvmin=0.6,bvmax=0.8,subaxi=False)


# In[18]:


figsize(6,4)
tindx= indx*(bv > .8)*(bv < 1.)
print fit_sample(tindx,plots=True,bvmin=0.8,bvmax=1.,disp=True,subaxi=False)


# All of the good bins:

# In[19]:


figsize(6,16)
gs= gridspec.GridSpec(4,1,wspace=0.1,hspace=0.075)
subplot(gs[0])
tindx= indx*(bv > .2)*(bv < 0.4)
fit_sample(tindx,plots=True,disp=False,bvmin=0.2,bvmax=0.4,subaxi=False,noxlabel=True,addlegend=True)
subplot(gs[1])
tindx= indx*(bv > .4)*(bv < 0.6)
fit_sample(tindx,plots=True,disp=False,bvmin=0.4,bvmax=0.6,subaxi=False,noxlabel=True)
subplot(gs[2])
tindx= indx*(bv > .6)*(bv < 0.8)
fit_sample(tindx,plots=True,disp=False,bvmin=0.6,bvmax=0.8,subaxi=False,noxlabel=True)
subplot(gs[3])
tindx= indx*(bv > .8)*(bv < 1.0)
fit_sample(tindx,plots=True,disp=False,bvmin=0.8,bvmax=1.,subaxi=False)
if save_figures:
    plt.savefig(os.path.join(os.getenv('PAPERSDIR'),'2016-oort','pml.pdf'),bbox_inches='tight')


# A well populated bin for $\mu_b$

# In[20]:


figsize(6,4)
tindx= indx*(bv > .4)*(bv < 0.6)
fit_sample(tindx,plots=True,disp=False,bvmin=0.4,bvmax=0.6,subaxi=False,showmub=True,addlegend=True)
if save_figures:
    plt.savefig(os.path.join(os.getenv('PAPERSDIR'),'2016-oort','pmb.pdf'),bbox_inches='tight')


# An example MCMC run:

# In[21]:


figsize(12,4)
s= mcmc_sample(tindx,plots=True,nsamples=30000)


# ### Fit all color-bins as a function of $B-V$

# In[22]:


bv_min= numpy.arange(0.,1.21,0.2)
bv_max= numpy.arange(0.2,1.41,0.2)
nsamples= 30000
all_samples= numpy.zeros((len(bv_min),9,nsamples))
for ii,(mi,ma) in enumerate(zip(bv_min,bv_max)):
    tindx= indx*(bv > mi)*(bv < ma)
    print("Working on %i/%i; %i stars" % (ii+1,len(bv_min),numpy.sum(tindx)))
    all_samples[ii]= mcmc_sample(tindx,nsamples=nsamples,plots=False)


# In[23]:


def out(sa):
    return (numpy.mean(numpy.median(sa,axis=1)),numpy.std(numpy.median(sa,axis=1)))
bv_med= 0.5*(bv_min+bv_max)
figsize(14,4)
subplot(1,3,1)
bovy_plot.bovy_plot(bv_med,numpy.median(all_samples[:,0,:],axis=1),
                   'ko',gcf=True,
                    xrange=[-0.2,1.7],
                    yrange=[0.,30.],
                    xlabel=r'$B-V$',ylabel=r'$A\ (\mathrm{km\,s}^{-1}\,\mathrm{kpc}^{-1})$',)
errorbar(bv_med,numpy.median(all_samples[:,0,:],axis=1),
         yerr=numpy.std(all_samples[:,0,:],axis=1),
         color='k',
         marker='None',ls='none',capsize=0)
axhline(out(all_samples[1:5,0,:])[0],ls='--',lw=2.,color='k')
subplot(1,3,2)
bovy_plot.bovy_plot(bv_med,numpy.median(all_samples[:,1,:],axis=1),
                   'ko',gcf=True,
                    xrange=[-0.2,1.7],
                    yrange=[-30.,0.],
                    xlabel=r'$B-V$',ylabel=r'$B\ (\mathrm{km\,s}^{-1}\,\mathrm{kpc}^{-1})$',)
errorbar(bv_med,numpy.median(all_samples[:,1,:],axis=1),
         yerr=numpy.std(all_samples[:,1,:],axis=1),
         color='k',
         marker='None',ls='none',capsize=0)
axhline(out(all_samples[1:5,1,:])[0],ls='--',lw=2.,color='k')
subplot(1,3,3)
bovy_plot.bovy_plot(bv_med,numpy.median(all_samples[:,0,:]-all_samples[:,1,:],axis=1),
                   'ko',gcf=True,
                    xrange=[-0.2,1.7],
                    yrange=[20.,40.],
                    xlabel=r'$B-V$',ylabel=r'$A-B\ (\mathrm{km\,s}^{-1}\,\mathrm{kpc}^{-1})$',)
errorbar(bv_med,numpy.median(all_samples[:,0,:]-all_samples[:,1,:],axis=1),
         yerr=numpy.std(all_samples[:,0,:]-all_samples[:,1,:],axis=1),
         color='k',
         marker='None',ls='none',capsize=0)
axhline(out(all_samples[1:5,0,:]-all_samples[1:5,1,:])[0],ls='--',lw=2.,color='k')
tight_layout()
if save_figures:
    plt.savefig(os.path.join(os.getenv('PAPERSDIR'),'2016-oort','AB.pdf'),bbox_inches='tight')


# In[24]:


figsize(14,4)
subplot(1,3,1)
bovy_plot.bovy_plot(numpy.median(all_samples[1:5,5,:],axis=1)-12.24,
                    numpy.median(all_samples[1:5,0,:],axis=1),
                   'ko',gcf=True,
                    xrange=[-7.,20.],
                    yrange=[0.,20.],
                    xlabel=r'$v_a\,(\mathrm{km\,s}^{-1})$',
                    ylabel=r'$A\ (\mathrm{km\,s}^{-1}\,\mathrm{kpc}^{-1})$',)
errorbar(numpy.median(all_samples[1:5,5,:],axis=1)-12.24,
         numpy.median(all_samples[1:5,0,:],axis=1),
         xerr=numpy.std(all_samples[1:5,5,:]-12.24,axis=1),
         yerr=numpy.std(all_samples[1:5,0,:],axis=1),
         color='k',
         marker='None',ls='none',capsize=0)
print "Linear fit:"
print numpy.polyfit(numpy.median(all_samples[1:5,5,:],axis=1)-12.24,
                    numpy.median(all_samples[1:5,0,:],axis=1),1)


# In[25]:


bv_med= 0.5*(bv_min+bv_max)
figsize(14,4)
subplot(1,3,1)
bovy_plot.bovy_plot(bv_med,numpy.median(all_samples[:,2,:],axis=1),
                   'ko',gcf=True,
                    xrange=[-0.2,1.7],
                    yrange=[-15.,5.],
                    xlabel=r'$B-V$',ylabel=r'$C\ (\mathrm{km\,s}^{-1}\,\mathrm{kpc}^{-1})$',)
errorbar(bv_med,numpy.median(all_samples[:,2,:],axis=1),
         yerr=numpy.std(all_samples[:,2,:],axis=1),
         color='k',
         marker='None',ls='none',capsize=0)
axhline(out(all_samples[1:5,2,:])[0],ls='--',lw=2.,color='k')
subplot(1,3,2)
bovy_plot.bovy_plot(bv_med,numpy.median(all_samples[:,3,:],axis=1),
                   'ko',gcf=True,
                    xrange=[-0.2,1.7],
                    yrange=[-15.,5.],
                    xlabel=r'$B-V$',ylabel=r'$K\ (\mathrm{km\,s}^{-1}\,\mathrm{kpc}^{-1})$',)
errorbar(bv_med,numpy.median(all_samples[:,3,:],axis=1),
         yerr=numpy.std(all_samples[:,3,:],axis=1),
         color='k',
         marker='None',ls='none',capsize=0)
axhline(out(all_samples[1:5,3,:])[0],ls='--',lw=2.,color='k')
subplot(1,3,3)
bovy_plot.bovy_plot(bv_med,numpy.median(all_samples[:,2,:]+all_samples[:,3,:],axis=1),
                   'ko',gcf=True,
                    xrange=[-0.2,1.7],
                    yrange=[-20.,20.],
                    xlabel=r'$B-V$',
                    ylabel=r'$\partial \bar{v}_R / \partial R = K+C\ (\mathrm{km\,s}^{-1}\,\mathrm{kpc}^{-1})$',)
errorbar(bv_med,numpy.median(all_samples[:,2,:]+all_samples[:,3,:],axis=1),
         yerr=numpy.std(all_samples[:,2,:]+all_samples[:,3,:],axis=1),
         color='k',
         marker='None',ls='none',capsize=0)
axhline(out(all_samples[1:5,2,:]+all_samples[1:5,3,:])[0],ls='--',lw=2.,color='k')
tight_layout()
if save_figures:
    plt.savefig(os.path.join(os.getenv('PAPERSDIR'),'2016-oort','CK.pdf'),bbox_inches='tight')


# In[26]:


bv_med= 0.5*(bv_min+bv_max)
figsize(14,4)
subplot(1,3,1)
bovy_plot.bovy_plot(bv_med,numpy.median(all_samples[:,4,:],axis=1),
                   'ko',gcf=True,
                    xrange=[-0.2,1.7],
                    yrange=[0.,15.],
                    xlabel=r'$B-V$',ylabel=r'$u_0\ (\mathrm{km\,s}^{-1})$',)
errorbar(bv_med,numpy.median(all_samples[:,4,:],axis=1),
         yerr=numpy.std(all_samples[:,4,:],axis=1),
         color='k',
         marker='None',ls='none',capsize=0)
axhline(out(all_samples[1:5,4,:])[0],ls='--',lw=2.,color='k')
subplot(1,3,2)
bovy_plot.bovy_plot(bv_med,numpy.median(all_samples[:,5,:],axis=1),
                   'ko',gcf=True,
                    xrange=[-0.2,1.7],
                    yrange=[0.,25.],
                    xlabel=r'$B-V$',ylabel=r'$v_0\ (\mathrm{km\,s}^{-1})$',)
errorbar(bv_med,numpy.median(all_samples[:,5,:],axis=1),
         yerr=numpy.std(all_samples[:,5,:],axis=1),
         color='k',
         marker='None',ls='none',capsize=0)
subplot(1,3,3)
bovy_plot.bovy_plot(bv_med,numpy.median(all_samples[:,6,:],axis=1),
                   'ko',gcf=True,
                    xrange=[-0.2,1.7],
                    yrange=[0.,10.],
                    xlabel=r'$B-V$',
                    ylabel=r'$w_0\ (\mathrm{km\,s}^{-1})$',)
errorbar(bv_med,numpy.median(all_samples[:,6,:],axis=1),
         yerr=numpy.std(all_samples[:,6,:],axis=1),
         color='k',
         marker='None',ls='none',capsize=0)
axhline(out(all_samples[1:5,6,:])[0],ls='--',lw=2.,color='k')
tight_layout()


# In[27]:


print "Parameters:"
print "A = %.2f +/- %.2f km/s/kpc" % (out(all_samples[1:5,0,:]))
print "Ac = %.2f +/- %.2f km/s/kpc" % (out(all_samples[1:5,0,:]+(all_samples[1:5,5,:]-12.24)/8.))
print "B = %.2f +/- %.2f km/s/kpc" % (out(all_samples[1:5,1,:]))
print "C = %.2f +/- %.2f km/s/kpc" % (out(all_samples[1:5,2,:]))
print "K = %.2f +/- %.2f km/s/kpc" % (out(all_samples[1:5,3,:]))
print "u0 = %.2f +/- %.2f km/s" % (out(all_samples[1:5,4,:]))
print "w0 = %.2f +/- %.2f km/s" % (out(all_samples[1:5,6,:]))
print "Derived parameters:"
print "A-B = %.2f +/- %.2f km/s/kpc" % (out(all_samples[1:5,0,:]-all_samples[1:5,1,:]))
print "K+C = %.2f +/- %.2f km/s/kpc" % (out(all_samples[1:5,2,:]+all_samples[1:5,3,:]))
print "Vc = %.2f +/- %.2f km/s" % (out((all_samples[1:5,0,:]-all_samples[1:5,1,:])*8.1))
print "d Vc / dR = %.2f +/- %.2f km/s/kpc" % (out(-all_samples[1:5,0,:]-all_samples[1:5,1,:]))
print "R/Vc x d Vc / dR = %.2f +/- %.2f" % (out(-(all_samples[1:5,0,:]+all_samples[1:5,1,:])\
                                 /(all_samples[1:5,0,:]-all_samples[1:5,1,:])))


# The formal combination using the individual uncertainties gives smaller uncertainties:

# In[28]:


def out(sa):
    int_mean= numpy.mean(sa,axis=1)
    int_var = numpy.var (sa,axis=1)
    return (numpy.sum(int_mean/int_var)/numpy.sum(1./int_var),
            numpy.sqrt(1./numpy.sum(1./int_var)))
print "A = %.2f +/- %.2f km/s/kpc" % (out(all_samples[1:5,0,:]))
print "Ac = %.2f +/- %.2f km/s/kpc" % (out(all_samples[1:5,0,:]+(all_samples[1:5,5,:]-12.24)/8.))
print "B = %.2f +/- %.2f km/s/kpc" % (out(all_samples[1:5,1,:]))
print "C = %.2f +/- %.2f km/s/kpc" % (out(all_samples[1:5,2,:]))
print "K = %.2f +/- %.2f km/s/kpc" % (out(all_samples[1:5,3,:]))
print "u0 = %.2f +/- %.2f km/s" % (out(all_samples[1:5,4,:]))
print "w0 = %.2f +/- %.2f km/s" % (out(all_samples[1:5,6,:]))
print "Derived parameters:"
print "A-B = %.2f +/- %.2f km/s/kpc" % (out(all_samples[1:5,0,:]-all_samples[1:5,1,:]))
print "K+C = %.2f +/- %.2f km/s/kpc" % (out(all_samples[1:5,2,:]+all_samples[1:5,3,:]))
print "Vc = %.2f +/- %.2f km/s" % (out((all_samples[1:5,0,:]-all_samples[1:5,1,:])*8.1))
print "d Vc / dR = %.2f +/- %.2f km/s/kpc" % (out(-all_samples[1:5,0,:]-all_samples[1:5,1,:]))
print "R/Vc x d Vc / dR = %.2f +/- %.2f" % (out(-(all_samples[1:5,0,:]+all_samples[1:5,1,:])\
                                 /(all_samples[1:5,0,:]-all_samples[1:5,1,:])))


# What happens when we decorrelate the errors in $\mu_l$, $\mu_b$, and $\varpi$?

# In[29]:


bv_min= numpy.arange(0.,1.21,0.2)
bv_max= numpy.arange(0.2,1.41,0.2)
bf= numpy.zeros((len(bv_min),9))
bf_decorr= numpy.zeros((len(bv_min),9))
for ii,(mi,ma) in enumerate(zip(bv_min,bv_max)):
    tindx= indx*(bv > mi)*(bv < ma)
    print("Working on %i/%i; %i stars" % (ii+1,len(bv_min),numpy.sum(tindx)))
    bf[ii]= fit_sample(tindx,plots=False,disp=False)
    bf_decorr[ii,:]= fit_sample(tindx,plots=False,disp=False,decorrelate_errors=True)


# In[30]:


figsize(14,8)
subplot(2,2,1)
bovy_plot.bovy_plot(bv_med,bf_decorr[:,0]-bf[:,0],'ko',
                    gcf=True,
                    xrange=[-0.2,1.7],
                    yrange=[-2.,2.],
                    ylabel=r'$\Delta A\,(\mathrm{km\,s}^{-1})$')
subplot(2,2,2)
bovy_plot.bovy_plot(bv_med,bf_decorr[:,1]-bf[:,1],'ko',
                    gcf=True,
                    xrange=[-0.2,1.7],
                    yrange=[-10.,10.],
                    ylabel=r'$\Delta B\,(\mathrm{km\,s}^{-1})$')
subplot(2,2,3)
bovy_plot.bovy_plot(bv_med,bf_decorr[:,2]-bf[:,2],'ko',
                    gcf=True,
                    xrange=[-0.2,1.7],
                    yrange=[-2.,2.],
                    xlabel=r'$B-V$',ylabel=r'$\Delta C\,(\mathrm{km\,s}^{-1})$')
subplot(2,2,4)
bovy_plot.bovy_plot(bv_med,bf_decorr[:,3]-bf[:,3],'ko',
                    gcf=True,
                    xrange=[-0.2,1.7],
                    yrange=[-10.,10.],
                    xlabel=r'$B-V$',ylabel=r'$\Delta K\,(\mathrm{km\,s}^{-1})$')


# In[31]:


print("Delta A:",numpy.mean(bf_decorr[1:5,0]-bf[1:5,0]))
print("Delta B:",numpy.mean(bf_decorr[1:5,1]-bf[1:5,1]))
print("Delta C:",numpy.mean(bf_decorr[1:5,2]-bf[1:5,2]))
print("Delta K:",numpy.mean(bf_decorr[1:5,3]-bf[1:5,3]))


# ## What does non-axisymmetry do to the Oort constants?
# 
# Given that we find that $C$ and $K$ are significantly non-zero, let's investigate briefly what kinds of $C$ and $K$ we would get from non-axisymmetric Milky-Way models. For example, for the elliptical-disk model, Bovy (2015) showed that $C$ and $K$ always have the opposite sign, inconsistent with what we find here.
# 
# We first look into the predicted $C$ and $K$ for the standard bar model and then look at what kinds of deviations from their axisymmetric value we get for $A$ and $B$. The bar model that we use is that from Dehnen (2000) as used in Bovy (2010) and Bovy et al. (2015).
# 
# We set up a distribution-function of a warm disk evolved in a flat-rotation-curve+bar potential using $\texttt{galpy}$'s $\texttt{galpy.df.evolveddiskdf}$:

# In[32]:


from galpy.potential import LogarithmicHaloPotential, DehnenBarPotential
from galpy.df import dehnendf, evolveddiskdf
lp= LogarithmicHaloPotential(normalize=1.)
dp= DehnenBarPotential(alpha=0.015)
idfwarm= dehnendf(beta=0.,profileParams=(1./3.,1.,0.15))
edfwarm= evolveddiskdf(idfwarm,[lp,dp],to=dp.tform())


# and then we compute the Oort constants $C$ and $K$:

# In[33]:


gridpoints= 201
oortcwarm, gridwarm, gridrwarm, gridphiwarm= edfwarm.oortC(1.,phi=0.,deg=True,
                                                           t=0.,returnGrids=True,gridpoints=gridpoints,
                                                           derivGridpoints=gridpoints,grid=True,
                                                           derivphiGrid=True,derivRGrid=True)
oortkwarm= edfwarm.oortK(1.,phi=0.,deg=True,t=0.,grid=gridwarm,derivphiGrid=gridphiwarm,derivRGrid=gridrwarm)
print(oortcwarm*220./8.,oortkwarm*220./8.)


# In[ ]: