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)
xyzuvw_init, origin = \
syn.synthesiseXYZUVW(group_pars, form='sphere', return_group=True,
xyzuvw_savefile=rdir + xyzuvw_init_savefile,
group_savefile=rdir + group_savefile)
logging.info("Origin has values\n"
"\tage: {}\n"
"\tsph_dX: {}\n"
"\tdV: {}\n"
"\tnstars: {}".format(
origin.age,
origin.sphere_dx,
origin.dv,
origin.nstars,
))
xyzuvw_now_perf =\
torb.trace_many_cartesian_orbit(xyzuvw_init, origin.age, single_age=True,
savefile=rdir+xyzuvw_perf_file)
for prec in precs:
logging.info("Fitting to prec: {}".format(prec))
pdir = rdir + prec + '/'
mkpath(pdir)
# os.chdir(prec)
try:
dummy = np.load(pdir + group_savefile)
logging.info("Precision [{}] already fitted for".format(prec))
except IOError:
# convert XYZUVW data into astrometry
astro_table = chronostar.synthdata.measureXYZUVW(xyzuvw_now_perf,
prec_val[prec],
savefile=pdir +
astro_savefile)
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)
ubound += margin
lbound -= margin
nbg_stars = int(BG_DENS * np.prod(ubound - lbound))
# centre bg stars on mean of assoc stars
# centre = np.mean(all_xyzuvw_now_perf, axis=0)
# centre = 0.5 * (ubound + lbound)
# spread = ubound - lbound
file_stems = [scen + "_" + prec for scen in scenarios for prec in precs]
#fit_files = [res_dir + "xyzuvw_now_" + stem + ".fits" for stem in file_stems]
assert len(fits_files) == len(file_stems)
for i in range(len(fits_files)):
xyzuvw_dict = gf.loadXYZUVW(fits_files[i])
nstars = xyzuvw_dict['xyzuvw'].shape[0]
origin_file = res_dir + file_stems[i][:-5] + '_origins.npy'
group = np.load(origin_file).item()
true_age = group.age
ntimes = int(2 * true_age + 1)
times = -np.linspace(0, 2 * true_age, ntimes)
tb = torb.trace_many_cartesian_orbit(xyzuvw_dict['xyzuvw'], times=times,
single_age=False,
savefile=res_dir + 'tb_{}.npy'. \
format(int(true_age)))
# for each timestep, get mean of association and distance from mean
dists = np.zeros((nstars, ntimes))
stds = np.zeros(ntimes)
for tix in range(ntimes):
dists[:, tix] = get_euc_dist(tb[:, tix, :3],
tb[:, tix, :3].mean(axis=0))
plt.clf()
plt.plot(-times, dists.T, '#888888')
plt.plot(-times, np.median(dists.T, axis=1), 'r')
plt.savefig(res_dir + file_stems[i]+"_sep_from_mean.pdf")
import chronostar.retired2.datatool as dt
rdir = "../results/synth_fit/50_2_1_50/"
# rdir = "../results/synth_fit/30_5_2_100/"
xyzuvw_init_file = rdir + "xyzuvw_init.npy"
xyzuvw_init = np.load(xyzuvw_init_file)
origin = np.load(rdir + 'origins.npy').item()
origin = dt.loadGroups(rdir + 'origins.npy')
max_age = origin.age
ntimes = int(max_age) + 1
#ntimes = 3
times = np.linspace(1e-5, max_age, ntimes)
traceforward = torb.trace_many_cartesian_orbit(xyzuvw_init, times, False)
nstars = xyzuvw_init.shape[0]
def plot_subplot(traceforward, t_ix, dim1, dim2, ax):
flat_tf = traceforward.reshape(-1,6)
# mins = np.min(flat_tf, axis=0)
# maxs = np.max(flat_tf, axis=0)
labels = ['X [pc]', 'Y [pc]', 'Z [pc]',
'U [km/s]', 'V [km/s]', 'W [km/s]']
for i in range(nstars):
# plot orbits
ax.plot(traceforward[i, :t_ix, dim1],
traceforward[i, :t_ix, dim2],
'b',
alpha=0.1)
#for cnt, dummy_group_pars_ex in enumerate(dummy_groups):
mean = my_group.mean
cov = my_group.generateEllipticalCovMatrix()
stars = np.random.multivariate_normal(mean, cov, nstars)
if plotit:
plt.clf()
plt.plot(stars[:, 0], stars[:, 1], 'b.')
#plt.hist2d(stars[:,0], stars[:,1], bins=20)
chronostar.fitplotter.plotCovEllipse(cov[:2, :2],
mean,
color='b',
alpha=0.3)
new_stars = np.zeros(stars.shape)
new_stars = torb.trace_many_cartesian_orbit(stars, np.array(0, age))
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)
import chronostar.traceorbit as torb
import chronostar.retired2.converter as cv
import chronostar.fitplotter as fp
pdir = "../figures/paper1/"
ERROR = 1.0
# set X to be really large to encourage significant errors
origin_pars = np.array([50., 0., -50., 0., 10., 0., 10., 0.5, 10., 100])
xyzuvw_init, origin = syn.synthesiseXYZUVW(origin_pars,
form='sphere',
return_group=True,
internal=False)
xyzuvw_now_perf = torb.trace_many_cartesian_orbit(xyzuvw_init,
times=origin.age,
single_age=True)
astr_table = chronostar.synthdata.measureXYZUVW(xyzuvw_now_perf, ERROR)
star_pars = cv.convertMeasurementsToCartesian(astr_table)
dims = 'XYZUVW'
plt.clf()
fp.plot1DProjection(0, star_pars, [origin], np.array([origin.nstars]))
plt.savefig(pdir + "1dprojection.pdf")
plt.clf()
dim1, dim2 = 'xu'
fp.plotPane(dim1,
dim2,
groups=origin,
star_pars=star_pars,
group_then=True,
#get_ipython().magic(u'autoreload 2')
import numpy as np
import matplotlib.pyplot as plt
import sys
sys.path.insert(0, '../../..')
import chronostar.synthdata as syn
import chronostar.compfitter as gf
import chronostar.traceorbit as torb
xyzuvw_init = np.load('xyzuvw_init.npy')
origin = np.load('origins.npy').item()
max_age = origin.age
ntimes = int(max_age) + 1
times = np.linspace(1e-5, max_age, ntimes)
traceback = torb.trace_many_cartesian_orbit(xyzuvw_init, times, False)
nstars = xyzuvw_init.shape[0]
#times = np.array([1e-5])
for i in range(nstars):
plt.plot(traceback[i,:,0], traceback[i,:,1], 'b')
plt.xlabel('X [pc]')
plt.ylabel('Y [pc]')
plt.savefig("xy.png")
plt.clf()
# TRACEBACK
for t_ix in range(times.shape[0]):
plt.clf()
file_stems = [scen + "_" + prec for scen in scenarios for prec in precs]
#fit_files = [res_dir + "xyzuvw_now_" + stem + ".fits" for stem in file_stems]
assert len(fits_files) == len(file_stems)
for i in range(len(fits_files)):
xyzuvw_dict = gf.loadXYZUVW(fits_files[i])
nstars = xyzuvw_dict['xyzuvw'].shape[0]
origin_file = res_dir + file_stems[i][:-5] + '_origins.npy'
group = np.load(origin_file).item()
true_age = group.age
ntimes = int(2 * true_age + 1)
times = -np.linspace(0, 2 * true_age, ntimes)
tb = torb.trace_many_cartesian_orbit(xyzuvw_dict['xyzuvw'], times=times,
single_age=False,
savefile=res_dir + 'tb_{}.npy'. \
format(int(true_age)))
# for each timestep, get mean of association and distance from mean
dists = np.zeros((nstars, ntimes))
stds = np.zeros(ntimes)
for tix in range(ntimes):
dists[:, tix] = get_euc_dist(tb[:, tix, :3], tb[:,
tix, :3].mean(axis=0))
plt.clf()
plt.plot(-times, dists.T, '#888888')
plt.plot(-times, np.median(dists.T, axis=1), 'r')
plt.savefig(res_dir + file_stems[i] + "_sep_from_mean.pdf")
my_group = Component(dummy_groups[0], form=False)
#for cnt, dummy_group_pars_ex in enumerate(dummy_groups):
mean = my_group.mean
cov = my_group.generateEllipticalCovMatrix()
stars = np.random.multivariate_normal(mean, cov, nstars)
if plotit:
plt.clf()
plt.plot(stars[:,0], stars[:,1], 'b.')
#plt.hist2d(stars[:,0], stars[:,1], bins=20)
chronostar.fitplotter.plotCovEllipse(cov[:2, :2], mean, color='b', alpha=0.3)
new_stars = np.zeros(stars.shape)
new_stars = torb.trace_many_cartesian_orbit(stars, np.array(0, age))
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)
# synthesise perfect XYZUVW data
logging.info("Synthesising data")
xyzuvw_init, origin =\
syn.synthesise_xyzuvw(group_pars, sphere=True,
xyzuvw_savefile=xyzuvw_init_savefile,
return_group=True, group_savefile=group_savefile)
logging.info("Origin has values\n"
"\tage: {}\n"
"\tsph_dX: {}\n"
"\tdV: {}\n"
"\tnstars: {}".format(
origin.age, origin.sphere_dx, origin.dv, origin.nstars,
))
xyzuvw_now_perf =\
torb.trace_many_cartesian_orbit(xyzuvw_init, origin.age, single_age=True,
savefile=xyzuvw_perf_file)
logging.info(mpi_msg)
if not using_mpi:
logging.info("MPI available! - call this with e.g. mpirun -np 19"
" python perform_synth_fit.py")
for prec in precs:
logging.info("Fitting to prec: {}".format(prec))
mkpath(prec)
os.chdir(prec)
np.save(group_savefile, origin) # store in each directory, for hexplotter
try:
res = np.load(result_file)
logging.info("Precision [{}] already fitted for".format(prec))
except IOError:
logging.basicConfig(level=logging.INFO, stream=sys.stdout)
save_dir = 'temp_data/'
group_savefile = save_dir + 'origins.npy'
xyzuvw_init_savefile = save_dir + 'xyzuvw_init.npy'
astro_savefile = save_dir + 'astro_table.txt'
group_pars = [0., 0., 0., 0., 0., 0., 1., 1., 0.5, 20]
xyzuvw_init, group = syn.synthesiseXYZUVW(group_pars, form='sphere',
return_group=True,
xyzuvw_savefile=xyzuvw_init_savefile,
group_savefile=group_savefile)
logging.info("Age is: {} Myr".format(group.age))
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()
xyzuvw_init += mean_then
np.save(rdir+'xyzuvw_init_offset.npy', xyzuvw_init)
# xyzuvw_init[:,:3] += approx_final_pos
# fit an approximate Gaussian to initial distribution for reference
mean = np.mean(xyzuvw_init, axis=0)
dx, dy, dz = np.std(xyzuvw_init[:,:3], axis=0)
dv = np.prod(np.std(xyzuvw_init[:,3:], axis=0))**(1./3.)
group_pars = np.hstack((mean, dx, dy, dz, dv, 0., 0., 0., age))
origin = chronostar.component.Component(group_pars, form=False,
internal=False)
np.save(group_savefile, origin)
xyzuvw_now_perf =\
torb.trace_many_cartesian_orbit(xyzuvw_init, age, single_age=True,
savefile=rdir+xyzuvw_perf_file)
if not using_mpi:
logging.info("MPI available! - call this with e.g. mpirun -np 19"
" python perform_synth_fit.py")
# Performing fit for each precision
for prec in precs:
logging.info("Fitting to prec: {}".format(prec))
pdir = rdir + prec + '/'
mkpath(pdir)
np.save(pdir+group_savefile, origin) # store in each directory, for hexplotter
try:
best_group = dt.loadGroups(pdir + 'final_groups.npy')
logging.info("Precision [{}] already fitted for".format(prec))
except IOError:
logging.basicConfig(level=logging.INFO, stream=sys.stdout)
save_dir = 'temp_data/'
group_savefile = save_dir + 'origins.npy'
xyzuvw_init_savefile = save_dir + 'xyzuvw_init.npy'
astro_savefile = save_dir + 'astro_table.txt'
group_pars = [0., 0., 0., 0., 0., 0., 1., 1., 0.5, 20]
xyzuvw_init, group = syn.synthesiseXYZUVW(group_pars,
form='sphere',
return_group=True,
xyzuvw_savefile=xyzuvw_init_savefile,
group_savefile=group_savefile)
logging.info("Age is: {} Myr".format(group.age))
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)))
import chronostar.retired2.datatool as dt
rdir = "../results/synth_fit/50_2_1_50/"
# rdir = "../results/synth_fit/30_5_2_100/"
xyzuvw_init_file = rdir + "xyzuvw_init.npy"
xyzuvw_init = np.load(xyzuvw_init_file)
origin = np.load(rdir + 'origins.npy').item()
origin = dt.loadGroups(rdir + 'origins.npy')
max_age = origin.age
ntimes = int(max_age) + 1
#ntimes = 3
times = np.linspace(1e-5, max_age, ntimes)
traceforward = torb.trace_many_cartesian_orbit(xyzuvw_init, times, False)
nstars = xyzuvw_init.shape[0]
def plot_subplot(traceforward, t_ix, dim1, dim2, ax):
flat_tf = traceforward.reshape(-1, 6)
# mins = np.min(flat_tf, axis=0)
# maxs = np.max(flat_tf, axis=0)
labels = ['X [pc]', 'Y [pc]', 'Z [pc]', 'U [km/s]', 'V [km/s]', 'W [km/s]']
for i in range(nstars):
# plot orbits
ax.plot(traceforward[i, :t_ix, dim1],
traceforward[i, :t_ix, dim2],
'b',
alpha=0.1)
xyzuvw_init += offset
xyzuvw_init[:,:3] += approx_final_pos
# fit an approximate Gaussian to initial distribution
mean = np.mean(xyzuvw_init, axis=0)
dx, dy, dz = np.std(xyzuvw_init[:,:3], axis=0)
dv = np.prod(np.std(xyzuvw_init[:,3:], axis=0))**(1./3.)
group_pars = np.hstack((mean, dx, dy, dz, dv, 0., 0., 0., age))
pdb.set_trace()
origin = chronostar.component.Component(group_pars, form=False)
np.save(group_savefile, origin)
xyzuvw_now_perf =\
torb.trace_many_cartesian_orbit(xyzuvw_init, age, single_age=True,
savefile=xyzuvw_perf_file)
logging.info(mpi_msg)
if not using_mpi:
logging.info("MPI available! - call this with e.g. mpirun -np 19"
" python perform_synth_fit.py")
# Performing fit for each precision
for prec in precs:
logging.info("Fitting to prec: {}".format(prec))
mkpath(prec)
os.chdir(prec)
np.save(group_savefile, origin) # store in each directory, for hexplotter
try:
res = np.load(result_file)
logging.info("Precision [{}] already fitted for".format(prec))
