[Yt-svn] yt: Adding pydot which is a GraphViz abstraction layer,
hg at spacepope.org
hg at spacepope.org
Tue Jan 18 11:30:00 PST 2011
hg Repository: yt
details: yt/rev/ede49708108b
changeset: 3675:ede49708108b
user: Stephen Skory <stephenskory at yahoo.com>
date:
Tue Jan 18 11:28:35 2011 -0800
description:
Adding pydot which is a GraphViz abstraction layer,
and modifying the merger trees to use it for output.
diffstat:
yt/analysis_modules/halo_merger_tree/enzofof_merger_tree.py | 79 +-
yt/analysis_modules/halo_merger_tree/merger_tree.py | 116 +-
yt/utilities/pydot.py | 1995 +++++++++++
3 files changed, 2078 insertions(+), 112 deletions(-)
diffs (truncated from 2279 to 300 lines):
diff -r 5da01bbe81ad -r ede49708108b yt/analysis_modules/halo_merger_tree/enzofof_merger_tree.py
--- a/yt/analysis_modules/halo_merger_tree/enzofof_merger_tree.py Mon Jan 17 15:49:23 2011 -0500
+++ b/yt/analysis_modules/halo_merger_tree/enzofof_merger_tree.py Tue Jan 18 11:28:35 2011 -0800
@@ -50,6 +50,7 @@
from yt.funcs import *
from yt.utilities.pykdtree import KDTree
+import yt.utilities.pydot as pydot
# We don't currently use this, but we may again find a use for it in the
# future.
@@ -338,23 +339,29 @@
print "--> Halo %8.8d :: fraction = %g" % (c[0], c[1])
def write_dot(self, filename=None):
- r"""Writes merger tree to a GraphViz file.
-
- User is responsible for creating an image file from it, e.g.
- dot -Tpng tree_halo00000.dot > image.png
+ r"""Writes merger tree to a GraphViz or image file.
Parameters
----------
filename : str, optional
Filename to write the GraphViz file. Default will be
- tree_halo%05i.dat.
+ tree_halo%05i.gv, which is a text file in the GraphViz format.
+ If filename is an image (e.g. "MergerTree.png") the output will
+ be in the appropriate image format made by calling GraphViz
+ automatically. See GraphViz (e.g. "dot -v")
+ for a list of available output formats.
"""
- if filename == None: filename = "tree_halo%5.5d.dot" % self.halonum
- fp = open(filename, "w")
- fp.write("digraph G {\n")
- fp.write(" node [shape=rect];\n")
+ if filename == None: filename = "tree_halo%5.5d.gv" % self.halonum
+ # Create the pydot graph object.
+ self.graph = pydot.Dot('galaxy', graph_type='digraph')
+ self.halo_shape = "rect"
+ self.z_shape = "plaintext"
+ # Subgraphs to align levels
+ self.subgs = {}
+ for num in self.numbers:
+ self.subgs[num] = pydot.Subgraph('', rank = 'same')
+ self.graph.add_subgraph(self.subgs[num])
sorted_lvl = sorted(self.levels, reverse=True)
- printed = []
for ii,lvl in enumerate(sorted_lvl):
# Since we get the cycle number from the key, it won't
# exist for the last level, i.e. children of last level.
@@ -366,10 +373,12 @@
for br in self.levels[lvl]:
for c in br.children:
color = "red" if c[0] == br.progenitor else "black"
- line = " C%d_H%d -> C%d_H%d [color=%s];\n" % \
- (lvl, br.halo_id, next_lvl, c[0], color)
+ self.graph.add_edge(pydot.Edge("C%d_H%d" %(lvl, br.halo_id),
+ "C%d_H%d" % (next_lvl, c[0]), color=color))
+ #line = " C%d_H%d -> C%d_H%d [color=%s];\n" % \
+ # (lvl, br.halo_id, next_lvl, c[0], color)
- fp.write(line)
+ #fp.write(line)
last_level = (ii,lvl)
for ii,lvl in enumerate(sorted_lvl):
npart_max = 0
@@ -378,35 +387,21 @@
for br in self.levels[lvl]:
halo_str = "C%d_H%d" % (lvl, br.halo_id)
style = "filled" if br.npart == npart_max else "solid"
- line = "%s [label=\"Halo %d\\n%d particles\",style=%s]\n" % \
- (halo_str, br.halo_id, br.npart, style)
- fp.write(line)
- # Last level, annotate children because they have no associated branches
- npart_max = 0
- for br in self.levels[last_level[1]]:
- for c in br.children:
- if c[2] > npart_max: npart_max = c[2]
- for br in self.levels[last_level[1]]:
- for c in br.children:
- lvl = self.numbers[0]
- style = "filled" if c[2] == npart_max else "solid"
- halo_str = "C%d_H%d" % (lvl, c[0])
- line = "%s [label=\"Halo %d\\n%d particles\",style=%s]\n" % \
- (halo_str, c[0], c[2], style)
- fp.write(line)
- # Output redshifts
- fp.write("\n")
- fp.write("node [shape=plaintext]\n")
- fp.write("edge [style=invis]\n")
- line = ""
- for k in sorted(self.redshifts, reverse=True):
- line = line + "\"z = %0.3f\"" % (self.redshifts[k]) + " -> "
- if k == self.numbers[0]: break
- line = line[:-4] # Remove last arrow
- fp.write("\n%s\n" % line)
-
- fp.write("}\n")
- fp.close()
+ self.graph.add_node(pydot.Node(halo_str,
+ label = "Halo %d\\n%d particles" % (br.halo_id, br.npart),
+ style = style, shape = self.halo_shape))
+ # Add this node to the correct level subgraph.
+ self.subgs[lvl].add_node(pydot.Node(halo_str))
+ for lvl in self.numbers:
+ # Don't add the z if there are no halos already in the subgraph.
+ if len(self.subgs[lvl].get_node_list()) == 0: continue
+ self.subgs[lvl].add_node(pydot.Node("%1.5e" % self.redshifts[lvl],
+ shape = self.z_shape, label = "z=%0.3f" % self.redshifts[lvl]))
+ # Based on the suffix of the file name, write out the result to a file.
+ suffix = filename.split(".")[-1]
+ if suffix == "gv": suffix = "raw"
+ mylog.info("Writing %s format %s to disk." % (suffix, filename))
+ self.graph.write("%s" % filename, format=suffix)
def find_halo_relationships(output1_id, output2_id, output_basename = None,
radius = 0.10):
diff -r 5da01bbe81ad -r ede49708108b yt/analysis_modules/halo_merger_tree/merger_tree.py
--- a/yt/analysis_modules/halo_merger_tree/merger_tree.py Mon Jan 17 15:49:23 2011 -0500
+++ b/yt/analysis_modules/halo_merger_tree/merger_tree.py Tue Jan 18 11:28:35 2011 -0800
@@ -36,6 +36,7 @@
HaloProfiler
from yt.convenience import load
from yt.utilities.logger import ytLogger as mylog
+import yt.utilities.pydot as pydot
try:
from yt.utilities.kdtree import *
except ImportError:
@@ -939,7 +940,13 @@
database : String
The name of the database file. Default = 'halos.db'.
dotfile : String
- The name of the file to write to. Default = 'MergerTree.gv'
+ The name of the file to write to. Default = 'MergerTree.gv'.
+ The suffix of this name gives the format of the output file,
+ so 'MergerTree.jpg' would output a jpg file. "dot -v" (from the
+ command line) will print
+ a list of image formats supported on the system. The default
+ suffix '.gv' will output the results to a text file in the Graphviz
+ markup language.
current_time : Integer
The SnapCurrentTimeIdentifier for the snapshot for the halos to
be tracked. This is identical to the CurrentTimeIdentifier in
@@ -969,15 +976,18 @@
if current_time is not None:
halos = self._translate_haloIDs(halos, current_time)
newhalos = set(halos)
- # A key is the GlobalHaloID for this halo, and the content is a
- # Node object.
- self.nodes = {}
- # A key is the GlobalHaloID for the parent in the relationship,
- # and the content is a Link ojbect.
- self.links = defaultdict(Link)
- # Record which halos are at the same z level for convenience.
- # They key is a z value, and the content a list of co-leveled halo IDs.
- self.levels = defaultdict(list)
+ # Create the pydot graph object.
+ self.graph = pydot.Dot('galaxy', graph_type='digraph')
+ # Build some initially empty subgraphs, which are used to identify
+ # nodes that are on the same rank (redshift).
+ line = "SELECT DISTINCT SnapZ FROM Halos;"
+ self.cursor.execute(line)
+ self.subgs = {}
+ result = self.cursor.fetchone()
+ while result:
+ self.subgs[result[0]] = pydot.Subgraph('', rank = 'same')
+ self.graph.add_subgraph(self.subgs[result[0]])
+ result = self.cursor.fetchone()
# For the first set of halos.
self._add_nodes(newhalos)
# Recurse over parents.
@@ -985,13 +995,7 @@
mylog.info("Finding parents for %d children." % len(newhalos))
newhalos = self._find_parents(newhalos)
self._add_nodes(newhalos)
- mylog.info("Writing out %s to disk." % dotfile)
- self._open_dot(dotfile)
- self._write_nodes()
- self._write_links()
- self._write_levels()
- self._close_dot()
- self._close_database()
+ self._write_dotfile(dotfile)
return None
def _translate_haloIDs(self, halos, current_time):
@@ -1032,9 +1036,10 @@
if pair[1] <= self.link_min or pair[0] != halo:
continue
else:
- self.nodes[halo].parentIDs.append(pID)
- self.links[pID].childIDs.append(halo)
- self.links[pID].fractions.append(pair[1])
+ self.graph.add_edge(pydot.Edge(pID, halo,
+ label = "%3.2f%%" % float(pair[1]*100),
+ color = "blue",
+ fontsize = "10"))
newhalos.add(pID)
result = self.cursor.fetchone()
return newhalos
@@ -1063,57 +1068,28 @@
value = (halo,)
self.cursor.execute(line, value)
result = self.cursor.fetchone()
- self.nodes[halo] = Node(na.array([result[2],result[3],result[4]]),
- result[1], [], result[0], 1. - float(result[5])/(maxID+1)) #+1 to prevent /0
- self.levels[result[0]].append(halo)
+ # Add the node to the pydot graph.
+ color_float = 1. - float(result[5])/(maxID+1)
+ self.graph.add_node(pydot.Node(halo,
+ label = "{%1.3e\\n(%1.3f,%1.3f,%1.3f)}" % \
+ (result[1], result[2], result[3], result[4]),
+ shape = "record",
+ color = "%0.3f 1. %0.3f" % (color_float, color_float)))
+ # Add this node to the correct subgraph.
+ self.subgs[result[0]].add_node(pydot.Node(halo))
+ # If this was the first node added to this subgraph, also add
+ # the lone node for the redshift value.
+ if len(self.subgs[result[0]].get_node_list()) == 1:
+ self.subgs[result[0]].add_node(pydot.Node("%1.5e" % result[0],
+ label = "%1.5f" % result[0],
+ shape = "record", color = "green"))
- def _open_dot(self, dotfile):
- # Write out the opening stuff in the dotfile.
- self.dotfile=self._write_on_root(dotfile)
- line = 'digraph galaxy {size="10, 10";\n'
- line += 'node [style=bold, shape=record];\n'
- self.dotfile.write(line)
-
- def _close_dot(self):
- self.dotfile.write("\n};\n")
- self.dotfile.close()
-
- def _write_nodes(self):
- # Write out the nodes to the dot file.
- self.dotfile.write("{\n")
- for halo in self.nodes:
- this = self.nodes[halo]
- line = '"%d" [label="{%1.3e\\n(%1.3f,%1.3f,%1.3f)}", shape="record",' \
- % (halo, this.mass, this.CoM[0], this.CoM[1], this.CoM[2])
- line += ' color="%0.3f 1. %0.3f"];\n' % (this.color, this.color)
- self.dotfile.write(line)
- self.dotfile.write("};\n")
-
- def _write_links(self):
- # Write out the links to the dot file.
- self.dotfile.write("{\n")
- for parent in self.links:
- this = self.links[parent]
- for child,frac in zip(this.childIDs, this.fractions):
- if frac > self.link_min:
- line = '"%d"->"%d" [label="%3.2f%%", color="blue", fontsize=10];\n' \
- % (parent, child, frac*100.)
- self.dotfile.write(line)
- self.dotfile.write("};\n")
-
- def _write_levels(self):
- # Write out the co-leveled halos to the dot file.
- for z in self.levels:
- this = self.levels[z]
- self.dotfile.write("{ rank = same;\n")
- line = '"%1.5f"; ' % z
- for halo in this:
- line += '"%d"; ' % halo
- line += "\n};\n"
- self.dotfile.write(line)
- # Also write out the unlinked boxes for the redshifts.
- line = '{"%1.5f" [label="{%1.5f}", shape="record" color="green"];}\n' \
- % (z, z)
+ def _write_dotfile(self, dotfile):
+ # Based on the suffix of the file name, write out the result to a file.
+ suffix = dotfile.split(".")[-1]
+ if suffix == "gv": suffix = "raw"
+ mylog.info("Writing %s format %s to disk." % (suffix, dotfile))
+ self.graph.write("%s" % dotfile, format=suffix)
class MergerTreeTextOutput(DatabaseFunctions, ParallelAnalysisInterface):
def __init__(self, database='halos.db', outfile='MergerTreeDB.txt'):
diff -r 5da01bbe81ad -r ede49708108b yt/utilities/pydot.py
--- /dev/null Thu Jan 01 00:00:00 1970 +0000
+++ b/yt/utilities/pydot.py Tue Jan 18 11:28:35 2011 -0800
@@ -0,0 +1,1995 @@
+# -*- coding: Latin-1 -*-
+"""Graphviz's dot language Python interface.
+
+This module provides with a full interface to create handle modify
+and process graphs in Graphviz's dot language.
+
+References:
+
+pydot Homepage: http://code.google.com/p/pydot/
+Graphviz: http://www.graphviz.org/
+DOT Language: http://www.graphviz.org/doc/info/lang.html
+
+Programmed and tested with Graphviz 2.26.3 and Python 2.6 on OSX 10.6.4
+by Ero Carrera (c) 2004-2010 [ero at dkbza.org]
+
+Distributed under MIT license [http://opensource.org/licenses/mit-license.html].
More information about the yt-svn
mailing list