CBDDraw.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216
  1. from IPython.display import Image, display
  2. import os
  3. import subprocess
  4. try:
  5. import pydot
  6. has_pydot = True
  7. except ImportError:
  8. print("pydot not installed")
  9. has_pydot = False
  10. from cbd.CBDBlocks import *
  11. def draw_to_file(name, cbd, show=False):
  12. if not has_pydot:
  13. print("No pydot")
  14. return
  15. if cbd is None:
  16. print("graph_to_dot Error: Empty CBD")
  17. return
  18. vattr = ''
  19. nodes = {}
  20. graph = pydot.Dot(name, graph_type='digraph')
  21. for i, block in enumerate(cbd.getBlocks()):
  22. block_type = block.getBlockType()
  23. vattr = block.getBlockName()#block_type + str(i)
  24. fillcolor = get_fill_color(block_type)
  25. if "Constant" in block_type:
  26. try:
  27. vattr += "\\n Value= " + str(str(block.getValue()))
  28. except KeyError:
  29. pass
  30. elif "Gain" in block_type:
  31. try:
  32. vattr += "\\n Value= " + str(str(block.gain))
  33. except KeyError:
  34. pass
  35. elif "Test" in block_type:
  36. try:
  37. vattr += "\\n Test= " + str(block.expr) + " "+ str(block.switch_value)
  38. except KeyError:
  39. pass
  40. nodes[block.getBlockName()] = pydot.Node(vattr, style="filled", fillcolor=fillcolor)
  41. graph.add_node(nodes[block.getBlockName()])
  42. link_colours = {"paired_with": "#505050", "link_in": "#6E2229", "init_val": "#938344",
  43. "IC": "darkgoldenrod", "reset": "#B9512B"}
  44. for block in cbd.getBlocks():
  45. for other in block.linksIN:
  46. colour = link_colours["link_in"]
  47. graph.add_edge(pydot.Edge(nodes[other.getBlockName()], nodes[block.getBlockName()], color=colour))
  48. try:
  49. other = block.IC
  50. colour = link_colours["IC"]
  51. graph.add_edge(pydot.Edge(nodes[other.getBlockName()], nodes[block.getBlockName()], color=colour, label="IC"))
  52. except AttributeError:
  53. pass
  54. try:
  55. other = block.reset
  56. colour = link_colours["reset"]
  57. graph.add_edge(pydot.Edge(nodes[other.getBlockName()], nodes[block.getBlockName()], color=colour, label="reset"))
  58. except AttributeError:
  59. pass
  60. try:
  61. other = block.init_val
  62. colour = link_colours["init_val"]
  63. graph.add_edge(pydot.Edge(nodes[other.getBlockName()], nodes[block.getBlockName()], color=colour, label="init_val"))
  64. except AttributeError:
  65. pass
  66. try:
  67. other = block.TRUE_in
  68. colour = "#000000"
  69. graph.add_edge(pydot.Edge(nodes[other.getBlockName()], nodes[block.getBlockName()], color=colour, label="TRUE"))
  70. except AttributeError:
  71. pass
  72. try:
  73. other = block.FALSE_in
  74. colour = "#000000"
  75. graph.add_edge(pydot.Edge(nodes[other.getBlockName()], nodes[block.getBlockName()], color=colour, label="FALSE"))
  76. except AttributeError:
  77. pass
  78. #
  79. # for other in in_blocks:
  80. # name = other.getBlockName()
  81. # label = ""
  82. #
  83. # if not name.startswith("IN"):
  84. # label=name
  85. #
  86. # # if not name.startswith("OUT"):
  87. # # label = label + " / " + other.getBlockName()
  88. #
  89. # write("{a} -> {b} [label=\"{lbl}\"];\n".format(a=other.getBlockName(),
  90. # b=block.getBlockName(),
  91. #
  92. # for e in g.es:
  93. #
  94. # src, trgt = e.tuple
  95. # src = int(src)
  96. # trgt = int(trgt)
  97. #
  98. # if src in internal_links.keys():
  99. # internal_links[src]["target"] = trgt
  100. # continue
  101. # elif trgt in internal_links.keys():
  102. # internal_links[trgt]["source"] = src
  103. # continue
  104. #
  105. # else:
  106. #
  107. # color = "black"
  108. # if "MatchModel" in mms[src]:
  109. # color = link_colours["match_contains"]
  110. # elif "ApplyModel" in mms[src]:
  111. # color = link_colours["apply_contains"]
  112. #
  113. #
  114. png_filename = name + '.png'
  115. dot_filename = png_filename.replace(".png", ".dot")
  116. graph.write(dot_filename, prog='dot')
  117. command = "dot -Tpng " + dot_filename + " -o " + png_filename
  118. subprocess.call(command, shell=True)
  119. command = "rm " + dot_filename
  120. subprocess.call(command, shell=True)
  121. if show:
  122. im = Image(png_filename)
  123. display(im)
  124. def get_fill_color(node_type):
  125. nt = node_type
  126. fillcolor = "#9999FF"
  127. if nt == 'MatchModel':
  128. fillcolor = "#E15C34"
  129. elif "Gain" in nt:
  130. fillcolor = "#22DDFF"
  131. elif "Integrator" in nt:
  132. fillcolor = "lightgoldenrod"
  133. elif "Test" in nt:
  134. fillcolor = "coral"
  135. elif nt == 'ApplyModel':
  136. fillcolor = "#FED017"
  137. elif nt == 'Equation':
  138. fillcolor = "#66FF33"
  139. elif nt in ['leftExpr', "rightExpr"]:
  140. fillcolor = "#22DDFF"
  141. elif nt in ['hasAttr_S', 'MT_pre__hasAttr_S', 'MT_post__hasAttr_S', "hasAttribute_S"]:
  142. fillcolor = "#E34A0E"
  143. elif node_type == 'hasAttr_T' or "hasAttribute" in nt:
  144. fillcolor = "#C75C0D"
  145. elif nt == 'Attribute':
  146. fillcolor = "#FFCC00"
  147. elif "hasArgs" in nt:
  148. fillcolor = "#FFFF99"
  149. elif "Concat" in nt:
  150. fillcolor = "#9999FF"
  151. elif "Constant" in nt:
  152. fillcolor = "lightblue"
  153. elif nt in ['directLink_S', 'directLink_T']:
  154. fillcolor = "lightyellow"
  155. elif nt in ['indirectLink_S', 'indirectLink_T']:
  156. fillcolor = "lightgreen"
  157. elif node_type == 'backward_link':
  158. fillcolor = "coral"
  159. elif node_type == 'trace_link':
  160. fillcolor = "lightgoldenrod"
  161. return fillcolor