Repository: pyg-team/pytorch_geometric
Pull request #7203 (2023-04-18 11:12:16)

Extend dataset summary to create stats for each node/edge type
example output for hetero data:
```
FakeHeteroDataset (#graphs=20):
+------------+----------+----------+
|            |   #nodes |   #edges |
|------------+----------+----------|
| mean       |   3945.4 |  49149.2 |
| std        |    247.8 |   3481.2 |
| min        |   3300   |  40440   |
| quantile25 |   3849   |  47016.8 |
| median     |   3974   |  49820   |
| quantile75 |   4129   |  51367.8 |
| max        |   4207   |  53761   |
+------------+----------+----------+
Number of nodes per node type:
+------------+--------+--------+--------+--------+
|            |    #v0 |    #v1 |    #v2 |    #v3 |
|------------+--------+--------+--------+--------|
| mean       |  980.5 |  994.8 |  944.7 | 1025.4 |
| std        |  140.7 |  123.8 |  123.9 |  123.7 |
| min        |  765   |  770   |  754   |  753   |
| quantile25 |  862.2 |  895   |  863   |  936.2 |
| median     |  979   | 1005   |  915   | 1067   |
| quantile75 | 1068.8 | 1064.8 | 1036   | 1103.2 |
| max        | 1235   | 1197   | 1201   | 1225   |
+------------+--------+--------+--------+--------+
Number of edges per edge type:
+------------+-----------------------+-----------------------+-----------------------+-----------------------+-----------------------+
|            |   #('v0', 'e0', 'v2') |   #('v2', 'e0', 'v0') |   #('v1', 'e0', 'v0') |   #('v3', 'e0', 'v0') |   #('v1', 'e0', 'v3') |
|------------+-----------------------+-----------------------+-----------------------+-----------------------+-----------------------|
| mean       |                9749.5 |                9400.5 |                9897   |               10203   |                9899.2 |
| std        |                1395.8 |                1229.9 |                1230.2 |                1230.6 |                1228.8 |
| min        |                7613   |                7502   |                7650   |                7498   |                7657   |
| quantile25 |                8577.2 |                8589.5 |                8899.8 |                9315.2 |                8901.5 |
| median     |                9732   |                9108   |               10014   |               10611   |                9978   |
| quantile75 |               10624.5 |               10307.8 |               10584.8 |               10971.5 |               10597   |
| max        |               12278   |               11935   |               11909   |               12186   |               11908   |
+------------+-----------------------+-----------------------+-----------------------+-----------------------+-----------------------+
```

For datasets that haven't `node_types` or `edge_types` defined or calling `from_dataset` method with parameter `per_type_breakdown=False` output will be the same as before the change.

codecov[bot]: ## [Codecov](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7203?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7203](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7203?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (1604e86) into [master](https://codecov.io/gh/pyg-team/pytorch_geometric/commit/f0c9fed058038a3288185e1f583069e6d98a28b4?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (f0c9fed) will **decrease** coverage by `0.04%`.
> The diff coverage is `67.56%`.

