done logic constructing graph; validated with weaker testing of construction

mattport/nerf/graph.py

@@ -14,13 +14,24 @@ class GraphInputs:
     points: float
 
 
+@dataclass
+class Node:
+    """Node datastructure for graph composition."""
+
+    name: str
+    children: set
+
+    def __hash__(self):
+        return hash(self.name)
+
 class Graph(nn.Module):
     """_summary_"""
 
     def __init__(self, modules_config: dict) -> None:
         super().__init__()
+        self.modules_config = modules_config
+        # calculate input dimensions based on module dependencies
         self.modules = {}
-        self.module_order = self.get_module_order(modules_config)
         for module_name, module_dict in modules_config.items():
             module = getattr(importlib.import_module("mattport.nerf.modules"), module_dict["class_name"])
             if module_dict["class_name"] != "Encoding":
@@ -29,9 +40,39 @@ class Graph(nn.Module):
                     in_dim += modules_config[inputs]["meta_data"]["out_dim"]
                 module_dict["meta_data"]["in_dim"] = in_dim
             self.modules[module_name] = module(**module_dict["meta_data"])
+        # generate dependency ordering for module calls
+        self.module_order = self.get_module_order()
 
-    def get_module_order(self, modules_config: dict) -> list:
-        """Generates a graph to determine the order in which the modules should execute
+    def construct_graph(self) -> set:
+        """Constructs a dependency graph given the module configuration
+
+        Args:
+            modules_config (dict): module definitions that make up the network
+
+        Returns:
+            list: root nodes for the constructed dependency graph
+        """
+        processed_modules = {}
+        roots = set()
+        for module_name, module_dict in self.modules_config.items():
+            if not module_name in processed_modules:
+                curr_module = Node(name=module_name, children=set())
+                processed_modules[module_name] = curr_module
+            else:
+                curr_module = processed_modules[module_name]
+            inputs = module_dict["inputs"]
+            for input_module in inputs:
+                if not input_module in processed_modules:
+                    parent_module = Node(name=input_module, children=set([curr_module]))
+                    processed_modules[input_module] = parent_module
+                else:
+                    processed_modules[input_module].children.add(curr_module)
+                if input_module == "x":
+                    roots.add(curr_module)
+        return roots
+
+    def get_module_order(self) -> list:
+        """Generates a graph and determines order of module operations using topological sorting
 
         Args:
             modules_config (dict): module definitions that make up the network
@@ -39,7 +80,10 @@ class Graph(nn.Module):
         Returns:
             list: ordering of the module names that should be executed
         """
-        return modules_config.keys()
+
+        # call construct graph
+        # function to do topological sort to get ordering
+        return self.modules_config.keys()
 
     def forward(self, x):
         """_summary_
@@ -49,4 +93,6 @@ class Graph(nn.Module):
         """
         # for name in self.module_order:
         #     x = self.modules[name](x)
+
+        ## can index into previous results using dictionary
         raise NotImplementedError

tests/nerf/test_graph.py

@@ -4,12 +4,37 @@ from mattport.nerf.graph import Graph
 
 
 def create_basic_graph():
-    """Create basic test for creating a graph
+    """Create basic graph structure
+        a -> b -> c
 
     Returns:
         dict: module definitions for constructing a graph
     """
     basic_graph = {
+        "mlp_0": {
+            "class_name": "MLP",
+            "inputs": ["encoder_0"],
+            "meta_data": {"in_dim": 0, "out_dim": 1, "num_layers": 2, "layer_width": 3},
+        },
+        "encoder_0": {"class_name": "Encoding", "inputs": ["x"], "meta_data": {"out_dim": 8}},
+        "mlp_1": {
+            "class_name": "MLP",
+            "inputs": ["mlp_0"],
+            "meta_data": {"in_dim": 4, "out_dim": 5, "num_layers": 6, "layer_width": 7},
+        },
+    }
+    return basic_graph
+
+
+def create_repeat_graph():
+    """Create graph with repeating nodes
+        a -> b -> c
+        a ->
+
+    Returns:
+        dict: module definitions for constructing a graph
+    """
+    repeat_graph = {
         "encoder_0": {"class_name": "Encoding", "inputs": ["x"], "meta_data": {"out_dim": 8}},
         "mlp_0": {
             "class_name": "MLP",
@@ -22,22 +47,104 @@ def create_basic_graph():
             "meta_data": {"in_dim": 4, "out_dim": 5, "num_layers": 6, "layer_width": 7},
         },
     }
