def plot_sub_traceback(xyzuvw, xyzuvw_cov, times, dim1, dim2, ax, age=None):
"""Plot the 2D traceback of stars with their error ellipses
Parameters
----------
xyzuvw : [nstars, ntimes, 6] array
exact orbital position in XYZ (pc) and UVW (km/s)
xyzuvw_cov : [nstars, ntimes, 6, 6] array
covariance matrix for each star at each time step
times : [ntimes] array
times corresponding to each traceback step
dim1, dim2 : ints
Denotes which phase dimensions will be plotted [0,1,2,3,4,5] -> [XYZUVW]
ax : pyplot axes object
the axes to be plotted in
age : float
the true age of the group, if known
Notes
-----
TODO: Maybe adjust so it takes in "fixed times" as its time parameter
This way I can map fits from the fixed time fits onto the traceback plot
"""
labels = ['X [pc]', 'Y [pc]', 'Z [pc]', 'U [km/s]', 'V [km/s]', 'W [km/s]']
# X and U are negative as per convention
axis_ranges = [-300, 300, 300, -50, 50, 50]
nstars = len(xyzuvw)
# setting max time index
if age is None:
mt_ix = -1
else:
mt_ix = (np.abs(times - age)).argmin()
logging.info("Max time index is {}".format(mt_ix))
logging.info("--> approx age of {}, true_age: {}".\
format(times[mt_ix], age))
cov_ix1 = [[dim1, dim2], [dim1, dim2]]
cov_ix2 = [[dim1, dim1], [dim2, dim2]]
for i in range(nstars):
ax.plot(xyzuvw[i, :mt_ix, dim1], xyzuvw[i, :mt_ix, dim2], 'b-')
cov_start = xyzuvw_cov[i, 0, cov_ix1, cov_ix2]
cov_end = xyzuvw_cov[i, mt_ix, cov_ix1, cov_ix2]
ee.plot_cov_ellipse(
cov_start,
[xyzuvw[i, 0, dim1], xyzuvw[i, 0, dim2]],
color='g', alpha=0.1, ax=ax
)
ee.plot_cov_ellipse(
cov_end,
[xyzuvw[i, mt_ix, dim1], xyzuvw[i, mt_ix, dim2]],
color='r', alpha=0.1, ax=ax
)
ax.set(aspect='equal')
ax.set_xlabel(labels[dim1])
ax.set_ylabel(labels[dim2])
ax.set_xlim(-axis_ranges[dim1], axis_ranges[dim1]) # note inverse X axis
ax.set_ylim(-axis_ranges[dim2], axis_ranges[dim2])
def plot_sub_traceback(xyzuvw, xyzuvw_cov, times, dim1, dim2, ax, age=None):
"""Plot the 2D traceback of stars with their error ellipses
Parameters
----------
xyzuvw : [nstars, ntimes, 6] array
exact orbital position in XYZ (pc) and UVW (km/s)
xyzuvw_cov : [nstars, ntimes, 6, 6] array
covariance matrix for each star at each time step
times : [ntimes] array
times corresponding to each traceback step
dim1, dim2 : ints
Denotes which phase dimensions will be plotted [0,1,2,3,4,5] -> [XYZUVW]
ax : pyplot axes object
the axes to be plotted in
age : float
the true age of the group, if known
Notes
-----
TODO: Maybe adjust so it takes in "fixed times" as its time parameter
This way I can map fits from the fixed time fits onto the traceback plot
"""
labels = ['X [pc]', 'Y [pc]', 'Z [pc]', 'U [km/s]', 'V [km/s]', 'W [km/s]']
# X and U are negative as per convention
axis_ranges = [-300, 300, 300, -50, 50, 50]
nstars = len(xyzuvw)
# setting max time index
if age is None:
mt_ix = -1
else:
mt_ix = (np.abs(times - age)).argmin()
logging.info("Max time index is {}".format(mt_ix))
logging.info("--> approx age of {}, true_age: {}".\
format(times[mt_ix], age))
cov_ix1 = [[dim1, dim2], [dim1, dim2]]
cov_ix2 = [[dim1, dim1], [dim2, dim2]]
for i in range(nstars):
ax.plot(xyzuvw[i, :mt_ix, dim1], xyzuvw[i, :mt_ix, dim2], 'b-')
cov_start = xyzuvw_cov[i, 0, cov_ix1, cov_ix2]
cov_end = xyzuvw_cov[i, mt_ix, cov_ix1, cov_ix2]
ee.plot_cov_ellipse(cov_start,
[xyzuvw[i, 0, dim1], xyzuvw[i, 0, dim2]],
color='g',
alpha=0.1,
ax=ax)
ee.plot_cov_ellipse(cov_end,
[xyzuvw[i, mt_ix, dim1], xyzuvw[i, mt_ix, dim2]],
color='r',
alpha=0.1,
ax=ax)
ax.set(aspect='equal')
ax.set_xlabel(labels[dim1])
ax.set_ylabel(labels[dim2])
ax.set_xlim(-axis_ranges[dim1], axis_ranges[dim1]) # note inverse X axis
ax.set_ylim(-axis_ranges[dim2], axis_ranges[dim2])
def plot_2d(val1, val2, save=False):
col = 'rgbcmyk'
plt.clf()
xmin = mu[:, val1] - sigma[:, val1]*2.1
xmax = mu[:, val1] + sigma[:, val1]*2.1
ymin = mu[:, val2] - sigma[:, val2]*2.1
ymax = mu[:, val2] + sigma[:, val2]*2.1
plt.xlim(xmin.min(), xmax.max())
plt.ylim(ymin.min(), ymax.max())
plt.xlabel(labels[val1], fontsize=16)
plt.ylabel(labels[val2], fontsize=16)
ppi = (pi - pi.min() + 0.01) / (pi.max() - pi.min() + 0.01)
for j in xrange(len(pi)):
i = len(pi) - j - 1
sig = sigma2d[i][[val1, val2]][:, [val1, val2]]
plot_cov_ellipse(sig, mu[i][[val1, val2]], facecolor=col[i], alpha=ppi[i])
if save:
plt.savefig('%s_%s.png' % (labels[val1], labels[val2]))
else:
plt.show()
def plot_circle(sigma_fixed, centroids, vol=.10):
error.plot_cov_ellipse(sigma_fixed[0],
centroids[0],
volume=vol,
a=.9,
ec=[1, 0, 0])
error.plot_cov_ellipse(sigma_fixed[1],
centroids[1],
volume=vol,
a=.9,
ec=[1, 0, 0])
error.plot_cov_ellipse(sigma_fixed[2],
centroids[2],
volume=vol,
a=.9,
ec=[1, 0, 0])
def plot_error_ellipse(sigma, mu):
#error.plot_cov_ellipse(sigma[0],mu[0])
#error.plot_cov_ellipse(sigma[1],mu[1])
#error.plot_cov_ellipse(sigma[2],mu[2])
#error.plot_cov_ellipse(sigma[0],mu[0], volume=.25)
#error.plot_cov_ellipse(sigma[1],mu[1], volume=.25)
#error.plot_cov_ellipse(sigma[2],mu[2], volume=.25)
#error.plot_cov_ellipse(sigma[0],mu[0], volume=.75, fc='red')
#error.plot_cov_ellipse(sigma[1],mu[1], volume=.75, fc='yellow')
#error.plot_cov_ellipse(sigma[2],mu[2], volume=.75, fc='blue')
error.plot_cov_ellipse(sigma[0], mu[0], volume=.75, a=.9, ec=[1, 0, 0])
error.plot_cov_ellipse(sigma[1], mu[1], volume=.75, a=.9, ec=[1, 0, 0])
error.plot_cov_ellipse(sigma[2], mu[2], volume=.75, a=.9, ec=[1, 0, 0])
def traceback(stars,
times,
max_plot_error=50,
plotit=False,
savefile='',
dims=[1, 2],
xoffset=[],
yoffset=[],
text_ix=[],
axis_range=[],
plot_text=True,
return_tb=False):
"""Trace back stellar orbits
Parameters
----------
stars:
An astropy Table object that includes
the columns 'RAdeg', 'DEdeg', 'Plx', 'pmRA', 'pmDE', 'RV', 'Name',
and the uncertainties:
'plx_sig', 'pmRA_sig', 'pmDEC_sig', 'RV_sig'. Optionally, include
the correlations between parameters: 'c_plx_pmRA', 'c_plx_pmDEC' and
'c_pmDEC_pmRA'.
times: float array
Times to trace back, in Myr. Note that positive numbers are going backwards in time.
max_plot_error: float
Maximum positional error in pc to allow plotting.
dims: list with 2 ints
Dimensions to be plotted (out of xyzuvw)
xoffset, yoffset: nstars long list
Offsets for text position in star labels.
Returns (opt)
-------------
star_pars: dict
'stars' : astropy table
stellar astrometry
'times' : [ntimes] array
traceback times
'xyzuvw' : [nstars, 6] array
mean phase kinematics of each star
'xyzuvw_cov' : [nstrs, 6, 6] array
covariance matrix of kinematics for each star
"""
nstars = len(stars)
nts = len(times)
#Positions and velocities in the co-rotating solar reference frame.
xyzuvw = np.zeros((nstars, nts, 6))
#Derivative of the past xyzuvw with respect to the present xyzuvw
xyzuvw_jac = np.zeros((nstars, nts, 6, 6))
#Past covariance matrices
xyzuvw_cov = np.zeros((nstars, nts, 6, 6))
#Trace back the local standard of rest. done manually here so that it doesn't
#have to be recomputed many times.
lsr_orbit = get_lsr_orbit(times)
#Delta parameters for numerical derivativd
dp = 1e-3
if plotit:
plt.clf()
dim1 = dims[0]
dim2 = dims[1]
cov_ix1 = [[dim1, dim2], [dim1, dim2]]
cov_ix2 = [[dim1, dim1], [dim2, dim2]]
axis_titles = [
'X (pc)', 'Y (pc)', 'Z (pc)', 'U (km/s)', 'V (km/s)', 'W (km/s)'
]
#Put some defaults in
if len(xoffset) == 0:
xoffset = np.zeros(nstars)
yoffset = np.zeros(nstars)
if len(text_ix) == 0:
text_ix = range(nstars)
if len(axis_range) == 0:
axis_range = [-100, 100, -100, 100]
#WARNING: The following is messy, but different types of input are
#accepted for now...
try:
colnames = stars.names
except:
colnames = stars.columns
for i in range(nstars):
if False:
if (i + 1) % 100 == 0:
print("Done {0:d} of {1:d} stars.".format(i + 1, nstars))
star = stars[i]
#Some defaults for correlations.
parallax_pmra_corr = 0.0
parallax_pmdec_corr = 0.0
pmra_pmdec_corr = 0.0
#Are we using our joint HIPPARCOS/TGAS table?
if 'ra_adopt' in colnames:
RAdeg = star['ra_adopt']
DEdeg = star['dec_adopt']
RV = star['rv_adopt']
e_RV = star['rv_adopt_error']
#Do we have a TGAS entry?
tgas_entry = True
try:
if (stars['parallax_1'].mask)[i]:
tgas_entry = False
except:
if star['parallax_1'] != star['parallax_1']:
tgas_entry = False
if tgas_entry:
Plx = star['parallax_1']
e_Plx = star['parallax_error']
pmRA = star['pmra_1']
e_pmRA = star['pmra_error']
pmDE = star['pmdec']
e_pmDE = star['pmdec_error']
parallax_pmra_corr = star['parallax_pmra_corr']
parallax_pmdec_corr = star['parallax_pmdec_corr']
pmra_pmdec_corr = star['pmra_pmdec_corr']
else:
Plx = star['Plx']
e_Plx = star['e_Plx']
pmDE = star['pmDE']
try:
e_pmDE = star['e_pmDE']
except:
e_pmDE = 0.5
pmRA = star['pmRA_2']
try:
e_pmRA = star['e_pmRA']
except:
e_pmRA = 0.5
else:
DEdeg = star['DEdeg']
#RAVE Dataset
if 'HRV' in star.columns:
RAdeg = star['RAdeg']
RV = star['HRV']
e_RV = star['e_HRV']
Plx = star['plx']
e_Plx = star['e_plx']
pmRA = star['pmRAU4']
e_pmRA = star['e_pmRAU4']
pmDE = star['pmDEU4']
e_pmDE = star['e_pmDEU4']
#HIPPARCOS
else:
if 'RAdeg' in star.columns:
RAdeg = star['RAdeg']
else:
RAdeg = star['RAhour'] * 15.0
RV = star['RV']
e_RV = star['e_RV']
Plx = star['Plx']
e_Plx = star['e_Plx']
pmRA = star['pmRA']
e_pmRA = star['e_pmRA']
pmDE = star['pmDE']
e_pmDE = star['e_pmDE']
params = np.array([RAdeg, DEdeg, Plx, pmRA, pmDE, RV])
xyzuvw[i] = integrate_xyzuvw(params, times, lsr_orbit, MWPotential2014)
#if star['Name1'] == 'HIP 12545':
# pdb.set_trace()
#Create numerical derivatives
for j in range(6):
params_plus = params.copy()
params_plus[j] += dp
xyzuvw_plus = integrate_xyzuvw(params_plus, times, lsr_orbit,
MWPotential2014)
for k in range(nts):
xyzuvw_jac[i, k, j, :] = (xyzuvw_plus[k] - xyzuvw[i, k]) / dp
#Now that we've got the jacobian, find the modified covariance matrix.
cov_obs = np.zeros((6, 6))
cov_obs[
0,
0] = 1e-1**2 #Nominal 0.1 degree. !!! There is almost a floating underflow in fit_group due to this
cov_obs[
1,
1] = 1e-1**2 #Nominal 0.1 degree. !!! Need to think about this more
cov_obs[2, 2] = e_Plx**2
cov_obs[3, 3] = e_pmRA**2
cov_obs[4, 4] = e_pmDE**2
cov_obs[5, 5] = e_RV**2
#Add covariances.
cov_obs[2, 3] = e_Plx * e_pmRA * parallax_pmra_corr
cov_obs[2, 4] = e_Plx * e_pmDE * parallax_pmdec_corr
cov_obs[3, 4] = e_pmRA * e_pmDE * pmra_pmdec_corr
for (j, k) in [(2, 3), (2, 4), (3, 4)]:
cov_obs[k, j] = cov_obs[j, k]
#Create the covariance matrix from the Jacobian. See e.g.:
#https://en.wikipedia.org/wiki/Covariance#A_more_general_identity_for_covariance_matrices
#Think of the Jacobian as just a linear transformation between the present and the past in
#"local" coordinates.
for k in range(nts):
xyzuvw_cov[i, k] = np.dot(np.dot(xyzuvw_jac[i, k].T, cov_obs),
xyzuvw_jac[i, k])
# !!! TODO: ^^ seems to be executing J.T X J instead of J X J.T
#Test for problems...
if (np.linalg.eigvalsh(xyzuvw_cov[i, k])[0] < 0):
pdb.set_trace()
#Plot beginning and end points, plus their the uncertainties from the
#covariance matrix.
# plt.plot(xyzuvw[i,0,0],xyzuvw[i,0,1],'go')
# plt.plot(xyzuvw[i,-1,0],xyzuvw[i,-1,1],'ro')
#Only plot if uncertainties are low...
if plotit:
cov_end = xyzuvw_cov[i, -1, cov_ix1, cov_ix2]
if (np.sqrt(cov_end.trace()) < max_plot_error):
if plot_text:
if i in text_ix:
plt.text(xyzuvw[i, 0, dim1] * 1.1 + xoffset[i],
xyzuvw[i, 0, dim2] * 1.1 + yoffset[i],
star['Name'],
fontsize=11)
plt.plot(xyzuvw[i, :, dim1], xyzuvw[i, :, dim2], 'b-')
plot_cov_ellipse(xyzuvw_cov[i, 0, cov_ix1, cov_ix2],
[xyzuvw[i, 0, dim1], xyzuvw[i, 0, dim2]],
color='g',
alpha=1)
plot_cov_ellipse(cov_end,
[xyzuvw[i, -1, dim1], xyzuvw[i, -1, dim2]],
color='r',
alpha=0.2)
#else:
# print(star['Name'] + " not plottable (errors too high)")
if plotit:
plt.xlabel(axis_titles[dim1])
plt.ylabel(axis_titles[dim2])
plt.axis(axis_range)
plt.savefig("myplot.eps")
if len(savefile) > 0:
if savefile[-3:] == 'pkl':
with open(savefile, 'w') as fp:
pickle.dump((stars, times, xyzuvw, xyzuvw_cov), fp)
#Stars is an astropy.Table of stars
elif (savefile[-3:] == 'fit') or (savefile[-4:] == 'fits'):
hl = pyfits.HDUList()
hl.append(pyfits.PrimaryHDU())
hl.append(pyfits.BinTableHDU(stars))
hl.append(pyfits.ImageHDU(times))
hl.append(pyfits.ImageHDU(xyzuvw))
hl.append(pyfits.ImageHDU(xyzuvw_cov))
hl.writeto(savefile, clobber=True)
else:
print("Unknown File Type!")
raise UserWarning
else:
return xyzuvw
if return_tb:
star_pars = {
'stars': stars,
'times': times,
'xyzuvw': xyzuvw,
'xyzuvw_cov': xyzuvw_cov
}
return star_pars
def traceback(stars,times,max_plot_error=50,plotit=False, savefile='',
dims=[1,2],xoffset=[],yoffset=[],text_ix=[],axis_range=[],
plot_text=True, return_tb=False):
"""Trace back stellar orbits
Parameters
----------
stars:
An astropy Table object that includes
the columns 'RAdeg', 'DEdeg', 'Plx', 'pmRA', 'pmDE', 'RV', 'Name',
and the uncertainties:
'plx_sig', 'pmRA_sig', 'pmDEC_sig', 'RV_sig'. Optionally, include
the correlations between parameters: 'c_plx_pmRA', 'c_plx_pmDEC' and
'c_pmDEC_pmRA'.
times: float array
Times to trace back, in Myr. Note that positive numbers are going backwards in time.
max_plot_error: float
Maximum positional error in pc to allow plotting.
dims: list with 2 ints
Dimensions to be plotted (out of xyzuvw)
xoffset, yoffset: nstars long list
Offsets for text position in star labels.
Returns (opt)
-------------
star_pars: dict
'stars' : astropy table
stellar astrometry
'times' : [ntimes] array
traceback times
'xyzuvw' : [nstars, 6] array
mean phase kinematics of each star
'xyzuvw_cov' : [nstrs, 6, 6] array
covariance matrix of kinematics for each star
"""
nstars = len(stars)
nts = len(times)
#Positions and velocities in the co-rotating solar reference frame.
xyzuvw = np.zeros( (nstars,nts,6) )
#Derivative of the past xyzuvw with respect to the present xyzuvw
xyzuvw_jac = np.zeros( (nstars,nts,6,6) )
#Past covariance matrices
xyzuvw_cov = np.zeros( (nstars,nts,6,6) )
#Trace back the local standard of rest. done manually here so that it doesn't
#have to be recomputed many times.
lsr_orbit = get_lsr_orbit(times)
#Delta parameters for numerical derivativd
dp = 1e-3
if plotit:
plt.clf()
dim1=dims[0]
dim2=dims[1]
cov_ix1 = [[dim1,dim2],[dim1,dim2]]
cov_ix2 = [[dim1,dim1],[dim2,dim2]]
axis_titles=['X (pc)','Y (pc)','Z (pc)','U (km/s)','V (km/s)','W (km/s)']
#Put some defaults in
if len(xoffset)==0:
xoffset = np.zeros(nstars)
yoffset = np.zeros(nstars)
if len(text_ix)==0:
text_ix = range(nstars)
if len(axis_range)==0:
axis_range = [-100,100,-100,100]
#WARNING: The following is messy, but different types of input are
#accepted for now...
try:
colnames = stars.names
except:
colnames = stars.columns
for i in range(nstars):
if False:
if (i+1) % 100 ==0:
print("Done {0:d} of {1:d} stars.".format(i+1, nstars))
star = stars[i]
#Some defaults for correlations.
parallax_pmra_corr=0.0
parallax_pmdec_corr=0.0
pmra_pmdec_corr=0.0
#Are we using our joint HIPPARCOS/TGAS table?
if 'ra_adopt' in colnames:
RAdeg = star['ra_adopt']
DEdeg = star['dec_adopt']
RV = star['rv_adopt']
e_RV = star['rv_adopt_error']
#Do we have a TGAS entry?
tgas_entry = True
try:
if (stars['parallax_1'].mask)[i]:
tgas_entry=False
except:
if star['parallax_1']!=star['parallax_1']:
tgas_entry=False
if tgas_entry:
Plx = star['parallax_1']
e_Plx = star['parallax_error']
pmRA = star['pmra_1']
e_pmRA = star['pmra_error']
pmDE = star['pmdec']
e_pmDE = star['pmdec_error']
parallax_pmra_corr = star['parallax_pmra_corr']
parallax_pmdec_corr = star['parallax_pmdec_corr']
pmra_pmdec_corr = star['pmra_pmdec_corr']
else:
Plx = star['Plx']
e_Plx = star['e_Plx']
pmDE = star['pmDE']
try:
e_pmDE = star['e_pmDE']
except:
e_pmDE = 0.5
pmRA = star['pmRA_2']
try:
e_pmRA = star['e_pmRA']
except:
e_pmRA = 0.5
else:
DEdeg = star['DEdeg']
#RAVE Dataset
if 'HRV' in star.columns:
RAdeg = star['RAdeg']
RV = star['HRV']
e_RV = star['e_HRV']
Plx = star['plx']
e_Plx = star['e_plx']
pmRA = star['pmRAU4']
e_pmRA = star['e_pmRAU4']
pmDE = star['pmDEU4']
e_pmDE = star['e_pmDEU4']
#HIPPARCOS
else:
if 'RAdeg' in star.columns:
RAdeg = star['RAdeg']
else:
RAdeg = star['RAhour']*15.0
RV = star['RV']
e_RV = star['e_RV']
Plx = star['Plx']
e_Plx = star['e_Plx']
pmRA = star['pmRA']
e_pmRA = star['e_pmRA']
pmDE = star['pmDE']
e_pmDE = star['e_pmDE']
params = np.array([RAdeg,DEdeg,Plx,pmRA,pmDE,RV])
xyzuvw[i] = integrate_xyzuvw(params,times,lsr_orbit,MWPotential2014)
#if star['Name1'] == 'HIP 12545':
# pdb.set_trace()
#Create numerical derivatives
for j in range(6):
params_plus = params.copy()
params_plus[j] += dp
xyzuvw_plus = integrate_xyzuvw(params_plus,times,lsr_orbit,MWPotential2014)
for k in range(nts):
xyzuvw_jac[i,k,j,:] = (xyzuvw_plus[k] - xyzuvw[i,k])/dp
#Now that we've got the jacobian, find the modified covariance matrix.
cov_obs = np.zeros( (6,6) )
cov_obs[0,0] = 1e-1**2 #Nominal 0.1 degree. !!! There is almost a floating underflow in fit_group due to this
cov_obs[1,1] = 1e-1**2 #Nominal 0.1 degree. !!! Need to think about this more
cov_obs[2,2] = e_Plx**2
cov_obs[3,3] = e_pmRA**2
cov_obs[4,4] = e_pmDE**2
cov_obs[5,5] = e_RV**2
#Add covariances.
cov_obs[2,3] = e_Plx * e_pmRA * parallax_pmra_corr
cov_obs[2,4] = e_Plx * e_pmDE * parallax_pmdec_corr
cov_obs[3,4] = e_pmRA * e_pmDE * pmra_pmdec_corr
for (j,k) in [(2,3),(2,4),(3,4)]:
cov_obs[k,j] = cov_obs[j,k]
#Create the covariance matrix from the Jacobian. See e.g.:
#https://en.wikipedia.org/wiki/Covariance#A_more_general_identity_for_covariance_matrices
#Think of the Jacobian as just a linear transformation between the present and the past in
#"local" coordinates.
for k in range(nts):
xyzuvw_cov[i,k] = np.dot(np.dot(xyzuvw_jac[i,k].T,cov_obs),xyzuvw_jac[i,k])
# !!! TODO: ^^ seems to be executing J.T X J instead of J X J.T
#Test for problems...
if (np.linalg.eigvalsh(xyzuvw_cov[i,k])[0]<0):
pdb.set_trace()
#Plot beginning and end points, plus their the uncertainties from the
#covariance matrix.
# plt.plot(xyzuvw[i,0,0],xyzuvw[i,0,1],'go')
# plt.plot(xyzuvw[i,-1,0],xyzuvw[i,-1,1],'ro')
#Only plot if uncertainties are low...
if plotit:
cov_end = xyzuvw_cov[i,-1,cov_ix1,cov_ix2]
if (np.sqrt(cov_end.trace()) < max_plot_error):
if plot_text:
if i in text_ix:
plt.text(xyzuvw[i,0,dim1]*1.1 + xoffset[i],xyzuvw[i,0,dim2]*1.1 + yoffset[i],star['Name'],fontsize=11)
plt.plot(xyzuvw[i,:,dim1],xyzuvw[i,:,dim2],'b-')
plot_cov_ellipse(xyzuvw_cov[i,0,cov_ix1,cov_ix2], [xyzuvw[i,0,dim1],xyzuvw[i,0,dim2]],color='g',alpha=1)
plot_cov_ellipse(cov_end, [xyzuvw[i,-1,dim1],xyzuvw[i,-1,dim2]],color='r',alpha=0.2)
#else:
# print(star['Name'] + " not plottable (errors too high)")
if plotit:
plt.xlabel(axis_titles[dim1])
plt.ylabel(axis_titles[dim2])
plt.axis(axis_range)
plt.savefig("myplot.eps")
if len(savefile)>0:
if savefile[-3:] == 'pkl':
with open(savefile,'w') as fp:
pickle.dump((stars,times,xyzuvw,xyzuvw_cov),fp)
#Stars is an astropy.Table of stars
elif (savefile[-3:] == 'fit') or (savefile[-4:] == 'fits'):
hl = pyfits.HDUList()
hl.append(pyfits.PrimaryHDU())
hl.append(pyfits.BinTableHDU(stars))
hl.append(pyfits.ImageHDU(times))
hl.append(pyfits.ImageHDU(xyzuvw))
hl.append(pyfits.ImageHDU(xyzuvw_cov))
hl.writeto(savefile, clobber=True)
else:
print("Unknown File Type!")
raise UserWarning
else:
return xyzuvw
if return_tb:
star_pars = {
'stars':stars, 'times':times, 'xyzuvw':xyzuvw,
'xyzuvw_cov':xyzuvw_cov
}
return star_pars
st = MultivariateTDistribution(DIMS, xy.mean(axis=1)*1.1, np.diag(xy.var(axis=1)), 3)
da = xy.T #[:500]
ds = DataSet()
ds.fromArray(da)
m = MixtureModel(1, [1], [st])
print m
m.EM(ds, 60, 0.1)
print m
#import ipdb; ipdb.set_trace()
#m2 = mix.MixtureModel(2, [0.8, 0.2], [m.components[0].distList[0], d2], compFix=[0, 0])
for _ in xrange(6):
m = mixturate(m, xy.mean(axis=1)-10., np.diag(xy.var(axis=1)))
m.randMaxEM(ds, 3, 30, 0.1)
print m
import joblib
joblib.dump(m, 'test2.mix', compress=3)
pl.plotData(da[:, :2])
col = 'rgbcmyk'
icol = 0
for icomp, comp in enumerate(m.components):
dist = comp.distList[0]
print m.pi[icomp], dist.df, dist.mu, np.sqrt(np.diag(dist.sigma))
for comp in m.components:
model = comp.distList[0]
plot_cov_ellipse(model.sigma[:2, :2], model.mu[:2], facecolor=col[icol])
icol += 1
pylab.show()
cov_obs[3,3] = star['pmRA_sig']**2
cov_obs[4,4] = star['pmDEC_sig']**2
cov_obs[5,5] = star['RV_sig']**2
for k in range(nts):
xyzuvw_cov[i,k] = np.dot(np.dot(xyzuvw_jac[i,k].T,cov_obs),xyzuvw_jac[i,k])
#Plot beginning and end points, plus their the uncertainties from the
#covariance matrix.
plt.plot(xyzuvw[i,0,0],xyzuvw[i,0,1],'go') #%TC
plt.plot(xyzuvw[i,-1,0],xyzuvw[i,-1,1],'ro') #%TC
#Only plot if uncertainties are low...
cov_end = xyzuvw_cov[i,-1,cov_ix1,cov_ix2]
if (np.sqrt(cov_end.trace()) < error_lim):
if i in text_ix:
plt.text(xyzuvw[i,0,dim1]*1.1 + xoffset[i],xyzuvw[i,0,dim2]*1.1 + yoffset[i],star['Name'],fontsize=11)
plt.plot(xyzuvw[i,:,dim1],xyzuvw[i,:,dim2],'b-')
plot_cov_ellipse(xyzuvw_cov[i,0,cov_ix1,cov_ix2], [xyzuvw[i,0,dim1],xyzuvw[i,0,dim2]],color='g',alpha=1)
plot_cov_ellipse(cov_end, [xyzuvw[i,-1,dim1],xyzuvw[i,-1,dim2]],color='r',alpha=0.2)
# w,v = np.linalg.eig(cov_xz)
# pdb.set_trace()
print(time.time()-tic)
plt.xlabel(axis_titles[dim1])
plt.ylabel(axis_titles[dim2])
plt.axis(axis_range)
xyz_cov = np.array([[ 34.25840977, 35.33697325, 56.24666544],
[ 35.33697325, 46.18069795, 66.76389275],
[ 56.24666544, 66.76389275, 109.98883853]])
xyz = [ -6.221, 63.288, 23.408]
plot_cov_ellipse(xyz_cov[cov_ix1,cov_ix2],[xyz[dim1],xyz[dim2]],alpha=0.5,color='k')
