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