"""This module provides the :py:class:`.GraphPlot` class."""
import copy
import logging
import os
import warnings
from typing import Any, Callable, Dict, Sequence, Tuple, Union
import matplotlib as mpl
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
from ._mpl import ColorManager
log = logging.getLogger(__name__)
# -----------------------------------------------------------------------------
POSITIONING_MODELS_NETWORKX = {
"spring": nx.spring_layout,
"circular": nx.circular_layout,
"shell": nx.shell_layout,
"bipartite": nx.bipartite_layout,
"kamada_kawai": nx.kamada_kawai_layout,
"planar": nx.planar_layout,
"random": nx.random_layout,
"spectral": nx.spectral_layout,
"spiral": nx.spiral_layout,
"graphviz": nx.nx_agraph.graphviz_layout,
}
"""Available networkx node layout options"""
# -----------------------------------------------------------------------------
[docs]class GraphPlot:
"""This class provides an interface for visualizing a
:py:class`networkx.Graph` object or a graph created from a
:py:class:`~utopya.eval.groups.GraphGroup`.
"""
[docs] def __init__(
self,
g: nx.Graph,
*,
fig: "matplotlib.figure.Figure" = None,
ax: "matplotlib.axes.Axes" = None,
select: dict = None,
positions: dict = None,
nodes: dict = None,
edges: dict = None,
node_labels: dict = None,
edge_labels: dict = None,
mark_nodes: dict = None,
mark_edges: dict = None,
):
"""Initializes a ``GraphPlot``, which provides drawing utilities for a
fixed graph.
The drawing kwargs are stored and used when calling
:py:meth:`.draw`.
A ``GraphPlot`` can also be initialized from a
:py:class:`~utopya.eval.groups.GraphGroup` via the
:py:meth:`.from_group` classmethod.
If drawing multiple times from the same ``GraphPlot`` instance, be
aware that it only keeps track of the nodes/edges/labels/colorbars that
were last associated with it. Use
:py:meth:`.clear_plot` before
re-drawing on the same axis.
.. note::
For some graph drawing kwargs it is possible to configure an
automatized mapping from node/edge properties.
This *property mapping* has the following syntax:
.. code-block:: yaml
some_layout_property:
from_property: my_node_or_edge_property
scale_to_interval: [low_lim, up_lim]
The ``from_property`` specifies the node or edge property to be
mapped from. If ``scale_to_interval`` is given, the layout property
values are rescaled linearly the specified interval.
Args:
g (networkx.Graph): The networkx graph object
fig (matplotlib.figure.Figure, optional): The matplotlib figure
used for drawing
ax (matplotlib.axes.Axes, optional): The matplotlib axes used for
drawing
select (dict, optional):
Draw only a subgraph induced by a selection of nodes. Either
select a list of nodes by passing the ``nodelist`` argument or
do a radial node selection by specifying a ``center`` node and
the ``radius``. The following arguments can be passed
additionally:
open_edges (bool, optional):
Whether to draw the edges for which only one of source and
destination is in the set of selected nodes. Disabled by
default.
drop (bool, optional):
Whether to remove the non-selected nodes from the graph.
If False, *all* nodes are passed to the node positioning
model. Enabled by default.
positions (dict, optional):
Configuration for the node positioning. The following arguments
are available:
from_dict (dict, optional):
Node positions (2-tuples) keyed by node. If given, the
layouting algorithm given by the ``model`` argument will be
ignored.
model (Union[str, Callable], optional):
The layout model that is used to calculate the node
positions (default=``spring``). Available
`networkx layout models <https://networkx.github.io/documentation/stable/reference/drawing.html#module-networkx.drawing.layout>`_
are: ``spring``, ``circular``, ``shell``, ``bipartite``,
``kamada_kawai``, ``planar``, ``random``, ``spectral``,
``spiral``.
See :py:data:`.POSITIONING_MODELS_NETWORKX`.
If installed, `GraphViz <https://pypi.org/project/graphviz/>`_
models can be selected using the
``agraph_graphviz`` :py:func:`~networkx.drawing.nx_agraph.graphviz_layout`)
layout function, with the ``prog`` argument specifying the layout model.
Options depend on the ``GraphViz`` version but may include:
``dot``, ``neato``, ``fdp``, ``sfdp``, ``twopi``,
``circo``. Other arguments must be passed as a single string separated by a hyphen,
e.g. ``'-len=1000-overlap_scaling=100-sep=100'``.
If the argument is a callable, it is invoked with the graph
as the first positional argument and is expected to return
networkx-compatible node positions, i.e. a mapping from
nodes to a 2-tuple denoting the position.
further kwargs:
Passed on to the chosen layout model.
nodes (dict, optional):
Drawing configuration for the nodes. The following arguments
are available for property mapping:
``node_size``, ``node_color``, ``alpha``.
The following arguments are allowed:
node_size (scalar or sequence of scalars, optional):
The node size (default=300). Available for property
mapping. Can be mapped directly from the nodes' ``degree``,
``in_degree``, or ``out_degree`` by setting the
``from_property`` argument accordingly.
node_color (color or sequene of colors, optional):
Single color (string or RGB(A) tuple or numeric value)
or sequence of colors (default: '#1f78b4'). If numeric
values are specified they will be mapped to colors using
the cmap and vmin, vmax parameters.
If mapped from property it may contain an additional
``map_to_scalar``, which is a dict of numeric target values
keyed by property value. This allows to map from non-numeric
(e.g. categorical) properties.
cmap (optional):
The colormap. Passed as ``cmap`` to
:py:class:`~dantro.plot.utils.color_mngr.ColorManager`.
cmap_norm (optional):
The norm used for the color mapping. Passed as ``norm`` to
:py:class:`~dantro.plot.utils.color_mngr.ColorManager`.
Is overwritten, if a discrete colormap is specified in
``cmap``.
colorbar (dict, optional):
The node colorbar configuration. The following arguments
are allowed:
enabled (bool, optional):
Whether to plot a colorbar. Enabled by default if
``node_color`` is mapped from property.
labels (dict, optional):
Colorbar tick-labels keyed by tick position (see
:py:meth:`~dantro.plot.utils.color_mngr.ColorManager.create_cbar`).
tick_params (dict, optional):
Colorbar axis tick parameters
label (str, optional):
The axis label for the colorbar
label_kwargs (dict, optional):
Further keyword arguments to adjust the aesthetics of
the colorbar label
further kwargs:
Passed on to :py:meth:`~dantro.plot.utils.color_mngr.ColorManager.create_cbar`.
further kwargs:
Passed to `draw_networkx_nodes <https://networkx.github.io/documentation/stable/reference/generated/networkx.drawing.nx_pylab.draw_networkx_nodes.html>`_
when calling :py:meth:`.draw`.
edges (dict, optional):
Drawing configuration for the edges. The following arguments
are available for property mapping: ``edge_color``, ``width``.
The ``edge_color``, ``edge_cmap``, and ``colorbar`` arguments
behave analogously for the edges as nodes.node_color,
nodes.cmap, and nodes.colorbar for the nodes. Any further
kwargs are (after applying property mapping), passed on to
`draw_networkx_edges <https://networkx.github.io/documentation/stable/reference/generated/networkx.drawing.nx_pylab.draw_networkx_edges.html>`_
when calling :py:meth:`.draw`.
If arrows are to be drawn (i.e. for directed edges with
arrows=True), only norms of type matplotlib.colors.Normalize
are allowed.
node_labels (dict, optional):
Drawing configuration for the node labels. The following
arguments are allowed:
enabled (bool, optional):
Whether to draw node labels. Disabled by default. If
enabled, nodes are labeled by their index by default.
show_only (list, optional):
If given, labels are drawn only for the nodes in this list.
labels (dict, optional):
Custom text labels keyed by node. Available for property
mapping.
format (str, optional):
If ``labels`` are mapped from property this format string
containing a ``label`` key is used for all node labels.
decode (str, optional):
Decoding specifier which is applied to all property values
if ``format`` is used.
further kwargs:
Passed on to `draw_networkx_labels <https://networkx.github.io/documentation/stable/reference/generated/networkx.drawing.nx_pylab.draw_networkx_labels.html>`_
when calling :py:meth:`.draw`.
edge_labels (dict, optional):
Drawing configuration for the edge labels. The following
arguments are allowed:
enabled (bool, optional):
Whether to draw edge labels. Disabled by default. If
enabled, edges are labeled by their (source, destination)
pair by default.
show_only (list, optional):
If given, labels are drawn only for the edges (2-tuples)
in this list.
edge_labels (dict, optional):
Custom text labels keyed by edge (2-tuple). Available for
property mapping.
format (str, optional):
If ``edge_labels`` are mapped from property this format
string containing a ``label`` key is used for all edge
labels.
decode (str, optional):
Decoding specifier which is applied to all property values
if ``format`` is used.
further kwargs:
Passed on to `draw_networkx_edge_labels <https://networkx.github.io/documentation/stable/reference/generated/networkx.drawing.nx_pylab.draw_networkx_edge_labels.html>`_
when calling :py:meth:`.draw`.
mark_nodes (dict, optional):
Mark specific nodes by changing their edgecolor. Either specify
a ``color`` for a list of nodes (``nodelist``), or specify a
``colors`` dictionary of colors keyed by node. Updates an
existing ``nodes.edgecolors`` entry.
mark_edges(dict, optional):
Mark specific edges by changing their color. Either specify a
``color`` for a list of edges (``edgelist``), or specify a
``colors`` dictionary of colors keyed by edge (2-tuple).
Updates an existing ``edges.edge_color`` entry.
"""
# Set matplotlib figure and axis
self.fig = fig if fig is not None else plt.gcf()
self.ax = ax if ax is not None else self.fig.gca()
self._g = g.copy()
self._nodes_to_draw = None
self._edges_to_draw = None
self._nodes_to_shrink = None
self.positions = None
self._select_subgraph(**(select if select else {}))
self.parse_positions(**(positions if positions else {}))
# Drawing configurations
# TODO With networkx v2.6 only FancyArrowPatches will be used for
# the edges (https://github.com/networkx/networkx/pull/4360).
# Then, there will be no simple way of showing edge colorbars.
# Remove the distinction between directed graphs (with arrows
# enabled) and undirected graphs at various places.
self._node_colormanager = None
self._edge_colormanager = None
self._node_cbar_kwargs = {}
self._edge_cbar_kwargs = {}
self._show_node_cbar = False
self._show_edge_cbar = False
self._show_node_labels = False
self._show_edge_labels = False
self._node_kwargs = {}
self._edge_kwargs = {}
self._node_label_kwargs = {}
self._edge_label_kwargs = {}
self.parse_nodes(**(nodes if nodes else {}))
self.parse_edges(**(edges if edges else {}))
self.parse_node_labels(**(node_labels if node_labels else {}))
self.parse_edge_labels(**(edge_labels if edge_labels else {}))
self.mark_nodes(**(mark_nodes if mark_nodes else {}))
self.mark_edges(**(mark_edges if mark_edges else {}))
# matplotlib objects
self._mpl_nodes = None
self._mpl_edges = None
self._mpl_node_labels = None
self._mpl_edge_labels = None
self._mpl_node_cbar = None
self._mpl_edge_cbar = None
# .........................................................................
# Properties
@property
def g(self):
"""Get a deep copy of the graph associated with this instance.
Returns:
networkx.Graph: The graph object
"""
return self._g.copy()
# .........................................................................
# Drawing methods
[docs] def draw(
self,
*,
fig: "matplotlib.figure.Figure" = None,
ax: "matplotlib.axes.Axes" = None,
positions: dict = None,
nodes: dict = None,
edges: dict = None,
node_labels: dict = None,
edge_labels: dict = None,
mark_nodes: dict = None,
mark_edges: dict = None,
suppress_cbar: bool = False,
update_colormapping: bool = True,
**add_colorbars,
):
"""Draws the graph associated with the ``GraphPlot`` using the current
drawing configuration.
The current drawing configuration may be temporarily updated for this
plot. The respective arguments accept the same input as in
:py:meth:`.__init__`.
Args:
fig (matplotlib.figure.Figure, optional): matplotlib figure
ax (matplotlib.axes.Axes, optional): matplotlib axis
positions (dict, optional): Position configuration. If given,
the current positions are replaced. If using a node positioning
model the positions are recalculated.
nodes (dict, optional): Temporarily updates the node-kwargs
edges (dict, optional): Temporarily updates the edge-kwargs
node_labels (dict, optional): Temporarily updates the
node-label-kwargs
edge_labels (dict, optional): Temporarily updates the
edge-label-kwargs
mark_nodes (dict, optional): Temporarily mark nodes kwargs
mark_edges (dict, optional): Temporarily mark edges kwargs
suppress_cbar (bool, optional): Whether to suppress the drawing of
colorbars
update_colormapping (bool, optional): Whether to reconfigure the
nodes' and edges'
:py:class:`~dantro.plot.utils.color_mngr.ColorManager`
(default=True). If True, the respective configuration entries
are ignored. Set to False if doing repetitive plotting with
fixed colormapping.
**add_colorbars: Passed to
:py:meth:`.add_colorbars`
"""
# Cache the current drawing configuration
node_colormanager_cache = copy.deepcopy(self._node_colormanager)
edge_colormanager_cache = copy.deepcopy(self._edge_colormanager)
node_cbar_kwargs_cache = copy.deepcopy(self._node_cbar_kwargs)
edge_cbar_kwargs_cache = copy.deepcopy(self._edge_cbar_kwargs)
show_node_cbar_cache = self._show_node_cbar
show_edge_cbar_cache = self._show_edge_cbar
show_node_labels_cache = self._show_node_labels
show_edge_labels_cache = self._show_edge_labels
node_kwargs_cache = copy.deepcopy(self._node_kwargs)
edge_kwargs_cache = copy.deepcopy(self._edge_kwargs)
node_label_kwargs_cache = copy.deepcopy(self._node_label_kwargs)
edge_label_kwargs_cache = copy.deepcopy(self._edge_label_kwargs)
fig = fig if fig is not None else self.fig
ax = ax if ax is not None else self.ax
if positions is not None:
self.parse_positions(**positions)
self.parse_nodes(
update_colormapping=update_colormapping,
**(nodes if nodes else {}),
)
self.parse_edges(
update_colormapping=update_colormapping,
**(edges if edges else {}),
)
self.parse_node_labels(**(node_labels if node_labels else {}))
self.parse_edge_labels(**(edge_labels if edge_labels else {}))
self.mark_nodes(**(mark_nodes if mark_nodes else {}))
self.mark_edges(**(mark_edges if mark_edges else {}))
log.remark("Now drawing ...")
# Draw nodes and edges
self._mpl_nodes = nx.draw_networkx_nodes(
self._g, pos=self.positions, ax=ax, **self._node_kwargs
)
self._mpl_edges = nx.draw_networkx_edges(
self._g, pos=self.positions, ax=ax, **self._edge_kwargs
)
# NOTE networkx does not pass on the norms to the respective matplotlib
# functions. Hence, they need to be set manually. For the edges,
# the cmap also needs to be set manually. Can only be set for the
# edges if graph is undirected or `arrows=False`.
self._mpl_nodes.set_norm(self._node_colormanager.norm)
if not isinstance(self._mpl_edges, list):
self._mpl_edges.set_norm(self._edge_colormanager.norm)
self._mpl_edges.set_cmap(self._edge_colormanager.cmap)
# Draw node labels and edge labels
if self._show_node_labels:
self._mpl_node_labels = nx.draw_networkx_labels(
self._g,
pos=self.positions,
ax=ax,
**self._node_label_kwargs,
)
if self._show_edge_labels:
self._mpl_edge_labels = nx.draw_networkx_edge_labels(
self._g,
pos=self.positions,
ax=ax,
**self._edge_label_kwargs,
)
if not suppress_cbar:
self.add_colorbars(
show_node_cbar=self._show_node_cbar,
show_edge_cbar=self._show_edge_cbar,
fig=fig,
ax=ax,
**add_colorbars,
)
ax.axis("off")
# Restore the previous drawing configuration
self._node_colormanager = node_colormanager_cache
self._edge_colormanager = edge_colormanager_cache
self._node_cbar_kwargs = node_cbar_kwargs_cache
self._edge_cbar_kwargs = edge_cbar_kwargs_cache
self._show_node_cbar = show_node_cbar_cache
self._show_edge_cbar = show_edge_cbar_cache
self._show_node_labels = show_node_labels_cache
self._show_edge_labels = show_edge_labels_cache
self._node_kwargs = node_kwargs_cache
self._edge_kwargs = edge_kwargs_cache
self._node_label_kwargs = node_label_kwargs_cache
self._edge_label_kwargs = edge_label_kwargs_cache
[docs] def add_colorbars(
self,
*,
show_node_cbar=True,
show_edge_cbar=True,
fig=None,
ax=None,
remove_previous=True,
**update_cbar_kwargs,
):
"""Adds colorbars for the drawn nodes and edges.
Args:
show_node_cbar (bool, optional): Whether to create a colorbar for
the nodes
show_edge_cbar (bool, optional): Whether to create a colorbar for
the edges
fig (None, optional): matplotlib figure
ax (None, optional): matplotlib axis
remove_previous (bool, optional): Whether the colorbars which are
currently associated with the ``GraphPlot`` are removed.
If False, the GraphPlot still loses track of the colorbars,
they can not be removed via the GraphPlot afterwards.
**update_cbar_kwargs: Update both node and edge colorbar kwargs,
passed to :py:meth:`~dantro.plot.utils.color_mngr.ColorManager.create_cbar`.
"""
fig = fig if fig is not None else self.fig
ax = ax if ax is not None else self.ax
if show_node_cbar:
if remove_previous and self._mpl_node_cbar:
self._mpl_node_cbar.remove()
self._mpl_node_cbar = self._node_colormanager.create_cbar(
self._mpl_nodes,
fig=fig,
ax=ax,
**self._node_cbar_kwargs,
**update_cbar_kwargs,
)
if show_edge_cbar:
if isinstance(self._mpl_edges, list):
# When drawing arrows, draw_networkx_edges returns a list of
# FancyArrowPatches which can not be used directly in
# fig.colorbar, but needs conversion to a PatchCollection and
# manual transfer of the chosen colormap and normalization ...
edge_pc = mpl.collections.PatchCollection(self._mpl_edges)
edge_pc.set_norm(self._edge_colormanager.norm)
edge_pc.set_cmap(self._edge_colormanager.cmap)
else:
edge_pc = self._mpl_edges
if remove_previous and self._mpl_edge_cbar:
self._mpl_edge_cbar.remove()
self._mpl_edge_cbar = self._edge_colormanager.create_cbar(
edge_pc,
fig=fig,
ax=ax,
**self._edge_cbar_kwargs,
**update_cbar_kwargs,
)
[docs] def clear_plot(self, *, keep_colorbars: bool = False):
"""Removes all matplotlib objects associated with the GraphPlot from
the respective axis. The GraphPlot loses track of all those objects,
the respective class attributes are reset.
Args:
keep_colorbars (bool, optional): Whether to keep the node and edge
colorbars. If True, the GraphPlot still loses track of the
colorbars, they can not be removed via the GraphPlot
afterwards.
"""
# Remove colorbars. It is important to remove them before the
# associated mappable.
if not keep_colorbars:
if self._mpl_node_cbar:
try:
self._mpl_node_cbar.remove()
except:
pass
if self._mpl_edge_cbar:
try:
self._mpl_edge_cbar.remove()
except:
pass
# Remove nodes and edges
if self._mpl_nodes:
self._mpl_nodes.remove()
# Distinguish between FancyArrowPatch-list and LineCollection
if isinstance(self._mpl_edges, list):
for i in range(len(self._mpl_edges)):
self._mpl_edges[i].remove()
elif self._mpl_edges:
self._mpl_edges.remove()
# Remove labels
if self._mpl_node_labels:
for label in self._mpl_node_labels.values():
label.remove()
if self._mpl_edge_labels:
for label in self._mpl_edge_labels.values():
label.remove()
# Reset class attributes
self._mpl_node_cbar = None
self._mpl_edge_cbar = None
self._mpl_nodes = None
self._mpl_edges = None
self._mpl_node_labels = None
self._mpl_edge_labels = None
log.remark("Cleared GraphPlot.")
# .........................................................................
# Public helper methods
[docs] @classmethod
def from_group(
cls,
graph_group: "utopya.eval.groups.GraphGroup",
*,
graph_creation: dict = None,
register_property_maps: dict = None,
clear_existing_property_maps: bool = True,
**init_kwargs,
):
"""Initializes a :py:class:`.GraphPlot` from a
:py:class:`~utopya.eval.groups.GraphGroup`.
Args:
graph_group (utopya.eval.groups.GraphGroup): The graph group
graph_creation (dict, optional): Configuration of the graph
creation. Passed to :py:meth:`.create_graph_from_group`
register_property_maps (dict, optional): Properties to be
registered in the graph group before the graph creation keyed
by name
clear_existing_property_maps (bool, optional): Whether to clear any
existing property maps from the graph group
**init_kwargs: Passed to :py:meth:`.__init__`
"""
g = cls.create_graph_from_group(
graph_group,
register_property_maps=register_property_maps,
clear_existing_property_maps=clear_existing_property_maps,
**(graph_creation if graph_creation is not None else {}),
)
return cls(g=g, **init_kwargs)
[docs] @staticmethod
def create_graph_from_group(
graph_group: "utopya.eval.groups.GraphGroup",
*,
register_property_maps: dict = None,
clear_existing_property_maps: bool = True,
**graph_creation,
) -> nx.Graph:
"""Creates a :py:class:`networkx.Graph` from a
:py:class`utopya.eval.groups.GraphGroup`.
Additional property maps may be added to the group beforehand.
Args:
graph_group (utopya.eval.groups.GraphGroup): The group to create
the graph from
register_property_maps (dict, optional): Properties to be
registered in the graph group before the graph creation keyed
by name.
clear_existing_property_maps (bool, optional): Whether to clear any
existing property maps from the graph group.
**graph_creation: Configuration of the graph creation.
Passed on to the ``create_graph`` method implemented in dantro,
:py:meth:`~dantro.groups.graph.GraphGroup.create_graph`.
Returns:
networkx.Graph: The created graph object.
"""
# Register external property data
if register_property_maps:
# Clear existing property maps in order to not have side effects if
# plotting multiple times, e.g. in interactive mode. This is
# important because the graph_group most probably is a reference.
if clear_existing_property_maps:
graph_group.property_maps.clear()
# Can register now
for tag, pmap in register_property_maps.items():
graph_group.register_property_map(tag, pmap)
g = graph_group.create_graph(**graph_creation)
return g
[docs] def parse_positions(
self,
*,
from_dict: Dict[Any, Tuple[float, float]] = None,
model: Union[str, Callable] = None,
**kwargs,
):
"""Parses the node positioning configuration. If a node positioning
model is to be used, (re)calculates the positions.
Args:
from_dict (dict, optional): Explicit node positions (2-tuples keyed
by node). If given, the ``model`` argument will be ignored.
model (Union[str, Callable], optional): The model used for node
positioning. If it is a string, it is looked up from the
available networkx positioning models. If None, the spring
model is used.
If it is callable, it will be called with the graph as first
positional argument.
**kwargs: Passed to the node positioning routine
Raises:
ModuleNotFoundError: If a graphviz model was chosen but pygraphviz
was not importable (via networkx).
"""
# Set spring-layout as default if nothing else is specified
if from_dict is None and model is None:
model = "spring"
if from_dict is not None:
if model is not None:
warnings.warn(
"Node positions were specified *both* via a positioning "
"model and explicitly via the `from_dict` argument. The "
"specified model will be ignored. To remove this warning, "
"set the graph_drawing.positions.model entry to None.",
UserWarning,
)
self.positions = copy.deepcopy(from_dict)
else:
log.remark(
"Calculating the node positions using a positioning model ..."
)
if callable(model):
self.positions = model(self._g, **kwargs)
# else: is a networkx positioning model
else:
self.positions = POSITIONING_MODELS_NETWORKX[model](
self._g, **kwargs
)
[docs] def parse_nodes(
self,
*,
node_size=None,
node_color=None,
alpha=None,
cmap=None,
cmap_norm=None,
vmin: float = None,
vmax: float = None,
edgecolors=None,
colorbar: dict = None,
update_colormapping: bool = True,
**kwargs,
):
"""Parses the node layout configuration and updates the node kwargs of
the GraphPlot.
The following arguments are available for property mapping:
``node_size``, ``node_color``, ``alpha``.
Args:
node_size (None, optional): Size of nodes (default=300). Available
for property mapping. Can be mapped directly from the nodes'
``degree``, ``in_degree``, or ``out_degree`` by setting the
``from_property`` argument accordingly.
node_color (None, optional): Single color (string or RGB(A) tuple
or numeric value) or sequence of colors.
If numeric values are specified they will be mapped to colors
using the cmap and vmin, vmax parameters.
If mapped from property it may contain an additional
``map_to_scalar``, which is a dict of numeric target values
keyed by property value. This allows to map from non-numeric
(e.g. categorical) properties.
alpha (None, optional): The node transparency
cmap (None, optional): The colormap. Passed as ``cmap`` to
:py:class:`~dantro.plot.utils.color_mngr.ColorManager`.
cmap_norm (None, optional): The norm used for the color mapping.
Passed as ``norm`` to
:py:class:`~dantro.plot.utils.color_mngr.ColorManager`.
Is overwritten, if a discrete colormap is specified in
``cmap``.
vmin (float, optional): Minimum for the colormap scaling
vmax (float, optional): Maximum for the colormap scaling
edgecolors (optional): Colors of node borders. The default is
'none', i.e. no node border is drawn.
colorbar (dict, optional): The node colorbar configuration.
The following keys are available:
enabled (bool, optional):
Whether to plot a colorbar. Enabled by default if
``node_color`` is mapped from property.
labels (dict, optional):
Colorbar tick-labels keyed by tick position (see
:py:class:`~dantro.plot.utils.color_mngr.ColorManager`).
further kwargs: Passed on to
:py:meth:`~dantro.plot.utils.color_mngr.ColorManager.create_cbar`
update_colormapping (bool, optional): Whether to reconfigure the
nodes' :py:class:`~dantro.plot.utils.color_mngr.ColorManager`
(default: True).
If False, the respective arguments are ignored.
Set to False if doing repetitive plotting with fixed
colormapping.
**kwargs: Update the node kwargs. Passed to
:py:func:`~networkx.drawing.nx_pylab.draw_networkx_nodes`
when calling :py:meth:`.draw`
"""
if kwargs.pop("nodelist", None) is not None:
warnings.warn(
"The 'nodelist' argument will be ignored. To draw a subset of "
"nodes use the 'select.nodelist' argument instead.",
UserWarning,
)
# Update node kwargs with simple kwargs. All kwargs that might need
# extra treatment are caught explicitly and handled below.
self._node_kwargs.update(kwargs)
self._node_kwargs["nodelist"] = self._nodes_to_draw
# Do the property mapping ...
# Node sizes
if isinstance(node_size, dict) and "from_property" in node_size:
prop = node_size["from_property"]
interval = node_size.get("scale_to_interval", None)
try:
_node_sizes = np.array(
[self._g.nodes[n][prop] for n in self._nodes_to_draw]
)
except KeyError as err:
if prop == "degree":
_node_sizes = node_sizes = np.array(
[self._g.degree[n] for n in self._nodes_to_draw]
)
elif prop == "in_degree":
_node_sizes = node_sizes = np.array(
[self._g.in_degree[n] for n in self._nodes_to_draw]
)
elif prop == "out_degree":
_node_sizes = node_sizes = np.array(
[self._g.out_degree[n] for n in self._nodes_to_draw]
)
else:
raise ValueError(
f"No property {prop} found. Make sure the property "
"exists for all nodes to draw. Additional options: "
"degree, in_degree, out_degree"
) from err
# If there are nodes to be shrinked, set their size to zero
to_shrink = np.isin(self._nodes_to_draw, self._nodes_to_shrink)
_node_sizes[to_shrink] = 0
_node_sizes[~to_shrink] = self._scale_to_interval(
_node_sizes[~to_shrink], interval
)
self._node_kwargs["node_size"] = list(_node_sizes)
elif node_size is not None:
self._node_kwargs["node_size"] = node_size
# Node colors
if isinstance(node_color, dict) and "from_property" in node_color:
prop = node_color["from_property"]
interval = node_color.get("scale_to_interval", None)
# If provided a mapping, map the property values to scalar values
if "map_to_scalar" in node_color:
map_to_scalar = np.vectorize(node_color["map_to_scalar"].get)
_node_colors = list(map_to_scalar(_node_colors))
else:
_node_colors = self._scale_to_interval(
[self._g.nodes[n][prop] for n in self._nodes_to_draw],
interval,
)
self._node_kwargs["node_color"] = _node_colors
self._show_node_cbar = True
elif node_color is not None:
self._node_kwargs["node_color"] = node_color
# Node transparency
if isinstance(alpha, dict) and "from_property" in alpha:
prop = alpha["from_property"]
interval = alpha.get("scale_to_interval", None)
self._node_kwargs["alpha"] = self._scale_to_interval(
[self._g.nodes[n][prop] for n in self._nodes_to_draw], interval
)
elif alpha is not None:
self._node_kwargs["alpha"] = alpha
# ... property mapping done.
# If a single or no value was set for the node size, create an explicit
# node size mapping if there are nodes to be shrinked to size zero.
# If no size was given, use the default size 300.
if self._nodes_to_shrink and not isinstance(
self._node_kwargs.get("node_size", None), list
):
_size = self._node_kwargs.get("node_size", 300)
self._node_kwargs["node_size"] = [
0 if n in self._nodes_to_shrink else _size
for n in self._nodes_to_draw
]
# Update the edgecolors of the nodes. Set the edgecolor to 'none'
# (=transparent) if it has not been set yet.
if edgecolors is not None:
self._node_kwargs["edgecolors"] = edgecolors
elif "edgecolors" not in self._node_kwargs:
self._node_kwargs["edgecolors"] = "none"
# For directed graphs (arrows enabled) individual FancyArrowPatches are
# drawn. Make sure that the arrows end right at the node borders.
if self._g.is_directed():
self._configure_node_patch_sizes()
# Get colorbar configuration and update the existing kwargs
colorbar = colorbar if colorbar is not None else {}
cbar_labels = colorbar.pop("labels", None)
self._show_node_cbar = colorbar.pop("enabled", self._show_node_cbar)
self._node_cbar_kwargs.update(colorbar)
# (Re-)create the colormanager. Only do so if there is no colormanager
# yet or if any kwargs are given.
if update_colormapping and (
self._node_colormanager is None
or any(
[
arg is not None
for arg in (cmap, cmap_norm, vmin, vmax, cbar_labels)
]
)
):
# Replace all kwargs that are None with the current configuration.
# That way the current configuration is not lost but updated.
if vmin is None:
vmin = self._node_kwargs.get("vmin", None)
if vmax is None:
vmax = self._node_kwargs.get("vmax", None)
if self._node_colormanager is not None:
if cmap is None:
cmap = self._node_colormanager.cmap
if cmap_norm is None:
cmap_norm = self._node_colormanager.norm
if cbar_labels is None:
cbar_labels = self._node_colormanager.labels
# Set up the ColorManager
self._node_colormanager = ColorManager(
cmap=cmap,
norm=cmap_norm,
labels=cbar_labels,
vmin=vmin,
vmax=vmax,
)
self._node_kwargs["cmap"] = self._node_colormanager.cmap
[docs] def parse_edges(
self,
*,
width=None,
edge_color=None,
edge_cmap=None,
cmap_norm=None,
alpha=None,
edge_vmin: float = None,
edge_vmax: float = None,
colorbar: dict = None,
update_colormapping: bool = True,
**kwargs,
):
"""Parses the edge layout configuration and updates the edge kwargs of
the GraphPlot.
The following arguments are available for property mapping:
``width``, ``edge_color``.
Args:
width (None, optional): Line width of edges
edge_color (None, optional): Single color (string or RGB(A) tuple
or numeric value) or sequence of colors (default='k').
If numeric values are specified they will be mapped to colors
using the edge_cmap and edge_vmin, edge_vmax parameters.
If mapped from property it may contain an additional
``map_to_scalar``, which is a dict of numeric target values
keyed by property value. This allows to map from non-numeric
(e.g. categorical) properties.
edge_cmap (None, optional): The colormap. Passed as ``cmap`` to
:py:class:`~dantro.plot.utils.color_mngr.ColorManager`.
cmap_norm (None, optional): The norm used for the color mapping.
Passed as ``norm`` to
:py:class:`~dantro.plot.utils.color_mngr.ColorManager`.
Is overwritten, if a discrete colormap is specified in
``edge_cmap``.
If arrows are to be drawn (i.e. for directed edges with
arrows=True), only norms of type matplotlib.colors.Normalize
are allowed.
edge_vmin (float, optional): Minimum for the colormap scaling
edge_vmax (float, optional): Maximum for the colormap scaling
colorbar (dict, optional): The edge colorbar configuration.
The following keys are available:
enabled (bool, optional):
Whether to plot a colorbar. Enabled by default if
``edge_color`` is mapped from property.
labels (dict, optional):
Colorbar tick-labels keyed by tick position (see
:py:class:`~dantro.plot.utils.color_mngr.ColorManager`).
further kwargs: Passed on to
:py:meth:`~dantro.plot.utils.color_mngr.ColorManager.create_cbar`
update_colormapping (bool, optional): Whether to reconfigure the
edges' :py:class:`~dantro.plot.utils.color_mngr.ColorManager`
(default=True).
If False, the respective arguments are ignored.
Set to False if doing repetitive plotting with fixed
colormapping.
**kwargs: Update the edge kwargs. Passed to
:py:func:`~networkx.drawing.nx_pylab.draw_networkx_nodes`
when calling :py:meth:`.draw`.
Raises:
TypeError: On norm type other than
:py:class:`matplotlib.colors.Normalize` and if arrows are to be
drawn.
"""
if kwargs.pop("edgelist", None) is not None:
warnings.warn(
"The 'edgelist' argument will be ignored. Use the select "
"configuration to specify a subgraph to be drawn.",
UserWarning,
)
self._edge_kwargs["edgelist"] = self._edges_to_draw
# Update edge kwargs with simple kwargs. All kwargs that might need
# extra treatment are caught explicitly and handled below.
self._edge_kwargs.update(kwargs)
# Do the property mapping ...
# Edge width
if isinstance(width, dict) and "from_property" in width:
prop = width["from_property"]
interval = width.get("scale_to_interval", None)
self._edge_kwargs["width"] = self._scale_to_interval(
[self._g.edges[n][prop] for n in self._edges_to_draw], interval
)
elif width is not None:
self._edge_kwargs["width"] = width
# Edge colors
if isinstance(edge_color, dict) and "from_property" in edge_color:
prop = edge_color["from_property"]
interval = edge_color.get("scale_to_interval", None)
# If provided a mapping, map the property values to scalar values
if "map_to_scalar" in edge_color:
map_to_scalar = np.vectorize(edge_color["map_to_scalar"].get)
_edge_colors = list(map_to_scalar(_edge_colors))
else:
_edge_colors = self._scale_to_interval(
[self._g.edges[n][prop] for n in self._edges_to_draw],
interval,
)
self._edge_kwargs["edge_color"] = _edge_colors
self._show_edge_cbar = True
elif edge_color is not None:
self._edge_kwargs["edge_color"] = edge_color
# Edge alpha
if isinstance(alpha, dict) and "from_property" in alpha:
prop = alpha["from_property"]
interval = alpha.get("scale_to_interval", None)
self._edge_kwargs["alpha"] = self._scale_to_interval(
[self._g.edges[n][prop] for n in self._edges_to_draw], interval
)
elif alpha is not None:
self._edge_kwargs["alpha"] = alpha
# ... property mapping done.
# Get colorbar configuration and update the existing kwargs
colorbar = colorbar if colorbar is not None else {}
cbar_labels = colorbar.pop("labels", None)
self._show_edge_cbar = colorbar.pop("enabled", self._show_edge_cbar)
self._edge_cbar_kwargs.update(colorbar)
# (Re-)create the colormanager. Only do so if there is no colormanager
# yet or if any kwargs are given.
if update_colormapping and (
self._edge_colormanager is None
or any(
[
arg is not None
for arg in (
edge_cmap,
cmap_norm,
edge_vmin,
edge_vmax,
cbar_labels,
)
]
)
):
# Replace all kwargs that are None with the current configuration.
# That way the current configuration is not lost but updated.
if edge_vmin is None:
edge_vmin = self._edge_kwargs.get("edge_vmin", None)
if edge_vmax is None:
edge_vmax = self._edge_kwargs.get("edge_vmax", None)
if self._edge_colormanager is not None:
if edge_cmap is None:
edge_cmap = self._edge_colormanager.cmap
if cmap_norm is None:
cmap_norm = self._edge_colormanager.norm
if cbar_labels is None:
cbar_labels = self._edge_colormanager.labels
# Set up the ColorManager
self._edge_colormanager = ColorManager(
cmap=edge_cmap,
norm=cmap_norm,
labels=cbar_labels,
vmin=edge_vmin,
vmax=edge_vmax,
)
self._edge_kwargs["edge_cmap"] = self._edge_colormanager.cmap
# NOTE In draw_networkx_edges, the Normalize norm is applied
# explicitly. Since the norm cannot be updated later (as edges
# with arrows are FancyArrowPatches), other norms than Normalize
# are forbidden here. When mapping the colors using the
# Colormanager a colorbar drawn from the edges mappable would
# still be wrong.
if (
isinstance(self._g, nx.DiGraph)
and self._edge_kwargs.get("arrows", True)
and type(self._edge_colormanager.norm) != mpl.colors.Normalize
):
raise TypeError(
"Received invalid norm type: "
f"{type(self._edge_colormanager.norm)}. For directed edges"
" with 'arrows=True', only the matplotlib.colors.Normalize"
" base class is supported."
)
[docs] def parse_node_labels(
self,
*,
enabled: bool = False,
show_only: list = None,
labels: dict = None,
format: str = "{label}",
decode: str = None,
**kwargs,
):
"""Parses the node labels configuration and updates the node label
kwargs of the :py:class:`.GraphPlot`.
Args:
enabled (bool, optional): Whether to draw node labels.
show_only (list, optional): If given, labels are drawn only for the
nodes in this list
labels (dict, optional): Custom text labels keyed by node.
Available for property mapping.
format (str, optional): If ``labels`` are mapped from property this
format string containing a ``label`` key is used for all node
labels.
decode (str, optional): Decoding specifier which is applied to all
property values if ``format`` is used.
**kwargs: Update the node label kwargs. Passed to
nx.draw_networkx_labels when calling
:py:meth:`.draw`.
"""
if enabled:
self._show_node_labels = True
# Update node label kwargs with simple kwargs. All kwargs that might
# need extra treatment are caught explicitly and handled below.
self._node_label_kwargs.update(kwargs)
if not enabled and not labels:
# Don't update the node_labels.labels -> can stop here
return
def to_show(n):
if show_only is None:
return n not in self._nodes_to_shrink
else:
return n in show_only and n not in self._nodes_to_shrink
if labels:
if "from_property" in labels:
prop = labels["from_property"]
_labels = {
n: format.format(
label=(
self._g.nodes[n][prop]
if decode is None
else self._g.nodes[n][prop].decode(decode)
)
)
for n in self._nodes_to_draw
if to_show(n)
}
else:
_labels = copy.deepcopy(labels)
if show_only is not None:
# show_only takes precedence over the provided node labels.
# Keep only those that are in show_only. nodes in show_only
# for which no label is given are labeled with their index.
for n in labels.keys():
if not to_show(n):
del _labels[n]
for n in show_only:
if (
n not in _labels.keys()
and n not in self._nodes_to_shrink
):
_labels[n] = n
elif enabled:
# If enabled but no labels given, label nodes with their index.
_labels = {n: n for n in self._nodes_to_draw if to_show(n)}
self._node_label_kwargs["labels"] = _labels
[docs] def parse_edge_labels(
self,
*,
enabled: bool = False,
show_only: list = None,
edge_labels: dict = None,
format: str = "{label}",
decode: str = None,
**kwargs,
):
"""Parses the edge labels configuration and updates the edge label
kwargs of the GraphPlot.
Args:
enabled (bool, optional): Whether to draw edge labels.
show_only (list, optional): If given, labels are drawn only for the
edges (2-tuples) in this list
edge_labels (dict, optional): Custom text labels keyed by edge
(2-tuple). Available for property mapping.
format (str, optional): If ``edge_labels`` are mapped from property
this format string containing a ``label`` key is used for all
edge labels.
decode (str, optional): Decoding specifier which is applied to all
property values if ``format`` is used.
**kwargs: Update the edge label kwargs. Passed to
nx.draw_networkx_edge_labels when calling
:py:meth:`.draw`.
"""
# Catch a pitfall: There is no 'labels' argument for the edge labels
# (as there is for the node labels), here it is named 'edge_labels'.
if "labels" in kwargs:
raise ValueError(
"Received invalid 'labels' key in edge label configuration. "
"For specifying an edge label dict, use the key 'edge_labels'."
)
if enabled:
self._show_edge_labels = True
# Update edge label kwargs with simple kwargs. All kwargs that might
# need extra treatment are caught explicitly and handled below.
self._edge_label_kwargs.update(kwargs)
if not enabled and not edge_labels:
# Don't update the edge_labels.labels -> can stop here
return
def to_show(e):
return True if show_only is None else e[:2] in show_only
if show_only is not None:
show_only = [tuple(e) for e in show_only]
if edge_labels:
if "from_property" in edge_labels:
prop = edge_labels["from_property"]
_labels = {
e[:2]: format.format(
label=(
self._g.edges[e][prop]
if decode is None
else self._g.edges[e][prop].decode(decode)
)
)
for e in self._edges_to_draw
if to_show(e)
}
else:
_labels = copy.deepcopy(edge_labels)
if show_only is not None:
# show_only takes precedence over the provided edge labels.
# Keep only those that are in show_only. edges in show_only
# for which no label is given are labeled with their
# (source, destination) pair.
for e in edge_labels.keys():
if not to_show(e):
del _labels[e]
for e in show_only:
if e not in _labels.keys():
_labels[e] = e
elif enabled:
# If enabled but no labels given, label edges with their
# (source, destination) pair.
_labels = {e[:2]: e[:2] for e in self._edges_to_draw if to_show(e)}
self._edge_label_kwargs["edge_labels"] = _labels
[docs] def mark_nodes(
self, *, nodelist: list = None, color=None, colors: dict = None
):
"""Mark specific nodes by changing their edgecolor.
.. note::
This function overwrites the ``edgecolors`` entry in the node
kwargs. Thus it might overwrite an existing ``edgecolors`` entry
specified via
:py:meth:`.parse_nodes`
(and vice versa).
Args:
nodelist (list, optional): Nodes to mark with the color specified
via ``color``
color (None, optional): Single edgecolor to use for the nodes in
``nodelist``
colors (dict, optional): Edgecolors keyed by node to mark. Must be
None if ``nodelist`` is given.
Raises:
ValueError: On ambiguous or missing mark_nodes configuration
"""
if nodelist is None and color is None and colors is None:
return
if colors is not None:
if not (nodelist is None and color is None):
raise ValueError(
"Received invalid 'mark_nodes' kwargs. Provide _either_ a "
"'colors' dict only _or_ an 'nodelist' together with a "
"single 'color'."
)
elif (nodelist is None) != (color is None):
_missing_arg = "nodelist" if nodelist is None else "color"
raise ValueError(
f"Missing argument '{_missing_arg}' in 'mark_nodes' kwargs."
)
if self._node_kwargs.get("edgecolors", None) is None:
# There is no edgecolors entry yet. Use the current node colors
# as base colors. If there are none, use the default networkx node
# color.
base_color = self._node_kwargs.get("node_color", "none")
else:
# Use the existing edgecolors entry.
base_color = self._node_kwargs["edgecolors"]
# From the base color(s) create dict of colors keyed by node
if mpl.colors.is_color_like(base_color):
_colors = {n: base_color for n in self._nodes_to_draw}
else:
# ... must be a list of colors.
# If base_color contains numeric values, they need to be
# transformed via the specified colormap.
if not mpl.colors.is_color_like(base_color[0]):
base_color = self._node_colormanager.map_to_color(base_color)
_colors = {
n: base_color[i] for i, n in enumerate(self._nodes_to_draw)
}
# Update the color dict with the values from the mark configuration
if nodelist:
for n in nodelist:
_colors[n] = color
else:
_colors.update(colors)
# The color values are aligned with 'self._nodes_to_draw' since
# dictionaries don't change their ordering.
self._node_kwargs["edgecolors"] = list(_colors.values())
# Reconfigure the node patch sizes which increase if a node boundary
# is drawn.
if self._g.is_directed():
self._configure_node_patch_sizes()
[docs] def mark_edges(
self, *, edgelist: list = None, color=None, colors: dict = None
):
"""Mark specific edges by changing their color.
.. note::
This function overwrites the ``edge_color`` entry in the edge
kwargs. Thus it might overwrite an existing ``edge_color`` entry
specified via
:py:meth:`.parse_edges`
(and vice versa).
Args:
edgelist (list, optional): Edges to mark with the color specified
via ``color``
color (None, optional): Single color to use for the edges in
``edgelist``
colors (dict, optional): Colors keyed by edge (2-tuple) to mark.
Must be None if ``edgelist`` is given.
Raises:
ValueError: On ambiguous or missing mark_edges configuration
"""
if edgelist is None and color is None and colors is None:
return
if colors is not None:
if not (edgelist is None and color is None):
raise ValueError(
"Received invalid 'mark_edges' kwargs. Provide _either_ a "
"'colors' dict only _or_ an 'edgelist' together with a "
"single 'color'."
)
elif (edgelist is None) != (color is None):
_missing_arg = "edgelist" if edgelist is None else "color"
raise ValueError(
f"Missing argument '{_missing_arg}' in 'mark_edges' kwargs."
)
# Create dict of colors keyed by node based on the current edge_color.
# If there is none, use the default networkx edge color.
base_color = self._edge_kwargs.get("edge_color", "k")
if mpl.colors.is_color_like(base_color):
_colors = {e[:2]: base_color for e in self._edges_to_draw}
else:
# ... must be a list of colors. Transform to color-like if needed.
if not mpl.colors.is_color_like(base_color[0]):
base_color = self._edge_colormanager.map_to_color(base_color)
_colors = {
e[:2]: base_color[i] for i, e in enumerate(self._edges_to_draw)
}
# Update the color dict with the values from the mark configuration
if edgelist:
for e in edgelist:
e = tuple(e)
if not isinstance(self._g, nx.DiGraph) and e not in _colors:
e = e[::-1]
_colors[e] = color
else:
for e, c in colors.items():
if not isinstance(self._g, nx.DiGraph) and e not in _colors:
e = e[::-1]
_colors[e] = c
self._edge_kwargs["edge_color"] = list(_colors.values())
# .........................................................................
# Private helper methods
[docs] @staticmethod
def _scale_to_interval(data: list, interval=None) -> list:
"""Rescales the data linearly to the given interval. If no interval is
given the data is returned as it is.
Args:
data (list): data that is rescaled linearly to the given interval
interval (optional): The target interval
Returns:
list: rescaled data
Raises:
ValueError: On invalid interval specification
"""
if interval is None:
return data
try:
lim_low, lim_up = interval
except ValueError as err:
raise ValueError(
"'interval' must be a 2-tuple or list of length 2! Received: "
f"{interval}"
) from err
data = np.array(data)
max_val = np.max(data)
min_val = np.min(data)
if max_val > min_val:
rescaled_data = (data - min_val) / (max_val - min_val) * (
lim_up - lim_low
) + lim_low
else:
# If all values are equal, set them to the mean of the interval
rescaled_data = np.zeros_like(data) + (lim_up - lim_low) / 2.0
return list(rescaled_data)
[docs] def _select_subgraph(
self, nodelist: list = None, drop: bool = True, **kwargs
):
"""Select a subgraph to draw. Sets the lists of nodes and edges to draw
and the nodes to shrink. Either a list of nodes is selected or radial
selection is done.
Args:
nodelist (list, optional): If given, select nodes from list
drop (bool, optional): Whether to remove the non-selected nodes and
edges from the graph
**kwargs: Passed to the selection routine
"""
def select_from_list(*, nodelist: list, open_edges: bool = False):
"""Given a list of nodes, selects all nodes and edges needed for
the graph drawing. If ``open_edges=False``, only those edges are
selected for which both ends are in ``nodes``.
Args:
nodelist (list): Nodes to be selected
open_edges (bool, optional): Whether to draw loose edges
Returns:
Tuple containing list of selected nodes, list of selected
edges, and list of nodes to be shrinked to size zero.
"""
subgraph = nx.induced_subgraph(self._g, nodelist)
if open_edges:
# Create outer subgraph from given nodes and all neighbors
node_selection = set(nodelist)
outer_nodes = set()
for n in nodelist:
outer_nodes.update(nx.all_neighbors(self._g, n))
outer_nodes -= node_selection
node_selection = node_selection.union(outer_nodes)
subgraph_outer = nx.induced_subgraph(self._g, node_selection)
# Set of nodes to shrink is the difference of the two node sets
nodes_to_shrink = list(subgraph_outer.nodes - set(nodelist))
# From the outer subgraph remove edges between outer nodes
edges_to_plot = (
subgraph_outer.edges
- nx.induced_subgraph(self._g, outer_nodes).edges
)
return (
list(subgraph_outer.nodes),
list(edges_to_plot),
nodes_to_shrink,
)
return list(subgraph.nodes), list(subgraph.edges), []
def select_radial(
*, center: int, radius: int, open_edges: bool = False
):
"""Selects all nodes around a given center within a given radius
(measured in numbers of neighborhoods). If ``open_edges=False``, those
edges are selected for which both ends are in the set of selected
nodes.
Args:
center (int): Index of the central node
radius (int): Selection radius
open_edges (bool, optional): Whether to draw loose edges
Returns:
Tuple containing list of selected nodes, list of selected
edges, and list of nodes to be shrinked to size zero.
"""
# After num_nodes-1 iterations (below), all nodes would be selected
if radius > self._g.number_of_nodes() - 1:
radius = self._g.number_of_nodes() - 1
# Identify nodes within the given radius around the central node.
# Start by adding the central nodes and all its neighbors to a set.
# Then, iteratively add all neighbors of the previously added nodes
# to the set, until the given radius is reached.
# TODO It might be worth testing the computational efficiency of this
# (also for large subgraphs) as each node is tried to be added to
# the set at least two times.
# Store the current node selection
node_selection = {center}
# Store the nodes added to the selection in the previous step
nbs_prev = {center}
# Store the new nodes to be selected
nbs_new = set()
for i in range(radius):
for n in nbs_prev:
nbs_new.update(nx.all_neighbors(self._g, n))
nbs_prev = nbs_new - node_selection
node_selection = node_selection.union(nbs_new)
nbs_new.clear()
if open_edges:
# Create inner subgraph from all nodes within the given radius
subgraph_inner = nx.induced_subgraph(self._g, node_selection)
# Create outer subgraph from all nodes within radius=radius+1
for n in nbs_prev:
nbs_new.update(nx.all_neighbors(self._g, n))
outer_nodes = nbs_new - node_selection
node_selection = node_selection.union(nbs_new)
subgraph_outer = nx.induced_subgraph(self._g, node_selection)
# Set of nodes to shrink is the difference of the two node sets
nodes_to_shrink = list(
subgraph_outer.nodes - subgraph_inner.nodes
)
# From the outer subgraph remove edges between outer nodes
edges_to_plot = (
subgraph_outer.edges
- nx.induced_subgraph(self._g, outer_nodes).edges
)
# Return the inner subgraph nodes and the outer subgraph edges
return (
list(subgraph_outer.nodes),
list(edges_to_plot),
nodes_to_shrink,
)
subgraph = nx.induced_subgraph(self._g, node_selection)
return list(subgraph.nodes), list(subgraph.edges), []
# Select the nodes and edges to be drawn. The selection is always based
# on the original graph self._g.
if nodelist is None and not kwargs:
# If no selection was specified, select all nodes and edges
nodes_to_draw, edges_to_draw, nodes_to_shrink = (
list(self._g.nodes),
list(self._g.edges),
[],
)
elif nodelist is not None:
# Select from list of nodes
nodes_to_draw, edges_to_draw, nodes_to_shrink = select_from_list(
nodelist=nodelist, **kwargs
)
else:
# Perform radial selection
nodes_to_draw, edges_to_draw, nodes_to_shrink = select_radial(
**kwargs
)
if drop:
# Remove the nodes that are not selected. This also removes all
# edges for which the source or destination was removed.
nodes_to_remove = self._g.nodes - set(nodes_to_draw)
if nodes_to_remove:
self._g.remove_nodes_from(nodes_to_remove)
self._nodes_to_draw = nodes_to_draw
self._edges_to_draw = edges_to_draw
self._nodes_to_shrink = nodes_to_shrink
