"""Brazil Band Plots."""
from collections import namedtuple
import matplotlib.pyplot as plt
import numpy as np
__all__ = [
"BrazilBandCollection",
"plot_brazil_band",
"plot_cls_components",
"plot_results",
]
def __dir__():
return __all__
[docs]
class BrazilBandCollection(
namedtuple(
"BrazilBandCollection",
(
"cls_obs",
"cls_exp",
"one_sigma_band",
"two_sigma_band",
"test_size",
"clsb",
"clb",
"axes",
),
)
):
r"""
:obj:`collections.namedtuple` containing the
:class:`matplotlib.artist.Artist` objects of the
"Brazil Band" and the observed :math:`\mathrm{CL}_{s+b}` and
:math:`\mathrm{CL}_{b}` --- the components of the
:math:`\mathrm{CL}_{s}` ratio.
Returned by :func:`~pyhf.contrib.viz.brazil.plot_results`.
:param cls_obs: The :class:`matplotlib.lines.Line2D` of the
:math:`\mathrm{CL}_{s,\mathrm{obs}}` line.
:param cls_exp: The :obj:`list` of :class:`matplotlib.lines.Line2D` of the
:math:`\mathrm{CL}_{s,\mathrm{exp}}` lines.
:param one_sigma_band: The :class:`matplotlib.collections.PolyCollection` of
the :math:`\mathrm{CL}_{s,\mathrm{exp}}` :math:`\pm1\sigma` band.
:param two_sigma_band: The :class:`matplotlib.collections.PolyCollection` of
the :math:`\mathrm{CL}_{s,\mathrm{exp}}` :math:`\pm2\sigma` band.
:param test_size: The :class:`matplotlib.lines.Line2D` of the test size line.
:param clsb: The :class:`matplotlib.lines.Line2D` of the observed
:math:`\mathrm{CL}_{s+b}` line.
:param clb: The :class:`matplotlib.lines.Line2D` of the observed
:math:`\mathrm{CL}_{b}` line.
:param axes: The :class:`matplotlib.axes.Axes` the artists are plotted on.
"""
[docs]
def plot_brazil_band(test_pois, cls_obs, cls_exp, test_size, ax, **kwargs):
r"""
Plot the values of :math:`\mathrm{CL}_{s,\mathrm{obs}}` and the
:math:`\mathrm{CL}_{s,\mathrm{exp}}` band (the "Brazil band") for a series
of hypothesis tests for various POI values.
Example:
:func:`plot_brazil_band` is generally meant to be used inside
:func:`~pyhf.contrib.viz.brazil.plot_results` but can be used by itself.
>>> import numpy as np
>>> import matplotlib.pyplot as plt
>>> import pyhf
>>> import pyhf.contrib.viz.brazil
>>> pyhf.set_backend("numpy")
>>> model = pyhf.simplemodels.uncorrelated_background(
... signal=[12.0, 11.0], bkg=[50.0, 52.0], bkg_uncertainty=[3.0, 7.0]
... )
>>> observations = [51, 48]
>>> data = observations + model.config.auxdata
>>> test_pois = np.linspace(0, 5, 41)
>>> results = [
... pyhf.infer.hypotest(test_poi, data, model, return_expected_set=True)
... for test_poi in test_pois
... ]
>>> cls_obs = np.array([test[0] for test in results]).flatten()
>>> cls_exp = [
... np.array([test[1][sigma_idx] for test in results]).flatten()
... for sigma_idx in range(5)
... ]
>>> test_size = 0.05
>>> fig, ax = plt.subplots()
>>> artists = pyhf.contrib.viz.brazil.plot_brazil_band(
... test_pois, cls_obs, cls_exp, test_size, ax
... )
Args:
test_pois (:obj:`list` or :obj:`array`): The values of the POI where the
hypothesis tests were performed.
cls_obs (:obj:`list` or :obj:`array`): The values of
:math:`\mathrm{CL}_{s,\mathrm{obs}}` for the POIs tested in ``test_pois``.
cls_exp (:obj:`list` or :obj:`array`): The values of the
:math:`\mathrm{CL}_{s,\mathrm{exp}}` band for the POIs tested in
``test_pois``.
test_size (:obj:`float`): The size, :math:`\alpha`, of the test.
ax (:obj:`matplotlib.axes.Axes`): The matplotlib axis object to plot on.
Returns:
:obj:`tuple`: The :obj:`matplotlib.artist` objects drawn.
"""
line_color = kwargs.pop("color", "black")
(cls_obs_line,) = ax.plot(
test_pois, cls_obs, color=line_color, label=r"$\mathrm{CL}_{s}$"
)
cls_exp_lines = []
for idx, color in zip(range(5), 5 * [line_color]):
(_cls_exp_line,) = ax.plot(
test_pois,
cls_exp[idx],
color=color,
linestyle="dotted" if idx != 2 else "dashed",
label=None if idx != 2 else r"$\mathrm{CL}_{s,\mathrm{exp}}$",
)
cls_exp_lines.append(_cls_exp_line)
two_sigma_band = ax.fill_between(
test_pois,
cls_exp[0],
cls_exp[-1],
facecolor="yellow",
label=r"$\pm2\sigma$ $\mathrm{CL}_{s,\mathrm{exp}}$",
)
one_sigma_band = ax.fill_between(
test_pois,
cls_exp[1],
cls_exp[-2],
facecolor="green",
label=r"$\pm1\sigma$ $\mathrm{CL}_{s,\mathrm{exp}}$",
)
test_size_color = kwargs.pop("test_size_color", "red")
test_size_linestyle = kwargs.pop("test_size_linestyle", "solid")
(test_size_line,) = ax.plot(
test_pois,
[test_size] * len(test_pois),
color=test_size_color,
linestyle=test_size_linestyle,
label=rf"$\alpha={test_size}$",
)
return cls_obs_line, cls_exp_lines, one_sigma_band, two_sigma_band, test_size_line
[docs]
def plot_cls_components(test_pois, tail_probs, ax, **kwargs):
r"""
Plot the values of :math:`\mathrm{CL}_{s+b}` and :math:`\mathrm{CL}_{b}`
--- the components of the :math:`\mathrm{CL}_{s}` ratio --- for a series of
hypothesis tests for various POI values.
Example:
:func:`plot_cls_components` is generally meant to be used inside
:func:`~pyhf.contrib.viz.brazil.plot_results` but can be used by itself.
>>> import numpy as np
>>> import matplotlib.pyplot as plt
>>> import pyhf
>>> import pyhf.contrib.viz.brazil
>>> pyhf.set_backend("numpy")
>>> model = pyhf.simplemodels.uncorrelated_background(
... signal=[12.0, 11.0], bkg=[50.0, 52.0], bkg_uncertainty=[3.0, 7.0]
... )
>>> observations = [51, 48]
>>> data = observations + model.config.auxdata
>>> test_pois = np.linspace(0, 5, 41)
>>> results = [
... pyhf.infer.hypotest(
... test_poi, data, model, return_expected_set=True, return_tail_probs=True
... )
... for test_poi in test_pois
... ]
>>> tail_probs = np.array([test[1] for test in results])
>>> fig, ax = plt.subplots()
>>> artists = pyhf.contrib.viz.brazil.plot_cls_components(test_pois, tail_probs, ax)
Args:
test_pois (:obj:`list` or :obj:`array`): The values of the POI where the
hypothesis tests were performed.
tail_probs (:obj:`list` or :obj:`array`): The values of
:math:`\mathrm{CL}_{s+b}` and :math:`\mathrm{CL}_{b}` for the POIs
tested in ``test_pois``.
ax (:obj:`matplotlib.axes.Axes`): The matplotlib axis object to plot on.
Keywords:
* ``no_clb`` (:obj:`bool`): Bool for not plotting the
:math:`\mathrm{CL}_{b}` component.
* ``no_clsb`` (:obj:`bool`): Bool for not plotting the
:math:`\mathrm{CL}_{s+b}` component.
Returns:
:obj:`tuple`: The :obj:`matplotlib.lines.Line2D` artists drawn.
"""
clsb_obs = np.array([tail_prob[0] for tail_prob in tail_probs])
clb_obs = np.array([tail_prob[1] for tail_prob in tail_probs])
linewidth = kwargs.pop("linewidth", 2)
no_clsb = kwargs.pop("no_clsb", False)
no_clb = kwargs.pop("no_clb", False)
clsb_obs_line_artist = None
if not no_clsb:
clsb_color = kwargs.pop("clsb_color", "red")
(clsb_obs_line_artist,) = ax.plot(
test_pois,
clsb_obs,
color=clsb_color,
linewidth=linewidth,
label=r"$\mathrm{CL}_{s+b}$",
)
clb_obs_line_artist = None
if not no_clb:
clb_color = kwargs.pop("clb_color", "blue")
(clb_obs_line_artist,) = ax.plot(
test_pois,
clb_obs,
color=clb_color,
linewidth=linewidth,
label=r"$\mathrm{CL}_{b}$",
)
return clsb_obs_line_artist, clb_obs_line_artist
[docs]
def plot_results(test_pois, tests, test_size=0.05, ax=None, **kwargs):
r"""
Plot a series of hypothesis tests for various POI values.
For more detail on use of keywords see
:func:`~pyhf.contrib.viz.brazil.plot_brazil_band` and
:func:`~pyhf.contrib.viz.brazil.plot_cls_components`.
Example:
A Brazil band plot.
>>> import numpy as np
>>> import matplotlib.pyplot as plt
>>> import pyhf
>>> import pyhf.contrib.viz.brazil
>>> pyhf.set_backend("numpy")
>>> model = pyhf.simplemodels.uncorrelated_background(
... signal=[12.0, 11.0], bkg=[50.0, 52.0], bkg_uncertainty=[3.0, 7.0]
... )
>>> observations = [51, 48]
>>> data = observations + model.config.auxdata
>>> test_pois = np.linspace(0, 5, 41)
>>> results = [
... pyhf.infer.hypotest(test_poi, data, model, return_expected_set=True)
... for test_poi in test_pois
... ]
>>> fig, ax = plt.subplots()
>>> artists = pyhf.contrib.viz.brazil.plot_results(test_pois, results, ax=ax)
A Brazil band plot with the components of the :math:`\mathrm{CL}_{s}`
ratio drawn on top.
>>> import numpy as np
>>> import matplotlib.pyplot as plt
>>> import pyhf
>>> import pyhf.contrib.viz.brazil
>>> pyhf.set_backend("numpy")
>>> model = pyhf.simplemodels.uncorrelated_background(
... signal=[12.0, 11.0], bkg=[50.0, 52.0], bkg_uncertainty=[3.0, 7.0]
... )
>>> observations = [51, 48]
>>> data = observations + model.config.auxdata
>>> test_pois = np.linspace(0, 5, 41)
>>> results = [
... pyhf.infer.hypotest(
... test_poi, data, model, return_expected_set=True, return_tail_probs=True
... )
... for test_poi in test_pois
... ]
>>> fig, ax = plt.subplots()
>>> artists = pyhf.contrib.viz.brazil.plot_results(
... test_pois, results, ax=ax, components=True
... )
Args:
test_pois (:obj:`list` or :obj:`array`): The values of the POI where the
hypothesis tests were performed.
tests (:obj:`list` or :obj:`array`): The collection of :math:`p`-value-like
values (:math:`\mathrm{CL}_{s}` values or tail probabilities) from
the hypothesis tests.
If the ``components`` keyword argument is ``True``, ``tests`` is required
to have the same structure as :func:`pyhf.infer.hypotest`'s return
when using ``return_expected_set=True`` and ``return_tail_probs=True``:
a tuple of :math:`\mathrm{CL}_{s}`,
:math:`\left[\mathrm{CL}_{s+b}, \mathrm{CL}_{b}\right]`,
:math:`\mathrm{CL}_{s,\mathrm{exp}}` band.
test_size (:obj:`float`): The size, :math:`\alpha`, of the test.
ax (:obj:`matplotlib.axes.Axes`): The matplotlib axis object to plot on.
Returns:
:class:`BrazilBandCollection`: Artist containing the :obj:`matplotlib.artist`
objects drawn.
"""
if ax is None:
ax = plt.gca()
plot_components = kwargs.pop("components", False)
no_cls = kwargs.pop("no_cls", False)
if plot_components and len(tests[0]) != 3:
raise ValueError(
f"The components of 'tests' should have len of 3 to visualize the CLs components but have len {len(tests[0])}."
+ "\n'tests' should have format of: [CLs_obs, [CLsb, CLb], [CLs_exp band]]"
)
cls_obs = np.array([test[0] for test in tests]).flatten()
if len(tests[0]) == 3:
# split into components
tail_probs = np.array([test[1] for test in tests])
CLs_exp_set = np.array([test[2] for test in tests])
else:
CLs_exp_set = np.array([test[1] for test in tests])
cls_exp = [
np.array([exp_set[sigma_idx] for exp_set in CLs_exp_set]).flatten()
for sigma_idx in range(5)
]
brazil_band_artists = (
plot_brazil_band(test_pois, cls_obs, cls_exp, test_size, ax, **kwargs)
if not no_cls
else (None,) * 5
)
clsb, clb = (
plot_cls_components(test_pois, tail_probs, ax, **kwargs)
if plot_components
else (None, None)
)
x_label = kwargs.pop("xlabel", r"$\mu$ (POI)")
y_label = kwargs.pop("ylabel", r"$\mathrm{CL}_{s}$")
ax.set_xlabel(x_label)
ax.set_ylabel(y_label)
if ax.get_yscale() == "log":
ax.set_ylim(test_size * 0.1, 1)
else:
ax.set_ylim(0, 1)
# Order legend: ensure CLs expected band and test size are last in legend
handles, labels = ax.get_legend_handles_labels()
if not no_cls:
for label_part in ["exp", "pm1", "pm2", "alpha"]:
label_idx = next(
idx for idx, label in enumerate(labels) if label_part in label
)
handles.append(handles.pop(label_idx))
labels.append(labels.pop(label_idx))
ax.legend(handles, labels, loc="best")
return BrazilBandCollection(*brazil_band_artists, clsb, clb, ax)