"""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
from arviz_base.validate import (
validate_ci_prob,
validate_dict_argument,
validate_or_use_rcparam,
validate_sample_dims,
)
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 (
ecdf_line,
fill_between_y,
hist,
labelled_title,
line_x,
line_xy,
point_estimate_text,
remove_axis,
scatter_x,
scatter_xy,
step_hist,
)
[docs]
def plot_dist(
dt,
*,
var_names=None,
filter_vars=None,
group="posterior",
coords=None,
sample_dims=None,
kind=None,
point_estimate=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",
"title",
"rug",
],
Sequence[str],
] = None,
visuals: Mapping[
Literal[
"dist",
"face",
"credible_interval",
"point_estimate",
"point_estimate_text",
"title",
"rug",
"remove_axis",
],
Mapping[str, Any] | bool,
] = None,
stats: Mapping[
Literal["dist", "credible_interval", "point_estimate"], 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`
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.
* 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"]``.
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))
.. 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 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,
)
if (
(density_kwargs is not False or face_kwargs is not False)
and ("model" in distribution)
and (plot_collection.coords is None)
):
y_ds = plot_collection.get_aes_as_dataset("y")["mapping"]
density_ys = density.sel(
plot_axis=[
coord for coord in density["plot_axis"].values if coord in ("y", "histogram")
]
)
density_ys_max = density_ys.max(
[dim for dim in density_ys.dims if dim not in plot_collection.facet_dims]
)
y_ds = 0.15 * y_ds * density_ys_max
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 density_kwargs is False and face_kwargs is False:
point_y = xr.full_like(point, 0.05)
else:
density_ys = density.sel(
plot_axis=[
coord for coord in density["plot_axis"].values if coord in ("y", "histogram")
]
)
point_y = 0.1 * density_ys.max(
[dim for dim in density_ys.dims if dim not in point.dims]
)
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,
)
# 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
