"""dist plot code."""
from collections.abc import Mapping, Sequence
from importlib import import_module
from typing import Any, Literal
import arviz_stats
import xarray as xr
from arviz_base import rcParams, references_to_dataset
from arviz_base.validate import (
validate_ci_prob,
validate_dict_argument,
validate_or_use_rcparam,
validate_sample_dims,
)
from arviz_stats import ci_in_rope
from arviz_plots.plot_collection import PlotCollection
from arviz_plots.plots.utils import (
_compute_func,
compute_dist,
filter_aes,
filter_aes_full,
get_visual_kwargs,
process_group_variables_coords,
set_wrap_layout,
)
from arviz_plots.visuals import (
annotate_xy,
ecdf_line,
fill_between_y,
hist,
labelled_title,
line_x,
line_xy,
point_estimate_text,
remove_axis,
scatter_x,
scatter_xy,
step_hist,
vspan,
)
[docs]
def plot_dist(
dt,
*,
var_names=None,
filter_vars=None,
group="posterior",
coords=None,
sample_dims=None,
kind=None,
point_estimate=None,
rope=None,
ci_kind=None,
ci_prob=None,
plot_collection=None,
backend=None,
labeller=None,
aes_by_visuals: Mapping[
Literal[
"dist",
"face",
"credible_interval",
"point_estimate",
"point_estimate_text",
"rope",
"rope_text",
"title",
"rug",
],
Sequence[str],
] = None,
visuals: Mapping[
Literal[
"dist",
"face",
"credible_interval",
"point_estimate",
"point_estimate_text",
"rope",
"rope_text",
"rug",
"remove_axis",
"title",
],
Mapping[str, Any] | bool,
] = None,
stats: Mapping[
Literal["dist", "credible_interval", "point_estimate", "rope"],
Mapping[str, Any] | xr.Dataset,
] = None,
**pc_kwargs,
):
"""Plot 1D marginal densities in the style of John K. Kruschke’s book [1]_.
Generate :term:`faceted` :term:`plots` with: a graphical representation of 1D marginal
densities (as KDE, histogram, ECDF or dotplot), a credible interval and a point estimate.
Parameters
----------
dt : DataTree or dict of {str : DataTree}
Input data. In case of dictionary input, the keys are taken to be model names.
In such cases, a dimension "model" is generated and can be used to map to aesthetics.
var_names : str or list of str, optional
One or more variables to be plotted.
Prefix the variables by ~ when you want to exclude them from the plot.
filter_vars : {None, "like", "regex"}, default=None
If None, interpret var_names as the real variables names.
If “like”, interpret var_names as substrings of the real variables names.
If “regex”, interpret var_names as regular expressions on the real variables names.
group : str, default "posterior"
Group to be plotted.
coords : dict, optional
sample_dims : str or sequence of hashable, optional
Dimensions to reduce unless mapped to an aesthetic.
Defaults to ``rcParams["data.sample_dims"]``
kind : {"auto", "kde", "hist", "dot", "ecdf"}, optional
How to represent the marginal density.
Defaults to ``rcParams["plot.density_kind"]``
point_estimate : {"mean", "median", "mode"}, optional
Which point estimate to plot. Defaults to rcParam :data:`stats.point_estimate`
rope : tuple of float, or dict of {str : tuple of float}, optional
If provided, a region of practical equivalence (ROPE) is plotted as a shaded area.
The annotation shows the % of the credible interval that falls within the ROPE.
For the % of the full posterior that falls within the ROPE pass
``stats={"rope": {"ci_prob": 1}},``
ci_kind : {"eti", "hdi"}, optional
Which credible interval to use. Defaults to ``rcParams["stats.ci_kind"]``
ci_prob : float, optional
Indicates the probability that should be contained within the plotted credible interval.
Defaults to ``rcParams["stats.ci_prob"]``
plot_collection : PlotCollection, optional
backend : {"matplotlib", "bokeh"}, optional
labeller : labeller, optional
aes_by_visuals : mapping of {str : sequence of str}, optional
Mapping of visuals to aesthetics that should use their mapping in `plot_collection`
when plotted. Valid keys are the same as for `visuals`.
With a single model, no aesthetic mappings are generated by default,
each variable+coord combination gets a :term:`plot` but they all look the same,
unless there are user provided aesthetic mappings.
With multiple models, ``plot_dist`` maps "color" and "y" to the "model" dimension.
By default, all aesthetics but "y" are mapped to the density representation,
and if multiple models are present, "color" and "y" are mapped to the
credible interval and the point estimate.
When "point_estimate" key is provided but "point_estimate_text" isn't,
the values assigned to the first are also used for the second.
visuals : mapping of {str : mapping or bool}, optional
Valid keys are:
* dist -> depending on the value of `kind` passed to:
* "kde" -> passed to :func:`~arviz_plots.visuals.line_xy`
* "ecdf" -> passed to :func:`~arviz_plots.visuals.ecdf_line`
* "hist" -> passed to :func: `~arviz_plots.visuals.step_hist`
* "dot" -> passed to :func:`~arviz_plots.visuals.scatter_xy`
* face -> :term:`visual` that fills the area under the marginal distribution representation.
Defaults to False. Depending on the value of `kind` it is passed to:
* "kde", "ecdf" or "dot" -> passed to :func:`~arviz_plots.visuals.fill_between_y`
* "hist" -> passed to :func:`~arviz_plots.visuals.hist`
* credible_interval -> passed to :func:`~arviz_plots.visuals.line_x`
* point_estimate -> passed to :func:`~arviz_plots.visuals.scatter_x`
* point_estimate_text -> passed to :func:`~arviz_plots.visuals.point_estimate_text`
* title -> passed to :func:`~arviz_plots.visuals.labelled_title`
* rug -> passed to :func:`~arviz_plots.visuals.scatter_x`. Defaults to False.
* rope -> passed to :func:`~arviz_plots.visuals.vspan`. Defaults to True.
* rope_text -> passed to :func:`~arviz_plots.visuals.annotate_xy`. Defaults to True.
Both rope and rope_text are ignored if the main rope argument is None.
* remove_axis -> not passed anywhere, can only be ``False`` to skip calling this function
stats : mapping of {str : mapping or Dataset}, optional
Valid keys are:
* dist -> passed to :func:`~arviz_stats.kde`, :func:`~arviz_stats.histogram`,
:func:`~arviz_stats.ecdf`, or :func:`~arviz_stats.qds` depending on `kind`
* credible_interval -> passed to :func:`~arviz_stats.eti` or :func:`arviz_stats.hdi`
* point_estimate -> passed to mean, median or mode. Defaults to
round the result according to ``rcParams["stats.round_to"]``.
* rope -> passed to :func:`~arviz_stats.ci_in_rope`
In case a :class:`~xarray.Dataset` is provided, it will be interpreted
as pre-computed values for that statistic.
**pc_kwargs
Passed to :class:`arviz_plots.PlotCollection.wrap`
Returns
-------
PlotCollection
See Also
--------
:ref:`plots_intro` :
General introduction to batteries-included plotting functions, common use and logic overview
Examples
--------
Map the color to the variable, and have the mapping apply
to the title too instead of only the density representation:
.. plot::
:context: close-figs
>>> from arviz_plots import plot_dist, style
>>> style.use("arviz-variat")
>>> from arviz_base import load_arviz_data
>>> non_centered = load_arviz_data('non_centered_eight')
>>> pc = plot_dist(
>>> non_centered,
>>> coords={"school": ["Choate", "Deerfield", "Hotchkiss"]},
>>> aes={"color": ["__variable__"]},
>>> aes_by_visuals={"title": ["color"]},
>>> )
Faceting and aesthetics mappings happen on unique coordinate values. If there are repeated
coordinate values they will be grouped and reduced along with `sample_dims`.
.. plot::
:context: close-figs
>>> post = non_centered.posterior.to_dataset()
>>> repeated_coords = ["a", "a", "a", "b", "b", "b", "b", "c"]
>>> pc = plot_dist(post.assign_coords(school=repeated_coords))
Add a region of practical equivalence (ROPE) and show the percentage of
the credible interval that falls within it:
.. plot::
:context: close-figs
>>> pc = plot_dist(
>>> non_centered,
>>> var_names=["mu", "tau"],
>>> rope={"mu": (-0.5, 0.9), "tau": (0.2, 9.4)},
>>> )
.. minigallery:: plot_dist
References
----------
.. [1] Kruschke. Doing Bayesian Data Analysis, Second Edition: A Tutorial with R,
JAGS, and Stan. Academic Press, 2014. ISBN 978-0-12-405888-0.
https://www.sciencedirect.com/book/9780124058880
"""
kind = validate_or_use_rcparam(kind, "plot.density_kind")
point_estimate = validate_or_use_rcparam(point_estimate, "stats.point_estimate")
ci_kind = validate_or_use_rcparam(ci_kind, "stats.ci_kind")
ci_prob = validate_ci_prob(ci_prob)
aes_by_visuals = validate_dict_argument(aes_by_visuals, (plot_dist, "aes_by_visuals"))
visuals = validate_dict_argument(visuals, (plot_dist, "visuals"))
stats = validate_dict_argument(stats, (plot_dist, "stats"))
if kind == "ecdf":
visuals.setdefault("remove_axis", False)
distribution = process_group_variables_coords(
dt, group=group, var_names=var_names, filter_vars=filter_vars, coords=coords
)
sample_dims = validate_sample_dims(sample_dims, data=distribution)
if backend is None:
if plot_collection is None:
backend = rcParams["plot.backend"]
else:
backend = plot_collection.backend
plot_bknd = import_module(f".backend.{backend}", package="arviz_plots")
if rope:
rope_dataset = references_to_dataset(
rope, distribution, sample_dims=sample_dims, ref_dim="band_dim"
)
ci_in_rope_kwargs = stats.get("rope", {}).copy()
ci_in_rope_kwargs.setdefault("ci_prob", ci_prob)
ci_in_rope_kwargs.setdefault("ci_kind", ci_kind)
rope_pct = ci_in_rope(distribution, rope=rope, **ci_in_rope_kwargs)
if plot_collection is None:
pc_kwargs["figure_kwargs"] = pc_kwargs.get("figure_kwargs", {}).copy()
if backend is None:
backend = rcParams["plot.backend"]
pc_kwargs.setdefault(
"cols",
["__variable__"]
+ [dim for dim in distribution.dims if dim not in {"model"}.union(sample_dims)],
)
if "model" in distribution:
pc_kwargs["aes"] = pc_kwargs.get("aes", {}).copy()
pc_kwargs["aes"].setdefault("color", ["model"])
pc_kwargs["aes"].setdefault("y", ["model"])
pc_kwargs = set_wrap_layout(pc_kwargs, plot_bknd, distribution)
plot_collection = PlotCollection.wrap(
distribution,
backend=backend,
**pc_kwargs,
)
face_kwargs = get_visual_kwargs(visuals, "face", False)
density_kwargs = get_visual_kwargs(visuals, "dist")
aes_by_visuals.setdefault("dist", plot_collection.aes_set.difference("y"))
if face_kwargs is not False:
aes_by_visuals.setdefault("face", set(aes_by_visuals["dist"]).difference({"linestyle"}))
if "model" in distribution:
aes_by_visuals.setdefault("credible_interval", ["color", "y"])
aes_by_visuals.setdefault("point_estimate", ["color", "y"])
if "point_estimate" in aes_by_visuals and "point_estimate_text" not in aes_by_visuals:
aes_by_visuals["point_estimate_text"] = aes_by_visuals["point_estimate"]
density_reduce, density_active, density_aes, density_ignore = filter_aes_full(
plot_collection, aes_by_visuals, "dist", sample_dims
)
density = compute_dist(distribution, density_reduce, density_active, kind=kind, stats=stats)
kind_var_map = {
kind_i: [k for k, da in density.items() if kind_i == da.attrs["kind"]]
for kind_i in ("dot", "ecdf", "hist", "kde")
}
kind_var_map = {k: v for k, v in kind_var_map.items() if v}
# filled face (should go under the dist visual if both present)
if face_kwargs is not False:
_, face_aes, face_ignore = filter_aes(plot_collection, aes_by_visuals, "face", sample_dims)
if "color" not in face_aes:
face_kwargs.setdefault("color", "C0")
if "alpha" not in face_aes:
face_kwargs.setdefault("alpha", 0.4)
if "dot" in kind_var_map:
kwargs = stats.get("dist", {}).copy()
kwargs["top_only"] = True
top_only_qds = distribution[kind_var_map["dot"]].azstats.qds(
dim=density_reduce, **kwargs
)
density = xr.merge((top_only_qds, density), compat="override")
if any(kind_i in kind_var_map for kind_i in ("kde", "ecdf", "dot")):
fill_between_vars = [
var_name
for kind_i in ("kde", "ecdf", "dot")
for var_name in kind_var_map.get(kind_i, [])
]
face_density = (
density[fill_between_vars]
.rename(plot_axis="kwarg")
.sel(kwarg=["x", "y"])
.pad(kwarg=(0, 1), constant_values=0)
.assign_coords(kwarg=["x", "y_top", "y_bottom"])
)
plot_collection.map(
fill_between_y,
"face",
data=face_density,
ignore_aes=face_ignore,
**face_kwargs,
)
if "hist" in kind_var_map:
plot_collection.map(
hist,
"face",
data=density[kind_var_map["hist"]],
ignore_aes=face_ignore,
**face_kwargs,
)
# density
if density_kwargs is not False:
if "color" not in density_aes:
density_kwargs.setdefault("color", "C0")
if "kde" in kind_var_map:
plot_collection.map(
line_xy,
"dist",
data=density[kind_var_map["kde"]],
ignore_aes=density_ignore,
**density_kwargs,
)
if "ecdf" in kind_var_map:
plot_collection.map(
ecdf_line,
"dist",
data=density[kind_var_map["ecdf"]],
ignore_aes=density_ignore,
**density_kwargs,
)
if "hist" in kind_var_map:
plot_collection.map(
step_hist,
"dist",
data=density[kind_var_map["hist"]],
ignore_aes=density_ignore,
**density_kwargs,
)
if "dot" in kind_var_map:
plot_collection.map(
scatter_xy,
"dist",
data=density[kind_var_map["dot"]],
ignore_aes=density_ignore,
**density_kwargs,
)
# rug
rug_kwargs = get_visual_kwargs(visuals, "rug", False)
if rug_kwargs is not False:
if not isinstance(rug_kwargs, dict):
rug_kwargs = {}
_, rug_aes, rug_ignore = filter_aes(plot_collection, aes_by_visuals, "rug", sample_dims)
if "color" not in rug_aes:
rug_kwargs.setdefault("color", "B1")
if "marker" not in rug_aes:
rug_kwargs.setdefault("marker", "|")
if "size" not in rug_aes:
rug_kwargs.setdefault("size", 15)
plot_collection.map(
scatter_x,
"rug",
data=distribution,
ignore_aes=rug_ignore,
**rug_kwargs,
)
plot_y_max = None
if density_kwargs is not False or face_kwargs is not False:
density_ys = density.sel(
plot_axis=[
coord for coord in density["plot_axis"].values if coord in ("y", "histogram")
]
)
plot_y_max = density_ys.max(
[dim for dim in density_ys.dims if dim not in plot_collection.facet_dims]
)
if (
(density_kwargs is not False or face_kwargs is not False)
and ("y" in plot_collection.aes_set)
and ("model" in distribution)
and (plot_collection.coords is None)
):
y_ds = plot_collection.get_aes_as_dataset("y")["mapping"] * plot_y_max * 0.15
plot_collection.update_aes_from_dataset("y", y_ds)
# credible interval
ci_kwargs = get_visual_kwargs(visuals, "credible_interval")
if ci_kwargs is not False:
ci_reduce, ci_active, ci_aes, ci_ignore = filter_aes_full(
plot_collection, aes_by_visuals, "credible_interval", sample_dims
)
ci_stats_value = stats.get("credible_interval", {})
if isinstance(ci_stats_value, xr.Dataset):
ci = ci_stats_value
else:
ci_stats_kwargs = ci_stats_value.copy()
ci_stats_kwargs["prob"] = ci_prob
if ci_kind == "eti":
ci = _compute_func(
arviz_stats.eti,
distribution,
active_dims=ci_active,
reduce_dims=ci_reduce,
kwargs=ci_stats_kwargs,
)
elif ci_kind == "hdi":
ci = _compute_func(
arviz_stats.hdi,
distribution,
active_dims=ci_active,
reduce_dims=ci_reduce,
kwargs=ci_stats_kwargs,
)
if "color" not in ci_aes:
ci_kwargs.setdefault("color", "B2")
plot_collection.map(line_x, "credible_interval", data=ci, ignore_aes=ci_ignore, **ci_kwargs)
# point estimate
pe_kwargs = get_visual_kwargs(visuals, "point_estimate")
pet_kwargs = get_visual_kwargs(visuals, "point_estimate_text")
if (pe_kwargs is not False) or (pet_kwargs is not False):
pe_reduce, pe_active, pe_aes, pe_ignore = filter_aes_full(
plot_collection, aes_by_visuals, "point_estimate", sample_dims
)
pe_stats_value = stats.get("point_estimate", {})
if isinstance(pe_stats_value, xr.Dataset):
point = pe_stats_value
else:
pe_stats_kwargs = pe_stats_value.copy()
pe_stats_kwargs["round_to"] = "none"
pe_func = {
"mean": arviz_stats.mean,
"median": arviz_stats.median,
"mode": arviz_stats.mode,
}[point_estimate]
point = _compute_func(
pe_func,
distribution,
active_dims=pe_active,
reduce_dims=pe_reduce,
kwargs=pe_stats_kwargs,
)
if pe_kwargs is not False:
if "color" not in pe_aes:
pe_kwargs.setdefault("color", "B2")
plot_collection.map(
scatter_x,
"point_estimate",
data=point,
ignore_aes=pe_ignore,
**pe_kwargs,
)
# point estimate text
if pet_kwargs is not False:
if plot_y_max is None:
point_y = xr.full_like(point, 0.05)
else:
point_y = 0.1 * plot_y_max
point = xr.concat((point, point_y), dim="plot_axis").assign_coords(plot_axis=["x", "y"])
_, pet_aes, pet_ignore = filter_aes(
plot_collection, aes_by_visuals, "point_estimate_text", sample_dims
)
if "color" not in pet_aes:
pet_kwargs.setdefault("color", "B2")
pet_kwargs.setdefault("horizontal_align", "center")
pet_kwargs.setdefault("point_label", "x")
plot_collection.map(
point_estimate_text,
"point_estimate_text",
data=point,
point_estimate=point_estimate,
ignore_aes=pet_ignore,
**pet_kwargs,
)
# rope
if rope:
# rope band
rope_kwargs = get_visual_kwargs(visuals, "rope", True)
if rope_kwargs is not False:
_, rope_aes, rope_ignore = filter_aes(
plot_collection, aes_by_visuals, "rope", sample_dims
)
if "color" not in rope_aes:
rope_kwargs.setdefault("color", "B1")
if "alpha" not in rope_aes:
rope_kwargs.setdefault("alpha", 0.2)
plot_collection.map(
vspan,
"rope",
data=rope_dataset,
ignore_aes=rope_ignore,
**rope_kwargs,
)
# rope text annotation
rope_text_kwargs = get_visual_kwargs(visuals, "rope_text", True)
if rope_text_kwargs is not False:
rope_mid = rope_dataset.mean("band_dim")
if plot_y_max is None:
rope_y = xr.full_like(rope_mid, 0.02)
else:
rope_y = 0.35 * plot_y_max
_, rope_text_aes, rope_text_ignore = filter_aes(
plot_collection, aes_by_visuals, "rope_text", sample_dims
)
if "color" not in rope_text_aes:
rope_text_kwargs.setdefault("color", "B1")
rope_text_kwargs.setdefault("horizontal_align", "center")
plot_collection.map(
annotate_xy,
"rope_text",
data=rope_pct,
ignore_aes=rope_text_ignore,
x=rope_mid,
y=rope_y,
text=lambda pct: f"{pct:.1f}% in ROPE",
**rope_text_kwargs,
)
# aesthetics
title_kwargs = get_visual_kwargs(visuals, "title")
if title_kwargs is not False:
_, title_aes, title_ignore = filter_aes(
plot_collection, aes_by_visuals, "title", sample_dims
)
if "color" not in title_aes:
title_kwargs.setdefault("color", "B1")
plot_collection.facet_map(
labelled_title,
"title",
ignore_aes=title_ignore,
subset_info=True,
labeller=labeller,
**title_kwargs,
)
if visuals.get("remove_axis", True) is not False:
plot_collection.facet_map(
remove_axis,
store_artist=backend == "none",
axis="y",
)
return plot_collection
