"""Plot fractional rank."""
from collections.abc import Mapping, Sequence
from importlib import import_module
from typing import Any, Literal
import numpy as np
import xarray as xr
from arviz_base import rcParams
from arviz_base.labels import BaseLabeller
from arviz_base.validate import (
validate_dict_argument,
validate_or_use_rcparam,
validate_sample_dims,
)
from arviz_stats.ecdf_utils import ecdf_pit
from arviz_plots.plot_collection import PlotCollection
from arviz_plots.plots.utils import (
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,
labelled_title,
labelled_x,
remove_axis,
scatter_xy,
set_ylim,
)
[docs]
def plot_rank(
dt,
*,
var_names=None,
filter_vars=None,
group="posterior",
coords=None,
sample_dims=None,
envelope_prob=None,
method="mtc_c",
thin=None,
plot_collection=None,
backend=None,
labeller=None,
aes_by_visuals: Mapping[
Literal[
"ecdf_lines",
"credible_interval",
"xlabel",
"title",
],
Sequence[str],
] = None,
visuals: Mapping[
Literal[
"ecdf_lines",
"credible_interval",
"xlabel",
"title",
"remove_axis",
],
Mapping[str, Any] | bool,
] = None,
stats: Mapping[
Literal["ecdf_pit", "mtc_c", "thin"],
Mapping[str, Any] | xr.Dataset,
] = None,
**pc_kwargs,
):
"""Fractional rank Δ-ECDF plots.
Rank plots are built by replacing the posterior draws by their ranking computed over all chains.
Then each chain is plotted independently. If all of the chains are targeting the same posterior,
we expect the ranks in each chain to be uniformly distributed.
To simplify comparison we compute the ordered fractional ranks, which are distributed
uniformly in [0, 1]. Additionally, we plot the Δ-ECDF, that is, the difference between the
expected CDF from the observed ECDF.
The points that contribute the most to deviations from uniformity are computed as described
in [1]_.
Parameters
----------
dt : DataTree
Input data
var_names : str or list of str, optional
One or more variables to be plotted. Currently only one variable is supported.
Prefix the variables by ~ when you want to exclude them from the plot.
filter_vars : {None, “like”, “regex”}, optional, default=None
If None (default), 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, optional
Which group to use. Defaults to "posterior".
coords : dict, optional
Coordinates to plot.
sample_dims : str or sequence of hashable, optional
Dimensions to reduce unless mapped to an aesthetic.
Defaults to ``rcParams["data.sample_dims"]``
envelope_prob : float, optional
Indicates the probability that should be contained within the envelope.
Defaults to ``rcParams["stats.envelope_prob"]``.
method : {"mtc_c", "envelope"}, default "mtc_c"
Method to use for the rank plot. If "mtc_c", the multi-chain test is performed and
suspicious points are highlighted. If "envelope", the envelope is computed and plotted.
thin : bool, default None
Whether to thin the data before plotting. Defaults to None, which means that it is set
to True if "method" is "envelope" and False otherwise.
plot_collection : PlotCollection, optional
backend : {"matplotlib", "bokeh", "plotly"}, 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`.
visuals : mapping of {str : mapping or bool}, optional
Valid keys are:
* ecdf_lines -> passed to :func:`~arviz_plots.visuals.ecdf_line`
* credible_interval -> passed to :func:`~arviz_plots.visuals.fill_between_y`
* xlabel -> passed to :func:`~arviz_plots.visuals.labelled_x`
* title -> passed to :func:`~arviz_plots.visuals.labelled_title`
* remove_axis -> not passed anywhere, can only be ``False`` to skip calling this function
stats : mapping, optional
Valid keys are:
* ecdf_pit -> passed to :func:`~arviz_stats.ecdf_utils.ecdf_pit`. Default is
``{"n_simulations": 1000}``.
* mtc_c -> passed to :func:`~arviz_stats.mchain_uniformity_test`.
* thin -> passed to :func:`~arviz_stats.thin`
**pc_kwargs
Passed to :class:`arviz_plots.PlotCollection.wrap`
Returns
-------
PlotCollection
Notes
-----
The preferred method is `mtc_c` as it takes into account the autocorrelation in the rank values
as described in [1]_. The "envelope" method is not longer recommended and it will likely be
removed in a future release.
Examples
--------
Rank plot for the crabs hurdle-negative-binomial dataset.
.. plot::
:context: close-figs
>>> from arviz_plots import plot_rank, style
>>> style.use("arviz-variat")
>>> from arviz_base import load_arviz_data
>>> dt = load_arviz_data('crabs_hurdle_nb')
>>> plot_rank(dt, var_names=["~mu"])
.. minigallery:: plot_rank
References
----------
.. [1] Tesso et al. *LOO-PIT predictive model checking* arXiv:2603.02928 (2026).
.. [2] Säilynoja et al. *Graphical test for discrete uniformity and
its applications in goodness-of-fit evaluation and multiple sample comparison*.
Statistics and Computing 32(32). (2022) https://doi.org/10.1007/s11222-022-10090-6
"""
envelope_prob = validate_or_use_rcparam(envelope_prob, "stats.envelope_prob")
aes_by_visuals = validate_dict_argument(aes_by_visuals, (plot_rank, "aes_by_visuals"))
visuals = validate_dict_argument(visuals, (plot_rank, "visuals"))
visuals.setdefault("remove_axis", True)
stats = validate_dict_argument(stats, (plot_rank, "stats"))
if method not in ["mtc_c", "envelope"]:
raise ValueError(f"Invalid method {method}. Valid options are 'mtc_c' and 'envelope'.")
if backend is None:
if plot_collection is None:
backend = rcParams["plot.backend"]
else:
backend = plot_collection.backend
if labeller is None:
labeller = BaseLabeller()
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)
ecdf_pit_kwargs = stats.get("ecdf_pit", {}).copy()
if method == "envelope":
ecdf_pit_kwargs.setdefault("n_simulations", 1000)
ecdf_pit_kwargs.setdefault("n_chains", distribution.sizes["chain"])
else:
ecdf_pit_kwargs.setdefault("gamma", 0)
ecdf_dims = ["draw"]
if thin is None:
thin = method == "envelope"
if thin:
distribution = distribution.azstats.thin(sample_dims=ecdf_dims, **stats.get("thin", {}))
sample_size = np.prod([len(distribution[dims]) for dims in ecdf_dims])
# Compute ranks
dt_ecdf_ranks = distribution.azstats.compute_ranks(dim=sample_dims)
# Compute ECDF
dt_ecdf = dt_ecdf_ranks.azstats.ecdf(dim=ecdf_dims, pit=True, npoints=sample_size)
dt_ecdf = dt_ecdf.rename(
{
dim: "ecdf_dim"
for dim in dt_ecdf.dims
if dim not in dt_ecdf_ranks.dims and dim != "plot_axis"
}
)
# Compute multi-chain test p-values
if method == "mtc_c":
alpha = 1 - envelope_prob
gamma = stats.get("ecdf_pit", {}).get("gamma", 0)
dt_ranks_rel = distribution.azstats.compute_ranks(dim=sample_dims, relative=True)
mtc_c_kwargs = stats.get("mtc_c", {}).copy()
p_values, b_shapley, w_shapley = dt_ranks_rel.azstats.mchain_uniformity_test(
dim=sample_dims, **mtc_c_kwargs
)
highlight = ((b_shapley > gamma) * (w_shapley > gamma)) & (p_values < alpha)
suspicious_mask = highlight.rename({"pit_dim": "ecdf_dim"})
# use the Dvoretzky-Kiefer-Wolfowitz inequality plus a small padding
# to get the default y-limits for the plot.
expected_max = np.sqrt(np.log(2 / alpha) / (2 * sample_size)) * 1.3
actual_max = np.max(np.abs(dt_ecdf.sel(plot_axis="y").to_array())).item()
epsilon = max(expected_max, actual_max)
else:
p_values = None
# Compute envelope
if method == "mtc_c":
x_ci = lower_ci = upper_ci = None
else:
dummy_vals = np.linspace(0, 1, sample_size)
x_ci, _, lower_ci, upper_ci = ecdf_pit(dummy_vals, envelope_prob, **ecdf_pit_kwargs)
lower_ci = lower_ci - x_ci
upper_ci = upper_ci - x_ci
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()
pc_kwargs["aes"] = pc_kwargs.get("aes", {}).copy()
pc_kwargs.setdefault("col_wrap", 4)
pc_kwargs.setdefault(
"cols", ["__variable__"] + [dim for dim in dt_ecdf_ranks.dims if dim not in sample_dims]
)
if "chain" in distribution:
pc_kwargs["aes"].setdefault("color", ["chain"])
pc_kwargs["aes"].setdefault("overlay", ["chain"])
pc_kwargs = set_wrap_layout(pc_kwargs, plot_bknd, dt_ecdf_ranks)
pc_kwargs["figure_kwargs"].setdefault("sharex", True)
plot_collection = PlotCollection.wrap(
dt_ecdf_ranks,
backend=backend,
**pc_kwargs,
)
aes_by_visuals.setdefault("ecdf_lines", plot_collection.aes_set)
## ecdf_line
ecdf_ls_kwargs = get_visual_kwargs(visuals, "ecdf_lines")
if ecdf_ls_kwargs is not False:
_, _, ecdf_ls_ignore = filter_aes(
plot_collection, aes_by_visuals, "ecdf_lines", sample_dims
)
plot_collection.map(
ecdf_line,
"ecdf_lines",
data=dt_ecdf,
ignore_aes=ecdf_ls_ignore,
**ecdf_ls_kwargs,
)
ci_kwargs = get_visual_kwargs(visuals, "credible_interval")
_, _, ci_ignore = filter_aes(plot_collection, aes_by_visuals, "credible_interval", sample_dims)
if method == "envelope":
if ci_kwargs is not False:
ci_kwargs.setdefault("color", "B1")
ci_kwargs.setdefault("alpha", 0.1)
plot_collection.map(
fill_between_y,
"credible_interval",
data=dt_ecdf,
x=x_ci,
y_bottom=lower_ci,
y_top=upper_ci,
step=True,
ignore_aes=ci_ignore,
**ci_kwargs,
)
else:
suspicious_kwargs = get_visual_kwargs(visuals, "suspicious_points")
_, suspicious_aes, suspicious_ignore = filter_aes(
plot_collection, aes_by_visuals, "suspicious_points", sample_dims
)
if suspicious_kwargs is not False:
if "color" not in suspicious_aes:
suspicious_kwargs.setdefault("color", "B1")
if "marker" not in suspicious_aes:
suspicious_kwargs.setdefault("marker", "C6")
plot_collection.map(
scatter_xy,
"suspicious_points",
data=dt_ecdf,
mask=suspicious_mask,
ignore_aes=suspicious_ignore,
**suspicious_kwargs,
)
plot_collection.map(
set_ylim,
"ylim",
limits=(-epsilon, epsilon),
store_artist=False,
ignore_aes="all",
)
# add p-values as annotations
p_value_kwargs = get_visual_kwargs(visuals, "p_value_text")
if p_value_kwargs is not False:
_, p_value_loop_dims, _, p_value_ignore = filter_aes_full(
plot_collection, aes_by_visuals, "p_value_text", sample_dims
)
# Only annotate variables whose p-value is scalar per subplot
annot_vars = [v for v, da in p_values.items() if set(da.dims) <= p_value_loop_dims]
if annot_vars:
p_value_kwargs.setdefault("text", lambda p: f"p={p:.2f}(α={alpha:.2f}) ")
p_value_kwargs.setdefault("x", 0)
p_value_kwargs.setdefault("y", 0.85 * epsilon)
p_value_kwargs.setdefault("horizontal_align", "left")
plot_collection.map(
annotate_xy,
"p_value_text",
data=p_values[annot_vars],
ignore_aes=p_value_ignore,
store_artist=backend == "none",
**p_value_kwargs,
)
# set xlabel
_, xlabels_aes, xlabels_ignore = filter_aes(
plot_collection, aes_by_visuals, "xlabel", sample_dims
)
xlabel_kwargs = get_visual_kwargs(visuals, "xlabel")
if xlabel_kwargs is not False:
if "color" not in xlabels_aes:
xlabel_kwargs.setdefault("color", "B1")
xlabel_kwargs.setdefault("text", "Fractional ranks")
plot_collection.map(
labelled_x,
"xlabel",
ignore_aes=xlabels_ignore,
subset_info=True,
**xlabel_kwargs,
)
# title
title_kwargs = get_visual_kwargs(visuals, "title")
_, _, title_ignore = filter_aes(plot_collection, aes_by_visuals, "title", sample_dims)
if title_kwargs is not False:
plot_collection.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.map(
remove_axis,
store_artist=backend == "none",
axis="y",
ignore_aes=plot_collection.aes_set,
)
return plot_collection
