项目文件夹

文件
Zihao Ye af61e2fbb4 [Feature] Support nn modules for bipartite graphs. (#1392)
* init gat

* fix

* gin

* 7 nn modules

* rename & lint

* upd

* upd

* fix lint

* upd test

* upd

* lint

* shape check

* upd

* lint

* address comments

* update tensorflow

Co-authored-by: Quan Gan <coin2028@hotmail.com>
Co-authored-by: Jinjing Zhou <VoVAllen@users.noreply.github.com>
Co-authored-by: Minjie Wang <wmjlyjemaine@gmail.com>
2020-03-29 16:48:38 +08:00

4436 行
157 KiB
Python

"""Classes for heterogeneous graphs."""
#pylint: disable= too-many-lines
from collections import defaultdict
from contextlib import contextmanager
import networkx as nx
import numpy as np
from . import graph_index
from . import heterograph_index
from . import utils
from . import backend as F
from . import init
from .runtime import ir, scheduler, Runtime, GraphAdapter
from .frame import Frame, FrameRef, frame_like, sync_frame_initializer
from .view import HeteroNodeView, HeteroNodeDataView, HeteroEdgeView, HeteroEdgeDataView
from .base import ALL, SLICE_FULL, NTYPE, NID, ETYPE, EID, is_all, DGLError, dgl_warning
from .udf import NodeBatch, EdgeBatch
__all__ = ['DGLHeteroGraph', 'combine_names']
class DGLHeteroGraph(object):
"""Base heterogeneous graph class.
**Do NOT instantiate from this class directly; use** :mod:`conversion methods
<dgl.convert>` **instead.**
A Heterogeneous graph is defined as a graph with node types and edge
types.
If two edges share the same edge type, then their source nodes, as well
as their destination nodes, also have the same type (the source node
types don't have to be the same as the destination node types).
Examples
--------
Suppose that we want to construct the following heterogeneous graph:
.. graphviz::
digraph G {
Alice -> Bob [label=follows]
Bob -> Carol [label=follows]
Alice -> Tetris [label=plays]
Bob -> Tetris [label=plays]
Bob -> Minecraft [label=plays]
Carol -> Minecraft [label=plays]
Nintendo -> Tetris [label=develops]
Mojang -> Minecraft [label=develops]
{rank=source; Alice; Bob; Carol}
{rank=sink; Nintendo; Mojang}
}
And suppose that one maps the users, games and developers to the following
IDs:
========= ===== === =====
User name Alice Bob Carol
========= ===== === =====
User ID 0 1 2
========= ===== === =====
========= ====== =========
Game name Tetris Minecraft
========= ====== =========
Game ID 0 1
========= ====== =========
============== ======== ======
Developer name Nintendo Mojang
============== ======== ======
Developer ID 0 1
============== ======== ======
One can construct the graph as follows:
>>> follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
>>> plays_g = dgl.bipartite([(0, 0), (1, 0), (1, 1), (2, 1)], 'user', 'plays', 'game')
>>> devs_g = dgl.bipartite([(0, 0), (1, 1)], 'developer', 'develops', 'game')
>>> g = dgl.hetero_from_relations([follows_g, plays_g, devs_g])
Or equivalently
>>> g = dgl.heterograph({
... ('user', 'follows', 'user'): [(0, 1), (1, 2)],
... ('user', 'plays', 'game'): [(0, 0), (1, 0), (1, 1), (2, 1)],
... ('developer', 'develops', 'game'): [(0, 0), (1, 1)],
... })
:func:`dgl.graph` and :func:`dgl.bipartite` can create a graph from a variety of
data types including:
* edge list
* edge tuples
* networkx graph
* scipy sparse matrix
Click the function names for more details.
Then one can query the graph structure by specifying the ``ntype`` or ``etype`` arguments:
>>> g.number_of_nodes('user')
3
>>> g.number_of_edges('plays')
4
>>> g.out_degrees(etype='develops') # out-degrees of source nodes of 'develops' relation
tensor([1, 1])
>>> g.in_edges(0, etype='develops') # in-edges of destination node 0 of 'develops' relation
(tensor([0]), tensor([0]))
Or on the sliced graph for an edge type:
>>> g['plays'].number_of_edges()
4
>>> g['develops'].out_degrees()
tensor([1, 1])
>>> g['develops'].in_edges(0)
(tensor([0]), tensor([0]))
Node type names must be distinct (no two types have the same name). Edge types could
have the same name but they must be distinguishable by the ``(src_type, edge_type, dst_type)``
triplet (called *canonical edge type*).
For example, suppose a graph that has two types of relation "user-watches-movie"
and "user-watches-TV" as follows:
>>> g0 = dgl.bipartite([(0, 1), (1, 0), (1, 1)], 'user', 'watches', 'movie')
>>> g1 = dgl.bipartite([(0, 0), (1, 1)], 'user', 'watches', 'TV')
>>> GG = dgl.hetero_from_relations([g0, g1]) # Merge the two graphs
To distinguish between the two "watches" edge type, one must specify a full triplet:
>>> GG.number_of_edges(('user', 'watches', 'movie'))
3
>>> GG.number_of_edges(('user', 'watches', 'TV'))
2
>>> GG['user', 'watches', 'movie'].out_degrees()
tensor([1, 2])
Using only one single edge type string "watches" is ambiguous and will cause error:
>>> GG.number_of_edges('watches') # AMBIGUOUS!!
In many cases, there is only one type of nodes or one type of edges, and the ``ntype``
and ``etype`` argument could be omitted. This is very common when using the sliced
graph, which usually contains only one edge type, and sometimes only one node type:
>>> g['follows'].number_of_nodes() # OK!! because g['follows'] only has one node type 'user'
3
>>> g['plays'].number_of_nodes() # ERROR!! There are two types 'user' and 'game'.
>>> g['plays'].number_of_edges() # OK!! because there is only one edge type 'plays'
TODO(minjie): docstring about uni-directional bipartite graph
Metagraph
---------
For each heterogeneous graph, one can often infer the *metagraph*, the template of
edge connections showing how many types of nodes and edges exist in the graph, and
how each edge type could connect between node types.
One can analyze the example gameplay graph above and figure out the metagraph as
follows:
.. graphviz::
digraph G {
User -> User [label=follows]
User -> Game [label=plays]
Developer -> Game [label=develops]
}
Parameters
----------
gidx : HeteroGraphIndex
Graph index object.
ntypes : list of str, pair of list of str
Node type list. ``ntypes[i]`` stores the name of node type i.
If a pair is given, the graph created is a uni-directional bipartite graph,
and its SRC node types and DST node types are given as in the pair.
etypes : list of str
Edge type list. ``etypes[i]`` stores the name of edge type i.
node_frames : list of FrameRef, optional
Node feature storage. If None, empty frame is created.
Otherwise, ``node_frames[i]`` stores the node features
of node type i. (default: None)
edge_frames : list of FrameRef, optional
Edge feature storage. If None, empty frame is created.
Otherwise, ``edge_frames[i]`` stores the edge features
of edge type i. (default: None)
"""
# pylint: disable=unused-argument
def __init__(self,
gidx,
ntypes,
etypes,
node_frames=None,
edge_frames=None):
self._init(gidx, ntypes, etypes, node_frames, edge_frames)
def _init(self, gidx, ntypes, etypes, node_frames, edge_frames):
"""Init internal states."""
self._graph = gidx
# Handle node types
if isinstance(ntypes, tuple):
if len(ntypes) != 2:
errmsg = 'Invalid input. Expect a pair (srctypes, dsttypes) but got {}'.format(
ntypes)
raise TypeError(errmsg)
if not is_unibipartite(self._graph.metagraph):
raise ValueError('Invalid input. The metagraph must be a uni-directional'
' bipartite graph.')
self._ntypes = ntypes[0] + ntypes[1]
self._srctypes_invmap = {t : i for i, t in enumerate(ntypes[0])}
self._dsttypes_invmap = {t : i + len(ntypes[0]) for i, t in enumerate(ntypes[1])}
self._is_unibipartite = True
else:
self._ntypes = ntypes
src_dst_map = find_src_dst_ntypes(self._ntypes, self._graph.metagraph)
self._is_unibipartite = (src_dst_map is not None)
if self._is_unibipartite:
self._srctypes_invmap, self._dsttypes_invmap = src_dst_map
else:
self._srctypes_invmap = {t : i for i, t in enumerate(self._ntypes)}
self._dsttypes_invmap = self._srctypes_invmap
# Handle edge types
self._etypes = etypes
self._canonical_etypes = make_canonical_etypes(
self._etypes, self._ntypes, self._graph.metagraph)
# An internal map from etype to canonical etype tuple.
# If two etypes have the same name, an empty tuple is stored instead to indicate
# ambiguity.
self._etype2canonical = {}
for i, ety in enumerate(self._etypes):
if ety in self._etype2canonical:
self._etype2canonical[ety] = tuple()
else:
self._etype2canonical[ety] = self._canonical_etypes[i]
self._etypes_invmap = {t : i for i, t in enumerate(self._canonical_etypes)}
# Cached metagraph in networkx
self._nx_metagraph = None
# node and edge frame
if node_frames is None:
node_frames = [None] * len(self._ntypes)
node_frames = [FrameRef(Frame(num_rows=self._graph.number_of_nodes(i)))
if frame is None else frame
for i, frame in enumerate(node_frames)]
self._node_frames = node_frames
if edge_frames is None:
edge_frames = [None] * len(self._etypes)
edge_frames = [FrameRef(Frame(num_rows=self._graph.number_of_edges(i)))
if frame is None else frame
for i, frame in enumerate(edge_frames)]
self._edge_frames = edge_frames
# message indicators
self._msg_indices = [None] * len(self._etypes)
self._msg_frames = []
for i in range(len(self._etypes)):
frame = FrameRef(Frame(num_rows=self._graph.number_of_edges(i)))
frame.set_initializer(init.zero_initializer)
self._msg_frames.append(frame)
def __getstate__(self):
return self._graph, self._ntypes, self._etypes, self._node_frames, self._edge_frames
def __setstate__(self, state):
self._init(*state)
def _get_msg_index(self, etid):
"""Internal function for getting the message index array of the given edge type id."""
if self._msg_indices[etid] is None:
self._msg_indices[etid] = utils.zero_index(
size=self._graph.number_of_edges(etid))
return self._msg_indices[etid]
def _set_msg_index(self, etid, index):
self._msg_indices[etid] = index
def __repr__(self):
if len(self.ntypes) == 1 and len(self.etypes) == 1:
ret = ('Graph(num_nodes={node}, num_edges={edge},\n'
' ndata_schemes={ndata}\n'
' edata_schemes={edata})')
return ret.format(node=self.number_of_nodes(), edge=self.number_of_edges(),
ndata=str(self.node_attr_schemes()),
edata=str(self.edge_attr_schemes()))
else:
ret = ('Graph(num_nodes={node},\n'
' num_edges={edge},\n'
' metagraph={meta})')
nnode_dict = {self.ntypes[i] : self._graph.number_of_nodes(i)
for i in range(len(self.ntypes))}
nedge_dict = {self.canonical_etypes[i] : self._graph.number_of_edges(i)
for i in range(len(self.etypes))}
meta = str(self.metagraph.edges())
return ret.format(node=nnode_dict, edge=nedge_dict, meta=meta)
#################################################################
# Mutation operations
#################################################################
def add_nodes(self, num, data=None, ntype=None):
"""Add multiple new nodes of the same node type
Currently not supported.
"""
raise DGLError('Mutation is not supported in heterograph.')
def add_edge(self, u, v, data=None, etype=None):
"""Add an edge of ``etype`` between u of the source node type, and v
of the destination node type..
Currently not supported.
"""
raise DGLError('Mutation is not supported in heterograph.')
def add_edges(self, u, v, data=None, etype=None):
"""Add multiple edges of ``etype`` between list of source nodes ``u``
and list of destination nodes ``v`` of type ``vtype``. A single edge
is added between every pair of ``u[i]`` and ``v[i]``.
Currently not supported.
"""
raise DGLError('Mutation is not supported in heterograph.')
#################################################################
# Metagraph query
#################################################################
@property
def is_unibipartite(self):
"""Return whether the graph is a uni-bipartite graph.
A uni-bipartite heterograph can further divide its node types into two sets:
SRC and DST. All edges are from nodes in SRC to nodes in DST. The following APIs
can be used to get the nodes and types that belong to SRC and DST sets:
* :func:`srctype` and :func:`dsttype`
* :func:`srcdata` and :func:`dstdata`
* :func:`srcnodes` and :func:`dstnodes`
Note that we allow two node types to have the same name as long as one
belongs to SRC while the other belongs to DST. To distinguish them, prepend
the name with ``"SRC/"`` or ``"DST/"`` when specifying a node type.
"""
return self._is_unibipartite
@property
def ntypes(self):
"""Return the list of node types of this graph.
Returns
-------
list of str
Examples
--------
>>> follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
>>> plays_g = dgl.bipartite([(0, 0), (1, 0), (1, 1), (2, 1)], 'user', 'plays', 'game')
>>> g = dgl.hetero_from_relations([follows_g, plays_g])
>>> g.ntypes
['user', 'game']
"""
return self._ntypes
@property
def etypes(self):
"""Return the list of edge types of this graph.
Returns
-------
list of str
Examples
--------
>>> follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
>>> plays_g = dgl.bipartite([(0, 0), (1, 0), (1, 1), (2, 1)], 'user', 'plays', 'game')
>>> g = dgl.hetero_from_relations([follows_g, plays_g])
>>> g.etypes
['follows', 'plays']
"""
return self._etypes
@property
def canonical_etypes(self):
"""Return the list of canonical edge types of this graph.
A canonical edge type is a tuple of string (src_type, edge_type, dst_type).
Returns
-------
list of 3-tuples
Examples
--------
>>> follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
>>> plays_g = dgl.bipartite([(0, 0), (1, 0), (1, 1), (2, 1)], 'user', 'plays', 'game')
>>> g = dgl.hetero_from_relations([follows_g, plays_g])
>>> g.canonical_etypes
[('user', 'follows', 'user'), ('user', 'plays', 'game')]
"""
return self._canonical_etypes
@property
def srctypes(self):
"""Return the node types in the SRC category. Return :attr:``ntypes`` if
the graph is not a uni-bipartite graph.
"""
if self.is_unibipartite:
return sorted(list(self._srctypes_invmap.keys()))
else:
return self.ntypes
@property
def dsttypes(self):
"""Return the node types in the DST category. Return :attr:``ntypes`` if
the graph is not a uni-bipartite graph.
"""
if self.is_unibipartite:
return sorted(list(self._dsttypes_invmap.keys()))
else:
return self.ntypes
@property
def metagraph(self):
"""Return the metagraph as networkx.MultiDiGraph.
The nodes are labeled with node type names.
The edges have their keys holding the edge type names.
Returns
-------
networkx.MultiDiGraph
Examples
--------
>>> follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
>>> plays_g = dgl.bipartite([(0, 0), (1, 0), (1, 1), (2, 1)], 'user', 'plays', 'game')
>>> g = dgl.hetero_from_relations([follows_g, plays_g])
>>> meta_g = g.metagraph
The metagraph then has two nodes and two edges.
>>> meta_g.nodes()
NodeView(('user', 'game'))
>>> meta_g.number_of_nodes()
2
>>> meta_g.edges()
OutMultiEdgeDataView([('user', 'user'), ('user', 'game')])
>>> meta_g.number_of_edges()
2
"""
if self._nx_metagraph is None:
nx_graph = self._graph.metagraph.to_networkx()
self._nx_metagraph = nx.MultiDiGraph()
for u_v in nx_graph.edges:
srctype, etype, dsttype = self.canonical_etypes[nx_graph.edges[u_v]['id']]
self._nx_metagraph.add_edge(srctype, dsttype, etype)
return self._nx_metagraph
def to_canonical_etype(self, etype):
"""Convert edge type to canonical etype: (srctype, etype, dsttype).
The input can already be a canonical tuple.
Parameters
----------
etype : str or tuple of str
Edge type
Returns
-------
tuple of str
Examples
--------
Instantiate a heterograph.
>>> g1 = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
>>> g2 = dgl.bipartite([(0, 0), (1, 0), (1, 1), (2, 1)], 'user', 'plays', 'game')
>>> g3 = dgl.bipartite([(0, 0), (1, 1)], 'developer', 'follows', 'game')
>>> g = dgl.hetero_from_relations([g1, g2, g3])
Get canonical edge types.
>>> g.to_canonical_etype('plays')
('user', 'plays', 'game')
>>> g.to_canonical_etype(('user', 'plays', 'game'))
('user', 'plays', 'game')
>>> g.to_canonical_etype('follows')
DGLError: Edge type "follows" is ambiguous.
Please use canonical etype type in the form of (srctype, etype, dsttype)
"""
if isinstance(etype, tuple):
return etype
else:
ret = self._etype2canonical.get(etype, None)
if ret is None:
raise DGLError('Edge type "{}" does not exist.'.format(etype))
if len(ret) == 0:
raise DGLError('Edge type "%s" is ambiguous. Please use canonical etype '
'type in the form of (srctype, etype, dsttype)' % etype)
return ret
def get_ntype_id(self, ntype):
"""Return the id of the given node type.
ntype can also be None. If so, there should be only one node type in the
graph.
Parameters
----------
ntype : str
Node type
Returns
-------
int
"""
if self.is_unibipartite and ntype is not None:
# Only check 'SRC/' and 'DST/' prefix when is_unibipartite graph is True.
if ntype.startswith('SRC/'):
return self.get_ntype_id_from_src(ntype[4:])
elif ntype.startswith('DST/'):
return self.get_ntype_id_from_dst(ntype[4:])
# If there is no prefix, fallback to normal lookup.
# Lookup both SRC and DST
if ntype is None:
if self.is_unibipartite or len(self._srctypes_invmap) != 1:
raise DGLError('Node type name must be specified if there are more than one '
'node types.')
return 0
ntid = self._srctypes_invmap.get(ntype, self._dsttypes_invmap.get(ntype, None))
if ntid is None:
raise DGLError('Node type "{}" does not exist.'.format(ntype))
return ntid
def get_ntype_id_from_src(self, ntype):
"""Return the id of the given SRC node type.
ntype can also be None. If so, there should be only one node type in the
SRC category. Callable even when the self graph is not uni-bipartite.
Parameters
----------
ntype : str
Node type
Returns
-------
int
"""
if ntype is None:
if len(self._srctypes_invmap) != 1:
raise DGLError('SRC node type name must be specified if there are more than one '
'SRC node types.')
return next(iter(self._srctypes_invmap.values()))
ntid = self._srctypes_invmap.get(ntype, None)
if ntid is None:
raise DGLError('SRC node type "{}" does not exist.'.format(ntype))
return ntid
def get_ntype_id_from_dst(self, ntype):
"""Return the id of the given DST node type.
ntype can also be None. If so, there should be only one node type in the
DST category. Callable even when the self graph is not uni-bipartite.
Parameters
----------
ntype : str
Node type
Returns
-------
int
"""
if ntype is None:
if len(self._dsttypes_invmap) != 1:
raise DGLError('DST node type name must be specified if there are more than one '
'DST node types.')
return next(iter(self._dsttypes_invmap.values()))
ntid = self._dsttypes_invmap.get(ntype, None)
if ntid is None:
raise DGLError('DST node type "{}" does not exist.'.format(ntype))
return ntid
def get_etype_id(self, etype):
"""Return the id of the given edge type.
etype can also be None. If so, there should be only one edge type in the
graph.
Parameters
----------
etype : str or tuple of str
Edge type
Returns
-------
int
"""
if etype is None:
if self._graph.number_of_etypes() != 1:
raise DGLError('Edge type name must be specified if there are more than one '
'edge types.')
return 0
etid = self._etypes_invmap.get(self.to_canonical_etype(etype), None)
if etid is None:
raise DGLError('Edge type "{}" does not exist.'.format(etype))
return etid
#################################################################
# View
#################################################################
@property
def nodes(self):
"""Return a node view that can be used to set/get feature
data of a single node type.
Examples
--------
The following example uses PyTorch backend.
To set features of all users
>>> g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
>>> g.nodes['user'].data['h'] = torch.zeros(3, 5)
See Also
--------
ndata
"""
return HeteroNodeView(self, self.get_ntype_id)
@property
def srcnodes(self):
"""Return a SRC node view that can be used to set/get feature
data of a single node type.
Examples
--------
The following example uses PyTorch backend.
To set features of all users
>>> g = dgl.biparite([(0, 1), (1, 2)], 'user', 'plays', 'game')
>>> g.srcnodes['user'].data['h'] = torch.zeros(2, 5)
See Also
--------
srcdata
"""
return HeteroNodeView(self, self.get_ntype_id_from_src)
@property
def dstnodes(self):
"""Return a DST node view that can be used to set/get feature
data of a single node type.
Examples
--------
The following example uses PyTorch backend.
To set features of all games
>>> g = dgl.biparite([(0, 1), (1, 2)], 'user', 'plays', 'game')
>>> g.dstnodes['game'].data['h'] = torch.zeros(3, 5)
See Also
--------
dstdata
"""
return HeteroNodeView(self, self.get_ntype_id_from_dst)
@property
def ndata(self):
"""Return the data view of all the nodes.
**Only works if the graph has one node type.**
Examples
--------
The following example uses PyTorch backend.
To set features of all nodes in a heterogeneous graph
with only one node type:
>>> g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
>>> g.ndata['h'] = torch.zeros(3, 5)
See Also
--------
nodes
"""
ntid = self.get_ntype_id(None)
ntype = self.ntypes[0]
return HeteroNodeDataView(self, ntype, ntid, ALL)
@property
def srcdata(self):
"""Return the data view of all nodes in the SRC category.
Only works if the graph is either
* Uni-bipartite and has one node type in the SRC category.
* Non-uni-bipartite and has only one node type (in this case identical to
:any:`DGLHeteroGraph.ndata`)
Examples
--------
The following example uses PyTorch backend.
To set features of all source nodes in a graph with only one edge type:
>>> g = dgl.bipartite([(0, 1), (1, 2)], 'user', 'plays', 'game')
>>> g.srcdata['h'] = torch.zeros(2, 5)
This is equivalent to
>>> g.nodes['user'].data['h'] = torch.zeros(2, 5)
Also work on more complex uni-bipartite graph
>>> g = dgl.heterograph({
... ('user', 'plays', 'game'), [(0, 1), (1, 2)],
... ('user', 'reads', 'book'), [(0, 1), (1, 0)],
... })
>>> print(g.is_unibipartite)
True
>>> g.srcdata['h'] = torch.zeros(2, 5)
Notes
-----
This is identical to :any:`DGLHeteroGraph.ndata` if the graph is homogeneous.
See Also
--------
nodes
"""
err_msg = (
'srcdata is only allowed when there is only one %s type.' %
('SRC' if self.is_unibipartite else 'node'))
assert len(self.srctypes) == 1, err_msg
ntype = self.srctypes[0]
ntid = self.get_ntype_id_from_src(ntype)
return HeteroNodeDataView(self, ntype, ntid, ALL)
@property
def dstdata(self):
"""Return the data view of all destination nodes.
Only works if the graph is either
* Uni-bipartite and has one node type in the SRC category.
* Non-uni-bipartite and has only one node type (in this case identical to
:any:`DGLHeteroGraph.ndata`)
Examples
--------
The following example uses PyTorch backend.
To set features of all source nodes in a graph with only one edge type:
>>> g = dgl.bipartite([(0, 1), (1, 2)], 'user', 'plays', 'game')
>>> g.dstdata['h'] = torch.zeros(3, 5)
This is equivalent to
>>> g.nodes['game'].data['h'] = torch.zeros(3, 5)
Also work on more complex uni-bipartite graph
>>> g = dgl.heterograph({
... ('user', 'plays', 'game'), [(0, 1), (1, 2)],
... ('store', 'sells', 'game'), [(0, 1), (1, 0)],
... })
>>> print(g.is_unibipartite)
True
>>> g.dstdata['h'] = torch.zeros(3, 5)
Notes
-----
This is identical to :any:`DGLHeteroGraph.ndata` if the graph is homogeneous.
See Also
--------
nodes
"""
err_msg = (
'dstdata is only allowed when there is only one %s type.' %
('DST' if self.is_unibipartite else 'node'))
assert len(self.dsttypes) == 1, err_msg
ntype = self.dsttypes[0]
ntid = self.get_ntype_id_from_dst(ntype)
return HeteroNodeDataView(self, ntype, ntid, ALL)
@property
def edges(self):
"""Return an edge view that can be used to set/get feature
data of a single edge type.
Examples
--------
The following example uses PyTorch backend.
To set features of all "play" relationships:
>>> g = dgl.bipartite([(0, 0), (1, 0), (1, 2)], 'user', 'plays', 'game')
>>> g.edges['plays'].data['h'] = torch.zeros(3, 4)
See Also
--------
edata
"""
return HeteroEdgeView(self)
@property
def edata(self):
"""Return the data view of all the edges.
**Only works if the graph has one edge type.**
Examples
--------
The following example uses PyTorch backend.
To set features of all edges in a heterogeneous graph
with only one edge type:
>>> g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
>>> g.edata['h'] = torch.zeros(2, 5)
See Also
--------
edges
"""
return HeteroEdgeDataView(self, None, ALL)
def _find_etypes(self, key):
etypes = [
i for i, (srctype, etype, dsttype) in enumerate(self._canonical_etypes) if
(key[0] == SLICE_FULL or key[0] == srctype) and
(key[1] == SLICE_FULL or key[1] == etype) and
(key[2] == SLICE_FULL or key[2] == dsttype)]
return etypes
def __getitem__(self, key):
"""Return the relation slice of this graph.
A relation slice is accessed with ``self[srctype, etype, dsttype]``, where
``srctype``, ``etype``, and ``dsttype`` can be either a string or a full
slice (``:``) representing wildcard (i.e. any source/edge/destination type).
A relation slice is a homogeneous (with one node type and one edge type) or
bipartite (with two node types and one edge type) graph, transformed from
the original heterogeneous graph.
If there is only one canonical edge type found, then the returned relation
slice would be a subgraph induced from the original graph. That is, it is
equivalent to ``self.edge_type_subgraph(etype)``. The node and edge features
of the returned graph would be shared with thew original graph.
If there are multiple canonical edge type found, then the source/edge/destination
node types would be a *concatenation* of original node/edge types. The
new source/destination node type would have the concatenation determined by
:func:`dgl.combine_names() <dgl.combine_names>` called on original source/destination
types as its name. The source/destination node would be formed by concatenating the
common features of the original source/destination types, therefore they are not
shared with the original graph. Edge type is similar.
"""
err_msg = "Invalid slice syntax. Use G['etype'] or G['srctype', 'etype', 'dsttype'] " +\
"to get view of one relation type. Use : to slice multiple types (e.g. " +\
"G['srctype', :, 'dsttype'])."
orig_key = key
if not isinstance(key, tuple):
key = (SLICE_FULL, key, SLICE_FULL)
if len(key) != 3:
raise DGLError(err_msg)
etypes = self._find_etypes(key)
if len(etypes) == 0:
raise DGLError('Invalid key "{}". Must be one of the edge types.'.format(orig_key))
if len(etypes) == 1:
# no ambiguity: return the unitgraph itself
srctype, etype, dsttype = self._canonical_etypes[etypes[0]]
stid = self.get_ntype_id_from_src(srctype)
etid = self.get_etype_id((srctype, etype, dsttype))
dtid = self.get_ntype_id_from_dst(dsttype)
new_g = self._graph.get_relation_graph(etid)
if stid == dtid:
new_ntypes = [srctype]
new_nframes = [self._node_frames[stid]]
else:
new_ntypes = [srctype, dsttype]
new_nframes = [self._node_frames[stid], self._node_frames[dtid]]
new_etypes = [etype]
new_eframes = [self._edge_frames[etid]]
return DGLHeteroGraph(new_g, new_ntypes, new_etypes, new_nframes, new_eframes)
else:
flat = self._graph.flatten_relations(etypes)
new_g = flat.graph
# merge frames
stids = flat.induced_srctype_set.asnumpy()
dtids = flat.induced_dsttype_set.asnumpy()
etids = flat.induced_etype_set.asnumpy()
new_ntypes = [combine_names(self.ntypes, stids)]
if new_g.number_of_ntypes() == 2:
new_ntypes.append(combine_names(self.ntypes, dtids))
new_nframes = [
combine_frames(self._node_frames, stids),
combine_frames(self._node_frames, dtids)]
else:
assert np.array_equal(stids, dtids)
new_nframes = [combine_frames(self._node_frames, stids)]
new_etypes = [combine_names(self.etypes, etids)]
new_eframes = [combine_frames(self._edge_frames, etids)]
# create new heterograph
new_hg = DGLHeteroGraph(new_g, new_ntypes, new_etypes, new_nframes, new_eframes)
src = new_ntypes[0]
dst = new_ntypes[1] if new_g.number_of_ntypes() == 2 else src
# put the parent node/edge type and IDs
new_hg.nodes[src].data[NTYPE] = F.zerocopy_from_dgl_ndarray(flat.induced_srctype)
new_hg.nodes[src].data[NID] = F.zerocopy_from_dgl_ndarray(flat.induced_srcid)
new_hg.nodes[dst].data[NTYPE] = F.zerocopy_from_dgl_ndarray(flat.induced_dsttype)
new_hg.nodes[dst].data[NID] = F.zerocopy_from_dgl_ndarray(flat.induced_dstid)
new_hg.edata[ETYPE] = F.zerocopy_from_dgl_ndarray(flat.induced_etype)
new_hg.edata[EID] = F.zerocopy_from_dgl_ndarray(flat.induced_eid)
return new_hg
#################################################################
# Graph query
#################################################################
def number_of_nodes(self, ntype=None):
"""Return the number of nodes of the given type in the heterograph.
Parameters
----------
ntype : str, optional
The node type. Can be omitted if there is only one node type
in the graph. (Default: None)
Returns
-------
int
The number of nodes
Examples
--------
>>> g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
>>> g.number_of_nodes('user')
3
>>> g.number_of_nodes()
3
"""
return self._graph.number_of_nodes(self.get_ntype_id(ntype))
def number_of_src_nodes(self, ntype=None):
"""Return the number of nodes of the given SRC node type in the heterograph.
The heterograph is usually a unidirectional bipartite graph.
Parameters
----------
ntype : str, optional
Node type.
If omitted, there should be only one node type in the SRC category.
Returns
-------
int
The number of nodes
Examples
--------
>>> g = dgl.bipartite([(0, 1), (1, 2)], 'user', 'plays', 'game')
>>> g.number_of_src_nodes('user')
2
>>> g.number_of_src_nodes()
2
>>> g.number_of_nodes('user')
2
"""
return self._graph.number_of_nodes(self.get_ntype_id_from_src(ntype))
def number_of_dst_nodes(self, ntype=None):
"""Return the number of nodes of the given DST node type in the heterograph.
The heterograph is usually a unidirectional bipartite graph.
Parameters
----------
ntype : str, optional
Node type.
If omitted, there should be only one node type in the DST category.
Returns
-------
int
The number of nodes
Examples
--------
>>> g = dgl.bipartite([(0, 1), (1, 2)], 'user', 'plays', 'game')
>>> g.number_of_dst_nodes('game')
3
>>> g.number_of_dst_nodes()
3
>>> g.number_of_nodes('game')
3
"""
return self._graph.number_of_nodes(self.get_ntype_id_from_dst(ntype))
def number_of_edges(self, etype=None):
"""Return the number of edges of the given type in the heterograph.
Parameters
----------
etype : str or tuple of str, optional
The edge type. Can be omitted if there is only one edge type
in the graph.
Returns
-------
int
The number of edges
Examples
--------
>>> g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
>>> g.number_of_edges(('user', 'follows', 'user'))
>>> g.number_of_edges('follows')
2
>>> g.number_of_edges()
2
"""
return self._graph.number_of_edges(self.get_etype_id(etype))
@property
def is_multigraph(self):
"""Whether the graph is a multigraph
Returns
-------
bool
True if the graph is a multigraph, False otherwise.
"""
return self._graph.is_multigraph()
@property
def is_readonly(self):
"""Whether the graph is readonly
Returns
-------
bool
True if the graph is readonly, False otherwise.
"""
return self._graph.is_readonly()
def has_node(self, vid, ntype=None):
"""Whether the graph has a node with a particular id and type.
Parameters
----------
vid : int
The node ID.
ntype : str, optional
The node type. Can be omitted if there is only one node type
in the graph. (Default: None)
Returns
-------
bool
True if the node exists, False otherwise
Examples
--------
>>> g.has_node(0, 'user')
True
>>> g.has_node(4, 'user')
False
See Also
--------
has_nodes
"""
return self._graph.has_node(self.get_ntype_id(ntype), vid)
def has_nodes(self, vids, ntype=None):
"""Whether the graph has nodes with ids and a particular type.
Parameters
----------
vid : list or tensor
The array of node IDs.
ntype : str, optional
The node type. Can be omitted if there is only one node type
in the graph.
Returns
-------
a : tensor
Binary tensor indicating the existence of nodes with the specified ids and type.
``a[i]=1`` if the graph contains node ``vids[i]`` of type ``ntype``, 0 otherwise.
Examples
--------
The following example uses PyTorch backend.
>>> g.has_nodes([0, 1, 2, 3, 4], 'user')
tensor([1, 1, 1, 0, 0])
See Also
--------
has_node
"""
vids = utils.toindex(vids)
rst = self._graph.has_nodes(self.get_ntype_id(ntype), vids)
return rst.tousertensor()
def has_edge_between(self, u, v, etype=None):
"""Whether the graph has an edge (u, v) of type ``etype``.
Parameters
----------
u : int
The node ID of source type.
v : int
The node ID of destination type.
etype : str or tuple of str, optional
The edge type. Can be omitted if there is only one edge type
in the graph.
Returns
-------
bool
True if the edge is in the graph, False otherwise.
Examples
--------
>>> g.has_edge_between(0, 1, ('user', 'plays', 'game'))
True
>>> g.has_edge_between(0, 2, ('user', 'plays', 'game'))
False
See Also
--------
has_edges_between
"""
return self._graph.has_edge_between(self.get_etype_id(etype), u, v)
def has_edges_between(self, u, v, etype=None):
"""Whether the graph has edges of type ``etype``.
Parameters
----------
u : list, tensor
The node ID array of source type.
v : list, tensor
The node ID array of destination type.
etype : str or tuple of str, optional
The edge type. Can be omitted if there is only one edge type
in the graph.
Returns
-------
a : tensor
Binary tensor indicating the existence of edges. ``a[i]=1`` if the graph
contains edge ``(u[i], v[i])`` of type ``etype``, 0 otherwise.
Examples
--------
The following example uses PyTorch backend.
>>> g.has_edges_between([0, 0], [1, 2], ('user', 'plays', 'game'))
tensor([1, 0])
See Also
--------
has_edge_between
"""
u = utils.toindex(u)
v = utils.toindex(v)
rst = self._graph.has_edges_between(self.get_etype_id(etype), u, v)
return rst.tousertensor()
def predecessors(self, v, etype=None):
"""Return the predecessors of node `v` in the graph with the specified
edge type.
Node `u` is a predecessor of `v` if an edge `(u, v)` with type `etype`
exists in the graph.
Parameters
----------
v : int
The destination node.
etype : str or tuple of str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
Returns
-------
tensor
Array of predecessor node IDs with the specified edge type.
Examples
--------
The following example uses PyTorch backend.
>>> plays_g = dgl.bipartite([(0, 0), (1, 0), (1, 1), (2, 1)], 'user', 'plays', 'game')
>>> devs_g = dgl.bipartite([(0, 0), (1, 1)], 'developer', 'develops', 'game')
>>> g = dgl.hetero_from_relations([plays_g, devs_g])
>>> g.predecessors(0, 'plays')
tensor([0, 1])
>>> g.predecessors(0, 'develops')
tensor([0])
See Also
--------
successors
"""
return self._graph.predecessors(self.get_etype_id(etype), v).tousertensor()
def successors(self, v, etype=None):
"""Return the successors of node `v` in the graph with the specified edge
type.
Node `u` is a successor of `v` if an edge `(v, u)` with type `etype` exists
in the graph.
Parameters
----------
v : int
The source node.
etype : str or tuple of str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
Returns
-------
tensor
Array of successor node IDs with the specified edge type.
Examples
--------
The following example uses PyTorch backend.
>>> plays_g = dgl.bipartite([(0, 0), (1, 0), (1, 1), (2, 1)], 'user', 'plays', 'game')
>>> follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
>>> g = dgl.hetero_from_relations([plays_g, follows_g])
>>> g.successors(0, 'plays')
tensor([0])
>>> g.successors(0, 'follows')
tensor([1])
See Also
--------
predecessors
"""
return self._graph.successors(self.get_etype_id(etype), v).tousertensor()
def edge_id(self, u, v, force_multi=None, return_array=False, etype=None):
"""Return the edge ID, or an array of edge IDs, between source node
`u` and destination node `v`, with the specified edge type
Parameters
----------
u : int
The node ID of source type.
v : int
The node ID of destination type.
force_multi : bool, optional
Deprecated (Will be deleted in the future).
If False, will return a single edge ID.
If True, will always return an array. (Default: False)
return_array : bool, optional
If False, will return a single edge ID.
If True, will always return an array. (Default: False)
etype : str or tuple of str, optional
The edge type. Can be omitted if there is only one edge type
in the graph.
Returns
-------
int or tensor
The edge ID if ``return_array == False``.
The edge ID array otherwise.
Notes
-----
If multiply edges exist between `u` and `v` and return_array is False,
the result is undefined.
Examples
--------
The following example uses PyTorch backend.
Instantiate a heterograph.
>>> plays_g = dgl.bipartite([(0, 0), (1, 0), (1, 2), (2, 1)], 'user', 'plays', 'game')
>>> follows_g = dgl.graph([(0, 1), (1, 2), (1, 2)], 'user', 'follows')
>>> g = dgl.hetero_from_relations([plays_g, follows_g])
Query for edge id.
>>> plays_g.edge_id(1, 2, etype=('user', 'plays', 'game'))
2
>>> g.edge_id(1, 2, return_array=True, etype=('user', 'follows', 'user'))
tensor([1, 2])
See Also
--------
edge_ids
"""
idx = self._graph.edge_id(self.get_etype_id(etype), u, v)
if force_multi is not None:
dgl_warning("force_multi will be deprecated." \
"Please use return_array instead")
return_array = force_multi
if return_array:
return idx.tousertensor()
else:
assert len(idx) == 1, "For return_array=False, there should be one and " \
"only one edge between u and v, but get {} edges. " \
"Please use return_array=True instead".format(len(idx))
return idx[0]
def edge_ids(self, u, v, force_multi=None, return_uv=False, etype=None):
"""Return all edge IDs between source node array `u` and destination
node array `v` with the specified edge type.
Parameters
----------
u : list, tensor
The node ID array of source type.
v : list, tensor
The node ID array of destination type.
force_multi : bool, optional
Deprecated (Will be deleted in the future).
Whether to always treat the graph as a multigraph. See the
"Returns" for their effects. (Default: False)
return_uv : bool
See the "Returns" for their effects. (Default: False)
etype : str or tuple of str, optional
The edge type. Can be omitted if there is only one edge type
in the graph.
Returns
-------
tensor, or (tensor, tensor, tensor)
* If ``return_uv=False``, return a single edge ID array ``e``.
``e[i]`` is the edge ID between ``u[i]`` and ``v[i]``.
* Otherwise, return three arrays ``(eu, ev, e)``. ``e[i]`` is the ID
of an edge between ``eu[i]`` and ``ev[i]``. All edges between ``u[i]``
and ``v[i]`` are returned.
Notes
-----
If the graph is a simple graph, ``return_uv=False``, and no edge
exists between some pairs of ``u[i]`` and ``v[i]``, the result is undefined
and an empty tensor is returned.
If the graph is a multi graph, ``return_uv=False``, and multi edges
exist between some pairs of `u[i]` and `v[i]`, the result is undefined.
Examples
--------
The following example uses PyTorch backend.
Instantiate a heterograph.
>>> plays_g = dgl.bipartite([(0, 0), (1, 0), (1, 2), (2, 1)], 'user', 'plays', 'game')
>>> follows_g = dgl.graph([(0, 1), (1, 2), (1, 2)], 'user', 'follows')
>>> g = dgl.hetero_from_relations([plays_g, follows_g])
Query for edge ids.
>>> plays_g.edge_ids([0], [2], etype=('user', 'plays', 'game'))
tensor([], dtype=torch.int64)
>>> plays_g.edge_ids([1], [2], etype=('user', 'plays', 'game'))
tensor([2])
>>> g.edge_ids([1], [2], return_uv=True, etype=('user', 'follows', 'user'))
(tensor([1, 1]), tensor([2, 2]), tensor([1, 2]))
See Also
--------
edge_id
"""
u = utils.toindex(u)
v = utils.toindex(v)
src, dst, eid = self._graph.edge_ids(self.get_etype_id(etype), u, v)
if force_multi is not None:
dgl_warning("force_multi will be deprecated, " \
"Please use return_uv instead")
return_uv = force_multi
if return_uv:
return src.tousertensor(), dst.tousertensor(), eid.tousertensor()
else:
assert len(eid) == max(len(u), len(v)), "If return_uv=False, there should be one and " \
"only one edge between each u and v, expect {} edges but get {}. " \
"Please use return_uv=True instead".format(max(len(u), len(v)), len(eid))
return eid.tousertensor()
def find_edges(self, eid, etype=None):
"""Given an edge ID array with the specified type, return the source
and destination node ID array ``s`` and ``d``. ``s[i]`` and ``d[i]``
are source and destination node ID for edge ``eid[i]``.
Parameters
----------
eid : list, tensor
The edge ID array.
etype : str or tuple of str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
Returns
-------
tensor
The source node ID array.
tensor
The destination node ID array.
Examples
--------
The following example uses PyTorch backend.
>>> g = dgl.bipartite([(0, 0), (1, 0), (1, 2)], 'user', 'plays', 'game')
>>> g.find_edges([0, 2], ('user', 'plays', 'game'))
(tensor([0, 1]), tensor([0, 2]))
>>> g.find_edges([0, 2])
(tensor([0, 1]), tensor([0, 2]))
"""
eid = utils.toindex(eid)
src, dst, _ = self._graph.find_edges(self.get_etype_id(etype), eid)
return src.tousertensor(), dst.tousertensor()
def in_edges(self, v, form='uv', etype=None):
"""Return the inbound edges of the node(s) with the specified type.
Parameters
----------
v : int, list, tensor
The node id(s) of destination type.
form : str, optional
The return form. Currently support:
- ``'eid'`` : one eid tensor
- ``'all'`` : a tuple ``(u, v, eid)``
- ``'uv'`` : a pair ``(u, v)``, default
etype : str or tuple of str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
Returns
-------
tensor or (tensor, tensor, tensor) or (tensor, tensor)
All inbound edges to ``v`` are returned.
* If ``form='eid'``, return a tensor for the ids of the
inbound edges of the nodes with the specified type.
* If ``form='all'``, return a 3-tuple of tensors
``(eu, ev, eid)``. ``eid[i]`` gives the ID of the
edge from ``eu[i]`` to ``ev[i]``.
* If ``form='uv'``, return a 2-tuple of tensors ``(eu, ev)``.
``eu[i]`` is the source node of an edge to ``ev[i]``.
Examples
--------
The following example uses PyTorch backend.
>>> g = dgl.bipartite([(0, 0), (1, 1), (1, 2)], 'user', 'plays', 'game')
>>> g.in_edges([0, 2], form='eid')
tensor([0, 2])
>>> g.in_edges([0, 2], form='all')
(tensor([0, 1]), tensor([0, 2]), tensor([0, 2]))
>>> g.in_edges([0, 2], form='uv')
(tensor([0, 1]), tensor([0, 2]))
"""
v = utils.toindex(v)
src, dst, eid = self._graph.in_edges(self.get_etype_id(etype), v)
if form == 'all':
return (src.tousertensor(), dst.tousertensor(), eid.tousertensor())
elif form == 'uv':
return (src.tousertensor(), dst.tousertensor())
elif form == 'eid':
return eid.tousertensor()
else:
raise DGLError('Invalid form:', form)
def out_edges(self, u, form='uv', etype=None):
"""Return the outbound edges of the node(s) with the specified type.
Parameters
----------
u : int, list, tensor
The node id(s) of source type.
form : str, optional
The return form. Currently support:
- ``'eid'`` : one eid tensor
- ``'all'`` : a tuple ``(u, v, eid)``
- ``'uv'`` : a pair ``(u, v)``, default
etype : str or tuple of str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
Returns
-------
tensor or (tensor, tensor, tensor) or (tensor, tensor)
All outbound edges from ``u`` are returned.
* If ``form='eid'``, return a tensor for the ids of the outbound edges
of the nodes with the specified type.
* If ``form='all'``, return a 3-tuple of tensors ``(eu, ev, eid)``.
``eid[i]`` gives the ID of the edge from ``eu[i]`` to ``ev[i]``.
* If ``form='uv'``, return a 2-tuple of tensors ``(eu, ev)``.
``ev[i]`` is the destination node of the edge from ``eu[i]``.
Examples
--------
>>> g = dgl.bipartite([(0, 0), (1, 1), (1, 2)], 'user', 'plays', 'game')
>>> g.out_edges([0, 1], form='eid')
tensor([0, 1, 2])
>>> g.out_edges([0, 1], form='all')
(tensor([0, 1, 1]), tensor([0, 1, 2]), tensor([0, 1, 2]))
>>> g.out_edges([0, 1], form='uv')
(tensor([0, 1, 1]), tensor([0, 1, 2]))
"""
u = utils.toindex(u)
src, dst, eid = self._graph.out_edges(self.get_etype_id(etype), u)
if form == 'all':
return (src.tousertensor(), dst.tousertensor(), eid.tousertensor())
elif form == 'uv':
return (src.tousertensor(), dst.tousertensor())
elif form == 'eid':
return eid.tousertensor()
else:
raise DGLError('Invalid form:', form)
def all_edges(self, form='uv', order=None, etype=None):
"""Return all edges with the specified type.
Parameters
----------
form : str, optional
The return form. Currently support:
- ``'eid'`` : one eid tensor
- ``'all'`` : a tuple ``(u, v, eid)``
- ``'uv'`` : a pair ``(u, v)``, default
order : str or None
The order of the returned edges. Currently support:
- ``'srcdst'`` : sorted by their src and dst ids.
- ``'eid'`` : sorted by edge Ids.
- ``None`` : arbitrary order, default
etype : str or tuple of str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
Returns
-------
tensor or (tensor, tensor, tensor) or (tensor, tensor)
* If ``form='eid'``, return a tensor for the ids of all edges
with the specified type.
* If ``form='all'``, return a 3-tuple of tensors ``(eu, ev, eid)``.
``eid[i]`` gives the ID of the edge from ``eu[i]`` to ``ev[i]``.
* If ``form='uv'``, return a 2-tuple of tensors ``(eu, ev)``.
``ev[i]`` is the destination node of the edge from ``eu[i]``.
Examples
--------
The following example uses PyTorch backend.
>>> g = dgl.bipartite([(1, 1), (0, 0), (1, 2)], 'user', 'plays', 'game')
>>> g.all_edges(form='eid', order='srcdst')
tensor([1, 0, 2])
>>> g.all_edges(form='all', order='srcdst')
(tensor([0, 1, 1]), tensor([0, 1, 2]), tensor([1, 0, 2]))
>>> g.all_edges(form='uv', order='eid')
(tensor([1, 0, 1]), tensor([1, 0, 2]))
"""
src, dst, eid = self._graph.edges(self.get_etype_id(etype), order)
if form == 'all':
return (src.tousertensor(), dst.tousertensor(), eid.tousertensor())
elif form == 'uv':
return (src.tousertensor(), dst.tousertensor())
elif form == 'eid':
return eid.tousertensor()
else:
raise DGLError('Invalid form:', form)
def in_degree(self, v, etype=None):
"""Return the in-degree of node ``v`` with edges of type ``etype``.
Parameters
----------
v : int
The node ID of destination type.
etype : str or tuple of str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
Returns
-------
int
The in-degree.
Examples
--------
Instantiate a heterograph.
>>> plays_g = dgl.bipartite([(0, 0), (1, 0), (1, 2), (2, 1)], 'user', 'plays', 'game')
>>> follows_g = dgl.graph([(0, 1), (1, 2), (1, 2)], 'user', 'follows')
>>> g = dgl.hetero_from_relations([plays_g, follows_g])
Query for node degree.
>>> g.in_degree(0, 'plays')
2
>>> g.in_degree(0, 'follows')
0
See Also
--------
in_degrees
"""
return self._graph.in_degree(self.get_etype_id(etype), v)
def in_degrees(self, v=ALL, etype=None):
"""Return the in-degrees of nodes v with edges of type ``etype``.
Parameters
----------
v : list, tensor, optional.
The node ID array of the destination type. Default is to return the
degrees of all nodes.
etype : str or tuple of str or None, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
Returns
-------
d : tensor
The in-degree array. ``d[i]`` gives the in-degree of node ``v[i]``
with edges of type ``etype``.
Examples
--------
The following example uses PyTorch backend.
Instantiate a heterograph.
>>> plays_g = dgl.bipartite([(0, 0), (1, 0), (1, 2), (2, 1)], 'user', 'plays', 'game')
>>> follows_g = dgl.graph([(0, 1), (1, 2), (1, 2)], 'user', 'follows')
>>> g = dgl.hetero_from_relations([plays_g, follows_g])
Query for node degree.
>>> g.in_degrees(0, 'plays')
tensor([2])
>>> g.in_degrees(etype='follows')
tensor([0, 1, 2])
See Also
--------
in_degree
"""
etid = self.get_etype_id(etype)
_, dtid = self._graph.metagraph.find_edge(etid)
if is_all(v):
v = utils.toindex(slice(0, self._graph.number_of_nodes(dtid)))
else:
v = utils.toindex(v)
return self._graph.in_degrees(etid, v).tousertensor()
def out_degree(self, u, etype=None):
"""Return the out-degree of node `u` with edges of type ``etype``.
Parameters
----------
u : int
The node ID of source type.
etype : str or tuple of str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
Returns
-------
int
The out-degree of node `u` with edges of type ``etype``.
Examples
--------
Instantiate a heterograph.
>>> plays_g = dgl.bipartite([(0, 0), (1, 0), (1, 2), (2, 1)], 'user', 'plays', 'game')
>>> follows_g = dgl.graph([(0, 1), (1, 2), (1, 2)], 'user', 'follows')
>>> g = dgl.hetero_from_relations([plays_g, follows_g])
Query for node degree.
>>> g.out_degree(0, 'plays')
1
>>> g.out_degree(1, 'follows')
2
See Also
--------
out_degrees
"""
return self._graph.out_degree(self.get_etype_id(etype), u)
def out_degrees(self, u=ALL, etype=None):
"""Return the out-degrees of nodes u with edges of type ``etype``.
Parameters
----------
u : list, tensor
The node ID array of source type. Default is to return the degrees
of all the nodes.
etype : str or tuple of str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
Returns
-------
d : tensor
The out-degree array. ``d[i]`` gives the out-degree of node ``u[i]``
with edges of type ``etype``.
Examples
--------
The following example uses PyTorch backend.
Instantiate a heterograph.
>>> plays_g = dgl.bipartite([(0, 0), (1, 0), (1, 2), (2, 1)], 'user', 'plays', 'game')
>>> follows_g = dgl.graph([(0, 1), (1, 2), (1, 2)], 'user', 'follows')
>>> g = dgl.hetero_from_relations([plays_g, follows_g])
Query for node degree.
>>> g.out_degrees(0, 'plays')
tensor([1])
>>> g.out_degrees(etype='follows')
tensor([1, 2, 0])
See Also
--------
out_degree
"""
etid = self.get_etype_id(etype)
stid, _ = self._graph.metagraph.find_edge(etid)
if is_all(u):
u = utils.toindex(slice(0, self._graph.number_of_nodes(stid)))
else:
u = utils.toindex(u)
return self._graph.out_degrees(etid, u).tousertensor()
def _create_hetero_subgraph(self, sgi, induced_nodes, induced_edges):
"""Internal function to create a subgraph."""
node_frames = [
FrameRef(Frame(
self._node_frames[i][induced_nodes_of_ntype],
num_rows=len(induced_nodes_of_ntype)))
for i, induced_nodes_of_ntype in enumerate(induced_nodes)]
edge_frames = [
FrameRef(Frame(
self._edge_frames[i][induced_edges_of_etype],
num_rows=len(induced_edges_of_etype)))
for i, induced_edges_of_etype in enumerate(induced_edges)]
hsg = DGLHeteroGraph(sgi.graph, self._ntypes, self._etypes, node_frames, edge_frames)
hsg.is_subgraph = True
for ntype, induced_nid in zip(self.ntypes, induced_nodes):
hsg.nodes[ntype].data[NID] = induced_nid.tousertensor()
for etype, induced_eid in zip(self.canonical_etypes, induced_edges):
hsg.edges[etype].data[EID] = induced_eid.tousertensor()
return hsg
def subgraph(self, nodes):
"""Return the subgraph induced on given nodes.
The metagraph of the returned subgraph is the same as the parent graph.
Features are copied from the original graph.
Parameters
----------
nodes : dict[str->list or iterable]
A dictionary mapping node types to node ID array for constructing
subgraph. All nodes must exist in the graph.
Returns
-------
G : DGLHeteroGraph
The subgraph.
The nodes and edges in the subgraph are relabeled using consecutive
integers from 0.
One can retrieve the mapping from subgraph node/edge ID to parent
node/edge ID via ``dgl.NID`` and ``dgl.EID`` node/edge features of the
subgraph.
Examples
--------
The following example uses PyTorch backend.
Instantiate a heterograph.
>>> plays_g = dgl.bipartite([(0, 0), (1, 0), (1, 2), (2, 1)], 'user', 'plays', 'game')
>>> follows_g = dgl.graph([(0, 1), (1, 2), (1, 2)], 'user', 'follows')
>>> g = dgl.hetero_from_relations([plays_g, follows_g])
>>> # Set node features
>>> g.nodes['user'].data['h'] = torch.tensor([[0.], [1.], [2.]])
Get subgraphs.
>>> g.subgraph({'user': [4, 5]})
An error occurs as these nodes do not exist.
>>> sub_g = g.subgraph({'user': [1, 2]})
>>> print(sub_g)
Graph(num_nodes={'user': 2, 'game': 0},
num_edges={('user', 'plays', 'game'): 0, ('user', 'follows', 'user'): 2},
metagraph=[('user', 'game'), ('user', 'user')])
Get the original node/edge indices.
>>> sub_g['follows'].ndata[dgl.NID] # Get the node indices in the raw graph
tensor([1, 2])
>>> sub_g['follows'].edata[dgl.EID] # Get the edge indices in the raw graph
tensor([1, 2])
Get the copied node features.
>>> sub_g.nodes['user'].data['h']
tensor([[1.],
[2.]])
>>> sub_g.nodes['user'].data['h'] += 1
>>> g.nodes['user'].data['h'] # Features are not shared.
tensor([[0.],
[1.],
[2.]])
See Also
--------
edge_subgraph
"""
induced_nodes = [utils.toindex(nodes.get(ntype, [])) for ntype in self.ntypes]
sgi = self._graph.node_subgraph(induced_nodes)
induced_edges = sgi.induced_edges
return self._create_hetero_subgraph(sgi, induced_nodes, induced_edges)
def edge_subgraph(self, edges, preserve_nodes=False):
"""Return the subgraph induced on given edges.
The metagraph of the returned subgraph is the same as the parent graph.
Features are copied from the original graph.
Parameters
----------
edges : dict[str->list or iterable]
A dictionary mapping edge types to edge ID array for constructing
subgraph. All edges must exist in the subgraph.
The edge types are characterized by triplets of
``(src type, etype, dst type)``.
preserve_nodes : bool
Whether to preserve all nodes or not. If false, all nodes
without edges will be removed. (Default: False)
Returns
-------
G : DGLHeteroGraph
The subgraph.
The nodes and edges are relabeled using consecutive integers from 0.
One can retrieve the mapping from subgraph node/edge ID to parent
node/edge ID via ``dgl.NID`` and ``dgl.EID`` node/edge features of the
subgraph.
Examples
--------
The following example uses PyTorch backend.
Instantiate a heterograph.
>>> plays_g = dgl.bipartite([(0, 0), (1, 0), (1, 2), (2, 1)], 'user', 'plays', 'game')
>>> follows_g = dgl.graph([(0, 1), (1, 2), (1, 2)], 'user', 'follows')
>>> g = dgl.hetero_from_relations([plays_g, follows_g])
>>> # Set edge features
>>> g.edges['follows'].data['h'] = torch.tensor([[0.], [1.], [2.]])
Get subgraphs.
>>> g.edge_subgraph({('user', 'follows', 'user'): [5, 6]})
An error occurs as these edges do not exist.
>>> sub_g = g.edge_subgraph({('user', 'follows', 'user'): [1, 2],
>>> ('user', 'plays', 'game'): [2]})
>>> print(sub_g)
Graph(num_nodes={'user': 2, 'game': 1},
num_edges={('user', 'plays', 'game'): 1, ('user', 'follows', 'user'): 2},
metagraph=[('user', 'game'), ('user', 'user')])
Get the original node/edge indices.
>>> sub_g['follows'].ndata[dgl.NID] # Get the node indices in the raw graph
tensor([1, 2])
>>> sub_g['plays'].edata[dgl.EID] # Get the edge indices in the raw graph
tensor([2])
Get the copied node features.
>>> sub_g.edges['follows'].data['h']
tensor([[1.],
[2.]])
>>> sub_g.edges['follows'].data['h'] += 1
>>> g.edges['follows'].data['h'] # Features are not shared.
tensor([[0.],
[1.],
[2.]])
See Also
--------
subgraph
"""
edges = {self.to_canonical_etype(etype): e for etype, e in edges.items()}
induced_edges = [
utils.toindex(edges.get(canonical_etype, []))
for canonical_etype in self.canonical_etypes]
sgi = self._graph.edge_subgraph(induced_edges, preserve_nodes)
induced_nodes = sgi.induced_nodes
return self._create_hetero_subgraph(sgi, induced_nodes, induced_edges)
def node_type_subgraph(self, ntypes):
"""Return the subgraph induced on given node types.
The metagraph of the returned subgraph is the subgraph of the original
metagraph induced from the node types.
Features are shared with the original graph.
Parameters
----------
ntypes : list[str]
The node types
Returns
-------
G : DGLHeteroGraph
The subgraph.
Examples
--------
The following example uses PyTorch backend.
Instantiate a heterograph.
>>> plays_g = dgl.bipartite([(0, 0), (1, 0), (1, 2), (2, 1)], 'user', 'plays', 'game')
>>> follows_g = dgl.graph([(0, 1), (1, 2), (1, 2)], 'user', 'follows')
>>> g = dgl.hetero_from_relations([plays_g, follows_g])
>>> # Set node features
>>> g.nodes['user'].data['h'] = torch.tensor([[0.], [1.], [2.]])
Get subgraphs.
>>> sub_g = g.node_type_subgraph(['user'])
>>> print(sub_g)
Graph(num_nodes=3, num_edges=3,
ndata_schemes={'h': Scheme(shape=(1,), dtype=torch.float32)}
edata_schemes={})
Get the shared node features.
>>> sub_g.nodes['user'].data['h']
tensor([[0.],
[1.],
[2.]])
>>> sub_g.nodes['user'].data['h'] += 1
>>> g.nodes['user'].data['h'] # Features are shared.
tensor([[1.],
[2.],
[3.]])
See Also
--------
edge_type_subgraph
"""
rel_graphs = []
meta_edges = []
induced_etypes = []
node_frames = [self._node_frames[self.get_ntype_id(ntype)] for ntype in ntypes]
edge_frames = []
num_nodes_per_type = [self.number_of_nodes(ntype) for ntype in ntypes]
ntypes_invmap = {ntype: i for i, ntype in enumerate(ntypes)}
srctype_id, dsttype_id, _ = self._graph.metagraph.edges('eid')
for i in range(len(self._etypes)):
srctype = self._ntypes[srctype_id[i]]
dsttype = self._ntypes[dsttype_id[i]]
if srctype in ntypes and dsttype in ntypes:
meta_edges.append((ntypes_invmap[srctype], ntypes_invmap[dsttype]))
rel_graphs.append(self._graph.get_relation_graph(i))
induced_etypes.append(self.etypes[i])
edge_frames.append(self._edge_frames[i])
metagraph = graph_index.from_edge_list(meta_edges, True)
hgidx = heterograph_index.create_heterograph_from_relations(
metagraph, rel_graphs, utils.toindex(num_nodes_per_type))
hg = DGLHeteroGraph(hgidx, ntypes, induced_etypes, node_frames, edge_frames)
return hg
def edge_type_subgraph(self, etypes):
"""Return the subgraph induced on given edge types.
The metagraph of the returned subgraph is the subgraph of the original metagraph
induced from the edge types.
Features are shared with the original graph.
Parameters
----------
etypes : list[str or tuple]
The edge types
Returns
-------
G : DGLHeteroGraph
The subgraph.
Examples
--------
The following example uses PyTorch backend.
Instantiate a heterograph.
>>> plays_g = dgl.bipartite([(0, 0), (1, 0), (1, 2), (2, 1)], 'user', 'plays', 'game')
>>> follows_g = dgl.graph([(0, 1), (1, 2), (1, 2)], 'user', 'follows')
>>> g = dgl.hetero_from_relations([plays_g, follows_g])
>>> # Set edge features
>>> g.edges['follows'].data['h'] = torch.tensor([[0.], [1.], [2.]])
Get subgraphs.
>>> sub_g = g.edge_type_subgraph(['follows'])
>>> print(sub_g)
Graph(num_nodes=3, num_edges=3,
ndata_schemes={}
edata_schemes={'h': Scheme(shape=(1,), dtype=torch.float32)})
Get the shared edge features.
>>> sub_g.edges['follows'].data['h']
tensor([[0.],
[1.],
[2.]])
>>> sub_g.edges['follows'].data['h'] += 1
>>> g.edges['follows'].data['h'] # Features are shared.
tensor([[1.],
[2.],
[3.]])
See Also
--------
node_type_subgraph
"""
etype_ids = [self.get_etype_id(etype) for etype in etypes]
meta_src, meta_dst, _ = self._graph.metagraph.find_edges(utils.toindex(etype_ids))
rel_graphs = [self._graph.get_relation_graph(i) for i in etype_ids]
meta_src = meta_src.tonumpy()
meta_dst = meta_dst.tonumpy()
ntypes_invmap = {n: i for i, n in enumerate(set(meta_src) | set(meta_dst))}
mapped_meta_src = [ntypes_invmap[v] for v in meta_src]
mapped_meta_dst = [ntypes_invmap[v] for v in meta_dst]
node_frames = [self._node_frames[i] for i in ntypes_invmap]
edge_frames = [self._edge_frames[i] for i in etype_ids]
induced_ntypes = [self._ntypes[i] for i in ntypes_invmap]
induced_etypes = [self._etypes[i] for i in etype_ids] # get the "name" of edge type
num_nodes_per_induced_type = [self.number_of_nodes(ntype) for ntype in induced_ntypes]
metagraph = graph_index.from_edge_list((mapped_meta_src, mapped_meta_dst), True)
hgidx = heterograph_index.create_heterograph_from_relations(
metagraph, rel_graphs, utils.toindex(num_nodes_per_induced_type))
hg = DGLHeteroGraph(hgidx, induced_ntypes, induced_etypes, node_frames, edge_frames)
return hg
def adjacency_matrix(self, transpose=None, ctx=F.cpu(), scipy_fmt=None, etype=None):
"""Return the adjacency matrix of edges of the given edge type.
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
----------
transpose : bool, optional
A flag to transpose the returned adjacency matrix. (Default: False)
ctx : context, optional
The context of returned adjacency matrix. (Default: cpu)
scipy_fmt : str, optional
If specified, return a scipy sparse matrix in the given format.
Otherwise, return a backend dependent sparse tensor. (Default: None)
etype : str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
Returns
-------
SparseTensor or scipy.sparse.spmatrix
Adjacency matrix.
Examples
--------
Instantiate a heterogeneous graph.
>>> follows_g = dgl.graph([(0, 0), (1, 1)], 'user', 'follows')
>>> devs_g = dgl.bipartite([(0, 0), (1, 2)], 'developer', 'develops', 'game')
>>> g = dgl.hetero_from_relations([follows_g, devs_g])
Get a backend dependent sparse tensor. Here we use PyTorch for example.
>>> g.adjacency_matrix(etype='develops')
tensor(indices=tensor([[0, 2],
[0, 1]]),
values=tensor([1., 1.]),
size=(3, 2), nnz=2, layout=torch.sparse_coo)
Get a scipy coo sparse matrix.
>>> g.adjacency_matrix(scipy_fmt='coo', etype='develops')
<3x2 sparse matrix of type '<class 'numpy.int64'>'
with 2 stored elements in COOrdinate format>
"""
if transpose is None:
dgl_warning(
"Currently adjacency_matrix() returns a matrix with destination as rows"
" by default. In 0.5 the result will have source as rows"
" (i.e. transpose=True)")
transpose = False
etid = self.get_etype_id(etype)
if scipy_fmt is None:
return self._graph.adjacency_matrix(etid, transpose, ctx)[0]
else:
return self._graph.adjacency_matrix_scipy(etid, transpose, scipy_fmt, False)
# Alias of ``adjacency_matrix``
adj = adjacency_matrix
def incidence_matrix(self, typestr, ctx=F.cpu(), etype=None):
"""Return the incidence matrix representation of edges with the given
edge type.
An incidence matrix is an n-by-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 incidence matrices :math:`I`:
* ``in``:
- :math:`I[v, e] = 1` if :math:`e` is the in-edge of :math:`v`
(or :math:`v` is the dst node of :math:`e`);
- :math:`I[v, e] = 0` otherwise.
* ``out``:
- :math:`I[v, e] = 1` if :math:`e` is the out-edge of :math:`v`
(or :math:`v` is the src node of :math:`e`);
- :math:`I[v, e] = 0` otherwise.
* ``both`` (only if source and destination node type are the same):
- :math:`I[v, e] = 1` if :math:`e` is the in-edge of :math:`v`;
- :math:`I[v, e] = -1` if :math:`e` is the out-edge of :math:`v`;
- :math:`I[v, e] = 0` otherwise (including self-loop).
Parameters
----------
typestr : str
Can be either ``in``, ``out`` or ``both``
ctx : context, optional
The context of returned incidence matrix. (Default: cpu)
etype : str, optional
The edge type. Can be omitted if there is only one edge type
in the graph.
Returns
-------
Framework SparseTensor
The incidence matrix.
Examples
--------
>>> g = dgl.graph([(0, 0), (1, 2)], 'user', 'follows')
>>> g.incidence_matrix('in')
tensor(indices=tensor([[0, 2],
[0, 1]]),
values=tensor([1., 1.]),
size=(3, 2), nnz=2, layout=torch.sparse_coo)
>>> g.incidence_matrix('out')
tensor(indices=tensor([[0, 1],
[0, 1]]),
values=tensor([1., 1.]),
size=(3, 2), nnz=2, layout=torch.sparse_coo)
>>> g.incidence_matrix('both')
tensor(indices=tensor([[1, 2],
[1, 1]]),
values=tensor([-1., 1.]),
size=(3, 2), nnz=2, layout=torch.sparse_coo)
"""
etid = self.get_etype_id(etype)
return self._graph.incidence_matrix(etid, typestr, ctx)[0]
# Alias of ``incidence_matrix``
inc = incidence_matrix
#################################################################
# Features
#################################################################
def node_attr_schemes(self, ntype=None):
"""Return the node feature schemes for the specified type.
Each feature scheme is a named tuple that stores the shape and data type
of the node feature.
Parameters
----------
ntype : str, optional
The node type. Can be omitted if there is only one node
type in the graph. Error will be raised otherwise.
(Default: None)
Returns
-------
dict of str to schemes
The schemes of node feature columns.
Examples
--------
The following uses PyTorch backend.
>>> g = dgl.graph([(0, 0), (1, 2)], 'user', 'follows')
>>> g.nodes['user'].data['h'] = torch.randn(3, 4)
>>> g.node_attr_schemes('user')
{'h': Scheme(shape=(4,), dtype=torch.float32)}
See Also
--------
edge_attr_schemes
"""
return self._node_frames[self.get_ntype_id(ntype)].schemes
def edge_attr_schemes(self, etype=None):
"""Return the edge feature schemes for the specified type.
Each feature scheme is a named tuple that stores the shape and data type
of the edge feature.
Parameters
----------
etype : str or tuple of str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
Returns
-------
dict of str to schemes
The schemes of edge feature columns.
Examples
--------
The following uses PyTorch backend.
>>> g = dgl.bipartite([(0, 0), (1, 0), (1, 2), (2, 1)], 'user', 'plays', 'game')
>>> g.edges['user', 'plays', 'game'].data['h'] = torch.randn(4, 4)
>>> g.edge_attr_schemes(('user', 'plays', 'game'))
{'h': Scheme(shape=(4,), dtype=torch.float32)}
See Also
--------
node_attr_schemes
"""
return self._edge_frames[self.get_etype_id(etype)].schemes
def set_n_initializer(self, initializer, field=None, ntype=None):
"""Set the initializer for empty node features.
Initializer is a callable that returns a tensor given the shape, data type
and device context.
When a subset of the nodes are assigned a new feature, initializer is
used to create feature for the rest of the nodes.
Parameters
----------
initializer : callable
The initializer, mapping (shape, data type, context) to tensor.
field : str, optional
The feature field name. Default is to set an initializer for all the
feature fields.
ntype : str, optional
The node type. Can be omitted if there is only one node
type in the graph. Error will be raised otherwise.
(Default: None)
Note
-----
User defined initializer must follow the signature of
:func:`dgl.init.base_initializer() <dgl.init.base_initializer>`
See Also
--------
set_e_initializer
"""
ntid = self.get_ntype_id(ntype)
self._node_frames[ntid].set_initializer(initializer, field)
def set_e_initializer(self, initializer, field=None, etype=None):
"""Set the initializer for empty edge features.
Initializer is a callable that returns a tensor given the shape, data
type and device context.
When a subset of the edges are assigned a new feature, initializer is
used to create feature for rest of the edges.
Parameters
----------
initializer : callable
The initializer, mapping (shape, data type, context) to tensor.
field : str, optional
The feature field name. Default is set an initializer for all the
feature fields.
etype : str or tuple of str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. Error will be raised otherwise.
(Default: None)
Note
-----
User defined initializer must follow the signature of
:func:`dgl.init.base_initializer() <dgl.init.base_initializer>`
See Also
--------
set_n_initializer
"""
etid = self.get_etype_id(etype)
self._edge_frames[etid].set_initializer(initializer, field)
def _set_n_repr(self, ntid, u, data, inplace=False):
"""Internal API to set node features.
`data` is a dictionary from the feature name to feature tensor. Each tensor
is of shape (B, D1, D2, ...), where B is the number of nodes to be updated,
and (D1, D2, ...) be the shape of the node representation tensor. The
length of the given node ids must match B (i.e, len(u) == B).
All update will be done out of place to work with autograd unless the
inplace flag is true.
Parameters
----------
ntid : int
Node type id.
u : node, container or tensor
The node(s).
data : dict of tensor
Node representation.
inplace : bool, optional
If True, update will be done in place, but autograd will break.
(Default: False)
"""
if is_all(u):
num_nodes = self._graph.number_of_nodes(ntid)
else:
u = utils.toindex(u)
num_nodes = len(u)
for key, val in data.items():
nfeats = F.shape(val)[0]
if nfeats != num_nodes:
raise DGLError('Expect number of features to match number of nodes (len(u)).'
' Got %d and %d instead.' % (nfeats, num_nodes))
if is_all(u):
for key, val in data.items():
self._node_frames[ntid][key] = val
else:
self._node_frames[ntid].update_rows(u, data, inplace=inplace)
def _get_n_repr(self, ntid, u):
"""Get node(s) representation of a single node type.
The returned feature tensor batches multiple node features on the first dimension.
Parameters
----------
ntid : int
Node type id.
u : node, container or tensor
The node(s).
Returns
-------
dict
Representation dict from feature name to feature tensor.
"""
if is_all(u):
return dict(self._node_frames[ntid])
else:
u = utils.toindex(u)
return self._node_frames[ntid].select_rows(u)
def _pop_n_repr(self, ntid, key):
"""Internal API to get and remove the specified node feature.
Parameters
----------
ntid : int
Node type id.
key : str
The attribute name.
Returns
-------
Tensor
The popped representation
"""
return self._node_frames[ntid].pop(key)
def _set_e_repr(self, etid, edges, data, inplace=False):
"""Internal API to set edge(s) features.
`data` is a dictionary from the feature name to feature tensor. Each tensor
is of shape (B, D1, D2, ...), where B is the number of edges to be updated,
and (D1, D2, ...) be the shape of the edge representation tensor.
All update will be done out of place to work with autograd unless the
inplace flag is true.
Parameters
----------
etid : int
Edge type id.
edges : edges
Edges can be either
* A pair of endpoint nodes (u, v), where u is the node ID of source
node type and v is that of destination node type.
* A tensor of edge ids of the given type.
The default value is all the edges.
data : tensor or dict of tensor
Edge representation.
inplace : bool, optional
If True, update will be done in place, but autograd will break.
(Default: False)
"""
# parse argument
if is_all(edges):
eid = ALL
elif isinstance(edges, tuple):
u, v = edges
u = utils.toindex(u)
v = utils.toindex(v)
# Rewrite u, v to handle edge broadcasting and multigraph.
_, _, eid = self._graph.edge_ids(etid, u, v)
else:
eid = utils.toindex(edges)
# sanity check
if not utils.is_dict_like(data):
raise DGLError('Expect dictionary type for feature data.'
' Got "%s" instead.' % type(data))
if is_all(eid):
num_edges = self._graph.number_of_edges(etid)
else:
eid = utils.toindex(eid)
num_edges = len(eid)
for key, val in data.items():
nfeats = F.shape(val)[0]
if nfeats != num_edges:
raise DGLError('Expect number of features to match number of edges.'
' Got %d and %d instead.' % (nfeats, num_edges))
# set
if is_all(eid):
# update column
for key, val in data.items():
self._edge_frames[etid][key] = val
else:
# update row
self._edge_frames[etid].update_rows(eid, data, inplace=inplace)
def _get_e_repr(self, etid, edges):
"""Internal API to get edge features.
Parameters
----------
etid : int
Edge type id.
edges : edges
Edges can be a pair of endpoint nodes (u, v), or a
tensor of edge ids. The default value is all the edges.
Returns
-------
dict
Representation dict
"""
# parse argument
if is_all(edges):
eid = ALL
elif isinstance(edges, tuple):
u, v = edges
u = utils.toindex(u)
v = utils.toindex(v)
# Rewrite u, v to handle edge broadcasting and multigraph.
_, _, eid = self._graph.edge_ids(etid, u, v)
else:
eid = utils.toindex(edges)
if is_all(eid):
return dict(self._edge_frames[etid])
else:
eid = utils.toindex(eid)
return self._edge_frames[etid].select_rows(eid)
def _pop_e_repr(self, etid, key):
"""Get and remove the specified edge repr of a single edge type.
Parameters
----------
etid : int
Edge type id.
key : str
The attribute name.
Returns
-------
Tensor
The popped representation
"""
self._edge_frames[etid].pop(key)
#################################################################
# Message passing
#################################################################
def apply_nodes(self, func, v=ALL, ntype=None, inplace=False):
"""Apply the function on the nodes with the same type to update their
features.
If None is provided for ``func``, nothing will happen.
Parameters
----------
func : callable or None
Apply function on the nodes. The function should be
a :mod:`Node UDF <dgl.udf>`.
v : int or iterable of int or tensor, optional
The (type-specific) node (ids) on which to apply ``func``. (Default: ALL)
ntype : str, optional
The node type. Can be omitted if there is only one node type
in the graph. (Default: None)
inplace : bool, optional
If True, update will be done in place, but autograd will break.
(Default: False)
Examples
--------
>>> g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
>>> g.nodes['user'].data['h'] = torch.ones(3, 5)
>>> g.apply_nodes(lambda nodes: {'h': nodes.data['h'] * 2}, ntype='user')
>>> g.nodes['user'].data['h']
tensor([[2., 2., 2., 2., 2.],
[2., 2., 2., 2., 2.],
[2., 2., 2., 2., 2.]])
See Also
--------
apply_edges
"""
ntid = self.get_ntype_id(ntype)
if is_all(v):
v_ntype = utils.toindex(slice(0, self.number_of_nodes(ntype)))
else:
v_ntype = utils.toindex(v)
with ir.prog() as prog:
scheduler.schedule_apply_nodes(v_ntype, func, self._node_frames[ntid],
inplace=inplace, ntype=self._ntypes[ntid])
Runtime.run(prog)
def apply_edges(self, func, edges=ALL, etype=None, inplace=False):
"""Apply the function on the edges with the same type to update their
features.
If None is provided for ``func``, nothing will happen.
Parameters
----------
func : callable or None
Apply function on the edge. The function should be
an :mod:`Edge UDF <dgl.udf>`.
edges : optional
Edges on which to apply ``func``. See :func:`send` for valid
edge specification. (Default: ALL)
etype : str or tuple of str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
inplace: bool, optional
If True, update will be done in place, but autograd will break.
(Default: False)
Examples
--------
>>> g = dgl.bipartite([(0, 0), (1, 0), (1, 2), (2, 1)], 'user', 'plays', 'game')
>>> g.edges[('user', 'plays', 'game')].data['h'] = torch.ones(4, 5)
>>> g.apply_edges(lambda edges: {'h': edges.data['h'] * 2})
>>> g.edges[('user', 'plays', 'game')].data['h']
tensor([[2., 2., 2., 2., 2.],
[2., 2., 2., 2., 2.],
[2., 2., 2., 2., 2.],
[2., 2., 2., 2., 2.]])
See Also
--------
apply_nodes
group_apply_edges
"""
etid = self.get_etype_id(etype)
stid, dtid = self._graph.metagraph.find_edge(etid)
if is_all(edges):
u, v, _ = self._graph.edges(etid, 'eid')
eid = utils.toindex(slice(0, self.number_of_edges(etype)))
elif isinstance(edges, tuple):
u, v = edges
u = utils.toindex(u)
v = utils.toindex(v)
# Rewrite u, v to handle edge broadcasting and multigraph.
u, v, eid = self._graph.edge_ids(etid, u, v)
else:
eid = utils.toindex(edges)
u, v, _ = self._graph.find_edges(etid, eid)
with ir.prog() as prog:
scheduler.schedule_apply_edges(
AdaptedHeteroGraph(self, stid, dtid, etid),
u, v, eid, func, inplace=inplace)
Runtime.run(prog)
def group_apply_edges(self, group_by, func, edges=ALL, etype=None, inplace=False):
"""Group the edges by nodes and apply the function of the grouped
edges to update their features. The edges are of the same edge type
(hence having the same source and destination node type).
Parameters
----------
group_by : str
Specify how to group edges. Expected to be either ``'src'`` or ``'dst'``
func : callable
Apply function on the edge. The function should be an
:mod:`Edge UDF <dgl.udf>`. The input of `Edge UDF` should be
(bucket_size, degrees, *feature_shape), and return the dict
with values of the same shapes.
edges : optional
Edges on which to group and apply ``func``. See :func:`send` for valid
edge specification. Default is all the edges.
etype : str or tuple of str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
inplace: bool, optional
If True, update will be done in place, but autograd will break.
(Default: False)
Examples
--------
>>> g = dgl.graph([(0, 1), (0, 2), (1, 2)], 'user', 'follows')
>>> g.edata['feat'] = torch.randn((g.number_of_edges(), 1))
>>> def softmax_feat(edges):
>>> return {'norm_feat': th.softmax(edges.data['feat'], dim=1)}
>>> g.group_apply_edges(group_by='src', func=softmax_feat)
>>> g.edata['norm_feat']
tensor([[0.3796],
[0.6204],
[1.0000]])
See Also
--------
apply_edges
"""
if group_by not in ('src', 'dst'):
raise DGLError("Group_by should be either src or dst")
etid = self.get_etype_id(etype)
stid, dtid = self._graph.metagraph.find_edge(etid)
if is_all(edges):
u, v, _ = self._graph.edges(etid, 'eid')
eid = utils.toindex(slice(0, self.number_of_edges(etype)))
elif isinstance(edges, tuple):
u, v = edges
u = utils.toindex(u)
v = utils.toindex(v)
# Rewrite u, v to handle edge broadcasting and multigraph.
u, v, eid = self._graph.edge_ids(etid, u, v)
else:
eid = utils.toindex(edges)
u, v, _ = self._graph.find_edges(etid, eid)
with ir.prog() as prog:
scheduler.schedule_group_apply_edge(
AdaptedHeteroGraph(self, stid, dtid, etid),
u, v, eid,
func, group_by,
inplace=inplace)
Runtime.run(prog)
def send(self, edges, message_func, etype=None):
"""Send messages along the given edges with the same edge type.
``edges`` can be any of the following types:
* ``int`` : Specify one edge using its edge id (of the given edge type).
* ``pair of int`` : Specify one edge using its endpoints (of source node type
and destination node type respectively).
* ``int iterable`` / ``tensor`` : Specify multiple edges using their edge ids.
* ``pair of int iterable`` / ``pair of tensors`` :
Specify multiple edges using their endpoints.
**Only works if the graph has one edge type.** For multiple types, use
.. code::
g['edgetype'].send(edges, message_func)
The UDF returns messages on the edges and can be later fetched in
the destination node's ``mailbox``. Receiving will consume the messages.
See :func:`recv` for example.
If multiple ``send`` are triggered on the same edge without ``recv``. Messages
generated by the later ``send`` will overwrite previous messages.
Parameters
----------
edges : optional
Edges on which to apply ``message_func``.
message_func : callable
Message function on the edges. The function should be
an :mod:`Edge UDF <dgl.udf>`.
Notes
-----
On multigraphs, if :math:`u` and :math:`v` are specified, then the messages will be sent
along all edges between :math:`u` and :math:`v`.
Examples
--------
>>> import dgl.function as fn
>>> import torch
>>> g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
>>> g.nodes['user'].data['h'] = torch.tensor([[0.], [1.], [2.]])
Different ways for sending messages.
>>> # Send the feature of source nodes along all edges
>>> g.send(g.edges(), fn.copy_src('h', 'm'))
>>> # Send the feature of source node along one edge specified by its id
>>> g.send(0, fn.copy_src('h', 'm'))
>>> # Send the feature of source node along one edge specified by its end points
>>> g.send((0, 1), fn.copy_src('h', 'm'))
>>> # Send the feature of source nodes along multiple edges specified by their ids
>>> g.send([0, 1], fn.copy_src('h', 'm'))
>>> # Send the feature of source nodes along multiple edges specified by their end points
>>> g.send(([0, 1], [1, 2]), fn.copy_src('h', 'm'))
"""
assert message_func is not None
etid = self.get_etype_id(etype)
stid, dtid = self._graph.metagraph.find_edge(etid)
if is_all(edges):
eid = utils.toindex(slice(0, self._graph.number_of_edges(etid)))
u, v, _ = self._graph.edges(etid, 'eid')
elif isinstance(edges, tuple):
u, v = edges
u = utils.toindex(u)
v = utils.toindex(v)
# Rewrite u, v to handle edge broadcasting and multigraph.
u, v, eid = self._graph.edge_ids(etid, u, v)
else:
eid = utils.toindex(edges)
u, v, _ = self._graph.find_edges(etid, eid)
if len(eid) == 0:
# no edge to be triggered
return
with ir.prog() as prog:
scheduler.schedule_send(
AdaptedHeteroGraph(self, stid, dtid, etid),
u, v, eid,
message_func)
Runtime.run(prog)
def recv(self,
v,
reduce_func,
apply_node_func=None,
etype=None,
inplace=False):
r"""Receive and reduce incoming messages and update the features of node(s) :math:`v`.
It calculates:
.. math::
h_v^{new} = \sigma(f(\{m_{uv} | u\in\mathcal{N}_{t}(v)\}))
where :math:`\mathcal{N}_t(v)` defines the predecessors of node(s) :math:`v` connected by
edges of type :math:`t`, and :math:`m_{uv}` is the message on edge :math:`(u,v)`.
* ``reduce_func`` specifies :math:`f`, e.g. summation or average.
* ``apply_node_func`` specifies :math:`\sigma`, e.g. ReLU activation.
Other notes:
* `reduce_func` will be skipped for nodes with no incoming message.
* If all ``v`` have no incoming message, this will downgrade to an :func:`apply_nodes`.
* If some ``v`` have no incoming message, their new feature value will be calculated
by the column initializer (see :func:`set_n_initializer`). The feature shapes and
dtypes will be inferred.
* The node features will be updated by the result of the ``reduce_func``.
* Messages are consumed once received.
* The provided UDF may be called multiple times so it is recommended to provide
function with no side effect.
Parameters
----------
v : int, container or tensor
The node(s) to be updated.
reduce_func : callable
Reduce function on the node. The function should be
a :mod:`Node UDF <dgl.udf>`.
apply_node_func : callable
Apply function on the nodes. The function should be
a :mod:`Node UDF <dgl.udf>`. (Default: None)
etype : str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
inplace: bool, optional
If True, update will be done in place, but autograd will break.
(Default: False)
Examples
--------
>>> import dgl
>>> import dgl.function as fn
>>> import torch
Instantiate a heterograph.
>>> follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
>>> plays_g = dgl.bipartite([(0, 0), (1, 0), (1, 1), (2, 1)], 'user', 'plays', 'game')
>>> g = dgl.hetero_from_relations([follows_g, plays_g])
>>> g.nodes['user'].data['h'] = torch.tensor([[0.], [1.], [2.]])
Send and receive.
>>> g.send(g['follows'].edges(), fn.copy_src('h', 'm'), etype='follows')
>>> g.recv(g.nodes('user'), fn.sum('m', 'h'), etype='follows')
>>> g.nodes['user'].data['h']
tensor([[0.],
[0.],
[1.]])
"""
etid = self.get_etype_id(etype)
stid, dtid = self._graph.metagraph.find_edge(etid)
if is_all(v):
v = F.arange(0, self.number_of_nodes(dtid))
elif isinstance(v, int):
v = [v]
v = utils.toindex(v)
if len(v) == 0:
# no vertex to be triggered.
return
with ir.prog() as prog:
scheduler.schedule_recv(AdaptedHeteroGraph(self, stid, dtid, etid),
v, reduce_func, apply_node_func,
inplace=inplace)
Runtime.run(prog)
def multi_recv(self, v, reducer_dict, cross_reducer, apply_node_func=None, inplace=False):
r"""Receive messages from multiple edge types and perform aggregation.
It calculates:
.. math::
\begin{align}
h_{v, t}^{new} &= f\left(\left\{m_{uv} | u\in\mathcal{N}_{t}(v)\right\}\right)\\
h_v^{new} &= \sigma\left(g\left(\left\{h_{v, t}^{new} | t\in T_e\right\}\right)\right)
\end{align}
* ``per_type_reducer`` is a dictionary mapping edge type (str or tuple of str) to
reduce functions :math:`f` of each type.
* ``cross_reducer`` specifies :math:`g`.
* ``apply_node_func`` specifies :math:`\sigma`.
Parameters
----------
v : int, container or tensor
The node(s) to be updated.
reducer_dict : dict of callable
Mapping edge type (str or tuple of str) to reduce function (:mod:`Node UDF <dgl.udf>`).
cross_reducer : str
Cross type reducer. One of ``"sum"``, ``"min"``, ``"max"``, ``"mean"``, ``"stack"``.
apply_node_func : callable
Apply function on the nodes. The function should be
a :mod:`Node UDF <dgl.udf>`. (Default: None)
inplace: bool, optional
If True, update will be done in place, but autograd will break.
(Default: False)
Examples
--------
>>> import dgl
>>> import dgl.function as fn
>>> import torch
Instantiate a heterograph.
>>> g1 = dgl.graph([(0, 1)], 'user', 'follows')
>>> g2 = dgl.bipartite([(0, 1)], 'game', 'attracts', 'user')
>>> g = dgl.hetero_from_relations([g1, g2])
>>> g.nodes['user'].data['h'] = torch.tensor([[1.], [2.]])
>>> g.nodes['game'].data['h'] = torch.tensor([[1.]])
Send and receive.
>>> g.send(g['follows'].edges(), fn.copy_src('h', 'm'), etype='follows')
>>> g.send(g['attracts'].edges(), fn.copy_src('h', 'm'), etype='attracts')
>>> g.multi_recv(g.nodes('user'), {'follows': fn.sum('m', 'h'),
>>> 'attracts': fn.sum('m', 'h')}, "sum")
>>> g.nodes['user'].data['h']
tensor([[0.],
[2.]])
"""
# infer receive node type
ntype = infer_ntype_from_dict(self, reducer_dict)
ntid = self.get_ntype_id_from_dst(ntype)
if is_all(v):
v = F.arange(0, self.number_of_nodes(ntid))
elif isinstance(v, int):
v = [v]
v = utils.toindex(v)
if len(v) == 0:
return
# TODO(minjie): currently loop over each edge type and reuse the old schedule.
# Should replace it with fused kernel.
all_out = []
merge_order = []
with ir.prog() as prog:
for ety, args in reducer_dict.items():
outframe = FrameRef(frame_like(self._node_frames[ntid]._frame))
args = pad_tuple(args, 2)
if args is None:
raise DGLError('Invalid per-type arguments. Should be either '
'(1) reduce_func or (2) (reduce_func, apply_node_func)')
rfunc, afunc = args
etid = self.get_etype_id(ety)
stid, dtid = self._graph.metagraph.find_edge(etid)
scheduler.schedule_recv(AdaptedHeteroGraph(self, stid, dtid, etid),
v, rfunc, afunc,
inplace=inplace, outframe=outframe)
all_out.append(outframe)
merge_order.append(etid) # use edge type id as merge order hint
Runtime.run(prog)
# merge by cross_reducer
self._node_frames[ntid].update(merge_frames(all_out, cross_reducer, merge_order))
# apply
if apply_node_func is not None:
self.apply_nodes(apply_node_func, v, ntype, inplace)
def send_and_recv(self,
edges,
message_func,
reduce_func,
apply_node_func=None,
etype=None,
inplace=False):
"""Send messages along edges of the specified type, and let destinations
receive them.
Optionally, apply a function to update the node features after "receive".
This is a convenient combination for performing
:mod:`send <dgl.DGLHeteroGraph.send>` along the ``edges`` and
:mod:`recv <dgl.DGLHeteroGraph.recv>` for the destinations of the ``edges``.
**Only works if the graph has one edge type.** For multiple types, use
.. code::
g['edgetype'].send_and_recv(edges, message_func, reduce_func,
apply_node_func, inplace=inplace)
Parameters
----------
edges : See :func:`send` for valid edge specification.
Edges on which to apply ``func``.
message_func : callable
Message function on the edges. The function should be
an :mod:`Edge UDF <dgl.udf>`.
reduce_func : callable
Reduce function on the node. The function should be
a :mod:`Node UDF <dgl.udf>`.
apply_node_func : callable, optional
Apply function on the nodes. The function should be
a :mod:`Node UDF <dgl.udf>`. (Default: None)
etype : str or tuple of str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
inplace: bool, optional
If True, update will be done in place, but autograd will break.
(Default: False)
Examples
--------
>>> import dgl
>>> import dgl.function as fn
>>> import torch
Instantiate a heterograph.
>>> follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
>>> plays_g = dgl.bipartite([(0, 0), (1, 0), (1, 1), (2, 1)], 'user', 'plays', 'game')
>>> g = dgl.hetero_from_relations([follows_g, plays_g])
Trigger "send" and "receive" separately.
>>> g.nodes['user'].data['h'] = torch.tensor([[0.], [1.], [2.]])
>>> g.send(g['follows'].edges(), fn.copy_src('h', 'm'), etype='follows')
>>> g.recv(g.nodes('user'), fn.sum('m', 'h'), etype='follows')
>>> g.nodes['user'].data['h']
tensor([[0.],
[0.],
[1.]])
Trigger "send" and "receive" in one call.
>>> g.nodes['user'].data['h'] = torch.tensor([[0.], [1.], [2.]])
>>> g.send_and_recv(g['follows'].edges(), fn.copy_src('h', 'm'),
>>> fn.sum('m', 'h'), etype='follows')
>>> g.nodes['user'].data['h']
tensor([[0.],
[0.],
[1.]])
"""
etid = self.get_etype_id(etype)
stid, dtid = self._graph.metagraph.find_edge(etid)
if isinstance(edges, tuple):
u, v = edges
u = utils.toindex(u)
v = utils.toindex(v)
# Rewrite u, v to handle edge broadcasting and multigraph.
u, v, eid = self._graph.edge_ids(etid, u, v)
else:
eid = utils.toindex(edges)
u, v, _ = self._graph.find_edges(etid, eid)
if len(u) == 0:
# no edges to be triggered
return
with ir.prog() as prog:
scheduler.schedule_snr(AdaptedHeteroGraph(self, stid, dtid, etid),
(u, v, eid),
message_func, reduce_func, apply_node_func,
inplace=inplace)
Runtime.run(prog)
def multi_send_and_recv(self, etype_dict, cross_reducer, apply_node_func=None, inplace=False):
r"""Send and receive messages along multiple edge types and perform aggregation.
Optionally, apply a function to update the node features after "receive".
This is a convenient combination for performing multiple
:mod:`send <dgl.DGLHeteroGraph.send>` along edges of different types and
:mod:`multi_recv <dgl.DGLHeteroGraph.multi_recv>` for the destinations of all edges.
Parameters
----------
etype_dict : dict
Mapping an edge type (str or tuple of str) to the type specific
configuration (4-tuples). Each 4-tuple represents
(edges, msg_func, reduce_func, apply_node_func):
* edges: See send() for valid edge specification.
Edges on which to pass messages.
* msg_func: callable
Message function on the edges. The function should be
an :mod:`Edge UDF <dgl.udf>`.
* reduce_func: callable
Reduce function on the node. The function should be
a :mod:`Node UDF <dgl.udf>`.
* apply_node_func : callable, optional
Apply function on the nodes. The function should be
a :mod:`Node UDF <dgl.udf>`. (Default: None)
cross_reducer : str
Cross type reducer. One of ``"sum"``, ``"min"``, ``"max"``, ``"mean"``, ``"stack"``.
apply_node_func : callable
Apply function on the nodes. The function should be
a :mod:`Node UDF <dgl.udf>`. (Default: None)
inplace: bool, optional
If True, update will be done in place, but autograd will break.
(Default: False)
Examples
--------
>>> import dgl
>>> import dgl.function as fn
>>> import torch
Instantiate a heterograph.
>>> g1 = dgl.graph([(0, 1)], 'user', 'follows')
>>> g2 = dgl.bipartite([(0, 1)], 'game', 'attracts', 'user')
>>> g = dgl.hetero_from_relations([g1, g2])
Trigger send and recv separately.
>>> g.nodes['user'].data['h'] = torch.tensor([[1.], [2.]])
>>> g.nodes['game'].data['h'] = torch.tensor([[1.]])
>>> g.send(g['follows'].edges(), fn.copy_src('h', 'm'), etype='follows')
>>> g.send(g['attracts'].edges(), fn.copy_src('h', 'm'), etype='attracts')
>>> g.multi_recv(g.nodes('user'),
>>> {'follows': fn.sum('m', 'h'), 'attracts': fn.sum('m', 'h')}, "sum")
>>> g.nodes['user'].data['h']
tensor([[0.],
[2.]])
Trigger “send” and “receive” in one call.
>>> g.nodes['user'].data['h'] = torch.tensor([[1.], [2.]])
>>> g.nodes['game'].data['h'] = torch.tensor([[1.]])
>>> g.multi_send_and_recv(
>>> {'follows': (g['follows'].edges(), fn.copy_src('h', 'm'), fn.sum('m', 'h')),
>>> 'attracts': (g['attracts'].edges(), fn.copy_src('h', 'm'), fn.sum('m', 'h'))},
>>> "sum")
>>> g.nodes['user'].data['h']
tensor([[0.],
[2.]])
"""
# infer receive node type
ntype = infer_ntype_from_dict(self, etype_dict)
dtid = self.get_ntype_id_from_dst(ntype)
# TODO(minjie): currently loop over each edge type and reuse the old schedule.
# Should replace it with fused kernel.
all_out = []
all_vs = []
merge_order = []
with ir.prog() as prog:
for etype, args in etype_dict.items():
etid = self.get_etype_id(etype)
stid, _ = self._graph.metagraph.find_edge(etid)
outframe = FrameRef(frame_like(self._node_frames[dtid]._frame))
args = pad_tuple(args, 4)
if args is None:
raise DGLError('Invalid per-type arguments. Should be '
'(edges, msg_func, reduce_func, [apply_node_func])')
edges, mfunc, rfunc, afunc = args
if isinstance(edges, tuple):
u, v = edges
u = utils.toindex(u)
v = utils.toindex(v)
# Rewrite u, v to handle edge broadcasting and multigraph.
u, v, eid = self._graph.edge_ids(etid, u, v)
else:
eid = utils.toindex(edges)
u, v, _ = self._graph.find_edges(etid, eid)
all_vs.append(v)
if len(u) == 0:
# no edges to be triggered
continue
scheduler.schedule_snr(AdaptedHeteroGraph(self, stid, dtid, etid),
(u, v, eid),
mfunc, rfunc, afunc,
inplace=inplace, outframe=outframe)
all_out.append(outframe)
merge_order.append(etid) # use edge type id as merge order hint
Runtime.run(prog)
# merge by cross_reducer
self._node_frames[dtid].update(merge_frames(all_out, cross_reducer, merge_order))
# apply
if apply_node_func is not None:
dstnodes = F.unique(F.cat([x.tousertensor() for x in all_vs], 0))
self.apply_nodes(apply_node_func, dstnodes, ntype, inplace)
def pull(self,
v,
message_func,
reduce_func,
apply_node_func=None,
etype=None,
inplace=False):
"""Pull messages from the node(s)' predecessors and then update their features.
Optionally, apply a function to update the node features after receive.
This is equivalent to :mod:`send_and_recv <dgl.DGLHeteroGraph.send_and_recv>`
on the incoming edges of ``v`` with the specified type.
Other notes:
* `reduce_func` will be skipped for nodes with no incoming messages.
* If all ``v`` have no incoming message, this will downgrade to an :func:`apply_nodes`.
* If some ``v`` have no incoming message, their new feature value will be calculated
by the column initializer (see :func:`set_n_initializer`). The feature shapes and
dtypes will be inferred.
**Only works if the graph has one edge type.** For multiple types, use
.. code::
g['edgetype'].pull(v, message_func, reduce_func, apply_node_func, inplace=inplace)
Parameters
----------
v : int, container or tensor, optional
The node(s) to be updated.
message_func : callable
Message function on the edges. The function should be
an :mod:`Edge UDF <dgl.udf>`.
reduce_func : callable
Reduce function on the node. The function should be
a :mod:`Node UDF <dgl.udf>`.
apply_node_func : callable, optional
Apply function on the nodes. The function should be
a :mod:`Node UDF <dgl.udf>`. (Default: None)
etype : str or tuple of str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
inplace: bool, optional
If True, update will be done in place, but autograd will break.
(Default: False)
Examples
--------
>>> import dgl
>>> import dgl.function as fn
>>> import torch
Instantiate a heterograph.
>>> follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows')
>>> plays_g = dgl.bipartite([(0, 0), (2, 1)], 'user', 'plays', 'game')
>>> g = dgl.hetero_from_relations([follows_g, plays_g])
>>> g.nodes['user'].data['h'] = torch.tensor([[0.], [1.], [2.]])
Pull.
>>> g['follows'].pull(2, fn.copy_src('h', 'm'), fn.sum('m', 'h'), etype='follows')
>>> g.nodes['user'].data['h']
tensor([[0.],
[1.],
[1.]])
"""
# only one type of edges
etid = self.get_etype_id(etype)
stid, dtid = self._graph.metagraph.find_edge(etid)
v = utils.toindex(v)
if len(v) == 0:
return
with ir.prog() as prog:
scheduler.schedule_pull(AdaptedHeteroGraph(self, stid, dtid, etid),
v,
message_func, reduce_func, apply_node_func,
inplace=inplace)
Runtime.run(prog)
def multi_pull(self, v, etype_dict, cross_reducer, apply_node_func=None, inplace=False):
r"""Pull and receive messages of the given nodes along multiple edge types
and perform aggregation.
This is equivalent to :mod:`multi_send_and_recv <dgl.DGLHeteroGraph.multi_send_and_recv>`
on the incoming edges of ``v`` with the specified types.
Parameters
----------
v : int, container or tensor
The node(s) to be updated.
etype_dict : dict
Mapping an edge type (str or tuple of str) to the type specific
configuration (3-tuples). Each 3-tuple represents
(msg_func, reduce_func, apply_node_func):
* msg_func: callable
Message function on the edges. The function should be
an :mod:`Edge UDF <dgl.udf>`.
* reduce_func: callable
Reduce function on the nodes. The function should be
a :mod:`Node UDF <dgl.udf>`.
* apply_node_func : callable, optional
Apply function on the nodes. The function should be
a :mod:`Node UDF <dgl.udf>`. (Default: None)
cross_reducer : str
Cross type reducer. One of ``"sum"``, ``"min"``, ``"max"``, ``"mean"``, ``"stack"``.
apply_node_func : callable
Apply function on the nodes. The function should be
a :mod:`Node UDF <dgl.udf>`. (Default: None)
inplace: bool, optional
If True, update will be done in place, but autograd will break.
(Default: False)
Examples
--------
>>> import dgl
>>> import dgl.function as fn
>>> import torch
Instantiate a heterograph.
>>> g1 = dgl.graph([(1, 1), (1, 0)], 'user', 'follows')
>>> g2 = dgl.bipartite([(0, 1)], 'game', 'attracts', 'user')
>>> g = dgl.hetero_from_relations([g1, g2])
Pull.
>>> g.nodes['user'].data['h'] = torch.tensor([[1.], [2.]])
>>> g.nodes['game'].data['h'] = torch.tensor([[1.]])
>>> g.multi_pull(1,
>>> {'follows': (fn.copy_src('h', 'm'), fn.sum('m', 'h')),
>>> 'attracts': (fn.copy_src('h', 'm'), fn.sum('m', 'h'))},
>>> "sum")
>>> g.nodes['user'].data['h']
tensor([[0.],
[3.]])
"""
v = utils.toindex(v)
if len(v) == 0:
return
# infer receive node type
ntype = infer_ntype_from_dict(self, etype_dict)
dtid = self.get_ntype_id_from_dst(ntype)
# TODO(minjie): currently loop over each edge type and reuse the old schedule.
# Should replace it with fused kernel.
all_out = []
merge_order = []
with ir.prog() as prog:
for etype, args in etype_dict.items():
etid = self.get_etype_id(etype)
stid, _ = self._graph.metagraph.find_edge(etid)
outframe = FrameRef(frame_like(self._node_frames[dtid]._frame))
args = pad_tuple(args, 3)
if args is None:
raise DGLError('Invalid per-type arguments. Should be '
'(msg_func, reduce_func, [apply_node_func])')
mfunc, rfunc, afunc = args
scheduler.schedule_pull(AdaptedHeteroGraph(self, stid, dtid, etid),
v,
mfunc, rfunc, afunc,
inplace=inplace, outframe=outframe)
all_out.append(outframe)
merge_order.append(etid) # use edge type id as merge order hint
Runtime.run(prog)
# merge by cross_reducer
self._node_frames[dtid].update(merge_frames(all_out, cross_reducer, merge_order))
# apply
if apply_node_func is not None:
self.apply_nodes(apply_node_func, v, ntype, inplace)
def push(self,
u,
message_func,
reduce_func,
apply_node_func=None,
etype=None,
inplace=False):
"""Send message from the node(s) to their successors and update them.
This is equivalent to performing
:mod:`send_and_recv <DGLHeteroGraph.send_and_recv>` along the outbound
edges from ``u``.
**Only works if the graph has one edge type.** For multiple types, use
.. code::
g['edgetype'].push(u, message_func, reduce_func, apply_node_func, inplace=inplace)
Parameters
----------
u : int, container or tensor
The node(s) to push out messages.
message_func : callable
Message function on the edges. The function should be
an :mod:`Edge UDF <dgl.udf>`.
reduce_func : callable
Reduce function on the node. The function should be
a :mod:`Node UDF <dgl.udf>`.
apply_node_func : callable, optional
Apply function on the nodes. The function should be
a :mod:`Node UDF <dgl.udf>`. (Default: None)
etype : str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
inplace: bool, optional
If True, update will be done in place, but autograd will break.
(Default: False)
Examples
--------
>>> import dgl
>>> import dgl.function as fn
>>> import torch
Instantiate a heterograph.
>>> g = dgl.graph([(0, 1), (0, 2)], 'user', 'follows')
>>> g.nodes['user'].data['h'] = torch.tensor([[0.], [1.], [2.]])
Push.
>>> g['follows'].push(0, fn.copy_src('h', 'm'), fn.sum('m', 'h'), etype='follows')
>>> g.nodes['user'].data['h']
tensor([[0.],
[0.],
[0.]])
"""
# only one type of edges
etid = self.get_etype_id(etype)
stid, dtid = self._graph.metagraph.find_edge(etid)
u = utils.toindex(u)
if len(u) == 0:
return
with ir.prog() as prog:
scheduler.schedule_push(AdaptedHeteroGraph(self, stid, dtid, etid),
u,
message_func, reduce_func, apply_node_func,
inplace=inplace)
Runtime.run(prog)
def update_all(self,
message_func,
reduce_func,
apply_node_func=None,
etype=None):
"""Send messages through all edges and update all nodes.
Optionally, apply a function to update the node features after receive.
This is equivalent to
:mod:`send_and_recv <dgl.DGLHeteroGraph.send_and_recv>` over all edges
of the specified type.
**Only works if the graph has one edge type.** For multiple types, use
.. code::
g['edgetype'].update_all(message_func, reduce_func, apply_node_func)
Parameters
----------
message_func : callable
Message function on the edges. The function should be
an :mod:`Edge UDF <dgl.udf>`.
reduce_func : callable
Reduce function on the node. The function should be
a :mod:`Node UDF <dgl.udf>`.
apply_node_func : callable, optional
Apply function on the nodes. The function should be
a :mod:`Node UDF <dgl.udf>`. (Default: None)
etype : str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
Examples
--------
>>> import torch
>>> import dgl
>>> import dgl.function as fn
Instantiate a heterograph.
>>> g = dgl.graph([(0, 1), (1, 2), (2, 2)], 'user', 'follows')
Update all.
>>> g.nodes['user'].data['h'] = torch.tensor([[0.], [1.], [2.]])
>>> g['follows'].update_all(fn.copy_src('h', 'm'), fn.sum('m', 'h'), etype='follows')
>>> g.nodes['user'].data['h']
tensor([[0.],
[0.],
[3.]])
"""
# only one type of edges
etid = self.get_etype_id(etype)
stid, dtid = self._graph.metagraph.find_edge(etid)
with ir.prog() as prog:
scheduler.schedule_update_all(AdaptedHeteroGraph(self, stid, dtid, etid),
message_func, reduce_func,
apply_node_func)
Runtime.run(prog)
def multi_update_all(self, etype_dict, cross_reducer, apply_node_func=None):
r"""Send and receive messages along all edges.
This is equivalent to
:mod:`multi_send_and_recv <dgl.DGLHeteroGraph.multi_send_and_recv>`
over all edges.
Parameters
----------
etype_dict : dict
Mapping an edge type (str or tuple of str) to the type specific
configuration (3-tuples). Each 3-tuple represents
(msg_func, reduce_func, apply_node_func):
* msg_func: callable
Message function on the edges. The function should be
an :mod:`Edge UDF <dgl.udf>`.
* reduce_func: callable
Reduce function on the nodes. The function should be
a :mod:`Node UDF <dgl.udf>`.
* apply_node_func : callable, optional
Apply function on the nodes. The function should be
a :mod:`Node UDF <dgl.udf>`. (Default: None)
cross_reducer : str
Cross type reducer. One of ``"sum"``, ``"min"``, ``"max"``, ``"mean"``, ``"stack"``.
apply_node_func : callable
Apply function on the nodes. The function should be
a :mod:`Node UDF <dgl.udf>`. (Default: None)
inplace: bool, optional
If True, update will be done in place, but autograd will break.
(Default: False)
etype_dict : dict of callable
``update_all`` arguments per edge type.
Examples
--------
>>> import dgl
>>> import dgl.function as fn
>>> import torch
Instantiate a heterograph.
>>> g1 = dgl.graph([(0, 1), (1, 1)], 'user', 'follows')
>>> g2 = dgl.bipartite([(0, 1)], 'game', 'attracts', 'user')
>>> g = dgl.hetero_from_relations([g1, g2])
>>> g.nodes['user'].data['h'] = torch.tensor([[1.], [2.]])
>>> g.nodes['game'].data['h'] = torch.tensor([[1.]])
Update all.
>>> g.multi_update_all(
>>> {'follows': (fn.copy_src('h', 'm'), fn.sum('m', 'h')),
>>> 'attracts': (fn.copy_src('h', 'm'), fn.sum('m', 'h'))},
>>> "sum")
>>> g.nodes['user'].data['h']
tensor([[0.],
[4.]])
"""
# TODO(minjie): currently loop over each edge type and reuse the old schedule.
# Should replace it with fused kernel.
all_out = defaultdict(list)
merge_order = defaultdict(list)
with ir.prog() as prog:
for etype, args in etype_dict.items():
etid = self.get_etype_id(etype)
stid, dtid = self._graph.metagraph.find_edge(etid)
outframe = FrameRef(frame_like(self._node_frames[dtid]._frame))
args = pad_tuple(args, 3)
if args is None:
raise DGLError('Invalid per-type arguments. Should be '
'(msg_func, reduce_func, [apply_node_func])')
mfunc, rfunc, afunc = args
scheduler.schedule_update_all(AdaptedHeteroGraph(self, stid, dtid, etid),
mfunc, rfunc, afunc,
outframe=outframe)
all_out[dtid].append(outframe)
merge_order[dtid].append(etid) # use edge type id as merge order hint
Runtime.run(prog)
for dtid, frames in all_out.items():
# merge by cross_reducer
self._node_frames[dtid].update(
merge_frames(frames, cross_reducer, merge_order[dtid]))
# apply
if apply_node_func is not None:
self.apply_nodes(apply_node_func, ALL, self.ntypes[dtid], inplace=False)
def prop_nodes(self,
nodes_generator,
message_func,
reduce_func,
apply_node_func=None,
etype=None):
"""Propagate messages using graph traversal by sequentially triggering
:func:`pull()` on nodes.
The traversal order is specified by the ``nodes_generator``. It generates
node frontiers, which is a list or a tensor of nodes. The nodes in the
same frontier will be triggered together, while nodes in different frontiers
will be triggered according to the generating order.
Parameters
----------
nodes_generator : iterable, each element is a list or a tensor of node ids
The generator of node frontiers. It specifies which nodes perform
:func:`pull` at each timestep.
message_func : callable
Message function on the edges. The function should be
an :mod:`Edge UDF <dgl.udf>`.
reduce_func : callable
Reduce function on the node. The function should be
a :mod:`Node UDF <dgl.udf>`.
apply_node_func : callable, optional
Apply function on the nodes. The function should be
a :mod:`Node UDF <dgl.udf>`. (Default: None)
etype : str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
Examples
--------
>>> import torch
>>> import dgl
>>> import dgl.function as fn
Instantiate a heterogrph and perform multiple rounds of message passing.
>>> g = dgl.graph(([0, 1, 2, 3], [2, 3, 4, 4]), 'user', 'follows')
>>> g.nodes['user'].data['h'] = torch.tensor([[1.], [2.], [3.], [4.], [5.]])
>>> g['follows'].prop_nodes([[2, 3], [4]], fn.copy_src('h', 'm'),
>>> fn.sum('m', 'h'), etype='follows')
tensor([[1.],
[2.],
[1.],
[2.],
[3.]])
See Also
--------
prop_edges
"""
for node_frontier in nodes_generator:
self.pull(node_frontier, message_func, reduce_func, apply_node_func, etype=etype)
def prop_edges(self,
edges_generator,
message_func,
reduce_func,
apply_node_func=None,
etype=None):
"""Propagate messages using graph traversal by sequentially triggering
:func:`send_and_recv()` on edges.
The traversal order is specified by the ``edges_generator``. It generates
edge frontiers. The edge frontiers should be of *valid edges type*.
See :func:`send` for more details.
Edges in the same frontier will be triggered together, and edges in
different frontiers will be triggered according to the generating order.
Parameters
----------
edges_generator : generator
The generator of edge frontiers.
message_func : callable
Message function on the edges. The function should be
an :mod:`Edge UDF <dgl.udf>`.
reduce_func : callable
Reduce function on the node. The function should be
a :mod:`Node UDF <dgl.udf>`.
apply_node_func : callable, optional
Apply function on the nodes. The function should be
a :mod:`Node UDF <dgl.udf>`. (Default: None)
etype : str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
Examples
--------
>>> import torch
>>> import dgl
>>> import dgl.function as fn
Instantiate a heterogrph and perform multiple rounds of message passing.
>>> g = dgl.graph(([0, 1, 2, 3], [2, 3, 4, 4]), 'user', 'follows')
>>> g.nodes['user'].data['h'] = torch.tensor([[1.], [2.], [3.], [4.], [5.]])
>>> g['follows'].prop_edges([[0, 1], [2, 3]], fn.copy_src('h', 'm'),
>>> fn.sum('m', 'h'), etype='follows')
>>> g.nodes['user'].data['h']
tensor([[1.],
[2.],
[1.],
[2.],
[3.]])
See Also
--------
prop_nodes
"""
for edge_frontier in edges_generator:
self.send_and_recv(edge_frontier, message_func, reduce_func,
apply_node_func, etype=etype)
#################################################################
# Misc
#################################################################
def to_networkx(self, node_attrs=None, edge_attrs=None):
"""Convert this graph to networkx graph.
The edge id will be saved as the 'id' edge attribute.
Parameters
----------
node_attrs : iterable of str, optional
The node attributes to be copied.
edge_attrs : iterable of str, optional
The edge attributes to be copied.
Returns
-------
networkx.DiGraph
The nx graph
Examples
--------
.. note:: Here we use pytorch syntax for demo. The general idea applies
to other frameworks with minor syntax change (e.g. replace
``torch.tensor`` with ``mxnet.ndarray``).
>>> import torch as th
>>> g = DGLGraph()
>>> g.add_nodes(5, {'n1': th.randn(5, 10)})
>>> g.add_edges([0,1,3,4], [2,4,0,3], {'e1': th.randn(4, 6)})
>>> nxg = g.to_networkx(node_attrs=['n1'], edge_attrs=['e1'])
See Also
--------
dgl.to_networkx
"""
# TODO(minjie): multi-type support
assert len(self.ntypes) == 1
assert len(self.etypes) == 1
src, dst = self.edges()
src = F.asnumpy(src)
dst = F.asnumpy(dst)
# xiangsx: Always treat graph as multigraph
nx_graph = nx.MultiDiGraph()
nx_graph.add_nodes_from(range(self.number_of_nodes()))
for eid, (u, v) in enumerate(zip(src, dst)):
nx_graph.add_edge(u, v, id=eid)
if node_attrs is not None:
for nid, attr in nx_graph.nodes(data=True):
feat_dict = self._get_n_repr(0, nid)
attr.update({key: F.squeeze(feat_dict[key], 0) for key in node_attrs})
if edge_attrs is not None:
for _, _, attr in nx_graph.edges(data=True):
eid = attr['id']
feat_dict = self._get_e_repr(0, eid)
attr.update({key: F.squeeze(feat_dict[key], 0) for key in edge_attrs})
return nx_graph
def filter_nodes(self, predicate, nodes=ALL, ntype=None):
"""Return a tensor of node IDs with the given node type that satisfy
the given predicate.
Parameters
----------
predicate : callable
A function of signature ``func(nodes) -> tensor``.
``nodes`` are :class:`NodeBatch` objects as in :mod:`~dgl.udf`.
The ``tensor`` returned should be a 1-D boolean tensor with
each element indicating whether the corresponding node in
the batch satisfies the predicate.
nodes : int, iterable or tensor of ints
The nodes to filter on. Default value is all the nodes.
ntype : str, optional
The node type. Can be omitted if there is only one node type
in the graph. (Default: None)
Returns
-------
tensor
Node ids indicating the nodes that satisfy the predicate.
Examples
--------
>>> import torch
>>> import dgl
>>> import dgl.function as fn
>>> g = dgl.graph([], 'user', 'follows', num_nodes=4)
>>> g.nodes['user'].data['h'] = torch.tensor([[0.], [1.], [1.], [0.]])
>>> g.filter_nodes(lambda nodes: (nodes.data['h'] == 1.).squeeze(1), ntype='user')
tensor([1, 2])
"""
ntid = self.get_ntype_id(ntype)
if is_all(nodes):
v = utils.toindex(slice(0, self._graph.number_of_nodes(ntid)))
else:
v = utils.toindex(nodes)
n_repr = self._get_n_repr(ntid, v)
nbatch = NodeBatch(v, n_repr, ntype=self.ntypes[ntid])
n_mask = F.copy_to(predicate(nbatch), F.cpu())
if is_all(nodes):
return F.nonzero_1d(n_mask)
else:
nodes = F.tensor(nodes)
return F.boolean_mask(nodes, n_mask)
def filter_edges(self, predicate, edges=ALL, etype=None):
"""Return a tensor of edge IDs with the given edge type that satisfy
the given predicate.
Parameters
----------
predicate : callable
A function of signature ``func(edges) -> tensor``.
``edges`` are :class:`EdgeBatch` objects as in :mod:`~dgl.udf`.
The ``tensor`` returned should be a 1-D boolean tensor with
each element indicating whether the corresponding edge in
the batch satisfies the predicate.
edges : valid edges type
Edges on which to apply ``func``. See :func:`send` for valid
edges type. Default value is all the edges.
etype : str, optional
The edge type. Can be omitted if there is only one edge type
in the graph. (Default: None)
Returns
-------
tensor
Edge ids indicating the edges that satisfy the predicate.
Examples
--------
>>> import torch
>>> import dgl
>>> import dgl.function as fn
>>> g = dgl.graph([(0, 0), (0, 1), (1, 2), (2, 3)], 'user', 'follows')
>>> g.edges['follows'].data['h'] = torch.tensor([[0.], [1.], [1.], [0.]])
>>> g.filter_edges(lambda edges: (edges.data['h'] == 1.).squeeze(1), etype='follows')
tensor([1, 2])
"""
etid = self.get_etype_id(etype)
stid, dtid = self._graph.metagraph.find_edge(etid)
if is_all(edges):
u, v, _ = self._graph.edges(etid, 'eid')
eid = utils.toindex(slice(0, self._graph.number_of_edges(etid)))
elif isinstance(edges, tuple):
u, v = edges
u = utils.toindex(u)
v = utils.toindex(v)
# Rewrite u, v to handle edge broadcasting and multigraph.
u, v, eid = self._graph.edge_ids(etid, u, v)
else:
eid = utils.toindex(edges)
u, v, _ = self._graph.find_edges(etid, eid)
src_data = self._get_n_repr(stid, u)
edge_data = self._get_e_repr(etid, eid)
dst_data = self._get_n_repr(dtid, v)
ebatch = EdgeBatch((u, v, eid), src_data, edge_data, dst_data,
canonical_etype=self.canonical_etypes[etid])
e_mask = F.copy_to(predicate(ebatch), F.cpu())
if is_all(edges):
return F.nonzero_1d(e_mask)
else:
edges = F.tensor(edges)
return F.boolean_mask(edges, e_mask)
def to(self, ctx): # pylint: disable=invalid-name
"""Move both ndata and edata to the targeted mode (cpu/gpu)
Framework agnostic
Parameters
----------
ctx : framework-specific context object
The context to move data to.
Returns
-------
g : DGLHeteroGraph
Moved DGLHeteroGraph of the targeted mode.
Examples
--------
The following example uses PyTorch backend.
>>> import torch
>>> g = dgl.bipartite([(0, 0), (1, 0), (1, 2), (2, 1)], 'user', 'plays', 'game')
>>> g.nodes['user'].data['h'] = torch.tensor([[0.], [1.], [2.]])
>>> g.edges['plays'].data['h'] = torch.tensor([[0.], [1.], [2.], [3.]])
>>> g = g.to(torch.device('cuda:0'))
"""
for i in range(len(self._node_frames)):
for k in self._node_frames[i].keys():
self._node_frames[i][k] = F.copy_to(self._node_frames[i][k], ctx)
for i in range(len(self._edge_frames)):
for k in self._edge_frames[i].keys():
self._edge_frames[i][k] = F.copy_to(self._edge_frames[i][k], ctx)
return self
def local_var(self):
"""Return a heterograph object that can be used in a local function scope.
The returned graph object shares the feature data and graph structure of this graph.
However, any out-place mutation to the feature data will not reflect to this graph,
thus making it easier to use in a function scope.
If set, the local graph object will use same initializers for node features and
edge features.
Returns
-------
DGLHeteroGraph
The graph object that can be used as a local variable.
Notes
-----
Internally, the returned graph shares the same feature tensors, but construct a new
dictionary structure (aka. Frame) so adding/removing feature tensors from the returned
graph will not reflect to the original graph. However, inplace operations do change
the shared tensor values, so will be reflected to the original graph. This function
also has little overhead when the number of feature tensors in this graph is small.
Examples
--------
The following example uses PyTorch backend.
Avoid accidentally overriding existing feature data. This is quite common when
implementing a NN module:
>>> def foo(g):
>>> g = g.local_var()
>>> g.edata['h'] = torch.ones((g.number_of_edges(), 3))
>>> return g.edata['h']
>>>
>>> g = dgl.bipartite([(0, 0), (1, 0), (1, 2)], 'user', 'plays', 'game')
>>> g.edata['h'] = torch.zeros((g.number_of_edges(), 3))
>>> newh = foo(g) # get tensor of all ones
>>> print(g.edata['h']) # still get tensor of all zeros
Automatically garbage collect locally-defined tensors without the need to manually
``pop`` the tensors.
>>> def foo(g):
>>> g = g.local_var()
>>> # This 'h' feature will stay local and be GCed when the function exits
>>> g.edata['h'] = torch.ones((g.number_of_edges(), 3))
>>> return g.edata['h']
>>>
>>> g = dgl.bipartite([(0, 0), (1, 0), (1, 2)], 'user', 'plays', 'game')
>>> h = foo(g)
>>> print('h' in g.edata)
False
See Also
--------
local_var
"""
local_node_frames = [FrameRef(Frame(fr._frame)) for fr in self._node_frames]
local_edge_frames = [FrameRef(Frame(fr._frame)) for fr in self._edge_frames]
# Use same per-column initializers and default initializer.
# If registered, a column (based on key) initializer will be used first,
# otherwise the default initializer will be used.
for fr1, fr2 in zip(local_node_frames, self._node_frames):
sync_frame_initializer(fr1._frame, fr2._frame)
for fr1, fr2 in zip(local_edge_frames, self._edge_frames):
sync_frame_initializer(fr1._frame, fr2._frame)
return DGLHeteroGraph(self._graph, self.ntypes, self.etypes,
local_node_frames,
local_edge_frames)
@contextmanager
def local_scope(self):
"""Enter a local scope context for this graph.
By entering a local scope, any out-place mutation to the feature data will
not reflect to the original graph, thus making it easier to use in a function scope.
If set, the local scope will use same initializers for node features and
edge features.
Examples
--------
The following example uses PyTorch backend.
Avoid accidentally overriding existing feature data. This is quite common when
implementing a NN module:
>>> def foo(g):
>>> with g.local_scope():
>>> g.edata['h'] = torch.ones((g.number_of_edges(), 3))
>>> return g.edata['h']
>>>
>>> g = dgl.bipartite([(0, 0), (1, 0), (1, 2)], 'user', 'plays', 'game')
>>> g.edata['h'] = torch.zeros((g.number_of_edges(), 3))
>>> newh = foo(g) # get tensor of all ones
>>> print(g.edata['h']) # still get tensor of all zeros
Automatically garbage collect locally-defined tensors without the need to manually
``pop`` the tensors.
>>> def foo(g):
>>> with g.local_scope():
>>> # This 'h' feature will stay local and be GCed when the function exits
>>> g.edata['h'] = torch.ones((g.number_of_edges(), 3))
>>> return g.edata['h']
>>>
>>> g = dgl.bipartite([(0, 0), (1, 0), (1, 2)], 'user', 'plays', 'game')
>>> h = foo(g)
>>> print('h' in g.edata)
False
See Also
--------
local_var
"""
old_nframes = self._node_frames
old_eframes = self._edge_frames
self._node_frames = [FrameRef(Frame(fr._frame)) for fr in self._node_frames]
self._edge_frames = [FrameRef(Frame(fr._frame)) for fr in self._edge_frames]
# Use same per-column initializers and default initializer.
# If registered, a column (based on key) initializer will be used first,
# otherwise the default initializer will be used.
for fr1, fr2 in zip(self._node_frames, old_nframes):
sync_frame_initializer(fr1._frame, fr2._frame)
for fr1, fr2 in zip(self._edge_frames, old_eframes):
sync_frame_initializer(fr1._frame, fr2._frame)
yield
self._node_frames = old_nframes
self._edge_frames = old_eframes
def is_homograph(self):
"""Return if the graph is homogeneous."""
return len(self.ntypes) == 1 and len(self.etypes) == 1
############################################################
# Internal APIs
############################################################
def make_canonical_etypes(etypes, ntypes, metagraph):
"""Internal function to convert etype name to (srctype, etype, dsttype)
Parameters
----------
etypes : list of str
Edge type list
ntypes : list of str
Node type list
metagraph : GraphIndex
Meta graph.
Returns
-------
list of tuples (srctype, etype, dsttype)
"""
# sanity check
if len(etypes) != metagraph.number_of_edges():
raise DGLError('Length of edge type list must match the number of '
'edges in the metagraph. {} vs {}'.format(
len(etypes), metagraph.number_of_edges()))
if len(ntypes) != metagraph.number_of_nodes():
raise DGLError('Length of nodes type list must match the number of '
'nodes in the metagraph. {} vs {}'.format(
len(ntypes), metagraph.number_of_nodes()))
src, dst, eid = metagraph.edges()
rst = [(ntypes[sid], etypes[eid], ntypes[did]) for sid, did, eid in zip(src, dst, eid)]
return rst
def is_unibipartite(graph):
"""Internal function that returns whether the given graph is a uni-directional
bipartite graph.
Parameters
----------
graph : GraphIndex
Input graph
Returns
-------
bool
True if the graph is a uni-bipartite.
"""
src, dst, _ = graph.edges()
return set(src.tonumpy()).isdisjoint(set(dst.tonumpy()))
def find_src_dst_ntypes(ntypes, metagraph):
"""Internal function to split ntypes into SRC and DST categories.
If the metagraph is not a uni-bipartite graph (so that the SRC and DST categories
are not well-defined), return None.
For node types that are isolated (i.e, no relation is associated with it), they
are assigned to the SRC category.
Parameters
----------
ntypes : list of str
Node type list
metagraph : GraphIndex
Meta graph.
Returns
-------
(dict[int, str], dict[int, str]) or None
Node types belonging to SRC and DST categories. Types are stored in
a dictionary from type name to type id. Return None if the graph is
not uni-bipartite.
"""
src, dst, _ = metagraph.edges()
if set(src.tonumpy()).isdisjoint(set(dst.tonumpy())):
srctypes = {ntypes[tid] : tid for tid in src}
dsttypes = {ntypes[tid] : tid for tid in dst}
# handle isolated node types
for ntid, ntype in enumerate(ntypes):
if ntype not in srctypes and ntype not in dsttypes:
srctypes[ntype] = ntid
return srctypes, dsttypes
else:
return None
def infer_ntype_from_dict(graph, etype_dict):
"""Infer node type from dictionary of edge type to values.
All the edge types in the dict must share the same destination node type
and the node type will be returned. Otherwise, throw error.
Parameters
----------
graph : DGLHeteroGraph
Graph
etype_dict : dict
Dictionary whose key is edge type
Returns
-------
str
Node type
"""
ntype = None
for ety in etype_dict:
_, _, dty = graph.to_canonical_etype(ety)
if ntype is None:
ntype = dty
if ntype != dty:
raise DGLError("Cannot infer destination node type from the dictionary. "
"A valid specification must make sure that all the edge "
"type keys share the same destination node type.")
return ntype
def pad_tuple(tup, length, pad_val=None):
"""Pad the given tuple to the given length.
If the input is not a tuple, convert it to a tuple of length one.
Return None if pad fails.
"""
if not isinstance(tup, tuple):
tup = (tup, )
if len(tup) > length:
return None
elif len(tup) == length:
return tup
else:
return tup + (pad_val,) * (length - len(tup))
def merge_frames(frames, reducer, order=None):
"""Merge input frames into one. Resolve conflict fields using reducer.
Parameters
----------
frames : list[FrameRef]
Input frames
reducer : str
One of "sum", "max", "min", "mean", "stack"
order : list[Int], optional
Merge order hint. Useful for "stack" reducer.
If provided, each integer indicates the relative order
of the ``frames`` list. Frames are sorted according to this list
in ascending order. Tie is not handled so make sure the order values
are distinct.
Returns
-------
FrameRef
Merged frame
"""
if len(frames) == 1 and reducer != 'stack':
# Directly return the only one input. Stack reducer requires
# modifying tensor shape.
return frames[0]
if reducer == 'stack':
# Stack order does not matter. However, it must be consistent!
if order:
assert len(order) == len(frames)
sorted_with_key = sorted(zip(frames, order), key=lambda x: x[1])
frames = list(zip(*sorted_with_key))[0]
def merger(flist):
return F.stack(flist, 1)
else:
redfn = getattr(F, reducer, None)
if redfn is None:
raise DGLError('Invalid cross type reducer. Must be one of '
'"sum", "max", "min", "mean" or "stack".')
def merger(flist):
return redfn(F.stack(flist, 0), 0) if len(flist) > 1 else flist[0]
ret = FrameRef(frame_like(frames[0]._frame))
keys = set()
for frm in frames:
keys.update(frm.keys())
for k in keys:
flist = []
for frm in frames:
if k in frm:
flist.append(frm[k])
ret[k] = merger(flist)
return ret
def combine_frames(frames, ids):
"""Merge the frames into one frame, taking the common columns.
Return None if there is no common columns.
Parameters
----------
frames : List[FrameRef]
List of frames
ids : List[int]
List of frame IDs
Returns
-------
FrameRef
The resulting frame
"""
# find common columns and check if their schemes match
schemes = {key: scheme for key, scheme in frames[ids[0]].schemes.items()}
for frame_id in ids:
frame = frames[frame_id]
for key, scheme in list(schemes.items()):
if key in frame.schemes:
if frame.schemes[key] != scheme:
raise DGLError('Cannot concatenate column %s with shape %s and shape %s' %
(key, frame.schemes[key], scheme))
else:
del schemes[key]
if len(schemes) == 0:
return None
# concatenate the columns
to_cat = lambda key: [frames[i][key] for i in ids if frames[i].num_rows > 0]
cols = {key: F.cat(to_cat(key), dim=0) for key in schemes}
return FrameRef(Frame(cols))
def combine_names(names, ids=None):
"""Combine the selected names into one new name.
Parameters
----------
names : list of str
String names
ids : numpy.ndarray, optional
Selected index
Returns
-------
str
"""
if ids is None:
return '+'.join(sorted(names))
else:
selected = sorted([names[i] for i in ids])
return '+'.join(selected)
class AdaptedHeteroGraph(GraphAdapter):
"""Adapt DGLGraph to interface required by scheduler.
Parameters
----------
graph : DGLHeteroGraph
Graph
stid : int
Source node type id
dtid : int
Destination node type id
etid : int
Edge type id
"""
def __init__(self, graph, stid, dtid, etid):
self.graph = graph
self.stid = stid
self.dtid = dtid
self.etid = etid
@property
def gidx(self):
return self.graph._graph
def num_src(self):
"""Number of source nodes."""
return self.graph._graph.number_of_nodes(self.stid)
def num_dst(self):
"""Number of destination nodes."""
return self.graph._graph.number_of_nodes(self.dtid)
def num_edges(self):
"""Number of edges."""
return self.graph._graph.number_of_edges(self.etid)
@property
def srcframe(self):
"""Frame to store source node features."""
return self.graph._node_frames[self.stid]
@property
def dstframe(self):
"""Frame to store source node features."""
return self.graph._node_frames[self.dtid]
@property
def edgeframe(self):
"""Frame to store edge features."""
return self.graph._edge_frames[self.etid]
@property
def msgframe(self):
"""Frame to store messages."""
return self.graph._msg_frames[self.etid]
@property
def msgindicator(self):
"""Message indicator tensor."""
return self.graph._get_msg_index(self.etid)
@msgindicator.setter
def msgindicator(self, val):
"""Set new message indicator tensor."""
self.graph._set_msg_index(self.etid, val)
def in_edges(self, nodes):
return self.graph._graph.in_edges(self.etid, nodes)
def out_edges(self, nodes):
return self.graph._graph.out_edges(self.etid, nodes)
def edges(self, form):
return self.graph._graph.edges(self.etid, form)
def get_immutable_gidx(self, ctx):
return self.graph._graph.get_unitgraph(self.etid, ctx)
def bits_needed(self):
return self.graph._graph.bits_needed(self.etid)
@property
def canonical_etype(self):
"""Canonical edge type."""
return self.graph.canonical_etypes[self.etid]