def stop():
# calculate the new mean and cov
new_mean = trace_forward(mean, age)
new_cov = tf.transform_covmatrix(cov, trace_forward, mean, dim=6, args=(age,))
new_eigvals = np.linalg.eigvalsh(new_cov)
estimated_cov = np.cov(new_stars.T)
estimated_eigvals = np.linalg.eigvalsh(estimated_cov)
if plotit:
#plt.clf()
plt.plot(new_stars[:,0], new_stars[:,1], 'r.')
#plt.hist2d(stars[:,0], stars[:,1], bins=20)
ymin, ymax = plt.ylim()
xmin, xmax = plt.xlim()
upper = max(xmax, ymax)
lower = min(xmin, ymin)
plt.xlim(upper, lower)
plt.ylim(lower, upper)
chronostar.fitplotter.plotCovEllipse(
new_cov[:2,:2], new_mean, color='r', alpha=0.1
)
plt.savefig("temp_plots/trace_forward{}.png".format(cnt))
plt.show()
assert np.allclose(new_eigvals, estimated_eigvals, rtol=.5)
def stop():
# calculate the new mean and cov
new_mean = trace_forward(mean, age)
new_cov = tf.transform_covmatrix(cov,
trace_forward,
mean,
dim=6,
args=(age, ))
new_eigvals = np.linalg.eigvalsh(new_cov)
estimated_cov = np.cov(new_stars.T)
estimated_eigvals = np.linalg.eigvalsh(estimated_cov)
if plotit:
#plt.clf()
plt.plot(new_stars[:, 0], new_stars[:, 1], 'r.')
#plt.hist2d(stars[:,0], stars[:,1], bins=20)
ymin, ymax = plt.ylim()
xmin, xmax = plt.xlim()
upper = max(xmax, ymax)
lower = min(xmin, ymin)
plt.xlim(upper, lower)
plt.ylim(lower, upper)
chronostar.fitplotter.plotCovEllipse(new_cov[:2, :2],
new_mean,
color='r',
alpha=0.1)
plt.savefig("temp_plots/trace_forward{}.png".format(cnt))
plt.show()
assert np.allclose(new_eigvals, estimated_eigvals, rtol=.5)
def dataGatherer(res_dir='', save_dir='', data_dir='', xyzuvw_file='',
title='', file_stem=''):
"""
Provided with a results directory, tries to find all she needs, then
plots
Parameters
----------
"""
covs = {}
means = {}
star_pars = {}
chain_file = res_dir + "final_chain.npy"
lnprob_file = res_dir + "final_lnprob.npy"
origin_file = res_dir + "origins.npy"
if not xyzuvw_file:
logging.info("No xyzuvw filename provided. Must be synth fit yes?")
xyzuvw_file = res_dir + "xyzuvw_now.fits"
chain = np.load(chain_file)
chain = np.array([chain])
lnprob = np.load(lnprob_file)
best_sample = dt.getBestSample(chain, lnprob)
best_group = chronostar.component.Component(best_sample,
form=len(best_sample) == 9,
internal=True)
star_pars['xyzuvw'] = fits.getdata(xyzuvw_file, 1)
star_pars['xyzuvw_cov'] = fits.getdata(xyzuvw_file, 2)
try:
origins = np.load(origin_file).item()
means['origin_then'] = np.array([origins.mean])
covs['origin_then'] = np.array([origins.generateCovMatrix()])
except IOError:
logging.info("No origins file: {}".format(origin_file))
means['fitted_then'] = np.array([best_group.mean])
means['fitted_now'] =\
np.array([torb.trace_cartesian_orbit(best_group.mean, best_group.age)])
covs['fitted_then'] = np.array([best_group.generateCovMatrix()])
covs['fitted_now'] =\
np.array([
tf.transform_covmatrix(covs['fitted_then'][0], torb.trace_cartesian_orbit,
means['fitted_then'][0],
args=(best_group.age,True)
)
])
plot_hexplot(star_pars, means, covs, chain, iter_count=0,
save_dir=save_dir, file_stem=file_stem, title=title)
plotXYandZW(star_pars, means, covs, chain, iter_count=0,
save_dir=save_dir, file_stem=file_stem, title=title)
def get_stable_expons(comp, rtol=1e-2, atol=1e-6, lo_exp=-10, hi_exp=2):
expons = np.arange(lo_exp, hi_exp)
# ref_cov_now = comp.get_covmatrix_now()
ref_cov_now = transform.transform_covmatrix(comp.get_covmatrix(),
comp.trace_orbit_func,
loc=comp.get_mean(),
h=1. * 10.**-3,
args=(comp.get_age(), ))
covs_now = []
for expon in expons:
covs_now.append(
transform.transform_covmatrix(comp.get_covmatrix(),
comp.trace_orbit_func,
loc=comp.get_mean(),
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
])]
def lnlike(pars, star_pars, z=None, return_lnols=False):
"""Computes the log-likelihood for a fit to a group.
The emcee parameters encode the modelled origin point of the stars. Using
the parameters, a mean and covariance in 6D space are constructed as well as
an age. The kinematics are then projected forward to the current age and
compared with the current stars' XYZUVW values (and uncertainties)
Parameters
----------
pars
Parameters describing the group model being fitted
star_pars
traceback data being fitted to
(retired) z
array of weights [0.0 - 1.0] for each star, describing how likely
they are members of group to be fitted.
Returns
-------
lnlike
the logarithm of the likelihood of the fit
"""
# convert pars into covariance matrix
mean_then = pars[0:6]
cov_then = generate_cov(pars)
age = pars[8]
mean_now = trace_forward(mean_then, age)
cov_now = tf.transform_covmatrix(cov_then,
trace_forward,
mean_then,
dim=6,
args=(age, ))
star_covs = star_pars['xyzuvw_cov'][:, 0]
star_mns = star_pars['xyzuvw'][:, 0]
nstars = star_mns.shape[0]
lnols = get_lnoverlaps(cov_now, mean_now, star_covs, star_mns, nstars)
if return_lnols:
return lnols
# prior on covariance matrix incorporated into parametrisation of dX and dV
return np.sum(lnols * z)
def lnlike(pars, star_pars, z=None, return_lnols=False):
"""Computes the log-likelihood for a fit to a group.
The emcee parameters encode the modelled origin point of the stars. Using
the parameters, a mean and covariance in 6D space are constructed as well as
an age. The kinematics are then projected forward to the current age and
compared with the current stars' XYZUVW values (and uncertainties)
Parameters
----------
pars
Parameters describing the group model being fitted
star_pars
traceback data being fitted to
(retired) z
array of weights [0.0 - 1.0] for each star, describing how likely
they are members of group to be fitted.
Returns
-------
lnlike
the logarithm of the likelihood of the fit
"""
# convert pars into covariance matrix
mean_then = pars[0:6]
cov_then = generate_cov(pars)
age = pars[8]
mean_now = trace_forward(mean_then, age)
cov_now = tf.transform_covmatrix(
cov_then, trace_forward, mean_then, dim=6, args=(age,)
)
star_covs = star_pars['xyzuvw_cov'][:,0]
star_mns = star_pars['xyzuvw'][:,0]
nstars = star_mns.shape[0]
lnols = get_lnoverlaps(
cov_now, mean_now, star_covs, star_mns, nstars
)
if return_lnols:
return lnols
# prior on covariance matrix incorporated into parametrisation of dX and dV
return np.sum(lnols * z)
def test_polar():
"""Generates a collection of points in polar space from a covariance matrix
Transforms covariance matrix to cartesian, and compares to a numpy
covariance matrix fitted to the points in cartesian space"""
logging.basicConfig(level=logging.DEBUG, stream=sys.stdout)
logging.info("Beginning polar demo")
# setting up polar points
r_mean = np.sqrt(10**2 + 10**2)
theta_mean = 14 * np.pi / 6.
r_std = 1.
theta_std = np.pi / 24.
C_rt = 0.0
pol_mean = np.array([r_mean, theta_mean])
pol_cov = np.array([
[r_std**2, C_rt * r_std * theta_std],
[C_rt * r_std * theta_std, theta_std**2],
])
nsamples = 100000
pol_samples = np.random.multivariate_normal(pol_mean, pol_cov, nsamples)
# converting to cartesian manually for comparison
res = convertManyPolarToCartesian(pol_samples) #[:, 0], pol_samples[:, 1])
cart_mean = convertPolarToCartesian(pol_mean)
logging.debug("Covariance matrix:\n{}".format(pol_cov))
logging.debug("Polar mean: {}".format(pol_mean))
cart_cov = tf.transform_covmatrix(pol_cov,
convertPolarToCartesian,
pol_mean,
dim=2)
cart_samples = np.random.multivariate_normal(cart_mean, cart_cov, nsamples)
estimated_mean = np.mean(res, axis=0)
estimated_cov = np.cov(res.T)
assert np.allclose(estimated_mean, cart_mean, rtol=1e-1)
assert np.allclose(estimated_cov, cart_cov, rtol=1e-1)
def test_polar():
"""Generates a collection of points in polar space from a covariance matrix
Transforms covariance matrix to cartesian, and compares to a numpy
covariance matrix fitted to the points in cartesian space"""
logging.basicConfig(level=logging.DEBUG, stream=sys.stdout)
logging.info("Beginning polar demo")
# setting up polar points
r_mean = np.sqrt(10 ** 2 + 10 ** 2)
theta_mean = 14 * np.pi / 6.
r_std = 1.
theta_std = np.pi / 24.
C_rt = 0.0
pol_mean = np.array([r_mean, theta_mean])
pol_cov = np.array([
[r_std ** 2, C_rt * r_std * theta_std],
[C_rt * r_std * theta_std, theta_std ** 2],
])
nsamples = 100000
pol_samples = np.random.multivariate_normal(pol_mean, pol_cov, nsamples)
# converting to cartesian manually for comparison
res = convertManyPolarToCartesian(pol_samples) #[:, 0], pol_samples[:, 1])
cart_mean = convertPolarToCartesian(pol_mean)
logging.debug("Covariance matrix:\n{}".format(pol_cov))
logging.debug("Polar mean: {}".format(pol_mean))
cart_cov = tf.transform_covmatrix(pol_cov, convertPolarToCartesian,
pol_mean, dim=2)
cart_samples = np.random.multivariate_normal(cart_mean, cart_cov, nsamples)
estimated_mean = np.mean(res, axis=0)
estimated_cov = np.cov(res.T)
assert np.allclose(estimated_mean, cart_mean, rtol=1e-1)
assert np.allclose(estimated_cov, cart_cov, rtol=1e-1)
def transformAstrCovsToCartesian(astr_covs, astr_arr):
"""
Converts a covariance matrix from astrometric coords to LSR XYZUVW
Parameters
----------
astr_covs: ([nstars, 6, 6] array)
values in the diagaonal are the squared errors of
(ra, dec, plx, pm_ra, pm_dec, rv), with the offdiagonals the product
of the correlation (valued between -1 and 1) and the two
intersecting coordinates.
astr_arr: ([nstars, 6] array)
the measured (mean) astrometric values
(ra, dec, plx, pm_ra, pm-dec, rv)
"""
nstars = astr_arr.shape[0]
xyzuvw_covs = np.zeros((nstars, 6, 6))
for ix in range(nstars):
xyzuvw_covs[ix] = transform.transform_covmatrix(
astr_covs[ix], coordinate.convert_astrometry2lsrxyzuvw, astr_arr[ix],
dim=6
)
return xyzuvw_covs
def convertRowToCartesian(row, row_ix=None, nrows=None):
dim=6
cart_col_names = ['X', 'Y', 'Z', 'U', 'V', 'W',
'dX', 'dY', 'dZ', 'dU', 'dV', 'dW',
'c_XY', 'c_XZ', 'c_XU', 'c_XV', 'c_XW',
'c_YZ', 'c_YU', 'c_YV', 'c_YW',
'c_ZU', 'c_ZV', 'c_ZW',
'c_UV', 'c_UW',
'c_VW']
try:
if row_ix % 100 == 0:
print("{:010.7f}% done".format(row_ix / float(nrows) * 100.))
except TypeError:
pass
astr_mean, astr_cov = datatool.convertRecToArray(row)
xyzuvw_mean = coordinate.convert_astrometry2lsrxyzuvw(astr_mean)
xyzuvw_cov = transform.transform_covmatrix(
astr_cov,
coordinate.convert_astrometry2lsrxyzuvw,
astr_mean,
dim=dim,
)
# fill in cartesian mean
for col_ix, col_name in enumerate(cart_col_names[:6]):
row[col_name] = xyzuvw_mean[col_ix]
# fill in standard deviations
xyzuvw_stds = np.sqrt(xyzuvw_cov[np.diag_indices(dim)])
for col_ix, col_name in enumerate(cart_col_names[6:12]):
row[col_name] = xyzuvw_stds[col_ix]
correl_matrix = xyzuvw_cov / xyzuvw_stds / xyzuvw_stds.reshape(6, 1)
# fill in correlations
for col_ix, col_name in enumerate(cart_col_names[12:]):
row[col_name] = correl_matrix[
np.triu_indices(dim, k=1)[0][col_ix],
np.triu_indices(dim, k=1)[1][col_ix]
]
print("With {:.2} stars accounted for by background"\
.format(np.sum(z[:,-1])))
means = {}
means['fitted_then'] = [g.mean for g in groups]
means['fitted_now'] = [
torb.trace_cartesian_orbit(g.mean, g.age, single_age=True)
for g in groups
]
covs = {}
covs['fitted_then'] = np.array([
g.generateSphericalCovMatrix() for g in groups
])
covs['fitted_now'] = np.array([
tf.transform_covmatrix(covs['fitted_then'][0], torb.trace_cartesian_orbit,
means['fitted_then'][0],
args=(g.age,True)
)
for g in groups
])
hp.plotNewQuad(star_pars, means, covs, None, 'final',
save_dir='temp_plots/',
file_stem='XY-UV_ncomps_{}_'.format(ncomps), z=z,
title="BIC: {}, likelihood: {}".format(bic, overall_lnlike),
top_dims=(0,1), bot_dims=(3,4))
hp.plotNewQuad(star_pars, means, covs, None, 'final',
save_dir='temp_plots/',
file_stem='XU-YV_ncomps_{}_'.format(ncomps), z=z,
title="BIC: {}, likelihood: {}".format(bic, overall_lnlike),
top_dims=(0,3), bot_dims=(1,4))
try:
cd_chains[label] = np.load(chaindir + 'current_day_chain.npy')
except IOError:
origin_chain = np.load(chaindir + 'final_chain.npy').reshape(-1,9)
npars = 6 + 6 + 15 # + 1 (don't need the age in current day fit)
current_day_chain = np.zeros((origin_chain.shape[0],npars))
for sample_ix, sample in enumerate(origin_chain):
if sample_ix % 100 == 0:
print("Done {:6} of {}".format(sample_ix, current_day_chain.shape[0]))
mean = sample[:6]
group = chronostar.component.Component(sample, internal=True)
cd_mean = torb.trace_cartesian_orbit(group.mean, group.age)
cd_cov = tf.transform_covmatrix(group.generateSphericalCovMatrix(), torb.trace_cartesian_orbit,
group.mean, args=(group.age, True))
# fill in cartesian mean
current_day_chain[sample_ix,0:6] = cd_mean
# fill in standard deviations
cd_stds = np.sqrt(cd_cov[np.diag_indices(DIM)])
current_day_chain[sample_ix,6:12] = cd_stds
correl_matrix = cd_cov / cd_stds / cd_stds.reshape(DIM, 1)
# fill in correlations
for col_ix in range(15):
current_day_chain[sample_ix,12+col_ix] = correl_matrix[
np.triu_indices(DIM, k=1)[0][col_ix],
np.triu_indices(DIM, k=1)[1][col_ix]
]
print("BIC is: {}".format(bic))
print("With {:.2} stars accounted for by background"\
.format(np.sum(z[:,-1])))
means = {}
means['fitted_then'] = [g.mean for g in groups]
means['fitted_now'] = [
torb.trace_cartesian_orbit(g.mean, g.age, single_age=True)
for g in groups
]
covs = {}
covs['fitted_then'] = np.array(
[g.generateSphericalCovMatrix() for g in groups])
covs['fitted_now'] = np.array([
tf.transform_covmatrix(covs['fitted_then'][0],
torb.trace_cartesian_orbit,
means['fitted_then'][0],
args=(g.age, True)) for g in groups
])
hp.plotNewQuad(star_pars,
means,
covs,
None,
'final',
save_dir='temp_plots/',
file_stem='XY-UV_ncomps_{}_'.format(ncomps),
z=z,
title="BIC: {}, likelihood: {}".format(
bic, overall_lnlike),
top_dims=(0, 1),
bot_dims=(3, 4))
astro_savefile)
astr_arr, err_arr = ms.convertTableToArray(star_table)
nstars = len(star_table)
astr_covs = cv.convertAstroErrsToCovs(err_arr)
xyzuvw_now = cc.convert_many_astrometry2lsrxyzuvw(astr_arr, mas=True)
logging.info("Mean of retrieved stars: {}".format(np.mean(xyzuvw_now, axis=0)))
if plot_it:
plt.clf()
plt.plot(xyzuvw_now[:, 1], xyzuvw_now[:, 2], '.')
chronostar.fitplotter.plotCovEllipse(
group.generateEllipticalCovMatrix()[1:3, 1:3],
group.mean[1:3],
with_line=True)
xyzuvw_covs = np.zeros((nstars, 6, 6))
for ix in range(nstars):
xyzuvw_covs[ix] = tf.transform_covmatrix(astr_covs[ix],
cc.convert_astrometry2lsrxyzuvw,
astr_arr[ix],
dim=6)
if plot_it:
for ix in range(nstars):
chronostar.fitplotter.plotCovEllipse(xyzuvw_covs[ix, 1:3, 1:3],
xyzuvw_now[ix, 1:3],
color='blue')
plt.savefig("temp_plots/demo_run.png")
except IOError:
origin_chain = np.load(chaindir + 'final_chain.npy').reshape(-1, 9)
npars = 6 + 6 + 15 # + 1 (don't need the age in current day fit)
current_day_chain = np.zeros((origin_chain.shape[0], npars))
for sample_ix, sample in enumerate(origin_chain):
if sample_ix % 100 == 0:
print("Done {:6} of {}".format(sample_ix,
current_day_chain.shape[0]))
mean = sample[:6]
group = chronostar.component.Component(sample, internal=True)
cd_mean = torb.trace_cartesian_orbit(group.mean, group.age)
cd_cov = tf.transform_covmatrix(group.generateSphericalCovMatrix(),
torb.trace_cartesian_orbit,
group.mean,
args=(group.age, True))
# fill in cartesian mean
current_day_chain[sample_ix, 0:6] = cd_mean
# fill in standard deviations
cd_stds = np.sqrt(cd_cov[np.diag_indices(DIM)])
current_day_chain[sample_ix, 6:12] = cd_stds
correl_matrix = cd_cov / cd_stds / cd_stds.reshape(DIM, 1)
# fill in correlations
for col_ix in range(15):
current_day_chain[sample_ix, 12 + col_ix] = correl_matrix[
np.triu_indices(DIM, k=1)[0][col_ix],
np.triu_indices(DIM, k=1)[1][col_ix]]
def polar_demo():
"""
Demonstrating the use of the transform module in covariance transformations
between coordinate systems
"""
logging.info("Beginning polar demo")
plotit = True
# setting up polar points
r_mean = np.sqrt(10**2 + 10**2)
theta_mean = 14 * np.pi / 6.
r_std = 1.
theta_std = np.pi / 24.
C_rt = 0.0
pol_mean = np.array([r_mean, theta_mean])
pol_cov = np.array([
[r_std**2, C_rt * r_std * theta_std],
[C_rt * r_std * theta_std, theta_std**2],
])
nsamples = 100000
pol_samples = np.random.multivariate_normal(pol_mean, pol_cov, nsamples)
# converting to cartesian manually for comparison
res = transform_ptcs(pol_samples[:, 0], pol_samples[:, 1])
cart_mean = transform_ptc(pol_mean)
cart_cov = tf.transform_covmatrix(pol_cov, transform_ptc, pol_mean, dim=2)
cart_samples = np.random.multivariate_normal(cart_mean, cart_cov, nsamples)
estimated_mean = np.mean(res, axis=0)
estimated_cov = np.cov(res.T)
# plotting
if plotit:
nbins = 100
plt.clf()
plt.hist2d(res[:, 0],
res[:, 1],
bins=nbins,
range=((-20, 20), (-20, 20)))
plt.savefig("temp_plots/jac_hist.png")
plt.clf()
plt.plot(res[:, 0], res[:, 1], '.')
plt.ylim(-20, 20)
plt.xlim(-20, 20)
plt.savefig("temp_plots/jac.png")
plt.clf()
plt.plot(cart_samples[:, 0], cart_samples[:, 1], '.')
plt.ylim(-20, 20)
plt.xlim(-20, 20)
plt.savefig("temp_plots/jac_cart.png")
plt.clf()
plt.hist2d(cart_samples[:, 0],
cart_samples[:, 1],
bins=nbins,
range=((-20, 20), (-20, 20)))
plt.savefig("temp_plots/jac_cart_hist.png")
# tolerance is so high because dealing with non-linearity and
assert np.allclose(estimated_mean, cart_mean, rtol=1e-1)
assert np.allclose(estimated_cov, cart_cov, rtol=1e-1)
logging.info("Ending polar demo")
def dataGathererEM(ngroups,
iter_count,
res_dir='',
save_dir='',
data_dir='',
xyzuvw_file='',
title='',
file_stem='',
groups_file=''):
"""
Provided with a results directory, tries to find all she needs, then
plots
Parameters
----------
ngroups: int
number of groups
"""
covs = {}
means = {}
star_pars = {}
if not groups_file:
groups_file = "best_group_fit.npy"
chain_file = "final_chain.npy"
# lnprob_file = "final_lnprob.npy"
# origin_file = res_dir + "origins.npy"
if not xyzuvw_file:
logging.info("No xyzuvw filename provided. Must be synth fit yes?")
xyzuvw_file = res_dir + "../xyzuvw_now.fits"
try:
print("Trying to load from {}".format(xyzuvw_file))
star_pars['xyzuvw'] = fits.getdata(xyzuvw_file, 1)
star_pars['xyzuvw_cov'] = fits.getdata(xyzuvw_file, 2)
print("Successfully loaded from xyzuvw_file")
except:
import chronostar.retired.groupfitter as rgf
old_star_pars = rgf.read_stars(res_dir + "../perf_tb_file.pkl")
star_pars = {
'xyzuvw': old_star_pars['xyzuvw'][:, 0],
'xyzuvw_cov': old_star_pars['xyzuvw_cov'][:, 0]
}
try:
origins = np.load(res_dir + '../origins.npy')
except:
origins = None
z = np.load(res_dir + '../memberships.npy')
fitted_then_mns = []
fitted_then_covs = []
fitted_now_mns = []
fitted_now_covs = []
origin_then_mns = []
origin_then_covs = []
all_chains = []
for group_ix in range(ngroups):
gdir = res_dir + "group{}/".format(group_ix)
chain = np.load(gdir + chain_file)
all_chains.append(chain)
best_group = np.load(gdir + groups_file).item()
fitted_then_mns.append(best_group.mean)
fitted_then_covs.append(best_group.generateCovMatrix())
fitted_now_mn = torb.trace_cartesian_orbit(fitted_then_mns[group_ix],
best_group.age,
single_age=True)
fitted_now_cov =\
tf.transform_covmatrix(fitted_then_covs[group_ix],
torb.trace_cartesian_orbit,
fitted_then_mns[group_ix],
args=(best_group.age,))
fitted_now_mns.append(fitted_now_mn)
fitted_now_covs.append(fitted_now_cov)
if origins is not None:
origin_then_mns.append(origins[group_ix].mean)
origin_then_covs.append(origins[group_ix].generateCovMatrix())
means = {
'origin_then': origin_then_mns,
'fitted_then': fitted_then_mns,
'fitted_now': fitted_now_mns,
}
covs = {
'origin_then': np.array(origin_then_covs),
'fitted_then': np.array(fitted_then_covs),
'fitted_now': np.array(fitted_now_covs),
}
plotNewQuad(star_pars,
means,
covs,
all_chains,
iter_count=iter_count,
save_dir=save_dir,
file_stem=file_stem,
title=title,
z=z)
plotNewHex(star_pars,
means,
covs,
all_chains,
iter_count=iter_count,
save_dir=save_dir,
file_stem=file_stem,
title=title,
z=z)
def dataGathererEM(ngroups, iter_count, res_dir='', save_dir='', data_dir='',
xyzuvw_file='', title='', file_stem='', groups_file=''):
"""
Provided with a results directory, tries to find all she needs, then
plots
Parameters
----------
ngroups: int
number of groups
"""
covs = {}
means = {}
star_pars = {}
if not groups_file:
groups_file = "best_group_fit.npy"
chain_file = "final_chain.npy"
# lnprob_file = "final_lnprob.npy"
# origin_file = res_dir + "origins.npy"
if not xyzuvw_file:
logging.info("No xyzuvw filename provided. Must be synth fit yes?")
xyzuvw_file = res_dir + "../xyzuvw_now.fits"
try:
print("Trying to load from {}".format(xyzuvw_file))
star_pars['xyzuvw'] = fits.getdata(xyzuvw_file, 1)
star_pars['xyzuvw_cov'] = fits.getdata(xyzuvw_file, 2)
print("Successfully loaded from xyzuvw_file")
except:
import chronostar.retired.groupfitter as rgf
old_star_pars = rgf.read_stars(res_dir + "../perf_tb_file.pkl")
star_pars = {'xyzuvw':old_star_pars['xyzuvw'][:,0],
'xyzuvw_cov':old_star_pars['xyzuvw_cov'][:,0]}
try:
origins = np.load(res_dir + '../origins.npy')
except:
origins = None
z = np.load(res_dir + '../memberships.npy')
fitted_then_mns = []
fitted_then_covs = []
fitted_now_mns = []
fitted_now_covs = []
origin_then_mns = []
origin_then_covs = []
all_chains = []
for group_ix in range(ngroups):
gdir = res_dir + "group{}/".format(group_ix)
chain = np.load(gdir + chain_file)
all_chains.append(chain)
best_group = np.load(gdir + groups_file).item()
fitted_then_mns.append(best_group.mean)
fitted_then_covs.append(best_group.generateCovMatrix())
fitted_now_mn = torb.trace_cartesian_orbit(fitted_then_mns[group_ix],
best_group.age,
single_age=True)
fitted_now_cov =\
tf.transform_covmatrix(fitted_then_covs[group_ix],
torb.trace_cartesian_orbit,
fitted_then_mns[group_ix],
args=(best_group.age,))
fitted_now_mns.append(fitted_now_mn)
fitted_now_covs.append(fitted_now_cov)
if origins is not None:
origin_then_mns.append(origins[group_ix].mean)
origin_then_covs.append(origins[group_ix].generateCovMatrix())
means = {
'origin_then':origin_then_mns,
'fitted_then':fitted_then_mns,
'fitted_now':fitted_now_mns,
}
covs = {
'origin_then':np.array(origin_then_covs),
'fitted_then':np.array(fitted_then_covs),
'fitted_now':np.array(fitted_now_covs),
}
plotNewQuad(star_pars, means, covs, all_chains, iter_count=iter_count,
save_dir=save_dir, file_stem=file_stem, title=title,
z=z)
plotNewHex(star_pars, means, covs, all_chains, iter_count=iter_count,
save_dir=save_dir, file_stem=file_stem, title=title,
z=z)
def calc_mns_covs(new_gps, ngroups, origins=None):
all_origin_mn_then = [None] * ngroups
all_origin_cov_then = [None] * ngroups
all_origin_mn_now = [None] * ngroups
all_origin_cov_now = [None] * ngroups
all_fitted_mn_then = [None] * ngroups
all_fitted_cov_then = [None] * ngroups
all_fitted_mn_now = [None] * ngroups
all_fitted_cov_now = [None] * ngroups
for i in range(ngroups):
if origins:
all_origin_mn_then[i] = origins[i][:6]
all_origin_cov_then[i] = utils.generate_cov(
utils.internalise_pars(origins[i])
)
all_origin_mn_now[i] = tb.trace_forward(all_origin_mn_then[i],
origins[i][-2])
all_origin_cov_now[i] = tf.transform_covmatrix(
all_origin_cov_then[i], tb.trace_forward, all_origin_mn_then[i],
dim=6, args=(origins[i][-2],)
)
all_fitted_mn_then[i] = new_gps[i][:6]
all_fitted_cov_then[i] = tfgf.generate_cov(new_gps[i])
all_fitted_mn_now[i] = tb.trace_forward(all_fitted_mn_then[i],
new_gps[i][-1])
all_fitted_cov_now[i] = tf.transform_covmatrix(
all_fitted_cov_then[i], tb.trace_forward, all_fitted_mn_then[i],
dim=6, args=(new_gps[i][-1],)
)
if origins:
all_origin_mn_then = np.array(all_origin_mn_then )
all_origin_cov_then = np.array(all_origin_cov_then)
all_origin_mn_now = np.array(all_origin_mn_now )
all_origin_cov_now = np.array(all_origin_cov_now )
all_fitted_mn_then = np.array(all_fitted_mn_then )
all_fitted_cov_then = np.array(all_fitted_cov_then)
all_fitted_mn_now = np.array(all_fitted_mn_now )
all_fitted_cov_now = np.array(all_fitted_cov_now )
all_means = {
'fitted_then' : all_fitted_mn_then,
'fitted_now' : all_fitted_mn_now ,
}
all_covs = {
'fitted_then' : all_fitted_cov_then,
'fitted_now' : all_fitted_cov_now ,
}
if origins:
all_means['origin_then'] = all_origin_mn_then
all_means['origin_now'] = all_origin_mn_now
all_covs['origin_then'] = all_origin_cov_then
all_covs['origin_now'] = all_origin_cov_now
np.save("means.npy", all_means)
np.save("covs.npy", all_covs)
return all_means, all_covs
def traceforward_group(mean, cov, age):
"""
Trace a group's Gaussian description from some time in the past
to present day.
:param mean: [6] array;
the mean XYZUVW value
:param cov: [6,6] array;
Covariance matrix
:param age: float;
Age of "birth" of the group
:return:
- current_mean: [6] array;
the mean phase position at current day
- current_cov: [6,6] array;
the cov matrix at current day
"""
new_mean = trace_forward(mean, age)
new_cov = tf.transform_covmatrix(cov,
trace_forward,
mean,
dim=6,
args=(age, ))
return new_mean, new_cov
# ------ everything in this block is rubbish I think -------
#def transform_cov_mat(xyzuvw, xyzuvw_cov, tstep):
# dp = 1e-3
# J = np.zeros((6,6))
#
# for coord in range(6):
# xyzuvw_plus = np.array(map(np.float, xyzuvw.copy()))
# xyzuvw_neg = np.array(map(np.float, xyzuvw.copy()))
# xyzuvw_plus[coord] += dp
# xyzuvw_neg[coord] -= dp
#
# # filling out the coordth column of the jacobian matrix
# J[:, coord] =\
# (trace_forward(xyzuvw_plus, tstep, solarmotion=None)
# - trace_forward(xyzuvw_neg, tstep, solarmotion=None)) / (2*dp)
#
## J[:, coord] =\
## (trace_forward(xyzuvw_plus, tstep, solarmotion=None)
## - trace_forward(xyzuvw, tstep, solarmotion=None)) / dp
#
# # pdb.set_trace()
#
# new_cov = np.dot(J, np.dot(xyzuvw_cov, J.T))
#
# return new_cov
#
#def trace_forward_group(xyzuvw, xyzuvw_cov, time, nsteps):
# """Take a group's origin in galactic coordinates and project forward
#
# Parameters
# ----------
# xyzuvw : [6] float array
# initial mean
# xyzuvw_cov : [6,6] float array
# initial covariance matrix
# time : float
# age of group in Myr
# nsteps : int
# number of steps to iterate through. The transformation of the
# covariance matrix is only valid over a linear regime, so need to
# restrict the step size in order to maintain accuracy
#
# Returns
# -------
# xyzuvw_current : [6] float array
# current age xyzuvw mean
# xyzuvw_cov_current : [6,6] float array
# current age covariance matrix
# """
tstep = time / nsteps
old_xyzuvw = xyzuvw
old_cov = xyzuvw_cov
for i in range(nsteps):
new_cov = transform_cov_mat(old_xyzuvw, old_cov,
tstep) # do we need current xyzuvw too?
new_xyzuvw = trace_forward(old_xyzuvw, tstep, solarmotion=None)
old_xyzuvw = new_xyzuvw
old_cov = new_cov
return new_xyzuvw, new_cov
origin_file = rdir + 'origins.npy'
perf_xyzuvw_file = rdir + 'perf_xyzuvw.npy'
memberships_file = rdir + 'memberships.npy'
final_groups_file = rdir + 'final_groups.npy'
xyzuvw_file = rdir + 'xyzuvw_now.fits'
origins = np.load(origin_file)
perf_xyzuvw_now = np.load(perf_xyzuvw_file)
z_final = np.load(memberships_file)
final_groups = np.load(final_groups_file)
xyzuvw_dict = gf.loadXYZUVW(xyzuvw_file)
ass_cov_now = tf.transform_covmatrix(final_groups[0].generateCovMatrix(),
torb.trace_cartesian_orbit,
final_groups[0].mean,
args=(final_groups[0].age,))
ass_mn_now = torb.trace_cartesian_orbit(final_groups[0].mean,
final_groups[0].age)
simple_cov_now = np.copy(ass_cov_now)
simple_cov_now[3:6,:3] = 0
simple_cov_now[:3,3:6] = 0
members_mask = np.arange(50)
field_mask = np.arange(50,1050)
nstars = z_final.shape[0]
precs = ['perf', 'half', 'gaia', 'double']
prec_val = {'perf':1e-5, 'half':0.5, 'gaia':1.0, 'double':2.0}
prec = 'gaia'
def dataGatherer(res_dir='',
save_dir='',
data_dir='',
xyzuvw_file='',
title='',
file_stem=''):
"""
Provided with a results directory, tries to find all she needs, then
plots
Parameters
----------
"""
covs = {}
means = {}
star_pars = {}
chain_file = res_dir + "final_chain.npy"
lnprob_file = res_dir + "final_lnprob.npy"
origin_file = res_dir + "origins.npy"
if not xyzuvw_file:
logging.info("No xyzuvw filename provided. Must be synth fit yes?")
xyzuvw_file = res_dir + "xyzuvw_now.fits"
chain = np.load(chain_file)
chain = np.array([chain])
lnprob = np.load(lnprob_file)
best_sample = dt.getBestSample(chain, lnprob)
best_group = chronostar.component.Component(best_sample,
form=len(best_sample) == 9,
internal=True)
star_pars['xyzuvw'] = fits.getdata(xyzuvw_file, 1)
star_pars['xyzuvw_cov'] = fits.getdata(xyzuvw_file, 2)
try:
origins = np.load(origin_file).item()
means['origin_then'] = np.array([origins.mean])
covs['origin_then'] = np.array([origins.generateCovMatrix()])
except IOError:
logging.info("No origins file: {}".format(origin_file))
means['fitted_then'] = np.array([best_group.mean])
means['fitted_now'] =\
np.array([torb.trace_cartesian_orbit(best_group.mean, best_group.age)])
covs['fitted_then'] = np.array([best_group.generateCovMatrix()])
covs['fitted_now'] =\
np.array([
tf.transform_covmatrix(covs['fitted_then'][0], torb.trace_cartesian_orbit,
means['fitted_then'][0],
args=(best_group.age,True)
)
])
plot_hexplot(star_pars,
means,
covs,
chain,
iter_count=0,
save_dir=save_dir,
file_stem=file_stem,
title=title)
plotXYandZW(star_pars,
means,
covs,
chain,
iter_count=0,
save_dir=save_dir,
file_stem=file_stem,
title=title)
z_final = np.load(final_membership_file)
bp_hdul = fits.open(bp_xyzuvw_file)
gaia_hdul = fits.open(gaia_astr_file)
bp_xyzuvw = bp_hdul[1].data
bp_xyzuvw_cov = bp_hdul[2].data
bp_core_mask = np.where(z_final[:,0] > 0.75)
bp_core_xyzuvw = bp_xyzuvw[bp_core_mask]
bpmg_then = np.load(final_groups_file)[0]
bpmg_mean_now = torb.trace_cartesian_orbit(bpmg_then.mean, bpmg_then.age,
single_age=True)
bpmg_cov_now = tf.transform_covmatrix(
bpmg_then.generateCovMatrix(), torb.trace_cartesian_orbit,
bpmg_then.mean, dim=6,
args=(bpmg_then.age, True)
)
ngaia_stars = gaia_xyzuvw.shape[0]
try:
gaia_bpmg_evals = np.load(gaia_bpmg_evals_file)
except IOError:
print("Evaluating gaia stars at BPMG current MVGauss distribution")
gaia_bpmg_evals = np.zeros(ngaia_stars)
bpmg_invcov_now = np.linalg.inv(bpmg_cov_now)
for i, gaia_star in enumerate(gaia_xyzuvw):
if (i % 100000 == 0):
print("{:10} of {:10}... {:6.1f}%".format(i, ngaia_stars,
i / ngaia_stars*100))
'c_ZU', 'c_ZV', 'c_ZW', 'c_UV', 'c_UW', 'c_VW'
]
# insert empty columns
for col_name in cart_col_names:
gt[col_name] = empty_col
for row_ix, gt_row in enumerate(gt):
dim = 6
if row_ix % 10 == 0:
print("{:02.2f}% done".format(row_ix / float(nrows) * 100.))
astr_mean, astr_cov = gc.convertRecToArray(gt_row)
xyzuvw_mean = coord.convert_astrometry2lsrxyzuvw(astr_mean)
xyzuvw_cov = tf.transform_covmatrix(
astr_cov,
coord.convert_astrometry2lsrxyzuvw,
astr_mean,
dim=dim,
)
# fill in cartesian mean
for col_ix, col_name in enumerate(cart_col_names[:6]):
gt_row[col_name] = xyzuvw_mean[col_ix]
# fill in standard deviations
xyzuvw_stds = np.sqrt(xyzuvw_cov[np.diag_indices(dim)])
for col_ix, col_name in enumerate(cart_col_names[6:12]):
gt_row[col_name] = xyzuvw_stds[col_ix]
correl_matrix = xyzuvw_cov / xyzuvw_stds / xyzuvw_stds.reshape(6, 1)
# fill in correlations
for col_ix, col_name in enumerate(cart_col_names[12:]):
gt_row[col_name] = correl_matrix[
def traceforward_group(mean, cov, age):
"""
Trace a group's Gaussian description from some time in the past
to present day.
:param mean: [6] array;
the mean XYZUVW value
:param cov: [6,6] array;
Covariance matrix
:param age: float;
Age of "birth" of the group
:return:
- current_mean: [6] array;
the mean phase position at current day
- current_cov: [6,6] array;
the cov matrix at current day
"""
new_mean = trace_forward(mean, age)
new_cov = tf.transform_covmatrix(
cov, trace_forward, mean, dim=6, args=(age,)
)
return new_mean, new_cov
# ------ everything in this block is rubbish I think -------
#def transform_cov_mat(xyzuvw, xyzuvw_cov, tstep):
# dp = 1e-3
# J = np.zeros((6,6))
#
# for coord in range(6):
# xyzuvw_plus = np.array(map(np.float, xyzuvw.copy()))
# xyzuvw_neg = np.array(map(np.float, xyzuvw.copy()))
# xyzuvw_plus[coord] += dp
# xyzuvw_neg[coord] -= dp
#
# # filling out the coordth column of the jacobian matrix
# J[:, coord] =\
# (trace_forward(xyzuvw_plus, tstep, solarmotion=None)
# - trace_forward(xyzuvw_neg, tstep, solarmotion=None)) / (2*dp)
#
## J[:, coord] =\
## (trace_forward(xyzuvw_plus, tstep, solarmotion=None)
## - trace_forward(xyzuvw, tstep, solarmotion=None)) / dp
#
# # pdb.set_trace()
#
# new_cov = np.dot(J, np.dot(xyzuvw_cov, J.T))
#
# return new_cov
#
#def trace_forward_group(xyzuvw, xyzuvw_cov, time, nsteps):
# """Take a group's origin in galactic coordinates and project forward
#
# Parameters
# ----------
# xyzuvw : [6] float array
# initial mean
# xyzuvw_cov : [6,6] float array
# initial covariance matrix
# time : float
# age of group in Myr
# nsteps : int
# number of steps to iterate through. The transformation of the
# covariance matrix is only valid over a linear regime, so need to
# restrict the step size in order to maintain accuracy
#
# Returns
# -------
# xyzuvw_current : [6] float array
# current age xyzuvw mean
# xyzuvw_cov_current : [6,6] float array
# current age covariance matrix
# """
tstep = time / nsteps
old_xyzuvw = xyzuvw
old_cov = xyzuvw_cov
for i in range(nsteps):
new_cov = transform_cov_mat(old_xyzuvw, old_cov, tstep) # do we need current xyzuvw too?
new_xyzuvw = trace_forward(old_xyzuvw, tstep, solarmotion=None)
old_xyzuvw = new_xyzuvw
old_cov = new_cov
return new_xyzuvw, new_cov
def polar_demo():
"""
Demonstrating the use of the transform module in covariance transformations
between coordinate systems
"""
logging.info("Beginning polar demo")
plotit = True
# setting up polar points
r_mean = np.sqrt(10**2+ 10**2)
theta_mean = 14*np.pi / 6.
r_std = 1.
theta_std = np.pi / 24.
C_rt = 0.0
pol_mean = np.array([r_mean, theta_mean])
pol_cov = np.array([
[r_std**2, C_rt * r_std * theta_std],
[C_rt * r_std * theta_std, theta_std**2],
])
nsamples = 100000
pol_samples = np.random.multivariate_normal(pol_mean, pol_cov, nsamples)
# converting to cartesian manually for comparison
res = transform_ptcs(pol_samples[:, 0], pol_samples[:, 1])
cart_mean = transform_ptc(pol_mean)
cart_cov = tf.transform_covmatrix(pol_cov, transform_ptc, pol_mean, dim=2)
cart_samples = np.random.multivariate_normal(cart_mean, cart_cov, nsamples)
estimated_mean = np.mean(res, axis=0)
estimated_cov = np.cov(res.T)
# plotting
if plotit:
nbins = 100
plt.clf()
plt.hist2d(res[:,0],res[:,1],bins=nbins,range=((-20,20),(-20,20)))
plt.savefig("temp_plots/jac_hist.png")
plt.clf()
plt.plot(res[:,0], res[:,1], '.')
plt.ylim(-20,20)
plt.xlim(-20,20)
plt.savefig("temp_plots/jac.png")
plt.clf()
plt.plot(cart_samples[:,0], cart_samples[:,1], '.')
plt.ylim(-20,20)
plt.xlim(-20,20)
plt.savefig("temp_plots/jac_cart.png")
plt.clf()
plt.hist2d(cart_samples[:,0],cart_samples[:,1],bins=nbins,range=((-20,20),(-20,20)))
plt.savefig("temp_plots/jac_cart_hist.png")
# tolerance is so high because dealing with non-linearity and
assert np.allclose(estimated_mean, cart_mean, rtol=1e-1)
assert np.allclose(estimated_cov, cart_cov, rtol=1e-1)
logging.info("Ending polar demo")
xyzuvw_now_true = to.trace_many_cartesian_orbit(xyzuvw_init, np.array([0., group.age]))[:, 1]
#assert np.allclose(np.mean(xyzuvw_now, axis=0), group.mean, rtol=1e-1)
logging.info("Mean of initial stars: {}".format(np.mean(xyzuvw_init, axis=0)))
logging.info("Mean of final stars: {}".format(np.mean(xyzuvw_now_true, axis=0)))
star_table = chronostar.synthdata.measureXYZUVW(xyzuvw_now_true, 20.0, astro_savefile)
astr_arr, err_arr = ms.convertTableToArray(star_table)
nstars = len(star_table)
astr_covs = cv.convertAstroErrsToCovs(err_arr)
xyzuvw_now = cc.convert_many_astrometry2lsrxyzuvw(astr_arr, mas=True)
logging.info("Mean of retrieved stars: {}".format(np.mean(xyzuvw_now, axis=0)))
if plot_it:
plt.clf()
plt.plot(xyzuvw_now[:,1], xyzuvw_now[:,2], '.')
chronostar.fitplotter.plotCovEllipse(group.generateEllipticalCovMatrix()[1:3, 1:3],
group.mean[1:3], with_line=True)
xyzuvw_covs = np.zeros((nstars,6,6))
for ix in range(nstars):
xyzuvw_covs[ix] = tf.transform_covmatrix(
astr_covs[ix], cc.convert_astrometry2lsrxyzuvw, astr_arr[ix], dim=6
)
if plot_it:
for ix in range(nstars):
chronostar.fitplotter.plotCovEllipse(xyzuvw_covs[ix, 1:3, 1:3], xyzuvw_now[ix, 1:3], color='blue')
plt.savefig("temp_plots/demo_run.png")
'c_UV', 'c_UW',
'c_VW']
# insert empty columns
for col_name in cart_col_names:
gt[col_name] = empty_col
for row_ix, gt_row in enumerate(gt):
dim = 6
if row_ix%10 == 0:
print("{:02.2f}% done".format(row_ix / float(nrows) * 100.))
astr_mean, astr_cov = gc.convertRecToArray(gt_row)
xyzuvw_mean = coord.convert_astrometry2lsrxyzuvw(astr_mean)
xyzuvw_cov = tf.transform_covmatrix(
astr_cov,
coord.convert_astrometry2lsrxyzuvw,
astr_mean,
dim=dim,
)
# fill in cartesian mean
for col_ix, col_name in enumerate(cart_col_names[:6]):
gt_row[col_name] = xyzuvw_mean[col_ix]
# fill in standard deviations
xyzuvw_stds = np.sqrt(xyzuvw_cov[np.diag_indices(dim)])
for col_ix, col_name in enumerate(cart_col_names[6:12]):
gt_row[col_name] = xyzuvw_stds[col_ix]
correl_matrix = xyzuvw_cov / xyzuvw_stds / xyzuvw_stds.reshape(6,1)
# fill in correlations
for col_ix, col_name in enumerate(cart_col_names[12:]):
sdir = 'temp_data/'
origin_file = rdir + 'origins.npy'
perf_xyzuvw_file = rdir + 'perf_xyzuvw.npy'
memberships_file = rdir + 'memberships.npy'
final_groups_file = rdir + 'final_groups.npy'
xyzuvw_file = rdir + 'xyzuvw_now.fits'
origins = np.load(origin_file)
perf_xyzuvw_now = np.load(perf_xyzuvw_file)
z_final = np.load(memberships_file)
final_groups = np.load(final_groups_file)
xyzuvw_dict = gf.loadXYZUVW(xyzuvw_file)
ass_cov_now = tf.transform_covmatrix(final_groups[0].generateCovMatrix(),
torb.trace_cartesian_orbit,
final_groups[0].mean,
args=(final_groups[0].age, ))
ass_mn_now = torb.trace_cartesian_orbit(final_groups[0].mean,
final_groups[0].age)
simple_cov_now = np.copy(ass_cov_now)
simple_cov_now[3:6, :3] = 0
simple_cov_now[:3, 3:6] = 0
members_mask = np.arange(50)
field_mask = np.arange(50, 1050)
nstars = z_final.shape[0]
precs = ['perf', 'half', 'gaia', 'double']
prec_val = {'perf': 1e-5, 'half': 0.5, 'gaia': 1.0, 'double': 2.0}
prec = 'gaia'
