9552846961

Committed 17 Jun 2024 06:28PM UTC coverage: 93.341% (+0.02%) from 93.318%

Build # 9552846961

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #359 from pyiron/plot

Add the option to plot the task graph

Run Details

61 of 65 new or added lines in 4 files covered. (93.85%)

855 of 916 relevant lines covered (93.34%)

0.93 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

97.44

/pympipool/shared/plot.py

from concurrent.futures import Future
from typing import Tuple

import cloudpickle


def generate_nodes_and_edges(
    task_hash_dict: dict, future_hash_inverse_dict: dict
) -> Tuple[list]:
    node_lst, edge_lst = [], []
    hash_id_dict = {}
    for k, v in task_hash_dict.items():
        hash_id_dict[k] = len(node_lst)
        node_lst.append({"name": v["fn"].__name__, "id": hash_id_dict[k]})
    for k, task_dict in task_hash_dict.items():
        for arg in task_dict["args"]:
            if not isinstance(arg, Future):
                node_id = len(node_lst)
                node_lst.append({"name": str(arg), "id": node_id})
                edge_lst.append({"start": node_id, "end": hash_id_dict[k], "label": ""})
            else:
                edge_lst.append(
                    {
                        "start": hash_id_dict[future_hash_inverse_dict[arg]],
                        "end": hash_id_dict[k],
                        "label": "",
                    }
                )
        for kw, v in task_dict["kwargs"].items():
            if not isinstance(v, Future):
                node_id = len(node_lst)
                node_lst.append({"name": str(v), "id": node_id})
                edge_lst.append(
                    {"start": node_id, "end": hash_id_dict[k], "label": str(kw)}
                )
            else:
                edge_lst.append(
                    {
                        "start": hash_id_dict[future_hash_inverse_dict[v]],
                        "end": hash_id_dict[k],
                        "label": str(kw),
                    }
                )
    return node_lst, edge_lst


def generate_task_hash(task_dict: dict, future_hash_inverse_dict: dict) -> bytes:
    args_for_hash = [
        arg if not isinstance(arg, Future) else future_hash_inverse_dict[arg]
        for arg in task_dict["args"]
    ]
    kwargs_for_hash = {
        k: v if not isinstance(v, Future) else future_hash_inverse_dict[v]
        for k, v in task_dict["kwargs"].items()
    }
    return cloudpickle.dumps(
        {"fn": task_dict["fn"], "args": args_for_hash, "kwargs": kwargs_for_hash}
    )


def draw(node_lst: list, edge_lst: list):
    from IPython.display import SVG, display  # noqa
    import matplotlib.pyplot as plt  # noqa
    import networkx as nx  # noqa

    graph = nx.DiGraph()
    for node in node_lst:
        graph.add_node(node["id"], label=node["name"])
    for edge in edge_lst:
        graph.add_edge(edge["start"], edge["end"], label=edge["label"])
    svg = nx.nx_agraph.to_agraph(graph).draw(prog="dot", format="svg")
    display(SVG(svg))
    plt.show()

1	from concurrent.futures import Future	1✔
2	from typing import Tuple	1✔
3
4	import cloudpickle	1✔
5
6
7	def generate_nodes_and_edges(	1✔
8	task_hash_dict: dict, future_hash_inverse_dict: dict
9	) -> Tuple[list]:
10	node_lst, edge_lst = [], []	1✔
11	hash_id_dict = {}	1✔
12	for k, v in task_hash_dict.items():	1✔
13	hash_id_dict[k] = len(node_lst)	1✔
14	node_lst.append({"name": v["fn"].__name__, "id": hash_id_dict[k]})	1✔
15	for k, task_dict in task_hash_dict.items():	1✔
16	for arg in task_dict["args"]:	1✔
17	if not isinstance(arg, Future):	1✔
18	node_id = len(node_lst)	1✔
19	node_lst.append({"name": str(arg), "id": node_id})	1✔
20	edge_lst.append({"start": node_id, "end": hash_id_dict[k], "label": ""})	1✔
21	else:
NEW 22	edge_lst.append(	×
23	{
24	"start": hash_id_dict[future_hash_inverse_dict[arg]],
25	"end": hash_id_dict[k],
26	"label": "",
27	}
28	)
29	for kw, v in task_dict["kwargs"].items():	1✔
30	if not isinstance(v, Future):	1✔
31	node_id = len(node_lst)	1✔
32	node_lst.append({"name": str(v), "id": node_id})	1✔
33	edge_lst.append(	1✔
34	{"start": node_id, "end": hash_id_dict[k], "label": str(kw)}
35	)
36	else:
37	edge_lst.append(	1✔
38	{
39	"start": hash_id_dict[future_hash_inverse_dict[v]],
40	"end": hash_id_dict[k],
41	"label": str(kw),
42	}
43	)
44	return node_lst, edge_lst	1✔
45
46
47	def generate_task_hash(task_dict: dict, future_hash_inverse_dict: dict) -> bytes:	1✔
48	args_for_hash = [	1✔
49	arg if not isinstance(arg, Future) else future_hash_inverse_dict[arg]
50	for arg in task_dict["args"]
51	]
52	kwargs_for_hash = {	1✔
53	k: v if not isinstance(v, Future) else future_hash_inverse_dict[v]
54	for k, v in task_dict["kwargs"].items()
55	}
56	return cloudpickle.dumps(	1✔
57	{"fn": task_dict["fn"], "args": args_for_hash, "kwargs": kwargs_for_hash}
58	)
59
60
61	def draw(node_lst: list, edge_lst: list):	1✔
62	from IPython.display import SVG, display # noqa	1✔
63	import matplotlib.pyplot as plt # noqa	1✔
64	import networkx as nx # noqa	1✔
65
66	graph = nx.DiGraph()	1✔
67	for node in node_lst:	1✔
68	graph.add_node(node["id"], label=node["name"])	1✔
69	for edge in edge_lst:	1✔
70	graph.add_edge(edge["start"], edge["end"], label=edge["label"])	1✔
71	svg = nx.nx_agraph.to_agraph(graph).draw(prog="dot", format="svg")	1✔
72	display(SVG(svg))	1✔
73	plt.show()	1✔

pyiron / pympipool / 9552846961

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous