# Licensed under a 3-clause BSD style license - see LICENSE.rst
import matplotlib.pyplot as plt
import numpy as np
import astropy.units as u
from scipy.stats import gaussian_kde, norm
[docs]
def plot_samples_violin_vs_energy(
energy_edges,
samples_per_band,
weights_per_band=None,
energy_power=0,
bw_method="scott",
grid_size=200,
ax=None,
errorbar_kwargs=None,
errorbar_ul_kwargs=None,
violin_kwargs=None,
violin_clip=None,
y_label="dN/dE",
):
"""Plot sample violin distributions per energy bin on log–log axes.
Unlike standard error bars which only shows summary statistics,
a violin plot shows the full probability density for each energy bin.
This function draws one violin per energy interval, using weighted
kernel-density estimation (KDE) in log-space. The median and
1σ-equivalent quantiles are overlaid using Gammapy-style flux-point
error bars. Users can optionally apply an ``E_center**p`` scaling to the
flux values and clip violin tails to a specified containment interval.
Samples of other energy-dependent parameters can also be plotted,
in that case the ``y_label`` should be changed accordingly.
Parameters
----------
energy_edges : `~numpy.ndarray` or `~astropy.units.Quantity`
Monotonically increasing energy bin edges of length ``nbins + 1``.
Must be positive and finite for log scaling.
samples_per_band : list of `~astropy.units.Quantity`
List of sample arrays, one per energy bin. Each array contains
draws from the flux posterior (or other distribution) in that bin.
weights_per_band : list of `~numpy.ndarray`, optional
Per-sample weights for each bin, same length as ``samples_per_band``.
If omitted, uniform weights are assumed.
energy_power : float, optional
Power of energy to multiply y-axis with. Default is 0.
bw_method : str or float, optional
Bandwidth selection passed to `~scipy.stats.gaussian_kde`.
Default is "scott".
grid_size : int, optional
Number of evaluation points in log-flux space for the KDE grid.
Default is 200.
ax : `~matplotlib.axes.Axes`, optional
Matplotlib axes to draw the violins and error bars. Default is None.
errorbar_kwargs : dict, optional
Keyword arguments forwarded to `~matplotlib.pyplot.errorbar` for the quantile bars.
errorbar_ul_kwargs : dict, optional
Keyword arguments forwarded to `~matplotlib.pyplot.errorbar` for the upper limit bars.
violin_kwargs : dict, optional
Keyword arguments forwarded to `~matplotlib.axes.Axes.fill` for violin plot.
violin_clip : (float, float), optional
Lower and upper containment fractions (in [0, 1]) used to clip
the violin tails in log-space. If omitted, defaults to
``(norm.cdf(-4), norm.cdf(4))``.
y_label : str, optional
Base label for the y-axis, before unit and scaling prefixes.
Default is "dN/dE".
Returns
-------
ax : `~matplotlib.axes.Axes`
Matplotlib axes.
Notes
-----
- KDE is performed on the transformed variable ``log10(E^p * F)``,
where ``p = energy_power``. The violin polygon is constructed by
mirroring the normalized log-density around the bin center in
log-energy.
- Violin clipping is applied only to the KDE grid, not to the
quantile-based error bars.
- Negative flux samples are replaced by a small positive surrogate
(10% of the smallest positive sample) to allow log-space evaluation.
- The plotted error bars correspond to the 16%, 50%, and 84%
weighted percentiles of the *unclipped* samples, unless the median
is non-positive, in which case an upper limit marker is drawn
at the 98% percentile. For a gaussian distribution these values corresponds to
1σ errors and 2σ upper limit plotted by default by the `~gammapy.estimators.FluxPoints.plot` method.
"""
if ax is None:
ax = plt.gca()
ax.set_xscale("log")
ax.set_yscale("log")
# Energy edges
energy_edges = u.Quantity(energy_edges)
unit_x = f"{energy_edges.unit}"
edges = np.asarray(energy_edges.to_value())
_validate_energy_edges(edges)
unit_y = f"{samples_per_band[0].unit}"
nbins = len(edges) - 1
samples_per_band, weights_per_band = _validate_inputs(
samples_per_band, weights_per_band, nbins
)
# clip defaults
if violin_clip is None:
violin_clip = (norm.cdf(-4), norm.cdf(4))
lc, uc = violin_clip
if not (0 <= lc < uc <= 1):
raise ValueError("violin_clip must satisfy 0 <= low < high <= 1.")
# Errorbar styles
errorbar_kwargs = errorbar_kwargs or {}
errorbar_kwargs_defaults = dict(
marker="o",
ms=4.5,
mec="black",
mfc="white",
color="black",
capsize=2.5,
elinewidth=1.2,
lw=1.2,
)
for key in errorbar_kwargs_defaults.keys():
errorbar_kwargs.setdefault(key, errorbar_kwargs_defaults[key])
errorbar_ul_kwargs = errorbar_ul_kwargs or {}
errorbar_ul_kwargs_defaults = dict(
marker="v",
ms=7,
mec="black",
mfc="white",
color="black",
capsize=2.5,
elinewidth=1.2,
lw=1.2,
)
for key in errorbar_ul_kwargs_defaults.keys():
errorbar_ul_kwargs.setdefault(key, errorbar_ul_kwargs_defaults[key])
if violin_kwargs is None:
violin_kwargs = dict()
violin_kwargs_defaults = dict(
alpha=0.45,
color="C0",
edgecolor="black",
lw=0.8,
)
for key in violin_kwargs_defaults.keys():
violin_kwargs.setdefault(key, violin_kwargs_defaults[key])
quantiles = [100 * norm.cdf(-1), 50, 100 * norm.cdf(1), 100 * norm.cdf(2)]
# Precompute bin geometry
emins = edges[:-1]
emaxs = edges[1:]
xlog_min = np.log10(emins)
xlog_max = np.log10(emaxs)
centers = 0.5 * (xlog_min + xlog_max)
halfwidths = 0.5 * (xlog_max - xlog_min) * 0.99
for samples, weights, xlog_c, hwlog, emin, emax in zip(
samples_per_band, weights_per_band, centers, halfwidths, emins, emaxs
):
samples = samples.to_value(unit_y)
s, smin_pos = _sanitize_samples(samples)
w = _sanitize_weights(weights, len(s))
if w is None or s.size == 0:
continue
# transform
ylog, scale = _compute_log_transformed(s, xlog_c, energy_power)
# KDE
grid_full, dens_full = _kde_evaluate(ylog, w, grid_size, bw_method)
# clipping interval
y_high = np.percentile(s, uc * 100, weights=w, method="inverted_cdf") * scale
y_low = np.percentile(s, lc * 100, weights=w, method="inverted_cdf") * scale
y_low = max(smin_pos * scale, y_low)
if y_high > y_low > 0:
ymin = max(grid_full.min(), np.log10(y_low))
ymax = min(grid_full.max(), np.log10(y_high))
mask = (grid_full >= ymin) & (grid_full <= ymax)
ygrid_log, dens = grid_full[mask], dens_full[mask]
else:
ygrid_log, dens = grid_full, dens_full
# violin polygon
_draw_violin(ax, xlog_c, hwlog, ygrid_log, dens, violin_kwargs)
# quantile bars
_draw_errorbar(
ax,
xlog_c,
emin,
emax,
samples,
weights,
quantiles,
scale,
errorbar_kwargs,
errorbar_ul_kwargs,
)
# Labels
ax.set_xlabel(f"Energy [{unit_x}]")
if energy_power:
p = f"{energy_power:g}"
unit_str = rf"[{unit_x}$^{p}$ × {unit_y}]"
ax.set_ylabel(rf"$E^{p}\,\times$ {y_label} {unit_str}")
else:
ax.set_ylabel(f"{y_label} [{unit_y}]")
return ax
def _validate_inputs(samples_per_band, weights_per_band, nbins):
"""Validate that samples and weights."""
if len(samples_per_band) != nbins:
raise ValueError("samples_per_band must match number of bins.")
if weights_per_band is None:
weights_per_band = [None] * nbins
elif len(weights_per_band) != nbins:
raise ValueError("weights_per_band must match number of bins.")
return samples_per_band, weights_per_band
def _validate_energy_edges(edges):
"""Validate that energy_edges is 1D, positive, finite, and strictly increasing."""
if edges.ndim != 1 or edges.size < 2:
raise ValueError("energy_edges must be 1D and contain at least two values.")
if (
not np.all(np.isfinite(edges))
or np.any(edges <= 0)
or not np.all(np.diff(edges) > 0)
):
raise ValueError("energy_edges must be strictly increasing, finite, and > 0.")
def _sanitize_samples(samples):
"""Return finite samples with negative values adjusted for log-scale use."""
s = samples.copy()
mask = np.isfinite(s)
s = s[mask]
s_pos = s[s > 0]
if np.any(s < 0):
if s_pos.size > 0:
s[s < 0] = 0.1 * s_pos.min()
return s, s_pos.min()
def _sanitize_weights(weights, size):
"""Return normalised finite non-negative weights or None if empty."""
if weights is None:
w = np.ones(size, float)
else:
w = np.asarray(weights.copy(), float)
w = w[np.isfinite(w) & (w >= 0)]
total = w.sum()
return w / total if total > 0 else None
def _compute_log_transformed(s, xlog_center, energy_power):
"""Transform samples to log10(E^p × F) and return log-values and linear scale."""
if energy_power:
shift = energy_power * xlog_center
scale = (10**xlog_center) ** energy_power
else:
shift = 0.0
scale = 1.0
return np.log10(s) + shift, scale
def _kde_evaluate(ylog, w, grid_size, bw_method):
"""Evaluate weighted KDE on a fixed log-grid and normalise density."""
grid = np.linspace(ylog.min(), ylog.max(), int(grid_size))
kde = gaussian_kde(ylog, weights=w, bw_method=bw_method)
dens = kde(grid)
peak = dens.max()
return grid, dens / peak if peak > 0 else dens
def _draw_violin(ax, x_center, hwlog, ygrid_log, dens, violin_kwargs):
"""Draw a symmetric violin polygon in log-space."""
xlog_left = x_center - hwlog * dens
xlog_right = x_center + hwlog * dens
x_left = 10**xlog_left
x_right = 10**xlog_right
y_lin = 10**ygrid_log
x_poly = np.concatenate([x_left, x_right[::-1]])
y_poly = np.concatenate([y_lin, y_lin[::-1]])
return ax.fill(x_poly, y_poly, zorder=2, **violin_kwargs)
def _draw_errorbar(ax, xlog_c, emin, emax, s, w, quantiles, scale, err_kw, err_ul_kw):
"""Draw Gammapy-style flux-point error bars or upper limits."""
y_qs = [
np.percentile(s, q, weights=w, method="inverted_cdf") * scale for q in quantiles
]
y_lo, y_med, y_hi, y_ul = y_qs[0], y_qs[1], y_qs[2], y_qs[3]
if y_med > 0:
yerr = np.array([[max(0.0, y_med - y_lo)], [max(0.0, y_hi - y_med)]])
kwargs = err_kw
else:
yerr = None
y_med = y_ul
kwargs = err_ul_kw
x_c = 10 ** (xlog_c)
xerr = np.array([[x_c - emin], [emax - x_c]])
return ax.errorbar(x_c, y_med, xerr=xerr, yerr=yerr, **(kwargs or {}), zorder=5)
