Refactor n2v walk (#81)

* Use id indexing in wv

* Refactor _random_walks

* Update deps, use numpy typing

* Fix typing

Author: RemyLau

pyproject.toml

@@ -5,3 +5,6 @@ build-backend = "setuptools.build_meta"
 [tool.mypy]
 ignore_missing_imports = true
 follow_imports = "skip"
+plugins = [
+    "numpy.typing.mypy_plugin",
+]

requirements.txt

@@ -1,4 +1,5 @@
     gensim==4.1.2
     numpy==1.21.5
-    numba==0.55.0
+    numba==0.55.1
     numba-progress==0.0.2
+    nptyping==1.4.4

src/pecanpy/graph.py

@@ -425,8 +425,10 @@ def read_npz(self, path: str, weighted: bool):
         raw = np.load(path)
         self.set_node_ids(raw["IDs"].tolist())
         self.data = raw["data"]
-        if not weighted:  # overwrite edge weights with constant
-            self.data[:] = 1.0
+        if self.data is None:
+            raise ValueError("Adjacency matrix data not found.")
+        elif not weighted:
+            self.data[:] = 1.0  # overwrite edge weights with constant
         self.indptr = raw["indptr"]
         self.indices = raw["indices"]
 
@@ -523,7 +525,7 @@ def data(self) -> Optional[np.ndarray]:
     def data(self, data: np.ndarray):
         """Set adjacency matrix and the corresponding nonzero matrix."""
         self._data = data.astype(float)
-        self._nonzero = self._data != 0
+        self._nonzero = np.array(self._data != 0, dtype=bool)
 
     @property
     def nonzero(self) -> Optional[np.ndarray]:

src/pecanpy/pecanpy.py

@@ -1,9 +1,12 @@
 """Different strategies for generating node2vec walks."""
+from typing import Any
+from typing import Callable
 from typing import List
 from typing import Optional
 
 import numpy as np
 from gensim.models import Word2Vec
+from nptyping import NDArray
 from numba import njit
 from numba import prange
 from numba.np.ufunc.parallel import _get_thread_id
@@ -14,6 +17,9 @@
 from .rw import SparseRWGraph
 from .wrappers import Timer
 
+HasNbrs = Callable[[np.uint32], bool]
+MoveForward = Callable[..., np.uint32]
+
 
 class Base(BaseGraph):
     """Base node2vec object.
@@ -137,51 +143,68 @@ def simulate_walks(
         has_nbrs = self.get_has_nbrs()
         verbose = self.verbose
 
-        @njit(parallel=True, nogil=True)
-        def node2vec_walks(num_iter, progress_proxy):
-            """Simulate a random walk starting from start node."""
-            # Seed the random number generator
-            if random_state is not None:
-                np.random.seed(random_state + _get_thread_id())
-
-            # use the last entry of each walk index array to keep track of the
-            # effective walk length
-            walk_idx_mat = np.zeros((num_iter, walk_length + 2), dtype=np.uint32)
-            walk_idx_mat[:, 0] = start_node_idx_ary  # initialize seeds
-            walk_idx_mat[:, -1] = walk_length + 1  # set to full walk length by default
-
-            for i in prange(num_iter):
-                # initialize first step as normal random walk
-                start_node_idx = walk_idx_mat[i, 0]
-                if has_nbrs(start_node_idx):
-                    walk_idx_mat[i, 1] = move_forward(start_node_idx)
-                else:
-                    walk_idx_mat[i, -1] = 1
-                    continue
-
-                # start bias random walk
-                for j in range(2, walk_length + 1):
-                    cur_idx = walk_idx_mat[i, j - 1]
-                    if has_nbrs(cur_idx):
-                        prev_idx = walk_idx_mat[i, j - 2]
-                        walk_idx_mat[i, j] = move_forward(cur_idx, prev_idx)
-                    else:
-                        walk_idx_mat[i, -1] = j
-                        break
-
-                progress_proxy.update(1)
-
-            return walk_idx_mat
-
         # Acquire numba progress proxy for displaying the progress bar
         with ProgressBar(total=tot_num_jobs, disable=not verbose) as progress:
-            walk_idx_mat = node2vec_walks(tot_num_jobs, progress)
+            walk_idx_mat = self._random_walks(
+                tot_num_jobs,
+                walk_length,
+                random_state,
+                start_node_idx_ary,
+                has_nbrs,
+                move_forward,
+                progress,
+            )
 
         # Map node index back to node ID
         walks = [self._map_walk(walk_idx_ary) for walk_idx_ary in walk_idx_mat]
 
         return walks
 
+    @staticmethod
+    @njit(parallel=True, nogil=True)
+    def _random_walks(
+        tot_num_jobs: int,
+        walk_length: int,
+        random_state: Optional[int],
+        start_node_idx_ary: NDArray[(Any,), np.uint32],
+        has_nbrs: HasNbrs,
+        move_forward: MoveForward,
+        progress_proxy: ProgressBar,
+    ):
+        """Simulate a random walk starting from start node."""
+        # Seed the random number generator
+        if random_state is not None:
+            np.random.seed(random_state + _get_thread_id())
+
+        # use the last entry of each walk index array to keep track of the
+        # effective walk length
+        walk_idx_mat = np.zeros((tot_num_jobs, walk_length + 2), dtype=np.uint32)
+        walk_idx_mat[:, 0] = start_node_idx_ary  # initialize seeds
+        walk_idx_mat[:, -1] = walk_length + 1  # set to full walk length by default
+
+        for i in prange(tot_num_jobs):
+            # initialize first step as normal random walk
+            start_node_idx = walk_idx_mat[i, 0]
+            if has_nbrs(start_node_idx):
+                walk_idx_mat[i, 1] = move_forward(start_node_idx)
+            else:
+                walk_idx_mat[i, -1] = 1
+                continue
+
+            # start bias random walk
+            for j in range(2, walk_length + 1):
+                cur_idx = walk_idx_mat[i, j - 1]
+                if has_nbrs(cur_idx):
+                    prev_idx = walk_idx_mat[i, j - 2]
+                    walk_idx_mat[i, j] = move_forward(cur_idx, prev_idx)
+                else:
+                    walk_idx_mat[i, -1] = j
+                    break
+
+            progress_proxy.update(1)
+
+        return walk_idx_mat
+
     def setup_get_normalized_probs(self):
         """Transition probability computation setup.
 
@@ -260,10 +283,7 @@ def embed(
             seed=self.random_state,
         )
 
-        # index mapping back to node IDs
-        idx_list = [w2v.wv.get_index(i) for i in self.nodes]
-
-        return w2v.wv.vectors[idx_list]
+        return w2v.wv[self.nodes]
 
 
 class FirstOrderUnweighted(Base, SparseRWGraph):

tox.ini

@@ -19,7 +19,9 @@ commands =
 
 [testenv:mypy]
 skip_install = true
-deps = mypy
+deps =
+    mypy
+    numpy
 commands = mypy src/pecanpy
 
 [testenv:flake8]