fix: tests

Author: rememberYou

examples/mutag.py

@@ -8,6 +8,7 @@
 from pyrdf2vec import RDF2VecTransformer
 from pyrdf2vec.embedders import Word2Vec
 from pyrdf2vec.graphs import KG
+from pyrdf2vec.samplers import WideSampler
 from pyrdf2vec.walkers import HALKWalker
 
 # Ensure the determinism of this script by initializing a pseudo-random number.
@@ -31,8 +32,17 @@
     Word2Vec(workers=1, epochs=10),
     # Extract all walks with a maximum depth of 2 for each entity using two
     # processes and use a random state to ensure that the same walks are
-    # generated for the entities.
-    walkers=[HALKWalker(2, None, n_jobs=2, random_state=RANDOM_STATE)],
+    # generated for the entities without hashing as MUTAG is a short KG.
+    walkers=[
+        HALKWalker(
+            2,
+            None,
+            n_jobs=2,
+            sampler=WideSampler(),
+            random_state=RANDOM_STATE,
+            md5_bytes=None,
+        )
+    ],
     verbose=1,
 ).fit_transform(
     KG(

pyrdf2vec/samplers/frequency.py

@@ -73,7 +73,7 @@ def get_weight(self, hop: Hop) -> int:
                 without the sampling strategy having been trained.
 
         """
-        if len(self._counts) == 0:
+        if not self._counts:
             raise ValueError(
                 "You must call the `fit(kg)` function before get the weight of"
                 + " a hop."
@@ -144,7 +144,7 @@ def get_weight(self, hop: Hop) -> int:
                 without the sampling strategy having been trained.
 
         """
-        if self._counts:
+        if not self._counts:
             raise ValueError(
                 "You must call the `fit(kg)` method before get the weight of"
                 + " a hop."
@@ -219,7 +219,7 @@ def get_weight(self, hop: Hop) -> int:
                 without the sampling strategy having been trained.
 
         """
-        if self._counts:
+        if not self._counts:
             raise ValueError(
                 "You must call the `fit(kg)` method before get the weight of"
                 + " a hop."

pyrdf2vec/samplers/wide.py

@@ -90,7 +90,9 @@ def get_weight(self, hop: Hop) -> float:
                 without the sampling strategy having been trained.
 
         """
-        if self._pred_degs or self._obj_degs or self._neighbor_counts:
+        if not (
+            self._pred_degs or self._obj_degs or not self._neighbor_counts
+        ):
             raise ValueError(
                 "You must call the `fit(kg)` method before get the weight of"
                 + " a hop."

pyrdf2vec/walkers/halk.py

@@ -141,13 +141,17 @@ def _post_extract(self, res: List[EntityWalks]) -> List[List[SWalk]]:
         for rare_predicates in pred_thresholds:
             for entity_walks in conv_res:
                 canonical_walks = []
+                curr_entity = entity_walks[0][0]
                 for walk in entity_walks:
-                    canonical_walk = [walk[0]]
+                    canonical_walk = [curr_entity]
                     for i, vertex in enumerate(walk[1::2], 2):
                         if vertex not in rare_predicates:
                             obj = walk[i] if i % 2 == 0 else walk[i + 1]
                             canonical_walk += [vertex, obj]
-                    if len(canonical_walk) > 1:
+                    if len(canonical_walk) >= 3:
                         canonical_walks.append(tuple(canonical_walk))
-                res_halk.append(canonical_walks)
+                if canonical_walks:
+                    res_halk.append(canonical_walks)
+                else:
+                    res_halk.append([(curr_entity,)])
         return res_halk

tests/walkers/test_halk.py

@@ -63,17 +63,17 @@ def test_extract(
         walks = HALKWalker(
             max_depth,
             max_walks,
-            freq_thresholds=[0.01],
+            freq_thresholds=[0.001],
             with_reverse=with_reverse,
             random_state=42,
-        )._extract(kg, Vertex(root))[root]
+        ).extract(kg, [root])
+
         if max_walks is not None:
-            if with_reverse:
-                assert len(walks) <= max_walks * max_walks
-            else:
-                assert len(walks) <= max_walks
-        for walk in walks:
-            if not with_reverse:
-                assert walk[0] == root
-            for pred_or_obj in walk[1:]:
-                assert pred_or_obj.startswith("b'")
+            assert len(walks) == 1
+
+        for entity_walks in walks:
+            for walk in entity_walks:
+                if not with_reverse:
+                    assert walk[0] == root
+                for obj in walk[2::2]:
+                    assert obj.startswith("b'")