refactor: walkers

Author: rememberYou

pyrdf2vec/walkers/anonymous.py

@@ -3,14 +3,15 @@
 import attr
 
 from pyrdf2vec.graphs import KG, Vertex
-from pyrdf2vec.typings import EntityWalks, SWalk
+from pyrdf2vec.typings import EntityWalks, List, SWalk
 from pyrdf2vec.walkers import RandomWalker
 
 
 @attr.s
 class AnonymousWalker(RandomWalker):
-    """Walker that transforms label information into positional information in
-    order to anonymize the random walks.
+    """Anonymous walking strategy which transforms each vertex name other than
+    the root node, into positional information, in order to anonymize the
+    randomly extracted walks.
 
     Attributes:
         _is_support_remote: True if the walking strategy can be used with a
@@ -33,27 +34,24 @@ class AnonymousWalker(RandomWalker):
 
     """
 
-    def _extract(self, kg: KG, instance: Vertex) -> EntityWalks:
-        """Extracts walks rooted at the provided entities which are then each
-        transformed into a numerical representation.
+    def _extract(self, kg: KG, entity: Vertex) -> EntityWalks:
+        """Extracts random walks for an entity based on a Knowledge Graph.
 
         Args:
             kg: The Knowledge Graph.
-            instance: The instance to be extracted from the Knowledge Graph.
+            entity: The root node to extract walks.
 
         Returns:
-            The 2D matrix with its number of rows equal to the number of
-            provided entities; number of column equal to the embedding size.
+            A dictionary having the entity as key and a list of tuples as value
+            corresponding to the extracted walks.
 
         """
         canonical_walks: Set[SWalk] = set()
-        for walk in self.extract_walks(kg, instance):
-            canonical_walk = []
-            str_walk = [hop.name for hop in walk]
-            for i, hop in enumerate(walk):
-                if i == 0:
-                    canonical_walk.append(hop.name)
-                else:
-                    canonical_walk.append(str(str_walk.index(hop.name)))
+        for walk in self.extract_walks(kg, entity):
+            vertex_names = [vertex.name for vertex in walk]
+            canonical_walk: List[str] = [
+                vertex.name if i == 0 else str(vertex_names.index(vertex.name))
+                for i, vertex in enumerate(walk)
+            ]
             canonical_walks.add(tuple(canonical_walk))
-        return {instance.name: list(canonical_walks)}
+        return {entity.name: list(canonical_walks)}

pyrdf2vec/walkers/community.py

@@ -35,7 +35,11 @@ def sample_from_iterable(x):
 
 @attr.s
 class CommunityWalker(Walker):
-    """Defines the community walking strategy.
+    """Community walking strategy which groups vertices with similar properties
+    through probabilities and relations that are not explicitly modeled in a
+    Knowledge Graph. Similar to the Random walking strategy, the Depth First
+    Search (DFS) algorithm is used if a maximum number of walks is specified.
+    Otherwise, the Breath First Search (BFS) algorithm is chosen.
 
     Attributes:
         _is_support_remote: True if the walking strategy can be used with a
@@ -127,27 +131,23 @@ def _community_detection(self, kg: KG) -> None:
             self.labels_per_community[self.communities[node]].append(node)
 
     def _bfs(
-        self, kg: KG, root: Vertex, is_reverse: bool = False
+        self, kg: KG, entity: Vertex, is_reverse: bool = False
     ) -> List[Walk]:
-        """Extracts random walks of depth - 1 hops rooted in root with
-        Breadth-first search.
+        """Extracts random walks for an entity based on Knowledge Graph using
+        the Depth First Search (DFS) algorithm.
 
         Args:
             kg: The Knowledge Graph.
-
-                The graph from which the neighborhoods are extracted for the
-                provided entities.
-            root: The root node to extract walks.
+            entity: The root node to extract walks.
             is_reverse: True to get the parent neighbors instead of the child
                 neighbors, False otherwise.
                 Defaults to False.
 
         Returns:
