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