Repository: pyg-team/pytorch_geometric
Pull request: Add `HeteroData` support to `RemoveIsolatedNodes` 
This PR adds `HeteroData` support to the `RemoveIsolatedNodes` transform, closing https://github.com/pyg-team/pytorch_geometric/issues/4013.

This is achieved by looping through the different node types and removing them one at a time.

Problem statement:
AttributeError: 'BaseStorage' object has no attribute 'edge_index' for hetero graph
### 🐛 Describe the bug

I tried to remove isolated nodes in my hetero graph. My hetero graph is as below:

data
Out[117]: 
HeteroData(
  card={ x=[15910, 300] },
  user={ x=[12363, 300] },
  (user, solved, card)={ edge_index=[88552, 2] }
)

When I use 'has_isolated_nodes()', it shows True

data.has_isolated_nodes()
Out[119]: True

Then I tried to remove the isolated nodes, but I got an error. 

`import torch`
`from torch_geometric.data import HeteroData`
`import torch_geometric.transforms as T`
`data = T.RemoveIsolatedNodes()(data)`


```
Traceback (most recent call last):

  File "/home/hojun/anaconda3/envs/geometric/lib/python3.6/site-packages/torch_geometric/data/storage.py", line 48, in __getattr__
    return self[key]

  File "/home/hojun/anaconda3/envs/geometric/lib/python3.6/site-packages/torch_geometric/data/storage.py", line 68, in __getitem__
    return self._mapping[key]

KeyError: 'edge_index'


During handling of the above exception, another exception occurred:

Traceback (most recent call last):

  File "<ipython-input-120-2f26aa6d93fc>", line 1, in <module>
    data = T.RemoveIsolatedNodes()(data)

  File "/home/hojun/anaconda3/envs/geometric/lib/python3.6/site-packages/torch_geometric/transforms/remove_isolated_nodes.py", line 13, in __call__
    out = remove_isolated_nodes(data.edge_index, data.edge_attr, num_nodes)

  File "/home/hojun/anaconda3/envs/geometric/lib/python3.6/site-packages/torch_geometric/data/hetero_data.py", line 118, in __getattr__
    return getattr(self._global_store, key)

  File "/home/hojun/anaconda3/envs/geometric/lib/python3.6/site-packages/torch_geometric/data/storage.py", line 51, in __getattr__
    f"'{self.__class__.__name__}' object has no attribute '{key}'")

AttributeError: 'BaseStorage' object has no attribute 'edge_index'
```

As it is hetero graph, I do not have edge_index alone. My metadata() shows as below:

data.metadata()
Out[124]: (['card', 'user'], [('user', 'solved', 'card')])


How can I use RemoveIsolatedNodes in hetero graph?


### Environment

* PyG version: 2.0.3
* PyTorch version: 1.10.1+cu102
* OS: LINUX
* Python version: 3.6
* CUDA/cuDNN version: cuda 10.2
* How you installed PyTorch and PyG (`conda`, `pip`, source): pip
* Any other relevant information (*e.g.*, version of `torch-scatter`):

Hints from issue discussion:
The `RemoveIsolatedNodes` transform does not support yet `HeteroData` objects. That's why you are seeing this error.
Do you still think this is P0 @rusty1s , if so I don't mind taking a look at adding support into `RemoveIsolatedNodes`
It would be great to finally fix this. Thanks @Padarn. It shouldn‘t be too hard to fix as well. Let me know if there are any questions.

Commit message: hetero isolated support

Patch to review:
diff --git a/test/transforms/test_remove_isolated_nodes.py b/test/transforms/test_remove_isolated_nodes.py
index 5040a517d52c..103ca41a427b 100644
--- a/test/transforms/test_remove_isolated_nodes.py
+++ b/test/transforms/test_remove_isolated_nodes.py
@@ -1,6 +1,6 @@
 import torch
 
-from torch_geometric.data import Data
+from torch_geometric.data import Data, HeteroData
 from torch_geometric.transforms import RemoveIsolatedNodes
 
 
