Neo4j

bayesian/bbn.py

@@ -4,6 +4,8 @@
 import copy
 import heapq
 
+import py2neo
+
 from random import random, choice
 from StringIO import StringIO
 from itertools import combinations, product
@@ -52,10 +54,22 @@ def __init__(self, nodes_dict, name=None, domains={}):
         for variable_name, node in nodes_dict.items():
             node.variable_name = variable_name
 
+    def export(self, filename=None, format='graphviz'):
+        '''Export the graph in GraphViz dot language.'''
+        if filename:
+            fh = open(filename, 'w')
+        else:
+            fh = sys.stdout
+        if format != 'graphviz':
+            raise 'Unsupported Export Format.'
+
+        fh.write(self.get_graphviz_source())
+
     def get_graphviz_source(self):
         fh = StringIO()
         fh.write('digraph G {\n')
-        fh.write('  graph [ dpi = 300 bgcolor="transparent" rankdir="LR"];\n')
+        # fh.write('  graph [ dpi = 300 bgcolor="transparent" rankdir="LR"];\n')
+        fh.write('  graph [ bgcolor="transparent" rankdir="LR"];\n')
         edges = set()
         for node in sorted(self.nodes, key=lambda x: x.name):
             fh.write('  %s [ shape="ellipse" color="blue"];\n' % node.name)
@@ -67,6 +81,30 @@ def get_graphviz_source(self):
         fh.write('}\n')
         return fh.getvalue()
 
+    def loadToNeo(self, connection=None):
+        if connection:
+            neoGraph = py2neo.Graph(connection)
+        else:
+            neoGraph = py2neo.Graph()
+
+        # WARNING WARNING
+        neoGraph.delete_all()
+        vertices = dict() # used to capture the Neo Nodes
+        edges = set()
+        for node in sorted(self.nodes, key=lambda x: x.name):
+            # vertices[node.name] = py2neo.Node("Node", name=node.name, variable_name=node.variable_name)
+            vertices[node.name] = py2neo.Node("Node", name=node.name)
+            # Add to the edge list
+            for child in node.children:
+                edge = (node.name, child.name)
+                edges.add(edge)
+        # Iterate through the edges
+        for source, target in sorted(edges, key=lambda x: (x[0], x[1])):
+            print('  %s -> %s;\n' % (source, target))
+            edge = py2neo.Relationship(vertices[source], "CHILD", vertices[target])
+            # edges.add(edge)
+            neoGraph.create(edge)
+
     def build_join_tree(self):
         jt = build_join_tree(self)
         return jt

bayesian/factor_graph.py

@@ -7,6 +7,9 @@
 import inspect
 import random
 
+# from py2neo import Graph
+import py2neo
+
 from collections import defaultdict
 from itertools import product as iter_product
 from Queue import Queue
@@ -1104,6 +1107,36 @@ def query_by_external_samples(self, **kwds):
             [(k, v / len(samples)) for k, v in counts.items()])
         return normalized
 
+    def loadToNeo(self, connection=None):
+        if connection:
+            neoGraph = py2neo.Graph(connection)
+        else:
+            neoGraph = py2neo.Graph()
+
+        # WARNING WARNING
+        neoGraph.delete_all()
+
+        # edges = set()
+        vertices = dict() # used to capture the Neo Nodes
+        for node in self.nodes:
+            if isinstance(node, FactorNode):
+                vertices[node.name] = py2neo.Node("FactorNode", name=node.name)
+                # fh.write('  %s [ shape="rectangle" color="red"];\n' % node.name)
+            else:
+                vertices[node.name] = py2neo.Node("VariableNode", name=node.name)
+                # fh.write('  %s [ shape="ellipse" color="blue"];\n' % node.name)
+        # Build the edges
+        for node in self.nodes:
+            for neighbour in node.neighbours:
+                # edge = [node.name, neighbour.name]
+                # edge = tuple(sorted(edge))
+
+                edge = py2neo.Relationship(vertices[node.name], "INFLUENCES", vertices[neighbour.name], since=1999)
+                # edges.add(edge)
+                neoGraph.create(edge)
+
+        # for source, target in edges:
+        #     fh.write('  %s -- %s;\n' % (source, target))
 
     def export(self, filename=None, format='graphviz'):
         '''Export the graph in GraphViz dot language.'''