File contents

#
# Visualizes Bayes Net's CPT
# Copyright: Willi Richert, 2007
#
# Requires: pydot, graphviz
# 

from OpenBayes import BNet, BVertex, DirEdge, JoinTree, MCMCEngine


import pydot

def bnToImg(G):
    node_dict = {}
    dot = pydot.Dot()

    def getProbabilityString(n):
        ps = ""
        num = len(n.distribution.parents)
        
        if num==0: # no parents
            return "| P("+n.name+") = %.3f"%float(n.distribution[0])

        for i in xrange(0, 2**num): # bit mask denoting True/False for the parents
            ps += "| P("+n.name+" / "
            
            VARS = []
            d = {}
            for j in xrange(0, num):
                print 2**j, i, 2**j & i == 2**j
                
                if 2**j & i == 2**j:                    
                    VARS.append(n.distribution.names_list[j+1]) # +1, weil erstes n.name ist
                else:
                    VARS.append("-"+n.distribution.names_list[j+1]) # +1, weil erstes n.name ist
                d[n.distribution.names_list[j+1]] = 2**j & i == 2**j
            #ps += ",".join(VARS)+") = %.3f"%float(ie.Marginalise(n.name).cpt[0])+" / %.3f"%+float(ie.Marginalise(n.name).cpt[1])

            ps += ",".join(VARS)+") = %.3f"%float(n.distribution[d][0])
            
        print "ps="+ps
        return ps
    
    def getPydotNode(bn_node):
        if not bn_node.name in node_dict:
            print "name=",bn_node.name
            dot.add_node(pydot.Node(bn_node.name))
            pydot_node = dot.get_node(bn_node.name)
            probs = getProbabilityString(bn_node)
            pydot_node.name ="{ %s %s }"%(pydot_node.name, probs)
            pydot_node.color = "blue"
            pydot_node.shape = "record"
            print "pydot node:",pydot_node.name
            node_dict[bn_node.name] = pydot_node
        else:
            pydot_node = node_dict[bn_node.name]
        return pydot_node
                
    edges = []

   
    for edge_key, edge in G.e.iteritems():
        for s in edge.src_v:
            sn = getPydotNode(s)
                
            for t in edge.dest_v:
                tn = getPydotNode(t)
                dot.add_edge(pydot.Edge(sn.name, tn.name))
                
                
    dot.splines=True
    dot.set_color("blue")
    dot.shape="record"
    dot.write_dot('bayes_net.dot')
    dot.write_ps('bayes_net.ps') 
    print
    print dot.to_string()

if __name__=="__main__":
    G = BNet("Water")
    c, s, r, w = [G.add_v(BVertex(nm, True, 2)) for nm in 'c s r w'.split()]
    for ep in [(c, r), (c, s), (r, w), (s, w)]:
        G.add_e(DirEdge(len(G.e), *ep))


    G.InitDistributions()
    c.setDistributionParameters([.5, .5])
    s.setDistributionParameters([.5, .9, .5, .1])
    r.setDistributionParameters([.8, .2, .2, .8])

    w.distribution[{'s':0, 'r':0}] = [0.99, 0.01]
    w.distribution[{'s':0, 'r':1}] = [0.1, 0.9]
    w.distribution[{'s':1, 'r':0}] = [0.1, 0.9]
    w.distribution[{'s':1, 'r':1}] = [0.0, 1.0]

    bnToImg(G)