-            The list of walks for the root node according to the depth and
-            max_walks.
+            The list of unique walks for the provided entity.
 
         """
-        walks: Set[Walk] = {(root,)}
+        walks: Set[Walk] = {(entity,)}
         for i in range(self.max_depth):
             for walk in walks.copy():
                 if is_reverse:
@@ -201,31 +201,27 @@ def _bfs(
         return list(walks)
 
     def _dfs(
-        self, kg: KG, root: Vertex, is_reverse: bool = False
+        self, kg: KG, entity: Vertex, is_reverse: bool = False
     ) -> List[Walk]:
-        """Extracts a random limited number of walks of depth - 1 hops rooted
-        in root with Depth-first search.
+        """Extracts random walks for an entity based on Knowledge Graph using
+        the Depth First Search (DFS) algorithm.
 
         Args:
             kg: The Knowledge Graph.
-
-                The graph from which the neighborhoods are extracted for the
-                provided entities.
-            root: The root node to extract walks.
+            entity: The root node to extract walks.
             is_reverse: True to get the parent neighbors instead of the child
                 neighbors, False otherwise.
-                Defaults to False
+                Defaults to False.
 
         Returns:
-            The list of walks for the root node according to the depth and
-            max_walks.
+            The list of unique walks for the provided entity.
 
         """
         self.sampler.visited = set()
         walks: List[Walk] = []
         assert self.max_walks is not None
         while len(walks) < self.max_walks:
