Source code for naima.analysis

# Licensed under a 3-clause BSD style license - see LICENSE.rst
import os
from pathlib import Path

import astropy.units as u
import h5py
import numpy as np
from astropy import log
from import read_table_hdf5, write_table_hdf5
from astropy.table import Table

from .plot import find_ML

    import yaml  # noqa

    HAS_PYYAML = True
except ImportError:
    HAS_PYYAML = False

__all__ = [

[docs]def save_diagnostic_plots( outname, sampler, modelidxs=None, pdf=False, sed=True, blob_labels=None, last_step=False, dpi=100, ): """ Generate diagnostic plots. - A plot for each of the chain parameters showing walker progression, final sample distribution and several statistical measures of this distribution: ``outname_chain_parN.png`` (see `naima.plot_chain`). - A corner plot of sample density in the two dimensional parameter space of all parameter pairs of the run, with the Maximum Likelihood parameter vector indicated in blue: ``outname_corner.png`` (see `corner.corner`). - A plot for each of the models returned as blobs by the model function. The maximum likelihood model is shown, as well as the 1 and 3 sigma confidence level contours. The first model will be compared with observational data and residuals shown. ``outname_fit_modelN.png`` (see `naima.plot_fit` and `naima.plot_blob`). Parameters ---------- outname : str Name to be used to save diagnostic plot files. sampler : `emcee.EnsembleSampler` instance Sampler instance from which chains, blobs and data are read. modelidxs : iterable of integers, optional Model numbers to be plotted. Default: All returned in sampler.get_blobs blob_labels : list of strings, optional Label for each of the outputs of the model. They will be used as title for the corresponding plot. pdf : bool, optional Whether to save plots to multipage pdf. """ from matplotlib import pyplot as plt from .plot import plot_blob, plot_chain, plot_corner # This function should never be interactive old_interactive = plt.rcParams["interactive"] plt.rcParams["interactive"] = False if pdf: plt.rc("pdf", fonttype=42) "Saving diagnostic plots in file " "{0}_plots.pdf".format(outname) ) from matplotlib.backends.backend_pdf import PdfPages outpdf = PdfPages("{0}_plots.pdf".format(outname)) # Chains for par, label in zip( range(sampler.get_chain().shape[-1]), sampler.labels ): try:"Plotting chain of parameter {0}...".format(label)) f = plot_chain(sampler, par, last_step=last_step) if pdf: f.savefig(outpdf, format="pdf", dpi=dpi) else: if "log(" in label or "log10(" in label: label = label.split("(")[-1].split(")")[0] f.savefig("{0}_chain_{1}.png".format(outname, label), dpi=dpi) f.clf() plt.close(f) except Exception as e: log.warning( "plot_chain failed for paramter" " {0} ({1}): {2}".format(label, par, e) ) # Corner plot"Plotting corner plot...") f = plot_corner(sampler) if f is not None: if pdf: f.savefig(outpdf, format="pdf", dpi=dpi) else: f.savefig("{0}_corner.png".format(outname), dpi=dpi) f.clf() plt.close(f) # Fit if modelidxs is None: nmodels = len(sampler.get_blobs()[-1][0]) modelidxs = list(range(nmodels)) if isinstance(sed, bool): sed = [sed for idx in modelidxs] if blob_labels is None: blob_labels = ["Model output {0}".format(idx) for idx in modelidxs] elif len(modelidxs) == 1 and isinstance(blob_labels, str): blob_labels = [blob_labels] elif len(blob_labels) < len(modelidxs): # Add labels n = len(blob_labels) blob_labels += [ "Model output {0}".format(idx) for idx in modelidxs[n:] ] for modelidx, plot_sed, label in zip(modelidxs, sed, blob_labels): try:"Plotting {0}...".format(label)) f = plot_blob( sampler, blobidx=modelidx, label=label, sed=plot_sed, n_samples=100, last_step=last_step, ) if pdf: f.savefig(outpdf, format="pdf", dpi=dpi) else: f.savefig( "{0}_model{1}.png".format(outname, modelidx), dpi=dpi ) f.clf() plt.close(f) except Exception as e: log.warning("plot_blob failed for {0}: {1}".format(label, e)) if pdf: outpdf.close() # set interactive back to original plt.rcParams["interactive"] = old_interactive
[docs]def save_results_table( outname, sampler, format="ascii.ecsv", convert_log=True, last_step=False, include_blobs=True, ): """ Save an ASCII table with the results stored in the `~emcee.EnsembleSampler`. The table contains the median, 16th and 84th percentile confidence region (~1sigma) for each parameter. Parameters ---------- outname : str Root name to be used to save the table. ``_results.dat`` will be appended for the output filename. sampler : `emcee.EnsembleSampler` instance Sampler instance from which chains, blobs and data are read. format : str, optional Format of the saved table. Must be a format string accepted by `astropy.table.Table.write`, see the `astropy unified file read/write interface documentation <>`_. Only the ``ascii.ecsv`` and ``ascii.ipac`` formats are able to preserve all the information stored in the ``run_info`` dictionary of the sampler. Defaults to ``ascii.ecsv`` if available (only in astropy > v1.0), else ``ascii.ipac``. convert_log : bool, optional Whether to convert natural or base-10 logarithms into original values in addition to saving the logarithm value. last_step : bool, optional Whether to only use the positions in the final step of the run (True, default) or the whole chain (False). include_blobs : bool, optional Whether to save the distribution properties of the scalar blobs in the sampler. Default is True. Returns ------- table : `~astropy.table.Table` Table with the results. """ if not HAS_PYYAML and format == "ascii.ecsv": format = "ascii.ipac" log.warning( "PyYAML package is required for ECSV format," " falling back to {0}...".format(format) ) elif format not in ["ascii.ecsv", "ascii.ipac"]: log.warning( "The chosen table format does not have an astropy" " writer that suppports metadata writing, no run info" " will be saved to the file!" ) file_extension = "dat" if format == "ascii.ecsv": file_extension = "ecsv" "Saving results table in {0}_results.{1}".format( outname, file_extension ) ) labels = sampler.labels if last_step: dists = sampler.get_chain()[:, -1, :] else: dists = sampler.get_chain(flat=True) quant = [16, 50, 84] # Do we need more info on the distributions? t = Table( names=["label", "median", "unc_lo", "unc_hi"], dtype=["S72", "f8", "f8", "f8"], ) t["label"].description = "Name of the parameter" t["median"].description = "Median of the posterior distribution function" t["unc_lo"].description = ( "Difference between the median and the" " {0}th percentile of the pdf, ~1sigma lower uncertainty".format( quant[0] ) ) t["unc_hi"].description = ( "Difference between the {0}th percentile" " and the median of the pdf, ~1sigma upper uncertainty".format( quant[2] ) ) metadata = {} # Start with info from the distributions used for storing the results metadata["n_samples"] = dists.shape[0] # save ML parameter vector and best/median loglikelihood ML, MLp, MLerr, _ = find_ML(sampler, None) metadata["ML_pars"] = [float(p) for p in MLp] metadata["MaxLogLikelihood"] = float(ML) # compute and save BIC BIC = len(MLp) * np.log(len( - 2 * ML metadata["BIC"] = BIC # And add all info stored in the sampler.run_info dict if hasattr(sampler, "run_info"): metadata.update(sampler.run_info) for p, label in enumerate(labels): dist = dists[:, p] xquant = np.percentile(dist, quant) quantiles = dict(zip(quant, xquant)) med = quantiles[50] lo, hi = med - quantiles[16], quantiles[84] - med t.add_row((label, med, lo, hi)) if convert_log and ("log10(" in label or "log(" in label): nlabel = label.split("(")[-1].split(")")[0] ltype = label.split("(")[0] if ltype == "log10": new_dist = 10 ** dist elif ltype == "log": new_dist = np.exp(dist) quantiles = dict(zip(quant, np.percentile(new_dist, quant))) med = quantiles[50] lo, hi = med - quantiles[16], quantiles[84] - med t.add_row((nlabel, med, lo, hi)) if include_blobs: blobs = sampler.get_blobs() nblobs = len(blobs[-1][0]) for idx in range(nblobs): blob0 = blobs[-1][0][idx] IS_SCALAR = False if isinstance(blob0, u.Quantity): if blob0.size == 1: IS_SCALAR = True unit = blob0.unit elif np.isscalar(blob0): IS_SCALAR = True unit = None if IS_SCALAR: if last_step: blobl = [m[idx] for m in blobs[-1]] else: blobl = [] for step in blobs: for walkerblob in step: blobl.append(walkerblob[idx]) if unit: dist = np.array([b.value for b in blobl]) metadata["blob{0}_unit".format(idx)] = unit.to_string() else: dist = np.array(blobl) quantiles = dict(zip(quant, np.percentile(dist, quant))) med = quantiles[50] lo, hi = med - quantiles[16], quantiles[84] - med t.add_row(("blob{0}".format(idx), med, lo, hi)) if format == "ascii.ipac": # Only keywords are written to IPAC tables t.meta["keywords"] = {} for di in metadata.items(): t.meta["keywords"][di[0]] = {"value": di[1]} else: if format == "ascii.ecsv": # there can be no numpy arrays in the metadata (YAML doesn't like # them) for di in list(metadata.items()): if type(di[1]).__module__ == np.__name__: try: # convert arrays metadata[di[0]] = [a.item() for a in di[1]] except TypeError: # convert scalars metadata[di[0]] = di[1].item() # Save it directly in meta for readability in ECSV t.meta.update(metadata) t.write("{0}_results.{1}".format(outname, file_extension), format=format) return t
[docs]def save_run(filename, sampler, compression=True, clobber=False): """ Save the sampler chain, data table, parameter labels, metadata blobs, and run information to a hdf5 file. The data table and parameter labels stored in the sampler will also be saved to the hdf5 file. Parameters ---------- filename : str Filename for hdf5 file. If the filename extension is not 'h5' or 'hdf5', the suffix '_chain.h5' will be appended to the filename. sampler : `emcee.EnsembleSampler` instance Sampler instance for which chain and run information is saved. compression : bool, optional Whether gzip compression is applied to the dataset on write. Default is True. clobber : bool, optional Whether to overwrite the output filename if it exists. """ filename = Path(filename) if filename.suffix not in {".hdf5", ".h5"}: raise ValueError("Filename must end in .hdf5 or .h5 suffix") if os.path.exists(filename) and not clobber: log.warning( "Not writing file because file exists and clobber is False" ) return with h5py.File(filename, "w") as f: group = f.create_group("mcmc") group.create_dataset( "chain", data=sampler.get_chain(), compression=compression ) group.create_dataset( "log_prob", data=sampler.get_log_prob(), compression=compression, ) # blobs blob = sampler.get_blobs()[-1][0] for idx, item in enumerate(blob): if isinstance(item, u.Quantity): # scalar or array quantity units = [item.unit.to_string()] elif isinstance(item, float): units = [""] elif ( isinstance(item, tuple) or isinstance(item, list) ) and np.all([isinstance(x, np.ndarray) for x in item]): units = [] for x in item: if isinstance(x, u.Quantity): units.append(x.unit.to_string()) else: units.append("") else: log.warning( "blob number {0} has unknown format and cannot be saved " "in HDF5 file" ) continue # traverse blobs list. This will probably be slow and there should # be a better way blob = [] for step in sampler.get_blobs(): for walkerblob in step: blob.append(walkerblob[idx]) blob = u.Quantity(blob).value blobdataset = group.create_dataset( "blob{0}".format(idx), data=blob, compression=compression ) if len(units) > 1: for j, unit in enumerate(units): blobdataset.attrs["unit{0}".format(j)] = unit else: blobdataset.attrs["unit"] = units[0] write_table_hdf5(, group, path="data", serialize_meta=True, compression=compression, ) # add all run info to group attributes if hasattr(sampler, "run_info"): for key in sampler.run_info.keys(): val = sampler.run_info[key] try: group.attrs[key] = val except TypeError: group.attrs[key] = str(val) # add other sampler info to the attrs group.attrs["acceptance_fraction"] = np.mean( sampler.acceptance_fraction ) # add labels as individual attrs (there might be a better way) for i, label in enumerate(sampler.labels): group.attrs["label{0}".format(i)] = label
class _result: """ Minimal emcee.EnsembleSampler like container for chain results """ def get_value(self, name, flat=False): v = getattr(self, name) if flat: s = list(v.shape[1:]) s[0] =[:2]) return v.reshape(s) return v def get_chain(self, **kwargs): return self.get_value("chain", **kwargs) def get_log_prob(self, **kwargs): return self.get_value("log_prob", **kwargs) def get_blobs(self, **kwargs): return self.get_value("_blobs", **kwargs)
[docs]def read_run(filename, modelfn=None): """ Read chain from a hdf5 saved with `save_run`. This function will also read the labels, data table, and metadata blobs stored in the original sampler. If you want to use the result object with `plot_fit` and setting the ``e_range`` parameter, you must provide the model function with the `modelfn` argument given that functions cannot be serialized in hdf5 files. Parameters ---------- filename : str Filename of the hdf5 containing the chain, lnprobability, and blob arrays in the group 'sampler' modelfn : function, optional Model function to be attached to the returned sampler Returns ------- result : class Container object with same properties as an `emcee.EnsembleSampler` resulting from a sampling run. This object can be passed onto `~naima.plot_fit`, `~naima.plot_chain`, and `~naima.plot_corner` for analysis as you would do with a `emcee.EnsembleSampler` instance. """ # initialize empty sampler class to return result = _result() result.modelfn = modelfn result.run_info = {} f = h5py.File(filename, "r") # chain and lnprobability result.chain = np.array(f["mcmc/chain"]) result.log_prob = np.array(f["mcmc/log_prob"]) # blobs result_blobs = [] nsteps, nwalkers, nblobs = result.chain.shape blobs = [] blobrank = [] for i in range(100): # first read each of the blobs and convert to Quantities try: ds = f["mcmc/blob{0}".format(i)] rank = np.ndim(ds[0]) blobrank.append(rank) if rank <= 1: blobs.append(u.Quantity(ds[()], unit=ds.attrs["unit"])) else: blob = [] for j in range(np.ndim(ds[0])): blob.append( u.Quantity( ds[:, j, :], unit=ds.attrs["unit{0}".format(j)] ) ) blobs.append(blob) except KeyError: break # Now organize in an awful list of lists of arrays for step in range(nsteps): steplist = [] for walker in range(nwalkers): n = step * nwalkers + walker walkerblob = [] for j in range(len(blobs)): if blobrank[j] <= 1: walkerblob.append(blobs[j][n]) else: blob = [] for k in range(blobrank[j]): blob.append(blobs[j][k][n]) walkerblob.append(blob) steplist.append(walkerblob) result_blobs.append(steplist) result._blobs = np.array(result_blobs, dtype=np.dtype("object")) # run info result.run_info = dict(f["mcmc"].attrs) result.acceptance_fraction = f["mcmc"].attrs["acceptance_fraction"] # labels result.labels = [] for i in range(nblobs): result.labels.append(f["mcmc"].attrs["label{0}".format(i)]) # data = read_table_hdf5(f["mcmc"], "data") return result