def test_traceforwardThenBack(): """Check that tracing a point forward then back for the same time step returns initial position """ return ABS_TOLERANCE = 1e-3 xyzuvws = np.array([ [0., 0., 25., 0., 0., 0.], # [10.,0.,-50.,0.,0.,0.], # [0.,0.,0.,10.,25.,30.,], ]) age = 100. times = np.linspace(0, 100, 1001) for xyzuvw_start in xyzuvws: galpy_start = None xyzuvw_end = torb.trace_cartesian_orbit( xyzuvw_start, times=age, single_age=True, ) xyzuvw_start_again = torb.trace_cartesian_orbit( xyzuvw_end, times=-age, single_age=True, ) assert np.allclose(xyzuvw_start, xyzuvw_start_again, atol=ABS_TOLERANCE)
def test_traceforwardThenBack(): """Check that tracing a point forward then back for the same time step returns initial position """ return ABS_TOLERANCE = 1e-3 xyzuvws = np.array([ [0.,0.,25.,0.,0.,0.], # [10.,0.,-50.,0.,0.,0.], # [0.,0.,0.,10.,25.,30.,], ]) age = 100. times = np.linspace(0,100,1001) for xyzuvw_start in xyzuvws: galpy_start = None xyzuvw_end = torb.trace_cartesian_orbit(xyzuvw_start, times=age, single_age=True, ) xyzuvw_start_again = torb.trace_cartesian_orbit(xyzuvw_end, times=-age, single_age=True, ) assert np.allclose(xyzuvw_start, xyzuvw_start_again, atol=ABS_TOLERANCE)
def test_singleTime(): """Test usage where we only provide the desired age, and not an array Good demo of how to traceback 2 stars forward through time, either with an array of time steps, or a single age """ xyzuvw_1 = [0., 0., 25., 0., 0., 0.] xyzuvw_2 = [0., 0., 0., 0., -10., 0.] xyzuvws = np.vstack((xyzuvw_1, xyzuvw_2)) age = 10. times = np.linspace(0., age, 2) # get position for each time in times xyzuvws_both = torb.trace_many_cartesian_orbit(xyzuvws, times, single_age=False) # get position for *age* only xyzuvws_now = torb.trace_many_cartesian_orbit(xyzuvws, age, single_age=True) assert np.allclose(xyzuvws_both[:, 1], xyzuvws_now) xyzuvw_both = torb.trace_cartesian_orbit(xyzuvws[0], times, single_age=False) xyzuvw_now = torb.trace_cartesian_orbit(xyzuvws[0], age, single_age=True) assert np.allclose(xyzuvw_both[1], xyzuvw_now)
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 project_stars(self): """Project stars from xyzuvw then to xyzuvw now based on their age""" for star in self.table: mean_then = self.extract_data_as_array( table=star, colnames=[dim + '0' for dim in self.cart_labels], ) xyzuvw_now = traceorbit.trace_cartesian_orbit(mean_then, times=star['age']) for ix, dim in enumerate(self.cart_labels): star[dim + '_now'] = xyzuvw_now[ix]
def project_stars(self): """Project stars from xyzuvw then to xyzuvw now based on their age""" for star in self.table: mean_then = self.extract_data_as_array( table=star, colnames=[dim+'0' for dim in self.cart_labels], ) xyzuvw_now = traceorbit.trace_cartesian_orbit(mean_then, times=star['age']) for ix, dim in enumerate(self.cart_labels): star[dim+'_now'] = xyzuvw_now[ix]
def test_interval_tracing_orig(): np.random.seed(0) start = np.random.rand(6) * 20 - 10 time_steps = [3., -10., -3., 10] current_pos = start for time_step in time_steps: current_pos = torb.trace_cartesian_orbit( current_pos, times=time_step, single_age=True, ) assert np.allclose(start, current_pos, 1e-3)
def test_trace_epicyclic_orbit(): """ Take a point, trace it forward with epicyclic approx. and galpy and compare results. TODO: Marusa left incomplete test """ if False: # xyzuvw_start = np.ones(6) * 1000.0 xyzuvw_start = np.array([10., 10., 10., 3., -2., 1.]) # times = np.array([100]) # Myr # TC: trace_cartesian_orbit expects a single float for times times = 100. galpy_end = torb.trace_cartesian_orbit(xyzuvw_start, times=times) epi_end = torb.trace_epicyclic_orbit(xyzuvw_start, times=times)
def test_interval_tracing_orig(): np.random.seed(0) start = np.random.rand(6) * 20 - 10 time_steps = [3.,-10.,-3.,10] current_pos = start for time_step in time_steps: current_pos = torb.trace_cartesian_orbit( current_pos, times=time_step, single_age=True, ) assert np.allclose(start, current_pos, 1e-3)
def test_LSR(): """ Check that LSR remains constant in our frame of reference. Since our frame of reference is **centred** on the LSR, then the LSR should remain at the origin. """ xyzuvw_lsr = [0.,0.,0.,0.,0.,0.] times = np.linspace(0,100,101) for method, atol in method_atols.items(): xyzuvws = torb.trace_cartesian_orbit(xyzuvw_lsr, times, single_age=False, method=method) assert np.allclose(xyzuvws[0,:5],xyzuvws[-1,:5], atol=atol), 'Method {}'.format(method)
def test_lcc_like(): """ Takes about 40 mins """ mean_now = np.array([50., -100., 25., 1.1, -7.76, 2.25]) age = 10. mean = trace_cartesian_orbit(mean_now, times=-age) dx = 5. dv = 2. covmatrix = np.identity(6) covmatrix[:3,:3] *= dx**2 covmatrix[3:,3:] *= dv**2 true_comp = SphereComponent(attributes={ 'mean':mean, 'covmatrix':covmatrix, 'age':age, }) nstars = 100 tiny_measurement_error = 1e-10 short_burnin_step = 200 best_comp, chain, lnprob = run_fit_helper( true_comp=true_comp, starcounts=nstars, measurement_error=tiny_measurement_error, burnin_step=short_burnin_step, run_name='lcc_like', ) np.save('temp_data/{}_compfitter_lcc_like_'\ 'true_and_best_comp.npy'.format(PY_VERS), [true_comp, best_comp],) assert np.allclose(true_comp.get_mean(), best_comp.get_mean(), atol=3.0) assert np.allclose(true_comp.get_age(), best_comp.get_age(), atol=1.0) assert np.allclose(true_comp.get_covmatrix(), best_comp.get_covmatrix(), atol=5.0)
def test_rotatedLSR(): """ Check that LSRs with different azimuthal positions also remain constant """ # method_atols = {'odeint':1e-11, 'dopr54_c':1e-6} for method, atol in method_atols.items(): rot_lsr_gp_coords = np.array([1., 0., 1., 0., 0., np.pi]) xyzuvw_rot_lsr = torb.convert_galpycoords2cart(rot_lsr_gp_coords) ### xyzuvw_rot_lsr = [16000., 0, 0, 0, 0, 0,] times = np.linspace(0,100,101) xyzuvws = torb.trace_cartesian_orbit(xyzuvw_rot_lsr, times, single_age=False, method=method) # On a circular orbit, same radius as LSR, so shouldn't vary at all # assert np.allclose(xyzuvws[0,:5],xyzuvws[-1,:5]) assert np.allclose(xyzuvws[0],xyzuvws[-1], atol=atol), 'Method {}'.format(method) # Should be initialised on opposite side of galaxy (X = 16kpc) assert np.allclose(16000., xyzuvws[0,0])
def test_lcc_like(): """ Takes about 40 mins """ mean_now = np.array([50., -100., 25., 1.1, -7.76, 2.25]) age = 10. mean = trace_cartesian_orbit(mean_now, times=-age) dx = 5. dv = 2. covmatrix = np.identity(6) covmatrix[:3, :3] *= dx**2 covmatrix[3:, 3:] *= dv**2 true_comp = SphereComponent(attributes={ 'mean': mean, 'covmatrix': covmatrix, 'age': age, }) nstars = 100 tiny_measurement_error = 1e-10 short_burnin_step = 200 best_comp, chain, lnprob = run_fit_helper( true_comp=true_comp, starcounts=nstars, measurement_error=tiny_measurement_error, burnin_step=short_burnin_step, run_name='lcc_like', ) np.save('temp_data/{}_compfitter_lcc_like_'\ 'true_and_best_comp.npy'.format(PY_VERS), [true_comp, best_comp],) assert np.allclose(true_comp.get_mean(), best_comp.get_mean(), atol=3.0) assert np.allclose(true_comp.get_age(), best_comp.get_age(), atol=1.0) assert np.allclose(true_comp.get_covmatrix(), best_comp.get_covmatrix(), atol=5.0)
def test_traceback_and_forward(): """The test that shows things are broken""" time = 10. # Myr times = np.array([0., 10.]) init_pos_chron = np.array([10., 0., 30., 0., 0., 0.]) # init_pos_chron = np.zeros(6) init_pos_chron = np.array([4000, 8000. * np.sqrt(3) / 2, 0, np.sin(np.pi / 3) * 220., -np.cos(np.pi / 3) * 220., 0]) final_pos_chron = torb.trace_cartesian_orbit( init_pos_chron, times=times, single_age=False)[-1] final_pos_chron = torb.traceforward_from_now( init_pos_chron, time=time, ) assert np.allclose( init_pos_chron, torb.traceback_to_now( final_pos_chron, time, ), atol=1e-5 )
def test_evolve_chronspace(): """ Compare orbits evolved with both galpy and epicyclic, comparing in chron space """ time = 50. #Myr btime = torb.convert_myr2bovytime(time) chron_start = np.array([0., 0., 10., -2., 0., 0.]) galpy_start = torb.convert_cart2galpycoords(chron_start) epi_start = eg.convert_galpy2epi(galpy_start) # Just make sure the starting point can be transformed back and forth assert np.allclose(galpy_start, eg.convert_epi2galpy(epi_start)) assert np.allclose( chron_start, torb.convert_galpycoords2cart(eg.convert_epi2galpy(epi_start))) epi_end = eg.evolve_epi(epi_start, time)[0] galpy_end = torb.trace_galpy_orbit(galpy_start, times=time, single_age=True) chron_end = torb.trace_cartesian_orbit(chron_start, times=time, single_age=True) # Chronostar orbit end point, in galpy units chron_end_gu = torb.convert_cart2galpycoords(chron_end, ts=time) # This should be exact, because algorithmically it's the same thing. assert np.allclose(galpy_end, chron_end_gu) # Epicyclic orbit end point, in chronostar units epi_end_chron = torb.convert_galpycoords2cart( eg.convert_epi2galpy(epi_end), ts=btime) # assert position accurate within 15 pc # NOTE this is quite large... offset mainly in X. Maybe something to # investigate. assert np.allclose(chron_end[:3], epi_end_chron[:3], atol=15.) # assert velocity accurate within 0.5 km/s assert np.allclose(chron_end[3:], epi_end_chron[3:], atol=.5)
def test_traceback_and_forward(): """The test that shows things are broken""" time = 10. #Myr times = np.array([0., 10.]) init_pos_chron = np.array([10.,0.,30.,0.,0.,0.]) # init_pos_chron = np.zeros(6) init_pos_chron = np.array([4000, 8000.*np.sqrt(3)/2, 0, np.sin(np.pi/3)*220., -np.cos(np.pi/3)*220., 0]) final_pos_chron = torb.trace_cartesian_orbit( init_pos_chron, times=times, single_age=False)[-1] final_pos_chron = torb.traceforward_from_now( init_pos_chron, time=time, ) assert np.allclose( init_pos_chron, torb.traceback_to_now( final_pos_chron, time, ), atol=1e-5 )
bg_hists = np.load(bg_hists_file) bg_ln_ols = em.backgroundLogOverlaps(star_pars['xyzuvw'], bg_hists) overall_lnlike, z = em.get_overall_lnlikelihood(star_pars, groups, bg_ln_ols, return_memb_probs=True) print("overall_lnlike with {} comps is: {:.5}".format(ncomps, overall_lnlike)) bic = calcBIC(star_pars, ncomps, overall_lnlike) 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/',
for chaindir, label in zip(chaindirs, labels): 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],
overall_lnlike, z = em.get_overall_lnlikelihood(star_pars, groups, bg_ln_ols, return_memb_probs=True) print("overall_lnlike with {} comps is: {:.5}".format( ncomps, overall_lnlike)) bic = calcBIC(star_pars, ncomps, overall_lnlike) 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,
max_time = 1000 orbit_times = np.linspace(0,max_time,100) niters = 100 print('Integrating up to {} Myr'.format(max_time)) print('Iterating {} times'.format(niters)) print('----------- Check accuracy ----------') xyzuvw_lsr = [0.,0.,0.,0.,0.,0.] traceback_age = 100. #Myr for method in integ_methods: print('_____ Using {} _____'.format(method)) xyzuvw_final = trace_cartesian_orbit(xyzuvw_lsr, traceback_age, single_age=True, method=method) diff = xyzuvw_final - xyzuvw_lsr pos_error = np.sqrt(np.sum(np.square(diff[:3]))) vel_error = np.sqrt(np.sum(np.square(diff[3:]))) print('Position error: {:.3} pc'.format(pos_error)) print('Velocity error: {:.3} km/s'.format(vel_error)) print('') print('----------- Check timings ----------') for method in integ_methods: print('_____ Using {} _____'.format(method)) duration_times = []
logging.info("Synth data exists! .....") print("Synth data exists") raise UserWarning except IOError: all_xyzuvw_init = np.zeros((0, 6)) all_xyzuvw_now_perf = np.zeros((0, 6)) origins = [] for i in range(ngroups): logging.info(" generating from group {}".format(i)) # MANUALLY SEPARATE CURRENT DAY DISTROS IN DIMENSION X mean_now_w_offset = mean_now.copy() # mean_now_w_offset[0] += i * 50 mean_now_w_offset += offsets[i] mean_then = torb.trace_cartesian_orbit(mean_now_w_offset, -extra_pars[i, -2], single_age=True) group_pars = np.hstack((mean_then, extra_pars[i])) xyzuvw_init, origin = syn.synthesiseXYZUVW(group_pars, form='sphere', return_group=True, internal=False) origins.append(origin) all_xyzuvw_init = np.vstack((all_xyzuvw_init, xyzuvw_init)) xyzuvw_now_perf = torb.trace_many_cartesian_orbit(xyzuvw_init, times=origin.age, single_age=True) all_xyzuvw_now_perf = np.vstack((all_xyzuvw_now_perf, xyzuvw_now_perf)) # insert 'background stars' with density `BG_DENS` [pc km/s]^-3 ubound = np.max(all_xyzuvw_now_perf, axis=0)
gaia_cov_file = andir + "gaia_cov.npy" bpmg_memb_probs_file = andir + "bpmg_memb_probs.npy" gaia_xyzuvw = np.load(gaia_xyzuvw_mean_file) 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):
def plot_orbit(self, dim1, dim2, ax, ntimes=50, with_arrow=False, annotate=False, color=None, alpha=0.3, **kwargs): """ For traceback use negative age Parameters ---------- pos_now: [6] array, known position of object dim1: integer, x-axis dimension dim2: integer, y-axis dimension ax: axes object, axes on which to plot line end_age: non-zero number, time to orbit till. Negative value --> traceback Positive value --> trace forward ntimes: integer {50], number of timesteps to calculate group_ix: index of group being plotted (for coloring reasons) with_arrow: (bool) {False}, whether to include arrows along orbit annotate: (bool) {False}, whether to include text """ # alpha=0.3 if color is None: color = 'black' # if group_ix is None: # color = COLORS[0] # else: # color = COLORS[group_ix] # orb_alpha = 0.1 comp_orb = trace_cartesian_orbit(self.get_mean(), times=np.linspace( 0, self.get_age(), ntimes), single_age=False) line_obj = ax.plot(comp_orb[:, dim1], comp_orb[:, dim2], ls='-', alpha=alpha, color=color, zorder=1, **kwargs) indices = [int(ntimes / 3), int(2 * ntimes / 3)] if with_arrow: # make sure arrow is always pointing forwards through time direction = 'right' if self.get_age() > 0 else 'left' self.add_arrow(line_obj[0], indices=indices, direction=direction, color=color, alpha=alpha, zorder=1) if annotate: ax.annotate("Orbital trajectory", (comp_orb[int(ntimes / 2), dim1], comp_orb[int(ntimes / 2), dim2]), color=color)
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)
import matplotlib as mpl mpl.use('Agg') import matplotlib.pyplot as plt from chronostar.component import SphereComponent from chronostar.traceorbit import trace_cartesian_orbit import numpy as np mean_now = np.array([0., 0., 30., 5., 5., 5.]) init_dx = 5. init_dv = 1. age = 100. mean_then = trace_cartesian_orbit(mean_now, times=-age) pars1 = np.hstack((mean_then, init_dx, init_dv, age)) comp1 = SphereComponent(pars1) print(comp1.get_pars()) labels = 'XYZUVW' units = 3 * ['pc'] + 3 * ['km/s'] for dim1, dim2 in [(0, 3), (1, 4), (2, 5)]: plt.clf() comp1.plot(dim1=dim1, dim2=dim2, comp_now=True, comp_then=True, comp_orbit=True) plt.xlabel('{} [{}]'.format(labels[dim1], units[dim1]))
pool.wait() sys.exit(0) print("Only one thread is master") print("Master should be working in the directory:\n{}".format(rdir)) logging.info("Performing fit with") logging.info("{} burnin steps".format(BURNIN_STEPS)) logging.info("{} sampling steps".format(SAMPLING_STEPS)) logging.info("{} tolerance".format(C_TOL)) logging.info("In the directory: {}".format(rdir)) ## Destination: (inspired by LCC) mean_now_lsr = np.array([50., -100., 25., 1.1, -7.76, 2.25]) # Calculate appropriate starting point mean_then = torb.trace_cartesian_orbit(mean_now_lsr, -age) # gather inputs # group_pars = np.hstack((mean_then, dX, dV, age, nstars)) # Setting up perfect current xyzuvw values try: xyzuvw_now_perf = np.load(rdir+xyzuvw_perf_file) # origin = np.load(group_savefile) origin = dt.loadGroups(rdir+group_savefile)[0] logging.basicConfig( level=logging.INFO, filemode='a', filename='my_investigator_demo.log', ) logging.info("appending to previous attempt") except IOError: logging.basicConfig(
y_pos_now = bugged_comp.get_mean_now()[1] print('Mean then: {}'.format(bugged_comp.get_mean())) print('Internally calculated mean_now: {}'.format(bugged_comp.get_mean_now())) mean_then = bugged_comp.get_mean() age = bugged_comp.get_age() y_pos_linear = mean_then[1] + mean_then[4] * age print('Linear Y motion: {}'.format(y_pos_linear)) print('Difference of {}'.format(y_pos_now - y_pos_linear)) ts = np.linspace(0, bugged_comp.get_age(), 50) for method in [ 'odeint', 'symplec4_c', 'rk4_c', 'dopr54_c', 'rk6_c', ]: manual_single_age_mean_now = trace_cartesian_orbit( bugged_comp.get_mean(), times=bugged_comp.get_age(), method=method) manual_multi_age_mean_now = trace_cartesian_orbit(bugged_comp.get_mean(), times=ts, method=method, single_age=False) print(method) print('Manual Y now (single age): {}'.format( manual_single_age_mean_now[1])) print('Manual Y now (multi age): {}'.format(manual_multi_age_mean_now[-1, 1]))
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' # astro_table = ms.measureXYZUVW(perf_xyzuvw_now, prec_val[prec]) # xyzuvw_dict = cv.convertMeasurementsToCartesian(astro_table)
dt.loadXYZUVW(xyzuvw_conv_savefile) logging.info("Synth data exists! .....") print("Synth data exists") raise UserWarning except IOError: all_xyzuvw_init = np.zeros((0,6)) all_xyzuvw_now_perf = np.zeros((0,6)) origins = [] for i in range(ngroups): logging.info(" generating from group {}".format(i)) # MANUALLY SEPARATE CURRENT DAY DISTROS IN DIMENSION X mean_now_w_offset = mean_now.copy() # mean_now_w_offset[0] += i * 50 mean_now_w_offset += offsets[i] mean_then = torb.trace_cartesian_orbit(mean_now_w_offset, -extra_pars[i, -2], single_age=True) group_pars = np.hstack((mean_then, extra_pars[i])) xyzuvw_init, origin = syn.synthesiseXYZUVW(group_pars, form='sphere', return_group=True, internal=False) origins.append(origin) all_xyzuvw_init = np.vstack((all_xyzuvw_init, xyzuvw_init)) xyzuvw_now_perf = torb.trace_many_cartesian_orbit(xyzuvw_init, times=origin.age, single_age=True) all_xyzuvw_now_perf = np.vstack((all_xyzuvw_now_perf, xyzuvw_now_perf)) # insert 'background stars' with density `BG_DENS` [pc km/s]^-3 ubound = np.max(all_xyzuvw_now_perf, axis=0) lbound = np.min(all_xyzuvw_now_perf, axis=0) margin = 0.5 * (ubound - lbound)
def test_different_potential(): miya_pot = MiyamotoNagaiPotential(a=0.5, b=0.0375, amp=1., normalize=1.) miya_trace_cartesian_orbit = torb.trace_orbit_builder(miya_pot) ABS_TOLERANCE = 1e-3 xyzuvws = np.array([ [0., 0., 25., 0., 0., 0.], [10., 0., -50., 0., 0., 0.], [10., 0., -50., 0., 0., -5.], [0., 0., 0., 10., 25., 30.,], ]) age = 100. times = np.linspace(0, 100, 1001) for xyzuvw_start in xyzuvws: xyzuvw_end = miya_trace_cartesian_orbit(xyzuvw_start, times=age, single_age=True, ) xyzuvw_start_again = miya_trace_cartesian_orbit(xyzuvw_end, times=-age, single_age=True, ) assert np.allclose(xyzuvw_start, xyzuvw_start_again, atol=ABS_TOLERANCE) # Confirm that tracing forward with one potential but back with another # gives different starting position xyzuvw_start_but_with_diff_pot = torb.trace_cartesian_orbit(xyzuvw_end, times=-age,) assert not np.allclose(xyzuvw_start, xyzuvw_start_but_with_diff_pot) def test_traceforwardThenBack(): """Check that tracing a point forward then back for the same time step returns initial position """ return ABS_TOLERANCE = 1e-3 xyzuvws = np.array([ [0., 0., 25., 0., 0., 0.], # [10.,0.,-50.,0.,0.,0.], # [0.,0.,0.,10.,25.,30.,], ]) age = 100. times = np.linspace(0, 100, 1001) for xyzuvw_start in xyzuvws: galpy_start = None xyzuvw_end = torb.trace_cartesian_orbit(xyzuvw_start, times=age, single_age=True, ) xyzuvw_start_again = torb.trace_cartesian_orbit(xyzuvw_end, times=-age, single_age=True, ) assert np.allclose(xyzuvw_start, xyzuvw_start_again, atol=ABS_TOLERANCE) def test_galpy_stationary_conversions(): """Check if gaply conversions behave as expected where everything is at time 0""" # Test LSR lsr_chron = np.zeros(6) lsr_galpy = np.array([1., 0, 1, 0, 0, 0]) assert np.allclose(lsr_chron, torb.convert_galpycoords2cart(lsr_galpy, ts=0.)) assert np.allclose(lsr_galpy, torb.convert_cart2galpycoords(lsr_chron, ts=0.)) # Test galactic centre gc_chron = np.array([8000., 0, 0, 0, -220., 0, ]) gc_galpy = np.ones(6) * 1e-15 assert np.allclose(gc_chron, torb.convert_galpycoords2cart(gc_galpy, ts=0.)) assert np.allclose(gc_galpy, torb.convert_cart2galpycoords(gc_chron, ts=0.)) # Test simple, off origin point off_chron = np.array([4000, 8000. * np.sqrt(3) / 2, 0, np.sin(np.pi / 3) * 220., -np.cos(np.pi / 3) * 220., 0]) off_galpy = np.array([1., 0, 1, 0, 0, np.pi / 3.]) assert np.allclose(off_galpy, torb.convert_cart2galpycoords(off_chron, ts=0.)) assert np.allclose(off_chron, torb.convert_galpycoords2cart(off_galpy, ts=0.)) # Test random positions SPREAD = 100000 NSAMPLES = int(1e6) many_pos_chron = (np.random.rand(NSAMPLES, 6) - 0.5) * SPREAD # uniform between -10 and 10 many_pos_galpy = torb.convert_cart2galpycoords(many_pos_chron, ts=0.) assert np.allclose(many_pos_chron, torb.convert_galpycoords2cart(many_pos_galpy, ts=0.), atol=1e-2) def test_galpy_moving_conversions(): """Check if gaply conversions behave as expected where time is allowed to vary.""" lsr_chron = np.zeros(6) lsr_galpy = np.array([1., 0, 1, 0, 0, 0]) # Incorporate positive time into lsr position checks NSTEPS = 10 galpy_times = np.linspace(0., 2 * np.pi, NSTEPS) lsrs_chron = np.repeat(lsr_chron, NSTEPS).reshape(6, -1).T lsrs_galpy = np.repeat(lsr_galpy, NSTEPS).reshape(6, -1).T lsrs_galpy[:, -1] = galpy_times chron_times = torb.convert_bovytime2myr(galpy_times) assert np.allclose( lsrs_chron, torb.convert_galpycoords2cart(lsrs_galpy, ts=galpy_times)) assert np.allclose( lsrs_galpy, torb.convert_cart2galpycoords(lsrs_chron, ts=chron_times) ) # Incorporate negative time into lsr position checks galpy_times = np.linspace(0., -2 * np.pi, NSTEPS) lsrs_chron = np.repeat(lsr_chron, NSTEPS).reshape(6, -1).T lsrs_galpy = np.repeat(lsr_galpy, NSTEPS).reshape(6, -1).T lsrs_galpy[:, -1] = galpy_times chron_times = torb.convert_bovytime2myr(galpy_times) assert np.allclose( lsrs_chron, torb.convert_galpycoords2cart(lsrs_galpy, ts=galpy_times)) assert np.allclose( lsrs_galpy, torb.convert_cart2galpycoords(lsrs_chron, ts=chron_times) ) # Test random positions with random times SPREAD = int(1e4) # pc NSAMPLES = 100 many_pos_chron = (np.random.rand(NSAMPLES, 6) - 0.5) * SPREAD # uniform between -10 and 10 many_chron_times = np.random.rand(NSAMPLES) * 100 # Myr many_pos_galpy = torb.convert_cart2galpycoords( many_pos_chron, ts=many_chron_times ) many_galpy_times = torb.convert_myr2bovytime(many_chron_times) for i in range(NSAMPLES): assert np.allclose(many_pos_chron[i], torb.convert_galpycoords2cart( many_pos_galpy[i], ts=many_galpy_times[i] ), atol=1e-2) def test_careful_traceback_and_forward(): """Step by step, project orbit forward, then backward""" bovy_times = np.array([0., np.pi / 3.]) chron_times = torb.convert_bovytime2myr(bovy_times) init_pos_chron = np.array([ 4000, 8000. * np.sqrt(3) / 2, 0, np.sin(np.pi / 3) * 220., -np.cos(np.pi / 3) * 220., 0 ]) init_pos_galpy = torb.convert_cart2galpycoords(init_pos_chron, ts=0.) assert np.allclose(np.array([1., 0, 1, 0, 0, np.pi / 3.]), init_pos_galpy) o = Orbit(vxvv=init_pos_galpy, ro=8., vo=220.) o.integrate(bovy_times, MWPotential2014, method='odeint') orbit_galpy = o.getOrbit() assert np.allclose(init_pos_galpy, orbit_galpy[0]) assert np.allclose(init_pos_galpy + np.array([0., 0., 0., 0., 0., bovy_times[-1]]), orbit_galpy[-1]) orbit_chron = torb.convert_galpycoords2cart(orbit_galpy, ts=bovy_times) assert np.allclose(init_pos_chron, orbit_chron[0]) assert np.allclose(init_pos_chron, orbit_chron[-1]) # Setup for backwards time integration # Currently at time of PI/3 back_init_pos_chron = orbit_chron[-1] back_init_pos_galpy = torb.convert_cart2galpycoords( back_init_pos_chron, bovy_times=bovy_times[-1], ) assert np.allclose(back_init_pos_galpy, torb.convert_cart2galpycoords( back_init_pos_chron, bovy_times=bovy_times[-1] )) back_o = Orbit(vxvv=back_init_pos_galpy, ro=8., vo=220.) back_o.integrate(-1 * bovy_times, MWPotential2014, method='odeint') back_orbit_galpy = back_o.getOrbit() assert np.allclose(back_init_pos_galpy, back_orbit_galpy[0]) assert np.allclose(back_init_pos_galpy - np.array([0., 0., 0., 0., 0., bovy_times[-1]]), back_orbit_galpy[-1]) assert np.allclose(init_pos_galpy, back_orbit_galpy[-1]) back_orbit_chron = torb.convert_galpycoords2cart( back_orbit_galpy, ts=bovy_times[::-1], ) assert np.allclose(init_pos_chron, back_orbit_chron[-1]) def test_traceback_and_forward(): """The test that shows things are broken""" time = 10. # Myr times = np.array([0., 10.]) init_pos_chron = np.array([10., 0., 30., 0., 0., 0.]) # init_pos_chron = np.zeros(6) init_pos_chron = np.array([4000, 8000. * np.sqrt(3) / 2, 0, np.sin(np.pi / 3) * 220., -np.cos(np.pi / 3) * 220., 0]) final_pos_chron = torb.trace_cartesian_orbit( init_pos_chron, times=times, single_age=False)[-1] final_pos_chron = torb.traceforward_from_now( init_pos_chron, time=time, ) assert np.allclose( init_pos_chron, torb.traceback_to_now( final_pos_chron, time, ), atol=1e-5 ) def test_multi_traceback_and_forward(): np.random.seed(0) NPOSITIONS = 10 init_positions = np.random.rand(NPOSITIONS, 6) * 20 - 10 time_spans = np.random.rand(NPOSITIONS) * 30 for pos, time in zip(init_positions, time_spans): final_pos = torb.traceforward_from_now(pos, time) init_pos = torb.traceback_to_now(final_pos, time) assert np.allclose(pos, init_pos, atol=1e-3) for pos, time in zip(init_positions, time_spans): print('time: {}'.format(time)) final_pos = torb.traceback_from_now(pos, time) init_pos = torb.traceforward_to_now(final_pos, time) assert np.allclose(pos, init_pos, atol=1e-3) def test_interval_tracing(): np.random.seed(0) start = np.random.rand(6) * 20 - 10 time_steps = [3., -10., -3., 10] current_pos = start for time_step in time_steps: current_pos = torb.base_trace_cartesian_orbit( current_pos, end_time=time_step, ) assert np.allclose(start, current_pos, atol=1e-3) def test_interval_tracing_orig(): np.random.seed(0) start = np.random.rand(6) * 20 - 10 time_steps = [3., -10., -3., 10] current_pos = start for time_step in time_steps: current_pos = torb.trace_cartesian_orbit( current_pos, times=time_step, single_age=True, ) assert np.allclose(start, current_pos, 1e-3)
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' # astro_table = ms.measureXYZUVW(perf_xyzuvw_now, prec_val[prec]) # xyzuvw_dict = cv.convertMeasurementsToCartesian(astro_table)
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):
logging.info("Performing fit with") logging.info("{} burnin steps".format(BURNIN_STEPS)) logging.info("{} sampling steps".format(SAMPLING_STEPS)) logging.info("{} tolerance".format(C_TOL)) logging.info("In the directory: {}".format(rdir)) # Destination: (inspired by LCC) mean_now = np.array([50., -100., 0., -10., -20., -5.]) # !!!!!!! ^^^^^ IS WRONG!!!!!!!! # ^^^^^^ IS LCC IN HELIOCENTRIC COORDINATES. IN LSR COORDINATES IT IS: mean_now_lsr = cc.convert_helio2lsr(np.array([50., -100., 0., -10., -20., -5.])) # !!!!! WHICH EQUALS [50., -100., 25., 1.1, -7.76, 2.25], much smaller vels # Calculate appropriate starting point mean_then = torb.trace_cartesian_orbit(mean_now, -age) # gather inputs group_pars = np.hstack((mean_then, dX, dV, age, nstars)) try: xyzuvw_now_perf = np.load(rdir + xyzuvw_perf_file) origin = np.load(rdir + group_savefile).item() logging.info("appending to previous attempt") except IOError: logging.info("Beginning fresh run:") logging.info("Input arguments: {}".format(sys.argv[1:])) logging.info("\n" "\tage: {}\n" "\tdX: {}\n" "\tdV: {}\n" "\tnstars: {}\n"
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)