🐛 Fix edge weight None in multi graphs will default to 1

eliorc · Aug 1, 2022 · 05a39e1 · 05a39e1
1 parent 3fc2f3f
commit 05a39e1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/node2vec/node2vec.py b/node2vec/node2vec.py
@@ -108,11 +108,11 @@ def _precompute_probabilities(self):
 
                     try:
                         if self.graph[current_node][destination].get(self.weight_key):
-                            weight = self.graph[current_node][destination].get(self.weight_key)
+                            weight = self.graph[current_node][destination].get(self.weight_key, 1)
                         else: 
                             ## Example : AtlasView({0: {'type': 1, 'weight':0.1}})- when we have edge weight
                             edge = list(self.graph[current_node][destination])[-1]
-                            weight = self.graph[current_node][destination][edge].get(self.weight_key)
+                            weight = self.graph[current_node][destination][edge].get(self.weight_key, 1)
 
                     except:
                         weight = 1