# Licensed under a 3-clause BSD style license - see LICENSE.rst
import astropy.units as u
import numpy as np
from astropy import log
from emcee import autocorr
from .extern.interruptible_pool import InterruptiblePool as Pool
from .extern.validator import validate_array
from .utils import sed_conversion, validate_data_table
__all__ = ["plot_chain", "plot_fit", "plot_data", "plot_blob", "plot_corner"]
marker_cycle = ["o", "s", "d", "p", "*"]
# from seaborn: sns.color_palette('deep',6)
color_cycle = [
(0.2980392156862745, 0.4470588235294118, 0.6901960784313725),
(0.3333333333333333, 0.6588235294117647, 0.40784313725490196),
(0.7686274509803922, 0.3058823529411765, 0.3215686274509804),
(0.5058823529411764, 0.4470588235294118, 0.6980392156862745),
(0.8, 0.7254901960784313, 0.4549019607843137),
(0.39215686274509803, 0.7098039215686275, 0.803921568627451),
]
[docs]def plot_chain(sampler, p=None, **kwargs):
"""Generate a diagnostic plot of the sampler chains.
Parameters
----------
sampler : `emcee.EnsembleSampler`
Sampler containing the chains to be plotted.
p : int (optional)
Index of the parameter to plot. If omitted, all chains are plotted.
last_step : bool (optional)
Whether to plot the last step of the chain or the complete chain
(default).
Returns
-------
figure : `matplotlib.figure.Figure`
Figure
"""
if p is None:
npars = sampler.get_chain().shape[-1]
for pp in range(npars):
_plot_chain_func(sampler, pp, **kwargs)
fig = None
else:
fig = _plot_chain_func(sampler, p, **kwargs)
return fig
def _latex_float(f, format=".3g"):
"""http://stackoverflow.com/a/13490601"""
float_str = "{{0:{0}}}".format(format).format(f)
if "e" in float_str:
base, exponent = float_str.split("e")
return r"{0}\times 10^{{{1}}}".format(base, int(exponent))
else:
return float_str
def round2(x, n):
y = round(x, n)
if n < 1:
y = str(int(y))
else:
# preserve trailing zeroes
y = ("{{0:.{0}f}}".format(n)).format(x)
return y
def _latex_value_error(val, elo, ehi=0, tol=0.25):
order = int(np.log10(abs(val)))
if order > 2 or order < -2:
val /= 10 ** order
elo /= 10 ** order
ehi /= 10 ** order
else:
order = 0
nlo = -int(np.floor(np.log10(elo)))
if elo * 10 ** nlo < 2:
nlo += 1
if ehi:
# elo = round(elo,nlo)
nhi = -int(np.floor(np.log10(ehi)))
if ehi * 10 ** nhi < 2:
nhi += 1
# ehi = round(ehi,nhi)
if np.abs(elo - ehi) / ((elo + ehi) / 2.0) > tol:
n = max(nlo, nhi)
string = "{0}^{{+{1}}}_{{-{2}}}".format(
*[
round2(x, nn)
for x, nn in zip([val, ehi, elo], [n, nhi, nlo])
]
)
else:
e = (elo + ehi) / 2.0
n = -int(np.floor(np.log10(e)))
if e * 10 ** n < 2:
n += 1
string = "{0} \\pm {1}".format(*[round2(x, n) for x in [val, e]])
else:
string = "{0} \\pm {1}".format(*[round2(x, nlo) for x in [val, elo]])
if order != 0:
string = "(" + string + r")\times10^{{{0}}}".format(order)
return "$" + string + "$"
def _plot_chain_func(sampler, p, last_step=False):
chain = sampler.get_chain()
label = sampler.labels[p]
import matplotlib.pyplot as plt
from scipy import stats
if len(chain.shape) > 2:
# transpose from (step, walker) to (walker, step)
traces = chain[:, :, p].T
if last_step:
# keep only last step
dist = traces[:, -1]
else:
# convert chain to flatchain
dist = traces.flatten()
else:
log.warning(
"we need the full chain to plot the traces, not a flatchain!"
)
return None
nwalkers = traces.shape[0]
nsteps = traces.shape[1]
f = plt.figure()
ax1 = f.add_subplot(221)
ax2 = f.add_subplot(122)
f.subplots_adjust(left=0.1, bottom=0.15, right=0.95, top=0.9)
# plot five percent of the traces darker
if nwalkers < 60:
thresh = 1 - 3.0 / nwalkers
else:
thresh = 0.95
red = np.arange(nwalkers) / float(nwalkers) >= thresh
ax1.set_rasterization_zorder(1)
for t in traces[~red]: # range(nwalkers):
ax1.plot(t, color=(0.1,) * 3, lw=1.0, alpha=0.25, zorder=0)
for t in traces[red]:
ax1.plot(t, color=color_cycle[0], lw=1.5, alpha=0.75, zorder=0)
ax1.set_xlabel("step number")
# [l.set_rotation(45) for l in ax1.get_yticklabels()]
ax1.set_ylabel(label)
ax1.yaxis.set_label_coords(-0.15, 0.5)
ax1.set_title("Walker traces")
nbins = min(max(25, int(len(dist) / 100.0)), 100)
xlabel = label
n, x, _ = ax2.hist(
dist,
nbins,
histtype="stepfilled",
color=color_cycle[0],
lw=0,
density=True,
)
kde = stats.kde.gaussian_kde(dist)
ax2.plot(x, kde(x), color="k", label="KDE")
quant = [16, 50, 84]
xquant = np.percentile(dist, quant)
quantiles = dict(zip(quant, xquant))
ax2.axvline(
quantiles[50],
ls="--",
color="k",
alpha=0.5,
lw=2,
label="50% quantile",
)
ax2.axvspan(
quantiles[16],
quantiles[84],
color=(0.5,) * 3,
alpha=0.25,
label="68% CI",
lw=0,
)
# ax2.legend()
for xticklabel in ax2.get_xticklabels():
xticklabel.set_rotation(45)
ax2.set_xlabel(xlabel)
ax2.xaxis.set_label_coords(0.5, -0.1)
ax2.set_title("posterior distribution")
ax2.set_ylim(top=n.max() * 1.05)
# Print distribution parameters on lower-left
try:
autocorr_message = "{0:.1f}".format(autocorr.integrated_time(chain)[p])
except autocorr.AutocorrError:
# Raised when chain is too short for meaningful auto-correlation
# estimation
autocorr_message = None
if last_step:
clen = "last ensemble"
else:
clen = "whole chain"
chain_props = "Walkers: {0} \nSteps in chain: {1} \n".format(
nwalkers, nsteps
)
if autocorr_message is not None:
chain_props += "Autocorrelation time: {0}\n".format(autocorr_message)
chain_props += "Mean acceptance fraction: {0:.3f}\n".format(
np.mean(sampler.acceptance_fraction)
) + "Distribution properties for the {clen}:\n \
$-$ median: ${median}$, std: ${std}$ \n \
$-$ median with uncertainties based on \n \
the 16th and 84th percentiles ($\\sim$1$\\sigma$):\n".format(
median=_latex_float(quantiles[50]),
std=_latex_float(np.std(dist)),
clen=clen,
)
info_line = (
" " * 10
+ label
+ " = "
+ _latex_value_error(
quantiles[50],
quantiles[50] - quantiles[16],
quantiles[84] - quantiles[50],
)
)
chain_props += info_line
if "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)
quant = [16, 50, 84]
quantiles = dict(zip(quant, np.percentile(new_dist, quant)))
label_template = "\n" + " " * 10 + "{{label:>{0}}}".format(len(label))
new_line = label_template.format(label=nlabel)
new_line += " = " + _latex_value_error(
quantiles[50],
quantiles[50] - quantiles[16],
quantiles[84] - quantiles[50],
)
chain_props += new_line
info_line += new_line
log.info("{0:-^50}\n".format(label) + info_line)
f.text(0.05, 0.45, chain_props, ha="left", va="top")
return f
def _process_blob(sampler, modelidx, last_step=False, energy=None):
"""
Process binary blob in sampler. If blob in position modelidx is:
- a Quantity array of len(blob[i])=len(data['energy']: use blob as model,
data['energy'] as modelx
- a tuple: use first item as modelx, second as model
- a Quantity scalar: return array of scalars
"""
# Allow process blob to be used by _calc_samples and _calc_ML by sending
# only blobs, not full sampler
try:
blobs = sampler.get_blobs()
blob0 = blobs[-1][0][modelidx]
energy = sampler.data["energy"]
except AttributeError:
blobs = [sampler]
blob0 = sampler[0][modelidx]
last_step = True
if isinstance(blob0, u.Quantity):
if blob0.size == energy.size:
# Energy array for blob is not provided, use data['energy']
modelx = energy
elif blob0.size == 1:
modelx = None
if last_step:
model = u.Quantity([m[modelidx] for m in blobs[-1]])
else:
model = []
for step in blobs:
for walkerblob in step:
model.append(walkerblob[modelidx])
model = u.Quantity(model)
elif np.isscalar(blob0):
modelx = None
if last_step:
model = u.Quantity([m[modelidx] for m in blobs[-1]])
else:
model = []
for step in blobs:
for walkerblob in step:
model.append(walkerblob[modelidx])
model = u.Quantity(model)
elif isinstance(blob0, list) or isinstance(blob0, tuple):
if (
len(blob0) == 2
and isinstance(blob0[0], u.Quantity)
and isinstance(blob0[1], u.Quantity)
):
# Energy array for model is item 0 in blob, model flux is item 1
modelx = blob0[0]
if last_step:
model = u.Quantity([m[modelidx][1] for m in blobs[-1]])
else:
model = []
for step in blobs:
for walkerblob in step:
model.append(walkerblob[modelidx][1])
model = u.Quantity(model)
else:
raise TypeError("Model {0} has wrong blob format".format(modelidx))
else:
raise TypeError("Model {0} has wrong blob format".format(modelidx))
return modelx, model
def _read_or_calc_samples(
sampler,
modelidx=0,
n_samples=100,
last_step=False,
e_range=None,
e_npoints=100,
threads=None,
):
"""Get samples from blob or compute them from chain and sampler.modelfn"""
if e_range is None:
# return the results saved in blobs
modelx, model = _process_blob(sampler, modelidx, last_step=last_step)
else:
# prepare bogus data for calculation
e_range = validate_array(
"e_range", u.Quantity(e_range), physical_type="energy"
)
e_unit = e_range.unit
energy = (
np.logspace(
np.log10(e_range[0].value),
np.log10(e_range[1].value),
e_npoints,
)
* e_unit
)
data = {
"energy": energy,
"flux": np.zeros(energy.shape) * sampler.data["flux"].unit,
}
# init pool and select parameters
chain = (
sampler.get_chain()[-1]
if last_step
else sampler.get_chain(flat=True)
)
pars = chain[np.random.randint(len(chain), size=n_samples)]
args = ((p, data) for p in pars)
blobs = []
pool = Pool(threads)
modelouts = pool.starmap(sampler.modelfn, args)
pool.close()
pool.terminate()
for modelout in modelouts:
if isinstance(modelout, np.ndarray):
blobs.append([modelout])
else:
blobs.append(modelout)
modelx, model = _process_blob(
blobs, modelidx=modelidx, energy=data["energy"]
)
return modelx, model
def _calc_ML(sampler, modelidx=0, e_range=None, e_npoints=100):
"""Get ML model from blob or compute them from chain and sampler.modelfn"""
ML, MLp, MLerr, ML_model = find_ML(sampler, modelidx)
if e_range is not None:
# prepare bogus data for calculation
e_range = validate_array(
"e_range", u.Quantity(e_range), physical_type="energy"
)
e_unit = e_range.unit
energy = (
np.logspace(
np.log10(e_range[0].value),
np.log10(e_range[1].value),
e_npoints,
)
* e_unit
)
data = {
"energy": energy,
"flux": np.zeros(energy.shape) * sampler.data["flux"].unit,
}
modelout = sampler.modelfn(MLp, data)
if isinstance(modelout, np.ndarray):
blob = modelout
else:
blob = modelout[modelidx]
if isinstance(blob, u.Quantity):
modelx = data["energy"].copy()
model_ML = blob.copy()
elif len(blob) == 2:
modelx = blob[0].copy()
model_ML = blob[1].copy()
else:
raise TypeError("Model {0} has wrong blob format".format(modelidx))
ML_model = (modelx, model_ML)
return ML, MLp, MLerr, ML_model
def _calc_CI(
sampler,
modelidx=0,
confs=[3, 1],
last_step=False,
e_range=None,
e_npoints=100,
threads=None,
):
"""Calculate confidence interval."""
from scipy import stats
# If we are computing the samples for the confidence intervals, we need at
# least one sample to constrain the highest confidence band
# 1 sigma -> 6 samples
# 2 sigma -> 43 samples
# 3 sigma -> 740 samples
# 4 sigma -> 31574 samples
# 5 sigma -> 3488555 samples
# We limit it to 1000 samples and warn that it might not be enough
if e_range is not None:
maxconf = np.max(confs)
minsamples = min(100, int(1 / stats.norm.cdf(-maxconf) + 1))
if minsamples > 1000:
log.warning(
"In order to sample the confidence band for {0} sigma,"
" {1} new samples need to be computed, but we are limiting"
" it to 1000 samples, so the confidence band might not be"
" well constrained."
" Consider reducing the maximum"
" confidence significance or using the samples stored in"
" the sampler by setting e_range"
" to None".format(maxconf, minsamples)
)
minsamples = 1000
else:
minsamples = None
modelx, model = _read_or_calc_samples(
sampler,
modelidx,
last_step=last_step,
e_range=e_range,
e_npoints=e_npoints,
n_samples=minsamples,
threads=threads,
)
nwalkers = len(model) - 1
CI = []
for conf in confs:
fmin = stats.norm.cdf(-conf)
fmax = stats.norm.cdf(conf)
ymin, ymax = [], []
for fr, y in ((fmin, ymin), (fmax, ymax)):
nf = int((fr * nwalkers))
for i in range(len(modelx)):
ysort = np.sort(model[:, i])
y.append(ysort[nf])
# create an array from lists ymin and ymax preserving units
CI.append((u.Quantity(ymin), u.Quantity(ymax)))
return modelx, CI
def _plot_MLmodel(ax, sampler, modelidx, e_range, e_npoints, e_unit, sed):
"""compute and plot ML model"""
ML, MLp, MLerr, ML_model = _calc_ML(
sampler, modelidx, e_range=e_range, e_npoints=e_npoints
)
f_unit, sedf = sed_conversion(ML_model[0], ML_model[1].unit, sed)
ax.loglog(
ML_model[0].to(e_unit).value,
(ML_model[1] * sedf).to(f_unit).value,
color="k",
lw=2,
alpha=0.8,
)
def plot_CI(
ax,
sampler,
modelidx=0,
sed=True,
confs=[3, 1, 0.5],
e_unit=u.eV,
label=None,
e_range=None,
e_npoints=100,
threads=None,
last_step=False,
):
"""Plot confidence interval.
Parameters
----------
ax : `matplotlib.Axes`
Axes to plot on.
sampler : `emcee.EnsembleSampler`
Sampler
modelidx : int, optional
Model index. Default is 0
sed : bool, optional
Whether to plot SED or differential spectrum. If `None`, the units of
the observed spectrum will be used.
confs : list, optional
List of confidence levels (in sigma) to use for generating the
confidence intervals. Default is `[3,1,0.5]`
e_unit : :class:`~astropy.units.Unit` or str parseable to unit
Unit in which to plot energy axis.
e_npoints : int, optional
How many points to compute for the model samples and ML model if
`e_range` is set.
threads : int, optional
How many parallel processing threads to use when computing the samples.
Defaults to the number of available cores.
last_step : bool, optional
Whether to only use the positions in the final step of the run (True,
default) or the whole chain (False).
"""
confs.sort(reverse=True)
modelx, CI = _calc_CI(
sampler,
modelidx=modelidx,
confs=confs,
e_range=e_range,
e_npoints=e_npoints,
last_step=last_step,
threads=threads,
)
# pick first confidence interval curve for units
f_unit, sedf = sed_conversion(modelx, CI[0][0].unit, sed)
for (ymin, ymax), conf in zip(CI, confs):
color = np.log(conf) / np.log(20) + 0.4
ax.fill_between(
modelx.to(e_unit).value,
(ymax * sedf).to(f_unit).value,
(ymin * sedf).to(f_unit).value,
lw=0.001,
color=(color,) * 3,
alpha=0.6,
zorder=-10,
)
_plot_MLmodel(ax, sampler, modelidx, e_range, e_npoints, e_unit, sed)
if label is not None:
ax.set_ylabel(
"{0} [{1}]".format(label, f_unit.to_string("latex_inline"))
)
def plot_samples(
ax,
sampler,
modelidx=0,
sed=True,
n_samples=100,
e_unit=u.eV,
e_range=None,
e_npoints=100,
threads=None,
label=None,
last_step=False,
):
"""Plot a number of samples from the sampler chain.
Parameters
----------
ax : `matplotlib.Axes`
Axes to plot on.
sampler : `emcee.EnsembleSampler`
Sampler
modelidx : int, optional
Model index. Default is 0
sed : bool, optional
Whether to plot SED or differential spectrum. If `None`, the units of
the observed spectrum will be used.
n_samples : int, optional
Number of samples to plot. Default is 100.
e_unit : :class:`~astropy.units.Unit` or str parseable to unit
Unit in which to plot energy axis.
e_range : list of `~astropy.units.Quantity`, length 2, optional
Limits in energy for the computation of the model samples and ML model.
Note that setting this parameter will mean that the samples for the
model are recomputed and depending on the model speed might be quite
slow.
e_npoints : int, optional
How many points to compute for the model samples and ML model if
`e_range` is set.
threads : int, optional
How many parallel processing threads to use when computing the samples.
Defaults to the number of available cores.
last_step : bool, optional
Whether to only use the positions in the final step of the run (True,
default) or the whole chain (False).
"""
modelx, model = _read_or_calc_samples(
sampler,
modelidx,
last_step=last_step,
e_range=e_range,
e_npoints=e_npoints,
threads=threads,
)
# pick first model sample for units
f_unit, sedf = sed_conversion(modelx, model[0].unit, sed)
sample_alpha = min(5.0 / n_samples, 0.5)
for my in model[np.random.randint(len(model), size=n_samples)]:
ax.loglog(
modelx.to(e_unit).value,
(my * sedf).to(f_unit).value,
color=(0.1,) * 3,
alpha=sample_alpha,
lw=1.0,
)
_plot_MLmodel(ax, sampler, modelidx, e_range, e_npoints, e_unit, sed)
if label is not None:
ax.set_ylabel(
"{0} [{1}]".format(label, f_unit.to_string("latex_inline"))
)
def find_ML(sampler, modelidx):
"""
Find Maximum Likelihood parameters as those in the chain with a highest log
probability.
"""
lnprobability = sampler.get_log_prob()
index = np.unravel_index(np.argmax(lnprobability), lnprobability.shape)
MLp = sampler.get_chain()[index]
blobs = sampler.get_blobs()
if modelidx is not None and blobs is not None:
blob = blobs[index][modelidx]
if isinstance(blob, u.Quantity):
modelx = sampler.data["energy"].copy()
model_ML = blob.copy()
elif len(blob) == 2:
modelx = blob[0].copy()
model_ML = blob[1].copy()
else:
raise TypeError("Model {0} has wrong blob format".format(modelidx))
elif modelidx is not None and hasattr(sampler, "modelfn"):
blob = _process_blob(
[sampler.modelfn(MLp, sampler.data)],
modelidx,
energy=sampler.data["energy"],
)
modelx, model_ML = blob[0], blob[1][0]
else:
modelx, model_ML = None, None
MLerr = []
for dist in sampler.get_chain(flat=True).T:
hilo = np.percentile(dist, [16.0, 84.0])
MLerr.append((hilo[1] - hilo[0]) / 2.0)
ML = lnprobability[index]
return ML, MLp, MLerr, (modelx, model_ML)
[docs]def plot_blob(
sampler, blobidx=0, label=None, last_step=False, figure=None, **kwargs
):
"""
Plot a metadata blob as a fit to spectral data or value distribution
Additional ``kwargs`` are passed to `plot_fit`.
Parameters
----------
sampler : `emcee.EnsembleSampler`
Sampler with a stored chain.
blobidx : int, optional
Metadata blob index to plot.
label : str, optional
Label for the value distribution. Labels for the fit plot can be passed
as ``xlabel`` and ``ylabel`` and will be passed to `plot_fit`.
Returns
-------
figure : `matplotlib.pyplot.Figure`
`matplotlib` figure instance containing the plot.
"""
modelx, model = _process_blob(sampler, blobidx, last_step)
if label is None:
label = "Model output {0}".format(blobidx)
if modelx is None:
# Blob is scalar, plot distribution
f = plot_distribution(model, label, figure=figure)
else:
f = plot_fit(
sampler,
modelidx=blobidx,
last_step=last_step,
label=label,
figure=figure,
**kwargs
)
return f
[docs]def plot_fit(
sampler,
modelidx=0,
label=None,
sed=True,
last_step=False,
n_samples=100,
confs=None,
ML_info=False,
figure=None,
plotdata=None,
plotresiduals=None,
e_unit=None,
e_range=None,
e_npoints=100,
threads=None,
xlabel=None,
ylabel=None,
ulim_opts={},
errorbar_opts={},
):
"""
Plot data with fit confidence regions.
Parameters
----------
sampler : `emcee.EnsembleSampler`
Sampler with a stored chain.
modelidx : int, optional
Model index to plot.
label : str, optional
Label for the title of the plot.
sed : bool, optional
Whether to plot SED or differential spectrum.
last_step : bool, optional
Whether to use only the samples of the last step in the run when
showing either the model samples or the confidence intervals.
n_samples : int, optional
If not ``None``, number of sample models to plot. If ``None``,
confidence bands will be plotted instead of samples. Default is 100.
confs : list, optional
List of confidence levels (in sigma) to use for generating the
confidence intervals. Default is to plot sample models instead of
confidence bands.
ML_info : bool, optional
Whether to plot information about the maximum likelihood parameters and
the standard deviation of their distributions. Default is True.
figure : `matplotlib.figure.Figure`, optional
`matplotlib` figure to plot on. If omitted a new one will be generated.
plotdata : bool, optional
Wheter to plot data on top of model confidence intervals. Default is
True if the physical types of the data and the model match.
plotresiduals : bool, optional
Wheter to plot the residuals with respect to the maximum likelihood
model. Default is True if ``plotdata`` is True and either ``confs`` or
``n_samples`` are set.
e_unit : `~astropy.units.Unit`, optional
Units for the energy axis of the plot. The default is to use the units
of the energy array of the observed data.
e_range : list of `~astropy.units.Quantity`, length 2, optional
Limits in energy for the computation of the model samples and ML model.
Note that setting this parameter will mean that the samples for the
model are recomputed and depending on the model speed might be quite
slow.
e_npoints : int, optional
How many points to compute for the model samples and ML model if
`e_range` is set.
threads : int, optional
How many parallel processing threads to use when computing the samples.
Defaults to the number of available cores.
xlabel : str, optional
Label for the ``x`` axis of the plot.
ylabel : str, optional
Label for the ``y`` axis of the plot.
ulim_opts : dict
Option for upper-limit plotting. Available options are capsize (arrow
width) and height_fraction (arrow length in fraction of flux value).
errorbar_opts : dict
Addtional options to pass to `matplotlib.plt.errorbar` for plotting the
spectral flux points.
"""
import matplotlib.pyplot as plt
ML, MLp, MLerr, model_ML = find_ML(sampler, modelidx)
infostr = "Maximum log probability: {0:.3g}\n".format(ML)
infostr += "Maximum Likelihood values:\n"
maxlen = np.max([len(ilabel) for ilabel in sampler.labels])
vartemplate = "{{2:>{0}}}: {{0:>8.3g}} +/- {{1:<8.3g}}\n".format(maxlen)
for p, v, ilabel in zip(MLp, MLerr, sampler.labels):
infostr += vartemplate.format(p, v, ilabel)
# log.info(infostr)
data = sampler.data
if e_range is None and not hasattr(sampler, "blobs"):
e_range = data["energy"][[0, -1]] * np.array((1.0 / 3.0, 3.0))
if plotdata is None and len(model_ML[0]) == len(data["energy"]):
model_unit, _ = sed_conversion(model_ML[0], model_ML[1].unit, sed)
data_unit, _ = sed_conversion(data["energy"], data["flux"].unit, sed)
plotdata = model_unit.is_equivalent(data_unit)
elif plotdata is None:
plotdata = False
if plotresiduals is None and plotdata and (confs is not None or n_samples):
plotresiduals = True
if confs is None and not n_samples and plotdata and not plotresiduals:
# We actually only want to plot the data, so let's go there
return plot_data(
sampler.data,
xlabel=xlabel,
ylabel=ylabel,
sed=sed,
figure=figure,
e_unit=e_unit,
ulim_opts=ulim_opts,
errorbar_opts=errorbar_opts,
)
if figure is None:
f = plt.figure()
else:
f = figure
if plotdata and plotresiduals:
ax1 = plt.subplot2grid((4, 1), (0, 0), rowspan=3)
ax2 = plt.subplot2grid((4, 1), (3, 0), sharex=ax1)
for subp in [ax1, ax2]:
f.add_subplot(subp)
else:
ax1 = f.add_subplot(111)
if e_unit is None:
e_unit = data["energy"].unit
if confs is not None:
plot_CI(
ax1,
sampler,
modelidx,
sed=sed,
confs=confs,
e_unit=e_unit,
label=label,
e_range=e_range,
e_npoints=e_npoints,
last_step=last_step,
threads=threads,
)
elif n_samples:
plot_samples(
ax1,
sampler,
modelidx,
sed=sed,
n_samples=n_samples,
e_unit=e_unit,
label=label,
e_range=e_range,
e_npoints=e_npoints,
last_step=last_step,
threads=threads,
)
else:
# plot only ML model
_plot_MLmodel(ax1, sampler, modelidx, e_range, e_npoints, e_unit, sed)
xlaxis = ax1
if plotdata:
_plot_data_to_ax(
data,
ax1,
e_unit=e_unit,
sed=sed,
ylabel=ylabel,
ulim_opts=ulim_opts,
errorbar_opts=errorbar_opts,
)
if plotresiduals:
_plot_residuals_to_ax(
data,
model_ML,
ax2,
e_unit=e_unit,
sed=sed,
errorbar_opts=errorbar_opts,
)
xlaxis = ax2
for tl in ax1.get_xticklabels():
tl.set_visible(False)
xmin = 10 ** np.floor(
np.log10(
np.min(data["energy"] - data["energy_error_lo"])
.to(e_unit)
.value
)
)
xmax = 10 ** np.ceil(
np.log10(
np.max(data["energy"] + data["energy_error_hi"])
.to(e_unit)
.value
)
)
ax1.set_xlim(xmin, xmax)
else:
ax1.set_xscale("log")
ax1.set_yscale("log")
if sed:
ndecades = 10
else:
ndecades = 20
# restrict y axis to ndecades to avoid autoscaling deep exponentials
xmin, xmax, ymin, ymax = ax1.axis()
ymin = max(ymin, ymax / 10 ** ndecades)
ax1.set_ylim(bottom=ymin)
# scale x axis to largest model_ML x point within ndecades decades of
# maximum
f_unit, sedf = sed_conversion(model_ML[0], model_ML[1].unit, sed)
hi = np.where((model_ML[1] * sedf).to(f_unit).value > ymin)
xmax = np.max(model_ML[0][hi])
ax1.set_xlim(right=10 ** np.ceil(np.log10(xmax.to(e_unit).value)))
if e_range is not None:
# ensure that xmin/xmax contains e_range
xmin, xmax, ymin, ymax = ax1.axis()
xmin = min(xmin, e_range[0].to(e_unit).value)
xmax = max(xmax, e_range[1].to(e_unit).value)
ax1.set_xlim(xmin, xmax)
if ML_info and (confs is not None or n_samples):
ax1.text(
0.05,
0.05,
infostr,
ha="left",
va="bottom",
transform=ax1.transAxes,
family="monospace",
)
if label is not None:
ax1.set_title(label)
if xlabel is None:
xlaxis.set_xlabel(
"Energy [{0}]".format(e_unit.to_string("latex_inline"))
)
else:
xlaxis.set_xlabel(xlabel)
f.subplots_adjust(hspace=0)
return f
def _plot_ulims(
ax, x, y, xerr, color, capsize=5, height_fraction=0.25, elinewidth=2
):
"""
Plot upper limits as arrows with cap at value of upper limit.
uplim behaviour has been fixed in matplotlib 1.4
"""
ax.errorbar(
x, y, xerr=xerr, ls="", color=color, elinewidth=elinewidth, capsize=0
)
from distutils.version import LooseVersion
import matplotlib
mpl_version = LooseVersion(matplotlib.__version__)
if mpl_version >= LooseVersion("1.4.0"):
ax.errorbar(
x,
y,
yerr=height_fraction * y,
ls="",
uplims=True,
color=color,
elinewidth=elinewidth,
capsize=capsize,
zorder=10,
)
else:
ax.errorbar(
x,
(1 - height_fraction) * y,
yerr=height_fraction * y,
ls="",
lolims=True,
color=color,
elinewidth=elinewidth,
capsize=capsize,
zorder=10,
)
def _plot_data_to_ax(
data_all,
ax1,
e_unit=None,
sed=True,
ylabel=None,
ulim_opts={},
errorbar_opts={},
):
"""Plots data errorbars and upper limits onto ax.
X label is left to plot_data and plot_fit because they depend on whether
residuals are plotted.
"""
if e_unit is None:
e_unit = data_all["energy"].unit
f_unit, sedf = sed_conversion(
data_all["energy"], data_all["flux"].unit, sed
)
if "group" not in data_all.keys():
data_all["group"] = np.zeros(len(data_all))
groups = np.unique(data_all["group"])
for g in groups:
data = data_all[np.where(data_all["group"] == g)]
_, sedfg = sed_conversion(data["energy"], data["flux"].unit, sed)
# wrap around color and marker cycles
color = color_cycle[int(g) % len(color_cycle)]
marker = marker_cycle[int(g) % len(marker_cycle)]
ul = data["ul"]
notul = ~ul
# Hack to show y errors compatible with 0 in loglog plot
yerr_lo = data["flux_error_lo"][notul]
y = data["flux"][notul].to(yerr_lo.unit)
bad_err = np.where((y - yerr_lo) <= 0.0)
yerr_lo[bad_err] = y[bad_err] * (1.0 - 1e-7)
yerr = u.Quantity((yerr_lo, data["flux_error_hi"][notul]))
xerr = u.Quantity((data["energy_error_lo"], data["energy_error_hi"]))
opts = dict(
zorder=100,
marker=marker,
ls="",
elinewidth=2,
capsize=0,
mec=color,
mew=0.1,
ms=5,
color=color,
)
opts.update(**errorbar_opts)
ax1.errorbar(
data["energy"][notul].to(e_unit).value,
(data["flux"][notul] * sedfg[notul]).to(f_unit).value,
yerr=(yerr * sedfg[notul]).to(f_unit).value,
xerr=xerr[:, notul].to(e_unit).value,
**opts
)
if np.any(ul):
if "elinewidth" in errorbar_opts:
ulim_opts["elinewidth"] = errorbar_opts["elinewidth"]
_plot_ulims(
ax1,
data["energy"][ul].to(e_unit).value,
(data["flux"][ul] * sedfg[ul]).to(f_unit).value,
(xerr[:, ul]).to(e_unit).value,
color,
**ulim_opts
)
ax1.set_xscale("log")
ax1.set_yscale("log")
xmin = 10 ** np.floor(
np.log10(
np.min(data["energy"] - data["energy_error_lo"]).to(e_unit).value
)
)
xmax = 10 ** np.ceil(
np.log10(
np.max(data["energy"] + data["energy_error_hi"]).to(e_unit).value
)
)
ax1.set_xlim(xmin, xmax)
# avoid autoscaling to errorbars to 0
notul = ~data_all["ul"]
if np.any(data_all["flux_error_lo"][notul] >= data_all["flux"][notul]):
elo = (data_all["flux"][notul] * sedf[notul]).to(f_unit).value - (
data_all["flux_error_lo"][notul] * sedf[notul]
).to(f_unit).value
gooderr = np.where(
data_all["flux_error_lo"][notul] < data_all["flux"][notul]
)
ymin = 10 ** np.floor(np.log10(np.min(elo[gooderr])))
ax1.set_ylim(bottom=ymin)
if ylabel is None:
if sed:
ax1.set_ylabel(
r"$E^2\mathrm{{d}}N/\mathrm{{d}}E$"
" [{0}]".format(u.Unit(f_unit).to_string("latex_inline"))
)
else:
ax1.set_ylabel(
r"$\mathrm{{d}}N/\mathrm{{d}}E$"
" [{0}]".format(u.Unit(f_unit).to_string("latex_inline"))
)
else:
ax1.set_ylabel(ylabel)
def _plot_residuals_to_ax(
data_all, model_ML, ax, e_unit=u.eV, sed=True, errorbar_opts={}
):
"""Function to compute and plot residuals in units of the uncertainty"""
if "group" not in data_all.keys():
data_all["group"] = np.zeros(len(data_all))
groups = np.unique(data_all["group"])
MLf_unit, MLsedf = sed_conversion(model_ML[0], model_ML[1].unit, sed)
MLene = model_ML[0].to(e_unit)
MLflux = (model_ML[1] * MLsedf).to(MLf_unit)
ax.axhline(0, color="k", lw=1, ls="--")
interp = False
if data_all["energy"].size != MLene.size or not np.allclose(
data_all["energy"].value, MLene.value
):
interp = True
from scipy.interpolate import interp1d
modelfunc = interp1d(MLene.value, MLflux.value, bounds_error=False)
for g in groups:
groupidx = np.where(data_all["group"] == g)
data = data_all[groupidx]
notul = ~data["ul"]
df_unit, dsedf = sed_conversion(data["energy"], data["flux"].unit, sed)
ene = data["energy"].to(e_unit)
xerr = u.Quantity((data["energy_error_lo"], data["energy_error_hi"]))
flux = (data["flux"] * dsedf).to(df_unit)
dflux = (data["flux_error_lo"] + data["flux_error_hi"]) / 2.0
dflux = (dflux * dsedf).to(df_unit)[notul]
if interp:
difference = flux[notul] - modelfunc(ene[notul]) * flux.unit
else:
difference = flux[notul] - MLflux[groupidx][notul]
# wrap around color and marker cycles
color = color_cycle[int(g) % len(color_cycle)]
marker = marker_cycle[int(g) % len(marker_cycle)]
opts = dict(
zorder=100,
marker=marker,
ls="",
elinewidth=2,
capsize=0,
mec=color,
mew=0.1,
ms=6,
color=color,
)
opts.update(errorbar_opts)
ax.errorbar(
ene[notul].value,
(difference / dflux).decompose().value,
yerr=(dflux / dflux).decompose().value,
xerr=xerr[:, notul].to(e_unit).value,
**opts
)
from matplotlib.ticker import MaxNLocator
ax.yaxis.set_major_locator(
MaxNLocator(5, integer="True", prune="upper", symmetric=True)
)
ax.set_ylabel(r"$\Delta\sigma$")
ax.set_xscale("log")
[docs]def plot_data(
input_data,
xlabel=None,
ylabel=None,
sed=True,
figure=None,
e_unit=None,
ulim_opts={},
errorbar_opts={},
):
"""
Plot spectral data.
Parameters
----------
input_data : `emcee.EnsembleSampler`, `astropy.table.Table`, or `dict`
Spectral data to plot. Can be given as a data table, a dict generated
with `validate_data_table` or a `emcee.EnsembleSampler` with a data
property.
xlabel : str, optional
Label for the ``x`` axis of the plot.
ylabel : str, optional
Label for the ``y`` axis of the plot.
sed : bool, optional
Whether to plot SED or differential spectrum.
figure : `matplotlib.figure.Figure`, optional
`matplotlib` figure to plot on. If omitted a new one will be generated.
e_unit : `astropy.unit.Unit`, optional
Units for energy axis. Defaults to those of the data.
ulim_opts : dict
Options for upper-limit plotting. Available options are capsize (arrow
width) and height_fraction (arrow length in fraction of flux value).
errorbar_opts : dict
Addtional options to pass to `matplotlib.plt.errorbar` for plotting the
spectral flux points.
"""
import matplotlib.pyplot as plt
try:
data = validate_data_table(input_data)
except TypeError as exc:
if hasattr(input_data, "data"):
data = input_data.data
elif isinstance(input_data, dict) and "energy" in input_data.keys():
data = input_data
else:
log.warning(
"input_data format unknown, no plotting data! "
"Data loading exception: {}".format(exc)
)
raise
if figure is None:
f = plt.figure()
else:
f = figure
if len(f.axes) > 0:
ax1 = f.axes[0]
else:
ax1 = f.add_subplot(111)
# try to get units from previous plot in figure
try:
old_e_unit = u.Unit(ax1.get_xlabel().split("[")[-1].split("]")[0])
except ValueError:
old_e_unit = u.Unit("")
if e_unit is None and old_e_unit.physical_type == "energy":
e_unit = old_e_unit
elif e_unit is None:
e_unit = data["energy"].unit
_plot_data_to_ax(
data,
ax1,
e_unit=e_unit,
sed=sed,
ylabel=ylabel,
ulim_opts=ulim_opts,
errorbar_opts=errorbar_opts,
)
if xlabel is not None:
ax1.set_xlabel(xlabel)
elif xlabel is None and ax1.get_xlabel() == "":
ax1.set_xlabel(
r"$\mathrm{Energy}$"
+ " [{0}]".format(e_unit.to_string("latex_inline"))
)
ax1.autoscale()
return f
def plot_distribution(samples, label, figure=None):
"""Plot a distribution and print statistics about it"""
import matplotlib.pyplot as plt
from scipy import stats
quant = [16, 50, 84]
quantiles = dict(zip(quant, np.percentile(samples, quant)))
std = np.std(samples)
if isinstance(samples[0], u.Quantity):
unit = samples[0].unit
std = std.value
quantiles = {k: v.value for k, v in quantiles.items()}
else:
unit = ""
dist_props = "{label} distribution properties:\n \
$-$ median: ${median}$ {unit}, std: ${std}$ {unit}\n \
$-$ Median with uncertainties based on \n \
the 16th and 84th percentiles ($\\sim$1$\\sigma$):\n\
{label} = {value_error} {unit}".format(
label=label,
median=_latex_float(quantiles[50]),
std=_latex_float(std),
value_error=_latex_value_error(
quantiles[50],
quantiles[50] - quantiles[16],
quantiles[84] - quantiles[50],
),
unit=unit,
)
if figure is None:
f = plt.figure()
else:
f = figure
ax = f.add_subplot(111)
f.subplots_adjust(bottom=0.40, top=0.93, left=0.06, right=0.95)
f.text(0.2, 0.27, dist_props, ha="left", va="top")
histnbins = min(max(25, int(len(samples) / 100.0)), 100)
xlabel = "" if label is None else label
if isinstance(samples, u.Quantity):
samples_nounit = samples.value
else:
samples_nounit = samples
n, x, _ = ax.hist(
samples_nounit,
histnbins,
histtype="stepfilled",
color=color_cycle[0],
lw=0,
density=True,
)
kde = stats.kde.gaussian_kde(samples_nounit)
ax.plot(x, kde(x), color="k", label="KDE")
ax.axvline(
quantiles[50],
ls="--",
color="k",
alpha=0.5,
lw=2,
label="50% quantile",
)
ax.axvspan(
quantiles[16],
quantiles[84],
color=(0.5,) * 3,
alpha=0.25,
label="68% CI",
lw=0,
)
for xticklabel in ax.get_xticklabels():
xticklabel.set_rotation(45)
if unit != "":
xlabel += " [{0}]".format(unit)
ax.set_xlabel(xlabel)
ax.set_title("Posterior distribution of {0}".format(label))
ax.set_ylim(top=n.max() * 1.05)
return f
[docs]def plot_corner(sampler, show_ML=True, **kwargs):
"""
A plot that summarizes the parameter samples by showing them as individual
histograms and 2D histograms against each other. The maximum likelihood
parameter vector is indicated by a cross.
This function is a thin wrapper around `corner.corner`, found at
https://github.com/corner/corner.py.
Parameters
----------
sampler : `emcee.EnsembleSampler`
Sampler with a stored chain.
show_ML : bool, optional
Whether to show the maximum likelihood parameter vector as a cross on
the 2D histograms.
"""
import matplotlib.pyplot as plt
oldlw = plt.rcParams["lines.linewidth"]
plt.rcParams["lines.linewidth"] = 0.7
try:
from corner import corner
if show_ML:
_, MLp, _, _ = find_ML(sampler, 0)
else:
MLp = None
corner_opts = {
"labels": sampler.labels,
"truths": MLp,
"quantiles": [0.16, 0.5, 0.84],
"verbose": False,
"truth_color": color_cycle[0],
}
corner_opts.update(kwargs)
f = corner(sampler.get_chain(flat=True), **corner_opts)
except ImportError:
log.warning(
"The corner package is not installed;" " corner plot not available"
)
f = None
plt.rcParams["lines.linewidth"] = oldlw
return f
