"""This module provides graph plotting functions.
.. default-domain:: utopya.eval.plots
"""
import copy
import logging
import warnings
from itertools import chain, product
from typing import Callable, Sequence, Union
import networkx as nx
import numpy as np
import xarray as xr
from dantro.tools import adjusted_log_levels as _adjusted_log_levels
from . import PlotHelper, is_plot_func
from ._graph import GraphPlot
from ._mpl import ColorManager
log = logging.getLogger(__name__)
# -----------------------------------------------------------------------------
# Helper Methods
[docs]def graph_array_from_group(
graph_group: "utopya.eval.groups.GraphGroup",
*,
graph_creation: dict = None,
register_property_maps: dict = None,
clear_existing_property_maps: bool = True,
times: dict = None,
sel: dict = None,
isel: dict = None,
) -> xr.DataArray:
"""From a :py:class:`utopya.eval.groups.GraphGroup` creates a
:py:class:`xarray.DataArray` with ``object`` dtype, containing the
:py:class:`networkx.Graph` objects created from the graph group at the
specified points in the group's coordinate space.
From all coordinates provided via the selection kwargs the cartesian
product is taken. Each of those points represents one entry in the returned
:py:class:`xarray.DataArray`.
The selection kwargs in ``graph_creation`` are ignored silently.
Args:
graph_group (utopya.eval.groups.GraphGroup): The graph group to create
the graph array from
graph_creation (dict, optional): Graph creation configuration
register_property_maps (dict, optional): Property maps to be registered
before graph creation
clear_existing_property_maps (bool, optional): Whether to remove any
existing property map at first
times (dict, optional): Equivalent to a ``sel.time`` entry.
(**Deprecated!**)
sel (dict, optional): Select by value
isel (dict, optional): Select by index
Returns:
xarray.DataArray: :py:class:`networkx.Graph` with the respective graph
group coordinates
"""
# Remove all selection kwargs from the configuration
graph_creation = (
copy.deepcopy(graph_creation) if graph_creation is not None else {}
)
graph_creation.pop("at_time", None)
graph_creation.pop("at_time_idx", None)
graph_creation.pop("sel", None)
graph_creation.pop("isel", None)
# Parse the selectors
sel = sel if sel else {}
isel = isel if isel else {}
# Any ``times`` entry is transformed into the respective sel/isel entry.
# TODO The ``times`` argument is deprecated. Remove it eventually.
if times is not None:
warnings.warn(
"The 'times' argument is deprecated and will be removed. Use the "
"'sel' and 'isel' arguments instead to specify a time selection.",
DeprecationWarning,
)
if len(times) > 1:
raise TypeError(
"Received ambiguous animation time specifications. Need "
"_one_ of: from_property, sel , isel"
)
if "from_property" in times:
sel["time"] = list(
graph_group._get_item_or_pmap(times["from_property"])
.coords["time"]
.values
)
elif "sel" in times:
sel["time"] = times["sel"]
elif "isel" in times:
isel["time"] = times["isel"]
# If needed, extract coordinates from property data
for dim in sel:
if isinstance(sel[dim], dict) and "from_property" in sel[dim]:
_k = sel[dim]["from_property"]
log.note(
"Extracting graph animation selection for dimension '%s' "
"from property '%s' ...",
dim,
_k,
)
_vals = graph_group._get_item_or_pmap(_k).coords[dim]
# subselection
_isel = sel[dim].get("isel") # for subselection
_sel = sel[dim].get("sel") # for subselection
if _isel:
log.remark("Performing index-based subselection: %s", _isel)
_vals = _vals.isel(_isel)
if _sel:
log.remark("Performing value-based subselection: %s", _sel)
_vals = _vals.sel(_sel)
# assign selector values
log.remark("Selector for '%s' dimension:\n%s", dim, _vals)
sel[dim] = list(_vals.values)
# With the given selectors, can now set up the graph-DataArray and create
# a graph from each coordinate combination.
# In order to take the cartesion product of all (sel and isel) selectors,
# collect them all and separate them again afterwards.
dims = list(sel.keys()) + list(isel.keys())
coords = {d: [] for d in dims}
all_selectors = [
c if isinstance(c, (list, tuple, np.ndarray)) else [c]
for c in chain(sel.values(), isel.values())
]
graphs = np.empty(tuple(len(c) for c in all_selectors), dtype="object")
indices = [[i for i in range(len(c))] for c in all_selectors]
N = len(list(product(*indices)))
log.note("Creating %d graphs for animation ...", N)
for i, (idxs, selector) in enumerate(
zip(product(*indices), product(*all_selectors))
):
sel_coords = selector[: len(sel)]
isel_coords = selector[len(sel) :]
# Prepare the selectors for the single points in coordinate space
select = {d: c for d, c in zip(sel.keys(), sel_coords)}
iselect = {d: c for d, c in zip(isel.keys(), isel_coords)}
with _adjusted_log_levels(("dantro.groups.graph", 21)):
g = GraphPlot.create_graph_from_group(
graph_group=graph_group,
register_property_maps=register_property_maps,
clear_existing_property_maps=clear_existing_property_maps,
sel=select,
isel=iselect,
**graph_creation,
)
graphs[idxs] = g
# Extract the coordinate values from the networkx graph attributes
# NOTE This assumes that there are no additional graph attributes that
# also match a dimension name ...
for d, c in g.graph.items():
if d not in coords:
continue
if coords[d] == [] or coords[d][-1] != c:
coords[d].append(c)
print(f" {i+1:4d} / {N}: {g}", end="\r")
coords = {d: coords[d] for d in dims}
# Combine the data and coordinate information in a DataArray
graphs = xr.DataArray(graphs, dims=dims, coords=coords)
return graphs
[docs]def graph_animation_update(
*,
hlpr: PlotHelper,
graphs: xr.DataArray = None,
graph_group=None,
graph_creation: dict = None,
register_property_maps: dict = None,
clear_existing_property_maps: bool = True,
positions: dict = None,
animation_kwargs: dict = None,
suptitle_kwargs: dict = None,
suptitle_formatter: Callable = None,
**drawing_kwargs,
):
"""Graph animation frame generator. Yields whenever the plot helper may
grab the current frame.
If ``graphs`` is given, the networkx graphs in the array are used to create
the frames.
Otherwise, use a graph group. The frames are defined via the selectors in
``animation_kwargs``. From all provided coordinates the cartesian product
is taken. Each of those points defines one graph and thus one frame.
The selection kwargs in ``graph_creation`` are ignored silently.
Args:
hlpr (PlotHelper): The plot helper
graphs (xarray.DataArray, optional): Networkx graphs to draw. The array
will be flattened beforehand.
graph_group (None, optional): Required if ``graphs`` is None. The
GraphGroup from which to generate the animation frames as specified
via sel and isel in ``animation_kwargs``.
graph_creation (dict, optional): Graph creation configuration. Passed
to :py:meth:`~utopya.eval.plots._graph.GraphPlot.create_graph_from_group`
if ``graph_group`` is given.
register_property_maps (dict, optional): Passed to
:py:meth:`~utopya.eval.plots._graph.GraphPlot.create_graph_from_group`
if ``graph_group`` is given.
clear_existing_property_maps (bool, optional): Passed to
:py:meth:`~utopya.eval.plots._graph.GraphPlot.create_graph_from_group`
if ``graph_group`` is given.
positions (dict, optional): The node position configuration.
If ``update_positions`` is True the positions are reconfigured for
each frame.
animation_kwargs (dict, optional): Animation configuration. The
following arguments are allowed:
times (dict, optional):
*Deprecated*: Equivaluent to a ``sel.time`` entry.
sel (dict, optional):
Select by value. Coordinate values (or ``from_property`` entry)
keyed by dimension name.
isel (dict, optional):
Select by index. Coordinate indices keyed by dimension. May be
given together with ``sel`` if no key appears in both.
update_positions (bool, optional):
Whether to reconfigure the node positions for each frame
(default=False).
update_colormapping (bool, optional):
Whether to reconfigure the nodes' and edges'
:py:class:`~dantro.plot.utils.color_mngr.ColorManager` for
each frame (default=False). If False, the colormapping (and the
colorbar) is configured with the first frame and then fixed.
skip_empty_frames (bool, optional):
Whether to skip the frames where the selected graph is missing
or of a type different than ``nx.Graph`` (default=False).
If False, such frames are empty.
suptitle_kwargs (dict, optional): Passed on to the
:py:class:`~utopya.eval.plothelper.PlotHelper` and its
``set_suptitle`` helper function. Only used in animation mode.
The ``title`` can be a format string containing a placeholder with
the dimension name as key for each dimension along which selection
is done. The format string is updated for each frame of the
animation. The default is ``<dim-name> = {<dim-name>}`` for each
dimension.
**drawing_kwargs: Passed to
:py:class:`~utopya.eval.plots._graph.GraphPlot`
"""
suptitle_kwargs = (
copy.deepcopy(suptitle_kwargs) if suptitle_kwargs is not None else {}
)
animation_kwargs = (
copy.deepcopy(animation_kwargs) if animation_kwargs is not None else {}
)
update_positions = animation_kwargs.pop("update_positions", False)
update_colormapping = animation_kwargs.pop("update_colormapping", False)
skip_empty_frames = animation_kwargs.pop("skip_empty_frames", False)
sel = animation_kwargs.pop("sel", None)
isel = animation_kwargs.pop("isel", None)
if graphs is None:
graphs = graph_array_from_group(
graph_group=graph_group,
graph_creation=graph_creation,
register_property_maps=register_property_maps,
clear_existing_property_maps=clear_existing_property_maps,
sel=sel,
isel=isel,
**animation_kwargs,
)
else:
# Apply selectors to the graphs DataArray
if sel is not None:
graphs = graphs.sel(**sel)
if isel is not None:
graphs = graphs.isel(**isel)
# Prepare the suptitle format string once for all frames
if "title" not in suptitle_kwargs:
suptitle_kwargs["title"] = "; ".join(
[f"{dim}" + " = {" + dim + "}" for dim in graphs.coords]
)
def get_graph_and_coords(graphs: xr.DataArray):
"""Generator that yields the (graph, coordinates) pairs"""
dims = graphs.dims
if dims:
# Stack all dimensions of the xr.DataArray
graphs = graphs.stack(_flat=dims)
# For each entry, yield the graph and the coords
for el in graphs:
g = el.item()
coords = {
d: c for d, c in zip(dims, el.coords["_flat"].item())
}
# Also get the scalar coordinates which were not stacked
coords.update(
{d: c.item() for d, c in el.coords.items() if d != "_flat"}
)
yield g, coords
else:
# zero-dimensional: extract single graph item
g = graphs.item()
coords = {d: c.item() for d, c in graphs.coords.items()}
yield g, coords
# Don't show the axis. This is also done in GraphPlot.draw but need to do
# it here in case the first frame is empty.
hlpr.ax.axis("off")
# Indicator for things to be done only once after the first GraphPlot.draw
# call.
_first_graph = True
# Always configure the colormanagers on the first GraphPlot.draw call.
# Overwritten by the update_colormapping kwarg afterwards.
_update_colormapping = True
# Loop over all remaining graphs
for g, coords in get_graph_and_coords(graphs):
_missing_val = not isinstance(g, nx.Graph)
# On missing entries, skip the frame if skip_empty_frames=True.
# If skip_empty_frames=False, the suptitle helper is applied but no
# graph is drawn.
if _missing_val and skip_empty_frames:
continue
if not _missing_val:
# Use a new GraphPlot for the next frame such that the `select`
# kwargs are re-evaluated.
with _adjusted_log_levels(
("dantro.plot.utils.color_mngr", 15),
("utopya.eval.plots._graph", 15),
):
log.remark(
"Drawing %s ... (%s)",
g,
", ".join(f"{k}: {v}" for k, v in coords.items()),
)
gp = GraphPlot(
g=g,
fig=hlpr.fig,
ax=hlpr.ax,
positions=positions,
**drawing_kwargs,
)
gp.draw(
suppress_cbar=not _update_colormapping,
update_colormapping=_update_colormapping,
)
# Update the suptitle format string and invoke the helper
st_kwargs = copy.deepcopy(suptitle_kwargs)
if suptitle_formatter is None:
suptitle_formatter = lambda fstr, **coords: fstr.format(**coords)
st_kwargs["title"] = suptitle_formatter(st_kwargs["title"], **coords)
hlpr.invoke_helper("set_suptitle", **st_kwargs)
if _first_graph:
# Fix the y-position of the suptitle after the first invocation
# since repetitive invocations of the set-suptitle helper would
# re-adjust the subplots size each time. Use the matplotlib default
# if not given.
if "y" not in suptitle_kwargs:
suptitle_kwargs["y"] = 0.98
_update_colormapping = update_colormapping
# Let the writer grab the current frame
yield
if not _missing_val:
# Clean up the frame
with _adjusted_log_levels(("utopya.eval.plots._graph", 15)):
gp.clear_plot(keep_colorbars=not _update_colormapping)
if _first_graph:
# If the positions should be fixed, overwrite the positions arg
if not update_positions:
positions = dict(from_dict=gp.positions)
# Done handling the first GraphPlot.draw call
_first_graph = False
# -----------------------------------------------------------------------------
# Plot functions
[docs]@is_plot_func(use_dag=True, supports_animation=True)
def draw_graph(
*,
hlpr: PlotHelper,
data: dict,
graph_group_tag: str = "graph_group",
graph: Union[nx.Graph, xr.DataArray] = None,
graph_creation: dict = None,
graph_drawing: dict = None,
graph_animation: dict = None,
register_property_maps: Sequence[str] = None,
clear_existing_property_maps: bool = True,
suptitle_kwargs: dict = None,
suptitle_formatter: Callable = None,
):
"""Draws a graph either from a :py:class:`~utopya.eval.groups.GraphGroup`
or directly from a :py:class:`networkx.Graph` using the
:py:class:`~utopya.eval.plots._graph.GraphPlot` class.
If the graph object is to be created from a graph group the latter needs to
be selected via the TransformationDAG. Additional property maps can also be
made available for plotting, see ``register_property_map`` argument.
Animations can be created either from a graph group by using the select
interface in ``graph_animation`` or by passing a DataArray of networkx
graphs via the ``graph`` argument.
For more information on how to use the transformation framework, refer to
the `dantro documentation <https://dantro.readthedocs.io/en/stable/plotting/plot_data_selection.html>`_
and its :py:class:`~dantro.dag.TransformationDAG` class.
For more information on how to configure the graph layout refer to the
:py:class:`~utopya.eval.plots._graph.GraphPlot` documentation.
Args:
hlpr (PlotHelper): The PlotHelper instance for this plot
data (dict): Data from TransformationDAG selection
graph_group_tag (str, optional): The TransformationDAG tag of the graph
group
graph (Union[networkx.Graph, xarray.DataArray], optional): If given,
the ``data`` and ``graph_creation`` arguments are ignored and this
graph is drawn directly.
If a DataArray of graphs is given, the first graph is drawn for a
single graph plot. In animation mode the (flattened) array
represents the animation frames.
graph_creation (dict, optional): Configuration of the graph creation.
Passed to ``GraphGroup.create_graph``.
graph_drawing (dict, optional): Configuration of the graph layout.
Passed to :py:class:`~utopya.eval.plots._graph.GraphPlot`.
graph_animation (dict, optional): Animation configuration. The
following arguments are allowed:
times (dict, optional):
*Deprecated*: Equivaluent to a ``sel.time`` entry.
sel (dict, optional):
Select by value. Dictionary with dimension names as keys. The
values may either be coordinate values or a dict with a
``from_property`` (str) entry which specifies a container
within the :py:class:`~utopya.eval.groups.GraphGroup` or
registered external data from which the coordinates are
extracted.
Additionally, there can be ``isel`` and ``sel`` keys beside the
``from_property`` key, performing a subselection on those
coordinates.
isel (dict, optional):
Select by index. Coordinate indices keyed by dimension. May be
given together with ``sel`` if no key appears in both.
update_positions (bool, optional):
Whether to update the node positions for each frame by
recalculating the layout with the parameters specified in
graph_drawing.positions. If this parameter is not given or
false, the positions are calculated once initially and then
fixed.
update_colormapping (bool, optional):
Whether to reconfigure the nodes' and edges'
:py:class:`~dantro.plot.utils.color_mngr.ColorManager` for
each frame (default=False). If False, the colormapping (and the
colorbar) is configured with the first frame and then fixed.
skip_empty_frames (bool, optional):
Whether to skip the frames where the selected graph is missing
or of a type different than :py:class:`networkx.Graph`
(default=False). If False, such frames are empty.
register_property_maps (Sequence[str], optional): Names of properties
to be registered in the graph group before the graph creation.
The property names must be valid TransformationDAG tags, i.e., be
available in ``data``. Note that the tags may not conflict with any
valid path reachable from inside the selected
:py:class:`~utopya.eval.groups.GraphGroup`.
clear_existing_property_maps (bool, optional): Whether to clear any
existing property maps from the selected
:py:class:`~utopya.eval.groups.GraphGroup`. This is enabled by
default to reduce side effects from previous plots. Set this to
False if you have property maps registered with the GraphGroup
that you would like to keep.
suptitle_kwargs (dict, optional): Passed on to the PlotHelper's
``set_suptitle`` helper function. Only used in animation mode.
The ``title`` can be a format string containing a placeholder with
the dimension name as key for each dimension along which selection
is done. The format string is updated for each frame of the
animation. The default is ``<dim-name> = {<dim-name>}`` for each
dimension.
suptitle_formatter (Callable, optional): If given, this should be a
function that takes coordinates and their values as keyword
arguments and returns the string that is to be used as the title.
Only used in animation mode.
Raises:
ValueError: On invalid or non-computed
:py:class:`~dantro.dag.TransformationDAG` tags in
``register_property_maps`` or invalid graph group tag.
"""
# Work on a copy such that the original configuration is not modified
graph_drawing = copy.deepcopy(graph_drawing if graph_drawing else {})
# Get the sub-configurations for the drawing of the graph
select = graph_drawing.get("select", {})
pos_kwargs = graph_drawing.get("positions", {})
node_kwargs = graph_drawing.get("nodes", {})
edge_kwargs = graph_drawing.get("edges", {})
node_label_kwargs = graph_drawing.get("node_labels", {})
edge_label_kwargs = graph_drawing.get("edge_labels", {})
mark_nodes_kwargs = graph_drawing.get("mark_nodes", {})
mark_edges_kwargs = graph_drawing.get("mark_edges", {})
def get_dag_data(tag):
try:
return data[tag]
except KeyError as err:
_available_tags = ", ".join(data.keys())
raise ValueError(
f"No tag '{tag}' found in the data selected by the DAG! Make "
"sure the tag is named correctly and is selected to be "
"computed; adjust the 'compute_only' argument if needed."
"\nThe following tags are available in the DAG results: "
f"{_available_tags}"
) from err
# Prepare graph group and external property data
graph_group = get_dag_data(graph_group_tag) if graph is None else None
property_maps = None
if register_property_maps:
property_maps = {}
for tag in register_property_maps:
property_maps[tag] = get_dag_data(tag)
# If not in animation mode, make a single graph plot
if not hlpr.animation_enabled:
# Set up a GraphPlot instance
if graph_group is not None:
# Create GraphPlot from graph group
gp = GraphPlot.from_group(
graph_group=graph_group,
graph_creation=graph_creation,
register_property_maps=property_maps,
clear_existing_property_maps=clear_existing_property_maps,
fig=hlpr.fig,
ax=hlpr.ax,
**graph_drawing,
)
else:
if graph_creation is not None:
warnings.warn(
"Received both a 'graph' argument and a 'graph_creation' "
"configuration. The latter will be ignored. To remove "
"this warning set graph_creation to None.",
UserWarning,
)
if isinstance(graph, xr.DataArray):
# Use the first array element for a single graph plot
g = graph.values.flat[0]
if graph.size > 1:
log.caution(
"Received a DataArray of size %d as 'graph' argument, "
"performing a single plot using the first entry. "
"Animations must be enabled by setting "
"animation.enabled to True.",
graph.size,
)
else:
g = graph
# Create a GraphPlot from a nx.Graph
gp = GraphPlot(g=g, fig=hlpr.fig, ax=hlpr.ax, **graph_drawing)
# Make the actual plot via the GraphPlot
gp.draw()
# In animation mode, register the animation frame generator
else:
# Prepare animation kwargs for the update routine
graph_animation = graph_animation if graph_animation else {}
def update():
"""The animation frames generator.
See :py:meth:`~utopya.eval.plots.dag.graph.graph_animation_update`.
"""
if graph_group is None and not isinstance(graph, xr.DataArray):
raise TypeError(
"Failed to create animation due to invalid type of the "
"'graph' argument. Required: xr.DataArray (filled "
f"with graph objects). Received: {type(graph)}"
)
yield from graph_animation_update(
hlpr=hlpr,
graphs=graph if isinstance(graph, xr.DataArray) else None,
graph_group=graph_group,
graph_creation=graph_creation,
register_property_maps=property_maps,
clear_existing_property_maps=clear_existing_property_maps,
suptitle_kwargs=suptitle_kwargs,
suptitle_formatter=suptitle_formatter,
animation_kwargs=graph_animation,
**graph_drawing,
)
hlpr.register_animation_update(update, invoke_helpers_before_grab=True)