```diff
@@            Coverage Diff             @@
##           master    #7203      +/-   ##
==========================================
- Coverage   91.68%   91.65%   -0.04%     
==========================================
  Files         436      436              
  Lines       24013    24048      +35     
==========================================
+ Hits        22017    22040      +23     
- Misses       1996     2008      +12     
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7203?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/data/summary.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7203?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvc3VtbWFyeS5weQ==) | `84.33% <67.56%> (-13.58%)` | :arrow_down: |

: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 e0a8404ee..03eaaec8b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -29,6 +29,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Extend dataset summary to create stats for each node/edge type ([#7203](https://github.com/pyg-team/pytorch_geometric/pull/7203))
 - Added an optional `batch_size` argument to `avg_pool_x` and `max_pool_x` ([#7216](https://github.com/pyg-team/pytorch_geometric/pull/7216))
 - Fixed `subgraph` on unordered inputs ([#7187](https://github.com/pyg-team/pytorch_geometric/pull/7187))
 - Allow missing node types in `HeteroDictLinear` ([#7185](https://github.com/pyg-team/pytorch_geometric/pull/7185))
diff --git a/torch_geometric/data/summary.py b/torch_geometric/data/summary.py
index c8570a319..4e1aa0eba 100644
--- a/torch_geometric/data/summary.py
+++ b/torch_geometric/data/summary.py
@@ -1,10 +1,11 @@
+from collections import defaultdict
 from dataclasses import dataclass
-from typing import List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 import torch
 from tqdm import tqdm
 
-from torch_geometric.data import Dataset
+from torch_geometric.data import Dataset, HeteroData
 
 
 @dataclass
@@ -40,12 +41,15 @@ class Summary:
     num_graphs: int
     num_nodes: Stats
     num_edges: Stats
+    num_nodes_per_type: Optional[Dict[str, Stats]] = None
+    num_edges_per_type: Optional[Dict[str, Stats]] = None
 
     @classmethod
     def from_dataset(
         cls,
         dataset: Dataset,
         progress_bar: Optional[bool] = None,
+        per_type: bool = True,
     ):
         r"""Creates a summary of a :class:`~torch_geometric.data.Dataset`
         object.
@@ -56,6 +60,9 @@ class Summary:
                 progress bar during stats computation. If set to :obj:`None`,
                 will automatically decide whether to show a progress bar based
                 on dataset size. (default: :obj:`None`)
+            per_type (bool, optional). If set to :obj:`True`, will separate
+                statistics per node and edge type (only applicable in
+                heterogeneous graph datasets). (default: :obj:`True`)
         """
         name = dataset.__class__.__name__
 
@@ -65,28 +72,86 @@ class Summary:
         if progress_bar:
             dataset = tqdm(dataset)
 
-        num_nodes_list, num_edges_list = [], []
+        num_nodes, num_edges = [], []
+        num_nodes_per_type = defaultdict(list)
+        num_edges_per_type = defaultdict(list)
+
         for data in dataset:
-            num_nodes_list.append(data.num_nodes)
-            num_edges_list.append(data.num_edges)
+            num_nodes.append(data.num_nodes)
+            num_edges.append(data.num_edges)
+
+            if per_type and isinstance(data, HeteroData):
+                for node_type in data.node_types:
+                    num_nodes_per_type[node_type].append(
+                        data[node_type].num_nodes)
+                for edge_type in data.edge_types:
+                    num_edges_per_type[edge_type].append(
+                        data[edge_type].num_edges)
+
+        if len(num_nodes_per_type) > 0:
+            num_nodes_per_type = {
+                node_type: Stats.from_data(num_nodes_list)
+                for node_type, num_nodes_list in num_nodes_per_type.items()
+            }
+        else:
+            num_nodes_per_type = None
+
+        if len(num_edges_per_type) > 0:
+            num_edges_per_type = {
+                edge_type: Stats.from_data(num_edges_list)
+                for edge_type, num_edges_list in num_edges_per_type.items()
+            }
+        else:
+            num_edges_per_type = None
 
         return cls(
             name=name,
             num_graphs=len(dataset),
-            num_nodes=Stats.from_data(num_nodes_list),
-            num_edges=Stats.from_data(num_edges_list),
+            num_nodes=Stats.from_data(num_nodes),
+            num_edges=Stats.from_data(num_edges),
+            num_nodes_per_type=num_nodes_per_type,
+            num_edges_per_type=num_edges_per_type,
         )
 
     def __repr__(self) -> str:
         from tabulate import tabulate
 
-        prefix = f'{self.name} (#graphs={self.num_graphs}):\n'
+        body = f'{self.name} (#graphs={self.num_graphs}):\n'
 
         content = [['', '#nodes', '#edges']]
         stats = [self.num_nodes, self.num_edges]
         for field in Stats.__dataclass_fields__:
             row = [field] + [f'{getattr(s, field):.1f}' for s in stats]
             content.append(row)
-        body = tabulate(content, headers='firstrow', tablefmt='psql')
-
-        return prefix + body
+        body += tabulate(content, headers='firstrow', tablefmt='psql')
+
+        if self.num_nodes_per_type is not None:
+            content = [['']]
+            content[0] += list(self.num_nodes_per_type.keys())
+
+            for field in Stats.__dataclass_fields__:
+                row = [field] + [
+                    f'{getattr(s, field):.1f}'
+                    for s in self.num_nodes_per_type.values()
+                ]
+                content.append(row)
+            body += "\nNumber of nodes per node type:\n"
+            body += tabulate(content, headers='firstrow', tablefmt='psql')
+
+        if self.num_edges_per_type is not None:
+            content = [['']]
+            content[0] += [
+                f"({', '.join(edge_type)})"
+                for edge_type in self.num_edges_per_type.keys()
+            ]
+
+            for field in Stats.__dataclass_fields__:
+                row = [field] + [
+                    f'{getattr(s, field):.1f}'
+                    for s in self.num_edges_per_type.values()
+                ]
+                content.append(row)
+            body += "\nNumber of edges per edge type:\n"
+            body += tabulate(content, headers='firstrow', tablefmt='psql')
+
+        return body


Test patch:
diff --git a/test/data/test_dataset_summary.py b/test/data/test_dataset_summary.py
index 4ebee67df..5eed8708d 100644
--- a/test/data/test_dataset_summary.py
+++ b/test/data/test_dataset_summary.py
@@ -1,10 +1,21 @@
 import torch
+from torch import Tensor
 
-from torch_geometric.data.summary import Summary
+from torch_geometric.data.summary import Stats, Summary
 from torch_geometric.datasets import FakeDataset, FakeHeteroDataset
 from torch_geometric.testing import withPackage
 
 
+def check_stats(stats: Stats, expected: Tensor):
+    assert stats.mean == float(expected.mean())
+    assert stats.std == float(expected.std())
+    assert stats.min == float(expected.min())
+    assert stats.quantile25 == float(expected.quantile(0.25))
+    assert stats.median == float(expected.median())
+    assert stats.quantile75 == float(expected.quantile(0.75))
+    assert stats.max == float(expected.max())
+
+
 def test_dataset_summary():
     dataset = FakeDataset(num_graphs=10)
     num_nodes = torch.Tensor([data.num_nodes for data in dataset])
@@ -15,27 +26,24 @@ def test_dataset_summary():
     assert summary.name == 'FakeDataset'
     assert summary.num_graphs == 10
 
-    assert summary.num_nodes.mean == num_nodes.mean().item()
-    assert summary.num_nodes.std == num_nodes.std().item()
-    assert summary.num_nodes.min == num_nodes.min().item()
-    assert summary.num_nodes.quantile25 == num_nodes.quantile(0.25).item()
-    assert summary.num_nodes.median == num_nodes.median().item()
-    assert summary.num_nodes.quantile75 == num_nodes.quantile(0.75).item()
-    assert summary.num_nodes.max == num_nodes.max().item()
+    check_stats(summary.num_nodes, num_nodes)
+    check_stats(summary.num_edges, num_edges)
+
+
+@withPackage('tabulate')
+def test_dataset_summary_representation():
+    dataset = FakeDataset(num_graphs=10)
+
+    summary1 = Summary.from_dataset(dataset, per_type=False)
+    summary2 = Summary.from_dataset(dataset, per_type=True)
 
-    assert summary.num_edges.mean == num_edges.mean().item()
-    assert summary.num_edges.std == num_edges.std().item()
-    assert summary.num_edges.min == num_edges.min().item()
-    assert summary.num_edges.quantile25 == num_edges.quantile(0.25).item()
-    assert summary.num_edges.median == num_edges.median().item()
-    assert summary.num_edges.quantile75 == num_edges.quantile(0.75).item()
-    assert summary.num_edges.max == num_edges.max().item()
+    assert str(summary1) == str(summary2)
 
 
 @withPackage('tabulate')
 def test_dataset_summary_hetero():
     dataset1 = FakeHeteroDataset(num_graphs=10)
-    summary1 = Summary.from_dataset(dataset1)
+    summary1 = Summary.from_dataset(dataset1, per_type=False)
 
     dataset2 = [data.to_homogeneous() for data in dataset1]
     summary2 = Summary.from_dataset(dataset2)
@@ -43,3 +51,48 @@ def test_dataset_summary_hetero():
 
     assert summary1 == summary2
     assert str(summary1) == str(summary2)
+
+
+@withPackage('tabulate')
+def test_dataset_summary_hetero_representation_length():
+    dataset = FakeHeteroDataset(num_graphs=10)
+    summary = Summary.from_dataset(dataset)
+    num_lines = len(str(summary).splitlines())
+
+    stats_len = len(Stats.__dataclass_fields__)
+    len_header_and_border = 5
+    num_tables = 3  # general, stats per node type, stats per edge type
+
+    assert num_lines == num_tables * (stats_len + len_header_and_border)
+
+
+def test_dataset_summary_hetero_per_type_check():
+    dataset = FakeHeteroDataset(num_graphs=10)
+    exp_num_nodes = torch.Tensor([data.num_nodes for data in dataset])
+    exp_num_edges = torch.Tensor([data.num_edges for data in dataset])
+
+    summary = dataset.get_summary()
+
+    assert summary.name == 'FakeHeteroDataset'
+    assert summary.num_graphs == 10
+
+    check_stats(summary.num_nodes, exp_num_nodes)
+    check_stats(summary.num_edges, exp_num_edges)
+
+    num_nodes_per_type = {}
+    for node_type in dataset.node_types:
+        num_nodes_per_type[node_type] = torch.Tensor(
+            [data[node_type].num_nodes for data in dataset])
+
+    assert len(summary.num_nodes_per_type) == len(dataset.node_types)
+    for node_type, stats in summary.num_nodes_per_type.items():
+        check_stats(stats, num_nodes_per_type[node_type])
+
+    num_edges_per_type = {}
+    for edge_type in dataset.edge_types:
+        num_edges_per_type[edge_type] = torch.Tensor(
+            [data[edge_type].num_edges for data in dataset])
+
+    assert len(summary.num_edges_per_type) == len(dataset.edge_types)
+    for edge_type, stats in summary.num_edges_per_type.items():
+        check_stats(stats, num_edges_per_type[edge_type])