Refactor typing (#82)

* Type NDArray

* Fix doc annotations

* Refactor typing

* Ignore type alias error

* Use alias type

* Use type alias

* Use type alias

Author: RemyLau

src/pecanpy/graph.py

@@ -1,12 +1,17 @@
 """Lite graph objects used by pecanpy."""
-from typing import Dict
-from typing import Iterator
-from typing import List
-from typing import Optional
-from typing import Tuple
-
 import numpy as np
 
+from .typing import AdjMat
+from .typing import AdjNonZeroMat
+from .typing import CSR
+from .typing import Dict
+from .typing import Float32Array
+from .typing import Iterator
+from .typing import List
+from .typing import Optional
+from .typing import Tuple
+from .typing import Uint32Array
+
 
 class BaseGraph:
     """Base Graph object.
@@ -78,8 +83,8 @@ class AdjlstGraph(BaseGraph):
         Python data structures like list and dict.
 
     Examples:
-        Read ``.edg`` file and create ``SparseGraph`` object using ``.read_edg``
-        method.
+        Read ``.edg`` file and create ``SparseGraph`` object using
+        ``.read_edg`` method.
 
         >>> from pecanpy.graph import AdjlstGraph
         >>>
@@ -277,7 +282,7 @@ def save(self, path: str, unweighted: bool = False, delimiter: str = "\t"):
                 terms = (h_id, t_id) if unweighted else (h_id, t_id, str(w))
                 f.write(f"{delimiter.join(terms)}\n")
 
-    def to_csr(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    def to_csr(self) -> CSR:
         """Construct compressed sparse row matrix."""
         indptr = np.zeros(len(self.nodes) + 1, dtype=np.uint32)
         for i, row_data in enumerate(self._data):
@@ -297,16 +302,16 @@ def to_csr(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
 
         return indptr, indices, data
 
-    def to_dense(self) -> np.ndarray:
+    def to_dense(self) -> AdjMat:
         """Construct dense adjacency matrix.
 
         Note:
             This method does not return DenseGraph object, but instead return
-            dense adjacency matrix as ``numpy.ndarray``, the index is the same
+            dense adjacency matrix as NDArray, the index is the same
             as that of ``nodes``.
 
         Return:
-            numpy.ndarray: Full adjacency matrix as 2d numpy array.
+            NDArray: Full adjacency matrix as 2d numpy array.
 
         """
         n_nodes = len(self.nodes)
@@ -319,11 +324,11 @@ def to_dense(self) -> np.ndarray:
         return mat
 
     @classmethod
-    def from_mat(cls, adj_mat: np.ndarray, node_ids: List[str], **kwargs):
+    def from_mat(cls, adj_mat: AdjMat, node_ids: List[str], **kwargs):
         """Construct graph using adjacency matrix and node IDs.
 
         Args:
-            adj_mat(:obj:`numpy.ndarray`): 2D numpy array of adjacency matrix
+            adj_mat(NDArray): 2D numpy array of adjacency matrix
             node_ids(:obj:`list` of str): node ID list
 
         Return:
@@ -347,8 +352,8 @@ class SparseGraph(BaseGraph):
     """Sparse Graph object that stores graph as adjacency list.
 
     Examples:
-        Read ``.edg`` file and create ``SparseGraph`` object using ``.read_edg``
-        method.
+        Read ``.edg`` file and create ``SparseGraph`` object using
+        ``.read_edg`` method.
 
         >>> from pecanpy.graph import SparseGraph
         >>>
@@ -366,9 +371,9 @@ class SparseGraph(BaseGraph):
     def __init__(self):
         """Initialize SparseGraph object."""
         super().__init__()
-        self.data: Optional[np.ndarray] = None
-        self.indptr: Optional[np.ndarray] = None
-        self.indices: Optional[np.ndarray] = None
+        self.data: Optional[Float32Array] = None
+        self.indptr: Optional[Uint32Array] = None
+        self.indices: Optional[Uint32Array] = None
 
     @property
     def num_edges(self) -> int:
@@ -457,14 +462,14 @@ def from_adjlst_graph(cls, adjlst_graph, **kwargs):
         return g
 
     @classmethod
-    def from_mat(cls, adj_mat: np.ndarray, node_ids: List[str], **kwargs):
+    def from_mat(cls, adj_mat: AdjMat, node_ids: List[str], **kwargs):
         """Construct csr graph using adjacency matrix and node IDs.
 
         Note:
             Only consider positive valued edges.
 
         Args:
-            adj_mat(:obj:`numpy.ndarray`): 2D numpy array of adjacency matrix
+            adj_mat(NDArray): 2D numpy array of adjacency matrix
             node_ids(:obj:`list` of str): node ID list
 
         """
@@ -479,15 +484,15 @@ class DenseGraph(BaseGraph):
     """Dense Graph object that stores graph as array.
 
     Examples:
-        Read ``.npz`` files and create ``DenseGraph`` object using ``read_npz``.
+        Read ``.npz`` files and create ``DenseGraph`` object using ``read_npz``
 
         >>> from pecanpy.graph import DenseGraph
         >>>
         >>> g = DenseGraph() # initialize DenseGraph object
         >>>
         >>> g.read_npz(paht_to_npz_file, weighted=True, directed=False)
 
-        Read ``.edg`` files and create ``DenseGraph`` object using ``read_edg``.
+        Read ``.edg`` files and create ``DenseGraph`` object using ``read_edg``
 
         >>> from pecanpy.graph import DenseGraph
         >>>
@@ -505,8 +510,8 @@ class DenseGraph(BaseGraph):
     def __init__(self):
         """Initialize DenseGraph object."""
         super().__init__()
-        self._data: Optional[np.ndarray] = None
-        self._nonzero: Optional[np.ndarray] = None
+        self._data: Optional[AdjMat] = None
+        self._nonzero: Optional[AdjNonZeroMat] = None
 
     @property
     def num_edges(self) -> int:
@@ -517,18 +522,18 @@ def num_edges(self) -> int:
             raise ValueError("Empty graph.")
 
     @property
-    def data(self) -> Optional[np.ndarray]:
+    def data(self) -> Optional[AdjMat]:
         """Return the adjacency matrix."""
         return self._data
 
     @data.setter
-    def data(self, data: np.ndarray):
+    def data(self, data: AdjMat):
         """Set adjacency matrix and the corresponding nonzero matrix."""
         self._data = data.astype(float)
         self._nonzero = np.array(self._data != 0, dtype=bool)
 
     @property
-    def nonzero(self) -> Optional[np.ndarray]:
+    def nonzero(self) -> Optional[AdjNonZeroMat]:
         """Return the nonzero mask for the adjacency matrix."""
         return self._nonzero
 
@@ -580,11 +585,11 @@ def from_adjlst_graph(cls, adjlst_graph, **kwargs):
         return g
 
     @classmethod
-    def from_mat(cls, adj_mat: np.ndarray, node_ids: List[str], **kwargs):
+    def from_mat(cls, adj_mat: AdjMat, node_ids: List[str], **kwargs):
         """Construct dense graph using adjacency matrix and node IDs.
 
         Args:
-            adj_mat(:obj:`numpy.ndarray`): 2D numpy array of adjacency matrix
+            adj_mat(NDArray): 2D numpy array of adjacency matrix
             node_ids(:obj:`list` of str): node ID list
 
         """

src/pecanpy/pecanpy.py

@@ -1,12 +1,6 @@
 """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
@@ -15,11 +9,16 @@
 from .graph import BaseGraph
 from .rw import DenseRWGraph
 from .rw import SparseRWGraph
+from .typing import Embeddings
+from .typing import Float32Array
+from .typing import HasNbrs
+from .typing import List
+from .typing import MoveForward
+from .typing import Optional
+from .typing import Uint32Array
+from .typing import Uint64Array
 from .wrappers import Timer
 
-HasNbrs = Callable[[np.uint32], bool]
-MoveForward = Callable[..., np.uint32]
-
 
 class Base(BaseGraph):
     """Base node2vec object.
@@ -99,7 +98,7 @@ def __init__(
         self.random_state = random_state
         self._preprocessed: bool = False
 
-    def _map_walk(self, walk_idx_ary: np.ndarray) -> List[str]:
+    def _map_walk(self, walk_idx_ary: Uint32Array) -> List[str]:
         """Map walk from node index to node ID.
 
         Note:
@@ -166,11 +165,11 @@ def _random_walks(
         tot_num_jobs: int,
         walk_length: int,
         random_state: Optional[int],
-        start_node_idx_ary: NDArray[(Any,), np.uint32],
+        start_node_idx_ary: Uint32Array,
         has_nbrs: HasNbrs,
         move_forward: MoveForward,
         progress_proxy: ProgressBar,
-    ):
+    ) -> Uint32Array:
         """Simulate a random walk starting from start node."""
         # Seed the random number generator
         if random_state is not None:
@@ -241,7 +240,7 @@ def embed(
         window_size: int = 10,
         epochs: int = 1,
         verbose: bool = False,
-    ) -> np.ndarray:
+    ) -> Embeddings:
         """Generate embeddings.
 
         This is a shortcut function that combines ``simulate_walks`` with
@@ -264,7 +263,7 @@ def embed(
                 skip-gram training if set to True
 
         Return:
-            numpy.ndarray: The embedding matrix, each row is a node embedding
+            Embeddings: The embedding matrix, each row is a node embedding
                 vector. The index is the same as that for the graph.
 
         """
@@ -375,7 +374,10 @@ class PreComp(Base, SparseRWGraph):
     def __init__(self, *args, **kwargs):
         """Initialize PreComp mode node2vec."""
         Base.__init__(self, *args, **kwargs)
-        self.alias_j = self.alias_q = self.alias_indptr = self.alias_dim = None
+        self.alias_dim: Optional[Uint32Array] = None
+        self.alias_j: Optional[Uint32Array] = None
+        self.alias_q: Optional[Float32Array] = None
+        self.alias_indptr: Optional[Uint64Array] = None
 
     def get_move_forward(self):
         """Wrap ``move_forward``.
@@ -624,7 +626,7 @@ def alias_setup(probs):
 
     Args:
         probs (list(float32)): normalized transition probabilities array, could
-            be in either list or numpy.ndarray, of float32 values.
+            be in either list or NDArray, of float32 values.
 
     """
     k = probs.size

src/pecanpy/rw/sparse_rw.py

@@ -154,9 +154,9 @@ def isnotin(ptr_ary1, ptr_ary2):
     returned to the ``get_normalized_probs``.
 
     Args:
-        ptr_ary1 (:obj:`numpy.ndarray` of :obj:`uint32`): array of pointers to
+        ptr_ary1 (Uint32Array): array of pointers to
             the neighbors of the current state
-        ptr_ary2 (:obj:`numpy.ndarray` of :obj:`uint32`): array of pointers to
+        ptr_ary2 (Uint32Array): array of pointers to
             the neighbors of the previous state
 
     Returns:
@@ -241,15 +241,13 @@ def isnotin_extended(ptr_ary1, ptr_ary2, wts_ary2, noise_thresholds):
     the previous state as out edges if the edge weight is below average.
 
     Args:
-        ptr_ary1 (:obj:`numpy.ndarray` of :obj:`uint32`): array of pointers to
-            the neighbors of the current state
-        ptr_ary2 (:obj:`numpy.ndarray` of :obj:`uint32`): array of pointers to
-            the neighbors of the previous state
-        wts_ary2 (:obj: `numpy.ndarray` of :obj:`float32`): array of edge
-            weights of the previous state
-        noise_thresholds (:obj: `numpy.ndarray` of :obj:`float32`): array of
-            noisy edge threshold computed based on the average and the std of
-            the edge weights of each node
+        ptr_ary1 (Uint32Array): array of pointers to the neighbors of the
+            current state
+        ptr_ary2 (Uint32Array): array of pointers to the neighbors of the
+            previous state
+        wts_ary2 (Float32Array): array of edge weights of the previous state
+        noise_thresholds (Float32Array): array of noisy edge threshold computed
+            based on the average and the std of the edge weights of each node
 
     Return:
         Indicator of whether a neighbor of the current state is considered as

src/pecanpy/typing.py

@@ -0,0 +1,45 @@
+"""Type annotations."""
+from typing import Any
+from typing import Callable
+from typing import Dict
+from typing import Iterator
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+import numpy as np
+from nptyping import NDArray
+from typing_extensions import TypeAlias
+
+# Callbacks ###################################################################
+HasNbrs = Callable[[np.uint32], bool]
+MoveForward = Callable[..., np.uint32]
+
+# Numpy array types ###########################################################
+# issue with type alias (https://stackoverflow.com/questions/62073473)
+Embeddings: TypeAlias = NDArray[(Any, Any), np.float32]  # type: ignore
+AdjMat: TypeAlias = NDArray[(Any, Any), Any]  # type: ignore
+AdjNonZeroMat: TypeAlias = NDArray[(Any, Any), bool]  # type: ignore
+Uint32Array: TypeAlias = NDArray[(Any,), np.uint32]  # type: ignore
+Uint64Array: TypeAlias = NDArray[(Any,), np.uint64]  # type: ignore
+Float32Array: TypeAlias = NDArray[(Any,), np.float32]  # type: ignore
+CSR = Tuple[Uint32Array, Uint32Array, Float32Array]
+
+__all__ = [
+    "Any",
+    "Callable",
+    "Dict",
+    "Iterator",
+    "List",
+    "Tuple",
+    "Optional",
+    "NDArray",
+    "HasNbrs",
+    "MoveForward",
+    "Embeddings",
+    "AdjMat",
+    "AdjNonZeroMat",
+    "Uint32Array",
+    "Float32Array",
+    "CSR",
+]