haspirater

commit 5d7d1b541867c51d93ae53f2257b6b7570574ff2
parent 7e4b18bf9ab421d22b00b9e8f29db4817c74390a
Author: Antoine Amarilli <a3nm@a3nm.net>
Date:   Sun, 24 Feb 2013 18:58:51 +0100

make trie2dot.py more flexible

Diffstat:
README
 | 
6
+++---
trie2dot.py
 | 
17
++++++++++-------

2 files changed, 13 insertions(+), 10 deletions(-)
diff --git a/README b/README
@@ -107,7 +107,7 @@ confidence values. We also drop useless leaf nodes there.
 
 You can use trie2dot.py to convert the output of buildtrie.py or
 compresstrie.py in the dot format which can be used to render a drawing
-of the trie. The result of such a drawing is given as haspirater.pdf
-(before majoritytrie.py: contains frequency info, but more nodes)
-and haspirater_majority.pdf (no frequency, less nodes).
+of the trie ("trie2dot.py h 0 1"). The result of such a drawing is given
+as haspirater.pdf (before majoritytrie.py: contains frequency info, but
+more nodes) and haspirater_majority.pdf (no frequency, less nodes).
 
diff --git a/trie2dot.py b/trie2dot.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 
-"""Read json trie in stdin, trim unneeded branches and output json dump
-to stdout"""
+"""Takes json as input with labels [value1, value2] and produces dot,
+usage: trie2dot.py prefix value1 value2"""
 
 import json
 import sys
@@ -18,9 +18,9 @@ def cget(d, k):
     else:
       return 0
   except AttributeError:
-    # we have only one value, not a dictionary
+    # we have a list, not a dictionary
     # this happens after majoritytrie.py
-    if k == d:
+    if k in d:
       return 1
     else:
       return 0
@@ -50,8 +50,8 @@ def to_dot(trie, prefix=''):
   values, children = trie
   my_id = free_id
   free_id += 1
-  count = cget(values, "0") + cget(values, "1")
-  fraction = cget(values, "1") / count
+  count = cget(values, v1) + cget(values, v2)
+  fraction = cget(values, v2) / count
 
   print("%d [label=\"%s\",color=\"#%s\",penwidth=%d]" % (my_id, prefix,
     fraction2rgb(fraction), 1+int(log(count))))
@@ -64,5 +64,8 @@ def to_dot(trie, prefix=''):
   return my_id
 
 print("digraph G {\naspect=\"1\"\n")
-to_dot(trie, 'h')
+prefix = sys.argv[1]
+v1 = sys.argv[2]
+v2 = sys.argv[3]
+to_dot(trie, prefix)
 print("}")