fix weighted_outdegree calculations and add tests

src/pathpyG/core/graph.py

@@ -441,8 +441,11 @@ def weighted_outdegrees(self) -> torch.Tensor:
         Returns:
             tensor: Weighted outdegrees of nodes.
         """
+        edge_weight = getattr(self.data, 'edge_weight', None)
+        if edge_weight is None:
+            edge_weight = torch.ones(self.data.num_edges, device=self.data.edge_index.device)
         weighted_outdegree = scatter(
-            self.data.edge_weight, self.data.edge_index[0], dim=0, dim_size=self.data.num_nodes, reduce="sum"
+            edge_weight, self.data.edge_index[0], dim=0, dim_size=self.data.num_nodes, reduce="sum"
         )
         return weighted_outdegree
 
@@ -455,7 +458,10 @@ def transition_probabilities(self) -> torch.Tensor:
         """
         weighted_outdegree = self.weighted_outdegrees()
         source_ids = self.data.edge_index[0]
-        return self.data.edge_weight / weighted_outdegree[source_ids]
+        edge_weight = getattr(self.data, 'edge_weight', None)
+        if edge_weight is None:
+            edge_weight = torch.ones(self.data.num_edges, device=self.data.edge_index.device)
+        return edge_weight / weighted_outdegree[source_ids]
 
     def laplacian(self, normalization: Any = None, edge_attr: Any = None) -> Any:
         """Return Laplacian matrix for a given graph.

tests/core/test_graph.py

@@ -241,6 +241,28 @@ def test_out_degrees(simple_graph):
     assert out_degrees["c"] == 0
 
 
+def test_weighted_outdegrees(simple_graph):
+    # Test on graph without defined weights
+    out_degrees = simple_graph.weighted_outdegrees()
+    assert out_degrees.equal(torch.tensor([2, 1, 0]))
+
+    # Test on graph with defined weights
+    simple_graph.data["edge_weight"] = torch.tensor([1, 3, 2])
+    out_degrees = simple_graph.weighted_outdegrees()
+    assert out_degrees.equal(torch.tensor([4, 2, 0]))
+
+
+def test_transition_probabilities(simple_graph):
+    # Test on graph without defined weights
+    transition_probs = simple_graph.transition_probabilities()
+    assert transition_probs.equal(torch.tensor([0.5, 0.5, 1]))
+
+    # Test on graph with defined weights
+    simple_graph.data["edge_weight"] = torch.tensor([1, 3, 2])
+    transition_probs = simple_graph.transition_probabilities()
+    assert transition_probs.equal(torch.tensor([0.25, 0.75, 1]))
+
+
 def test_laplacian(simple_graph):
     laplacian = simple_graph.laplacian()
     assert laplacian.shape == (3, 3)