wehub-resource-sync
179 lines
7.0 KiB
Python
179 lines
7.0 KiB
Python
"""Module for mapping between node/edge IDs and node/edge types."""
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from .. import backend as F, utils
|
|
|
|
from .._ffi.function import _init_api
|
|
|
|
|
|
__all__ = ["IdMap"]
|
|
|
|
|
|
class IdMap:
|
|
"""A map for converting node/edge IDs to their type IDs and type-wise IDs.
|
|
|
|
For a heterogeneous graph, DGL assigns an integer ID to each node/edge type;
|
|
node and edge of different types have independent IDs starting from zero.
|
|
Therefore, a node/edge can be uniquely identified by an ID pair,
|
|
``(type_id, type_wise_id)``. To make it convenient for distributed processing,
|
|
DGL further encodes the ID pair into one integer ID, which we refer to
|
|
as *homogeneous ID*.
|
|
|
|
DGL arranges nodes and edges so that all nodes of the same type have contiguous
|
|
homogeneous IDs. If the graph is partitioned, the nodes/edges of the same type
|
|
within a partition have contiguous homogeneous IDs.
|
|
|
|
Below is an example adjancency matrix of an unpartitioned heterogeneous graph
|
|
stored using the above ID assignment. Here, the graph has two types of nodes
|
|
(``T0`` and ``T1``), and four types of edges (``R0``, ``R1``, ``R2``, ``R3``).
|
|
There are a total of 400 nodes in the graph and each type has 200 nodes. Nodes
|
|
of type 0 have IDs in [0,200), while nodes of type 1 have IDs in [200, 400).
|
|
|
|
```
|
|
0 <- T0 -> 200 <- T1 -> 400
|
|
0 +-----------+------------+
|
|
| | |
|
|
^ | R0 | R1 |
|
|
T0 | | |
|
|
v | | |
|
|
200 +-----------+------------+
|
|
| | |
|
|
^ | R2 | R3 |
|
|
T1 | | |
|
|
v | | |
|
|
400 +-----------+------------+
|
|
```
|
|
|
|
Below shows the adjacency matrix after the graph is partitioned into two.
|
|
Note that each partition still has two node types and four edge types,
|
|
and nodes/edges of the same type have contiguous IDs.
|
|
|
|
```
|
|
partition 0 partition 1
|
|
|
|
0 <- T0 -> 100 <- T1 -> 200 <- T0 -> 300 <- T1 -> 400
|
|
0 +-----------+------------+-----------+------------+
|
|
| | | |
|
|
^ | R0 | R1 | |
|
|
T0 | | | |
|
|
v | | | |
|
|
100 +-----------+------------+ |
|
|
| | | |
|
|
^ | R2 | R3 | |
|
|
T1 | | | |
|
|
v | | | |
|
|
200 +-----------+------------+-----------+------------+
|
|
| | | |
|
|
^ | | R0 | R1 |
|
|
T0 | | | |
|
|
v | | | |
|
|
100 | +-----------+------------+
|
|
| | | |
|
|
^ | | R2 | R3 |
|
|
T1 | | | |
|
|
v | | | |
|
|
200 +-----------+------------+-----------+------------+
|
|
```
|
|
|
|
The following table is an alternative way to represent the above ID assignments.
|
|
It is easy to see that the homogeneous ID range [0, 100) is used for nodes of type 0
|
|
in partition 0, [100, 200) is used for nodes of type 1 in partition 0, and so on.
|
|
```
|
|
+---------+------+----------
|
|
range | type | partition
|
|
[0, 100) | 0 | 0
|
|
[100,200) | 1 | 0
|
|
[200,300) | 0 | 1
|
|
[300,400) | 1 | 1
|
|
```
|
|
|
|
The goal of this class is to, given a node's homogenous ID, convert it into the
|
|
ID pair ``(type_id, type_wise_id)``. For example, homogeneous node ID 90 is mapped
|
|
to (0, 90); homogeneous node ID 201 is mapped to (0, 101).
|
|
|
|
Parameters
|
|
----------
|
|
id_ranges : dict[str, Tensor].
|
|
Node ID ranges within partitions for each node type. The key is the node type
|
|
name in string. The value is a tensor of shape :math:`(K, 2)`, where :math:`K` is
|
|
the number of partitions. Each row has two integers: the starting and the ending IDs
|
|
for a particular node type in a partition. For example, all nodes of type ``"T"`` in
|
|
partition ``i`` has ID range ``id_ranges["T"][i][0]`` to ``id_ranges["T"][i][1]``.
|
|
It is the same as the `node_map` argument in `RangePartitionBook`.
|
|
"""
|
|
|
|
def __init__(self, id_ranges):
|
|
id_ranges_values = list(id_ranges.values())
|
|
assert isinstance(
|
|
id_ranges_values[0], np.ndarray
|
|
), "id_ranges should be a dict of numpy arrays."
|
|
self.num_parts = id_ranges_values[0].shape[0]
|
|
self.dtype = id_ranges_values[0].dtype
|
|
self.dtype_str = "int32" if self.dtype == np.int32 else "int64"
|
|
self.num_types = len(id_ranges)
|
|
ranges = np.zeros(
|
|
(self.num_parts * self.num_types, 2), dtype=self.dtype
|
|
)
|
|
typed_map = []
|
|
id_ranges = id_ranges_values
|
|
id_ranges.sort(key=lambda a: a[0, 0])
|
|
for i, id_range in enumerate(id_ranges):
|
|
ranges[i :: self.num_types] = id_range
|
|
map1 = np.cumsum(id_range[:, 1] - id_range[:, 0], dtype=self.dtype)
|
|
typed_map.append(map1)
|
|
|
|
assert np.all(np.diff(ranges[:, 0]) >= 0)
|
|
assert np.all(np.diff(ranges[:, 1]) >= 0)
|
|
self.range_start = utils.toindex(
|
|
np.ascontiguousarray(ranges[:, 0]), dtype=self.dtype_str
|
|
)
|
|
self.range_end = utils.toindex(
|
|
np.ascontiguousarray(ranges[:, 1]) - 1, dtype=self.dtype_str
|
|
)
|
|
self.typed_map = utils.toindex(
|
|
np.concatenate(typed_map), dtype=self.dtype_str
|
|
)
|
|
|
|
def __call__(self, ids):
|
|
"""Convert the homogeneous IDs to (type_id, type_wise_id).
|
|
|
|
Parameters
|
|
----------
|
|
ids : 1D tensor
|
|
The homogeneous ID.
|
|
|
|
Returns
|
|
-------
|
|
type_ids : Tensor
|
|
Type IDs
|
|
per_type_ids : Tensor
|
|
Type-wise IDs
|
|
"""
|
|
if self.num_types == 0:
|
|
return F.zeros((len(ids),), F.dtype(ids), F.cpu()), ids
|
|
if len(ids) == 0:
|
|
return ids, ids
|
|
|
|
ids = utils.toindex(ids, dtype=self.dtype_str)
|
|
ret = _CAPI_DGLHeteroMapIds(
|
|
ids.todgltensor(),
|
|
self.range_start.todgltensor(),
|
|
self.range_end.todgltensor(),
|
|
self.typed_map.todgltensor(),
|
|
self.num_parts,
|
|
self.num_types,
|
|
)
|
|
ret = utils.toindex(ret, dtype=self.dtype_str).tousertensor()
|
|
return ret[: len(ids)], ret[len(ids) :]
|
|
|
|
@property
|
|
def torch_dtype(self):
|
|
"""Return the data type of the ID map."""
|
|
# [TODO][Rui] Use torch instead of numpy.
|
|
return torch.int32 if self.dtype == np.int32 else torch.int64
|
|
|
|
|
|
_init_api("dgl.distributed.id_map")
|