Merge pull request #33 from /pull/32

Pull/32

Author: eliorc

README.md

@@ -4,15 +4,6 @@
 Python3 implementation of the node2vec algorithm Aditya Grover, Jure Leskovec and Vid Kocijan.
 [node2vec: Scalable Feature Learning for Networks. A. Grover, J. Leskovec. ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (KDD), 2016.](https://snap.stanford.edu/node2vec/)
 
-## Changes:
-
-New in `0.3.0` (right now on version `0.3.1`):
-
-Added support for big graphs which cannot be fit into memory during algorithm execution (causing OOM errors).
-
-Thanks  [`@pg2455`](https://github.com/pg2455) for the contribution of this feature.
-
-
 ## Installation
 
 `pip install node2vec`

node2vec/__init__.py

@@ -1,4 +1,4 @@
 from .node2vec import Node2Vec
 from . import edges
 
-__version__ = '0.3.1'
+__version__ = '0.3.2'

node2vec/node2vec.py

@@ -72,7 +72,6 @@ def _precompute_probabilities(self):
         """
 
         d_graph = self.d_graph
-        first_travel_done = set()
 
         nodes_generator = self.graph.nodes() if self.quiet \
             else tqdm(self.graph.nodes(), desc='Computing transition probabilities')
@@ -90,7 +89,6 @@ def _precompute_probabilities(self):
                     d_graph[current_node][self.PROBABILITIES_KEY] = dict()
 
                 unnormalized_weights = list()
-                first_travel_weights = list()
                 d_neighbors = list()
 
                 # Calculate unnormalized weights
@@ -110,23 +108,25 @@ def _precompute_probabilities(self):
 
                     # Assign the unnormalized sampling strategy weight, normalize during random walk
                     unnormalized_weights.append(ss_weight)
-                    if current_node not in first_travel_done:
-                        first_travel_weights.append(self.graph[current_node][destination].get(self.weight_key, 1))
                     d_neighbors.append(destination)
 
                 # Normalize
                 unnormalized_weights = np.array(unnormalized_weights)
                 d_graph[current_node][self.PROBABILITIES_KEY][
                     source] = unnormalized_weights / unnormalized_weights.sum()
 
-                if current_node not in first_travel_done:
-                    unnormalized_weights = np.array(first_travel_weights)
-                    d_graph[current_node][self.FIRST_TRAVEL_KEY] = unnormalized_weights / unnormalized_weights.sum()
-                    first_travel_done.add(current_node)
-
                 # Save neighbors
                 d_graph[current_node][self.NEIGHBORS_KEY] = d_neighbors
 
+            # Calculate first_travel weights for source
+            first_travel_weights = []
+
+            for destination in self.graph.neighbors(source):
+                first_travel_weights.append(self.graph[source][destination].get(self.weight_key, 1))
+
+            first_travel_weights = np.array(first_travel_weights)
+            d_graph[source][self.FIRST_TRAVEL_KEY] = first_travel_weights / first_travel_weights.sum()
+
     def _generate_walks(self) -> list:
         """
         Generates the random walks which will be used as the skip-gram input.

setup.py

@@ -4,7 +4,7 @@
 setup(
     name='node2vec',
     packages=['node2vec'],
-    version='0.3.1',
+    version='0.3.2',
     description='Implementation of the node2vec algorithm.',
     author='Elior Cohen',
     author_email='[email protected]',