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')