@@ -16,3 +16,30 @@ def test_remove_isolated_nodes():
     assert data.edge_index.tolist() == [[0, 1, 0], [1, 0, 0]]
     assert data.edge_attr.tolist() == [1, 2, 4]
     assert data.x.tolist() == [[1], [3]]
+
+    data = HeteroData()
+    data['paper'].x = torch.arange(6).type(torch.float)
+    data['author'].x = torch.arange(6).type(torch.float)
+
+    # first self loop to be isolated, second not, and third not self loop
+    data['paper', 'paper'].edge_index = torch.Tensor([[0, 1, 2],
+                                                      [0, 1,
+                                                       3]]).type(torch.long)
+
+    # remove isolation of first paper self loop, and add another
+    data['paper', 'author'].edge_index = torch.Tensor([[1, 3, 5],
+                                                       [0, 1,
+                                                        2]]).type(torch.long)
+
+    # only edge_attr on one type
+    data['paper', 'author'].edge_attr = torch.Tensor([1, 2,
+                                                      3]).type(torch.long)
+
+    # add a duplicate (in node index) edge
+    data['paper', 'cites',
+         'author'].edge_index = torch.Tensor([[5], [2]]).type(torch.long)
+
+    data = RemoveIsolatedNodes()(data)
+    assert data['paper'].num_nodes == 4
+    assert data['author'].num_nodes == 3
+    assert torch.allclose(data['paper'].x, torch.Tensor([1, 2, 3, 5]))
diff --git a/torch_geometric/transforms/remove_isolated_nodes.py b/torch_geometric/transforms/remove_isolated_nodes.py
index 8c6969bdeeb2..62a451d5a89a 100644
--- a/torch_geometric/transforms/remove_isolated_nodes.py
+++ b/torch_geometric/transforms/remove_isolated_nodes.py
@@ -1,7 +1,9 @@
 import re
+from typing import Union
 
 import torch
 
+from torch_geometric.data import Data, HeteroData
 from torch_geometric.data.datapipes import functional_transform
 from torch_geometric.transforms import BaseTransform
 from torch_geometric.utils import remove_isolated_nodes
@@ -11,18 +13,91 @@
 class RemoveIsolatedNodes(BaseTransform):
     r"""Removes isolated nodes from the graph
     (functional name: :obj:`remove_isolated_nodes`)."""
-    def __call__(self, data):
-        num_nodes = data.num_nodes
-        out = remove_isolated_nodes(data.edge_index, data.edge_attr, num_nodes)
-        data.edge_index, data.edge_attr, mask = out
+    def __call__(self, data: Union[Data,
+                                   HeteroData]) -> Union[Data, HeteroData]:
 
-        if hasattr(data, '__num_nodes__'):
-            data.num_nodes = int(mask.sum())
+        if isinstance(data, Data):
+            num_nodes = data.num_nodes
+            out = remove_isolated_nodes(data.edge_index, data.edge_attr,
+                                        num_nodes)
+            data.edge_index, data.edge_attr, mask = out
 
-        for key, item in data:
+            if hasattr(data, '__num_nodes__'):
+                data.num_nodes = int(mask.sum())
+
+            for key, item in data:
+                if bool(re.search('edge', key)):
+                    continue
+                if torch.is_tensor(item) and item.size(0) == num_nodes:
+                    data[key] = item[mask]
+
+            return data
+
+        elif isinstance(data, HeteroData):
+            return remove_hetero_isolated_nodes(data)
+
+        raise TypeError(f"`RemoveIsolatedNodes` invalid type: {type(data)}")
+
+
+def remove_hetero_isolated_nodes(data: HeteroData):
+    r"""Removes the isolated nodes from the heterogenous graph
+    given by :attr:`data`.
+    Self-loops are preserved onlt for non-isolated nodes. Nodes with only
+    a self loop are removed.
+
+    Args:
+        data (HeteroData): The graph to remove nodes from.
+
+    :rtype: HeteroData
+    """
+    device = data[data.edge_types[0]].edge_index.device
+
+    for node_type in data.node_types:
+
+        num_nodes = data[node_type].num_nodes
+        mask = torch.zeros(num_nodes, dtype=torch.bool, device=device)
+
+        for edge_type in data.edge_types:
+            if 'edge_index' in data[edge_type]:
+                edge_index = data[edge_type].edge_index
+                if edge_type[0] == edge_type[-1] == node_type:
+                    loop_mask = torch.where(edge_index[0] != edge_index[1])
+                    mask[edge_index[0][loop_mask[0]]] = 1
+                elif edge_type[0] == node_type:
+                    mask[edge_index[0]] = 1
+                elif edge_type[-1] == node_type:
+                    mask[edge_index[1]] = 1
+
+        assoc = torch.full((num_nodes, ), -1, dtype=torch.long,
+                           device=mask.device)
+
+        assoc[mask] = torch.arange(mask.sum(), device=assoc.device)
+
+        for key, values in data[node_type].items():
             if bool(re.search('edge', key)):
                 continue
