Repository: pyg-team/pytorch_geometric
Pull request #7713 (2023-07-08 09:49:26)

`HeteroData` support in `to_networkx`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7713?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7713](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7713?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (155e546) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/10f235097a6ea06ccc5d544b6702812b8bd48276?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (10f2350) will **decrease** coverage by `0.35%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master    #7713      +/-   ##
==========================================
- Coverage   91.92%   91.58%   -0.35%     
==========================================
  Files         452      452              
  Lines       25522    25534      +12     
==========================================
- Hits        23461    23385      -76     
- Misses       2061     2149      +88     
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7713?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/utils/geodesic.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7713?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL2dlb2Rlc2ljLnB5) | `15.78% <ø> (ø)` | |
| [torch\_geometric/data/data.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7713?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvZGF0YS5weQ==) | `89.71% <100.00%> (-0.25%)` | :arrow_down: |
| [torch\_geometric/utils/convert.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7713?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL2NvbnZlcnQucHk=) | `57.57% <100.00%> (+1.77%)` | :arrow_up: |

... and [21 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7713/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7e17dd25a..f0b0f6bba 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added `HeteroData` support in `to_networkx` ([#7713](https://github.com/pyg-team/pytorch_geometric/pull/7713))
 - Added `FlopsCount` support via `fvcore` ([#7693](https://github.com/pyg-team/pytorch_geometric/pull/7693))
 - Added back support for PyTorch >= 1.11.0 ([#7656](https://github.com/pyg-team/pytorch_geometric/pull/7656))
 - Added `Data.sort()` and `HeteroData.sort()` functionalities ([#7649](https://github.com/pyg-team/pytorch_geometric/pull/7649))
diff --git a/torch_geometric/data/data.py b/torch_geometric/data/data.py
index c6d36d23d..fd19b16db 100644
--- a/torch_geometric/data/data.py
+++ b/torch_geometric/data/data.py
@@ -194,6 +194,15 @@ class BaseData:
         r"""Returns all edge-level tensor attribute names."""
         return list(set(chain(*[s.edge_attrs() for s in self.edge_stores])))
 
+    @property
+    def node_offsets(self) -> Dict[NodeType, int]:
+        out: Dict[NodeType, int] = {}
+        offset: int = 0
+        for store in self.node_stores:
+            out[store._key] = offset
+            offset += store.num_nodes
+        return out
+
     def generate_ids(self):
         r"""Generates and sets :obj:`n_id` and :obj:`e_id` attributes to assign
         each node and edge to a continuously ascending and unique ID."""
diff --git a/torch_geometric/utils/convert.py b/torch_geometric/utils/convert.py
index 946d34412..eccb23051 100644
--- a/torch_geometric/utils/convert.py
+++ b/torch_geometric/utils/convert.py
@@ -1,5 +1,5 @@
 from collections import defaultdict
-from typing import Any, Iterable, List, Optional, Tuple, Union
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 import scipy.sparse
 import torch
@@ -78,7 +78,10 @@ def from_scipy_sparse_matrix(
 
 
 def to_networkx(
-    data: 'torch_geometric.data.Data',
+    data: Union[
+        'torch_geometric.data.Data',
+        'torch_geometric.data.HeteroData',
+    ],
     node_attrs: Optional[Iterable[str]] = None,
     edge_attrs: Optional[Iterable[str]] = None,
     graph_attrs: Optional[Iterable[str]] = None,
@@ -90,7 +93,8 @@ def to_networkx(
     a directed :obj:`networkx.DiGraph` otherwise.
 
     Args:
-        data (torch_geometric.data.Data): The data object.
+        data (torch_geometric.data.Data or torch_geometric.data.HeteroData): A
+            homogeneous or heterogeneous data object.
         node_attrs (iterable of str, optional): The node attributes to be
             copied. (default: :obj:`None`)
         edge_attrs (iterable of str, optional): The edge attributes to be
@@ -98,14 +102,15 @@ def to_networkx(
         graph_attrs (iterable of str, optional): The graph attributes to be
             copied. (default: :obj:`None`)
         to_undirected (bool or str, optional): If set to :obj:`True` or
-            "upper", will return a :obj:`networkx.Graph` instead of a
+            :obj:`"upper"`, will return a :obj:`networkx.Graph` instead of a
             :obj:`networkx.DiGraph`. The undirected graph will correspond to
             the upper triangle of the corresponding adjacency matrix.
-            Similarly, if set to "lower", the undirected graph will correspond
-            to the lower triangle of the adjacency matrix. (default:
-            :obj:`False`)
+            Similarly, if set to :obj:`"lower"`, the undirected graph will
+            correspond to the lower triangle of the adjacency matrix.
+            Only applicable in case the :obj:`data` object holds a homogeneous
+            graph. (default: :obj:`False`)
         remove_self_loops (bool, optional): If set to :obj:`True`, will not
-            include self loops in the resulting graph. (default: :obj:`False`)
+            include self-loops in the resulting graph. (default: :obj:`False`)
 
     Examples:
 
@@ -120,47 +125,53 @@ def to_networkx(
     """
     import networkx as nx
 
-    G = nx.Graph() if to_undirected else nx.DiGraph()
-
-    G.add_nodes_from(range(data.num_nodes))
-
-    node_attrs = node_attrs or []
-    edge_attrs = edge_attrs or []
-    graph_attrs = graph_attrs or []
-
-    values = {}
-    for key, value in data(*(node_attrs + edge_attrs + graph_attrs)):
-        if torch.is_tensor(value):
-            value = value if value.dim() <= 1 else value.squeeze(-1)
-            values[key] = value.tolist()
-        else:
-            values[key] = value
-
-    to_undirected = "upper" if to_undirected is True else to_undirected
-    to_undirected_upper = True if to_undirected == "upper" else False
-    to_undirected_lower = True if to_undirected == "lower" else False
-
-    for i, (u, v) in enumerate(data.edge_index.t().tolist()):
+    from torch_geometric.data import HeteroData
 
-        if to_undirected_upper and u > v:
-            continue
-        elif to_undirected_lower and u < v:
-            continue
+    to_undirected_upper: bool = to_undirected in {'upper', True}
+    to_undirected_lower: bool = to_undirected == 'lower'
+    to_undirected: bool = to_undirected_upper or to_undirected_lower
 
-        if remove_self_loops and u == v:
-            continue
+    if isinstance(data, HeteroData) and to_undirected:
+        raise ValueError("'to_undirected' is not supported in "
+                         "'to_networkx' for heterogeneous graphs")
 
-        G.add_edge(u, v)
+    G = nx.Graph() if to_undirected else nx.DiGraph()
 
-        for key in edge_attrs:
-            G[u][v][key] = values[key][i]
+    def to_networkx_value(value: Any) -> Any:
+        return value.tolist() if isinstance(value, Tensor) else value
+
+    for key in graph_attrs or []:
+        G.graph[key] = to_networkx_value(data[key])
+
+    node_offsets = data.node_offsets
+    for store in data.node_stores:
+        start = node_offsets[store._key]
+        for i in range(store.num_nodes):
+            attr: Dict[str, Any] = {}
+            if isinstance(data, HeteroData):
+                attr['type'] = store._key
+            for key in node_attrs or []:
+                attr[key] = to_networkx_value(store[key][i])
+            G.add_node(start + i, **attr)
+
+    for store in data.edge_stores:
+        for i, (v, w) in enumerate(store.edge_index.t().tolist()):
+            if to_undirected_upper and v > w:
+                continue
+            elif to_undirected_lower and v < w:
+                continue
+            elif remove_self_loops and v == w and not store.is_bipartite():
+                continue
 
-    for key in node_attrs:
-        for i, feat_dict in G.nodes(data=True):
-            feat_dict.update({key: values[key][i]})
+            attr: Dict[str, Any] = {}
+            if isinstance(data, HeteroData):
+                v = v + node_offsets[store._key[0]]
+                w = w + node_offsets[store._key[-1]]
+                attr['type'] = store._key
+            for key in edge_attrs or []:
+                attr[key] = to_networkx_value(store[key][i])
 
-    for key in graph_attrs:
-        G.graph[key] = values[key]
+            G.add_edge(v, w, **attr)
 
     return G
 
diff --git a/torch_geometric/utils/geodesic.py b/torch_geometric/utils/geodesic.py
index 889f805b2..e452984d6 100644
--- a/torch_geometric/utils/geodesic.py
+++ b/torch_geometric/utils/geodesic.py
@@ -66,9 +66,9 @@ def geodesic_distance(
 
     if 'dest' in kwargs:
         dst = kwargs['dest']
-        warnings.warn(
-            "'dest' attribute in 'geodesic_distance' is deprecated and will "
-            "be removed in a future release. Use 'dst' argument instead.")
+        warnings.warn("'dest' attribute in 'geodesic_distance' is deprecated "
+                      "and will be removed in a future release. Use the 'dst' "
+                      "argument instead.")
 
     max_distance = float('inf') if max_distance is None else max_distance
 


Test patch:
diff --git a/test/utils/test_convert.py b/test/utils/test_convert.py
index f5f58a325..93ea55d84 100644
--- a/test/utils/test_convert.py
+++ b/test/utils/test_convert.py
@@ -3,7 +3,7 @@ import scipy.sparse
 import torch
 
 from torch_geometric.data import Data, HeteroData
-from torch_geometric.testing import withPackage
+from torch_geometric.testing import get_random_edge_index, withPackage
 from torch_geometric.utils import (
     from_cugraph,
     from_dgl,
@@ -131,6 +131,43 @@ def test_to_networkx_undirected():
     assert nx.to_numpy_array(G).tolist() == [[3, 2], [2, 0]]
 
 
+@withPackage('networkx')
+def test_to_networkx_hetero():
+    edge_index = get_random_edge_index(5, 10, 20, coalesce=True)
+
+    data = HeteroData()
+    data['global_id'] = 0
+    data['author'].x = torch.arange(5)
+    data['paper'].x = torch.arange(10)
+    data['author', 'paper'].edge_index = edge_index
+    data['author', 'paper'].edge_attr = torch.arange(edge_index.size(1))
+
+    G = to_networkx(data, node_attrs=['x'], edge_attrs=['edge_attr'],
+                    graph_attrs=['global_id'])
+
+    assert G.number_of_nodes() == 15
+    assert G.number_of_edges() == edge_index.size(1)
+
+    assert G.graph == {'global_id': 0}
+
+    for i, (v, data) in enumerate(G.nodes(data=True)):
+        assert i == v
+        assert len(data) == 2
+        if i < 5:
+            assert data['x'] == i
+            assert data['type'] == 'author'
+        else:
+            assert data['x'] == i - 5
+            assert data['type'] == 'paper'
+
+    for i, (v, w, data) in enumerate(G.edges(data=True)):
+        assert v == int(edge_index[0, i])
+        assert w == int(edge_index[1, i]) + 5
+        assert len(data) == 2
+        assert data['type'] == ('author', 'to', 'paper')
+        assert data['edge_attr'] == i
+
+
 @withPackage('networkx')
 def test_from_networkx():
     x = torch.randn(2, 8)