tests and display improvements

bamt/display/__init__.py

@@ -1,4 +1,4 @@
-from display import Display
+from .display import Display
 
 
 def plot_(output, *args):

bamt/display/display.py

@@ -140,10 +140,7 @@ def build(self, nodes, edges, **kwargs):
 
         network.hrepulsion(node_distance=300, central_gravity=0.5)
 
-        if not (os.path.exists("visualization_result")):
-            os.mkdir("visualization_result")
-
-        return network.show(f"visualization_result/" + self.output)
+        return network.show(self.output)
 
     @staticmethod
     def get_info(bn, as_df):

bamt/logging.conf

@@ -1,5 +1,5 @@
 [loggers]
-keys=root, preprocessor, builder, nodes, network, py.warnings
+keys=root, preprocessor, builder, nodes, network, display, py.warnings
 
 [handlers]
 keys=consoleHandler, fileHandler
@@ -35,6 +35,12 @@ qualname=nodes
 handlers=consoleHandler, fileHandler
 propagate=0
 
+[logger_display]
+level=INFO
+qualname=display
+handlers=consoleHandler, fileHandler
+propagate=0
+
 
 [logger_py.warnings]
 level=INFO

bamt/networks/base.py

@@ -766,7 +766,9 @@ def find_family(
         with_nodes: Optional[List] = None,
         plot_to: Optional[str] = None,
     ):
-        structure = GraphUtils.GraphAnalyzer(self).find_family(node_name, height, depth, with_nodes)
+        structure = GraphUtils.GraphAnalyzer(self).find_family(
+            node_name, height, depth, with_nodes
+        )
 
         if plot_to:
             plot_(

bamt/utils/GraphUtils.py

@@ -104,7 +104,7 @@ def markov_blanket(self, node_name: str):
         nodes = parents + children + fremd_eltern + [node_name]
 
         edges = self._isolate_structure(nodes)
-        return {"nodes": nodes, "edges": edges}
+        return {"nodes": list(set(nodes)), "edges": edges}
 
     def _collect_height(self, node_name, height):
         nodes = []
@@ -143,12 +143,11 @@ def find_family(self, *args):
         else:
             with_nodes = list(with_nodes)
         nodes = (
-                self._collect_depth(node_name, depth)
-                + self._collect_height(node_name, height)
-                + [node_name]
+            self._collect_depth(node_name, depth)
+            + self._collect_height(node_name, height)
+            + [node_name]
         )
 
         nodes = list(set(nodes + with_nodes))
 
-        return {"nodes": nodes,
-                "edges": self._isolate_structure(nodes + with_nodes)}
+        return {"nodes": nodes, "edges": self._isolate_structure(nodes + with_nodes)}

tests/test_builders.py

@@ -1,6 +1,3 @@
-from contextlib import redirect_stdout, redirect_stderr
-import os
-
 import itertools
 import unittest
 
@@ -9,7 +6,7 @@
 import pandas as pd
 
 from bamt.builders.builders_base import StructureBuilder, VerticesDefiner
-from bamt.builders.hc_builder import HCStructureBuilder, HillClimbDefiner
+from bamt.builders.hc_builder import HillClimbDefiner
 from bamt.builders.evo_builder import EvoStructureBuilder
 
 from bamt.nodes.gaussian_node import GaussianNode

tests/test_graph_analyzer.py

@@ -0,0 +1,90 @@
+import unittest
+import logging
+
+from bamt.utils import GraphUtils
+from bamt.networks.discrete_bn import DiscreteBN
+from bamt.nodes.discrete_node import DiscreteNode
+from bamt.builders.builders_base import VerticesDefiner
+from bamt.display import plot_
+
+logging.getLogger("builder").setLevel(logging.CRITICAL)
+
+
+class TestGraphAnalyzer(unittest.TestCase):
+    def setUp(self):
+        self.bn = DiscreteBN()
+        definer = VerticesDefiner(descriptor={"types": {
+            "Node0": "disc",
+            "Node1": "disc",
+            "Node2": "disc",
+            "Node3": "disc_num",
+            "Node4": "disc",
+            "Node5": "disc",
+            "Node6": "disc",
+            "Node7": "disc",
+            "Node8": "disc",
+            "Node9": "disc",
+        },
+            "signs": {},
+        },
+            regressor=None)
+
+        definer.skeleton["V"] = [
+            DiscreteNode(name=f"Node{i}") for i in range(10)
+        ]
+        definer.skeleton["E"] = [
+            ("Node0", "Node1"),
+            ("Node0", "Node2"),
+            ("Node2", "Node3"),
+            ("Node4", "Node7"),
+            ("Node1", "Node5"),
+            ("Node5", "Node6"),
+            ("Node7", "Node0"),
+            ("Node8", "Node1"),
+            ("Node9", "Node2"),
+        ]
+        definer.get_family()
+
+        self.bn.nodes = definer.skeleton["V"]
+        self.bn.edges = definer.skeleton["E"]
+
+        self.analyzer = GraphUtils.GraphAnalyzer(self.bn)
+
+    def test_markov_blanket(self):
+        result = self.analyzer.markov_blanket("Node0")
+        result["nodes"] = sorted(result["nodes"])
+        self.assertEqual(
+            {'edges': [('Node0', 'Node1'),
+                       ('Node0', 'Node2'),
+                       ('Node7', 'Node0'),
+                       ('Node8', 'Node1'),
+                       ('Node9', 'Node2')],
+             'nodes': sorted(['Node0', 'Node1', 'Node2', 'Node7', 'Node8', 'Node9'])}
+            ,
+            result
+        )
+
+    def test_find_family(self):
+        without_parents = self.analyzer.find_family("Node0", 0, 2, None)
+        without_parents["nodes"] = sorted(without_parents["nodes"])
+        self.assertEqual(
+            {'nodes': sorted(['Node3', 'Node2', 'Node1', 'Node0', 'Node5']),
+             'edges': [('Node0', 'Node1'), ('Node0', 'Node2'), ('Node2', 'Node3'), ('Node1', 'Node5')]},
+            without_parents
+        )
+
+        without_children = self.analyzer.find_family("Node0", 2, 0, None)
+        without_children["nodes"] = sorted(without_children["nodes"])
+        self.assertEqual(
+            {'nodes': sorted(['Node4', 'Node7', 'Node0']),
+             'edges': [('Node4', 'Node7'), ('Node7', 'Node0')]},
+            without_children
+        )
+
+        plot_("here.html",
+              [self.bn[node] for node in without_children["nodes"]],
+              without_children["edges"])
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)