import numpy as np
import networkx as nx
from .utils.misc import _insert_line_breaks
from .utils.fp import lmap
from functools import partial
from networkx import DiGraph
from pygraphviz import AGraph
import pickle
from datetime import datetime
import matplotlib
import matplotlib.cm as cm
from matplotlib.colors import Normalize
from delphi.utils.misc import choose_font
FONT = choose_font()
# ==========================================================================
# Export
# ==========================================================================
def _process_datetime(indicator_dict: Dict):
time = indicator_dict.get("time")
indicator_dict["time"] = str(time)
return indicator_dict
def export_edge(e):
return {
"source": e[0],
def to_agraph(
self,
indicators: bool = False,
indicator_values: bool = False,
nodes_to_highlight: Optional[Union[str, List[str]]] = None,
**kwargs,
):
""" Exports the CAG as a pygraphviz AGraph for visualization.
Args:
indicators: Whether to display indicators in the AGraph
indicator_values: Whether to display indicator values in the AGraph
nodes_to_highlight: Nodes to highlight in the AGraph.
Returns:
A PyGraphviz AGraph representation of the AnalysisGraph.
"""
from delphi.utils.misc import choose_font
FONT = choose_font()
A = nx.nx_agraph.to_agraph(self)
A.graph_attr.update({
"dpi": 227,
"fontsize": 20,
"rankdir": kwargs.get("rankdir", "TB"),
"fontname": FONT,
"overlap": "scale",
"splines": True,
})
A.node_attr.update({
"shape": "rectangle",
"color": "black",
# "color": "#650021",
"style": "rounded",
"fontname": FONT,
})
nodes_with_indicators = [
n for n in self.nodes(data=True)
if n[1].get("indicators") is not None
]
n_max = max([
sum([len(s.evidence) for s in e[2]["InfluenceStatements"]])
for e in self.edges(data=True)
])
nodeset = {n.split("/")[-1] for n in self.nodes}
simplified_labels = len(nodeset) == len(self)
color_str = "#650021"
for n in self.nodes(data=True):
if kwargs.get("values"):
node_label = (n[0].capitalize().replace("_", " ") + " (" +
str(np.mean(n[1]["rv"].dataset)) + ")")
else:
node_label = (n[0].split("/")[-1].replace(
"_", " ").capitalize() if simplified_labels else n[0])
A.add_node(n[0], label=node_label)
if list(self.edges(data=True))[0][2].get("βs") is not None:
max_median_betas = max(
[abs(np.median(e[2]["βs"])) for e in self.edges(data=True)])
for e in self.edges(data=True):
# Calculate reinforcement (ad-hoc!)
sts = e[2]["InfluenceStatements"]
total_evidence_pieces = sum([len(s.evidence) for s in sts])
reinforcement = (sum(
[stmt.overall_polarity() * len(stmt.evidence)
for stmt in sts]) / total_evidence_pieces)
opacity = total_evidence_pieces / n_max
h = (opacity * 255).hex()
if list(self.edges(data=True))[0][2].get("βs") is not None:
penwidth = 3 * abs(np.median(e[2]["βs"]) / max_median_betas)
cmap = cm.Greens if reinforcement > 0 else cm.Reds
c_str = (matplotlib.colors.rgb2hex(cmap(abs(reinforcement))) +
h[4:6])
else:
penwidth = 1
c_str = "black"
A.add_edge(e[0], e[1], color=c_str, penwidth=penwidth)
# Drawing indicator variables
if indicators:
for n in nodes_with_indicators:
for indicator_name, ind in n[1]["indicators"].items():
node_label = _insert_line_breaks(
ind.name.replace("_", " "), 30)
if indicator_values:
if ind.unit is not None:
units = f" {ind.unit}"
else:
units = ""
if ind.mean is not None:
ind_value = "{:.2f}".format(ind.mean)
node_label = (
f"{node_label}\n{ind_value} {ind.unit}"
f"\nSource: {ind.source}"
f"\nAggregation axes: {ind.aggaxes}"
f"\nAggregation method: {ind.aggregation_method}"
)
A.add_node(
indicator_name,
style="rounded, filled",
fillcolor="lightblue",
label=node_label,
)
A.add_edge(n[0], indicator_name, color="royalblue4")
nodes_to_highlight = kwargs.get("nodes_to_highlight")
if nodes_to_highlight is not None:
if isinstance(nodes_to_highlight, list):
for n in nodes_to_highlight:
if n in A.nodes():
A.add_node(n, fontcolor="royalblue")
elif isinstance(nodes_to_highlight, str):
if nodes_to_highlight in A.nodes():
A.add_node(nodes_to_highlight, fontcolor="royalblue")
else:
pass
if kwargs.get("graph_label") is not None:
A.graph_attr["label"] = kwargs["graph_label"]
return A