wehub-resource-sync
1608 lines
47 KiB
Python
1608 lines
47 KiB
Python
"""Module for heterogeneous graph index class definition."""
|
|
from __future__ import absolute_import
|
|
|
|
import itertools
|
|
import sys
|
|
|
|
import numpy as np
|
|
import scipy
|
|
|
|
from . import backend as F, utils
|
|
from ._ffi.function import _init_api
|
|
from ._ffi.object import ObjectBase, register_object
|
|
from ._ffi.streams import to_dgl_stream_handle
|
|
from .base import dgl_warning, DGLError
|
|
from .graph_index import from_coo
|
|
|
|
|
|
@register_object("graph.HeteroGraph")
|
|
class HeteroGraphIndex(ObjectBase):
|
|
"""HeteroGraph index object.
|
|
|
|
Note
|
|
----
|
|
Do not create GraphIndex directly.
|
|
"""
|
|
|
|
def __new__(cls):
|
|
obj = ObjectBase.__new__(cls)
|
|
obj._cache = {}
|
|
return obj
|
|
|
|
def __getstate__(self):
|
|
"""Issue: https://github.com/pytorch/pytorch/issues/32351
|
|
Need to set the tensor created in the __getstate__ function
|
|
as object attribute to avoid potential bugs
|
|
"""
|
|
self._pk_state = _CAPI_DGLHeteroPickle(self)
|
|
return self._pk_state
|
|
|
|
def __setstate__(self, state):
|
|
self._cache = {}
|
|
|
|
# Pickle compatibility check
|
|
# TODO: we should store a storage version number in later releases.
|
|
if isinstance(state, HeteroPickleStates):
|
|
# post-0.4.3
|
|
self.__init_handle_by_constructor__(_CAPI_DGLHeteroUnpickle, state)
|
|
elif isinstance(state, tuple) and len(state) == 3:
|
|
# pre-0.4.2
|
|
metagraph, num_nodes, edges = state
|
|
|
|
self._cache = {}
|
|
# loop over etypes and recover unit graphs
|
|
rel_graphs = []
|
|
for i, edges_per_type in enumerate(edges):
|
|
src_ntype, dst_ntype = metagraph.find_edge(i)
|
|
num_src = num_nodes[src_ntype]
|
|
num_dst = num_nodes[dst_ntype]
|
|
src_id, dst_id, _ = edges_per_type
|
|
rel_graphs.append(
|
|
create_unitgraph_from_coo(
|
|
1 if src_ntype == dst_ntype else 2,
|
|
num_src,
|
|
num_dst,
|
|
src_id,
|
|
dst_id,
|
|
["coo", "csr", " csc"],
|
|
)
|
|
)
|
|
self.__init_handle_by_constructor__(
|
|
_CAPI_DGLHeteroCreateHeteroGraph, metagraph, rel_graphs
|
|
)
|
|
|
|
@property
|
|
def metagraph(self):
|
|
"""Meta graph
|
|
|
|
Returns
|
|
-------
|
|
GraphIndex
|
|
The meta graph.
|
|
"""
|
|
return _CAPI_DGLHeteroGetMetaGraph(self)
|
|
|
|
def is_metagraph_unibipartite(self):
|
|
"""Return whether or not the graph is unibiparite."""
|
|
return _CAPI_DGLHeteroIsMetaGraphUniBipartite(self)
|
|
|
|
def number_of_ntypes(self):
|
|
"""Return number of node types."""
|
|
return self.metagraph.num_nodes()
|
|
|
|
def number_of_etypes(self):
|
|
"""Return number of edge types."""
|
|
return self.metagraph.num_edges()
|
|
|
|
def get_relation_graph(self, etype):
|
|
"""Get the unitgraph graph of the given edge/relation type.
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
The edge/relation type.
|
|
|
|
Returns
|
|
-------
|
|
HeteroGraphIndex
|
|
The unitgraph graph.
|
|
"""
|
|
return _CAPI_DGLHeteroGetRelationGraph(self, int(etype))
|
|
|
|
def flatten_relations(self, etypes):
|
|
"""Convert the list of requested unitgraph graphs into a single unitgraph
|
|
graph.
|
|
|
|
Parameters
|
|
----------
|
|
etypes : list[int]
|
|
The edge/relation types.
|
|
|
|
Returns
|
|
-------
|
|
FlattenedHeteroGraph
|
|
A flattened heterograph object
|
|
"""
|
|
return _CAPI_DGLHeteroGetFlattenedGraph(self, etypes)
|
|
|
|
def add_nodes(self, ntype, num):
|
|
"""Add nodes.
|
|
|
|
Parameters
|
|
----------
|
|
ntype : int
|
|
Node type
|
|
num : int
|
|
Number of nodes to be added.
|
|
"""
|
|
_CAPI_DGLHeteroAddVertices(self, int(ntype), int(num))
|
|
self.clear_cache()
|
|
|
|
def add_edge(self, etype, u, v):
|
|
"""Add one edge.
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
Edge type
|
|
u : int
|
|
The src node.
|
|
v : int
|
|
The dst node.
|
|
"""
|
|
_CAPI_DGLHeteroAddEdge(self, int(etype), int(u), int(v))
|
|
self.clear_cache()
|
|
|
|
def add_edges(self, etype, u, v):
|
|
"""Add many edges.
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
Edge type
|
|
u : utils.Index
|
|
The src nodes.
|
|
v : utils.Index
|
|
The dst nodes.
|
|
"""
|
|
_CAPI_DGLHeteroAddEdges(
|
|
self, int(etype), u.todgltensor(), v.todgltensor()
|
|
)
|
|
self.clear_cache()
|
|
|
|
def clear(self):
|
|
"""Clear the graph."""
|
|
_CAPI_DGLHeteroClear(self)
|
|
self._cache.clear()
|
|
|
|
@property
|
|
def dtype(self):
|
|
"""Return the data type of this graph index.
|
|
|
|
Returns
|
|
-------
|
|
DGLDataType
|
|
The data type of the graph.
|
|
"""
|
|
return _CAPI_DGLHeteroDataType(self)
|
|
|
|
@property
|
|
def ctx(self):
|
|
"""Return the context of this graph index.
|
|
|
|
Returns
|
|
-------
|
|
DGLContext
|
|
The context of the graph.
|
|
"""
|
|
return _CAPI_DGLHeteroContext(self)
|
|
|
|
def bits_needed(self, etype):
|
|
"""Return the number of integer bits needed to represent the unitgraph graph.
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
The edge type.
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
The number of bits needed.
|
|
"""
|
|
stype, dtype = self.metagraph.find_edge(etype)
|
|
if (
|
|
self.num_edges(etype) >= 0x80000000
|
|
or self.num_nodes(stype) >= 0x80000000
|
|
or self.num_nodes(dtype) >= 0x80000000
|
|
):
|
|
return 64
|
|
else:
|
|
return 32
|
|
|
|
def asbits(self, bits):
|
|
"""Transform the graph to a new one with the given number of bits storage.
|
|
|
|
NOTE: this method only works for immutable graph index
|
|
|
|
Parameters
|
|
----------
|
|
bits : int
|
|
The number of integer bits (32 or 64)
|
|
|
|
Returns
|
|
-------
|
|
HeteroGraphIndex
|
|
The graph index stored using the given number of bits.
|
|
"""
|
|
return _CAPI_DGLHeteroAsNumBits(self, int(bits))
|
|
|
|
def copy_to(self, ctx):
|
|
"""Copy this immutable graph index to the given device context.
|
|
|
|
NOTE: this method only works for immutable graph index
|
|
|
|
Parameters
|
|
----------
|
|
ctx : DGLContext
|
|
The target device context.
|
|
|
|
Returns
|
|
-------
|
|
HeteroGraphIndex
|
|
The graph index on the given device context.
|
|
"""
|
|
return _CAPI_DGLHeteroCopyTo(self, ctx.device_type, ctx.device_id)
|
|
|
|
def pin_memory(self):
|
|
"""Copies the graph structure to pinned memory, if it's not already
|
|
pinned.
|
|
|
|
NOTE: This function is similar to PyTorch's Tensor.pin_memory(), but
|
|
tailored for graphs. It utilizes the same pin_memory allocator as
|
|
PyTorch, so the lifecycle of the graph is also managed by PyTorch.
|
|
If a batch includes a DGL graph object (HeteroGraphIndex),
|
|
PyTorch's DataLoader memory pinning logic will detect it and
|
|
automatically activate this function when pin_memory=True.
|
|
|
|
Returns
|
|
-------
|
|
HeteroGraphIndex
|
|
The pinned graph index.
|
|
"""
|
|
return _CAPI_DGLHeteroPinMemory(self)
|
|
|
|
def pin_memory_(self):
|
|
"""Pin this graph to the page-locked memory.
|
|
|
|
NOTE: This is an inplace method to pin the current graph index, i.e.,
|
|
it does not require new memory allocation but simply flags the
|
|
existing graph structure to be page-locked. The graph structure
|
|
must be on CPU to be pinned. If the graph struture is already
|
|
pinned, the function directly returns it.
|
|
|
|
Returns
|
|
-------
|
|
HeteroGraphIndex
|
|
The pinned graph index.
|
|
"""
|
|
return _CAPI_DGLHeteroPinMemory_(self)
|
|
|
|
def unpin_memory_(self):
|
|
"""Unpin this graph from the page-locked memory.
|
|
|
|
NOTE: this is an inplace method.
|
|
If the graph struture is not pinned, e.g., on CPU or GPU,
|
|
the function directly returns it.
|
|
|
|
Returns
|
|
-------
|
|
HeteroGraphIndex
|
|
The unpinned graph index.
|
|
"""
|
|
return _CAPI_DGLHeteroUnpinMemory_(self)
|
|
|
|
def is_pinned(self):
|
|
"""Check if this graph is pinned to the page-locked memory.
|
|
|
|
Returns
|
|
-------
|
|
bool
|
|
True if the graph is pinned.
|
|
"""
|
|
return bool(_CAPI_DGLHeteroIsPinned(self))
|
|
|
|
def record_stream(self, stream):
|
|
"""Record the stream that is using this graph.
|
|
|
|
Parameters
|
|
----------
|
|
stream : torch.cuda.Stream
|
|
The stream that is using this graph.
|
|
|
|
Returns
|
|
-------
|
|
HeteroGraphIndex
|
|
self.
|
|
"""
|
|
return _CAPI_DGLHeteroRecordStream(self, to_dgl_stream_handle(stream))
|
|
|
|
def shared_memory(
|
|
self, name, ntypes=None, etypes=None, formats=("coo", "csr", "csc")
|
|
):
|
|
"""Return a copy of this graph in shared memory
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The name of the shared memory.
|
|
ntypes : list of str
|
|
Name of node types
|
|
etypes : list of str
|
|
Name of edge types
|
|
format : list of str
|
|
Desired formats to be materialized.
|
|
|
|
Returns
|
|
-------
|
|
HeteroGraphIndex
|
|
The graph index in shared memory
|
|
"""
|
|
assert len(name) > 0, "The name of shared memory cannot be empty"
|
|
assert len(formats) > 0
|
|
for fmt in formats:
|
|
assert fmt in ("coo", "csr", "csc")
|
|
ntypes = [] if ntypes is None else ntypes
|
|
etypes = [] if etypes is None else etypes
|
|
return _CAPI_DGLHeteroCopyToSharedMem(
|
|
self, name, ntypes, etypes, formats
|
|
)
|
|
|
|
def is_multigraph(self):
|
|
"""Return whether the graph is a multigraph
|
|
The time cost will be O(E)
|
|
|
|
Returns
|
|
-------
|
|
bool
|
|
True if it is a multigraph, False otherwise.
|
|
"""
|
|
return bool(_CAPI_DGLHeteroIsMultigraph(self))
|
|
|
|
def is_readonly(self):
|
|
"""Return whether the graph index is read-only.
|
|
|
|
Returns
|
|
-------
|
|
bool
|
|
True if it is a read-only graph, False otherwise.
|
|
"""
|
|
return bool(_CAPI_DGLHeteroIsReadonly(self))
|
|
|
|
def num_nodes(self, ntype):
|
|
"""Return the number of nodes.
|
|
|
|
Parameters
|
|
----------
|
|
ntype : int
|
|
Node type.
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
The number of nodes.
|
|
"""
|
|
return _CAPI_DGLHeteroNumVertices(self, int(ntype))
|
|
|
|
def num_edges(self, etype):
|
|
"""Return the number of edges.
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
Edge type.
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
The number of edges.
|
|
"""
|
|
return _CAPI_DGLHeteroNumEdges(self, int(etype))
|
|
|
|
# TODO(#5485): remove this method.
|
|
def number_of_nodes(self, ntype):
|
|
"""Return the number of nodes.
|
|
|
|
Parameters
|
|
----------
|
|
ntype : int
|
|
Node type
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
The number of nodes
|
|
"""
|
|
return _CAPI_DGLHeteroNumVertices(self, int(ntype))
|
|
|
|
# TODO(#5485): remove this method.
|
|
def number_of_edges(self, etype):
|
|
"""Return the number of edges.
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
Edge type
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
The number of edges
|
|
"""
|
|
return _CAPI_DGLHeteroNumEdges(self, int(etype))
|
|
|
|
def has_nodes(self, ntype, vids):
|
|
"""Return true if the nodes exist.
|
|
|
|
Parameters
|
|
----------
|
|
ntype : int
|
|
Node type
|
|
vid : Tensor
|
|
Node IDs
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
0-1 array indicating existence
|
|
"""
|
|
return F.from_dgl_nd(
|
|
_CAPI_DGLHeteroHasVertices(self, int(ntype), F.to_dgl_nd(vids))
|
|
)
|
|
|
|
def has_edges_between(self, etype, u, v):
|
|
"""Return true if the edge exists.
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
Edge type
|
|
u : Tensor
|
|
Src node Ids.
|
|
v : Tensor
|
|
Dst node Ids.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
0-1 array indicating existence
|
|
"""
|
|
return F.from_dgl_nd(
|
|
_CAPI_DGLHeteroHasEdgesBetween(
|
|
self, int(etype), F.to_dgl_nd(u), F.to_dgl_nd(v)
|
|
)
|
|
)
|
|
|
|
def predecessors(self, etype, v):
|
|
"""Return the predecessors of the node.
|
|
|
|
Assume that node_type(v) == dst_type(etype). Thus, the ntype argument is omitted.
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
Edge type
|
|
v : int
|
|
The node.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
Array of predecessors
|
|
"""
|
|
return F.from_dgl_nd(
|
|
_CAPI_DGLHeteroPredecessors(self, int(etype), int(v))
|
|
)
|
|
|
|
def successors(self, etype, v):
|
|
"""Return the successors of the node.
|
|
|
|
Assume that node_type(v) == src_type(etype). Thus, the ntype argument is omitted.
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
Edge type
|
|
v : int
|
|
The node.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
Array of successors
|
|
"""
|
|
return F.from_dgl_nd(
|
|
_CAPI_DGLHeteroSuccessors(self, int(etype), int(v))
|
|
)
|
|
|
|
def edge_ids_all(self, etype, u, v):
|
|
"""Return a triplet of arrays that contains the edge IDs.
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
Edge type
|
|
u : Tensor
|
|
The src nodes.
|
|
v : Tensor
|
|
The dst nodes.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The src nodes.
|
|
Tensor
|
|
The dst nodes.
|
|
Tensor
|
|
The edge ids.
|
|
"""
|
|
edge_array = _CAPI_DGLHeteroEdgeIdsAll(
|
|
self, int(etype), F.to_dgl_nd(u), F.to_dgl_nd(v)
|
|
)
|
|
|
|
src = F.from_dgl_nd(edge_array(0))
|
|
dst = F.from_dgl_nd(edge_array(1))
|
|
eid = F.from_dgl_nd(edge_array(2))
|
|
|
|
return src, dst, eid
|
|
|
|
def edge_ids_one(self, etype, u, v):
|
|
"""Return an arrays of edge IDs.
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
Edge type
|
|
u : Tensor
|
|
The src nodes.
|
|
v : Tensor
|
|
The dst nodes.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The edge ids.
|
|
"""
|
|
eid = F.from_dgl_nd(
|
|
_CAPI_DGLHeteroEdgeIdsOne(
|
|
self, int(etype), F.to_dgl_nd(u), F.to_dgl_nd(v)
|
|
)
|
|
)
|
|
return eid
|
|
|
|
def find_edges(self, etype, eid):
|
|
"""Return a triplet of arrays that contains the edge IDs.
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
Edge type
|
|
eid : Tensor
|
|
Edge ids.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The src nodes.
|
|
Tensor
|
|
The dst nodes.
|
|
Tensor
|
|
The edge ids.
|
|
"""
|
|
edge_array = _CAPI_DGLHeteroFindEdges(
|
|
self, int(etype), F.to_dgl_nd(eid)
|
|
)
|
|
|
|
src = F.from_dgl_nd(edge_array(0))
|
|
dst = F.from_dgl_nd(edge_array(1))
|
|
eid = F.from_dgl_nd(edge_array(2))
|
|
|
|
return src, dst, eid
|
|
|
|
def in_edges(self, etype, v):
|
|
"""Return the in edges of the node(s).
|
|
|
|
Assume that node_type(v) == dst_type(etype). Thus, the ntype argument is omitted.
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
Edge type
|
|
v : Tensor
|
|
Node IDs.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The src nodes.
|
|
Tensor
|
|
The dst nodes.
|
|
Tensor
|
|
The edge ids.
|
|
"""
|
|
edge_array = _CAPI_DGLHeteroInEdges_2(self, int(etype), F.to_dgl_nd(v))
|
|
src = F.from_dgl_nd(edge_array(0))
|
|
dst = F.from_dgl_nd(edge_array(1))
|
|
eid = F.from_dgl_nd(edge_array(2))
|
|
return src, dst, eid
|
|
|
|
def out_edges(self, etype, v):
|
|
"""Return the out edges of the node(s).
|
|
|
|
Assume that node_type(v) == src_type(etype). Thus, the ntype argument is omitted.
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
Edge type
|
|
v : Tensor
|
|
Node IDs.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The src nodes.
|
|
Tensor
|
|
The dst nodes.
|
|
Tensor
|
|
The edge ids.
|
|
"""
|
|
edge_array = _CAPI_DGLHeteroOutEdges_2(self, int(etype), F.to_dgl_nd(v))
|
|
src = F.from_dgl_nd(edge_array(0))
|
|
dst = F.from_dgl_nd(edge_array(1))
|
|
eid = F.from_dgl_nd(edge_array(2))
|
|
return src, dst, eid
|
|
|
|
def edges(self, etype, order=None):
|
|
"""Return all the edges
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
Edge type
|
|
order : string
|
|
The order of the returned edges. Currently support:
|
|
|
|
- 'srcdst' : sorted by their src and dst ids.
|
|
- 'eid' : sorted by edge Ids.
|
|
- None : the arbitrary order.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The src nodes.
|
|
Tensor
|
|
The dst nodes.
|
|
Tensor
|
|
The edge ids.
|
|
"""
|
|
if order is None:
|
|
order = ""
|
|
elif order not in ["srcdst", "eid"]:
|
|
raise DGLError(
|
|
"Expect order to be one of None, 'srcdst', 'eid', "
|
|
"got {}".format(order)
|
|
)
|
|
edge_array = _CAPI_DGLHeteroEdges(self, int(etype), order)
|
|
src = F.from_dgl_nd(edge_array(0))
|
|
dst = F.from_dgl_nd(edge_array(1))
|
|
eid = F.from_dgl_nd(edge_array(2))
|
|
return src, dst, eid
|
|
|
|
def in_degrees(self, etype, v):
|
|
"""Return the in degrees of the nodes.
|
|
|
|
Assume that node_type(v) == dst_type(etype). Thus, the ntype argument is omitted.
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
Edge type
|
|
v : Tensor
|
|
The nodes.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The in degree array.
|
|
"""
|
|
return F.from_dgl_nd(
|
|
_CAPI_DGLHeteroInDegrees(self, int(etype), F.to_dgl_nd(v))
|
|
)
|
|
|
|
def out_degrees(self, etype, v):
|
|
"""Return the out degrees of the nodes.
|
|
|
|
Assume that node_type(v) == src_type(etype). Thus, the ntype argument is omitted.
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
Edge type
|
|
v : Tensor
|
|
The nodes.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The out degree array.
|
|
"""
|
|
return F.from_dgl_nd(
|
|
_CAPI_DGLHeteroOutDegrees(self, int(etype), F.to_dgl_nd(v))
|
|
)
|
|
|
|
def adjacency_matrix(self, etype, transpose, ctx):
|
|
"""Return the adjacency matrix representation of this graph.
|
|
|
|
By default, a row of returned adjacency matrix represents the source
|
|
of an edge and the column represents the destination.
|
|
|
|
When transpose is True, a row represents the destination and a column represents
|
|
the source.
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
Edge type
|
|
transpose : bool
|
|
A flag to transpose the returned adjacency matrix.
|
|
ctx : context
|
|
The context of the returned matrix.
|
|
|
|
Returns
|
|
-------
|
|
SparseTensor
|
|
The adjacency matrix.
|
|
Tensor
|
|
A index for data shuffling due to sparse format change. Return None
|
|
if shuffle is not required.
|
|
"""
|
|
if not isinstance(transpose, bool):
|
|
raise DGLError(
|
|
'Expect bool value for "transpose" arg,'
|
|
" but got %s." % (type(transpose))
|
|
)
|
|
fmt = F.get_preferred_sparse_format()
|
|
rst = _CAPI_DGLHeteroGetAdj(self, int(etype), transpose, fmt)
|
|
# convert to framework-specific sparse matrix
|
|
srctype, dsttype = self.metagraph.find_edge(etype)
|
|
nrows = (
|
|
self.num_nodes(dsttype) if transpose else self.num_nodes(srctype)
|
|
)
|
|
ncols = (
|
|
self.num_nodes(srctype) if transpose else self.num_nodes(dsttype)
|
|
)
|
|
nnz = self.num_edges(etype)
|
|
if fmt == "csr":
|
|
indptr = F.copy_to(F.from_dgl_nd(rst(0)), ctx)
|
|
indices = F.copy_to(F.from_dgl_nd(rst(1)), ctx)
|
|
shuffle = F.copy_to(F.from_dgl_nd(rst(2)), ctx)
|
|
dat = F.ones(
|
|
nnz, dtype=F.float32, ctx=ctx
|
|
) # FIXME(minjie): data type
|
|
spmat = F.sparse_matrix(
|
|
dat, ("csr", indices, indptr), (nrows, ncols)
|
|
)[0]
|
|
return spmat, shuffle
|
|
elif fmt == "coo":
|
|
idx = F.copy_to(F.from_dgl_nd(rst(0)), ctx)
|
|
idx = F.reshape(idx, (2, nnz))
|
|
dat = F.ones((nnz,), dtype=F.float32, ctx=ctx)
|
|
adj, shuffle_idx = F.sparse_matrix(
|
|
dat, ("coo", idx), (nrows, ncols)
|
|
)
|
|
return adj, shuffle_idx
|
|
else:
|
|
raise Exception("unknown format")
|
|
|
|
def adjacency_matrix_tensors(self, etype, transpose, fmt):
|
|
"""Return the adjacency matrix as a triplet of tensors.
|
|
|
|
By default, a row of returned adjacency matrix represents the source
|
|
of an edge and the column represents the destination.
|
|
|
|
When transpose is True, a row represents the destination and a column represents
|
|
the source.
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
Edge type
|
|
transpose : bool
|
|
A flag to transpose the returned adjacency matrix.
|
|
fmt : str
|
|
Indicates the format of returned adjacency matrix.
|
|
|
|
Returns
|
|
-------
|
|
tuple[int, int, Tensor, Tensor] or tuple[int, int, Tensor, Tensor, Tensor]
|
|
The number of rows and columns, followed by the adjacency matrix tensors
|
|
whose data type and device are the same as those of the graph.
|
|
|
|
If :attr:`fmt` is ``'coo'``, then the triplet will be
|
|
the row array and column array of the COO representation.
|
|
|
|
If :attr:`fmt` is ``'csr'``, then the triplet will be
|
|
the index pointer array (``indptr``), indices array, and data array
|
|
of the CSR representation. The data array will contain the edge ID for
|
|
each entry of the adjacency matrix. If the data array is empty, then it is
|
|
equivalent to a consecutive array from zero to the number of edges minus one.
|
|
"""
|
|
if not isinstance(transpose, bool):
|
|
raise DGLError(
|
|
'Expect bool value for "transpose" arg,'
|
|
" but got %s." % (type(transpose))
|
|
)
|
|
|
|
rst = _CAPI_DGLHeteroGetAdj(self, int(etype), transpose, fmt)
|
|
srctype, dsttype = self.metagraph.find_edge(etype)
|
|
nrows = (
|
|
self.num_nodes(dsttype) if transpose else self.num_nodes(srctype)
|
|
)
|
|
ncols = (
|
|
self.num_nodes(srctype) if transpose else self.num_nodes(dsttype)
|
|
)
|
|
nnz = self.num_edges(etype)
|
|
if fmt == "csr":
|
|
indptr = F.from_dgl_nd(rst(0))
|
|
indices = F.from_dgl_nd(rst(1))
|
|
data = F.from_dgl_nd(rst(2))
|
|
return nrows, ncols, indptr, indices, data
|
|
elif fmt == "coo":
|
|
idx = F.from_dgl_nd(rst(0))
|
|
row, col = F.reshape(idx, (2, nnz))
|
|
return nrows, ncols, row, col
|
|
else:
|
|
raise ValueError("unknown format")
|
|
|
|
def adjacency_matrix_scipy(
|
|
self, etype, transpose, fmt, return_edge_ids=None
|
|
):
|
|
"""Return the scipy adjacency matrix representation of this graph.
|
|
|
|
By default, a row of returned adjacency matrix represents the destination
|
|
of an edge and the column represents the source.
|
|
|
|
When transpose is True, a row represents the source and a column represents
|
|
a destination.
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
Edge type
|
|
transpose : bool
|
|
A flag to transpose the returned adjacency matrix.
|
|
fmt : str
|
|
Indicates the format of returned adjacency matrix.
|
|
return_edge_ids : bool
|
|
Indicates whether to return edge IDs or 1 as elements.
|
|
|
|
Returns
|
|
-------
|
|
scipy.sparse.spmatrix
|
|
The scipy representation of adjacency matrix.
|
|
"""
|
|
if return_edge_ids is None:
|
|
dgl_warning(
|
|
"Adjacency matrix by default currently returns edge IDs."
|
|
" As a result there is one 0 entry which is not eliminated."
|
|
" In the next release it will return 1s by default,"
|
|
" and 0 will be eliminated otherwise.",
|
|
FutureWarning,
|
|
)
|
|
return_edge_ids = True
|
|
|
|
if fmt == "csr":
|
|
nrows, ncols, indptr, indices, data = self.adjacency_matrix_tensors(
|
|
etype, transpose, fmt
|
|
)
|
|
indptr = F.asnumpy(indptr)
|
|
indices = F.asnumpy(indices)
|
|
data = F.asnumpy(data)
|
|
|
|
# Check if edge ID is omitted
|
|
if return_edge_ids and data.shape[0] == 0:
|
|
data = np.arange(self.num_edges(etype))
|
|
else:
|
|
data = np.ones_like(indices)
|
|
|
|
return scipy.sparse.csr_matrix(
|
|
(data, indices, indptr), shape=(nrows, ncols)
|
|
)
|
|
elif fmt == "coo":
|
|
nrows, ncols, row, col = self.adjacency_matrix_tensors(
|
|
etype, transpose, fmt
|
|
)
|
|
row = F.asnumpy(row)
|
|
col = F.asnumpy(col)
|
|
data = (
|
|
np.arange(self.num_edges(etype))
|
|
if return_edge_ids
|
|
else np.ones_like(row)
|
|
)
|
|
return scipy.sparse.coo_matrix(
|
|
(data, (row, col)), shape=(nrows, ncols)
|
|
)
|
|
else:
|
|
raise ValueError("unknown format")
|
|
|
|
def incidence_matrix(self, etype, typestr, ctx):
|
|
"""Return the incidence matrix representation of this graph.
|
|
|
|
An incidence matrix is an n x m sparse matrix, where n is
|
|
the number of nodes and m is the number of edges. Each nnz
|
|
value indicating whether the edge is incident to the node
|
|
or not.
|
|
|
|
There are three types of an incidence matrix `I`:
|
|
* "in":
|
|
- I[v, e] = 1 if e is the in-edge of v (or v is the dst node of e);
|
|
- I[v, e] = 0 otherwise.
|
|
* "out":
|
|
- I[v, e] = 1 if e is the out-edge of v (or v is the src node of e);
|
|
- I[v, e] = 0 otherwise.
|
|
* "both":
|
|
- I[v, e] = 1 if e is the in-edge of v;
|
|
- I[v, e] = -1 if e is the out-edge of v;
|
|
- I[v, e] = 0 otherwise (including self-loop).
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
Edge type
|
|
typestr : str
|
|
Can be either "in", "out" or "both"
|
|
ctx : context
|
|
The context of returned incidence matrix.
|
|
|
|
Returns
|
|
-------
|
|
SparseTensor
|
|
The incidence matrix.
|
|
utils.Index
|
|
A index for data shuffling due to sparse format change. Return None
|
|
if shuffle is not required.
|
|
"""
|
|
src, dst, eid = self.edges(etype)
|
|
srctype, dsttype = self.metagraph.find_edge(etype)
|
|
|
|
m = self.num_edges(etype)
|
|
if typestr == "in":
|
|
n = self.num_nodes(dsttype)
|
|
row = F.unsqueeze(dst, 0)
|
|
col = F.unsqueeze(eid, 0)
|
|
idx = F.copy_to(F.cat([row, col], dim=0), ctx)
|
|
# FIXME(minjie): data type
|
|
dat = F.ones((m,), dtype=F.float32, ctx=ctx)
|
|
inc, shuffle_idx = F.sparse_matrix(dat, ("coo", idx), (n, m))
|
|
elif typestr == "out":
|
|
n = self.num_nodes(srctype)
|
|
row = F.unsqueeze(src, 0)
|
|
col = F.unsqueeze(eid, 0)
|
|
idx = F.copy_to(F.cat([row, col], dim=0), ctx)
|
|
# FIXME(minjie): data type
|
|
dat = F.ones((m,), dtype=F.float32, ctx=ctx)
|
|
inc, shuffle_idx = F.sparse_matrix(dat, ("coo", idx), (n, m))
|
|
elif typestr == "both":
|
|
assert (
|
|
srctype == dsttype
|
|
), "'both' is supported only if source and destination type are the same"
|
|
n = self.num_nodes(srctype)
|
|
# first remove entries for self loops
|
|
mask = F.logical_not(F.equal(src, dst))
|
|
src = F.boolean_mask(src, mask)
|
|
dst = F.boolean_mask(dst, mask)
|
|
eid = F.boolean_mask(eid, mask)
|
|
n_entries = F.shape(src)[0]
|
|
# create index
|
|
row = F.unsqueeze(F.cat([src, dst], dim=0), 0)
|
|
col = F.unsqueeze(F.cat([eid, eid], dim=0), 0)
|
|
idx = F.copy_to(F.cat([row, col], dim=0), ctx)
|
|
# FIXME(minjie): data type
|
|
x = -F.ones((n_entries,), dtype=F.float32, ctx=ctx)
|
|
y = F.ones((n_entries,), dtype=F.float32, ctx=ctx)
|
|
dat = F.cat([x, y], dim=0)
|
|
inc, shuffle_idx = F.sparse_matrix(dat, ("coo", idx), (n, m))
|
|
else:
|
|
raise DGLError("Invalid incidence matrix type: %s" % str(typestr))
|
|
return inc, shuffle_idx
|
|
|
|
def node_subgraph(self, induced_nodes):
|
|
"""Return the induced node subgraph.
|
|
|
|
Parameters
|
|
----------
|
|
induced_nodes : list of utils.Index
|
|
Induced nodes. The length should be equal to the number of
|
|
node types in this heterograph.
|
|
|
|
Returns
|
|
-------
|
|
SubgraphIndex
|
|
The subgraph index.
|
|
"""
|
|
vids = [F.to_dgl_nd(nodes) for nodes in induced_nodes]
|
|
return _CAPI_DGLHeteroVertexSubgraph(self, vids)
|
|
|
|
def edge_subgraph(self, induced_edges, preserve_nodes):
|
|
"""Return the induced edge subgraph.
|
|
|
|
Parameters
|
|
----------
|
|
induced_edges : list of utils.Index
|
|
Induced edges. The length should be equal to the number of
|
|
edge types in this heterograph.
|
|
preserve_nodes : bool
|
|
Indicates whether to preserve all nodes or not.
|
|
If true, keep the nodes which have no edge connected in the subgraph;
|
|
If false, all nodes without edge connected to it would be removed.
|
|
|
|
Returns
|
|
-------
|
|
SubgraphIndex
|
|
The subgraph index.
|
|
"""
|
|
eids = [F.to_dgl_nd(edges) for edges in induced_edges]
|
|
return _CAPI_DGLHeteroEdgeSubgraph(self, eids, preserve_nodes)
|
|
|
|
def get_unitgraph(self, etype, ctx):
|
|
"""Create a unitgraph graph from given edge type and copy to the given device
|
|
context.
|
|
|
|
Note: this internal function is for DGL scheduler use only
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
If the graph index is a Bipartite graph index, this argument must be None.
|
|
Otherwise, it represents the edge type.
|
|
ctx : DGLContext
|
|
The context of the returned graph.
|
|
|
|
Returns
|
|
-------
|
|
HeteroGraphIndex
|
|
"""
|
|
g = self.get_relation_graph(etype)
|
|
return g.copy_to(ctx).asbits(self.bits_needed(etype or 0))
|
|
|
|
def get_csr_shuffle_order(self, etype):
|
|
"""Return the edge shuffling order when a coo graph is converted to csr format
|
|
|
|
Parameters
|
|
----------
|
|
etype : int
|
|
The edge type
|
|
|
|
Returns
|
|
-------
|
|
tuple of two utils.Index
|
|
The first element of the tuple is the shuffle order for outward graph
|
|
The second element of the tuple is the shuffle order for inward graph
|
|
"""
|
|
csr = _CAPI_DGLHeteroGetAdj(self, int(etype), False, "csr")
|
|
order = csr(2)
|
|
rev_csr = _CAPI_DGLHeteroGetAdj(self, int(etype), True, "csr")
|
|
rev_order = rev_csr(2)
|
|
return utils.toindex(order, self.dtype), utils.toindex(
|
|
rev_order, self.dtype
|
|
)
|
|
|
|
def formats(self, formats=None):
|
|
"""Get a graph index with the specified allowed sparse format(s) or
|
|
query for the usage status of sparse formats.
|
|
|
|
If the graph has multiple edge types, they will have the same
|
|
sparse format.
|
|
|
|
When ``formats`` is not None, if the intersection between `formats` and
|
|
the current graph's created sparse format(s) is not empty, the returned
|
|
cloned graph only retains all sparse format(s) in the intersection. If
|
|
the intersection is empty, a sparse format will be selected to be
|
|
created following the order of ``'coo' -> 'csr' -> 'csc'``.
|
|
|
|
Parameters
|
|
----------
|
|
formats : str or list of str or None
|
|
|
|
* If formats is None, return the usage status of sparse formats
|
|
* Otherwise, it can be ``'coo'``/``'csr'``/``'csc'`` or a sublist of
|
|
them, specifying the sparse formats to use.
|
|
|
|
Returns
|
|
-------
|
|
dict or GraphIndex
|
|
|
|
* If formats is None, the result will be a dict recording the usage
|
|
status of sparse formats.
|
|
* Otherwise, a GraphIndex will be returned, which is a clone of the
|
|
original graph with the specified allowed sparse format(s)
|
|
``formats``.
|
|
|
|
"""
|
|
formats_allowed = _CAPI_DGLHeteroGetAllowedFormats(self)
|
|
formats_created = _CAPI_DGLHeteroGetCreatedFormats(self)
|
|
created = []
|
|
not_created = []
|
|
if formats is None:
|
|
for fmt in ["coo", "csr", "csc"]:
|
|
if fmt in formats_allowed:
|
|
if fmt in formats_created:
|
|
created.append(fmt)
|
|
else:
|
|
not_created.append(fmt)
|
|
return {"created": created, "not created": not_created}
|
|
else:
|
|
if isinstance(formats, str):
|
|
formats = [formats]
|
|
return _CAPI_DGLHeteroGetFormatGraph(self, formats)
|
|
|
|
def create_formats_(self):
|
|
"""Create all sparse matrices allowed for the graph."""
|
|
return _CAPI_DGLHeteroCreateFormat(self)
|
|
|
|
def reverse(self):
|
|
"""Reverse the heterogeneous graph adjacency
|
|
|
|
The node types and edge types are not changed.
|
|
|
|
Returns
|
|
-------
|
|
A new graph index.
|
|
"""
|
|
return _CAPI_DGLHeteroReverse(self)
|
|
|
|
|
|
@register_object("graph.HeteroSubgraph")
|
|
class HeteroSubgraphIndex(ObjectBase):
|
|
"""Hetero-subgraph data structure"""
|
|
|
|
@property
|
|
def graph(self):
|
|
"""The subgraph structure
|
|
|
|
Returns
|
|
-------
|
|
HeteroGraphIndex
|
|
The subgraph
|
|
"""
|
|
return _CAPI_DGLHeteroSubgraphGetGraph(self)
|
|
|
|
@property
|
|
def induced_nodes(self):
|
|
"""Induced nodes for each node type. The return list
|
|
length should be equal to the number of node types.
|
|
|
|
Returns
|
|
-------
|
|
list of utils.Index
|
|
Induced nodes
|
|
"""
|
|
ret = _CAPI_DGLHeteroSubgraphGetInducedVertices(self)
|
|
return [F.from_dgl_nd(v) for v in ret]
|
|
|
|
@property
|
|
def induced_edges(self):
|
|
"""Induced edges for each edge type. The return list
|
|
length should be equal to the number of edge types.
|
|
|
|
Returns
|
|
-------
|
|
list of utils.Index
|
|
Induced edges
|
|
"""
|
|
ret = _CAPI_DGLHeteroSubgraphGetInducedEdges(self)
|
|
return [F.from_dgl_nd(v) for v in ret]
|
|
|
|
|
|
#################################################################
|
|
# Creators
|
|
#################################################################
|
|
|
|
|
|
def create_metagraph_index(ntypes, canonical_etypes):
|
|
"""Return a GraphIndex instance for a metagraph given the node types and canonical
|
|
edge types.
|
|
|
|
This function will reorder the node types and canonical edge types.
|
|
|
|
Parameters
|
|
----------
|
|
ntypes : Iterable[str]
|
|
The node types.
|
|
canonical_etypes : Iterable[tuple[str, str, str]]
|
|
The canonical edge types.
|
|
|
|
Returns
|
|
-------
|
|
GraphIndex
|
|
The index object for metagraph.
|
|
list[str]
|
|
The reordered node types for each node in the metagraph.
|
|
list[str]
|
|
The reordered edge types for each edge in the metagraph.
|
|
list[tuple[str, str, str]]
|
|
The reordered canonical edge types for each edge in the metagraph.
|
|
"""
|
|
# Sort the ntypes and relation tuples to have a deterministic order for the same set
|
|
# of type names.
|
|
ntypes = list(sorted(ntypes))
|
|
relations = list(sorted(canonical_etypes))
|
|
ntype_dict = {ntype: i for i, ntype in enumerate(ntypes)}
|
|
meta_edges_src = []
|
|
meta_edges_dst = []
|
|
etypes = []
|
|
for srctype, etype, dsttype in relations:
|
|
meta_edges_src.append(ntype_dict[srctype])
|
|
meta_edges_dst.append(ntype_dict[dsttype])
|
|
etypes.append(etype)
|
|
# metagraph is DGLGraph, currently still using int64 as index dtype
|
|
metagraph = from_coo(len(ntypes), meta_edges_src, meta_edges_dst, True)
|
|
return metagraph, ntypes, etypes, relations
|
|
|
|
|
|
def create_unitgraph_from_coo(
|
|
num_ntypes,
|
|
num_src,
|
|
num_dst,
|
|
row,
|
|
col,
|
|
formats,
|
|
row_sorted=False,
|
|
col_sorted=False,
|
|
):
|
|
"""Create a unitgraph graph index from COO format
|
|
|
|
Parameters
|
|
----------
|
|
num_ntypes : int
|
|
Number of node types (must be 1 or 2).
|
|
num_src : int
|
|
Number of nodes in the src type.
|
|
num_dst : int
|
|
Number of nodes in the dst type.
|
|
row : utils.Index
|
|
Row index.
|
|
col : utils.Index
|
|
Col index.
|
|
formats : list of str.
|
|
Restrict the storage formats allowed for the unit graph.
|
|
row_sorted : bool, optional
|
|
Whether or not the rows of the COO are in ascending order.
|
|
col_sorted : bool, optional
|
|
Whether or not the columns of the COO are in ascending order within
|
|
each row. This only has an effect when ``row_sorted`` is True.
|
|
|
|
Returns
|
|
-------
|
|
HeteroGraphIndex
|
|
"""
|
|
if isinstance(formats, str):
|
|
formats = [formats]
|
|
return _CAPI_DGLHeteroCreateUnitGraphFromCOO(
|
|
int(num_ntypes),
|
|
int(num_src),
|
|
int(num_dst),
|
|
F.to_dgl_nd(row),
|
|
F.to_dgl_nd(col),
|
|
formats,
|
|
row_sorted,
|
|
col_sorted,
|
|
)
|
|
|
|
|
|
def create_unitgraph_from_csr(
|
|
num_ntypes,
|
|
num_src,
|
|
num_dst,
|
|
indptr,
|
|
indices,
|
|
edge_ids,
|
|
formats,
|
|
transpose=False,
|
|
):
|
|
"""Create a unitgraph graph index from CSR format
|
|
|
|
Parameters
|
|
----------
|
|
num_ntypes : int
|
|
Number of node types (must be 1 or 2).
|
|
num_src : int
|
|
Number of nodes in the src type.
|
|
num_dst : int
|
|
Number of nodes in the dst type.
|
|
indptr : utils.Index
|
|
CSR indptr.
|
|
indices : utils.Index
|
|
CSR indices.
|
|
edge_ids : utils.Index
|
|
Edge shuffle id.
|
|
formats : str
|
|
Restrict the storage formats allowed for the unit graph.
|
|
transpose : bool, optional
|
|
If True, treats the input matrix as CSC.
|
|
|
|
Returns
|
|
-------
|
|
HeteroGraphIndex
|
|
"""
|
|
if isinstance(formats, str):
|
|
formats = [formats]
|
|
return _CAPI_DGLHeteroCreateUnitGraphFromCSR(
|
|
int(num_ntypes),
|
|
int(num_src),
|
|
int(num_dst),
|
|
F.to_dgl_nd(indptr),
|
|
F.to_dgl_nd(indices),
|
|
F.to_dgl_nd(edge_ids),
|
|
formats,
|
|
transpose,
|
|
)
|
|
|
|
|
|
def create_heterograph_from_relations(
|
|
metagraph, rel_graphs, num_nodes_per_type
|
|
):
|
|
"""Create a heterograph from metagraph and graphs of every relation.
|
|
|
|
Parameters
|
|
----------
|
|
metagraph : GraphIndex
|
|
Meta-graph.
|
|
rel_graphs : list of HeteroGraphIndex
|
|
Bipartite graph of each relation.
|
|
num_nodes_per_type : utils.Index, optional
|
|
Number of nodes per node type
|
|
|
|
Returns
|
|
-------
|
|
HeteroGraphIndex
|
|
"""
|
|
if num_nodes_per_type is None:
|
|
return _CAPI_DGLHeteroCreateHeteroGraph(metagraph, rel_graphs)
|
|
else:
|
|
return _CAPI_DGLHeteroCreateHeteroGraphWithNumNodes(
|
|
metagraph, rel_graphs, num_nodes_per_type.todgltensor()
|
|
)
|
|
|
|
|
|
def create_heterograph_from_shared_memory(name):
|
|
"""Create a heterograph from shared memory with the given name.
|
|
|
|
Paramaters
|
|
----------
|
|
name : str
|
|
The name of the share memory
|
|
|
|
Returns
|
|
-------
|
|
HeteroGraphIndex (in shared memory)
|
|
ntypes : list of str
|
|
Names of node types
|
|
etypes : list of str
|
|
Names of edge types
|
|
"""
|
|
g, ntypes, etypes = _CAPI_DGLHeteroCreateFromSharedMem(name)
|
|
return g, list(ntypes), list(etypes)
|
|
|
|
|
|
def joint_union(metagraph, gidx_list):
|
|
"""Return a joint union of the input heterographs.
|
|
|
|
Parameters
|
|
----------
|
|
metagraph : GraphIndex
|
|
Meta-graph.
|
|
gidx_list : list of HeteroGraphIndex
|
|
Heterographs to be joint_unioned.
|
|
|
|
Returns
|
|
-------
|
|
HeteroGraphIndex
|
|
joint_unioned Heterograph.
|
|
"""
|
|
return _CAPI_DGLHeteroJointUnion(metagraph, gidx_list)
|
|
|
|
|
|
def disjoint_union(metagraph, graphs):
|
|
"""Return a disjoint union of the input heterographs.
|
|
|
|
Parameters
|
|
----------
|
|
metagraph : GraphIndex
|
|
Meta-graph.
|
|
graphs : list of HeteroGraphIndex
|
|
Heterographs to be batched.
|
|
|
|
Returns
|
|
-------
|
|
HeteroGraphIndex
|
|
Batched Heterograph.
|
|
"""
|
|
return _CAPI_DGLHeteroDisjointUnion_v2(metagraph, graphs)
|
|
|
|
|
|
def disjoint_partition(graph, bnn_all_types, bne_all_types):
|
|
"""Partition the graph disjointly.
|
|
|
|
Parameters
|
|
----------
|
|
graph : HeteroGraphIndex
|
|
The graph to be partitioned.
|
|
bnn_all_types : list of list of int
|
|
bnn_all_types[t] gives the number of nodes with t-th type in the batch.
|
|
bne_all_types : list of list of int
|
|
bne_all_types[t] gives the number of edges with t-th type in the batch.
|
|
|
|
Returns
|
|
--------
|
|
list of HeteroGraphIndex
|
|
Heterographs unbatched.
|
|
"""
|
|
bnn_all_types = utils.toindex(
|
|
list(itertools.chain.from_iterable(bnn_all_types))
|
|
)
|
|
bne_all_types = utils.toindex(
|
|
list(itertools.chain.from_iterable(bne_all_types))
|
|
)
|
|
return _CAPI_DGLHeteroDisjointPartitionBySizes_v2(
|
|
graph, bnn_all_types.todgltensor(), bne_all_types.todgltensor()
|
|
)
|
|
|
|
|
|
def slice_gidx(graph, num_nodes, start_nid, num_edges, start_eid):
|
|
"""Slice a chunk of the graph.
|
|
|
|
Parameters
|
|
----------
|
|
graph : HeteroGraphIndex
|
|
The batched graph to slice.
|
|
num_nodes : utils.Index
|
|
Number of nodes per node type in the result graph.
|
|
start_nid : utils.Index
|
|
Start node ID per node type in the result graph.
|
|
num_edges : utils.Index
|
|
Number of edges per edge type in the result graph.
|
|
start_eid : utils.Index
|
|
Start edge ID per edge type in the result graph.
|
|
|
|
Returns
|
|
-------
|
|
HeteroGraphIndex
|
|
The sliced graph.
|
|
"""
|
|
return _CAPI_DGLHeteroSlice(
|
|
graph,
|
|
num_nodes.todgltensor(),
|
|
start_nid.todgltensor(),
|
|
num_edges.todgltensor(),
|
|
start_eid.todgltensor(),
|
|
)
|
|
|
|
|
|
#################################################################
|
|
# Data structure used by C APIs
|
|
#################################################################
|
|
|
|
|
|
@register_object("graph.FlattenedHeteroGraph")
|
|
class FlattenedHeteroGraph(ObjectBase):
|
|
"""FlattenedHeteroGraph object class in C++ backend."""
|
|
|
|
|
|
@register_object("graph.HeteroPickleStates")
|
|
class HeteroPickleStates(ObjectBase):
|
|
"""Pickle states object class in C++ backend."""
|
|
|
|
@property
|
|
def version(self):
|
|
"""Version number
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
version number
|
|
"""
|
|
return _CAPI_DGLHeteroPickleStatesGetVersion(self)
|
|
|
|
@property
|
|
def meta(self):
|
|
"""Meta info
|
|
|
|
Returns
|
|
-------
|
|
bytearray
|
|
Serialized meta info
|
|
"""
|
|
return bytearray(_CAPI_DGLHeteroPickleStatesGetMeta(self))
|
|
|
|
@property
|
|
def arrays(self):
|
|
"""Arrays representing the graph structure (COO or CSR)
|
|
|
|
Returns
|
|
-------
|
|
list of dgl.ndarray.NDArray
|
|
Arrays
|
|
"""
|
|
num_arr = _CAPI_DGLHeteroPickleStatesGetArraysNum(self)
|
|
arr_func = _CAPI_DGLHeteroPickleStatesGetArrays(self)
|
|
return [arr_func(i) for i in range(num_arr)]
|
|
|
|
def __getstate__(self):
|
|
"""Issue: https://github.com/pytorch/pytorch/issues/32351
|
|
Need to set the tensor created in the __getstate__ function
|
|
as object attribute to avoid potential bugs
|
|
"""
|
|
self._pk_arrays = [
|
|
F.zerocopy_from_dgl_ndarray(arr) for arr in self.arrays
|
|
]
|
|
return self.version, self.meta, self._pk_arrays
|
|
|
|
def __setstate__(self, state):
|
|
if isinstance(state[0], int):
|
|
version, meta, arrays = state
|
|
arrays = [F.zerocopy_to_dgl_ndarray(arr) for arr in arrays]
|
|
self.__init_handle_by_constructor__(
|
|
_CAPI_DGLCreateHeteroPickleStates, version, meta, arrays
|
|
)
|
|
else:
|
|
metagraph, num_nodes_per_type, adjs = state
|
|
num_nodes_per_type = F.zerocopy_to_dgl_ndarray(num_nodes_per_type)
|
|
self.__init_handle_by_constructor__(
|
|
_CAPI_DGLCreateHeteroPickleStatesOld,
|
|
metagraph,
|
|
num_nodes_per_type,
|
|
adjs,
|
|
)
|
|
|
|
|
|
def _forking_rebuild(pk_state):
|
|
version, meta, arrays = pk_state
|
|
arrays = [F.to_dgl_nd(arr) for arr in arrays]
|
|
states = _CAPI_DGLCreateHeteroPickleStates(version, meta, arrays)
|
|
graph_index = _CAPI_DGLHeteroForkingUnpickle(states)
|
|
graph_index._forking_pk_state = pk_state
|
|
return graph_index
|
|
|
|
|
|
def _forking_reduce(graph_index):
|
|
# Because F.from_dgl_nd(F.to_dgl_nd(x)) loses the information of shared memory
|
|
# file descriptor (because DLPack does not keep it), without caching the tensors
|
|
# PyTorch will allocate one shared memory region for every single worker.
|
|
# The downside is that if a graph_index is shared by forking and new formats are created
|
|
# afterwards, then sharing it again will not bring together the new formats. This case
|
|
# should be rare though because (1) DataLoader will create all the formats if num_workers > 0
|
|
# anyway, and (2) we require the users to explicitly create all formats before calling
|
|
# mp.spawn().
|
|
if hasattr(graph_index, "_forking_pk_state"):
|
|
return _forking_rebuild, (graph_index._forking_pk_state,)
|
|
states = _CAPI_DGLHeteroForkingPickle(graph_index)
|
|
arrays = [F.from_dgl_nd(arr) for arr in states.arrays]
|
|
# Similar to what being mentioned in HeteroGraphIndex.__getstate__, we need to save
|
|
# the tensors as an attribute of the original graph index object. Otherwise
|
|
# PyTorch will throw weird errors like bad value(s) in fds_to_keep or unable to
|
|
# resize file.
|
|
graph_index._forking_pk_state = (states.version, states.meta, arrays)
|
|
return _forking_rebuild, (graph_index._forking_pk_state,)
|
|
|
|
|
|
if not (F.get_preferred_backend() == "mxnet" and sys.version_info.minor <= 6):
|
|
# Python 3.6 MXNet crashes with the following statement; remove until we no longer support
|
|
# 3.6 (which is EOL anyway).
|
|
from multiprocessing.reduction import ForkingPickler
|
|
|
|
ForkingPickler.register(HeteroGraphIndex, _forking_reduce)
|
|
|
|
_init_api("dgl.heterograph_index")
|