-            if torch.is_tensor(item) and item.size(0) == num_nodes:
-                data[key] = item[mask]
+            if torch.is_tensor(values) and values.size(
+                    0) == data[node_type].num_nodes:
+                data[node_type][key] = values[mask]
+
+        for edge_type in data.edge_types:
+            if 'edge_index' in data[edge_type]:
+                edge_index = data[edge_type].edge_index
+
+                if edge_type[0] == node_type:
+                    edge_index[0] = assoc[edge_index[0]]
+                if edge_type[-1] == node_type:
+                    edge_index[1] = assoc[edge_index[1]]
+
+                data[edge_type].edge_index = edge_index
+
+            if edge_type[0] == edge_type[-1] == node_type:
+                loop_mask = edge_index[0] != -1
+                data[edge_type].edge_index = edge_index[:, loop_mask]
+                if 'edge_attr' in data[edge_type]:
+                    edge_attr = data[edge_type].edge_attr
+                    print(edge_attr.size())
+                    print(loop_mask.size())
+                    data[edge_type].edge_attr = edge_attr[loop_mask]
 
-        return data
+    return data


Code review comments:
diff_hunk: @@ -11,18 +13,91 @@
 class RemoveIsolatedNodes(BaseTransform):
     r"""Removes isolated nodes from the graph
     (functional name: :obj:`remove_isolated_nodes`)."""
-    def __call__(self, data):
-        num_nodes = data.num_nodes
-        out = remove_isolated_nodes(data.edge_index, data.edge_attr, num_nodes)
-        data.edge_index, data.edge_attr, mask = out
+    def __call__(self, data: Union[Data,
+                                   HeteroData]) -> Union[Data, HeteroData]:
 
-        if hasattr(data, '__num_nodes__'):
-            data.num_nodes = int(mask.sum())
+        if isinstance(data, Data):
+            num_nodes = data.num_nodes
+            out = remove_isolated_nodes(data.edge_index, data.edge_attr,
+                                        num_nodes)
+            data.edge_index, data.edge_attr, mask = out
 
-        for key, item in data:
+            if hasattr(data, '__num_nodes__'):
+                data.num_nodes = int(mask.sum())
+
+            for key, item in data:
+                if bool(re.search('edge', key)):
+                    continue
+                if torch.is_tensor(item) and item.size(0) == num_nodes:
+                    data[key] = item[mask]
+
+            return data
+
+        elif isinstance(data, HeteroData):
+            return remove_hetero_isolated_nodes(data)
+
+        raise TypeError(f"`RemoveIsolatedNodes` invalid type: {type(data)}")
+
+
+def remove_hetero_isolated_nodes(data: HeteroData):
path: torch_geometric/transforms/remove_isolated_nodes.py
text: @author:
This cannot go in 'torch_geometric/utils/isolated.py` because that class is imported by `torch_geometric/data/hetero_data.py`.

I thought about writing this taking a dictionary of edge_index, in the same way that `remove_isolated_nodes` works, I can adapt it to work like this if wanted, but I'd rather the `remove_isolated_nodes` were just depreciated in favour of one taking a `Data` object.

Any thoughts?

@user1:
Yeah, handling this with `data` only sounds like a better approach :)