Add if, elseif, and else.

graphkit/__init__.py

@@ -9,3 +9,4 @@
 # For backwards compatibility
 from .base import Operation
 from .network import Network
+from .control import If, ElseIf, Else

graphkit/base.py

@@ -54,6 +54,7 @@ def __init__(self, **kwargs):
         self.provides = kwargs.get('provides')
         self.params = kwargs.get('params', {})
         self.color = kwargs.get('color', None)
+        self.order = 0
 
         # call _after_init as final step of initialization
         self._after_init()
@@ -165,3 +166,23 @@ def __getstate__(self):
         state = Operation.__getstate__(self)
         state['net'] = self.__dict__['net']
         return state
+
+
+class Control(Operation):
+
+    def __init__(self, **kwargs):
+        super(Control, self).__init__(**kwargs)
+
+    def __repr__(self):
+        """
+        Display more informative names for the Operation class
+        """
+        if hasattr(self, 'condition_needs'):
+            return u"%s(name='%s', needs=%s, provides=%s, condition_needs=%s)" % \
+                (self.__class__.__name__,
+                 self.name,
+                 self.needs,
+                 self.provides,
+                 self.condition_needs)
+        else:
+            return super(Control, self).__repr__()

graphkit/control.py

@@ -0,0 +1,47 @@
+"""
+This sub-module contains statements that can be used for conditional evaluation of the graph.
+"""
+
+from .base import Control
+from .functional import compose
+
+
+class If(Control):
+
+    def __init__(self, condition_needs, condition, **kwargs):
+        super(If, self).__init__(**kwargs)
+        self.condition_needs = condition_needs
+        self.condition = condition
+        self.order = 1
+
+    def __call__(self, *args):
+        self.graph = compose(name=self.name)(*args)
+        return self
+
+    def _compute_condition(self, named_inputs):
+        inputs = [named_inputs[d] for d in self.condition_needs]
+        return self.condition(*inputs)
+
+    def _compute(self, named_inputs):
+        return self.graph(named_inputs)
+
+
+class ElseIf(If):
+
+    def __init__(self, condition_needs, condition, **kwargs):
+        super(ElseIf, self).__init__(condition_needs, condition, **kwargs)
+        self.order = 2
+
+
+class Else(Control):
+
+    def __init__(self, **kwargs):
+        super(Else, self).__init__(**kwargs)
+        self.order = 3
+
+    def __call__(self, *args):
+        self.graph = compose(name=self.name)(*args)
+        return self
+
+    def _compute(self, named_inputs):
+        return self.graph(named_inputs)

graphkit/network.py

@@ -7,7 +7,7 @@
 
 from io import StringIO
 
-from .base import Operation
+from .base import Operation, Control
 
 
 class DataPlaceholderNode(str):
@@ -83,6 +83,10 @@ def add_op(self, operation):
         for p in operation.provides:
             self.graph.add_edge(operation, DataPlaceholderNode(p))
 
+        if isinstance(operation, Control) and hasattr(operation, 'condition_needs'):
+            for n in operation.condition_needs:
+                self.graph.add_edge(DataPlaceholderNode(n), operation)
+
         # clear compiled steps (must recompile after adding new layers)
         self.steps = []
 
@@ -97,6 +101,8 @@ def show_layers(self):
             print("\t", "needs: ", step.needs)
             print("\t", "provides: ", step.provides)
             print("\t", "color: ", step.color)
+            if hasattr(step, 'condition_needs'):
+                print("\t", "condition needs: ", step.condition_needs)
             print("")
 
     def compile(self):
@@ -107,14 +113,37 @@ def compile(self):
         self.steps = []
 
         # create an execution order such that each layer's needs are provided.
-        ordered_nodes = list(nx.dag.topological_sort(self.graph))
+        try:
+            def key(node):
+
+                if hasattr(node, 'order'):
+                    return node.order
+                elif isinstance(node, DataPlaceholderNode):
+                    return float('-inf')
+                else:
+                    return 0
+
+            ordered_nodes = list(nx.dag.lexicographical_topological_sort(self.graph,
+                                                                         key=key))
+        except TypeError as e:
+            if self._debug:
+                print("Lexicographical topological sort failed! Falling back to topological sort.")
+
+            if not any(map(lambda node: isinstance(node, Control), self.graph.nodes)):
+                ordered_nodes = list(nx.dag.topological_sort(self.graph))
+            else:
+                print("Topological sort failed!")
+                raise e
 
         # add Operations evaluation steps, and instructions to free data.
         for i, node in enumerate(ordered_nodes):
 
             if isinstance(node, DataPlaceholderNode):
                 continue
 
+            elif isinstance(node, Control):
+                self.steps.append(node)
+
             elif isinstance(node, Operation):
 
                 # add layer to list of steps
@@ -256,11 +285,24 @@ def compute(self, outputs, named_inputs, color=None):
         # Find the subset of steps we need to run to get to the requested
         # outputs from the provided inputs.
         all_steps = self._find_necessary_steps(outputs, named_inputs, color)
-
+        # import pdb
         self.times = {}
+        if_true = False
         for step in all_steps:
 
-            if isinstance(step, Operation):
+            if isinstance(step, Control):
+                # pdb.set_trace()
+                if hasattr(step, 'condition'):
+                    if_true = step._compute_condition(cache)
+                    if if_true:
+                        layer_outputs = step._compute(cache)
+                        cache.update(layer_outputs)
+                elif not if_true:
+                    layer_outputs = step._compute(cache)
+                    cache.update(layer_outputs)
+                    if_true = False
+
+            elif isinstance(step, Operation):
 
                 if self._debug:
                     print("-"*32)