-            sub_walk: Walk = (root,)
+            sub_walk: Walk = (entity,)
             d = 1
             while d // 2 < self.max_depth:
                 pred_obj = self.sampler.sample_hop(
@@ -301,18 +297,15 @@ def extract(
         self._community_detection(kg)
         return super().extract(kg, entities, verbose)
 
-    def extract_walks(self, kg: KG, root: Vertex) -> List[Walk]:
+    def extract_walks(self, kg: KG, entity: Vertex) -> List[Walk]:
         """Extracts random walks of depth - 1 hops rooted in root.
 
         Args:
             kg: The Knowledge Graph.
-
-                The graph from which the neighborhoods are extracted for the
-                provided entities.
-            root: The root node to extract walks.
+            entity: The root node to extract walks.
 
         Returns:
-            The list of walks.
+            The list of unique walks for the provided entity.
 
         """
         if self.max_walks is None:
@@ -322,33 +315,30 @@ def extract_walks(self, kg: KG, root: Vertex) -> List[Walk]:
         if self.with_reverse:
             return [
                 r_walk[:-1] + walk
-                for walk in fct_search(kg, root)
-                for r_walk in fct_search(kg, root, is_reverse=True)
+                for walk in fct_search(kg, entity)
+                for r_walk in fct_search(kg, entity, is_reverse=True)
             ]
-        return [walk for walk in fct_search(kg, root)]
+        return [walk for walk in fct_search(kg, entity)]
 
-    def _extract(self, kg: KG, instance: Vertex) -> EntityWalks:
-        """Extracts walks rooted at the provided entities which are then each
-        transformed into a numerical representation.
+    def _extract(self, kg: KG, entity: Vertex) -> EntityWalks:
+        """Extracts random walks for an entity based on a Knowledge Graph.
 
         Args:
             kg: The Knowledge Graph.
-            instance: The instance to be extracted from the Knowledge Graph.
+            entity: The root node to extract walks.
 
         Returns:
-            The 2D matrix with its number of rows equal to the number of
-            provided entities; number of column equal to the embedding size.
+            A dictionary having the entity as key and a list of tuples as value
+            corresponding to the extracted walks.
 
         """
         canonical_walks: Set[SWalk] = set()
-        for walk in self.extract_walks(kg, instance):
-            canonical_walk: List[str] = []
-            for i, hop in enumerate(walk):
-                if i == 0 or i % 2 == 1 or self.md5_bytes is None:
-                    canonical_walk.append(hop.name)
-                else:
-                    canonical_walk.append(
-                        str(md5(hop.name.encode()).digest()[: self.md5_bytes])
-                    )
+        for walk in self.extract_walks(kg, entity):
+            canonical_walk: List[str] = [
+                vertex.name
+                if i == 0 or i % 2 == 1 or self.md5_bytes is None
+                else str(md5(vertex.name.encode()).digest()[: self.md5_bytes])
+                for i, vertex in enumerate(walk)
+            ]
             canonical_walks.add(tuple(canonical_walk))
-        return {instance.name: list(canonical_walks)}
+        return {entity.name: list(canonical_walks)}

pyrdf2vec/walkers/halk.py

@@ -1,19 +1,19 @@
 from collections import defaultdict
 from hashlib import md5
-from typing import List, Set
+from typing import DefaultDict, List, Set
 
 import attr
 
 from pyrdf2vec.graphs import KG, Vertex
-from pyrdf2vec.typings import EntityWalks, SWalk
+from pyrdf2vec.typings import EntityWalks, SWalk, Walk
 from pyrdf2vec.walkers import RandomWalker
 
 
 @attr.s
 class HALKWalker(RandomWalker):
-    """Walker that removes the rare entities from the random walks in order to
-    increase the quality of the generated embeddings while decreasing the
-    memory usage.
+    """HALK walking strategy which removes rare vertices from randomly
+    extracted walks, increasing the quality of the generated embeddings while
+    memory usage decreases.
 
     Attributes:
         _is_support_remote: True if the walking strategy can be used with a
@@ -54,47 +54,83 @@ class HALKWalker(RandomWalker):
 
     md5_bytes = attr.ib(kw_only=True, default=8, type=int, repr=False)
 
-    def _extract(self, kg: KG, instance: Vertex) -> EntityWalks:
-        """Extracts walks rooted at the provided entities which are then each
-        transformed into a numerical representation.
+    def build_dictionary(
+        self, walks: List[Walk]
+    ) -> DefaultDict[Vertex, Set[int]]:
+        """Builds a dictionary of vertices mapped to the extracted walk indices.
 
         Args:
-            kg: The Knowledge Graph.
+            walks: The walks to build the dictionary.
+
+        Returns:
+            The dictionary of vertex.
+
+        """
+        vertex_to_windices: DefaultDict[Vertex, Set[int]] = defaultdict(set)
+        for i in range(len(walks)):
+            for vertex in walks[i]:
+                vertex_to_windices[vertex].add(i)
+        return vertex_to_windices
+
+    def get_rare_vertices(
+        self,
+        vertex_to_windices: DefaultDict[Vertex, Set[int]],
+        walks: List[Walk],
+        freq_threshold: float,
+    ) -> Set[Vertex]:
+        """Gets vertices which doesn't reach a certain threshold of frequency
+        of occurrence.
+
+        Args:
+            vertex_to_windices: The dictionary of vertices mapped to the
+                extracted walk indices.
+            walks: The walks.
+            freq_threshold: The threshold frequency of occurrence.
 
-                The graph from which the neighborhoods are extracted for the
-                provided entities.
-            instance: The instance to be extracted from the Knowledge Graph.
         Returns:
-            The 2D matrix with its number of rows equal to the number of
-            provided entities; number of column equal to the embedding size.
+            the infrequent vertices.
 
         """
-        walks = self.extract_walks(kg, instance)
+        rare_vertices = set()
+        for vertex in vertex_to_windices:
+            if len(vertex_to_windices[vertex]) / len(walks) < freq_threshold:
+                rare_vertices.add(vertex)
+        return rare_vertices
 
+    def _extract(self, kg: KG, entity: Vertex) -> EntityWalks:
+        """Extracts random walks for an entity based on a Knowledge Graph.
+
+        Args:
+            kg: The Knowledge Graph.
+            entity: The root node to extract walks.
+
+        Returns:
+            A dictionary having the entity as key and a list of tuples as value
+            corresponding to the extracted walks.
+
+        """
+        walks = self.extract_walks(kg, entity)
+        vertex_to_windices = self.build_dictionary(walks)
         canonical_walks: Set[SWalk] = set()
-        hop_to_freq = defaultdict(set)
-        for i in range(len(walks)):
-            for hop in walks[i]:
-                hop_to_freq[hop].add(i)
 
         for freq_threshold in self.freq_thresholds:
-            uniformative_hops = set()
-            for hop in hop_to_freq:
-                if len(hop_to_freq[hop]) / len(walks) < freq_threshold:
-                    uniformative_hops.add(hop)
-
+            rare_vertices = self.get_rare_vertices(
+                vertex_to_windices, walks, freq_threshold
+            )
             for walk in walks:
                 canonical_walk = []
-                for i, hop in enumerate(walk):
-                    if i == 0 or self.md5_bytes is None:
-                        canonical_walk.append(hop.name)
-                    elif hop.name not in uniformative_hops:
+                for i, vertex in enumerate(walk):
+                    if i == 0 or (
+                        vertex not in rare_vertices and self.md5_bytes is None
+                    ):
+                        canonical_walk.append(vertex.name)
+                    elif vertex not in rare_vertices:
                         canonical_walk.append(
                             str(
-                                md5(hop.name.encode()).digest()[
+                                md5(vertex.name.encode()).digest()[
                                     : self.md5_bytes
                                 ]
                             )
                         )
                 canonical_walks.add(tuple(canonical_walk))
-        return {instance.name: list(canonical_walks)}
+        return {entity.name: list(canonical_walks)}