# pylint: disable=unused-argument
"""Utilities to generate labels from xarray objects."""
from collections.abc import Hashable, Iterable, Mapping, Sequence
from typing import Protocol
__all__ = [
"mix_labellers",
"Labeller",
"BaseLabeller",
"DimCoordLabeller",
"IdxLabeller",
"DimIdxLabeller",
"MapLabeller",
"NoVarLabeller",
]
class Labeller(Protocol):
"""Placeholder for type checking ``labeller`` used in docstrings."""
def make_label_vert( # noqa: D102
self,
var_name: str | None,
sel: Mapping[Hashable, Hashable],
isel: Mapping[Hashable, int | Sequence[int]],
) -> str: ...
def make_label_flat( # noqa: D102
self,
var_name: str | None,
sel: Mapping[Hashable, Hashable],
isel: Mapping[Hashable, int | Sequence[int]],
) -> str: ...
[docs]
def mix_labellers(labellers, class_name="MixtureLabeller"):
"""Combine Labeller classes dynamically.
The Labeller class aims to split plot labeling in ArviZ into atomic tasks to maximize
extensibility, and the few classes provided are designed with small deviations
from the base class, in many cases only one method is modified by the child class.
It is to be expected then to want to use multiple classes "at once".
This functions helps combine classes dynamically.
For a general overview of ArviZ label customization, including
``mix_labellers``, see the :ref:`label_guide` page.
Parameters
----------
labellers : iterable of type
Iterable of Labeller types to combine
class_name : str, optional
The name of the generated class
Returns
-------
type
Mixture class object. **It is not initialized**, and it should be
initialized before passing it to ArviZ functions.
Examples
--------
Combine the :class:`~arviz.labels.DimCoordLabeller` with the
:class:`~arviz.labels.MapLabeller` to generate labels in the style of the
``DimCoordLabeller`` but using the mappings defined by ``MapLabeller``.
Note that this works even though both modify the same methods because
``MapLabeller`` implements the mapping and then calls `super().method`.
.. jupyter-execute::
from arviz_base.labels import mix_labellers, DimCoordLabeller, MapLabeller
l1 = DimCoordLabeller()
sel = {"dim1": "a", "dim2": "top"}
print(f"Output of DimCoordLabeller alone > {l1.sel_to_str(sel, sel)}")
l2 = MapLabeller(dim_map={"dim1": "$d_1$", "dim2": r"$d_2$"})
print(f"Output of MapLabeller alone > {l2.sel_to_str(sel, sel)}")
l3 = mix_labellers(
(MapLabeller, DimCoordLabeller)
)(dim_map={"dim1": "$d_1$", "dim2": r"$d_2$"})
print(f"Output of mixture labeller > {l3.sel_to_str(sel, sel)}")
We can see how the mappings are taken into account as well as the dim+coord style. However,
he order in the ``labellers`` arg iterator is important! See for yourself:
.. jupyter-execute::
l4 = mix_labellers(
(DimCoordLabeller, MapLabeller)
)(dim_map={"dim1": "$d_1$", "dim2": r"$d_2$"})
print(f"Output of inverted mixture labeller > {l4.sel_to_str(sel, sel)}")
"""
return type(class_name, labellers, {})
[docs]
class BaseLabeller:
"""Base labeller class.
The default labels for "theta" variable on the subset corresponding to "chain" 0
and "school" "Name" are:
.. code-block::
:caption: Single line label
theta[0, Name]
.. code-block::
:caption: Multi-line label
theta
0, Name
See Also
--------
:ref:`label_guide`: Tutorial page on using labellers with ArviZ
"""
def dim_coord_to_str(self, dim, coord_val, coord_idx): # pylint: disable=no-self-use
"""Format a single dimension name, its value and positional indexes as a string.
Parameters
----------
dim : hashable
Dimension name
coord_val : hashable
Coordinate label
coord_idx : int or sequence of int
Positions along `dim` where its coordinate is `coord_val`
Returns
-------
str
"""
return f"{coord_val}"
def sel_to_str(self, sel, isel):
"""Format selection dictionaries as a string.
Parameters
----------
sel : mapping of {hashable : hashable}
isel : mapping of {hashable : int or sequence of int}
Returns
-------
str
"""
if sel:
return ", ".join(
[
self.dim_coord_to_str(dim, v, i)
for (dim, v), (_, i) in zip(sel.items(), isel.items())
]
)
return ""
def var_name_to_str(self, var_name): # pylint: disable=no-self-use
"""Format a variable name as a string.
Parameters
----------
var_name : str or None
The variable name. It should accept ``None``
Returns
-------
str or None
"""
return var_name
def var_pp_to_str(self, var_name, pp_var_name):
"""Format the corresponding variable names for observation and posterior predictive.
Parameters
----------
var_name : str or None
pp_var_name : str or None
Returns
-------
str or None
"""
var_name_str = self.var_name_to_str(var_name)
pp_var_name_str = self.var_name_to_str(pp_var_name)
if var_name_str is None and pp_var_name_str is None:
return None
if var_name_str == pp_var_name_str:
return var_name_str
return f"{var_name_str} / {pp_var_name_str}"
def make_label_vert(self, var_name, sel, isel):
"""Format variable name and corresponding subset as a multiline string.
Parameters
----------
var_name : str or None
sel : mapping of {hashable : hashable}
isel : mapping of {hashable : int or sequence of int}
Returns
-------
str
"""
var_name_str = self.var_name_to_str(var_name)
sel_str = self.sel_to_str(sel, isel)
if not sel_str:
return "" if var_name_str is None else var_name_str
if var_name_str is None:
return sel_str
return f"{var_name_str}\n{sel_str}"
def make_label_flat(self, var_name: str, sel: dict, isel: dict):
"""Format variable name and corresponding subset as a single line string.
Parameters
----------
var_name : str or None
sel : mapping of {hashable : hashable}
isel : mapping of {hashable : int or sequence of int}
Returns
-------
str
"""
var_name_str = self.var_name_to_str(var_name)
sel_str = self.sel_to_str(sel, isel)
if not sel_str:
return "" if var_name_str is None else var_name_str
if var_name_str is None:
return sel_str
return f"{var_name_str}[{sel_str}]"
def make_pp_label(self, var_name, pp_var_name, sel, isel):
"""Format obs+pp variable name plus corresponding subsets as a multiline string.
Parameters
----------
var_name, pp_var_name : str or None
sel : mapping of {hashable : hashable}
isel : mapping of {hashable : int or sequence of int}
Returns
-------
str
"""
names = self.var_pp_to_str(var_name, pp_var_name)
return self.make_label_vert(names, sel, isel)
[docs]
class DimCoordLabeller(BaseLabeller):
"""Labeller class to show both dimension and coordinate value information.
The default labels for "theta" variable on the subset corresponding to "chain" 0
and "school" "Name" are:
.. code-block::
:caption: Single line label
theta[chain: 0, school: Name]
.. code-block::
:caption: Multi-line label
theta
chain: 0, school: Name
See Also
--------
:ref:`label_guide`: Tutorial page on using labellers with ArviZ
"""
def dim_coord_to_str(self, dim, coord_val, coord_idx):
"""Format a single dimension name, its value and positional indexes as a string.
Parameters
----------
dim : hashable
Dimension name
coord_val : hashable
Coordinate label
coord_idx : int or sequence of int
Positions along `dim` where its coordinate is `coord_val`
Returns
-------
str
"""
return f"{dim}: {coord_val}"
[docs]
class IdxLabeller(BaseLabeller):
"""Labeller class to show positional index information.
The default labels for "theta" variable on the subset corresponding to "chain" 0
and "school" "Name" (3rd and 5th element in the school coordinate values) are:
.. code-block::
:caption: Single line label
theta[0, 2,4]
.. code-block::
:caption: Multi-line label
theta
0, 2,4
See Also
--------
:ref:`label_guide`: Tutorial page on using labellers with ArviZ
"""
def dim_coord_to_str(self, dim, coord_val, coord_idx):
"""Format a single dimension name, its value and positional indexes as a string.
Parameters
----------
dim : hashable
Dimension name
coord_val : hashable
Coordinate label
coord_idx : int or sequence of int
Positions along `dim` where its coordinate is `coord_val`
Returns
-------
str
"""
if not isinstance(coord_idx, Iterable):
return f"{coord_idx}"
return f"{','.join(str(idx) for idx in coord_idx)}"
[docs]
class DimIdxLabeller(BaseLabeller):
"""Labeller class to show both dimension and positional index information.
The default labels for "theta" variable on the subset corresponding to "chain" 0
and "school" "Name" (3rd and 5th element in the school coordinate values) are:
.. code-block::
:caption: Single line label
theta[chain#0, school#2,4]
.. code-block::
:caption: Multi-line label
theta
chain#0, school#2,4
See Also
--------
:ref:`label_guide`: Tutorial page on using labellers with ArviZ
"""
def dim_coord_to_str(self, dim, coord_val, coord_idx):
"""Format a single dimension name, its value and positional indexes as a string.
Parameters
----------
dim : hashable
Dimension name
coord_val : hashable
Coordinate label
coord_idx : int or sequence of int
Positions along `dim` where its coordinate is `coord_val`
Returns
-------
str
"""
if not isinstance(coord_idx, Iterable):
return f"{dim}#{coord_idx}"
return f"{dim}#{','.join(str(idx) for idx in coord_idx)}"
[docs]
class MapLabeller(BaseLabeller):
"""Labeller class to perform provided replacements to elements when converting to string.
It is a subclass of :class:`BaseLabeller` so the base behaviour is the same,
but we can define replacements through dictionaries when initializing the class.
The default labels for "theta" variable on the subset corresponding to "chain" 0
and "school" "Name" with a replacement mapping on the variable name ``{"theta": "θ"}``
and one on the coordinate values ``{"Name": "𝑁𝑎𝑚𝑒"}``
.. code-block::
:caption: Single line label
θ[0, 𝑁𝑎𝑚𝑒]
.. code-block::
:caption: Multi-line label
θ
0, 𝑁𝑎𝑚𝑒
See Also
--------
:ref:`label_guide`: Tutorial page on using labellers with ArviZ
"""
[docs]
def __init__(self, var_name_map=None, dim_map=None, coord_map=None):
"""Initialize a MapLabeller class.
Parameters
----------
var_name_map, dim_map : mapping of {hashable : hashable}, optional
Keys are existing names and values are their respective desired labels.
coord_map : mapping of {hashable : mapping of {hashable : hashable}}, optional
The keys of the first level dictionary are dimension names, the inner
dictionary has existing coord names as keys and their corresponding desired
label as values.
"""
self.var_name_map = {} if var_name_map is None else var_name_map
self.dim_map = {} if dim_map is None else dim_map
self.coord_map = {} if coord_map is None else coord_map
def dim_coord_to_str(self, dim, coord_val, coord_idx):
"""Format a single dimension name, its value and positional indexes as a string.
Parameters
----------
dim : hashable
Dimension name
coord_val : hashable
Coordinate label
coord_idx : int or sequence of int
Positions along `dim` where its coordinate is `coord_val`
Returns
-------
str
"""
dim_str = self.dim_map.get(dim, dim)
coord_str = self.coord_map.get(dim, {}).get(coord_val, coord_val)
return super().dim_coord_to_str(dim_str, coord_str, coord_idx)
def var_name_to_str(self, var_name):
"""Format a variable name as a string.
Parameters
----------
var_name : str or None
The variable name. It should accept ``None``
Returns
-------
str or None
"""
var_name_str = self.var_name_map.get(var_name, var_name)
return super().var_name_to_str(var_name_str)
[docs]
class NoVarLabeller(BaseLabeller):
"""Labeller class to exclude the variable name from the generated labels.
The default labels for "theta" variable on the subset corresponding to "chain" 0
and "school" "Name" are:
.. code-block::
:caption: Single line label
0, Name
.. code-block::
:caption: Multi-line label
0, Name
See Also
--------
:ref:`label_guide`: Tutorial page on using labellers with ArviZ
"""
def var_name_to_str(self, var_name):
"""Format a variable name as a string.
Parameters
----------
var_name : str or None
The variable name. It should accept ``None``
Returns
-------
str or None
"""
return None
