"""
Represent states and transitions as graph.
"""
from __future__ import annotations
import re
from collections.abc import Generator, Sequence
from typing import TYPE_CHECKING, Any
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
from matplotlib import rcParams, rcParamsDefault
if TYPE_CHECKING:
import pandas as pd
from matplotlib.axes import Axes as mplAxes
__all__: list[str] = [
"construct_state_graphs",
"construct_transition_graph",
"check_graph_suitable",
"determine_node_order",
"plot_graph",
"draw_networkx_curved_edge_labels",
]
[docs]
def construct_state_graphs(transition_df: pd.DataFrame) -> list[nx.MultiDiGraph]:
"""
Constructs graphs of states (nodes) and their transitions (edges). Each fluorophore
or fluorophore combination gets a separate graph.
Parameters
----------
transition_df
Dataframe of all given transitions with non-zero rate containing their id as
second level index and their other attributes as columns. Name of fluorophores
as first level index.
Returns
-------
graphs : list[nx.MultiDiGraph]
Contains objects of type nx.MultiDiGraph.
"""
graphs = []
grouped = transition_df.groupby(level=0)
for fluorophore, f_transitions in grouped:
G = nx.MultiDiGraph()
edges = []
for (_, _), transition in f_transitions.iterrows():
abbr = transition["abbreviation"]
if "dist" not in fluorophore:
source = transition["initial_state"].name
destination = transition["final_state"].name
edge = (
fluorophore + "_" + source,
fluorophore + "_" + destination,
{"w": abbr, "dist": ""},
)
edges.append(edge)
else:
pattern = re.compile(r"D: (\w+), A: (\w+), dist: ([\d.]+)")
d, a, dist = pattern.findall(fluorophore)[0]
source_1 = transition["initial_state"].value[0].name
source_2 = transition["initial_state"].value[1].name
edge = (
d + "_" + source_1,
a + "_" + source_2 + "(2)",
{"w": abbr, "dist": f"distance: {dist} nm"},
)
edges.append(edge)
G.add_edges_from(edges)
graphs.append(G)
return graphs
[docs]
def construct_transition_graph(transition_df: pd.DataFrame) -> nx.MultiDiGraph:
"""
Constructs a graph of transitions (nodes) and their involved states (edges).
Parameters
----------
transition_df
Dataframe of all given transitions with non-zero rate containing their id as
second level index and their other attributes as columns. Name of fluorophores
as first level index.
Returns
-------
G : nx.MultiDiGraph
Markov chain representation by nodes and edges.
"""
if transition_df.index.get_level_values(0).nunique() > 1:
raise ValueError(
"construct_transition_graph only available for single "
"fluorophore systems."
)
G = nx.MultiDiGraph()
edges = []
for (_, id_source), row in transition_df.iterrows():
final_state = row["final_state"]
for (_, id_destination), row in transition_df.iterrows():
if row["initial_state"] == final_state:
source = id_source
destination = id_destination
edge = (source, destination, {"w": f"{final_state.name}"})
edges.append(edge)
G.add_edges_from(edges)
return G
[docs]
def check_graph_suitable(
G: nx.MultiDiGraph, starting_node: int
) -> tuple[bool, list[Any]]:
"""
Checks whether a Markov chain is suitable for an approximation of its development
in time. This means being acyclic (except cycles that include the starting node).
Parameters
----------
G
Markov Chain representation by nodes and edges.
starting_node
Numeric value representing the starting node (i.e., state).
Returns
-------
graph_suited : bool
Whether the graph is suited for the algorithms.
cycles : list
Contains each simple cycle of G.
"""
# check for reversible reactions and loops that do not contain the starting node:
graph_suited = True
cycles = list(nx.simple_cycles(G))
for cycle in cycles:
if starting_node not in cycle:
graph_suited = False
return graph_suited, cycles
[docs]
def determine_node_order(
G: nx.MultiDiGraph, starting_node: int
) -> Generator[Any, None, None]:
"""
Determine the order of nodes of a graph such that each node that leads to another
node has been visited before the other node. Requires the graph to be a DAG
(directed acyclic graph). If the starting node is part of each cycle, it can be
removed to convert the graph to DAG.
Parameters
----------
G
Markov Chain representation by nodes and edges.
starting_node
Numeric value representing the starting node (i.e., state).
Returns
-------
node_order : generator
Yields the topological sort of the graph.
"""
G_mutated = G.copy()
edges_to_remove = []
for edge in G.edges:
if edge[1] == starting_node:
edges_to_remove.append(edge)
G_mutated.remove_edges_from(edges_to_remove)
node_order = nx.topological_sort(G_mutated)
return node_order
[docs]
def plot_graph(
G: nx.MultiDiGraph,
graph_type: str = "shell",
colors: Sequence[str] = None,
scale: float = 1,
ax: mplAxes | None = None,
) -> mplAxes:
"""
Plot graph.
Adapted from https://stackoverflow.com/questions/22785849/drawing-multiple-edges-
between-two-nodes-with-networkx.
Parameters
----------
G
Markov Chain representation by nodes and edges.
graph_type
Specifies network layout. One of 'shell', 'circular', 'planar' or 'kamada'.
colors
Contains two colors as Hex values of type str.
scale
Factor to scale the figure.
ax : mpl.Axes
The axes on which to show the image
Returns
-------
mpl.Axes
Axes object with the plot.
"""
if ax is None:
ax = plt.gca()
if colors is None:
colors = ["#ADD8E6", "#FFF0C8"]
rcParams["figure.dpi"] = rcParamsDefault["figure.dpi"] * scale
if graph_type == "circular":
pos = nx.circular_layout(G)
elif graph_type == "planar":
pos = nx.planar_layout(G)
elif graph_type == "shell":
pos = nx.shell_layout(G)
else:
pos = nx.kamada_kawai_layout(G)
labels = {}
colormap = []
for _, node in enumerate(G):
if isinstance(node, str) and "(2)" in node:
colormap.append(colors[1])
labels[node] = node.replace("(2)", "")
else:
colormap.append(colors[0])
labels[node] = node
nx.draw_networkx_nodes(G=G, pos=pos, ax=ax, node_color=colormap)
nx.draw_networkx_labels(G=G, pos=pos, ax=ax, labels=labels)
edge_weights = nx.get_edge_attributes(G, name="w")
straight_edges = []
arc_rad = 0
arc_rad_reversed = 0
for i, new_edge in enumerate(G.edges):
if i == 0:
distance = nx.get_edge_attributes(G, name="dist")[new_edge]
ax.set_title(distance)
nothing_found = True
for old_edge in straight_edges:
if new_edge[:2] == old_edge[:2]:
arc_rad += 0.25
nothing_found = False
nx.draw_networkx_edges(
G=G,
pos=pos,
ax=ax,
edgelist=[new_edge],
connectionstyle=f"arc3, rad = {arc_rad}",
)
draw_networkx_curved_edge_labels(
G=G,
pos=pos,
ax=ax,
edge_labels={new_edge: edge_weights[new_edge]},
rad=arc_rad,
)
break
elif list(reversed(new_edge[:2])) == list(old_edge[:2]):
arc_rad_reversed += 0.25
nothing_found = False
nx.draw_networkx_edges(
G=G,
pos=pos,
ax=ax,
edgelist=[new_edge],
connectionstyle=f"arc3, rad = {arc_rad_reversed}",
)
draw_networkx_curved_edge_labels(
G=G,
pos=pos,
ax=ax,
edge_labels={new_edge: edge_weights[new_edge]},
rad=arc_rad_reversed,
)
break
if nothing_found:
arc_rad = 0
arc_rad_reversed = 0
straight_edges.append(new_edge)
nx.draw_networkx_edges(G=G, pos=pos, ax=ax, edgelist=straight_edges)
straight_edge_labels = {edge: edge_weights[edge] for edge in straight_edges}
draw_networkx_curved_edge_labels(
G=G, pos=pos, ax=ax, edge_labels=straight_edge_labels, rad=0
)
return ax
[docs]
def draw_networkx_curved_edge_labels(
G: nx.Graph,
pos: dict[Any, Any],
ax: mplAxes | None = None,
edge_labels: dict[Any, Any] = None,
rad: float = 0,
) -> mplAxes:
"""
Draws labels to curved edges.
Adapted from https://stackoverflow.com/questions/22785849/drawing-multiple-edges-
between-two-nodes-with-networkx.
Parameters
----------
G
A networkx graph.
pos
Nodes as keys and positions as values.
ax
Axis to plot on.
edge_labels
Edges (tuples) as keys and labels as values.
rad
Rounding radius of curved edge.
Returns
-------
ax : mpl.Axes
"""
if ax is None:
ax = plt.gca()
if edge_labels is None:
labels = {(u, v): d for u, v, d in G.edges(data=True)}
else:
labels = edge_labels
text_items = {}
for (n1, n2, _), label in labels.items():
pos_1 = ax.transData.transform(np.array(pos[n1]))
pos_2 = ax.transData.transform(np.array(pos[n2]))
linear_mid = 0.5 * pos_1 + 0.5 * pos_2
d_pos = pos_2 - pos_1
rotation_matrix = np.array([(0, 1), (-1, 0)])
ctrl_1 = linear_mid + rad * rotation_matrix @ d_pos
ctrl_mid_1 = 0.5 * pos_1 + 0.5 * ctrl_1
ctrl_mid_2 = 0.5 * pos_2 + 0.5 * ctrl_1
bezier_mid = 0.5 * ctrl_mid_1 + 0.5 * ctrl_mid_2
x, y = ax.transData.inverted().transform(bezier_mid)
trans_angle = 0.0
# use default box of white with white border
bbox = dict(boxstyle="round", ec=(1.0, 1.0, 1.0), fc=(1.0, 1.0, 1.0))
if not isinstance(label, str):
label = str(label) # this makes "1" and 1 labeled the same
t = ax.text(
x,
y,
label,
size=10,
color="k",
family="sans-serif",
weight="normal",
alpha=None,
horizontalalignment="center",
verticalalignment="center",
rotation=trans_angle,
transform=ax.transData,
bbox=bbox,
zorder=1,
clip_on=True,
)
text_items[(n1, n2)] = t
ax.tick_params(
axis="both",
which="both",
bottom=False,
left=False,
labelbottom=False,
labelleft=False,
)
return ax