def stop(): # calculate the new mean and cov new_mean = trace_forward(mean, age) new_cov = tf.transform_covmatrix(cov, trace_forward, mean, dim=6, args=(age,)) new_eigvals = np.linalg.eigvalsh(new_cov) estimated_cov = np.cov(new_stars.T) estimated_eigvals = np.linalg.eigvalsh(estimated_cov) if plotit: #plt.clf() plt.plot(new_stars[:,0], new_stars[:,1], 'r.') #plt.hist2d(stars[:,0], stars[:,1], bins=20) ymin, ymax = plt.ylim() xmin, xmax = plt.xlim() upper = max(xmax, ymax) lower = min(xmin, ymin) plt.xlim(upper, lower) plt.ylim(lower, upper) chronostar.fitplotter.plotCovEllipse( new_cov[:2,:2], new_mean, color='r', alpha=0.1 ) plt.savefig("temp_plots/trace_forward{}.png".format(cnt)) plt.show() assert np.allclose(new_eigvals, estimated_eigvals, rtol=.5)
def stop(): # calculate the new mean and cov new_mean = trace_forward(mean, age) new_cov = tf.transform_covmatrix(cov, trace_forward, mean, dim=6, args=(age, )) new_eigvals = np.linalg.eigvalsh(new_cov) estimated_cov = np.cov(new_stars.T) estimated_eigvals = np.linalg.eigvalsh(estimated_cov) if plotit: #plt.clf() plt.plot(new_stars[:, 0], new_stars[:, 1], 'r.') #plt.hist2d(stars[:,0], stars[:,1], bins=20) ymin, ymax = plt.ylim() xmin, xmax = plt.xlim() upper = max(xmax, ymax) lower = min(xmin, ymin) plt.xlim(upper, lower) plt.ylim(lower, upper) chronostar.fitplotter.plotCovEllipse(new_cov[:2, :2], new_mean, color='r', alpha=0.1) plt.savefig("temp_plots/trace_forward{}.png".format(cnt)) plt.show() assert np.allclose(new_eigvals, estimated_eigvals, rtol=.5)
def dataGatherer(res_dir='', save_dir='', data_dir='', xyzuvw_file='', title='', file_stem=''): """ Provided with a results directory, tries to find all she needs, then plots Parameters ---------- """ covs = {} means = {} star_pars = {} chain_file = res_dir + "final_chain.npy" lnprob_file = res_dir + "final_lnprob.npy" origin_file = res_dir + "origins.npy" if not xyzuvw_file: logging.info("No xyzuvw filename provided. Must be synth fit yes?") xyzuvw_file = res_dir + "xyzuvw_now.fits" chain = np.load(chain_file) chain = np.array([chain]) lnprob = np.load(lnprob_file) best_sample = dt.getBestSample(chain, lnprob) best_group = chronostar.component.Component(best_sample, form=len(best_sample) == 9, internal=True) star_pars['xyzuvw'] = fits.getdata(xyzuvw_file, 1) star_pars['xyzuvw_cov'] = fits.getdata(xyzuvw_file, 2) try: origins = np.load(origin_file).item() means['origin_then'] = np.array([origins.mean]) covs['origin_then'] = np.array([origins.generateCovMatrix()]) except IOError: logging.info("No origins file: {}".format(origin_file)) means['fitted_then'] = np.array([best_group.mean]) means['fitted_now'] =\ np.array([torb.trace_cartesian_orbit(best_group.mean, best_group.age)]) covs['fitted_then'] = np.array([best_group.generateCovMatrix()]) covs['fitted_now'] =\ np.array([ tf.transform_covmatrix(covs['fitted_then'][0], torb.trace_cartesian_orbit, means['fitted_then'][0], args=(best_group.age,True) ) ]) plot_hexplot(star_pars, means, covs, chain, iter_count=0, save_dir=save_dir, file_stem=file_stem, title=title) plotXYandZW(star_pars, means, covs, chain, iter_count=0, save_dir=save_dir, file_stem=file_stem, title=title)
def get_stable_expons(comp, rtol=1e-2, atol=1e-6, lo_exp=-10, hi_exp=2): expons = np.arange(lo_exp, hi_exp) # ref_cov_now = comp.get_covmatrix_now() ref_cov_now = transform.transform_covmatrix(comp.get_covmatrix(), comp.trace_orbit_func, loc=comp.get_mean(), h=1. * 10.**-3, args=(comp.get_age(), )) covs_now = [] for expon in expons: covs_now.append( transform.transform_covmatrix(comp.get_covmatrix(), comp.trace_orbit_func, loc=comp.get_mean(), h=1. * 10.**expon, args=(comp.get_age(), ))) return expons[np.where([ np.allclose(cov_mat, ref_cov_now, rtol=rtol, atol=atol) for cov_mat in covs_now ])]
def lnlike(pars, star_pars, z=None, return_lnols=False): """Computes the log-likelihood for a fit to a group. The emcee parameters encode the modelled origin point of the stars. Using the parameters, a mean and covariance in 6D space are constructed as well as an age. The kinematics are then projected forward to the current age and compared with the current stars' XYZUVW values (and uncertainties) Parameters ---------- pars Parameters describing the group model being fitted star_pars traceback data being fitted to (retired) z array of weights [0.0 - 1.0] for each star, describing how likely they are members of group to be fitted. Returns ------- lnlike the logarithm of the likelihood of the fit """ # convert pars into covariance matrix mean_then = pars[0:6] cov_then = generate_cov(pars) age = pars[8] mean_now = trace_forward(mean_then, age) cov_now = tf.transform_covmatrix(cov_then, trace_forward, mean_then, dim=6, args=(age, )) star_covs = star_pars['xyzuvw_cov'][:, 0] star_mns = star_pars['xyzuvw'][:, 0] nstars = star_mns.shape[0] lnols = get_lnoverlaps(cov_now, mean_now, star_covs, star_mns, nstars) if return_lnols: return lnols # prior on covariance matrix incorporated into parametrisation of dX and dV return np.sum(lnols * z)
def lnlike(pars, star_pars, z=None, return_lnols=False): """Computes the log-likelihood for a fit to a group. The emcee parameters encode the modelled origin point of the stars. Using the parameters, a mean and covariance in 6D space are constructed as well as an age. The kinematics are then projected forward to the current age and compared with the current stars' XYZUVW values (and uncertainties) Parameters ---------- pars Parameters describing the group model being fitted star_pars traceback data being fitted to (retired) z array of weights [0.0 - 1.0] for each star, describing how likely they are members of group to be fitted. Returns ------- lnlike the logarithm of the likelihood of the fit """ # convert pars into covariance matrix mean_then = pars[0:6] cov_then = generate_cov(pars) age = pars[8] mean_now = trace_forward(mean_then, age) cov_now = tf.transform_covmatrix( cov_then, trace_forward, mean_then, dim=6, args=(age,) ) star_covs = star_pars['xyzuvw_cov'][:,0] star_mns = star_pars['xyzuvw'][:,0] nstars = star_mns.shape[0] lnols = get_lnoverlaps( cov_now, mean_now, star_covs, star_mns, nstars ) if return_lnols: return lnols # prior on covariance matrix incorporated into parametrisation of dX and dV return np.sum(lnols * z)
def test_polar(): """Generates a collection of points in polar space from a covariance matrix Transforms covariance matrix to cartesian, and compares to a numpy covariance matrix fitted to the points in cartesian space""" logging.basicConfig(level=logging.DEBUG, stream=sys.stdout) logging.info("Beginning polar demo") # setting up polar points r_mean = np.sqrt(10**2 + 10**2) theta_mean = 14 * np.pi / 6. r_std = 1. theta_std = np.pi / 24. C_rt = 0.0 pol_mean = np.array([r_mean, theta_mean]) pol_cov = np.array([ [r_std**2, C_rt * r_std * theta_std], [C_rt * r_std * theta_std, theta_std**2], ]) nsamples = 100000 pol_samples = np.random.multivariate_normal(pol_mean, pol_cov, nsamples) # converting to cartesian manually for comparison res = convertManyPolarToCartesian(pol_samples) #[:, 0], pol_samples[:, 1]) cart_mean = convertPolarToCartesian(pol_mean) logging.debug("Covariance matrix:\n{}".format(pol_cov)) logging.debug("Polar mean: {}".format(pol_mean)) cart_cov = tf.transform_covmatrix(pol_cov, convertPolarToCartesian, pol_mean, dim=2) cart_samples = np.random.multivariate_normal(cart_mean, cart_cov, nsamples) estimated_mean = np.mean(res, axis=0) estimated_cov = np.cov(res.T) assert np.allclose(estimated_mean, cart_mean, rtol=1e-1) assert np.allclose(estimated_cov, cart_cov, rtol=1e-1)
def test_polar(): """Generates a collection of points in polar space from a covariance matrix Transforms covariance matrix to cartesian, and compares to a numpy covariance matrix fitted to the points in cartesian space""" logging.basicConfig(level=logging.DEBUG, stream=sys.stdout) logging.info("Beginning polar demo") # setting up polar points r_mean = np.sqrt(10 ** 2 + 10 ** 2) theta_mean = 14 * np.pi / 6. r_std = 1. theta_std = np.pi / 24. C_rt = 0.0 pol_mean = np.array([r_mean, theta_mean]) pol_cov = np.array([ [r_std ** 2, C_rt * r_std * theta_std], [C_rt * r_std * theta_std, theta_std ** 2], ]) nsamples = 100000 pol_samples = np.random.multivariate_normal(pol_mean, pol_cov, nsamples) # converting to cartesian manually for comparison res = convertManyPolarToCartesian(pol_samples) #[:, 0], pol_samples[:, 1]) cart_mean = convertPolarToCartesian(pol_mean) logging.debug("Covariance matrix:\n{}".format(pol_cov)) logging.debug("Polar mean: {}".format(pol_mean)) cart_cov = tf.transform_covmatrix(pol_cov, convertPolarToCartesian, pol_mean, dim=2) cart_samples = np.random.multivariate_normal(cart_mean, cart_cov, nsamples) estimated_mean = np.mean(res, axis=0) estimated_cov = np.cov(res.T) assert np.allclose(estimated_mean, cart_mean, rtol=1e-1) assert np.allclose(estimated_cov, cart_cov, rtol=1e-1)
def transformAstrCovsToCartesian(astr_covs, astr_arr): """ Converts a covariance matrix from astrometric coords to LSR XYZUVW Parameters ---------- astr_covs: ([nstars, 6, 6] array) values in the diagaonal are the squared errors of (ra, dec, plx, pm_ra, pm_dec, rv), with the offdiagonals the product of the correlation (valued between -1 and 1) and the two intersecting coordinates. astr_arr: ([nstars, 6] array) the measured (mean) astrometric values (ra, dec, plx, pm_ra, pm-dec, rv) """ nstars = astr_arr.shape[0] xyzuvw_covs = np.zeros((nstars, 6, 6)) for ix in range(nstars): xyzuvw_covs[ix] = transform.transform_covmatrix( astr_covs[ix], coordinate.convert_astrometry2lsrxyzuvw, astr_arr[ix], dim=6 ) return xyzuvw_covs
def convertRowToCartesian(row, row_ix=None, nrows=None): dim=6 cart_col_names = ['X', 'Y', 'Z', 'U', 'V', 'W', 'dX', 'dY', 'dZ', 'dU', 'dV', 'dW', 'c_XY', 'c_XZ', 'c_XU', 'c_XV', 'c_XW', 'c_YZ', 'c_YU', 'c_YV', 'c_YW', 'c_ZU', 'c_ZV', 'c_ZW', 'c_UV', 'c_UW', 'c_VW'] try: if row_ix % 100 == 0: print("{:010.7f}% done".format(row_ix / float(nrows) * 100.)) except TypeError: pass astr_mean, astr_cov = datatool.convertRecToArray(row) xyzuvw_mean = coordinate.convert_astrometry2lsrxyzuvw(astr_mean) xyzuvw_cov = transform.transform_covmatrix( astr_cov, coordinate.convert_astrometry2lsrxyzuvw, astr_mean, dim=dim, ) # fill in cartesian mean for col_ix, col_name in enumerate(cart_col_names[:6]): row[col_name] = xyzuvw_mean[col_ix] # fill in standard deviations xyzuvw_stds = np.sqrt(xyzuvw_cov[np.diag_indices(dim)]) for col_ix, col_name in enumerate(cart_col_names[6:12]): row[col_name] = xyzuvw_stds[col_ix] correl_matrix = xyzuvw_cov / xyzuvw_stds / xyzuvw_stds.reshape(6, 1) # fill in correlations for col_ix, col_name in enumerate(cart_col_names[12:]): row[col_name] = correl_matrix[ np.triu_indices(dim, k=1)[0][col_ix], np.triu_indices(dim, k=1)[1][col_ix] ]
print("With {:.2} stars accounted for by background"\ .format(np.sum(z[:,-1]))) means = {} means['fitted_then'] = [g.mean for g in groups] means['fitted_now'] = [ torb.trace_cartesian_orbit(g.mean, g.age, single_age=True) for g in groups ] covs = {} covs['fitted_then'] = np.array([ g.generateSphericalCovMatrix() for g in groups ]) covs['fitted_now'] = np.array([ tf.transform_covmatrix(covs['fitted_then'][0], torb.trace_cartesian_orbit, means['fitted_then'][0], args=(g.age,True) ) for g in groups ]) hp.plotNewQuad(star_pars, means, covs, None, 'final', save_dir='temp_plots/', file_stem='XY-UV_ncomps_{}_'.format(ncomps), z=z, title="BIC: {}, likelihood: {}".format(bic, overall_lnlike), top_dims=(0,1), bot_dims=(3,4)) hp.plotNewQuad(star_pars, means, covs, None, 'final', save_dir='temp_plots/', file_stem='XU-YV_ncomps_{}_'.format(ncomps), z=z, title="BIC: {}, likelihood: {}".format(bic, overall_lnlike), top_dims=(0,3), bot_dims=(1,4))
try: cd_chains[label] = np.load(chaindir + 'current_day_chain.npy') except IOError: origin_chain = np.load(chaindir + 'final_chain.npy').reshape(-1,9) npars = 6 + 6 + 15 # + 1 (don't need the age in current day fit) current_day_chain = np.zeros((origin_chain.shape[0],npars)) for sample_ix, sample in enumerate(origin_chain): if sample_ix % 100 == 0: print("Done {:6} of {}".format(sample_ix, current_day_chain.shape[0])) mean = sample[:6] group = chronostar.component.Component(sample, internal=True) cd_mean = torb.trace_cartesian_orbit(group.mean, group.age) cd_cov = tf.transform_covmatrix(group.generateSphericalCovMatrix(), torb.trace_cartesian_orbit, group.mean, args=(group.age, True)) # fill in cartesian mean current_day_chain[sample_ix,0:6] = cd_mean # fill in standard deviations cd_stds = np.sqrt(cd_cov[np.diag_indices(DIM)]) current_day_chain[sample_ix,6:12] = cd_stds correl_matrix = cd_cov / cd_stds / cd_stds.reshape(DIM, 1) # fill in correlations for col_ix in range(15): current_day_chain[sample_ix,12+col_ix] = correl_matrix[ np.triu_indices(DIM, k=1)[0][col_ix], np.triu_indices(DIM, k=1)[1][col_ix] ]
print("BIC is: {}".format(bic)) print("With {:.2} stars accounted for by background"\ .format(np.sum(z[:,-1]))) means = {} means['fitted_then'] = [g.mean for g in groups] means['fitted_now'] = [ torb.trace_cartesian_orbit(g.mean, g.age, single_age=True) for g in groups ] covs = {} covs['fitted_then'] = np.array( [g.generateSphericalCovMatrix() for g in groups]) covs['fitted_now'] = np.array([ tf.transform_covmatrix(covs['fitted_then'][0], torb.trace_cartesian_orbit, means['fitted_then'][0], args=(g.age, True)) for g in groups ]) hp.plotNewQuad(star_pars, means, covs, None, 'final', save_dir='temp_plots/', file_stem='XY-UV_ncomps_{}_'.format(ncomps), z=z, title="BIC: {}, likelihood: {}".format( bic, overall_lnlike), top_dims=(0, 1), bot_dims=(3, 4))
astro_savefile) astr_arr, err_arr = ms.convertTableToArray(star_table) nstars = len(star_table) astr_covs = cv.convertAstroErrsToCovs(err_arr) xyzuvw_now = cc.convert_many_astrometry2lsrxyzuvw(astr_arr, mas=True) logging.info("Mean of retrieved stars: {}".format(np.mean(xyzuvw_now, axis=0))) if plot_it: plt.clf() plt.plot(xyzuvw_now[:, 1], xyzuvw_now[:, 2], '.') chronostar.fitplotter.plotCovEllipse( group.generateEllipticalCovMatrix()[1:3, 1:3], group.mean[1:3], with_line=True) xyzuvw_covs = np.zeros((nstars, 6, 6)) for ix in range(nstars): xyzuvw_covs[ix] = tf.transform_covmatrix(astr_covs[ix], cc.convert_astrometry2lsrxyzuvw, astr_arr[ix], dim=6) if plot_it: for ix in range(nstars): chronostar.fitplotter.plotCovEllipse(xyzuvw_covs[ix, 1:3, 1:3], xyzuvw_now[ix, 1:3], color='blue') plt.savefig("temp_plots/demo_run.png")
except IOError: origin_chain = np.load(chaindir + 'final_chain.npy').reshape(-1, 9) npars = 6 + 6 + 15 # + 1 (don't need the age in current day fit) current_day_chain = np.zeros((origin_chain.shape[0], npars)) for sample_ix, sample in enumerate(origin_chain): if sample_ix % 100 == 0: print("Done {:6} of {}".format(sample_ix, current_day_chain.shape[0])) mean = sample[:6] group = chronostar.component.Component(sample, internal=True) cd_mean = torb.trace_cartesian_orbit(group.mean, group.age) cd_cov = tf.transform_covmatrix(group.generateSphericalCovMatrix(), torb.trace_cartesian_orbit, group.mean, args=(group.age, True)) # fill in cartesian mean current_day_chain[sample_ix, 0:6] = cd_mean # fill in standard deviations cd_stds = np.sqrt(cd_cov[np.diag_indices(DIM)]) current_day_chain[sample_ix, 6:12] = cd_stds correl_matrix = cd_cov / cd_stds / cd_stds.reshape(DIM, 1) # fill in correlations for col_ix in range(15): current_day_chain[sample_ix, 12 + col_ix] = correl_matrix[ np.triu_indices(DIM, k=1)[0][col_ix], np.triu_indices(DIM, k=1)[1][col_ix]]
def polar_demo(): """ Demonstrating the use of the transform module in covariance transformations between coordinate systems """ logging.info("Beginning polar demo") plotit = True # setting up polar points r_mean = np.sqrt(10**2 + 10**2) theta_mean = 14 * np.pi / 6. r_std = 1. theta_std = np.pi / 24. C_rt = 0.0 pol_mean = np.array([r_mean, theta_mean]) pol_cov = np.array([ [r_std**2, C_rt * r_std * theta_std], [C_rt * r_std * theta_std, theta_std**2], ]) nsamples = 100000 pol_samples = np.random.multivariate_normal(pol_mean, pol_cov, nsamples) # converting to cartesian manually for comparison res = transform_ptcs(pol_samples[:, 0], pol_samples[:, 1]) cart_mean = transform_ptc(pol_mean) cart_cov = tf.transform_covmatrix(pol_cov, transform_ptc, pol_mean, dim=2) cart_samples = np.random.multivariate_normal(cart_mean, cart_cov, nsamples) estimated_mean = np.mean(res, axis=0) estimated_cov = np.cov(res.T) # plotting if plotit: nbins = 100 plt.clf() plt.hist2d(res[:, 0], res[:, 1], bins=nbins, range=((-20, 20), (-20, 20))) plt.savefig("temp_plots/jac_hist.png") plt.clf() plt.plot(res[:, 0], res[:, 1], '.') plt.ylim(-20, 20) plt.xlim(-20, 20) plt.savefig("temp_plots/jac.png") plt.clf() plt.plot(cart_samples[:, 0], cart_samples[:, 1], '.') plt.ylim(-20, 20) plt.xlim(-20, 20) plt.savefig("temp_plots/jac_cart.png") plt.clf() plt.hist2d(cart_samples[:, 0], cart_samples[:, 1], bins=nbins, range=((-20, 20), (-20, 20))) plt.savefig("temp_plots/jac_cart_hist.png") # tolerance is so high because dealing with non-linearity and assert np.allclose(estimated_mean, cart_mean, rtol=1e-1) assert np.allclose(estimated_cov, cart_cov, rtol=1e-1) logging.info("Ending polar demo")
def dataGathererEM(ngroups, iter_count, res_dir='', save_dir='', data_dir='', xyzuvw_file='', title='', file_stem='', groups_file=''): """ Provided with a results directory, tries to find all she needs, then plots Parameters ---------- ngroups: int number of groups """ covs = {} means = {} star_pars = {} if not groups_file: groups_file = "best_group_fit.npy" chain_file = "final_chain.npy" # lnprob_file = "final_lnprob.npy" # origin_file = res_dir + "origins.npy" if not xyzuvw_file: logging.info("No xyzuvw filename provided. Must be synth fit yes?") xyzuvw_file = res_dir + "../xyzuvw_now.fits" try: print("Trying to load from {}".format(xyzuvw_file)) star_pars['xyzuvw'] = fits.getdata(xyzuvw_file, 1) star_pars['xyzuvw_cov'] = fits.getdata(xyzuvw_file, 2) print("Successfully loaded from xyzuvw_file") except: import chronostar.retired.groupfitter as rgf old_star_pars = rgf.read_stars(res_dir + "../perf_tb_file.pkl") star_pars = { 'xyzuvw': old_star_pars['xyzuvw'][:, 0], 'xyzuvw_cov': old_star_pars['xyzuvw_cov'][:, 0] } try: origins = np.load(res_dir + '../origins.npy') except: origins = None z = np.load(res_dir + '../memberships.npy') fitted_then_mns = [] fitted_then_covs = [] fitted_now_mns = [] fitted_now_covs = [] origin_then_mns = [] origin_then_covs = [] all_chains = [] for group_ix in range(ngroups): gdir = res_dir + "group{}/".format(group_ix) chain = np.load(gdir + chain_file) all_chains.append(chain) best_group = np.load(gdir + groups_file).item() fitted_then_mns.append(best_group.mean) fitted_then_covs.append(best_group.generateCovMatrix()) fitted_now_mn = torb.trace_cartesian_orbit(fitted_then_mns[group_ix], best_group.age, single_age=True) fitted_now_cov =\ tf.transform_covmatrix(fitted_then_covs[group_ix], torb.trace_cartesian_orbit, fitted_then_mns[group_ix], args=(best_group.age,)) fitted_now_mns.append(fitted_now_mn) fitted_now_covs.append(fitted_now_cov) if origins is not None: origin_then_mns.append(origins[group_ix].mean) origin_then_covs.append(origins[group_ix].generateCovMatrix()) means = { 'origin_then': origin_then_mns, 'fitted_then': fitted_then_mns, 'fitted_now': fitted_now_mns, } covs = { 'origin_then': np.array(origin_then_covs), 'fitted_then': np.array(fitted_then_covs), 'fitted_now': np.array(fitted_now_covs), } plotNewQuad(star_pars, means, covs, all_chains, iter_count=iter_count, save_dir=save_dir, file_stem=file_stem, title=title, z=z) plotNewHex(star_pars, means, covs, all_chains, iter_count=iter_count, save_dir=save_dir, file_stem=file_stem, title=title, z=z)
def dataGathererEM(ngroups, iter_count, res_dir='', save_dir='', data_dir='', xyzuvw_file='', title='', file_stem='', groups_file=''): """ Provided with a results directory, tries to find all she needs, then plots Parameters ---------- ngroups: int number of groups """ covs = {} means = {} star_pars = {} if not groups_file: groups_file = "best_group_fit.npy" chain_file = "final_chain.npy" # lnprob_file = "final_lnprob.npy" # origin_file = res_dir + "origins.npy" if not xyzuvw_file: logging.info("No xyzuvw filename provided. Must be synth fit yes?") xyzuvw_file = res_dir + "../xyzuvw_now.fits" try: print("Trying to load from {}".format(xyzuvw_file)) star_pars['xyzuvw'] = fits.getdata(xyzuvw_file, 1) star_pars['xyzuvw_cov'] = fits.getdata(xyzuvw_file, 2) print("Successfully loaded from xyzuvw_file") except: import chronostar.retired.groupfitter as rgf old_star_pars = rgf.read_stars(res_dir + "../perf_tb_file.pkl") star_pars = {'xyzuvw':old_star_pars['xyzuvw'][:,0], 'xyzuvw_cov':old_star_pars['xyzuvw_cov'][:,0]} try: origins = np.load(res_dir + '../origins.npy') except: origins = None z = np.load(res_dir + '../memberships.npy') fitted_then_mns = [] fitted_then_covs = [] fitted_now_mns = [] fitted_now_covs = [] origin_then_mns = [] origin_then_covs = [] all_chains = [] for group_ix in range(ngroups): gdir = res_dir + "group{}/".format(group_ix) chain = np.load(gdir + chain_file) all_chains.append(chain) best_group = np.load(gdir + groups_file).item() fitted_then_mns.append(best_group.mean) fitted_then_covs.append(best_group.generateCovMatrix()) fitted_now_mn = torb.trace_cartesian_orbit(fitted_then_mns[group_ix], best_group.age, single_age=True) fitted_now_cov =\ tf.transform_covmatrix(fitted_then_covs[group_ix], torb.trace_cartesian_orbit, fitted_then_mns[group_ix], args=(best_group.age,)) fitted_now_mns.append(fitted_now_mn) fitted_now_covs.append(fitted_now_cov) if origins is not None: origin_then_mns.append(origins[group_ix].mean) origin_then_covs.append(origins[group_ix].generateCovMatrix()) means = { 'origin_then':origin_then_mns, 'fitted_then':fitted_then_mns, 'fitted_now':fitted_now_mns, } covs = { 'origin_then':np.array(origin_then_covs), 'fitted_then':np.array(fitted_then_covs), 'fitted_now':np.array(fitted_now_covs), } plotNewQuad(star_pars, means, covs, all_chains, iter_count=iter_count, save_dir=save_dir, file_stem=file_stem, title=title, z=z) plotNewHex(star_pars, means, covs, all_chains, iter_count=iter_count, save_dir=save_dir, file_stem=file_stem, title=title, z=z)
def calc_mns_covs(new_gps, ngroups, origins=None): all_origin_mn_then = [None] * ngroups all_origin_cov_then = [None] * ngroups all_origin_mn_now = [None] * ngroups all_origin_cov_now = [None] * ngroups all_fitted_mn_then = [None] * ngroups all_fitted_cov_then = [None] * ngroups all_fitted_mn_now = [None] * ngroups all_fitted_cov_now = [None] * ngroups for i in range(ngroups): if origins: all_origin_mn_then[i] = origins[i][:6] all_origin_cov_then[i] = utils.generate_cov( utils.internalise_pars(origins[i]) ) all_origin_mn_now[i] = tb.trace_forward(all_origin_mn_then[i], origins[i][-2]) all_origin_cov_now[i] = tf.transform_covmatrix( all_origin_cov_then[i], tb.trace_forward, all_origin_mn_then[i], dim=6, args=(origins[i][-2],) ) all_fitted_mn_then[i] = new_gps[i][:6] all_fitted_cov_then[i] = tfgf.generate_cov(new_gps[i]) all_fitted_mn_now[i] = tb.trace_forward(all_fitted_mn_then[i], new_gps[i][-1]) all_fitted_cov_now[i] = tf.transform_covmatrix( all_fitted_cov_then[i], tb.trace_forward, all_fitted_mn_then[i], dim=6, args=(new_gps[i][-1],) ) if origins: all_origin_mn_then = np.array(all_origin_mn_then ) all_origin_cov_then = np.array(all_origin_cov_then) all_origin_mn_now = np.array(all_origin_mn_now ) all_origin_cov_now = np.array(all_origin_cov_now ) all_fitted_mn_then = np.array(all_fitted_mn_then ) all_fitted_cov_then = np.array(all_fitted_cov_then) all_fitted_mn_now = np.array(all_fitted_mn_now ) all_fitted_cov_now = np.array(all_fitted_cov_now ) all_means = { 'fitted_then' : all_fitted_mn_then, 'fitted_now' : all_fitted_mn_now , } all_covs = { 'fitted_then' : all_fitted_cov_then, 'fitted_now' : all_fitted_cov_now , } if origins: all_means['origin_then'] = all_origin_mn_then all_means['origin_now'] = all_origin_mn_now all_covs['origin_then'] = all_origin_cov_then all_covs['origin_now'] = all_origin_cov_now np.save("means.npy", all_means) np.save("covs.npy", all_covs) return all_means, all_covs
def traceforward_group(mean, cov, age): """ Trace a group's Gaussian description from some time in the past to present day. :param mean: [6] array; the mean XYZUVW value :param cov: [6,6] array; Covariance matrix :param age: float; Age of "birth" of the group :return: - current_mean: [6] array; the mean phase position at current day - current_cov: [6,6] array; the cov matrix at current day """ new_mean = trace_forward(mean, age) new_cov = tf.transform_covmatrix(cov, trace_forward, mean, dim=6, args=(age, )) return new_mean, new_cov # ------ everything in this block is rubbish I think ------- #def transform_cov_mat(xyzuvw, xyzuvw_cov, tstep): # dp = 1e-3 # J = np.zeros((6,6)) # # for coord in range(6): # xyzuvw_plus = np.array(map(np.float, xyzuvw.copy())) # xyzuvw_neg = np.array(map(np.float, xyzuvw.copy())) # xyzuvw_plus[coord] += dp # xyzuvw_neg[coord] -= dp # # # filling out the coordth column of the jacobian matrix # J[:, coord] =\ # (trace_forward(xyzuvw_plus, tstep, solarmotion=None) # - trace_forward(xyzuvw_neg, tstep, solarmotion=None)) / (2*dp) # ## J[:, coord] =\ ## (trace_forward(xyzuvw_plus, tstep, solarmotion=None) ## - trace_forward(xyzuvw, tstep, solarmotion=None)) / dp # # # pdb.set_trace() # # new_cov = np.dot(J, np.dot(xyzuvw_cov, J.T)) # # return new_cov # #def trace_forward_group(xyzuvw, xyzuvw_cov, time, nsteps): # """Take a group's origin in galactic coordinates and project forward # # Parameters # ---------- # xyzuvw : [6] float array # initial mean # xyzuvw_cov : [6,6] float array # initial covariance matrix # time : float # age of group in Myr # nsteps : int # number of steps to iterate through. The transformation of the # covariance matrix is only valid over a linear regime, so need to # restrict the step size in order to maintain accuracy # # Returns # ------- # xyzuvw_current : [6] float array # current age xyzuvw mean # xyzuvw_cov_current : [6,6] float array # current age covariance matrix # """ tstep = time / nsteps old_xyzuvw = xyzuvw old_cov = xyzuvw_cov for i in range(nsteps): new_cov = transform_cov_mat(old_xyzuvw, old_cov, tstep) # do we need current xyzuvw too? new_xyzuvw = trace_forward(old_xyzuvw, tstep, solarmotion=None) old_xyzuvw = new_xyzuvw old_cov = new_cov return new_xyzuvw, new_cov
origin_file = rdir + 'origins.npy' perf_xyzuvw_file = rdir + 'perf_xyzuvw.npy' memberships_file = rdir + 'memberships.npy' final_groups_file = rdir + 'final_groups.npy' xyzuvw_file = rdir + 'xyzuvw_now.fits' origins = np.load(origin_file) perf_xyzuvw_now = np.load(perf_xyzuvw_file) z_final = np.load(memberships_file) final_groups = np.load(final_groups_file) xyzuvw_dict = gf.loadXYZUVW(xyzuvw_file) ass_cov_now = tf.transform_covmatrix(final_groups[0].generateCovMatrix(), torb.trace_cartesian_orbit, final_groups[0].mean, args=(final_groups[0].age,)) ass_mn_now = torb.trace_cartesian_orbit(final_groups[0].mean, final_groups[0].age) simple_cov_now = np.copy(ass_cov_now) simple_cov_now[3:6,:3] = 0 simple_cov_now[:3,3:6] = 0 members_mask = np.arange(50) field_mask = np.arange(50,1050) nstars = z_final.shape[0] precs = ['perf', 'half', 'gaia', 'double'] prec_val = {'perf':1e-5, 'half':0.5, 'gaia':1.0, 'double':2.0} prec = 'gaia'
z_final = np.load(final_membership_file) bp_hdul = fits.open(bp_xyzuvw_file) gaia_hdul = fits.open(gaia_astr_file) bp_xyzuvw = bp_hdul[1].data bp_xyzuvw_cov = bp_hdul[2].data bp_core_mask = np.where(z_final[:,0] > 0.75) bp_core_xyzuvw = bp_xyzuvw[bp_core_mask] bpmg_then = np.load(final_groups_file)[0] bpmg_mean_now = torb.trace_cartesian_orbit(bpmg_then.mean, bpmg_then.age, single_age=True) bpmg_cov_now = tf.transform_covmatrix( bpmg_then.generateCovMatrix(), torb.trace_cartesian_orbit, bpmg_then.mean, dim=6, args=(bpmg_then.age, True) ) ngaia_stars = gaia_xyzuvw.shape[0] try: gaia_bpmg_evals = np.load(gaia_bpmg_evals_file) except IOError: print("Evaluating gaia stars at BPMG current MVGauss distribution") gaia_bpmg_evals = np.zeros(ngaia_stars) bpmg_invcov_now = np.linalg.inv(bpmg_cov_now) for i, gaia_star in enumerate(gaia_xyzuvw): if (i % 100000 == 0): print("{:10} of {:10}... {:6.1f}%".format(i, ngaia_stars, i / ngaia_stars*100))
'c_ZU', 'c_ZV', 'c_ZW', 'c_UV', 'c_UW', 'c_VW' ] # insert empty columns for col_name in cart_col_names: gt[col_name] = empty_col for row_ix, gt_row in enumerate(gt): dim = 6 if row_ix % 10 == 0: print("{:02.2f}% done".format(row_ix / float(nrows) * 100.)) astr_mean, astr_cov = gc.convertRecToArray(gt_row) xyzuvw_mean = coord.convert_astrometry2lsrxyzuvw(astr_mean) xyzuvw_cov = tf.transform_covmatrix( astr_cov, coord.convert_astrometry2lsrxyzuvw, astr_mean, dim=dim, ) # fill in cartesian mean for col_ix, col_name in enumerate(cart_col_names[:6]): gt_row[col_name] = xyzuvw_mean[col_ix] # fill in standard deviations xyzuvw_stds = np.sqrt(xyzuvw_cov[np.diag_indices(dim)]) for col_ix, col_name in enumerate(cart_col_names[6:12]): gt_row[col_name] = xyzuvw_stds[col_ix] correl_matrix = xyzuvw_cov / xyzuvw_stds / xyzuvw_stds.reshape(6, 1) # fill in correlations for col_ix, col_name in enumerate(cart_col_names[12:]): gt_row[col_name] = correl_matrix[
def traceforward_group(mean, cov, age): """ Trace a group's Gaussian description from some time in the past to present day. :param mean: [6] array; the mean XYZUVW value :param cov: [6,6] array; Covariance matrix :param age: float; Age of "birth" of the group :return: - current_mean: [6] array; the mean phase position at current day - current_cov: [6,6] array; the cov matrix at current day """ new_mean = trace_forward(mean, age) new_cov = tf.transform_covmatrix( cov, trace_forward, mean, dim=6, args=(age,) ) return new_mean, new_cov # ------ everything in this block is rubbish I think ------- #def transform_cov_mat(xyzuvw, xyzuvw_cov, tstep): # dp = 1e-3 # J = np.zeros((6,6)) # # for coord in range(6): # xyzuvw_plus = np.array(map(np.float, xyzuvw.copy())) # xyzuvw_neg = np.array(map(np.float, xyzuvw.copy())) # xyzuvw_plus[coord] += dp # xyzuvw_neg[coord] -= dp # # # filling out the coordth column of the jacobian matrix # J[:, coord] =\ # (trace_forward(xyzuvw_plus, tstep, solarmotion=None) # - trace_forward(xyzuvw_neg, tstep, solarmotion=None)) / (2*dp) # ## J[:, coord] =\ ## (trace_forward(xyzuvw_plus, tstep, solarmotion=None) ## - trace_forward(xyzuvw, tstep, solarmotion=None)) / dp # # # pdb.set_trace() # # new_cov = np.dot(J, np.dot(xyzuvw_cov, J.T)) # # return new_cov # #def trace_forward_group(xyzuvw, xyzuvw_cov, time, nsteps): # """Take a group's origin in galactic coordinates and project forward # # Parameters # ---------- # xyzuvw : [6] float array # initial mean # xyzuvw_cov : [6,6] float array # initial covariance matrix # time : float # age of group in Myr # nsteps : int # number of steps to iterate through. The transformation of the # covariance matrix is only valid over a linear regime, so need to # restrict the step size in order to maintain accuracy # # Returns # ------- # xyzuvw_current : [6] float array # current age xyzuvw mean # xyzuvw_cov_current : [6,6] float array # current age covariance matrix # """ tstep = time / nsteps old_xyzuvw = xyzuvw old_cov = xyzuvw_cov for i in range(nsteps): new_cov = transform_cov_mat(old_xyzuvw, old_cov, tstep) # do we need current xyzuvw too? new_xyzuvw = trace_forward(old_xyzuvw, tstep, solarmotion=None) old_xyzuvw = new_xyzuvw old_cov = new_cov return new_xyzuvw, new_cov
def polar_demo(): """ Demonstrating the use of the transform module in covariance transformations between coordinate systems """ logging.info("Beginning polar demo") plotit = True # setting up polar points r_mean = np.sqrt(10**2+ 10**2) theta_mean = 14*np.pi / 6. r_std = 1. theta_std = np.pi / 24. C_rt = 0.0 pol_mean = np.array([r_mean, theta_mean]) pol_cov = np.array([ [r_std**2, C_rt * r_std * theta_std], [C_rt * r_std * theta_std, theta_std**2], ]) nsamples = 100000 pol_samples = np.random.multivariate_normal(pol_mean, pol_cov, nsamples) # converting to cartesian manually for comparison res = transform_ptcs(pol_samples[:, 0], pol_samples[:, 1]) cart_mean = transform_ptc(pol_mean) cart_cov = tf.transform_covmatrix(pol_cov, transform_ptc, pol_mean, dim=2) cart_samples = np.random.multivariate_normal(cart_mean, cart_cov, nsamples) estimated_mean = np.mean(res, axis=0) estimated_cov = np.cov(res.T) # plotting if plotit: nbins = 100 plt.clf() plt.hist2d(res[:,0],res[:,1],bins=nbins,range=((-20,20),(-20,20))) plt.savefig("temp_plots/jac_hist.png") plt.clf() plt.plot(res[:,0], res[:,1], '.') plt.ylim(-20,20) plt.xlim(-20,20) plt.savefig("temp_plots/jac.png") plt.clf() plt.plot(cart_samples[:,0], cart_samples[:,1], '.') plt.ylim(-20,20) plt.xlim(-20,20) plt.savefig("temp_plots/jac_cart.png") plt.clf() plt.hist2d(cart_samples[:,0],cart_samples[:,1],bins=nbins,range=((-20,20),(-20,20))) plt.savefig("temp_plots/jac_cart_hist.png") # tolerance is so high because dealing with non-linearity and assert np.allclose(estimated_mean, cart_mean, rtol=1e-1) assert np.allclose(estimated_cov, cart_cov, rtol=1e-1) logging.info("Ending polar demo")
xyzuvw_now_true = to.trace_many_cartesian_orbit(xyzuvw_init, np.array([0., group.age]))[:, 1] #assert np.allclose(np.mean(xyzuvw_now, axis=0), group.mean, rtol=1e-1) logging.info("Mean of initial stars: {}".format(np.mean(xyzuvw_init, axis=0))) logging.info("Mean of final stars: {}".format(np.mean(xyzuvw_now_true, axis=0))) star_table = chronostar.synthdata.measureXYZUVW(xyzuvw_now_true, 20.0, astro_savefile) astr_arr, err_arr = ms.convertTableToArray(star_table) nstars = len(star_table) astr_covs = cv.convertAstroErrsToCovs(err_arr) xyzuvw_now = cc.convert_many_astrometry2lsrxyzuvw(astr_arr, mas=True) logging.info("Mean of retrieved stars: {}".format(np.mean(xyzuvw_now, axis=0))) if plot_it: plt.clf() plt.plot(xyzuvw_now[:,1], xyzuvw_now[:,2], '.') chronostar.fitplotter.plotCovEllipse(group.generateEllipticalCovMatrix()[1:3, 1:3], group.mean[1:3], with_line=True) xyzuvw_covs = np.zeros((nstars,6,6)) for ix in range(nstars): xyzuvw_covs[ix] = tf.transform_covmatrix( astr_covs[ix], cc.convert_astrometry2lsrxyzuvw, astr_arr[ix], dim=6 ) if plot_it: for ix in range(nstars): chronostar.fitplotter.plotCovEllipse(xyzuvw_covs[ix, 1:3, 1:3], xyzuvw_now[ix, 1:3], color='blue') plt.savefig("temp_plots/demo_run.png")
'c_UV', 'c_UW', 'c_VW'] # insert empty columns for col_name in cart_col_names: gt[col_name] = empty_col for row_ix, gt_row in enumerate(gt): dim = 6 if row_ix%10 == 0: print("{:02.2f}% done".format(row_ix / float(nrows) * 100.)) astr_mean, astr_cov = gc.convertRecToArray(gt_row) xyzuvw_mean = coord.convert_astrometry2lsrxyzuvw(astr_mean) xyzuvw_cov = tf.transform_covmatrix( astr_cov, coord.convert_astrometry2lsrxyzuvw, astr_mean, dim=dim, ) # fill in cartesian mean for col_ix, col_name in enumerate(cart_col_names[:6]): gt_row[col_name] = xyzuvw_mean[col_ix] # fill in standard deviations xyzuvw_stds = np.sqrt(xyzuvw_cov[np.diag_indices(dim)]) for col_ix, col_name in enumerate(cart_col_names[6:12]): gt_row[col_name] = xyzuvw_stds[col_ix] correl_matrix = xyzuvw_cov / xyzuvw_stds / xyzuvw_stds.reshape(6,1) # fill in correlations for col_ix, col_name in enumerate(cart_col_names[12:]):
sdir = 'temp_data/' origin_file = rdir + 'origins.npy' perf_xyzuvw_file = rdir + 'perf_xyzuvw.npy' memberships_file = rdir + 'memberships.npy' final_groups_file = rdir + 'final_groups.npy' xyzuvw_file = rdir + 'xyzuvw_now.fits' origins = np.load(origin_file) perf_xyzuvw_now = np.load(perf_xyzuvw_file) z_final = np.load(memberships_file) final_groups = np.load(final_groups_file) xyzuvw_dict = gf.loadXYZUVW(xyzuvw_file) ass_cov_now = tf.transform_covmatrix(final_groups[0].generateCovMatrix(), torb.trace_cartesian_orbit, final_groups[0].mean, args=(final_groups[0].age, )) ass_mn_now = torb.trace_cartesian_orbit(final_groups[0].mean, final_groups[0].age) simple_cov_now = np.copy(ass_cov_now) simple_cov_now[3:6, :3] = 0 simple_cov_now[:3, 3:6] = 0 members_mask = np.arange(50) field_mask = np.arange(50, 1050) nstars = z_final.shape[0] precs = ['perf', 'half', 'gaia', 'double'] prec_val = {'perf': 1e-5, 'half': 0.5, 'gaia': 1.0, 'double': 2.0} prec = 'gaia'