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_expectation():
"""
Super basic, generates some association stars along
with some background stars and checks membership allocation
is correct
"""
age = 1e-5
ass_pars1 = np.array([0, 0, 0, 0, 0, 0, 5., 2., age])
comp1 = SphereComponent(ass_pars1)
ass_pars2 = np.array([100., 0, 0, 20, 0, 0, 5., 2., age])
comp2 = SphereComponent(ass_pars2)
starcounts = [100,100]
synth_data = SynthData(pars=[ass_pars1, ass_pars2],
starcounts=starcounts)
synth_data.synthesise_everything()
tabletool.convert_table_astro2cart(synth_data.table)
true_memb_probs = np.zeros((np.sum(starcounts), 2))
true_memb_probs[:starcounts[0], 0] = 1.
true_memb_probs[starcounts[0]:, 1] = 1.
# star_means, star_covs = tabletool.buildDataFromTable(synth_data.astr_table)
# all_lnols = em.getAllLnOverlaps(
# synth_data.astr_table, [comp1, comp2]
# )
fitted_memb_probs = em.expectation(
tabletool.build_data_dict_from_table(synth_data.table),
[comp1, comp2]
)
assert np.allclose(true_memb_probs, fitted_memb_probs, atol=1e-10)
def test_pythonFuncs():
"""
TODO: remove the requirements of file, have data stored in file?
"""
true_comp_mean = np.zeros(6)
true_comp_dx = 2.
true_comp_dv = 2.
true_comp_covmatrix = np.identity(6)
true_comp_covmatrix[:3, :3] *= true_comp_dx ** 2
true_comp_covmatrix[3:, 3:] *= true_comp_dv ** 2
true_comp_age = 1e-10
true_comp = SphereComponent(attributes={
'mean': true_comp_mean,
'covmatrix': true_comp_covmatrix,
'age': true_comp_age,
})
nstars = 100
synth_data = SynthData(pars=true_comp.get_pars(), starcounts=nstars)
synth_data.synthesise_everything()
tabletool.convert_table_astro2cart(synth_data.table)
star_data = tabletool.build_data_dict_from_table(synth_data.table)
# star_data['means'] = star_data['means']
# star_data['covs'] = star_data['covs']
group_mean = true_comp.get_mean()
group_cov = true_comp.get_covmatrix()
# Test overlap with true component
co1s = []
co2s = []
for i, (scov, smn) in enumerate(zip(star_data['covs'], star_data['means'])):
co1s.append(co1(group_cov, group_mean, scov, smn))
co2s.append(co2(group_cov, group_mean, scov, smn))
co1s = np.array(co1s)
co2s = np.array(co2s)
co3s = np.exp(p_lno(group_cov, group_mean,
star_data['covs'], star_data['means']))
assert np.allclose(co1s, co2s)
assert np.allclose(co2s, co3s)
assert np.allclose(co1s, co3s)
# Test overlap with neighbouring star (with the aim of testing
# tiny overlap values). Note that most overlaps go to 0, but the
# log overlaps retain the information
co1s = []
co2s = []
for i, (scov, smn) in enumerate(zip(star_data['covs'], star_data['means'])):
co1s.append(co1(star_data['covs'][15], star_data['means'][15],
scov, smn))
co2s.append(co2(star_data['covs'][15], star_data['means'][15],
scov, smn))
co1s = np.array(co1s)
co2s = np.array(co2s)
lnos = p_lno(star_data['covs'][15], star_data['means'][15],
star_data['covs'], star_data['means'])
co3s = np.exp(lnos)
assert np.allclose(co1s, co2s)
assert np.allclose(co2s, co3s)
assert np.allclose(co1s, co3s)
def test_generateInitXYZUVW():
"""Check that the mean of initial xyzuvw of stars matches that of the
initialising component"""
starcounts = (int(1e6), )
sd = SynthData(pars=PARS[:1], starcounts=starcounts, Components=COMPONENTS)
sd.generate_all_init_cartesian()
comp = SphereComponent(PARS[0])
init_xyzuvw = sd.extract_data_as_array([dim + '0' for dim in 'xyzuvw'])
assert np.allclose(comp.get_mean(), np.mean(init_xyzuvw, axis=0), atol=0.1)
def test_generateInitXYZUVW():
"""Check that the mean of initial xyzuvw of stars matches that of the
initialising component"""
starcounts = (int(1e6),)
sd = SynthData(pars=PARS[:1], starcounts=starcounts, Components=COMPONENTS)
sd.generate_all_init_cartesian()
comp = SphereComponent(PARS[0])
init_xyzuvw = sd.extract_data_as_array([dim + '0' for dim in 'xyzuvw'])
assert np.allclose(comp.get_mean(), np.mean(init_xyzuvw, axis=0),
atol=0.1)
def test_lnprob_func():
"""
Generates two components. Generates a synthetic data set based on the
first component. Confrims that the lnprob is larger for the first
component than the second.
"""
measurement_error = 1e-10
star_count = 500
tiny_age = 1e-10
dim = 6
comp_covmatrix = np.identity(dim)
comp_means = {
'comp1': np.zeros(dim),
'comp2': 10 * np.ones(dim)
}
comps = {}
data = {}
for comp_name in comp_means.keys():
comp = SphereComponent(attributes={
'mean':comp_means[comp_name],
'covmatrix':comp_covmatrix,
'age':tiny_age
})
synth_data = SynthData(pars=[comp.get_pars()], starcounts=star_count,
measurement_error=measurement_error)
synth_data.synthesise_everything()
tabletool.convert_table_astro2cart(synth_data.table)
data[comp_name] = tabletool.build_data_dict_from_table(synth_data.table)
comps[comp_name] = comp
lnprob_comp1_data1 = likelihood.lnprob_func(pars=comps['comp1'].get_pars(),
data=data['comp1'])
lnprob_comp2_data1 = likelihood.lnprob_func(pars=comps['comp2'].get_pars(),
data=data['comp1'])
lnprob_comp1_data2 = likelihood.lnprob_func(pars=comps['comp1'].get_pars(),
data=data['comp2'])
lnprob_comp2_data2 = likelihood.lnprob_func(pars=comps['comp2'].get_pars(),
data=data['comp2'])
print(lnprob_comp1_data1)
print(lnprob_comp2_data1)
print(lnprob_comp1_data2)
print(lnprob_comp2_data2)
assert lnprob_comp1_data1 > lnprob_comp2_data1
assert lnprob_comp2_data2 > lnprob_comp1_data2
# Check that the different realisations only differ by 20%
assert np.isclose(lnprob_comp1_data1, lnprob_comp2_data2, rtol=2e-1)
assert np.isclose(lnprob_comp1_data2, lnprob_comp2_data1, rtol=2e-1)
def test_stationary_component():
"""
Integrated test which fits a single component to a synthetic association.
Runtime on my mac (single thread) is ~ 20 mins. Check logs/compfitter.log
and temp_plots/*.png for progress.
Takes about 10 mins single thread with C implementation of overlap
or ~40 mins with python implementation of overlap
"""
# log_filename = 'logs/compfitter_stationary.log'
# synth_data_savefile = 'temp_data/compfitter_stationary_synthdata.fits'
short_burnin_step = 200
true_comp_mean = np.zeros(6)
true_comp_dx = 2.
true_comp_dv = 2.
true_comp_covmatrix = np.identity(6)
true_comp_covmatrix[:3,:3] *= true_comp_dx**2
true_comp_covmatrix[3:,3:] *= true_comp_dv**2
true_comp_age = 1e-10
true_comp = SphereComponent(attributes={
'mean':true_comp_mean,
'covmatrix':true_comp_covmatrix,
'age':true_comp_age,
})
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=short_burnin_step,
trace_orbit_func=dummy_trace_orbit_func,
)
np.save('temp_data/{}_compfitter_stationary_' \
'true_and_best_comp.npy'.format(PY_VERS),
[true_comp, best_comp],)
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_get_lnoverlaps():
"""
Confirms that star-component overlaps get smaller as stars get further
away.
First generates a component `sphere_comp`. Then generates three stars.
The first one is identical to `sphere_comp` in mean and covmatrix.
The other two share the same covmatrix yet are separated in X.
We check that the overlap integral is smaller for the more separated
stars.
"""
dim = 6
mean = np.zeros(dim)
covmatrix = np.identity(dim)
age = 1e-10
sphere_comp = SphereComponent(attributes={
'mean':mean,
'covmatrix':covmatrix,
'age':age,
})
dx_offsets = [0., 1., 10.]
star_comps = []
for dx_offset in dx_offsets:
star = SphereComponent(attributes={
'mean':sphere_comp.get_mean()+np.array([dx_offset,0.,0.,0.,0.,0.]),
'covmatrix':sphere_comp.get_covmatrix(),
'age':sphere_comp.get_age(),
})
star_comps.append(star)
nstars = len(star_comps)
dummy_table = Table(data=np.arange(nstars).reshape(nstars,1),
names=['name'])
tabletool.append_cart_cols_to_table(dummy_table)
for star_comp, row in zip(star_comps, dummy_table):
tabletool.insert_data_into_row(row,
star_comp.get_mean(),
star_comp.get_covmatrix(),
cartesian=True,
)
dummy_data = tabletool.build_data_dict_from_table(dummy_table)
ln_overlaps = likelihood.get_lnoverlaps(sphere_comp, data=dummy_data)
# Checks that ln_overlaps is descending
assert np.allclose(ln_overlaps, sorted(ln_overlaps)[::-1])
def test_stationary_component():
"""
Integrated test which fits a single component to a synthetic association.
Runtime on my mac (single thread) is ~ 20 mins. Check logs/compfitter.log
and temp_plots/*.png for progress.
Takes about 10 mins single thread with C implementation of overlap
or ~40 mins with python implementation of overlap
"""
# log_filename = 'logs/compfitter_stationary.log'
# synth_data_savefile = 'temp_data/compfitter_stationary_synthdata.fits'
short_burnin_step = 200
true_comp_mean = np.zeros(6)
true_comp_dx = 2.
true_comp_dv = 2.
true_comp_covmatrix = np.identity(6)
true_comp_covmatrix[:3, :3] *= true_comp_dx**2
true_comp_covmatrix[3:, 3:] *= true_comp_dv**2
true_comp_age = 1e-10
true_comp = SphereComponent(
attributes={
'mean': true_comp_mean,
'covmatrix': true_comp_covmatrix,
'age': true_comp_age,
})
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=short_burnin_step,
trace_orbit_func=dummy_trace_orbit_func,
)
np.save('temp_data/{}_compfitter_stationary_' \
'true_and_best_comp.npy'.format(PY_VERS),
[true_comp, best_comp],)
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_get_best_from_chain():
# Triplicate sphere pars (are copies)
# Represents a chain with 2 walkers and 3 steps
intern_sphere_pars = SphereComponent.internalise(SPHERE_PARS)
dummy_chain = np.array(
[[intern_sphere_pars, intern_sphere_pars, intern_sphere_pars],
[intern_sphere_pars, intern_sphere_pars, intern_sphere_pars]])
dummy_lnprob = np.zeros(dummy_chain.shape[:2])
# Incorporate identifying marker at desired index
true_best_ix = (1, 1)
dummy_chain[true_best_ix][0] = 10.
dummy_lnprob[true_best_ix] = 1.
best_comp = SphereComponent.get_best_from_chain(dummy_chain, dummy_lnprob)
assert np.allclose(dummy_chain[true_best_ix], best_comp.get_emcee_pars())
def test_lnprob_func():
"""
Generates two components. Generates a synthetic data set based on the
first component. Confrims that the lnprob is larger for the first
component than the second.
"""
measurement_error = 1e-10
star_count = 500
tiny_age = 1e-10
dim = 6
comp_covmatrix = np.identity(dim)
comp_means = {
'comp1': np.zeros(dim),
'comp2': 10 * np.ones(dim)
}
comps = {}
data = {}
for comp_name in comp_means.keys():
comp = SphereComponent(attributes={
'mean':comp_means[comp_name],
'covmatrix':comp_covmatrix,
'age':tiny_age
})
synth_data = SynthData(pars=[comp.get_pars()], starcounts=star_count,
measurement_error=measurement_error)
synth_data.synthesise_everything()
tabletool.convert_table_astro2cart(synth_data.table)
data[comp_name] = tabletool.build_data_dict_from_table(synth_data.table)
comps[comp_name] = comp
lnprob_comp1_data1 = likelihood.lnprob_func(pars=comps['comp1'].get_pars(),
data=data['comp1'])
lnprob_comp2_data1 = likelihood.lnprob_func(pars=comps['comp2'].get_pars(),
data=data['comp1'])
lnprob_comp1_data2 = likelihood.lnprob_func(pars=comps['comp1'].get_pars(),
data=data['comp2'])
lnprob_comp2_data2 = likelihood.lnprob_func(pars=comps['comp2'].get_pars(),
data=data['comp2'])
assert lnprob_comp1_data1 > lnprob_comp2_data1
assert lnprob_comp2_data2 > lnprob_comp1_data2
# Check that the different realisations only differ by 10%
assert np.isclose(lnprob_comp1_data1, lnprob_comp2_data2, rtol=1e-1)
assert np.isclose(lnprob_comp1_data2, lnprob_comp2_data1, rtol=1e-1)
def test_swigImplementation():
"""
Compares the swigged c implementation against the python one in
likelihood.py
"""
true_comp_mean = np.zeros(6)
true_comp_dx = 2.
true_comp_dv = 2.
true_comp_covmatrix = np.identity(6)
true_comp_covmatrix[:3, :3] *= true_comp_dx**2
true_comp_covmatrix[3:, 3:] *= true_comp_dv**2
true_comp_age = 1e-10
true_comp = SphereComponent(
attributes={
'mean': true_comp_mean,
'covmatrix': true_comp_covmatrix,
'age': true_comp_age,
})
nstars = 100
synth_data = SynthData(pars=true_comp.get_pars(), starcounts=nstars)
synth_data.synthesise_everything()
tabletool.convert_table_astro2cart(synth_data.table)
star_data = tabletool.build_data_dict_from_table(synth_data.table)
p_lnos = p_lno(true_comp.get_covmatrix(), true_comp.get_mean(),
star_data['covs'], star_data['means'])
c_lnos = c_lno(true_comp.get_covmatrix(), true_comp.get_mean(),
star_data['covs'], star_data['means'], nstars)
assert np.allclose(p_lnos, c_lnos)
assert np.isfinite(p_lnos).all()
assert np.isfinite(c_lnos).all()
def test_execution_simple_fit():
"""
Don't test for correctness, but check that everything actually executes
"""
run_name = 'quickdirty'
logging.info(60 * '-')
logging.info(15 * '-' + '{:^30}'.format('TEST: ' + run_name) + 15 * '-')
logging.info(60 * '-')
savedir = 'temp_data/{}_expectmax_{}/'.format(PY_VERS, run_name)
mkpath(savedir)
data_filename = savedir + '{}_expectmax_{}_data.fits'.format(
PY_VERS, run_name)
log_filename = 'temp_data/{}_expectmax_{}/log.log'.format(
PY_VERS, run_name)
logging.basicConfig(level=logging.INFO,
filemode='w',
filename=log_filename)
uniform_age = 1e-10
sphere_comp_pars = np.array([
# X, Y, Z, U, V, W, dX, dV, age,
[0, 0, 0, 0, 0, 0, 10., 5, uniform_age],
])
starcount = 100
background_density = 1e-9
ncomps = sphere_comp_pars.shape[0]
# true_memb_probs = np.zeros((starcount, ncomps))
# true_memb_probs[:,0] = 1.
synth_data = SynthData(
pars=sphere_comp_pars,
starcounts=[starcount],
Components=SphereComponent,
background_density=background_density,
)
synth_data.synthesise_everything()
tabletool.convert_table_astro2cart(synth_data.table)
background_count = len(synth_data.table) - starcount
# insert background densities
synth_data.table['background_log_overlap'] =\
len(synth_data.table) * [np.log(background_density)]
synth_data.table.write(data_filename, overwrite=True)
origins = [SphereComponent(pars) for pars in sphere_comp_pars]
best_comps, med_and_spans, memb_probs = \
expectmax.fit_many_comps(data=synth_data.table, ncomps=ncomps,
rdir=savedir, burnin=10, sampling_steps=10,
trace_orbit_func=dummy_trace_orbit_func,
use_background=True, ignore_stable_comps=False,
max_em_iterations=200)
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_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_multiple_synth_components():
"""Check initialising with multiple components works"""
age = 1e-10
dx = 5.
dv = 2.
ass_pars1 = np.array([10, 20, 30, 40, 50, 60, dx, dv, age])
comp1 = SphereComponent(ass_pars1)
ass_pars2 = np.array([0., 0., 0, 0, 0, 0, dx, dv, age])
comp2 = SphereComponent(ass_pars2)
starcounts = [100, 100]
try:
synth_data = SynthData(pars=[ass_pars1, ass_pars2],
starcounts=starcounts[0],
Components=SphereComponent)
raise UserWarning('AssertionError should have been thrown by synthdata')
except AssertionError:
pass
synth_data = SynthData(pars=[ass_pars1, ass_pars2],
starcounts=starcounts,
Components=SphereComponent)
synth_data.synthesise_everything()
assert len(synth_data.table) == np.sum(starcounts)
means = tabletool.build_data_dict_from_table(
synth_data.table,
main_colnames=[el+'0' for el in 'xyzuvw'],
only_means=True
)
assert np.allclose(comp2.get_mean(), means[starcounts[0]:].mean(axis=0),
atol=2.)
assert np.allclose(comp1.get_mean(), means[:starcounts[0]].mean(axis=0),
atol=2.)
def test_swigImplementation():
"""
Compares the swigged c implementation against the python one in
likelihood.py
"""
true_comp_mean = np.zeros(6)
true_comp_dx = 2.
true_comp_dv = 2.
true_comp_covmatrix = np.identity(6)
true_comp_covmatrix[:3,:3] *= true_comp_dx**2
true_comp_covmatrix[3:,3:] *= true_comp_dv**2
true_comp_age = 1e-10
true_comp = SphereComponent(attributes={
'mean':true_comp_mean,
'covmatrix':true_comp_covmatrix,
'age':true_comp_age,
})
nstars = 100
synth_data = SynthData(pars=true_comp.get_pars(), starcounts=nstars)
synth_data.synthesise_everything()
tabletool.convert_table_astro2cart(synth_data.table)
star_data = tabletool.build_data_dict_from_table(synth_data.table)
p_lnos = p_lno(true_comp.get_covmatrix(), true_comp.get_mean(),
star_data['covs'], star_data['means'])
c_lnos = c_lno(true_comp.get_covmatrix(), true_comp.get_mean(),
star_data['covs'], star_data['means'], nstars)
assert np.allclose(p_lnos, c_lnos)
assert np.isfinite(p_lnos).all()
assert np.isfinite(c_lnos).all()
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_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 test_get_lnoverlaps():
"""
Confirms that star-component overlaps get smaller as stars get further
away.
First generates a component `sphere_comp`. Then generates three stars.
The first one is identical to `sphere_comp` in mean and covmatrix.
The other two share the same covmatrix yet are separated in X.
We check that the overlap integral is smaller for the more separated
stars.
"""
dim = 6
mean = np.zeros(dim)
covmatrix = np.identity(dim)
age = 1e-10
sphere_comp = SphereComponent(attributes={
'mean': mean,
'covmatrix': covmatrix,
'age': age,
})
dx_offsets = [0., 1., 10.]
star_comps = []
for dx_offset in dx_offsets:
star = SphereComponent(
attributes={
'mean':
sphere_comp.get_mean() +
np.array([dx_offset, 0., 0., 0., 0., 0.]),
'covmatrix':
sphere_comp.get_covmatrix(),
'age':
sphere_comp.get_age(),
})
star_comps.append(star)
nstars = len(star_comps)
dummy_table = Table(data=np.arange(nstars).reshape(nstars, 1),
names=['name'])
tabletool.append_cart_cols_to_table(dummy_table)
for star_comp, row in zip(star_comps, dummy_table):
tabletool.insert_data_into_row(
row,
star_comp.get_mean(),
star_comp.get_covmatrix(),
cartesian=True,
)
dummy_data = tabletool.build_data_dict_from_table(dummy_table)
ln_overlaps = likelihood.get_lnoverlaps(sphere_comp, data=dummy_data)
# Checks that ln_overlaps is descending
assert np.allclose(ln_overlaps, sorted(ln_overlaps)[::-1])
def test_maximisation_gradient_descent_with_multiprocessing_tech():
"""
Added by MZ 2020 - 07 - 13
Test if maximisation works when using gradient descent and multiprocessing.
NOTE: this is not a test if maximisation returns appropriate results but
it only tests if the code runs withour errors. This is mainly to test
multiprocessing.
"""
age = 1e-5
ass_pars1 = np.array([0, 0, 0, 0, 0, 0, 5., 2., age])
comp1 = SphereComponent(ass_pars1)
starcounts = [100,]
synth_data = SynthData(pars=[ass_pars1,],
starcounts=starcounts)
synth_data.synthesise_everything()
tabletool.convert_table_astro2cart(synth_data.table)
true_memb_probs = np.zeros((np.sum(starcounts), 1))
true_memb_probs[:starcounts[0], 0] = 1.
#~ true_memb_probs[starcounts[0]:, 1] = 1.
ncomps = len(starcounts)
noise = np.random.rand(ass_pars1.shape[0])*5
all_init_pars = [ass_pars1 + noise]
new_comps, all_samples, _, all_init_pos, success_mask =\
expectmax.maximisation(synth_data.table, ncomps,
true_memb_probs, 100, 'iter00',
all_init_pars,
optimisation_method='Nelder-Mead',
nprocess_ncomp=True,
)
def test_pythonFuncs():
"""
TODO: remove the requirements of file, have data stored in file?
"""
true_comp_mean = np.zeros(6)
true_comp_dx = 2.
true_comp_dv = 2.
true_comp_covmatrix = np.identity(6)
true_comp_covmatrix[:3, :3] *= true_comp_dx**2
true_comp_covmatrix[3:, 3:] *= true_comp_dv**2
true_comp_age = 1e-10
true_comp = SphereComponent(
attributes={
'mean': true_comp_mean,
'covmatrix': true_comp_covmatrix,
'age': true_comp_age,
})
nstars = 100
synth_data = SynthData(pars=true_comp.get_pars(), starcounts=nstars)
synth_data.synthesise_everything()
tabletool.convert_table_astro2cart(synth_data.table)
star_data = tabletool.build_data_dict_from_table(synth_data.table)
# star_data['means'] = star_data['means']
# star_data['covs'] = star_data['covs']
group_mean = true_comp.get_mean()
group_cov = true_comp.get_covmatrix()
# Test overlap with true component
co1s = []
co2s = []
for i, (scov, smn) in enumerate(zip(star_data['covs'],
star_data['means'])):
co1s.append(co1(group_cov, group_mean, scov, smn))
co2s.append(co2(group_cov, group_mean, scov, smn))
co1s = np.array(co1s)
co2s = np.array(co2s)
co3s = np.exp(
p_lno(group_cov, group_mean, star_data['covs'], star_data['means']))
assert np.allclose(co1s, co2s)
assert np.allclose(co2s, co3s)
assert np.allclose(co1s, co3s)
# Test overlap with neighbouring star (with the aim of testing
# tiny overlap values). Note that most overlaps go to 0, but the
# log overlaps retain the information
co1s = []
co2s = []
for i, (scov, smn) in enumerate(zip(star_data['covs'],
star_data['means'])):
co1s.append(
co1(star_data['covs'][15], star_data['means'][15], scov, smn))
co2s.append(
co2(star_data['covs'][15], star_data['means'][15], scov, smn))
co1s = np.array(co1s)
co2s = np.array(co2s)
lnos = p_lno(star_data['covs'][15], star_data['means'][15],
star_data['covs'], star_data['means'])
co3s = np.exp(lnos)
assert np.allclose(co1s, co2s)
assert np.allclose(co2s, co3s)
assert np.allclose(co1s, co3s)
import matplotlib.pyplot as plt
import numpy as np
import sys
sys.path.insert(0, '..')
from chronostar.component import SphereComponent
from chronostar import tabletool
from chronostar import likelihood
from chronostar import expectmax
component_file = '../results/all_nonbg_scocen_comps.npy'
membership_file = '../results/all_scocen_total_membership.npy'
joined_table = '../data/scocen/joined_scocen_no_duplicates.fit'
star_pars = tabletool.build_data_dict_from_table(joined_table, historical=True)
all_comps = SphereComponent.load_raw_components(component_file)
init_z = np.load(membership_file)
# pop manually determined duplicates
if True:
all_comps.pop(9)
all_comps.pop(6)
init_z = init_z[(np.array([0, 1, 2, 3, 4, 5, 7, 8]), )]
print(len(all_comps))
print(len(init_z))
init_z.shape = (1, -1)
memberships = expectmax.expectation(star_pars,
all_comps,
old_memb_probs=init_z)
def test_fit_stability_mixed_comps():
"""
Have a fit with some iterations that have a mix of stable and
unstable comps.
TODO: Maybe give 2 similar comps tiny age but overlapping origins
"""
run_name = 'mixed_stability'
logging.info(60 * '-')
logging.info(15 * '-' + '{:^30}'.format('TEST: ' + run_name) + 15 * '-')
logging.info(60 * '-')
savedir = 'temp_data/{}_expectmax_{}/'.format(PY_VERS, run_name)
mkpath(savedir)
data_filename = savedir + '{}_expectmax_{}_data.fits'.format(
PY_VERS, run_name)
log_filename = 'temp_data/{}_expectmax_{}/log.log'.format(
PY_VERS, run_name)
logging.basicConfig(level=logging.INFO,
filemode='w',
filename=log_filename)
shared_cd_mean = np.zeros(6)
tiny_age = 0.1
medium_age = 10.
# origin_1 = traceorbit.trace_cartesian_orbit(shared_cd_mean, times=-medium_age)
# origin_2 = traceorbit.trace_cartesian_orbit(shared_cd_mean, times=-2*medium_age)
#
# cd_mean_3 = np.array([-200,200,0,0,50,0.])
# origin_3 = traceorbit.trace_cartesian_orbit(cd_mean_3, times=-tiny_age)
#
# sphere_comp_pars = np.array([
# # X, Y, Z, U, V, W, dX, dV, age,
# np.hstack((origin_1, 10., 5., medium_age)), # Next two comps share a current day origin
# np.hstack((origin_2, 10., 5., 2*medium_age)), # so hopefully will need several iterations to\
# # disentangle
# np.hstack((origin_3, 10., 5., tiny_age)), # a distinct comp that is stable quickly
# ])
uniform_age = 1e-10
sphere_comp_pars = np.array([
# X, Y, Z, U, V, W, dX, dV, age,
[50, 0, 0, 0, 50, 0, 10., 5,
uniform_age], # Very distant (and stable) comp
[0, -20, 0, 0, -5, 0, 10., 5, uniform_age], # Overlapping comp 1
[0, 20, 0, 0, 5, 0, 10., 5, uniform_age], # Overlapping comp 2
])
starcounts = [50, 100, 200]
ncomps = sphere_comp_pars.shape[0]
# initialise z appropriately
true_memb_probs = np.zeros((np.sum(starcounts), ncomps))
start = 0
for i in range(ncomps):
true_memb_probs[start:start + starcounts[i], i] = 1.0
start += starcounts[i]
# Initialise some random membership probablities
# which will serve as our starting guess
init_memb_probs = np.random.rand(np.sum(starcounts), ncomps)
# To aid a component in quickly becoming stable, initialse the memberships
# correclty for stars belonging to this component
init_memb_probs[:starcounts[0]] = 0.
init_memb_probs[:starcounts[0], 0] = 1.
init_memb_probs[starcounts[0]:, 0] = 0.
# Normalising such that each row sums to 1
init_memb_probs = (init_memb_probs.T / init_memb_probs.sum(axis=1)).T
synth_data = SynthData(
pars=sphere_comp_pars,
starcounts=starcounts,
Components=SphereComponent,
)
synth_data.synthesise_everything()
tabletool.convert_table_astro2cart(synth_data.table,
write_table=True,
filename=data_filename)
origins = [SphereComponent(pars) for pars in sphere_comp_pars]
SphereComponent.store_raw_components(savedir + 'origins.npy', origins)
best_comps, med_and_spans, memb_probs = \
expectmax.fit_many_comps(data=synth_data.table, ncomps=ncomps,
rdir=savedir, init_memb_probs=init_memb_probs,
trace_orbit_func=dummy_trace_orbit_func,
ignore_stable_comps=True)
perm = expectmax.get_best_permutation(memb_probs, true_memb_probs)
logging.info('Best permutation is: {}'.format(perm))
# Calculate the membership difference, we divide by 2 since
# incorrectly allocated stars are double counted
total_diff = 0.5 * np.sum(np.abs(true_memb_probs - memb_probs[:, perm]))
# Assert that expected membership is less than 10%
assert total_diff < 0.1 * np.sum(starcounts)
for origin, best_comp in zip(origins, np.array(best_comps)[perm, ]):
assert (isinstance(origin, SphereComponent)
and isinstance(best_comp, SphereComponent))
o_pars = origin.get_pars()
b_pars = best_comp.get_pars()
logging.info("origin pars: {}".format(o_pars))
logging.info("best fit pars: {}".format(b_pars))
assert np.allclose(origin.get_mean(), best_comp.get_mean(), atol=5.)
assert np.allclose(origin.get_sphere_dx(),
best_comp.get_sphere_dx(),
atol=2.)
assert np.allclose(origin.get_sphere_dv(),
best_comp.get_sphere_dv(),
atol=2.)
assert np.allclose(origin.get_age(), best_comp.get_age(), atol=1.)
def test_fit_many_comps():
"""
Synthesise a file with negligible error, retrieve initial
parameters
Takes a while... maybe this belongs in integration unit_tests
"""
run_name = 'stationary'
logging.info(60 * '-')
logging.info(15 * '-' + '{:^30}'.format('TEST: ' + run_name) + 15 * '-')
logging.info(60 * '-')
savedir = 'temp_data/{}_expectmax_{}/'.format(PY_VERS, run_name)
mkpath(savedir)
data_filename = savedir + '{}_expectmax_{}_data.fits'.format(
PY_VERS, run_name)
log_filename = 'temp_data/{}_expectmax_{}/log.log'.format(
PY_VERS, run_name)
logging.basicConfig(level=logging.INFO,
filemode='w',
filename=log_filename)
uniform_age = 1e-10
sphere_comp_pars = np.array([
# X, Y, Z, U, V, W, dX, dV, age,
[-50, -50, -50, 0, 0, 0, 10., 5, uniform_age],
[50, 50, 50, 0, 0, 0, 10., 5, uniform_age],
])
starcounts = [20, 50]
ncomps = sphere_comp_pars.shape[0]
# initialise z appropriately
true_memb_probs = np.zeros((np.sum(starcounts), ncomps))
start = 0
for i in range(ncomps):
true_memb_probs[start:start + starcounts[i], i] = 1.0
start += starcounts[i]
# Initialise some random membership probablities
# Normalising such that each row sums to 1
init_memb_probs = np.random.rand(np.sum(starcounts), ncomps)
init_memb_probs = (init_memb_probs.T / init_memb_probs.sum(axis=1)).T
synth_data = SynthData(
pars=sphere_comp_pars,
starcounts=starcounts,
Components=SphereComponent,
)
synth_data.synthesise_everything()
tabletool.convert_table_astro2cart(synth_data.table,
write_table=True,
filename=data_filename)
origins = [SphereComponent(pars) for pars in sphere_comp_pars]
best_comps, med_and_spans, memb_probs = \
expectmax.fit_many_comps(data=synth_data.table, ncomps=ncomps,
rdir=savedir, init_memb_probs=init_memb_probs,
trace_orbit_func=dummy_trace_orbit_func,
ignore_stable_comps=False)
perm = expectmax.get_best_permutation(memb_probs, true_memb_probs)
logging.info('Best permutation is: {}'.format(perm))
assert np.allclose(true_memb_probs, memb_probs[:, perm])
for origin, best_comp in zip(origins, np.array(best_comps)[perm, ]):
assert (isinstance(origin, SphereComponent)
and isinstance(best_comp, SphereComponent))
o_pars = origin.get_pars()
b_pars = best_comp.get_pars()
logging.info("origin pars: {}".format(o_pars))
logging.info("best fit pars: {}".format(b_pars))
assert np.allclose(origin.get_mean(), best_comp.get_mean(), atol=5.)
assert np.allclose(origin.get_sphere_dx(),
best_comp.get_sphere_dx(),
atol=2.)
assert np.allclose(origin.get_sphere_dv(),
best_comp.get_sphere_dv(),
atol=2.)
assert np.allclose(origin.get_age(), best_comp.get_age(), atol=1.)
def test_fit_one_comp_with_background():
"""
Synthesise a file with negligible error, retrieve initial
parameters
Takes a while...
Parameters
----------
"""
run_name = 'background'
logging.info(60 * '-')
logging.info(15 * '-' + '{:^30}'.format('TEST: ' + run_name) + 15 * '-')
logging.info(60 * '-')
savedir = 'temp_data/{}_expectmax_{}/'.format(PY_VERS, run_name)
mkpath(savedir)
data_filename = savedir + '{}_expectmax_{}_data.fits'.format(
PY_VERS, run_name)
log_filename = 'temp_data/{}_expectmax_{}/log.log'.format(
PY_VERS, run_name)
logging.basicConfig(level=logging.INFO,
filemode='w',
filename=log_filename)
uniform_age = 1e-10
sphere_comp_pars = np.array([
# X, Y, Z, U, V, W, dX, dV, age,
[0, 0, 0, 0, 0, 0, 10., 5, uniform_age],
])
starcount = 200
background_density = 1e-9
ncomps = sphere_comp_pars.shape[0]
# true_memb_probs = np.zeros((starcount, ncomps))
# true_memb_probs[:,0] = 1.
synth_data = SynthData(
pars=sphere_comp_pars,
starcounts=[starcount],
Components=SphereComponent,
background_density=background_density,
)
synth_data.synthesise_everything()
tabletool.convert_table_astro2cart(synth_data.table)
background_count = len(synth_data.table) - starcount
logging.info('Generated {} background stars'.format(background_count))
# insert background densities
synth_data.table['background_log_overlap'] =\
len(synth_data.table) * [np.log(background_density)]
synth_data.table.write(data_filename, overwrite=True)
origins = [SphereComponent(pars) for pars in sphere_comp_pars]
best_comps, med_and_spans, memb_probs = \
expectmax.fit_many_comps(data=synth_data.table, ncomps=ncomps,
rdir=savedir, burnin=500, sampling_steps=5000,
trace_orbit_func=dummy_trace_orbit_func,
use_background=True, ignore_stable_comps=False,
max_em_iterations=200)
# return best_comps, med_and_spans, memb_probs
# Check parameters are close
assert np.allclose(sphere_comp_pars, best_comps[0].get_pars(), atol=1.5)
# Check most assoc members are correctly classified
recovery_count_threshold = 0.95 * starcount
recovery_count_actual = np.sum(memb_probs[:starcount, 0] > 0.5)
assert recovery_count_threshold < recovery_count_actual
# Check most background stars are correctly classified
# Number of bg stars classified as members should be less than 5%
# of all background stars
contamination_count_threshold = 0.05 * len(memb_probs[starcount:])
contamination_count_actual = np.sum(memb_probs[starcount:, 0] > 0.5)
assert contamination_count_threshold > contamination_count_actual
# Check reported membership probabilities are consistent with recovery
# rate (within 5%)
mean_membership_confidence = np.mean(memb_probs[:starcount, 0])
assert np.isclose(recovery_count_actual / starcount,
mean_membership_confidence,
atol=0.05)
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]))
plt.ylabel('{} [{}]'.format(labels[dim2], units[dim2]))
plt.savefig('../plots/simple_plot_{}{}.pdf'.format(labels[dim1],
# Initialise all stars in dataset to be full members of first component
using_bg = config.config.get('include_background_distribution', False)
init_memb_probs = np.zeros((len(data_dict['means']), 1 + using_bg))
init_memb_probs[:, 0] = 1.
store_burnin_chains = config.advanced.get('store_burnin_chains', False)
if store_burnin_chains:
log_message(msg='Storing burnin chains', symbol='-')
# Try and recover any results from previous run
try:
prev_med_and_spans = np.load(run_dir + 'final/' + final_med_and_spans_file)
prev_memb_probs = np.load(run_dir + 'final/' + final_memb_probs_file)
try:
prev_comps = Component.load_raw_components(
str(run_dir + 'final/' + final_comps_file))
# Final comps are there, they just can't be read by current module
# so quickly fit them based on fixed prev membership probabilities
except AttributeError:
logging.info('Component class has been modified, reconstructing '
'from chain')
prev_comps = ncomps * [None]
for i in range(ncomps):
final_cdir = run_dir + 'final/comp{}/'.format(i)
chain = np.load(final_cdir + 'final_chain.npy')
lnprob = np.load(final_cdir + 'final_lnprob.npy')
npars = len(Component.PARAMETER_FORMAT)
best_ix = np.argmax(lnprob)
best_pars = chain.reshape(-1, npars)[best_ix]
prev_comps[i] = Component(emcee_pars=best_pars)
Component.store_raw_components(
'''
Noticed issue with calculating a component's current
day mean. Behaviour of trace_cartesian_orbit varies based
on whether a single age is provided vs a range of times
'''
import numpy as np
import sys
sys.path.insert(0, '..')
from chronostar.component import SphereComponent
from chronostar.traceorbit import trace_cartesian_orbit
bugged_comp = SphereComponent.load_raw_components('bugged_component.npy')[0]
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',
cd_med_and_spans[label] = np.load(chaindir + cd_file_stem)
print('loaded for {}'.format(label))
except IOError:
print('calculating for {}'.format(label))
# Convert chain of sampling origin to corresponding chain of current day
flat_final_chain = np.load(chaindir + 'final_chain.npy').reshape(-1, 9)
nsamples = len(flat_final_chain)
# initialise empty array
current_day_chain = np.zeros((nsamples, len(FreeComponent.PARAMETER_FORMAT)))
# One by one, get equivalent pars of current day
for ix, sample in enumerate(flat_final_chain):
if ix % 100 == 0:
print('{} of {} done'.format(ix, len(flat_final_chain)))
comp = SphereComponent(emcee_pars=sample)
cd_mean, cd_cov = comp.get_currentday_projection()
cd_comp = FreeComponent(attributes={'mean':cd_mean,
'covmatrix':cd_cov,
'age':0.})
current_day_chain[ix] = cd_comp.get_pars()
cd_med_and_spans[label] = compfitter.calc_med_and_span(current_day_chain)
np.save(chaindir + cd_file_stem, cd_med_and_spans[label])
# cd_med_and_spans[label] = np.load(chaindir + 'cd_med_and_span.npy')
# cd_med_and_spans[label][6:8] = np.exp(cd_med_and_spans[label][6:8])
origin_med_and_spans[label] = np.load(rdir + 'final_med_errs.npy')[comp_ix[label]]
origin_med_and_spans[label][6:8] = np.exp(origin_med_and_spans[label][6:8])
zs = np.load(rdir + 'final_membership.npy')
"""
import numpy as np
import sys
sys.path.insert(0, '..')
from chronostar.component import SphereComponent
from chronostar import tabletool
from chronostar import expectmax
from astropy.table import Table, vstack, join
# Create components
#c = np.load('all_nonbg_scocen_comps.npy') # including LCC
#comps = [SphereComponent(pars=x) for x in c]
comps = SphereComponent.load_raw_components('all_nonbg_scocen_comps.npy')
print('components', c.shape)
print('Are there duplicate components?')
#datafile = 'data_table_cartesian_with_bg_ols.fits'
datafile = 'data_table_cartesian_including_tims_stars_with_bg_ols.fits'
# Read Gaia data including both stars with known and missing radial velocities
#datafile = 'data_table_cartesian_100k.fits'
data_table = tabletool.read(datafile)
# This table is masked. Unmask:
data_table = data_table.filled()
print('DATA READ', len(data_table))
historical = 'c_XU' in data_table.colnames
return recent_lims
else:
try:
return [np.min((current_lims[0], recent_lims[0])),
np.max((current_lims[1], recent_lims[1]))]
except:
import pdb; pdb.set_trace()
labels = 'XYZUVW'
units = 3*['pc'] + 3*['km/s']
plt_dir = '../plots/scocen_formation/'
all_comps_file = '../results/all_nonbg_scocen_comps.npy'
# Reading in and tidying components
all_comps = SphereComponent.load_raw_components(all_comps_file)
# # Pop off some manually identified duplicates
# all_comps.pop(9)
# all_comps.pop(6)
print('ages of components: {}'.format([c.get_age() for c in all_comps]))
max_age = np.max([c.get_age() for c in all_comps])
ntimes = 100
times = np.linspace(max_age, 0, ntimes)
#for time in times:
time_ix = ntimes - ntimes/2
logging.basicConfig(filename=rdir + 'log.log', level=logging.INFO)
log_message('Beginning Chronostar run', symbol='_', surround=True)
log_message('Setting up', symbol='.', surround=True)
assert os.access(rdir, os.W_OK)
# ------------------------------------------------------------
# ----- SETTING UP ALL DATA PREP ---------------------------
# ------------------------------------------------------------
# Set up some filename constants (INPUT files) FOR COMPARISON and expectmax
final_comps_file_with_rv = '../data/no_rv_paper/beta_pic_sphere_component.npy'
final_memb_probs_file_with_rv = '../data/no_rv_paper/beta_memb_probs.npy'
bp_comp_with_rv = SphereComponent.load_components(final_comps_file_with_rv)
bp_probs_with_rv = np.load(final_memb_probs_file_with_rv)
print(bp_probs_with_rv)
# Set up some filename constants (OUTPUT files)
final_comps_file = 'final_comps.npy'
final_med_and_spans_file = 'final_med_and_spans.npy'
final_memb_probs_file = 'final_membership.npy'
# TODO: data_loadfile: mask out rows with nans
# First see if a data savefile path has been provided, and if
# so, then just assume this script has already been performed
# and the data prep has already been done
if (config.config['data_savefile'] != ''
and os.path.isfile(config.config['data_savefile'])):
def test_spherecomponent_initialisation():
sphere_comp = SphereComponent(pars=SPHERE_PARS)
assert np.allclose(SPHERE_PARS[:6], sphere_comp._mean)
assert np.allclose(AGE, sphere_comp._age)
assert np.isclose(DX, sphere_comp.get_sphere_dx())
assert np.isclose(DV, sphere_comp.get_sphere_dv())
def test_2comps_and_background():
"""
Synthesise a file with negligible error, retrieve initial
parameters
Takes a while... maybe this belongs in integration unit_tests
Performance of test is a bit tricky to callibrate. Since we are skipping
any temporal evolution for speed reasons, we model two
isotropic Gaussians. Now if these Gaussians are too far apart, NaiveFit
will gravitate to one of the Gaussians during the 1 component fit, and then
struggle to discover the second Gaussian.
If the Gaussians are too close, then both will be characteresied by the
1 component fit, and the BIC will decide two Gaussians components are
overkill.
I think I've addressed this by having the two groups have
large number of stars.
"""
using_bg = True
run_name = '2comps_and_background'
logging.info(60 * '-')
logging.info(15 * '-' + '{:^30}'.format('TEST: ' + run_name) + 15 * '-')
logging.info(60 * '-')
savedir = 'temp_data/{}_naive_{}/'.format(PY_VERS, run_name)
mkpath(savedir)
data_filename = savedir + '{}_naive_{}_data.fits'.format(PY_VERS, run_name)
log_filename = 'temp_data/{}_naive_{}/log.log'.format(PY_VERS, run_name)
logging.basicConfig(level=logging.INFO,
filemode='w',
filename=log_filename)
### INITIALISE SYNTHETIC DATA ###
# DON'T CHANGE THE AGE! BECAUSE THIS TEST DOESN'T USE ANY ORBIT INTEGRATION!!!
# Note: if peaks are too far apart, it will be difficult for
# chronostar to identify the 2nd when moving from a 1-component
# to a 2-component fit.
uniform_age = 1e-10
sphere_comp_pars = np.array([
# X, Y, Z, U, V, W, dX, dV, age,
[0, 0, 0, 0, 0, 0, 10., 5, uniform_age],
[30, 0, 0, 0, 5, 0, 10., 5, uniform_age],
])
starcounts = [100, 150]
ncomps = sphere_comp_pars.shape[0]
nstars = np.sum(starcounts)
background_density = 1e-9
# initialise z appropriately
true_memb_probs = np.zeros((np.sum(starcounts), ncomps))
start = 0
for i in range(ncomps):
true_memb_probs[start:start + starcounts[i], i] = 1.0
start += starcounts[i]
try:
# Check if the synth data has already been constructed
data_dict = tabletool.build_data_dict_from_table(data_filename)
except:
synth_data = SynthData(
pars=sphere_comp_pars,
starcounts=starcounts,
Components=SphereComponent,
background_density=background_density,
)
synth_data.synthesise_everything()
tabletool.convert_table_astro2cart(synth_data.table,
write_table=True,
filename=data_filename)
background_count = len(synth_data.table) - np.sum(starcounts)
# insert background densities
synth_data.table['background_log_overlap'] =\
len(synth_data.table) * [np.log(background_density)]
synth_data.table.write(data_filename, overwrite=True)
origins = [SphereComponent(pars) for pars in sphere_comp_pars]
### SET UP PARAMETER FILE ###
fit_pars = {
'results_dir': savedir,
'data_table': data_filename,
'trace_orbit_func': 'dummy_trace_orbit_func',
'return_results': True,
'par_log_file': savedir + 'fit_pars.log',
'overwrite_prev_run': True,
# 'nthreads':18,
'nthreads': 3,
}
### INITIALISE AND RUN A NAIVE FIT ###
naivefit = NaiveFit(fit_pars=fit_pars)
result, score = naivefit.run_fit()
best_comps = result['comps']
memb_probs = result['memb_probs']
# Check membership has ncomps + 1 (bg) columns
n_fitted_comps = memb_probs.shape[-1] - 1
assert ncomps == n_fitted_comps
### CHECK RESULT ###
# No guarantee of order, so check if result is permutated
# also we drop the bg memberships for permutation reasons
perm = expectmax.get_best_permutation(memb_probs[:nstars, :ncomps],
true_memb_probs)
memb_probs = memb_probs[:nstars]
logging.info('Best permutation is: {}'.format(perm))
n_misclassified_stars = np.sum(
np.abs(true_memb_probs - np.round(memb_probs[:, perm])))
# Check fewer than 15% of association stars are misclassified
try:
assert n_misclassified_stars / nstars * 100 < 15
except AssertionError:
import pdb
pdb.set_trace()
for origin, best_comp in zip(origins, np.array(best_comps)[perm, ]):
assert (isinstance(origin, SphereComponent)
and isinstance(best_comp, SphereComponent))
o_pars = origin.get_pars()
b_pars = best_comp.get_pars()
logging.info("origin pars: {}".format(o_pars))
logging.info("best fit pars: {}".format(b_pars))
assert np.allclose(origin.get_mean(), best_comp.get_mean(), atol=5.)
assert np.allclose(origin.get_sphere_dx(),
best_comp.get_sphere_dx(),
atol=2.5)
assert np.allclose(origin.get_sphere_dv(),
best_comp.get_sphere_dv(),
atol=2.5)
assert np.allclose(origin.get_age(), best_comp.get_age(), atol=1.)
log_message('Beginning Chronostar run',
symbol='_', surround=True)
log_message('Setting up', symbol='.', surround=True)
assert os.access(rdir, os.W_OK)
# ------------------------------------------------------------
# ----- SETTING UP ALL DATA PREP ---------------------------
# ------------------------------------------------------------
# Set up some filename constants (INPUT files) FOR COMPARISON and expectmax
final_comps_file_with_rv = '../data/no_rv_paper/beta_pic_sphere_component.npy'
final_memb_probs_file_with_rv = '../data/no_rv_paper/beta_memb_probs.npy'
bp_comp_with_rv = SphereComponent.load_components(final_comps_file_with_rv)
bp_probs_with_rv = np.load(final_memb_probs_file_with_rv)
print(bp_probs_with_rv)
# Set up some filename constants (OUTPUT files)
final_comps_file = 'final_comps.npy'
final_med_and_spans_file = 'final_med_and_spans.npy'
final_memb_probs_file = 'final_membership.npy'
# TODO: data_loadfile: mask out rows with nans
# First see if a data savefile path has been provided, and if
# so, then just assume this script has already been performed
# and the data prep has already been done
if (config.config['data_savefile'] != '' and
os.path.isfile(config.config['data_savefile'])):
historical = 'c_XU' in data_table.colnames
############################################################################
############ COMPONENT OVERLAPS ############################################
############################################################################
print('Create data dict')
# Create data dict
data_dict = tabletool.build_data_dict_from_table(
data_table,
get_background_overlaps=True,
historical=historical,
)
# Create components
comps = SphereComponent.load_raw_components(comps_filename)
# COMPONENT OVERLAPS
overlaps = expectmax.get_all_lnoverlaps(data_dict, comps)
print('overlaps.shape', overlaps.shape, len(comps))
# MEMBERSHIP PROBABILITIES
membership_probabilities = np.array(
[expectmax.calc_membership_probs(ol) for ol in overlaps])
# Create a table
for i in range(membership_probabilities.shape[1] - 1):
data_table['membership_%d' % (i + 1)] = membership_probabilities[:, i]
data_table['membership_bg'] = membership_probabilities[:, -1]
# Print data
def test_spherecomponent_initialisation():
sphere_comp = SphereComponent(pars=SPHERE_PARS)
assert np.allclose(SPHERE_PARS[:6], sphere_comp._mean)
assert np.allclose(AGE, sphere_comp._age)
assert np.isclose(DX, sphere_comp.get_sphere_dx())
assert np.isclose(DV, sphere_comp.get_sphere_dv())
h=1. * 10.**expon,
args=(comp.get_age(), )))
return expons[np.where([
np.allclose(cov_mat, ref_cov_now, rtol=rtol, atol=atol)
for cov_mat in covs_now
])]
mean = np.zeros(6) # centre at LSR
mean = np.array([20., -80., 25., -1.9, 11.76, 2.25])
dx = 10.
dv = 2.
age = 30.
pars = np.hstack((mean, dx, dv, age))
comp = SphereComponent(pars=pars)
results = {}
import matplotlib.pyplot as plt
for label, rtol in zip(['e-4', 'e-3', 'e-2', 'e-1'], [1e-4, 1e-3, 1e-2, 1e-1]):
all_stable_expons = []
lo_age = 0
hi_age = 100
lo_expon = -10
hi_expon = 2
for age in range(lo_age, hi_age):
pars = np.hstack((mean, dx, dv, age))
stable_expons = get_stable_expons(SphereComponent(pars),
rtol=rtol,
# Fit the first component
log_message(msg='FITTING {} COMPONENT'.format(ncomps),
symbol='*', surround=True)
run_dir = rdir + '{}/'.format(ncomps)
# Initialise all stars in dataset to be full members of first component
init_memb_probs = np.zeros((len(data_dict['means']),2))
init_memb_probs[:,0] = 1.
# Try and recover any results from previous run
try:
prev_med_and_spans = np.load(run_dir + 'final/'
+ final_med_and_spans_file)
prev_memb_probs = np.load(run_dir + 'final/' + final_memb_probs_file)
try:
prev_comps = Component.load_components(
str(run_dir+'final/'+final_comps_file))
# Final comps are there, they just can't be read by current module
# so quickly fit them based on fixed prev membership probabilities
except AttributeError:
logging.info('Component class has been modified, reconstructing '
'from chain')
prev_comps = ncomps * [None]
for i in range(ncomps):
final_cdir = run_dir + 'final/comp{}/'.format(i)
chain = np.load(final_cdir + 'final_chain.npy')
lnprob = np.load(final_cdir + 'final_lnprob.npy')
npars = len(Component.PARAMETER_FORMAT)
best_ix = np.argmax(lnprob)
best_pars = chain.reshape(-1,npars)[best_ix]
prev_comps[i] = Component(emcee_pars=best_pars)
np.save(str(run_dir+'final/'+final_comps_file), prev_comps)