Source code for utopya.eval.plots.graph

"""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)