def getZfromOrigins(origins, star_pars):
if type(origins) is str:
origins = dt.loadGroups(origins)
if type(star_pars) is str:
star_pars = dt.loadXYZUVW(star_pars)
nstars = star_pars['xyzuvw'].shape[0]
ngroups = len(origins)
nassoc_stars = np.sum([o.nstars for o in origins])
using_bg = nstars != nassoc_stars
z = np.zeros((nstars, ngroups + using_bg))
stars_so_far = 0
# set associaiton members memberships to 1
for i, o in enumerate(origins):
z[stars_so_far:stars_so_far+o.nstars, i] = 1.
stars_so_far += o.nstars
# set remaining stars as members of background
if using_bg:
z[stars_so_far:,-1] = 1.
return z
while ncomps < MAX_COMP:
if ncomps >= MAX_COMP:
logging.info("++++++++++++++++++++++++++++++++++++++++++++++++++")
logging.info("+++++++++++ REACHED MAX COMP LIMIT +++++++++++")
logging.info("++++++++++++++++++++++++++++++++++++++++++++++++++")
# handle special case of one component
if ncomps == 1:
logging.info("******************************************")
logging.info("********* FITTING 1 COMPONENT **********")
logging.info("******************************************")
run_dir = rdir + '{}/'.format(ncomps)
mkpath(run_dir)
try:
new_groups = dt.loadGroups(run_dir + 'final/final_groups.npy')
new_meds = np.load(run_dir + 'final/final_med_errs.npy')
new_z = np.load(run_dir + 'final/final_membership.npy')
logging.info("Loaded from previous run")
except IOError:
new_groups, new_meds, new_z = \
em.fit_many_comps(star_pars, ncomps, rdir=run_dir, pool=pool,
bg_ln_ols=bg_ln_ols) # kill component here
new_groups = np.array(new_groups)
new_lnlike = em.get_overall_lnlikelihood(star_pars, new_groups,
bg_ln_ols=bg_ln_ols)
new_lnpost = em.get_overall_lnlikelihood(star_pars, new_groups,
bg_ln_ols=bg_ln_ols,
inc_posterior=True)
new_BIC = em.calc_bic(star_pars, ncomps, new_lnlike)
import numpy as np
import matplotlib.pyplot as plt
import sys
sys.path.insert(0, '..')
import chronostar.traceorbit as torb
import chronostar.retired2.datatool as dt
rdir = "../results/synth_fit/50_2_1_50/"
# rdir = "../results/synth_fit/30_5_2_100/"
xyzuvw_init_file = rdir + "xyzuvw_init.npy"
xyzuvw_init = np.load(xyzuvw_init_file)
origin = np.load(rdir + 'origins.npy').item()
origin = dt.loadGroups(rdir + 'origins.npy')
max_age = origin.age
ntimes = int(max_age) + 1
#ntimes = 3
times = np.linspace(1e-5, max_age, ntimes)
traceforward = torb.trace_many_cartesian_orbit(xyzuvw_init, times, False)
nstars = xyzuvw_init.shape[0]
def plot_subplot(traceforward, t_ix, dim1, dim2, ax):
flat_tf = traceforward.reshape(-1,6)
# mins = np.min(flat_tf, axis=0)
# maxs = np.max(flat_tf, axis=0)
labels = ['X [pc]', 'Y [pc]', 'Z [pc]',
'U [km/s]', 'V [km/s]', 'W [km/s]']
fitted_z_stem = 'final_membership.npy'
med_errs_stem = 'final_med_errs.npy'
use_bg = False
for fit, order, suffix in zip(fits, orders, suffixs):
print(fit)
rdir = '../results/em_fit/{}'.format(fit) + suffix + '/'
star_pars_file = '../data/{}_xyzuvw.fits'.format(fit)
save_file_name = '../results/tables/{}_table.tex'.format(fit)
origin_file = rdir + origin_stem
fitted_z_file = rdir + fitted_z_stem
med_errs_file = rdir + med_errs_stem
origins = dt.loadGroups(origin_file)
true_z = dt.getZfromOrigins(origins, star_pars_file)
true_nstars = np.sum(true_z, axis=0)
fitted_z = np.load(fitted_z_file)
fitted_nstars = np.sum(fitted_z, axis=0)
med_errs = np.load(med_errs_file)
med_errs[:,6:8] = np.exp(med_errs[:,6:8])
# manually work out ordering of groups
origins = origins[(order,)]
if use_bg:
true_nstars = true_nstars[(order + [-1],)]
else:
true_nstars = true_nstars[(order,)]
# manually inform on presence of background or not
def loadFinalResults(fdir):
fgroups = dt.loadGroups(fdir + 'final_groups.npy')
fmembs = np.load(fdir + 'final_membership.npy')
fmed_errs = np.load(fdir + 'final_med_errs.npy')
return fgroups, fmembs, fmed_errs
import numpy as np
import matplotlib.pyplot as plt
import sys
sys.path.insert(0, '..')
import chronostar.traceorbit as torb
import chronostar.retired2.datatool as dt
rdir = "../results/synth_fit/50_2_1_50/"
# rdir = "../results/synth_fit/30_5_2_100/"
xyzuvw_init_file = rdir + "xyzuvw_init.npy"
xyzuvw_init = np.load(xyzuvw_init_file)
origin = np.load(rdir + 'origins.npy').item()
origin = dt.loadGroups(rdir + 'origins.npy')
max_age = origin.age
ntimes = int(max_age) + 1
#ntimes = 3
times = np.linspace(1e-5, max_age, ntimes)
traceforward = torb.trace_many_cartesian_orbit(xyzuvw_init, times, False)
nstars = xyzuvw_init.shape[0]
def plot_subplot(traceforward, t_ix, dim1, dim2, ax):
flat_tf = traceforward.reshape(-1, 6)
# mins = np.min(flat_tf, axis=0)
# maxs = np.max(flat_tf, axis=0)
labels = ['X [pc]', 'Y [pc]', 'Z [pc]', 'U [km/s]', 'V [km/s]', 'W [km/s]']
import matplotlib.pyplot as plt
import numpy as np
from scipy import stats
import sys
sys.path.insert(0, '..')
import chronostar.retired2.datatool as dt
import chronostar.expectmax as em
from chronostar._overlap import get_lnoverlaps
bpmg_group_file = 'final_groups.npy'
gaia_sep_hist_file = 'bg_hists.npy'
gaia_6d_hist_file = 'gaia_6d_hist.npy'
gaia_xyzuvw = np.load('../data/gaia_dr2_mean_xyzuvw.npy')
bpmg_group = dt.loadGroups(bpmg_group_file)[0]
ref_mean = bpmg_group.mean[:]
# ref_mean[-1] += 4. #SANITY CHECK, perform same analysis for W offset by 4
print("1 comp. fit to BPMG has mean {}".format(ref_mean))
print("We will calculate all densities at this point...")
gaia_sep_hist = np.load(gaia_sep_hist_file)
try:
gaia_6d_hist = np.load(gaia_6d_hist_file)
print("Loaded histogram with bins {}".format(gaia_6d_hist[0].shape[0]))
except IOError:
bins = 20
print("Generating histogram with bins {}".format(bins))
gaia_xyzuvw = np.load('../data/gaia_dr2_mean_xyzuvw.npy')
gaia_6d_hist = np.histogramdd(gaia_xyzuvw, bins=bins)
np.save(gaia_6d_hist_file, gaia_6d_hist)
if BG_DENS != 0.:
bg_ln_ols = np.log(np.zeros(nstars) + BG_DENS)
else:
bg_ln_ols = None
while ncomps < MAX_COMP:
# handle special case of one component
if ncomps == 1:
logging.info("******************************************")
logging.info("********* FITTING 1 COMPONENT **********")
logging.info("******************************************")
run_dir = rdir + '{}/'.format(ncomps)
mkpath(run_dir)
try:
new_groups = dt.loadGroups(run_dir + 'final/final_groups.npy')
new_meds = np.load(run_dir + 'final/final_med_errs.npy')
new_z = np.load(run_dir + 'final/final_membership.npy')
logging.info("Loaded from previous run")
except IOError:
new_groups, new_meds, new_z =\
em.fit_many_comps(star_pars, ncomps, rdir=run_dir, pool=pool,
bg_dens=BG_DENS,
)
new_groups = np.array(new_groups)
new_lnlike = em.get_overall_lnlikelihood(star_pars, new_groups,
bg_ln_ols=bg_ln_ols)
new_lnpost = em.get_overall_lnlikelihood(star_pars, new_groups,
bg_ln_ols=bg_ln_ols,
inc_posterior=True)
logging.info("{} tolerance".format(C_TOL))
logging.info("In the directory: {}".format(rdir))
## Destination: (inspired by LCC)
mean_now_lsr = np.array([50., -100., 25., 1.1, -7.76, 2.25])
# Calculate appropriate starting point
mean_then = torb.trace_cartesian_orbit(mean_now_lsr, -age)
# gather inputs
# group_pars = np.hstack((mean_then, dX, dV, age, nstars))
# Setting up perfect current xyzuvw values
try:
xyzuvw_now_perf = np.load(rdir+xyzuvw_perf_file)
# origin = np.load(group_savefile)
origin = dt.loadGroups(rdir+group_savefile)[0]
logging.basicConfig(
level=logging.INFO, filemode='a',
filename='my_investigator_demo.log',
)
logging.info("appending to previous attempt")
except IOError:
logging.basicConfig(
level=logging.INFO, filemode='w',
filename='my_investigator_demo.log',
)
logging.info("Beginning fresh run:")
logging.info("Input arguments: {}".format(sys.argv[1:]))
logging.info("\n"
"\tage: {}\n"
"\tprecs: {}".format(
ERROR = 1.0
group_pars = np.array([
[ 26.11, 38.20, 23.59, 0.34, -3.87, -0.36, 2.23, -0.05, 14.32, 35],
[ 35.25, 52.58, 24.00, -2.80, -2.85, 0.25, 3.07, -0.27, 16.72, 25],
])
ngroups = group_pars.shape[0]
# Get background stars
# get typical background density
# nbgstars = 10
try:
all_xyzuvw_now_perf = np.load(xyzuvw_perf_file)
origins = dt.loadGroups(groups_savefile)
star_pars = dt.loadXYZUVW(xyzuvw_conv_savefile)
logging.info("Loaded synth data from previous run")
except IOError:
all_xyzuvw_init = np.zeros((0,6))
all_xyzuvw_now_perf = np.zeros((0,6))
origins = []
for i in range(ngroups):
logging.info(" generating from group {}".format(i))
# MANUALLY SEPARATE CURRENT DAY DISTROS IN DIMENSION X
# mean_now_w_offset = mean_now.copy()
# mean_now_w_offset[0] += i * 50
# mean_then = torb.traceOrbitXYZUVW(mean_now_w_offset, -extra_pars[i,-2],
# single_age=True)
xyzuvw_init, origin = syn.synthesiseXYZUVW(group_pars[i], form='sphere',
# if dim1 == 0 and dim2 == 1 and debugging_circles:
# bpmg_range[1] = temp_range
# To ensure consistency, we now plot the BANYAN bpmg stars only,
# and use the ragnes from previous plot
# 2019-07-13 [TC]: not sure why this is here...
if False:
fit_name = 'banyan_bpmg'
rdir = '../../results/em_fit/beta_Pictoris/'
memb_file = rdir + 'final_membership.npy'
groups_file = rdir + 'final_best_groups.npy'
star_pars_file = '../../data/beta_Pictoris_with_gaia_small_xyzuvw.fits'
z = np.load(memb_file)
groups = dt.loadGroups(groups_file)
star_pars = dt.loadDictFromTable(star_pars_file, 'beta Pictoris')
nstars = len(star_pars['xyzuvw'])
# First do all, then just do possible membs of BPMG
for dim1, dim2 in DEFAULT_DIMS: #[(0,1), (0, 3), (1, 4), (2,5)]: #, (2, 5)]: # , 'yv', 'zw']:
# if dim1 == 0 and dim2 == 1 and debugging_circles:
# temp_range = bpmg_range[1]
# bpmg_range[1] = [-120, 80]
# import pdb; pdb.set_trace()
fp.plotPane(
dim1,
dim2,
groups=groups,
star_pars=star_pars,
group_now=True,
import sys
import chronostar.likelihood
sys.path.insert(0, '..')
import chronostar.retired2.datatool as dt
import chronostar.retired2.converter as cv
import chronostar.coordinate as cc
import chronostar.expectmax as em
filename = '../data/2M0249-05.fits'
star_tab = Table.read(filename)
final_groups = '../results/em_fit/beta_Pictoris_wgs_inv2_5B_res/' \
'final_groups.npy'
beta_fit = dt.loadGroups(final_groups)[0]
gaia_xyzuvw_file = '../data/gaia_dr2_mean_xyzuvw.npy'
if np.isnan(star_tab['radial_velocity']):
print("Its nan")
# from Shkolnik 2017
# star_tab['radial_velocity'] = 14.4
# star_tab['radial_velocity_error'] = 0.4
star_tab['radial_velocity'] = 16.44
star_tab['radial_velocity_error'] = 1.
if np.isnan(star_tab['radial_velocity']):
print("Its nan")
# extend proper motion uncertainty
np.save(bg_savefile, BG_DENS)
logging.info(" with error fraction {}".format(ERROR))
logging.info(" and background density {}".format(BG_DENS))
# Set a current-day location around which synth stars will end up
mean_now = np.array([50., -100., -0., -10., -20., -5.])
logging.info("Mean (now):\n{}".format(mean_now))
logging.info("Extra pars:\n{}".format(extra_pars))
logging.info("Offsets:\n{}".format(offsets))
try:
#all_xyzuvw_now_perf = np.load(xyzuvw_perf_file)
np.load(xyzuvw_perf_file)
#origins = dt.loadGroups(groups_savefile)
dt.loadGroups(groups_savefile)
#star_pars = dt.loadXYZUVW(xyzuvw_conv_savefile)
dt.loadXYZUVW(xyzuvw_conv_savefile)
logging.info("Synth data exists! .....")
print("Synth data exists")
raise UserWarning
except IOError:
all_xyzuvw_init = np.zeros((0, 6))
all_xyzuvw_now_perf = np.zeros((0, 6))
origins = []
for i in range(ngroups):
logging.info(" generating from group {}".format(i))
# MANUALLY SEPARATE CURRENT DAY DISTROS IN DIMENSION X
mean_now_w_offset = mean_now.copy()
# mean_now_w_offset[0] += i * 50
mean_now_w_offset += offsets[i]
Calculates the density of a pdf characterised by a histogram at point x
"""
# Check if handling 1D histogram
if len(bin_heights.shape) == 1:
raise UserWarning
dims = len(bin_heights.shape)
bin_widths = [bins[1] - bins[0] for bins in bin_edges]
bin_area = np.prod(bin_widths)
x_ix = tuple([np.digitize(x[dim], bin_edges[dim]) - 1
for dim in range(dims)])
return bin_heights[x_ix] / bin_area
rdir = 'example_comps/'
init_groups = dt.loadGroups(rdir+'iter00/best_groups.npy')
init_z = np.load(rdir+'iter00/membership.npy')
later_groups = dt.loadGroups(rdir + 'iter15/best_groups.npy')
later_z = np.load(rdir+'iter15/membership.npy')
xs = np.linspace(2, 100,10000)
plt.clf()
for sig in [.5, .75, 1]:
plt.plot(xs, gf.lnlognormal(xs, sig=sig), label='{}'.format(sig))
plt.legend(loc='best')
plt.savefig("lnlognormal.pdf")
print("Init alphas")
for i, (igroup, lgroup) in enumerate(zip(init_groups, later_groups)):
sig=1.
print('Group {}'.format(i))
plt_file = rdir + 'multi_plot_{}_{}_{}_{}_{}_{}.pdf'.format(*scenario)
print("Checking {}".format(plt_file))
if not os.path.isfile(plt_file):
print("Plotting {}".format(plt_file))
try:
star_pars_file = rdir + 'xyzuvw_now.fits'
chain_file = rdir + 'final_chain.npy'
origins_file = rdir + 'origins.npy'
lnprob_file = rdir + 'final_lnprob.npy'
chain = np.load(chain_file).reshape(-1, 9)
lnprob = np.load(lnprob_file)
best_pars = chain[np.argmax(lnprob_file)]
best_fit = chronostar.component.Component(best_pars, internal=True)
origins = dt.loadGroups(origins_file)
star_pars = dt.loadXYZUVW(star_pars_file)
fp.plotMultiPane(
['xy', 'xz', 'uv', 'xu', 'yv', 'zw'],
star_pars,
[best_fit],
origins=origins,
save_file=rdir +
'multi_plot_{}_{}_{}_{}_{}_{}.pdf'.format(*scenario),
title='{}Myr, {}pc, {}km/s, {} stars, {}, {}'.format(
*scenario),
)
print("done")
except:
def plotPane(dim1=0,
dim2=1,
ax=None,
groups=(),
star_pars=None,
origin_star_pars=None,
star_orbits=False,
origins=None,
group_then=False,
group_now=False,
group_orbit=False,
annotate=False,
membership=None,
true_memb=None,
savefile='',
with_bg=False,
markers=None,
group_bg=False,
marker_labels=None,
color_labels=None,
marker_style=None,
marker_legend=None,
color_legend=None,
star_pars_label=None,
origin_star_pars_label=None,
range_1=None,
range_2=None,
isotropic=False,
ordering=None,
no_bg_covs=False):
"""
Plots a single pane capturing kinematic info in any desired 2D plane
Uses global constants COLORS and HATCHES to inform consistent colour
scheme.
Can use this to plot different panes of one whole figure
Parameters
----------
dim1: x-axis, can either be integer 0-5 (inclusive) or a letter form
'xyzuvw' (either case)
dim2: y-axis, same conditions as dim1
ax: the axes object on which to plot (defaults to pyplots currnet axes)
groups: a list of (or just one) synthesiser.Group objects, corresponding
to the fit of the origin(s)
star_pars: dict object with keys 'xyzuvw' ([nstars,6] array of current
star means) and 'xyzuvw_cov' ([nstars,6,6] array of current
star covariance matrices)
star_orbits: (bool) plot the calculated stellar traceback orbits of
central estimate of measurements
group_then: (bool) plot the group's origin
group_now: (bool) plot the group's current day distribution
group_orbit: (bool) plot the trajectory of the group's mean
annotate: (bool) add text describing the figure's contents
with_bg: (bool) treat the last column in Z as members of background, and
color accordingly
no_bg_covs: (bool) ignore covariance matrices of stars fitted to background
Returns
-------
(nothing returned)
"""
labels = 'XYZUVW'
units = 3 * ['pc'] + 3 * ['km/s']
if savefile:
plt.clf()
# Tidying up inputs
if ax is None:
ax = plt.gca()
if type(dim1) is not int:
dim1 = labels.index(dim1.upper())
if type(dim2) is not int:
dim2 = labels.index(dim2.upper())
if type(star_pars) is str:
star_pars = dt.loadXYZUVW(star_pars)
if type(membership) is str:
membership = np.load(membership)
if type(groups) is str:
groups = dt.loadGroups(groups)
if marker_style is None:
marker_style = MARKERS[:]
# if type(origin_star_pars) is str:
# origin_star_pars = dt.loadXYZUVW(origin_star_pars)
legend_pts = []
legend_labels = []
# ensure groups is iterable
try:
len(groups)
except:
groups = [groups]
ngroups = len(groups)
if ordering is None:
ordering = range(len(marker_style))
# plot stellar data (positions with errors and optionally traceback
# orbits back to some ill-defined age
if star_pars:
nstars = star_pars['xyzuvw'].shape[0]
# apply default color and markers, to be overwritten if needed
pt_colors = np.array(nstars * [COLORS[0]])
if markers is None:
markers = np.array(nstars * ['.'])
# Incorporate fitted membership into colors of the pts
if membership is not None:
best_mship = np.argmax(membership[:, :ngroups + with_bg], axis=1)
pt_colors = np.array(COLORS[:ngroups] +
with_bg * ['xkcd:grey'])[best_mship]
# Incoporate "True" membership into pt markers
if true_memb is not None:
markers = np.array(MARKERS)[np.argmax(true_memb, axis=1)]
if with_bg:
true_bg_mask = np.where(true_memb[:, -1] == 1.)
markers[true_bg_mask] = '.'
all_mark_size = np.array(nstars * [MARK_SIZE])
# group_bg handles case where background is fitted to by final component
if with_bg:
all_mark_size[np.where(
np.argmax(membership, axis=1) == ngroups -
group_bg)] = BG_MARK_SIZE
mns = star_pars['xyzuvw']
try:
covs = np.copy(star_pars['xyzuvw_cov'])
# replace background cov matrices with None so as to avoid plotting
if with_bg and no_bg_covs:
print("Discarding background cov-mats")
# import pdb; pdb.set_trace()
covs[np.where(
np.argmax(membership, axis=1) == ngroups -
group_bg)] = None
except KeyError:
covs = len(mns) * [None]
star_pars['xyzuvw_cov'] = covs
st_count = 0
for star_mn, star_cov, marker, pt_color, m_size in zip(
mns, covs, markers, pt_colors, all_mark_size):
pt = ax.scatter(
star_mn[dim1],
star_mn[dim2],
s=m_size, #s=MARK_SIZE,
color=pt_color,
marker=marker,
alpha=PT_ALPHA,
linewidth=0.0,
)
# plot uncertainties
if star_cov is not None:
plotCovEllipse(
star_cov[np.ix_([dim1, dim2], [dim1, dim2])],
star_mn[np.ix_([dim1, dim2])],
ax=ax,
alpha=COV_ALPHA,
linewidth='0.1',
color=pt_color,
)
# plot traceback orbits for as long as oldest group (if known)
# else, 30 Myr
if star_orbits and st_count % 3 == 0:
try:
tb_limit = max([g.age for g in groups])
except:
tb_limit = 30
plotOrbit(star_mn,
dim1,
dim2,
ax,
end_age=-tb_limit,
color='xkcd:grey')
st_count += 1
if star_pars_label:
# ax.legend(numpoints=1)
legend_pts.append(pt)
legend_labels.append(star_pars_label)
if origin_star_pars is not None:
for star_mn, marker, pt_color, m_size in\
zip(origin_star_pars['xyzuvw'],
# origin_star_pars['xyzuvw_cov'],
markers, pt_colors, all_mark_size):
pt = ax.scatter(
star_mn[dim1],
star_mn[dim2],
s=0.5 * m_size,
# s=MARK_SIZE,
color=pt_color,
marker='s',
alpha=PT_ALPHA,
linewidth=0.0, #label=origin_star_pars_label,
)
# # plot uncertainties
# if star_cov is not None:
# ee.plotCovEllipse(
# star_cov[np.ix_([dim1, dim2], [dim1, dim2])],
# star_mn[np.ix_([dim1, dim2])],
# ax=ax, alpha=0.05, linewidth='0.1',
# color=pt_color,
# )
if origin_star_pars_label:
legend_pts.append(pt)
legend_labels.append(origin_star_pars_label)
# plot info for each group (fitted, or true synthetic origin)
for i, group in enumerate(groups):
cov_then = group.get_covmatrix()
mean_then = group.get_mean()
# plot group initial distribution
if group_then:
ax.plot(mean_then[dim1],
mean_then[dim2],
marker='+',
alpha=0.3,
color=COLORS[i])
plotCovEllipse(cov_then[np.ix_([dim1, dim2], [dim1, dim2])],
mean_then[np.ix_([dim1, dim2])],
with_line=True,
ax=ax,
alpha=0.3,
ls='--',
color=COLORS[i])
if annotate:
ax.annotate(r'$\mathbf{\mu}_0, \mathbf{\Sigma}_0$',
(mean_then[dim1], mean_then[dim2]),
color=COLORS[i])
# plot group current day distribution (should match well with stars)
if group_now:
mean_now = torb.trace_cartesian_orbit(mean_then,
group.get_age(),
single_age=True)
cov_now = tf.transform_covmatrix(cov_then,
torb.trace_cartesian_orbit,
mean_then,
args=[group.get_age()])
ax.plot(mean_now[dim1],
mean_now[dim2],
marker='+',
alpha=0.3,
color=COLORS[i])
plotCovEllipse(
cov_now[np.ix_([dim1, dim2], [dim1, dim2])],
mean_now[np.ix_([dim1, dim2])],
# with_line=True,
ax=ax,
alpha=0.4,
ls='-.',
ec=COLORS[i],
fill=False,
hatch=HATCHES[i],
color=COLORS[i])
if annotate:
ax.annotate(r'$\mathbf{\mu}_c, \mathbf{\Sigma}_c$',
(mean_now[dim1], mean_now[dim2]),
color=COLORS[i])
# plot orbit of mean of group
if group_orbit:
plotOrbit(mean_now,
dim1,
dim2,
ax,
-group.age,
group_ix=i,
with_arrow=True,
annotate=annotate)
if origins:
for origin in origins:
cov_then = origin.generateSphericalCovMatrix()
mean_then = origin.mean
# plot origin initial distribution
ax.plot(mean_then[dim1],
mean_then[dim2],
marker='+',
color='xkcd:grey')
plotCovEllipse(cov_then[np.ix_([dim1, dim2], [dim1, dim2])],
mean_then[np.ix_([dim1, dim2])],
with_line=True,
ax=ax,
alpha=0.1,
ls='--',
color='xkcd:grey')
ax.set_xlabel("{} [{}]".format(labels[dim1], units[dim1]))
ax.set_ylabel("{} [{}]".format(labels[dim2], units[dim2]))
# NOT QUITE....
# if marker_legend is not None and color_legend is not None:
# x_loc = np.mean(star_pars['xyzuvw'][:,dim1])
# y_loc = np.mean(star_pars['xyzuvw'][:,dim2])
# for label in marker_legend.keys():
# ax.plot(x_loc, y_loc, color=color_legend[label],
# marker=marker_legend[label], alpha=0, label=label)
# ax.legend(loc='best')
# if star_pars_label is not None:
# ax.legend(numpoints=1, loc='best')
# ax.legend(loc='best')
# if marker_order is not None:
# for label_ix, marker_ix in enumerate(marker_order):
# axleg.scatter(0,0,color='black',marker=MARKERS[marker_ix],
# label=MARKER_LABELS[label_ix])
# #
# if len(legend_pts) > 0:
# ax.legend(legend_pts, legend_labels)
# update fontsize
for item in ([ax.title, ax.xaxis.label, ax.yaxis.label] +
ax.get_xticklabels() + ax.get_yticklabels()):
item.set_fontsize(FONTSIZE)
if ax.get_legend() is not None:
for item in ax.get_legend().get_texts():
item.set_fontsize(FONTSIZE)
# ensure we have some handle on the ranges
# if range_1 is None:
# range_1 = ax.get_xlim()
# if range_2 is None:
# range_2 = ax.get_ylim()
if range_2:
ax.set_ylim(range_2)
if isotropic:
print("Setting isotropic for dims {} and {}".format(dim1, dim2))
# plt.gca().set_aspect('equal', adjustable='box')
# import pdb; pdb.set_trace()
plt.gca().set_aspect('equal', adjustable='datalim')
# manually calculate what the new xaxis must be...
figW, figH = ax.get_figure().get_size_inches()
xmid = (ax.get_xlim()[1] + ax.get_xlim()[0]) * 0.5
yspan = ax.get_ylim()[1] - ax.get_ylim()[0]
xspan = figW * yspan / figH
# check if this increases span
if xspan > ax.get_xlim()[1] - ax.get_xlim()[0]:
ax.set_xlim(xmid - 0.5 * xspan, xmid + 0.5 * xspan)
# if not, need to increase yspan
else:
ymid = (ax.get_ylim()[1] + ax.get_ylim()[0]) * 0.5
xspan = ax.get_xlim()[1] - ax.get_xlim()[0]
yspan = figH * xspan / figW
ax.set_ylim(ymid - 0.5 * yspan, ymid + 0.5 * yspan)
# import pdb; pdb.set_trace()
elif range_1:
ax.set_xlim(range_1)
if color_labels is not None:
xlim = ax.get_xlim()
ylim = ax.get_ylim()
for i, color_label in enumerate(color_labels):
ax.plot(1e10, 1e10, color=COLORS[i], label=color_label)
ax.set_xlim(xlim)
ax.set_ylim(ylim)
ax.legend(loc='best')
if marker_labels is not None:
xlim = ax.get_xlim()
ylim = ax.get_ylim()
# import pdb; pdb.set_trace()
for i, marker_label in enumerate(marker_labels):
ax.scatter(
1e10,
1e10,
c='black',
marker=np.array(marker_style)[ordering][i],
# marker=MARKERS[list(marker_labels).index(marker_label)],
label=marker_label)
if with_bg:
ax.scatter(1e10,
1e10,
c='xkcd:grey',
marker='.',
label='Background')
ax.set_xlim(xlim)
ax.set_ylim(ylim)
ax.legend(loc='best')
# if marker_legend is not None:
# xlim = ax.get_xlim()
# ylim = ax.get_ylim()
# # import pdb; pdb.set_trace()
# for k, v in marker_legend.items():
# ax.scatter(1e10, 1e10, c='black',
# marker=v, label=k)
# ax.set_xlim(xlim)
# ax.set_ylim(ylim)
# ax.legend(loc='best')
if color_legend is not None and marker_legend is not None:
xlim = ax.get_xlim()
ylim = ax.get_ylim()
# import pdb; pdb.set_trace()
for label in color_legend.keys():
ax.scatter(1e10,
1e10,
c=color_legend[label],
marker=marker_legend[label],
label=label)
ax.set_xlim(xlim)
ax.set_ylim(ylim)
ax.legend(loc='best')
if savefile:
# set_size(4,2,ax)
plt.savefig(savefile)
# import pdb; pdb.set_trace()
# return ax.get_window_extent(None).width, ax.get_window_extent(None).height
return ax.get_xlim(), ax.get_ylim()
right=True)
print("... saving")
plt.savefig(plot_name)
if PLOT_BPMG_REAL:
for iteration in ['5B']: #, '6C']:
star_pars_file = '../../data/beta_Pictoris_with_gaia_small_xyzuvw.fits'
star_pars = dt.loadXYZUVW(star_pars_file)
fit_name = 'bpmg_and_nearby'
rdir = '../../results/em_fit/beta_Pictoris_wgs_inv2_{}_res/'.format(iteration)
memb_file = rdir + 'final_membership.npy'
groups_file = rdir + 'final_groups.npy'
z = np.load(memb_file)
groups = dt.loadGroups(groups_file)
# Assign markers based on BANYAN membership
gt_sp = dt.loadDictFromTable('../../data/banyan_with_gaia_near_bpmg_xyzuvw.fits')
banyan_membership = len(star_pars['xyzuvw']) * ['N/A']
for i in range(len(star_pars['xyzuvw'])):
master_table_ix = np.where(gt_sp['table']['source_id']==star_pars['gaia_ids'][i])
banyan_membership[i] = gt_sp['table']['Moving group'][master_table_ix[0][0]]
# assign markers based on present moving groups, keep track of
# assoc -> marker relationship incase a legend is called for
banyan_membership=np.array(banyan_membership)
banyan_markers = np.array(len(banyan_membership) * ['.'])
banyan_memb_set = set(banyan_membership)
for bassoc in set(gt_sp['table']['Moving group']):
np.save(bg_savefile, BG_DENS)
logging.info(" with error fraction {}".format(ERROR))
logging.info(" and background density {}".format(BG_DENS))
# Set a current-day location around which synth stars will end up
mean_now = np.array([50., -100., -0., -10., -20., -5.])
logging.info("Mean (now):\n{}".format(mean_now))
logging.info("Extra pars:\n{}".format(extra_pars))
logging.info("Offsets:\n{}".format(offsets))
try:
#all_xyzuvw_now_perf = np.load(xyzuvw_perf_file)
np.load(xyzuvw_perf_file)
#origins = dt.loadGroups(groups_savefile)
dt.loadGroups(groups_savefile)
#star_pars = dt.loadXYZUVW(xyzuvw_conv_savefile)
dt.loadXYZUVW(xyzuvw_conv_savefile)
logging.info("Synth data exists! .....")
print("Synth data exists")
raise UserWarning
except IOError:
all_xyzuvw_init = np.zeros((0,6))
all_xyzuvw_now_perf = np.zeros((0,6))
origins = []
for i in range(ngroups):
logging.info(" generating from group {}".format(i))
# MANUALLY SEPARATE CURRENT DAY DISTROS IN DIMENSION X
mean_now_w_offset = mean_now.copy()
# mean_now_w_offset[0] += i * 50
mean_now_w_offset += offsets[i]
final_med_errs.append(fmed_errs)
else:
comp_fgroups = []
comp_fmembs = []
comp_fmed_errs = []
for i in range(ncomps-1):
subrdir = rdir + '{}/{}/final/'.format(ncomps, chr(ord('A') + i))
if os.path.isdir(subrdir):
fgroups, fmembs, fmed_errs = loadFinalResults(subrdir)
comp_fgroups.append(fgroups)
comp_fmembs.append(fmembs)
comp_fmed_errs.append(fmed_errs)
final_fits.append(comp_fgroups)
final_membs.append(comp_fmembs)
final_med_errs.append(comp_fmed_errs)
ncomps += 1
# load in the synthetic data
sdir = rdir + 'synth_data/'
if os.path.isdir(sdir):
origins = dt.loadGroups(sdir + 'origins.npy')
perf_xyzuvw = np.load(sdir + 'perf_xyzuvw.npy')
true_z = dt.getZfromOrigins(origins, star_pars)
nassoc_stars = np.sum([o.nstars for o in origins])
try:
bg_ln_ols = np.load(rdir + 'bg_ln_ols.npy')
except IOError:
print("couldn't find background ln overlaps...")
import sys
import chronostar.likelihood
sys.path.insert(0, '..')
import chronostar.retired2.datatool as dt
import chronostar.retired2.converter as cv
import chronostar.coordinate as cc
import chronostar.expectmax as em
filename = '../data/2M0249-05.fits'
star_tab = Table.read(filename)
final_groups = '../results/em_fit/beta_Pictoris_wgs_inv2_5B_res/' \
'final_groups.npy'
beta_fit = dt.loadGroups(final_groups)[0]
gaia_xyzuvw_file = '../data/gaia_dr2_mean_xyzuvw.npy'
if np.isnan(star_tab['radial_velocity']):
print("Its nan")
# from Shkolnik 2017
# star_tab['radial_velocity'] = 14.4
# star_tab['radial_velocity_error'] = 0.4
star_tab['radial_velocity'] = 16.44
star_tab['radial_velocity_error'] = 1.
if np.isnan(star_tab['radial_velocity']):
print("Its nan")
# extend proper motion uncertainty
def plotPane(dim1=0, dim2=1, ax=None, groups=(), star_pars=None,
origin_star_pars=None,
star_orbits=False, origins=None,
group_then=False, group_now=False, group_orbit=False,
annotate=False, membership=None, true_memb=None,
savefile='', with_bg=False, markers=None, group_bg=False,
marker_labels=None, color_labels=None,
marker_style=None,
marker_legend=None, color_legend=None,
star_pars_label=None, origin_star_pars_label=None,
range_1=None, range_2=None, isotropic=False,
ordering=None, no_bg_covs=False):
"""
Plots a single pane capturing kinematic info in any desired 2D plane
Uses global constants COLORS and HATCHES to inform consistent colour
scheme.
Can use this to plot different panes of one whole figure
Parameters
----------
dim1: x-axis, can either be integer 0-5 (inclusive) or a letter form
'xyzuvw' (either case)
dim2: y-axis, same conditions as dim1
ax: the axes object on which to plot (defaults to pyplots currnet axes)
groups: a list of (or just one) synthesiser.Group objects, corresponding
to the fit of the origin(s)
star_pars: dict object with keys 'xyzuvw' ([nstars,6] array of current
star means) and 'xyzuvw_cov' ([nstars,6,6] array of current
star covariance matrices)
star_orbits: (bool) plot the calculated stellar traceback orbits of
central estimate of measurements
group_then: (bool) plot the group's origin
group_now: (bool) plot the group's current day distribution
group_orbit: (bool) plot the trajectory of the group's mean
annotate: (bool) add text describing the figure's contents
with_bg: (bool) treat the last column in Z as members of background, and
color accordingly
no_bg_covs: (bool) ignore covariance matrices of stars fitted to background
Returns
-------
(nothing returned)
"""
labels = 'XYZUVW'
units = 3 * ['pc'] + 3 * ['km/s']
if savefile:
plt.clf()
# Tidying up inputs
if ax is None:
ax = plt.gca()
if type(dim1) is not int:
dim1 = labels.index(dim1.upper())
if type(dim2) is not int:
dim2 = labels.index(dim2.upper())
if type(star_pars) is str:
star_pars = dt.loadXYZUVW(star_pars)
if type(membership) is str:
membership = np.load(membership)
if type(groups) is str:
groups = dt.loadGroups(groups)
if marker_style is None:
marker_style = MARKERS[:]
# if type(origin_star_pars) is str:
# origin_star_pars = dt.loadXYZUVW(origin_star_pars)
legend_pts = []
legend_labels = []
# ensure groups is iterable
try:
len(groups)
except:
groups = [groups]
ngroups = len(groups)
if ordering is None:
ordering = range(len(marker_style))
# plot stellar data (positions with errors and optionally traceback
# orbits back to some ill-defined age
if star_pars:
nstars = star_pars['xyzuvw'].shape[0]
# apply default color and markers, to be overwritten if needed
pt_colors = np.array(nstars * [COLORS[0]])
if markers is None:
markers = np.array(nstars * ['.'])
# Incorporate fitted membership into colors of the pts
if membership is not None:
best_mship = np.argmax(membership[:,:ngroups+with_bg], axis=1)
pt_colors = np.array(COLORS[:ngroups] + with_bg*['xkcd:grey'])[best_mship]
# Incoporate "True" membership into pt markers
if true_memb is not None:
markers = np.array(MARKERS)[np.argmax(true_memb,
axis=1)]
if with_bg:
true_bg_mask = np.where(true_memb[:,-1] == 1.)
markers[true_bg_mask] = '.'
all_mark_size = np.array(nstars * [MARK_SIZE])
# group_bg handles case where background is fitted to by final component
if with_bg:
all_mark_size[np.where(np.argmax(membership, axis=1) == ngroups-group_bg)] = BG_MARK_SIZE
mns = star_pars['xyzuvw']
try:
covs = np.copy(star_pars['xyzuvw_cov'])
# replace background cov matrices with None so as to avoid plotting
if with_bg and no_bg_covs:
print("Discarding background cov-mats")
# import pdb; pdb.set_trace()
covs[np.where(np.argmax(membership, axis=1) == ngroups-group_bg)] = None
except KeyError:
covs = len(mns) * [None]
star_pars['xyzuvw_cov'] = covs
st_count = 0
for star_mn, star_cov, marker, pt_color, m_size in zip(mns, covs, markers, pt_colors,
all_mark_size):
pt = ax.scatter(star_mn[dim1], star_mn[dim2], s=m_size, #s=MARK_SIZE,
color=pt_color, marker=marker, alpha=PT_ALPHA,
linewidth=0.0,
)
# plot uncertainties
if star_cov is not None:
plotCovEllipse(star_cov[np.ix_([dim1, dim2], [dim1, dim2])],
star_mn[np.ix_([dim1, dim2])],
ax=ax, alpha=COV_ALPHA, linewidth='0.1',
color=pt_color,
)
# plot traceback orbits for as long as oldest group (if known)
# else, 30 Myr
if star_orbits and st_count%3==0:
try:
tb_limit = max([g.age for g in groups])
except:
tb_limit = 30
plotOrbit(star_mn, dim1, dim2, ax, end_age=-tb_limit,
color='xkcd:grey')
st_count += 1
if star_pars_label:
# ax.legend(numpoints=1)
legend_pts.append(pt)
legend_labels.append(star_pars_label)
if origin_star_pars is not None:
for star_mn, marker, pt_color, m_size in\
zip(origin_star_pars['xyzuvw'],
# origin_star_pars['xyzuvw_cov'],
markers, pt_colors, all_mark_size):
pt = ax.scatter(star_mn[dim1], star_mn[dim2], s=0.5*m_size,
# s=MARK_SIZE,
color=pt_color, marker='s', alpha=PT_ALPHA,
linewidth=0.0, #label=origin_star_pars_label,
)
# # plot uncertainties
# if star_cov is not None:
# ee.plotCovEllipse(
# star_cov[np.ix_([dim1, dim2], [dim1, dim2])],
# star_mn[np.ix_([dim1, dim2])],
# ax=ax, alpha=0.05, linewidth='0.1',
# color=pt_color,
# )
if origin_star_pars_label:
legend_pts.append(pt)
legend_labels.append(origin_star_pars_label)
# plot info for each group (fitted, or true synthetic origin)
for i, group in enumerate(groups):
cov_then = group.generateSphericalCovMatrix()
mean_then = group.mean
# plot group initial distribution
if group_then:
ax.plot(mean_then[dim1], mean_then[dim2], marker='+', alpha=0.3,
color=COLORS[i])
plotCovEllipse(cov_then[np.ix_([dim1, dim2], [dim1, dim2])],
mean_then[np.ix_([dim1,dim2])],
with_line=True,
ax=ax, alpha=0.3, ls='--',
color=COLORS[i])
if annotate:
ax.annotate(r'$\mathbf{\mu}_0, \mathbf{\Sigma}_0$',
(mean_then[dim1],
mean_then[dim2]),
color=COLORS[i])
# plot group current day distribution (should match well with stars)
if group_now:
mean_now = torb.trace_cartesian_orbit(mean_then, group.age,
single_age=True)
cov_now = tf.transform_covmatrix(cov_then, torb.trace_cartesian_orbit,
mean_then, args=[group.age])
ax.plot(mean_now[dim1], mean_now[dim2], marker='+', alpha=0.3,
color=COLORS[i])
plotCovEllipse(cov_now[np.ix_([dim1, dim2], [dim1, dim2])],
mean_now[np.ix_([dim1,dim2])],
# with_line=True,
ax=ax, alpha=0.4, ls='-.',
ec=COLORS[i], fill=False, hatch=HATCHES[i],
color=COLORS[i])
if annotate:
ax.annotate(r'$\mathbf{\mu}_c, \mathbf{\Sigma}_c$',
(mean_now[dim1],mean_now[dim2]),
color=COLORS[i])
# plot orbit of mean of group
if group_orbit:
plotOrbit(mean_now, dim1, dim2, ax, -group.age, group_ix=i,
with_arrow=True, annotate=annotate)
if origins:
for origin in origins:
cov_then = origin.generateSphericalCovMatrix()
mean_then = origin.mean
# plot origin initial distribution
ax.plot(mean_then[dim1], mean_then[dim2], marker='+',
color='xkcd:grey')
plotCovEllipse(
cov_then[np.ix_([dim1, dim2], [dim1, dim2])],
mean_then[np.ix_([dim1, dim2])],
with_line=True,
ax=ax, alpha=0.1, ls='--',
color='xkcd:grey')
ax.set_xlabel("{} [{}]".format(labels[dim1], units[dim1]))
ax.set_ylabel("{} [{}]".format(labels[dim2], units[dim2]))
# NOT QUITE....
# if marker_legend is not None and color_legend is not None:
# x_loc = np.mean(star_pars['xyzuvw'][:,dim1])
# y_loc = np.mean(star_pars['xyzuvw'][:,dim2])
# for label in marker_legend.keys():
# ax.plot(x_loc, y_loc, color=color_legend[label],
# marker=marker_legend[label], alpha=0, label=label)
# ax.legend(loc='best')
# if star_pars_label is not None:
# ax.legend(numpoints=1, loc='best')
# ax.legend(loc='best')
# if marker_order is not None:
# for label_ix, marker_ix in enumerate(marker_order):
# axleg.scatter(0,0,color='black',marker=MARKERS[marker_ix],
# label=MARKER_LABELS[label_ix])
# #
# if len(legend_pts) > 0:
# ax.legend(legend_pts, legend_labels)
# update fontsize
for item in ([ax.title, ax.xaxis.label, ax.yaxis.label] +
ax.get_xticklabels() + ax.get_yticklabels()):
item.set_fontsize(FONTSIZE)
if ax.get_legend() is not None:
for item in ax.get_legend().get_texts():
item.set_fontsize(FONTSIZE)
# ensure we have some handle on the ranges
# if range_1 is None:
# range_1 = ax.get_xlim()
# if range_2 is None:
# range_2 = ax.get_ylim()
if range_2:
ax.set_ylim(range_2)
if isotropic:
print("Setting isotropic for dims {} and {}".format(dim1, dim2))
# plt.gca().set_aspect('equal', adjustable='box')
# import pdb; pdb.set_trace()
plt.gca().set_aspect('equal', adjustable='datalim')
# manually calculate what the new xaxis must be...
figW, figH = ax.get_figure().get_size_inches()
xmid = (ax.get_xlim()[1] + ax.get_xlim()[0]) * 0.5
yspan = ax.get_ylim()[1] - ax.get_ylim()[0]
xspan = figW * yspan / figH
# check if this increases span
if xspan > ax.get_xlim()[1] - ax.get_xlim()[0]:
ax.set_xlim(xmid - 0.5 * xspan, xmid + 0.5 * xspan)
# if not, need to increase yspan
else:
ymid = (ax.get_ylim()[1] + ax.get_ylim()[0]) * 0.5
xspan = ax.get_xlim()[1] - ax.get_xlim()[0]
yspan = figH * xspan / figW
ax.set_ylim(ymid - 0.5*yspan, ymid + 0.5*yspan)
# import pdb; pdb.set_trace()
elif range_1:
ax.set_xlim(range_1)
if color_labels is not None:
xlim = ax.get_xlim()
ylim = ax.get_ylim()
for i, color_label in enumerate(color_labels):
ax.plot(1e10, 1e10, color=COLORS[i], label=color_label)
ax.set_xlim(xlim)
ax.set_ylim(ylim)
ax.legend(loc='best')
if marker_labels is not None:
xlim = ax.get_xlim()
ylim = ax.get_ylim()
# import pdb; pdb.set_trace()
for i, marker_label in enumerate(marker_labels):
ax.scatter(1e10, 1e10, c='black',
marker=np.array(marker_style)[ordering][i],
# marker=MARKERS[list(marker_labels).index(marker_label)],
label=marker_label)
if with_bg:
ax.scatter(1e10, 1e10, c='xkcd:grey',
marker='.', label='Background')
ax.set_xlim(xlim)
ax.set_ylim(ylim)
ax.legend(loc='best')
# if marker_legend is not None:
# xlim = ax.get_xlim()
# ylim = ax.get_ylim()
# # import pdb; pdb.set_trace()
# for k, v in marker_legend.items():
# ax.scatter(1e10, 1e10, c='black',
# marker=v, label=k)
# ax.set_xlim(xlim)
# ax.set_ylim(ylim)
# ax.legend(loc='best')
if color_legend is not None and marker_legend is not None:
xlim = ax.get_xlim()
ylim = ax.get_ylim()
# import pdb; pdb.set_trace()
for label in color_legend.keys():
ax.scatter(1e10, 1e10, c=color_legend[label],
marker=marker_legend[label], label=label)
ax.set_xlim(xlim)
ax.set_ylim(ylim)
ax.legend(loc='best')
if savefile:
# set_size(4,2,ax)
plt.savefig(savefile)
# import pdb; pdb.set_trace()
# return ax.get_window_extent(None).width, ax.get_window_extent(None).height
return ax.get_xlim(), ax.get_ylim()
