Skip to content

Commit

Permalink
Merge pull request #38 from gerritjandebruin/master
Browse files Browse the repository at this point in the history
Add seed parameter.
  • Loading branch information
eliorc authored Nov 28, 2020
2 parents 50952ef + 854c159 commit eb30a74
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 2 deletions.
3 changes: 2 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ graph = nx.fast_gnp_random_graph(n=100, p=0.5)
node2vec = Node2Vec(graph, dimensions=64, walk_length=30, num_walks=200, workers=4) # Use temp_folder for big graphs

# Embed nodes
model = node2vec.fit(window=10, min_count=1, batch_words=4) # Any keywords acceptable by gensim.Word2Vec can be passed, `diemnsions` and `workers` are automatically passed (from the Node2Vec constructor)
model = node2vec.fit(window=10, min_count=1, batch_words=4) # Any keywords acceptable by gensim.Word2Vec can be passed, `dimensions` and `workers` are automatically passed (from the Node2Vec constructor)

# Look for most similar nodes
model.wv.most_similar('2') # Output node names are always strings
Expand Down Expand Up @@ -73,6 +73,7 @@ edges_kv.save_word2vec_format(EDGES_EMBEDDING_FILENAME)
Use these keys exactly. If not set, will use the global ones which were passed on the object initialization`
10. `quiet`: Boolean controlling the verbosity. (default: False)
11. `temp_folder`: String path pointing to folder to save a shared memory copy of the graph - Supply when working on graphs that are too big to fit in memory during algorithm execution.
12. `seed`: Seed for the random number generator (default: None). Deterministic results can be obtained if seed is set and `workers=1`.

- `Node2Vec.fit` method:
Accepts any key word argument acceptable by gensim.Word2Vec
Expand Down
8 changes: 7 additions & 1 deletion node2vec/node2vec.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import random
import os
from collections import defaultdict

Expand All @@ -22,7 +23,7 @@ class Node2Vec:

def __init__(self, graph: nx.Graph, dimensions: int = 128, walk_length: int = 80, num_walks: int = 10, p: float = 1,
q: float = 1, weight_key: str = 'weight', workers: int = 1, sampling_strategy: dict = None,
quiet: bool = False, temp_folder: str = None):
quiet: bool = False, temp_folder: str = None, seed: int = None):
"""
Initiates the Node2Vec object, precomputes walking probabilities and generates the walks.
Expand All @@ -35,6 +36,7 @@ def __init__(self, graph: nx.Graph, dimensions: int = 128, walk_length: int = 80
:param weight_key: On weighted graphs, this is the key for the weight attribute (default: 'weight')
:param workers: Number of workers for parallel execution (default: 1)
:param sampling_strategy: Node specific sampling strategies, supports setting node specific 'q', 'p', 'num_walks' and 'walk_length'.
:param seed: Seed for the random number generator.
Use these keys exactly. If not set, will use the global ones which were passed on the object initialization
:param temp_folder: Path to folder with enough space to hold the memory map of self.d_graph (for big graphs); to be passed joblib.Parallel.temp_folder
"""
Expand Down Expand Up @@ -63,6 +65,10 @@ def __init__(self, graph: nx.Graph, dimensions: int = 128, walk_length: int = 80
self.temp_folder = temp_folder
self.require = "sharedmem"

if seed is not None:
random.seed(seed)
np.random.seed(seed)

self._precompute_probabilities()
self.walks = self._generate_walks()

Expand Down

0 comments on commit eb30a74

Please sign in to comment.