-
Notifications
You must be signed in to change notification settings - Fork 79
/
graph.py
173 lines (155 loc) · 4.99 KB
/
graph.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
# -*- coding: utf-8 -*-
"""
Convenience functions for representing reaction systems as graphs.
"""
import os
import subprocess
import shutil
import tempfile
def rsys2dot(
rsys,
tex=False,
rprefix="r",
rref0=1,
nodeparams='[label="{}",shape=diamond]',
colors=("maroon", "darkgreen"),
penwidths=None,
include_inactive=True,
):
"""
Returns list of lines of DOT (graph description language)
formatted graph.
Parameters
==========
rsys: ReactionSystem
tex: bool (default False)
If set True, output will be LaTeX formatted
(Substance need to have latex_name attribute set)
rprefix: string
Reaction enumeration prefix, default: r
rref0: integer
Reaction enumeration initial counter value, default: 1
nodeparams: string
DOT formatted param list, default: [label={} shape=diamond]
Returns
=======
list of lines of DOT representation of the graph representation.
"""
lines = ['digraph "' + str(rsys.name) + '" {\n']
ind = " " # indentation
if penwidths is None:
penwidths = [1.0] * rsys.nr
categories = rsys.categorize_substances(checks=())
def add_substance(key):
fc = "black"
if key in categories["depleted"]:
fc = colors[0]
if key in categories["accumulated"]:
fc = colors[1]
label = ("$%s$" if tex else "%s") % getattr(
rsys.substances[key], "latex_name" if tex else "name"
)
lines.append(
ind
+ '"{key}" [fontcolor={fc} label="{lbl}"];\n'.format(
key=key, fc=fc, lbl=label
)
)
for sk in rsys.substances:
add_substance(sk)
def add_vertex(key, num, reac, penwidth):
snum = str(num) if num > 1 else ""
fmt = ",".join(
['label="{}"'.format(snum)]
+ (["penwidth={}".format(penwidth)] if penwidth != 1 else [])
)
lines.append(
ind
+ '"{}" -> "{}" [color={},fontcolor={},{}];\n'.format(
*(
(key, rid, colors[0], colors[0], fmt)
if reac
else (rid, key, colors[1], colors[1], fmt)
)
)
)
if include_inactive:
reac_stoichs = rsys.all_reac_stoichs()
prod_stoichs = rsys.all_prod_stoichs()
else:
reac_stoichs = rsys.active_reac_stoichs()
prod_stoichs = rsys.active_prod_stoichs()
for ri, rxn in enumerate(rsys.rxns):
rid = rprefix + str(ri + rref0)
lines.append(ind + "{")
lines.append(ind * 2 + "node " + nodeparams.format(rxn.name or rid))
lines.append(ind * 2 + rid)
lines.append(ind + "}\n")
for idx, key in enumerate(rsys.substances):
num = reac_stoichs[ri, idx]
if num == 0:
continue
add_vertex(key, num, True, penwidths[ri])
for idx, key in enumerate(rsys.substances):
num = prod_stoichs[ri, idx]
if num == 0:
continue
add_vertex(key, num, False, penwidths[ri])
lines.append("}\n")
return lines
def rsys2graph(rsys, fname, output_dir=None, prog=None, save=False, **kwargs):
"""
Convenience function to call `rsys2dot` and write output to file
and render the graph
Parameters
----------
rsys : ReactionSystem
fname : str
filename
output_dir : str (optional)
path to directory (default: temporary directory)
prog : str (optional)
default: 'dot'
save : bool
removes temporary directory if False, default: False
\\*\\*kwargs :
Keyword arguments passed along to py:func:`rsys2dot`.
Returns
-------
str
Outpath
Examples
--------
>>> rsys2graph(rsys, sbstncs, '/tmp/out.png') # doctest: +SKIP
"""
lines = rsys2dot(rsys, **kwargs)
created_tempdir = False
try:
if output_dir is None:
output_dir = tempfile.mkdtemp()
created_tempdir = True
basename, ext = os.path.splitext(os.path.basename(fname))
outpath = os.path.join(output_dir, fname)
dotpath = os.path.join(output_dir, basename + ".dot")
with open(dotpath, "wt") as ofh:
ofh.writelines(lines)
if ext == ".tex":
cmds = [prog or "dot2tex"]
else:
cmds = [prog or "dot", "-T" + outpath.split(".")[-1]]
p = subprocess.Popen(cmds + [dotpath, "-o", outpath])
retcode = p.wait()
if retcode:
fmtstr = "{}\n returned with exit status {}"
raise RuntimeError(fmtstr.format(" ".join(cmds), retcode))
return outpath
finally:
if save is True or save == "True":
pass
else:
if save is False or save == "False":
if created_tempdir:
shutil.rmtree(output_dir)
else:
# interpret save as path to copy pdf to.
shutil.copy(outpath, save)