def test_different_component_forms():
"""Check component forms can be different"""
tiny_age = 1e-10
mean1 = np.zeros(6)
covmatrix1 = np.eye(6) * 4
comp1 = SphereComponent(attributes={
'mean':mean1,
'covmatrix':covmatrix1,
'age':tiny_age,
})
mean2 = np.zeros(6) + 10.
covmatrix2 = np.eye(6) * 9
comp2 = EllipComponent(attributes={
'mean':mean2,
'covmatrix':covmatrix2,
'age':tiny_age,
})
starcounts = [100,100]
synth_data = SynthData(pars=[comp1.get_pars(), comp2.get_pars()],
starcounts=starcounts,
Components=[SphereComponent, EllipComponent])
synth_data.synthesise_everything()
assert len(synth_data.table) == np.sum(starcounts)
def test_ellipcomponent_initialisation():
ellip_pars = np.copy(ELLIP_PARS)
# remove correlations for sphere_dx checks
ellip_pars[10:13] = 0.
ellip_comp = EllipComponent(pars=ellip_pars)
assert np.allclose(ellip_pars[:6], ellip_comp._mean)
assert np.allclose(AGE, ellip_comp._age)
sphere_dx = (DX * DY * DZ)**(1. / 3)
assert np.isclose(sphere_dx, ellip_comp.get_sphere_dx())
assert np.isclose(DV, ellip_comp.get_sphere_dv())
def test_ellipcomponent_initialisation():
ellip_pars = np.copy(ELLIP_PARS)
# remove correlations for sphere_dx checks
ellip_pars[10:13] = 0.
ellip_comp = EllipComponent(pars=ellip_pars)
assert np.allclose(ellip_pars[:6], ellip_comp._mean)
assert np.allclose(AGE, ellip_comp._age)
sphere_dx = (DX * DY * DZ)**(1./3)
assert np.isclose(sphere_dx, ellip_comp.get_sphere_dx())
assert np.isclose(DV, ellip_comp.get_sphere_dv())
def test_stationary_comp():
"""
Tests an ellip comp fit with no traceforward, i.e. no evolution of stars or
component
"""
# log_filename = 'logs/compfitter_stationary.log'
# synth_data_savefile = 'temp_data/compfitter_stationary_synthdata.fits'
burnin_step = 1000
true_comp_mean = np.zeros(6)
true_comp_dx = 10.
true_comp_dy = 5.
true_comp_du = 2.
true_comp_dv = 4.
true_comp_roll = 0.
true_comp_pitch = 0.
true_comp_yaw = 0.
true_comp_cov_xv = 0.4
true_comp_age = 1e-10
true_comp_pars = np.hstack([
true_comp_mean,
true_comp_dx,
true_comp_dy,
true_comp_du,
true_comp_dv,
true_comp_roll,
true_comp_pitch,
true_comp_yaw,
true_comp_cov_xv,
true_comp_age,
])
true_comp = EllipComponent(true_comp_pars)
nstars = 100
measurement_error = 1e-10
best_comp, chain, lnprob = run_fit_helper(
true_comp=true_comp,
starcounts=nstars,
measurement_error=measurement_error,
run_name='stationary',
burnin_step=burnin_step,
trace_orbit_func=dummy_trace_orbit_func,
)
assert np.allclose(true_comp.get_mean(), best_comp.get_mean(), atol=1.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=2.0)
def test_lnprior():
dim = 6
mean = np.zeros(dim)
covmatrix = np.identity(dim)
age = 10.
sphere_comp = SphereComponent(attributes={
'mean': mean,
'covmatrix': covmatrix,
'age': age,
})
memb_probs = np.ones(10)
assert np.isfinite(likelihood.lnprior(sphere_comp, memb_probs))
# Now increase age to something ridiculous
sphere_comp.update_attribute(attributes={
'age': 1e10,
})
assert np.isinf(likelihood.lnprior(sphere_comp, memb_probs))
# Try an EllipComponent with a non-symmetrical covariance matrix
covmatrix[0, 1] = 1.01
# covmatrix[1,0] = 100
ellip_comp = EllipComponent(attributes={
'mean': mean,
'covmatrix': covmatrix,
'age': age,
})
assert np.isinf(likelihood.lnprior(ellip_comp, memb_probs))
# Try an EllipComponent with a very broken correlation value
covmatrix[0, 1] = 1.01
covmatrix[1, 0] = 1.01
ellip_comp = EllipComponent(attributes={
'mean': mean,
'covmatrix': covmatrix,
'age': age,
})
assert np.isinf(likelihood.lnprior(ellip_comp, memb_probs))
def test_externalise_and_internalise_pars():
"""Check that pars are successfully converted from internal form (used by
emcee) to external form (interacted with by user) successfully"""
# Check SphereComponent
internal_sphere_pars = np.copy(SPHERE_PARS)
internal_sphere_pars[6:8] = np.log(internal_sphere_pars[6:8])
sphere_comp = SphereComponent(emcee_pars=internal_sphere_pars)
external_sphere_pars = sphere_comp.get_pars()
assert np.allclose(SPHERE_PARS, external_sphere_pars)
re_internal_sphere_pars = sphere_comp.internalise(external_sphere_pars)
assert np.allclose(internal_sphere_pars, re_internal_sphere_pars)
# Check EllipComponent
internal_ellip_pars = np.copy(ELLIP_PARS)
internal_ellip_pars[6:10] = np.log(internal_ellip_pars[6:10])
ellip_comp = EllipComponent(emcee_pars=internal_ellip_pars)
external_ellip_pars = ellip_comp.get_pars()
assert np.allclose(ELLIP_PARS, external_ellip_pars)
re_internal_ellip_pars = ellip_comp.internalise(external_ellip_pars)
assert np.allclose(internal_ellip_pars, re_internal_ellip_pars)
def test_externalise_and_internalise_pars():
"""Check that pars are successfully converted from internal form (used by
emcee) to external form (interacted with by user) successfully"""
# Check SphereComponent
internal_sphere_pars = np.copy(SPHERE_PARS)
internal_sphere_pars[6:8] = np.log(internal_sphere_pars[6:8])
sphere_comp = SphereComponent(emcee_pars=internal_sphere_pars)
external_sphere_pars = sphere_comp.get_pars()
assert np.allclose(SPHERE_PARS, external_sphere_pars)
re_internal_sphere_pars = sphere_comp.internalise(external_sphere_pars)
assert np.allclose(internal_sphere_pars, re_internal_sphere_pars)
# Check EllipComponent
internal_ellip_pars = np.copy(ELLIP_PARS)
internal_ellip_pars[6:10] = np.log(internal_ellip_pars[6:10])
ellip_comp = EllipComponent(emcee_pars=internal_ellip_pars)
external_ellip_pars = ellip_comp.get_pars()
assert np.allclose(ELLIP_PARS, external_ellip_pars)
re_internal_ellip_pars = ellip_comp.internalise(external_ellip_pars)
assert np.allclose(internal_ellip_pars, re_internal_ellip_pars)
def run_fit_helper(true_comp,
starcounts,
measurement_error,
burnin_step=None,
run_name='default',
trace_orbit_func=None,
Component=EllipComponent,
init_pars=None):
py_vers = sys.version[0]
save_dir = 'temp_data/%s_compfitter_%s/' % (py_vers, run_name)
data_filename = save_dir + 'synth_data.fits'.format(py_vers, run_name)
plot_dir = save_dir
print("---------", save_dir)
if not os.path.isdir(save_dir):
os.mkdir(save_dir)
log_filename = save_dir + 'log.log'.format(py_vers, run_name)
logging.basicConfig(level=logging.INFO,
filename=log_filename,
filemode='w')
synth_data = SynthData(pars=true_comp.get_pars(),
starcounts=starcounts,
measurement_error=measurement_error,
Components=Component)
synth_data.synthesise_everything()
tabletool.convert_table_astro2cart(synth_data.table,
write_table=True,
filename=data_filename)
print("newPars ------------------------------ \n", init_pars)
if init_pars is None:
internal_pars = None
else:
internal_pars = Component.internalise(init_pars)
res = cf.fit_comp(data=synth_data.table,
plot_it=True,
burnin_steps=burnin_step,
store_burnin_chains=True,
plot_dir=plot_dir,
save_dir=save_dir,
trace_orbit_func=trace_orbit_func,
optimisation_method='emcee',
Component=Component,
init_pars=internal_pars)
comps_filename = save_dir + 'true_and_best_comp.py'
best_comp = res[0]
EllipComponent.store_raw_components(comps_filename, [true_comp, best_comp])
star_pars = tabletool.build_data_dict_from_table(synth_data.table)
plot_results(true_comp,
best_fit_comp=res[0],
star_pars=star_pars,
plt_dir=save_dir)
return res
def test_any_age_comp(age=32):
"""
Tests an ellip component fit with age as a parameter
"""
# log_filename = 'logs/compfitter_stationary.log'
# synth_data_savefile = 'temp_data/compfitter_stationary_synthdata.fits'
burnin_step = 2000
true_comp_mean = [0., 0., 0., 2., 2., 2.]
true_comp_dx = 8.
true_comp_dy = 3.
true_comp_du = 6.
true_comp_dv = 3.
true_comp_roll = 0.
true_comp_pitch = 0.5
true_comp_yaw = 1.
true_comp_cov_xv = 2.5
true_comp_age = age
true_comp_pars = np.hstack([
true_comp_mean,
true_comp_dx,
true_comp_dy,
true_comp_du,
true_comp_dv,
true_comp_roll,
true_comp_pitch,
true_comp_yaw,
true_comp_cov_xv,
true_comp_age,
])
true_comp = EllipComponent(true_comp_pars)
nstars = 100
measurement_error = 1e-10
best_comp, chain, lnprob = run_fit_helper(
true_comp=true_comp,
starcounts=nstars,
measurement_error=measurement_error,
run_name='priorTest_11_age_%.1e' % true_comp_age,
burnin_step=burnin_step,
trace_orbit_func=trace_epicyclic_orbit,
)
print("Age: --~~~~~~~~~~~~~~~~", best_comp.get_age())
# old_pars = best_comp.get_pars()
# edited_init_pars = np.copy(old_pars)
# # edited_init_pars[10] = old_pars[10]+(np.pi/2)
# # edited_init_pars[11] = old_pars[11]+(np.pi/2)
# edited_init_pars[12] = old_pars[12]+(np.pi/2)
#
# edited_best_comp, edited_chain, edited_lnprob = run_fit_helper(
# true_comp=true_comp, starcounts=nstars,
# measurement_error=measurement_error,
# run_name='priorTest_01_Edited_age_%.1e'%true_comp_age,
# burnin_step=burnin_step,
# trace_orbit_func=trace_epicyclic_orbit,
# init_pars=edited_init_pars
# )
# if lnprob.max() < edited_lnprob.max():
# print("lnProb - max", lnprob.max())
# print("edited _ lnProb - max", edited_lnprob.max())
#
# print("-------------------------------------------------------")
# print("Pi/2 Error, A second fit was run and is set as best fit")
# print("-------------------------------------------------------")
# print("New Age: --~~~~~~~~~~~~~~~~", edited_best_comp.get_age())
assert np.allclose(true_comp.get_mean(), best_comp.get_mean(), atol=1.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=2.0)
print("Plotting fed stars)")
# Setting up file names
synth_fit = 'fed_stars'
rdir = '../../results/archive/fed_fits/20/gaia/'
origins_file = rdir + 'origins.npy'
origin_comp_file = rdir + 'origin_ellip_comp.npy'
chain_file = rdir + 'final_chain.npy'
lnprob_file = rdir + 'final_lnprob.npy'
star_pars_file = rdir + 'xyzuvw_now.fits'
init_xyzuvw_file = rdir + '../xyzuvw_init_offset.npy'
perf_xyzuvw_now = rdir + '../perf_xyzuvw.npy'
# loading in data
best_comp = SphereComponent.get_best_from_chain(chain_file, lnprob_file)
origin_comp = EllipComponent.load_raw_components(origin_comp_file)[0]
init_xyzuvw = np.load(init_xyzuvw_file)
star_pars = dt.loadXYZUVW(star_pars_file)
perf_mean_now = np.load(perf_xyzuvw_now)
original_origin = dt.loadGroups(origins_file)[0]
# assigning useful shorthands
mns = star_pars['xyzuvw']
covs = star_pars['xyzuvw_cov']
fed_xranges, fed_yranges = calcRanges(
{'xyzuvw': np.vstack((star_pars['xyzuvw'], init_xyzuvw))},
sep_axes=True,
)