-    return basic_graph
+    return repeat_graph
+
+
+def create_complex_graph():
+    """Create graph with varied depths, splits and re-joins
+        a -> b -> d
+               -> e
+          -> c
+          ------>
+
+    Returns:
+        dict: module definitions for constructing a graph
+    """
+    complex_graph = {
+        "encoder_0": {"class_name": "Encoding", "inputs": ["x"], "meta_data": {"out_dim": 8}},
+        "mlp_0": {
+            "class_name": "MLP",
+            "inputs": ["encoder_0"],
+            "meta_data": {"in_dim": 0, "out_dim": 1, "num_layers": 2, "layer_width": 3},
+        },
+        "mlp_1": {
+            "class_name": "MLP",
+            "inputs": ["encoder_0"],
+            "meta_data": {"in_dim": 4, "out_dim": 5, "num_layers": 6, "layer_width": 7},
+        },
+        "mlp_2": {
+            "class_name": "MLP",
+            "inputs": ["encoder_0", "mlp_0"],
+            "meta_data": {"in_dim": 4, "out_dim": 5, "num_layers": 6, "layer_width": 7},
+        },
+        "mlp_3": {
+            "class_name": "MLP",
+            "inputs": ["mlp_0"],
+            "meta_data": {"in_dim": 4, "out_dim": 5, "num_layers": 6, "layer_width": 7},
+        },
+    }
+    return complex_graph
+
+
+def check_consistency(root: "Node", targets: list):
+    """rudimentary check on predicted graph versus expected target ordering"""
+    curr_pointer = root
+    for target_name, target_children in targets:
+        assert curr_pointer.name == target_name
+        assert len(curr_pointer.children) == target_children
+        if len(curr_pointer.children) > 0:
+            # for now traversing smaller parts of graph.. need stabler checks
+            curr_pointer = list(curr_pointer.children)[0]
 
 
 def test_graph_init():
     """test modules dictionary is instantiated properly"""
     test_graph = Graph(create_basic_graph())
     assert test_graph.modules is not None
-    assert type(test_graph.modules['encoder_0']).__name__ == 'Encoding'
-    assert type(test_graph.modules['mlp_0']).__name__ == 'MLP'
-    assert type(test_graph.modules['mlp_1']).__name__ == 'MLP'
+    assert type(test_graph.modules["encoder_0"]).__name__ == "Encoding"
+    assert type(test_graph.modules["mlp_0"]).__name__ == "MLP"
+    assert type(test_graph.modules["mlp_1"]).__name__ == "MLP"
 
 
 def test_input_dimension():
     """test module input dimension calculation"""
+    test_graph = Graph(create_repeat_graph())
+    assert not hasattr(test_graph.modules["encoder_0"], "in_dim")
+    assert test_graph.modules["mlp_0"].in_dim == 16
+    assert test_graph.modules["mlp_1"].in_dim == 1
+
+
+def test_construct_graph_basic():
+    """test construction of dependency graph: ensure correct ordering"""
     test_graph = Graph(create_basic_graph())
-    assert not hasattr(test_graph.modules['encoder_0'], 'in_dim')
-    assert test_graph.modules['mlp_0'].in_dim == 16
-    assert test_graph.modules['mlp_1'].in_dim == 1
-    
+    roots = test_graph.construct_graph()
+    assert len(roots) == 1
+    roots = list(roots)
+    targets = [("encoder_0", 1), ("mlp_0", 1), ("mlp_1", 0)]
+    check_consistency(roots[0], targets)
+
+
+def test_construct_graph_repeat():
+    """test construction of dependency graph: ensure duplicate nodes are ok"""
+    test_graph = Graph(create_repeat_graph())
+    roots = test_graph.construct_graph()
+    assert len(roots) == 1
+    roots = list(roots)
+    targets = [("encoder_0", 1), ("mlp_0", 1), ("mlp_1", 0)]
+    check_consistency(roots[0], targets)
+
+    test_graph.modules_config["mlp_0"]["inputs"][1] = "encoder_1"
+    test_graph.modules_config["encoder_1"] = {"class_name": "Encoding", "inputs": ["x"], "meta_data": {"out_dim": 8}}
+    print(test_graph.modules_config["mlp_0"]["inputs"][1])
+    roots = test_graph.construct_graph()
+    assert len(roots) == 2
+
+
+def test_construct_graph_complex():
+    """test construction of dependency graph: ensure different depths are ok"""
+    test_graph = Graph(create_complex_graph())
+    roots = test_graph.construct_graph()
+    assert len(roots) == 1
+    roots = list(roots)
+    targets = [("encoder_0", 3), ("mlp_0", 2), ("mlp_3", 0)]
+    check_consistency(roots[0], targets)