Merge pull request #254 from neserende/edge-compartment-update

Fixed the error about edge_attr dict in cases where edge is defined a…

ndlib/models/compartments/EdgeCategoricalAttribute.py

@@ -23,24 +23,47 @@ def __init__(
 
     def execute(self, node, graph, status, status_map, *args, **kwargs):
         neighbors = list(graph.neighbors(node))
+        isDiGraph = False #true if the current graph is directed
+
         if isinstance(graph, nx.DiGraph):
             neighbors = list(graph.predecessors(node))
+            isDiGraph = False
 
         edge_attr = graph.get_edge_attributes(self.attribute)
 
         if self.trigger is not None:
-            triggered = [
-                v
-                for v in neighbors
-                if status[v] == status_map[self.trigger]
-                and edge_attr[(min([node, v]), max([node, v]))] == self.attribute_value
-            ]
+            if(isDiGraph):
+                triggered = [
+                    v
+                    for v in neighbors
+                    if status[v] == status_map[self.trigger]
+                    and edge_attr[(min([node, v]), max([node, v]))] == self.attribute_value
+                ]
+            else:
+                triggered = []
+                for v in neighbors:
+                    if((node, v) in edge_attr):
+                        check = edge_attr[(node, v)]
+                    else: #then graph.has_edge(v, node) is True because v is given as neighbor of node
+                        check = edge_attr[(v, node)]
+                    if status[v] == status_map[self.trigger] and check == self.attribute_value:
+                        triggered.append(v)
         else:
-            triggered = [
-                v
-                for v in neighbors
-                if edge_attr[(min([node, v]), max([node, v]))] == self.attribute_value
-            ]
+            if(isDiGraph):
+                triggered = [
+                    v
+                    for v in neighbors
+                    if edge_attr[(min([node, v]), max([node, v]))] == self.attribute_value
+                ]
+            else:
+                triggered = []
+                for v in neighbors:
+                    if((node, v) in edge_attr):
+                        check = edge_attr[(node, v)]
+                    else:
+                        check = edge_attr[(v, node)]
+                    if check == self.attribute_value:
+                        triggered.append(v)
 
         for _ in triggered:
             p = np.random.random_sample()

ndlib/models/compartments/EdgeNumericalAttribute.py

@@ -67,7 +67,10 @@ def execute(self, node, graph, status, status_map, *args, **kwargs):
         if self.trigger is not None:
             for v in neighbors:
                 if status[v] == status_map[self.trigger]:
-                    val = edge_attr[(min([node, v]), max([node, v]))]
+                    if((node, v) in edge_attr):
+                        val = edge_attr[(node, v)]
+                    else:
+                        val = edge_attr[(v, node)]
                     if self.operator == "IN":
                         if self.__available_operators[self.operator][0](
                             val, self.attribute_range[0]
@@ -82,7 +85,10 @@ def execute(self, node, graph, status, status_map, *args, **kwargs):
                             triggered.append(v)
         else:
             for v in neighbors:
-                val = edge_attr[(min([node, v]), max([node, v]))]
+                if((node, v) in edge_attr):
+                    val = edge_attr[(node, v)]
+                else:
+                    val = edge_attr[(v, node)]
                 if self.operator == "IN":
                     if self.__available_operators[self.operator][0](
                         val, self.attribute_range[0]