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_convertSynthTableToCart():
"""
Checks that current day measured cartesian values (with negligbile
measurement error) match the true current day cartesian values
"""
AGE = 40.
PARS = np.array([
[0., 0., 0., 0., 0., 0., 10., 5., AGE],
])
STARCOUNTS = [50] #, 30]
COMPONENTS = SphereComponent
MEASUREMENT_ERROR = 1e-10
# Generate synthetic data
synth_data = SynthData(pars=PARS, starcounts=STARCOUNTS,
Components=COMPONENTS,
measurement_error=MEASUREMENT_ERROR,
)
synth_data.synthesise_everything()
# Convert (inplace) astrometry to cartesian
tabletool.convert_table_astro2cart(synth_data.table)
# Check consistency between true current-day kinematics and measured
# current-day kinematics (with negliglbe error)
for dim in 'XYZUVW':
dim_now = dim.lower() + '_now'
assert np.allclose(synth_data.table[dim_now],
synth_data.table[dim])
def test_convertAstrTableToCart():
"""
Using a historical table, confirm that cartesian conversion yields
same results by comparing the cartesian means and covariance matrices
are identical.
Gets historical cartesian data from building data from table cart cols.
Gets updated cartesian data from building astro data from table cols,
converting to cartesian (stored back into table) then building data
from newly inserted table cart cols.
"""
# hist_filename = '../data/paper1/historical_beta_Pictoris_with_gaia_small_everything_final.fits'
hist_table = Table.read(HIST_FILE_NAME)
# curr_filename = '../data/paper1/beta_Pictoris_with_gaia_small_everything_final.fits'
curr_table = Table.read(CURR_FILE_NAME)
# Drop stars that have gone through any binary checking
# hist_table = Table(hist_table[100:300])
# curr_table = Table(curr_table[100:300])
# load in original means and covs
orig_cart_data =\
tabletool.build_data_dict_from_table(table=hist_table, cartesian=True,
historical=True)
tabletool.convert_table_astro2cart(table=curr_table, write_table=False)
cart_data = tabletool.build_data_dict_from_table(curr_table,
cartesian=True)
assert np.allclose(orig_cart_data['means'], cart_data['means'])
assert np.allclose(hist_table['dX'], curr_table['X_error'])
assert np.allclose(orig_cart_data['covs'], cart_data['covs'])
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 run_fit_helper(
true_comp,
starcounts,
measurement_error,
burnin_step=None,
run_name='default',
trace_orbit_func=None,
):
py_vers = sys.version[0]
data_filename = 'temp_data/{}_compfitter_{}.fits'.format(py_vers, run_name)
log_filename = 'logs/{}_compfitter_{}.log'.format(py_vers, run_name)
plot_dir = 'temp_plots/{}_compfitter_{}'.format(py_vers, run_name)
save_dir = 'temp_data/'
logging.basicConfig(level=logging.INFO,
filename=log_filename,
filemode='w')
synth_data = SynthData(pars=true_comp.get_pars(),
starcounts=starcounts,
measurement_error=measurement_error)
synth_data.synthesise_everything()
tabletool.convert_table_astro2cart(synth_data.table,
write_table=True,
filename=data_filename)
res = gf.fit_comp(
data=synth_data.table,
plot_it=True,
burnin_steps=burnin_step,
plot_dir=plot_dir,
save_dir=save_dir,
trace_orbit_func=trace_orbit_func,
)
return res
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 run_fit_helper(true_comp, starcounts, measurement_error,
burnin_step=None,
run_name='default',
trace_orbit_func=None,
):
py_vers = sys.version[0]
data_filename = 'temp_data/{}_compfitter_{}.fits'.format(py_vers, run_name)
log_filename = 'logs/{}_compfitter_{}.log'.format(py_vers, run_name)
plot_dir = 'temp_plots/{}_compfitter_{}'.format(py_vers, run_name)
save_dir = 'temp_data/'
logging.basicConfig(level=logging.INFO,
filename=log_filename,
filemode='w')
synth_data = SynthData(pars=true_comp.get_pars(),
starcounts=starcounts,
measurement_error=measurement_error)
synth_data.synthesise_everything()
tabletool.convert_table_astro2cart(synth_data.table,
write_table=True,
filename=data_filename)
res = gf.fit_comp(
data=synth_data.table,
plot_it=True,
burnin_steps=burnin_step,
plot_dir=plot_dir,
save_dir=save_dir,
trace_orbit_func=trace_orbit_func,
)
return res
def test_convertAstrTableToCart():
"""
Using a historical table, confirm that cartesian conversion yields
same results by comparing the cartesian means and covariance matrices
are identical.
Gets historical cartesian data from building data from table cart cols.
Gets updated cartesian data from building astro data from table cols,
converting to cartesian (stored back into table) then building data
from newly inserted table cart cols.
"""
hist_filename = '../data/paper1/historical_beta_Pictoris_with_gaia_small_everything_final.fits'
hist_table = Table.read(hist_filename)
curr_filename = '../data/paper1/beta_Pictoris_with_gaia_small_everything_final.fits'
curr_table = Table.read(curr_filename)
# Drop stars that have gone through any binary checking
hist_table = Table(hist_table[100:300])
curr_table = Table(curr_table[100:300])
# load in original means and covs
orig_cart_data =\
tabletool.build_data_dict_from_table(table=hist_table, cartesian=True,
historical=True)
tabletool.convert_table_astro2cart(table=curr_table, write_table=False)
cart_data = tabletool.build_data_dict_from_table(curr_table, cartesian=True)
assert np.allclose(orig_cart_data['means'], cart_data['means'])
assert np.allclose(hist_table['dX'], curr_table['X_error'])
assert np.allclose(orig_cart_data['covs'], cart_data['covs'])
def test_convertSynthTableToCart():
"""
Checks that current day measured cartesian values (with negligbile
measurement error) match the true current day cartesian values
"""
AGE = 40.
PARS = np.array([
[0., 0., 0., 0., 0., 0., 10., 5., AGE],
])
STARCOUNTS = [50] #, 30]
COMPONENTS = SphereComponent
MEASUREMENT_ERROR = 1e-10
# Generate synthetic data
synth_data = SynthData(
pars=PARS,
starcounts=STARCOUNTS,
Components=COMPONENTS,
measurement_error=MEASUREMENT_ERROR,
)
synth_data.synthesise_everything()
# Convert (inplace) astrometry to cartesian
tabletool.convert_table_astro2cart(synth_data.table)
# Check consistency between true current-day kinematics and measured
# current-day kinematics (with negliglbe error)
for dim in 'XYZUVW':
dim_now = dim.lower() + '_now'
assert np.allclose(synth_data.table[dim_now], synth_data.table[dim])
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 add_UVW_chronostar(tab):
tabletool.convert_table_astro2cart(table=tab,
main_colnames=None,
error_colnames=None,
corr_colnames=None,
return_table=True)
tab.write('ScoCen_box_result_with_kinematics.fits')
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_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_badColNames():
"""
Check that columns have consistent (or absent) units across measurements
and errors
First test comparing column with degrees to column with mas/yr raises
UserWarning
Then test comparing colum with degrees to column without units raises
no issue.
"""
main_colnames, error_colnames, corr_colnames = \
tabletool.get_colnames(cartesian=False)
# main_colnames[5] = 'radial_velocity_best'
# error_colnames[5] = 'radial_velocity_error_best'
# corrupt ordering of column names
corrupted_error_colnames = list(error_colnames)
corrupted_error_colnames[0], corrupted_error_colnames[3] =\
error_colnames[3], error_colnames[0]
filename = '../data/paper1/beta_Pictoris_with_gaia_small_everything_final.fits'
table = Table.read(filename)
# Only need a handful of rows
table = Table(table[:10])
# Catch when units are inconsistent
try:
tabletool.convert_table_astro2cart(
table,
main_colnames=main_colnames,
error_colnames=corrupted_error_colnames,
corr_colnames=corr_colnames
)
except Exception as e:
assert type(e) == exceptions.UserWarning
# In the case where units have not been provided, then just leave it be
try:
error_colnames[0] = 'ra_dec_corr'
tabletool.convert_table_astro2cart(table,
main_colnames=main_colnames,
error_colnames=error_colnames,
corr_colnames=corr_colnames)
except:
assert False
def test_badColNames():
"""
Check that columns have consistent (or absent) units across measurements
and errors
First test comparing column with degrees to column with mas/yr raises
UserWarning
Then test comparing colum with degrees to column without units raises
no issue.
"""
main_colnames, error_colnames, corr_colnames = \
tabletool.get_colnames(cartesian=False)
# main_colnames[5] = 'radial_velocity_best'
# error_colnames[5] = 'radial_velocity_error_best'
# corrupt ordering of column names
corrupted_error_colnames = list(error_colnames)
corrupted_error_colnames[0], corrupted_error_colnames[3] =\
error_colnames[3], error_colnames[0]
# filename = '../data/paper1/beta_Pictoris_with_gaia_small_everything_final.fits'
table = Table.read(CURR_FILE_NAME)
# Only need a handful of rows
table = Table(table[:10])
# Catch when units are inconsistent
try:
tabletool.convert_table_astro2cart(
table,
astr_main_colnames=main_colnames,
astr_error_colnames=corrupted_error_colnames,
astr_corr_colnames=corr_colnames)
except Exception as e:
assert type(e) == exceptions.UserWarning
# In the case where units have not been provided, then just leave it be
try:
error_colnames[0] = 'ra_dec_corr'
tabletool.convert_table_astro2cart(table,
astr_main_colnames=main_colnames,
astr_error_colnames=error_colnames,
astr_corr_colnames=corr_colnames)
except:
assert False
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_get_region():
"""
Test whether get_region applies data cut successfully.
Synthesise two data sets, one which made up of Component A, and the other
made up of Component A & B. Then, check if applying a get_region cut on
the combined data set, with the Component A set as reference, only returns
the Component A stars.
"""
data_a_filename = 'temp_data/test_get_region_A.fits'
synth_dataset_a = synthdata.SynthData(pars=PARS[0],
starcounts=STARCOUNTS[0])
np.random.seed(0)
synth_dataset_a.synthesise_everything(filename=data_a_filename,
overwrite=True)
tabletool.convert_table_astro2cart(synth_dataset_a.table,
write_table=True,
filename=data_a_filename)
data_both_filename = 'temp_data/test_get_region_both.fits'
synth_dataset_both = synthdata.SynthData(pars=PARS, starcounts=STARCOUNTS)
np.random.seed(0)
synth_dataset_both.synthesise_everything(filename=data_both_filename,
overwrite=True)
# Prepare .par file
par_file = 'temp_data/test_get_region.par'
with open(par_file, 'w') as fp:
fp.write('par_log_file = temp_data/test_get_region_pars.log\n')
fp.write('input_file = {}\n'.format(data_both_filename))
fp.write('convert_astrometry = True\n')
fp.write('apply_cart_cuts = True\n')
fp.write('cut_on_region = True\n')
fp.write('cut_ref_table = {}\n'.format(data_a_filename))
# fp.write('output_file = {}\n'.format())
fp.write('return_data_table = True\n')
# Apply datatool to synthetically generated dataset
data_table = datatool.prepare_data(par_file)
assert len(data_table) == len(synth_dataset_a.table)
def add_missing_tims_stars_to_my_set():
# My table
d = Table.read('data_table_cartesian_with_bg_ols.fits')
# Existing Tim's table
usco = Table.read('usco_res/usco_run_subset.fit')
ucl = Table.read('ucl_res/ucl_run_subset.fit')
lcc = Table.read('lcc_res/lcc_run_subset.fit')
tim = vstack([usco, ucl, lcc])
tim_existing = unique(tim, keys='source_id')
tim_missing = Table.read('missing_columns_for_tims_stars.fits')
tim = join(tim_existing, tim_missing, keys='source_id')
print tim.colnames
tim.remove_columns([
'c_VW', 'astrometric_primary_flag', 'c_XZ', 'c_XY', 'c_ZU', 'c_ZV',
'c_ZW', 'c_XV', 'c_XW', 'c_XU', 'c_UW', 'c_UV', 'c_YU', 'c_YW', 'c_YV',
'c_YZ', 'dZ', 'dX', 'dY', 'dV', 'dW', 'dU'
])
both = vstack([d, tim])
print len(both)
both = unique(both, keys='source_id')
print len(both)
# This table is masked. Unmask:
both = both.filled()
tabletool.convert_table_astro2cart(table=both, return_table=True)
# WRITE
both.write('data_table_cartesian_including_tims_stars_with_bg_ols.fits',
format='fits',
overwrite=True)
ct = set(tim.colnames)
cd = set(d.colnames)
print
print 'Colnames in mine and not in Tims set', cd.difference(ct)
print
print 'Colnames in Tims and not in my set', ct.difference(cd)
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_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)
"""
Add very large RV errors for stars with no known RVs.
Convert to cartesian.
"""
import numpy as np
import sys
sys.path.insert(0, '..')
from chronostar import tabletool
from astropy.table import Table
datafile = '../data/ScoCen_box_result.fits'
d = tabletool.read(datafile)
# Set missing radial velocities (nan) to 0
d['radial_velocity'] = np.nan_to_num(d['radial_velocity'])
# Set missing radial velocity errors (nan) to 1e+10
d['radial_velocity_error'][np.isnan(d['radial_velocity_error'])] = 1e+4
print('Convert to cartesian')
tabletool.convert_table_astro2cart(table=d, return_table=True)
d.write(
'/priv/mulga1/marusa/chronostar/data/ScoCen_box_result_15M_ready_for_bg_ols.fits'
)
print('Cartesian written.', len(d))
historical = False
log_message('Data table has {} rows'.format(len(data_table)))
# data_table['radial_velocity'] = data_table['radial_velocity_best']
# data_table['radial_velocity_error'] = data_table['radial_velocity_error_best']
#
# By the end of this, data will be a astropy table
# with cartesian data written in
# columns in default way.
if config.config['convert_to_cartesian']:
print('Converting to cartesian')
# Performs conversion in place (in memory) on `data_table`
tabletool.convert_table_astro2cart(
table=data_table,
main_colnames=config.astro_colnames.get('main_colnames', None),
error_colnames=config.astro_colnames.get('error_colnames', None),
corr_colnames=config.astro_colnames.get('corr_colnames', None),
return_table=True)
# Calculate background overlaps, storing in data
bg_lnol_colname = 'background_log_overlap'
if config.config['include_background_distribution']:
print("Calculating background overlaps")
# Only calculate if missing
if bg_lnol_colname not in data_table.colnames:
log_message('Calculating background densities')
# background_means = tabletool.build_data_dict_from_table(
# config.config['kernel_density_input_datafile'],
# only_means=True,
# )
# star_means = tabletool.build_data_dict_from_table(
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_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'
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 = [200,200]
ncomps = sphere_comp_pars.shape[0]
# initialise z appropriately
# start = 0
# for i in range(ngroups):
# nstars_in_group = int(group_pars[i,-1])
# z[start:start+nstars_in_group,i] = 1.0
# start += nstars_in_group
true_memb_probs = np.zeros((np.sum(starcounts), ncomps))
true_memb_probs[:200,0] = 1.
true_memb_probs[200:,1] = 1.
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,
trace_orbit_func=dummy_trace_orbit_func, )
# compare fit with input
try:
assert np.allclose(true_memb_probs, memb_probs)
except AssertionError:
# If not close, check if flipping component order fixes things
memb_probs = memb_probs[:,::-1]
best_comps = best_comps[::-1]
assert np.allclose(true_memb_probs, memb_probs)
for origin, best_comp in zip(origins, best_comps):
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... maybe this belongs in integration unit_tests
"""
run_name = 'background'
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,
write_table=True,
filename=data_filename)
background_count = len(synth_data.table) - starcount
# insert background densities
synth_data.table['background_log_overlap'] =\
len(synth_data.table) * [np.log(background_density)]
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,
trace_orbit_func=dummy_trace_orbit_func,
use_background=True)
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.)
# Check most assoc members are correctly classified
recovery_count_threshold = 0.95 * starcounts[0]
recovery_count_actual = np.sum(np.round(memb_probs[:starcount,0]))
assert recovery_count_threshold < recovery_count_actual
# Check most background stars are correctly classified
contamination_count_threshold = 0.05 * len(memb_probs[100:])
contamination_count_actual = np.sum(np.round(memb_probs[starcount:,0]))
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/100., mean_membership_confidence,
atol=0.05)
comp_then=True,
comp_now=True,
comp_orbit=True)
plt.xlabel('Y')
plt.ylabel('V')
plt.savefig('../scripts/synthData_plot_of_yv.png')
plt.clf()
xu_pos = means[:, np.array([0, 3])]
print('means')
print(means)
data_filename = '../scripts/synthData_ellip.fits'
tt.convert_table_astro2cart(my_synth_data.table,
write_table=True,
filename=data_filename)
# res = compfitter.fit_comp(
# data=my_synth_data.table,
# plot_it=True,
# burnin_steps=burnin_step,
# plot_dir=plot_dir,
# save_dir=save_dir,
# trace_orbit_func=trace_orbit_func,
# )
# my_table = my_synth_data.table
# print(len(my_table))
# print(len(means))
# print(my_table.colnames)
# my_table['X'] = means[:,0]
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
from chronostar import compfitter
if __name__ == '__main__':
logging.basicConfig(level=logging.INFO, filename='temp_logs/groupfitter.log')
save_dir = 'temp_data/'
group_savefile = save_dir + 'origins_stat.npy'
xyzuvw_init_savefile = save_dir + 'xyzuvw_init_stat.npy'
astro_savefile = save_dir + 'astro_table_stat.txt'
xyzuvw_conv_savefile = save_dir + 'xyzuvw_conv_stat.fits'
pars = np.array([0., 0., 0., 0., 0., 0., 5., 2., 1e-8])
starcount = 100
error_frac = 1.
synth_data = SynthData(pars=pars, starcounts=starcount)
synth_data.synthesise_everything()
tabletool.convert_table_astro2cart(synth_data.table)
data = tabletool.build_data_dict_from_table(synth_data.table)
stat_file = 'stat_dumps/groupfitter.stat'
# best_fit, chain, lnprob = \
cProfile.run(
"groupfitter.fit_comp(data=data, plot_it=True,"
"convergence_tol=2., burnin_steps=400, plot_dir='temp_plots/',"
"save_dir='temp_data/')",
stat_file,
)
stat = pstats.Stats(stat_file)
stat.sort_stats('cumtime')
stat.print_stats(0.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.)
from chronostar import compfitter
if __name__ == '__main__':
logging.basicConfig(level=logging.INFO, filename='compfitter.log')
save_dir = ''
group_savefile = save_dir + 'origins_stat.npy'
xyzuvw_init_savefile = save_dir + 'xyzuvw_init_stat.npy'
astro_savefile = save_dir + 'astro_table_stat.txt'
xyzuvw_conv_savefile = save_dir + 'xyzuvw_conv_stat.fits'
pars = np.array([0., 0., 0., 0., 0., 0., 5., 2., 1e-8])
starcount = 100
error_frac = 1.
synth_data = SynthData(pars=pars, starcounts=starcount)
synth_data.synthesise_everything()
tabletool.convert_table_astro2cart(synth_data.table)
data = tabletool.build_data_dict_from_table(synth_data.table)
stat_file = 'compfitter.stat'
# best_fit, chain, lnprob = \
cProfile.run(
"compfitter.fit_comp(data=data, plot_it=True,"
"convergence_tol=2., burnin_steps=400, plot_dir='',"
"save_dir='')",
stat_file,
)
stat = pstats.Stats(stat_file)
stat.sort_stats('cumtime')
stat.print_stats(0.3)
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_convert_astrometry():
"""
Use a synethetically generated set of astrometry, convert
to cartesian both manually, and via datatool
In order to compare results, the datatool cartesian conversion
will be stored in alternatively named columns
"""
synth_table_filename = 'temp_data/test_convert_astrometry_data.fits'
synth_dataset = synthdata.SynthData(pars=PARS, starcounts=STARCOUNTS)
synth_dataset.synthesise_everything(filename=synth_table_filename,
overwrite=True)
tabletool.convert_table_astro2cart(synth_table_filename,
write_table=True,
filename=synth_table_filename)
# Prepare a pars file
par_file = 'temp_data/test_convert_astrometry.par'
alt_cart_main_colnames = ['{}_alt'.format(dim) for dim in DIMS]
alt_cart_error_colnames = ['{}_error_alt'.format(dim) for dim in DIMS]
alt_cart_corr_colnames = []
for i, colname1 in enumerate(DIMS):
for colname2 in DIMS[i + 1:]:
alt_cart_corr_colnames.append('{}_{}_corr_alt'.format(
colname1, colname2))
with open(par_file, 'w') as fp:
fp.write('par_log_file = temp_data/test_convert_astrometry_pars.log\n')
fp.write('input_file = {}\n'.format(synth_table_filename))
fp.write('convert_astrometry = True\n')
fp.write('{} = {}\n'.format('cart_main_colnames',
alt_cart_main_colnames).replace("'", ''))
fp.write('{} = {}\n'.format('cart_error_colnames',
alt_cart_error_colnames).replace("'", ''))
fp.write('{} = {}\n'.format('cart_corr_colnames',
alt_cart_corr_colnames).replace("'", ''))
fp.write('overwrite_datafile = True\n')
fp.write('output_file = {}\n'.format(synth_table_filename))
fp.write('return_data_table = True\n')
# Apply datatool to synthetically generated dataset
data_table = datatool.prepare_data(par_file)
main_colnames, error_colnames, corr_colnames = tabletool.get_colnames(
cartesian=True)
for orig, alt in zip([main_colnames, error_colnames, corr_colnames], [
alt_cart_main_colnames, alt_cart_error_colnames,
alt_cart_corr_colnames
]):
for orig_colname, alt_colname in zip(orig, alt):
assert np.allclose(data_table[orig_colname],
data_table[alt_colname],
rtol=1e-5)
print(
np.max(
np.abs(data_table[alt_colname] -
data_table[orig_colname])))
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('Data table has {} rows'.format(len(data_table)))
# data_table['radial_velocity'] = data_table['radial_velocity_best']
# data_table['radial_velocity_error'] = data_table['radial_velocity_error_best']
#
# By the end of this, data will be a astropy table
# with cartesian data written in
# columns in default way.
if config.config['convert_to_cartesian']:
print('Converting to cartesian')
# Performs conversion in place (in memory) on `data_table`
tabletool.convert_table_astro2cart(
table=data_table,
main_colnames=config.astro_colnames.get('main_colnames', None),
error_colnames=config.astro_colnames.get('error_colnames', None),
corr_colnames=config.astro_colnames.get('corr_colnames', None),
return_table=True)
# Calculate background overlaps, storing in data
bg_lnol_colname = 'background_log_overlap'
if config.config['include_background_distribution']:
print("Calculating background overlaps")
# Only calculate if missing
if bg_lnol_colname not in data_table.colnames:
log_message('Calculating background densities')
# background_means = tabletool.build_data_dict_from_table(
# config.config['kernel_density_input_datafile'],
# only_means=True,
# )
