Source code for arviz.labels

# pylint: disable=unused-argument
"""Utilities to generate labels from xarray objects."""
from typing import Union

__all__ = [
    "mix_labellers",
    "BaseLabeller",
    "DimCoordLabeller",
    "DimIdxLabeller",
    "MapLabeller",
    "NoVarLabeller",
    "NoModelLabeller",
]


[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 types 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.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: """WIP."""
[docs] def dim_coord_to_str(self, dim, coord_val, coord_idx): """WIP.""" return f"{coord_val}"
[docs] def sel_to_str(self, sel: dict, isel: dict): """WIP.""" if sel: return ", ".join( [ self.dim_coord_to_str(dim, v, i) for (dim, v), (_, i) in zip(sel.items(), isel.items()) ] ) return ""
[docs] def var_name_to_str(self, var_name: Union[str, None]): """WIP.""" return var_name
[docs] def var_pp_to_str(self, var_name, pp_var_name): """WIP.""" 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 == pp_var_name_str: return f"{var_name_str}" return f"{var_name_str} / {pp_var_name_str}"
[docs] def model_name_to_str(self, model_name): """WIP.""" return model_name
[docs] def make_label_vert(self, var_name: Union[str, None], sel: dict, isel: dict): """WIP.""" 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}"
[docs] def make_label_flat(self, var_name: str, sel: dict, isel: dict): """WIP.""" 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}]"
[docs] def make_pp_label(self, var_name, pp_var_name, sel, isel): """WIP.""" names = self.var_pp_to_str(var_name, pp_var_name) return self.make_label_vert(names, sel, isel)
[docs] def make_model_label(self, model_name, label): """WIP.""" model_name_str = self.model_name_to_str(model_name) if model_name_str is None: return label if label is None or label == "": return model_name_str return f"{model_name_str}: {label}"
[docs] class DimCoordLabeller(BaseLabeller): """WIP."""
[docs] def dim_coord_to_str(self, dim, coord_val, coord_idx): """WIP.""" return f"{dim}: {coord_val}"
[docs] class IdxLabeller(BaseLabeller): """WIP."""
[docs] def dim_coord_to_str(self, dim, coord_val, coord_idx): """WIP.""" return f"{coord_idx}"
[docs] class DimIdxLabeller(BaseLabeller): """WIP."""
[docs] def dim_coord_to_str(self, dim, coord_val, coord_idx): """WIP.""" return f"{dim}#{coord_idx}"
[docs] class MapLabeller(BaseLabeller): """WIP."""
[docs] def __init__(self, var_name_map=None, dim_map=None, coord_map=None, model_name_map=None): """WIP.""" 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 self.model_name_map = {} if model_name_map is None else model_name_map
[docs] def dim_coord_to_str(self, dim, coord_val, coord_idx): """WIP.""" 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)
[docs] def var_name_to_str(self, var_name): """WIP.""" var_name_str = self.var_name_map.get(var_name, var_name) return super().var_name_to_str(var_name_str)
[docs] def model_name_to_str(self, model_name): """WIP.""" model_name_str = self.var_name_map.get(model_name, model_name) return super().model_name_to_str(model_name_str)
[docs] class NoVarLabeller(BaseLabeller): """WIP."""
[docs] def var_name_to_str(self, var_name): """WIP.""" return None
[docs] class NoModelLabeller(BaseLabeller): """WIP."""
[docs] def make_model_label(self, model_name, label): """WIP.""" return label