def get_separate_mcsamples(chain):
"""
Function that returns separate :class:`~getdist.mcsamples.MCSamples`
for each sampler chain.
:param chain: :class:`~getdist.mcsamples.MCSamples` the input chain.
:return: list of :class:`~getdist.mcsamples.MCSamples` with the separate
chains.
"""
# get separate chains:
_chains = chain.getSeparateChains()
# copy the param names and ranges:
_mc_samples = []
for ch in _chains:
temp = MCSamples()
temp.paramNames = chain.getParamNames()
temp.setSamples(ch.samples, weights=ch.weights, loglikes=ch.loglikes)
temp.sampler = chain.sampler
temp.ranges = chain.ranges
temp.updateBaseStatistics()
_mc_samples.append(temp.copy())
#
return _mc_samples
def plot_contours(chaintags,
legend_labels=None,
params_toplot=None,
nsteps=None,
colors=None,
legend_loc='upper center',
legend_fontsize=20,
weight_with_dynesty=False):
bounds = utils.get_emulator_bounds()
sample_arr = []
for i, chaintag in enumerate(chaintags):
resample_chains = False
if 'dynesty' in chaintag:
resample_chains = True
assert not (nsteps != None and resample_chains
), "if resampling, shouldn't be selecting nsteps!"
chain_fn = f'../chains/chains_{chaintag}.h5'
fw = h5py.File(chain_fn, 'r')
chain_dset = fw['chain']
chain = np.array(chain_dset)
print(chaintag, ':', chain.shape)
if nsteps:
chain = chain[:, :nsteps, :]
lnprob_dset = fw['lnprob']
param_names = fw.attrs['param_names']
true_values = fw.attrs['true_values']
if params_toplot is not None:
idxs = []
for pm in params_toplot:
idxs.append(np.where(param_names == pm))
idxs = np.array(idxs).flatten()
chain = chain[:, :, idxs]
param_names = params_toplot
true_values = true_values[idxs]
nwalkers, nchain, ndim = chain.shape
samples = chain.reshape(-1, chain.shape[-1])
labels = [param_labels[pn] for pn in param_names]
ranges = [bounds[pn] for pn in param_names]
if resample_chains:
lnweight_dset = fw['lnweight']
lnevidence_dset = fw['lnevidence']
#[-1] bc just care about final evidence value
lnweight = np.array(lnweight_dset)[0] #[0] bc an extra array dim
lnevidence = np.array(lnevidence_dset)[0]
weights = np.exp(lnweight - lnevidence[-1])
weights = weights.flatten()
samples = np.empty((nchain, ndim))
for nd in range(ndim):
cn = chain[:, :, nd].flatten()
if weight_with_dynesty:
samples[:, nd] = dynesty.utils.resample_equal(cn, weights)
else:
samples[:, nd] = cn
else:
weights = None
fw.close()
if weight_with_dynesty:
weights = None
# have not gotten ranges to work
samps = MCSamples(names=param_names, labels=labels, ranges=ranges)
# for some reason get slightly diff answer if don't use setSamples and pass them straight to MCSamples!
samps.setSamples(samples, weights=weights)
sample_arr.append(samps)
g = plots.get_subplot_plotter()
g.settings.alpha_filled_add = 0.4
g.settings.figure_legend_frame = False
g.settings.legend_fontsize = legend_fontsize
g.settings.axis_marker_lw = 1.0
g.settings.axis_marker_color = 'dimgrey'
g.triangle_plot(
sample_arr,
filled=True,
contour_colors=colors,
names=param_names,
legend_labels=legend_labels,
markers=true_values,
title_limit=1,
legend_loc=legend_loc,
#marker_args={'color': 'orange', 'lw':1.5}
)
return g
def plot_contours_dynesty(chaintags,
legend_labels=None,
params_toplot=None,
colors=None,
legend_loc='upper center',
legend_fontsize=20,
vertical_markers=None,
vertical_marker_color='grey'):
bounds = utils.get_emulator_bounds()
sample_arr = []
for i, chaintag in enumerate(chaintags):
chain_fn = f'../chains/chains_{chaintag}.h5'
fw = h5py.File(chain_fn, 'r')
param_names = fw.attrs['param_names']
if vertical_markers is None:
vertical_markers_toplot = fw.attrs['true_values']
else:
vertical_markers_toplot = vertical_markers
fw.close()
pickle_fn = f'{pickle_dir}/results_{chaintag}.pkl'
with open(pickle_fn, 'rb') as pf:
res = pickle.load(pf)
samples = res['samples']
lnweight = np.array(res['logwt'])
lnevidence = np.array(res['logz'])
if params_toplot is not None:
idxs = []
for pm in params_toplot:
idxs.append(np.where(param_names == pm))
idxs = np.array(idxs).flatten()
samples = samples[:, idxs]
# weight and evidence are for all parameters, shape (n_samps,)
# so don't need to slice
param_names = params_toplot
vertical_markers_toplot = vertical_markers_toplot[idxs]
labels = [param_labels[pn] for pn in param_names]
ranges = [bounds[pn] for pn in param_names]
#[-1] bc just care about final evidence value
weights = np.exp(lnweight - lnevidence[-1])
weights = weights.flatten()
samps = MCSamples(names=param_names, labels=labels)
samps.setSamples(samples, weights=weights)
sample_arr.append(samps)
marker_args = {'color': vertical_marker_color}
g = plots.get_subplot_plotter()
g.settings.alpha_filled_add = 0.4
g.settings.figure_legend_frame = False
g.settings.legend_fontsize = legend_fontsize
g.settings.axis_marker_lw = 1.0
g.settings.axis_marker_color = 'dimgrey'
g.triangle_plot(sample_arr,
filled=True,
contour_colors=colors,
names=param_names,
legend_labels=legend_labels,
markers=vertical_markers_toplot,
title_limit=1,
legend_loc=legend_loc,
marker_args=marker_args,
axis_marker_color='red')
return g
