heterograph.py 225 KB
Newer Older
Da Zheng's avatar
Da Zheng committed
1
"""Classes for heterogeneous graphs."""
2
#pylint: disable= too-many-lines
3
4
from collections import defaultdict
from collections.abc import Mapping, Iterable
Minjie Wang's avatar
Minjie Wang committed
5
from contextlib import contextmanager
6
import copy
7
import numbers
8
import networkx as nx
Minjie Wang's avatar
Minjie Wang committed
9
10
import numpy as np

11
12
13
from ._ffi.function import _init_api
from .base import ALL, SLICE_FULL, NTYPE, NID, ETYPE, EID, is_all, DGLError, dgl_warning
from . import core
Minjie Wang's avatar
Minjie Wang committed
14
15
from . import graph_index
from . import heterograph_index
16
17
from . import utils
from . import backend as F
18
from .frame import Frame
19
from .view import HeteroNodeView, HeteroNodeDataView, HeteroEdgeView, HeteroEdgeDataView
Minjie Wang's avatar
Minjie Wang committed
20
21
22
23

__all__ = ['DGLHeteroGraph', 'combine_names']

class DGLHeteroGraph(object):
24
25
    """Class for storing graph structure and node/edge feature data.

26
    There are a few ways to create a DGLGraph:
27
28
29
30
31
32
33
34

    * To create a homogeneous graph from Tensor data, use :func:`dgl.graph`.
    * To create a heterogeneous graph from Tensor data, use :func:`dgl.heterograph`.
    * To create a graph from other data sources, use ``dgl.*`` create ops. See
      :ref:`api-graph-create-ops`.

    Read the user guide chapter :ref:`guide-graph` for an in-depth explanation about its
    usage.
Minjie Wang's avatar
Minjie Wang committed
35
    """
36
37
    is_block = False

38
    # pylint: disable=unused-argument, dangerous-default-value
Minjie Wang's avatar
Minjie Wang committed
39
    def __init__(self,
40
41
42
                 gidx=[],
                 ntypes=['_U'],
                 etypes=['_V'],
Minjie Wang's avatar
Minjie Wang committed
43
                 node_frames=None,
44
45
                 edge_frames=None,
                 **deprecate_kwargs):
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
        """Internal constructor for creating a DGLGraph.

        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[Frame], 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[Frame], optional
            Edge feature storage. If None, empty frame is created.
            Otherwise, ``edge_frames[i]`` stores the edge features
            of edge type i. (default: None)
        """
67
68
69
70
71
72
73
74
75
76
77
        if isinstance(gidx, DGLHeteroGraph):
            raise DGLError('The input is already a DGLGraph. No need to create it again.')
        if not isinstance(gidx, heterograph_index.HeteroGraphIndex):
            dgl_warning('Recommend creating graphs by `dgl.graph(data)`'
                        ' instead of `dgl.DGLGraph(data)`.')
            u, v, num_src, num_dst = utils.graphdata2tensors(gidx)
            gidx = heterograph_index.create_unitgraph_from_coo(
                1, num_src, num_dst, u, v, ['coo', 'csr', 'csc'])
        if len(deprecate_kwargs) != 0:
            dgl_warning('Keyword arguments {} are deprecated in v0.5, and can be safely'
                        ' removed in all cases.'.format(list(deprecate_kwargs.keys())))
78
        self._init(gidx, ntypes, etypes, node_frames, edge_frames)
Da Zheng's avatar
Da Zheng committed
79

80
81
    def _init(self, gidx, ntypes, etypes, node_frames, edge_frames):
        """Init internal states."""
Minjie Wang's avatar
Minjie Wang committed
82
        self._graph = gidx
83
        self._canonical_etypes = None
84
85
        self._batch_num_nodes = None
        self._batch_num_edges = None
86
87
88
89
90
91
92
93
94
95
96
97
98
99

        # 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
100
101
            if len(ntypes[0]) == 1 and len(ntypes[1]) == 1 and len(etypes) == 1:
                self._canonical_etypes = [(ntypes[0][0], etypes[0], ntypes[1][0])]
102
103
        else:
            self._ntypes = ntypes
104
105
106
107
            if len(ntypes) == 1:
                src_dst_map = None
            else:
                src_dst_map = find_src_dst_ntypes(self._ntypes, self._graph.metagraph)
108
109
110
111
112
113
114
115
            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
116
        self._etypes = etypes
117
        if self._canonical_etypes is None:
118
119
120
121
122
            if (len(etypes) == 1 and len(ntypes) == 1):
                self._canonical_etypes = [(ntypes[0], etypes[0], ntypes[0])]
            else:
                self._canonical_etypes = make_canonical_etypes(
                    self._etypes, self._ntypes, self._graph.metagraph)
123

Minjie Wang's avatar
Minjie Wang committed
124
        # An internal map from etype to canonical etype tuple.
125
126
        # If two etypes have the same name, an empty tuple is stored instead to indicate
        # ambiguity.
Minjie Wang's avatar
Minjie Wang committed
127
        self._etype2canonical = {}
128
        for i, ety in enumerate(self._etypes):
Minjie Wang's avatar
Minjie Wang committed
129
130
131
132
133
            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)}
Da Zheng's avatar
Da Zheng committed
134

Minjie Wang's avatar
Minjie Wang committed
135
136
137
        # node and edge frame
        if node_frames is None:
            node_frames = [None] * len(self._ntypes)
138
        node_frames = [Frame(num_rows=self._graph.number_of_nodes(i))
Minjie Wang's avatar
Minjie Wang committed
139
140
141
                       if frame is None else frame
                       for i, frame in enumerate(node_frames)]
        self._node_frames = node_frames
Da Zheng's avatar
Da Zheng committed
142

Minjie Wang's avatar
Minjie Wang committed
143
144
        if edge_frames is None:
            edge_frames = [None] * len(self._etypes)
145
        edge_frames = [Frame(num_rows=self._graph.number_of_edges(i))
Minjie Wang's avatar
Minjie Wang committed
146
147
148
                       if frame is None else frame
                       for i, frame in enumerate(edge_frames)]
        self._edge_frames = edge_frames
Da Zheng's avatar
Da Zheng committed
149

150
    def __setstate__(self, state):
151
152
        # Compatibility check
        # TODO: version the storage
153
154
155
        if isinstance(state, dict):
            # Since 0.5 we use the default __dict__ method
            self.__dict__.update(state)
156
157
158
159
        elif isinstance(state, tuple) and len(state) == 5:
            # DGL == 0.4.3
            dgl_warning("The object is pickled with DGL == 0.4.3.  "
                        "Some of the original attributes are ignored.")
160
161
            self._init(*state)
        elif isinstance(state, dict):
162
163
            # DGL <= 0.4.2
            dgl_warning("The object is pickled with DGL <= 0.4.2.  "
164
165
166
167
168
                        "Some of the original attributes are ignored.")
            self._init(state['_graph'], state['_ntypes'], state['_etypes'], state['_node_frames'],
                       state['_edge_frames'])
        else:
            raise IOError("Unrecognized pickle format.")
Mufei Li's avatar
Mufei Li committed
169

Minjie Wang's avatar
Minjie Wang committed
170
171
172
173
174
175
176
177
178
179
180
181
182
183
    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))}
184
            nedge_dict = {self.canonical_etypes[i] : self._graph.number_of_edges(i)
Minjie Wang's avatar
Minjie Wang committed
185
                          for i in range(len(self.etypes))}
186
            meta = str(self.metagraph().edges(keys=True))
Minjie Wang's avatar
Minjie Wang committed
187
188
            return ret.format(node=nnode_dict, edge=nedge_dict, meta=meta)

189
190
191
192
193
    def __copy__(self):
        """Shallow copy implementation."""
        #TODO(minjie): too many states in python; should clean up and lower to C
        cls = type(self)
        obj = cls.__new__(cls)
194
        obj.__dict__.update(self.__dict__)
195
196
        return obj

Minjie Wang's avatar
Minjie Wang committed
197
198
199
200
201
    #################################################################
    # Mutation operations
    #################################################################

    def add_nodes(self, num, data=None, ntype=None):
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
        r"""Add new nodes of the same node type

        Parameters
        ----------
        num : int
            Number of nodes to add.
        data : dict, optional
            Feature data of the added nodes.
        ntype : str, optional
            The type of the new nodes. Can be omitted if there is
            only one node type in the graph.

        Notes
        -----

        * Inplace update is applied to the current graph.
        * If the key of ``data`` does not contain some existing feature fields,
219
220
          those features for the new nodes will be created by initializers
          defined with :func:`set_n_initializer` (default initializer fills zeros).
221
        * If the key of ``data`` contains new feature fields, those features for
222
223
224
225
226
227
          the old nodes will be created by initializers defined with
          :func:`set_n_initializer` (default initializer fills zeros).
        * This function discards the batch information. Please use
          :func:`dgl.DGLGraph.set_batch_num_nodes`
          and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
          to maintain the information.
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270

        Examples
        --------

        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        **Homogeneous Graphs or Heterogeneous Graphs with A Single Node Type**

        >>> g = dgl.graph((torch.tensor([0, 1]), torch.tensor([1, 2])))
        >>> g.num_nodes()
        3
        >>> g.add_nodes(2)
        >>> g.num_nodes()
        5

        If the graph has some node features and new nodes are added without
        features, their features will be created by initializers defined
        with :func:`set_n_initializer`.

        >>> g.ndata['h'] = torch.ones(5, 1)
        >>> g.add_nodes(1)
        >>> g.ndata['h']
        tensor([[1.], [1.], [1.], [1.], [1.], [0.]])

        We can also assign features for the new nodes in adding new nodes.

        >>> g.add_nodes(1, {'h': torch.ones(1, 1), 'w': torch.ones(1, 1)})
        >>> g.ndata['h']
        tensor([[1.], [1.], [1.], [1.], [1.], [0.], [1.]])

        Since ``data`` contains new feature fields, the features for old nodes
        will be created by initializers defined with :func:`set_n_initializer`.

        >>> g.ndata['w']
        tensor([[0.], [0.], [0.], [0.], [0.], [0.], [1.]])


        **Heterogeneous Graphs with Multiple Node Types**

        >>> g = dgl.heterograph({
271
272
273
274
275
        ...     ('user', 'plays', 'game'): (torch.tensor([0, 1, 1, 2]),
        ...                                 torch.tensor([0, 0, 1, 1])),
        ...     ('developer', 'develops', 'game'): (torch.tensor([0, 1]),
        ...                                         torch.tensor([0, 1]))
        ...     })
276
277
278
279
280
281
282
283
        >>> g.add_nodes(2)
        DGLError: Node type name must be specified
        if there are more than one node types.
        >>> g.num_nodes('user')
        3
        >>> g.add_nodes(2, ntype='user')
        >>> g.num_nodes('user')
        5
Minjie Wang's avatar
Minjie Wang committed
284

285
286
287
288
289
        See Also
        --------
        remove_nodes
        add_edges
        remove_edges
Minjie Wang's avatar
Minjie Wang committed
290
        """
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
        # TODO(xiangsx): block do not support add_nodes
        if ntype is None:
            if self._graph.number_of_ntypes() != 1:
                raise DGLError('Node type name must be specified if there are more than one '
                               'node types.')

        # nothing happen
        if num == 0:
            return

        assert num > 0, 'Number of new nodes should be larger than one.'
        ntid = self.get_ntype_id(ntype)
        # update graph idx
        metagraph = self._graph.metagraph
        num_nodes_per_type = []
        for c_ntype in self.ntypes:
            if self.get_ntype_id(c_ntype) == ntid:
                num_nodes_per_type.append(self.number_of_nodes(c_ntype) + num)
            else:
                num_nodes_per_type.append(self.number_of_nodes(c_ntype))

        relation_graphs = []
        for c_etype in self.canonical_etypes:
            # src or dst == ntype, update the relation graph
            if self.get_ntype_id(c_etype[0]) == ntid or self.get_ntype_id(c_etype[2]) == ntid:
                u, v = self.edges(form='uv', order='eid', etype=c_etype)
                hgidx = heterograph_index.create_unitgraph_from_coo(
                    1 if c_etype[0] == c_etype[2] else 2,
                    self.number_of_nodes(c_etype[0]) + \
                        (num if self.get_ntype_id(c_etype[0]) == ntid else 0),
                    self.number_of_nodes(c_etype[2]) + \
                        (num if self.get_ntype_id(c_etype[2]) == ntid else 0),
                    u,
                    v,
                    ['coo', 'csr', 'csc'])
                relation_graphs.append(hgidx)
            else:
                # do nothing
                relation_graphs.append(self._graph.get_relation_graph(self.get_etype_id(c_etype)))
        hgidx = heterograph_index.create_heterograph_from_relations(
            metagraph, relation_graphs, utils.toindex(num_nodes_per_type, "int64"))
        self._graph = hgidx

        # update data frames
        if data is None:
            # Initialize feature with :func:`set_n_initializer`
            self._node_frames[ntid].add_rows(num)
        else:
            self._node_frames[ntid].append(data)
        self._reset_cached_info()
Minjie Wang's avatar
Minjie Wang committed
341
342

    def add_edge(self, u, v, data=None, etype=None):
343
        """Add one edge to the graph.
Minjie Wang's avatar
Minjie Wang committed
344

345
        DEPRECATED: please use ``add_edges``.
Minjie Wang's avatar
Minjie Wang committed
346
        """
347
348
        dgl_warning("DGLGraph.add_edge is deprecated. Please use DGLGraph.add_edges")
        self.add_edges(u, v, data, etype)
Minjie Wang's avatar
Minjie Wang committed
349
350

    def add_edges(self, u, v, data=None, etype=None):
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
        r"""Add multiple new edges for the specified edge type

        The i-th new edge will be from ``u[i]`` to ``v[i]``.

        Parameters
        ----------
        u : int, tensor, numpy.ndarray, list
            Source node IDs, ``u[i]`` gives the source node for the i-th new edge.
        v : int, tensor, numpy.ndarray, list
            Destination node IDs, ``v[i]`` gives the destination node for the i-th new edge.
        data : dict, optional
            Feature data of the added edges. The i-th row of the feature data
            corresponds to the i-th new edge.
        etype : str or tuple of str, optional
            The type of the new edges. Can be omitted if there is
            only one edge type in the graph.

        Notes
        -----

        * Inplace update is applied to the current graph.
        * If end nodes of adding edges does not exists, add_nodes is invoked
373
374
375
376
          to add new nodes. The node features of the new nodes will be created
          by initializers defined with :func:`set_n_initializer` (default
          initializer fills zeros). In certain cases, it is recommanded to
          add_nodes first and then add_edges.
377
        * If the key of ``data`` does not contain some existing feature fields,
378
379
          those features for the new edges will be created by initializers
          defined with :func:`set_n_initializer` (default initializer fills zeros).
380
        * If the key of ``data`` contains new feature fields, those features for
381
382
383
384
385
386
          the old edges will be created by initializers defined with
          :func:`set_n_initializer` (default initializer fills zeros).
        * This function discards the batch information. Please use
          :func:`dgl.DGLGraph.set_batch_num_nodes`
          and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
          to maintain the information.
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421

        Examples
        --------

        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        **Homogeneous Graphs or Heterogeneous Graphs with A Single Edge Type**

        >>> g = dgl.graph((torch.tensor([0, 1]), torch.tensor([1, 2])))
        >>> g.num_edges()
        2
        >>> g.add_edges(torch.tensor([1, 3]), torch.tensor([0, 1]))
        >>> g.num_edges()
        4

        Since ``u`` or ``v`` contains a non-existing node ID, the nodes are
        added implicitly.
        >>> g.num_nodes()
        4

        If the graph has some edge features and new edges are added without
        features, their features will be created by initializers defined
        with :func:`set_n_initializer`.

        >>> g.edata['h'] = torch.ones(4, 1)
        >>> g.add_edges(torch.tensor([1]), torch.tensor([1]))
        >>> g.edata['h']
        tensor([[1.], [1.], [1.], [1.], [0.]])

        We can also assign features for the new edges in adding new edges.

        >>> g.add_edges(torch.tensor([0, 0]), torch.tensor([2, 2]),
422
        ...             {'h': torch.tensor([[1.], [2.]]), 'w': torch.ones(2, 1)})
423
424
425
426
427
428
429
430
431
432
433
434
        >>> g.edata['h']
        tensor([[1.], [1.], [1.], [1.], [0.], [1.], [2.]])

        Since ``data`` contains new feature fields, the features for old edges
        will be created by initializers defined with :func:`set_n_initializer`.

        >>> g.edata['w']
        tensor([[0.], [0.], [0.], [0.], [0.], [1.], [1.]])

        **Heterogeneous Graphs with Multiple Edge Types**

        >>> g = dgl.heterograph({
435
436
437
438
439
        ...     ('user', 'plays', 'game'): (torch.tensor([0, 1, 1, 2]),
        ...                                 torch.tensor([0, 0, 1, 1])),
        ...     ('developer', 'develops', 'game'): (torch.tensor([0, 1]),
        ...                                         torch.tensor([0, 1]))
        ...     })
440
441
442
443
444
        >>> g.add_edges(torch.tensor([3]), torch.tensor([3]))
        DGLError: Edge type name must be specified
        if there are more than one edge types.
        >>> g.number_of_edges('plays')
        4
445
        >>> g.add_edges(torch.tensor([3]), torch.tensor([3]), etype='plays')
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
        >>> g.number_of_edges('plays')
        5

        See Also
        --------
        add_nodes
        remove_nodes
        remove_edges
        """
        # TODO(xiangsx): block do not support add_edges
        u = utils.prepare_tensor(self, u, 'u')
        v = utils.prepare_tensor(self, v, 'v')

        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.')

        # nothing changed
        if len(u) == 0 or len(v) == 0:
            return

        assert len(u) == len(v) or len(u) == 1 or len(v) == 1, \
            'The number of source nodes and the number of destination nodes should be same, ' \
            'or either the number of source nodes or the number of destination nodes is 1.'

        if len(u) == 1 and len(v) > 1:
            u = F.full_1d(len(v), F.as_scalar(u), dtype=F.dtype(u), ctx=F.context(u))
        if len(v) == 1 and len(u) > 1:
            v = F.full_1d(len(u), F.as_scalar(v), dtype=F.dtype(v), ctx=F.context(v))

        u_type, e_type, v_type = self.to_canonical_etype(etype)
        # if end nodes of adding edges does not exists
        # use add_nodes to add new nodes first.
        num_of_u = self.number_of_nodes(u_type)
        num_of_v = self.number_of_nodes(v_type)
        u_max = F.as_scalar(F.max(u, dim=0)) + 1
        v_max = F.as_scalar(F.max(v, dim=0)) + 1

        if u_type == v_type:
            num_nodes = max(u_max, v_max)
            if num_nodes > num_of_u:
                self.add_nodes(num_nodes - num_of_u, ntype=u_type)
        else:
            if u_max > num_of_u:
                self.add_nodes(u_max - num_of_u, ntype=u_type)
            if v_max > num_of_v:
                self.add_nodes(v_max - num_of_v, ntype=v_type)

        # metagraph is not changed
        metagraph = self._graph.metagraph
        num_nodes_per_type = []
        for ntype in self.ntypes:
            num_nodes_per_type.append(self.number_of_nodes(ntype))
        # update graph idx
        relation_graphs = []
        for c_etype in self.canonical_etypes:
            # the target edge type
            if c_etype == (u_type, e_type, v_type):
                old_u, old_v = self.edges(form='uv', order='eid', etype=c_etype)
                hgidx = heterograph_index.create_unitgraph_from_coo(
                    1 if u_type == v_type else 2,
                    self.number_of_nodes(u_type),
                    self.number_of_nodes(v_type),
                    F.cat([old_u, u], dim=0),
                    F.cat([old_v, v], dim=0),
                    ['coo', 'csr', 'csc'])
                relation_graphs.append(hgidx)
            else:
                # do nothing
                # Note: node range change has been handled in add_nodes()
                relation_graphs.append(self._graph.get_relation_graph(self.get_etype_id(c_etype)))

        hgidx = heterograph_index.create_heterograph_from_relations(
            metagraph, relation_graphs, utils.toindex(num_nodes_per_type, "int64"))
        self._graph = hgidx

        # handle data
        etid = self.get_etype_id(etype)
        if data is None:
            self._edge_frames[etid].add_rows(len(u))
        else:
            self._edge_frames[etid].append(data)
        self._reset_cached_info()

531
    def remove_edges(self, eids, etype=None, store_ids=False):
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
        r"""Remove multiple edges with the specified edge type

        Nodes will not be removed. After removing edges, the rest
        edges will be re-indexed using consecutive integers from 0,
        with their relative order preserved.

        The features for the removed edges will be removed accordingly.

        Parameters
        ----------
        eids : int, tensor, numpy.ndarray, list
            IDs for the edges to remove.
        etype : str or tuple of str, optional
            The type of the edges to remove. Can be omitted if there is
            only one edge type in the graph.
547
548
549
550
        store_ids : bool, optional
            If True, it will store the raw IDs of the extracted nodes and edges in the ``ndata``
            and ``edata`` of the resulting graph under name ``dgl.NID`` and ``dgl.EID``,
            respectively.
551

552
553
554
555
556
557
558
559
        Notes
        -----

        This function discards the batch information. Please use
        :func:`dgl.DGLGraph.set_batch_num_nodes`
        and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
        to maintain the information.

560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
        Examples
        --------

        >>> import dgl
        >>> import torch

        **Homogeneous Graphs or Heterogeneous Graphs with A Single Edge Type**

        >>> g = dgl.graph((torch.tensor([0, 0, 2]), torch.tensor([0, 1, 2])))
        >>> g.edata['he'] = torch.arange(3).float().reshape(-1, 1)
        >>> g.remove_edges(torch.tensor([0, 1]))
        >>> g
        Graph(num_nodes=3, num_edges=1,
            ndata_schemes={}
            edata_schemes={'he': Scheme(shape=(1,), dtype=torch.float32)})
        >>> g.edges('all')
        (tensor([2]), tensor([2]), tensor([0]))
        >>> g.edata['he']
        tensor([[2.]])

        **Heterogeneous Graphs with Multiple Edge Types**

        >>> g = dgl.heterograph({
583
584
585
586
587
        ...     ('user', 'plays', 'game'): (torch.tensor([0, 1, 1, 2]),
        ...                                 torch.tensor([0, 0, 1, 1])),
        ...     ('developer', 'develops', 'game'): (torch.tensor([0, 1]),
        ...                                         torch.tensor([0, 1]))
        ...     })
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
        >>> g.remove_edges(torch.tensor([0, 1]))
        DGLError: Edge type name must be specified
        if there are more than one edge types.
        >>> g.remove_edges(torch.tensor([0, 1]), 'plays')
        >>> g.edges('all', etype='plays')
        (tensor([0, 1]), tensor([0, 0]), tensor([0, 1]))

        See Also
        --------
        add_nodes
        add_edges
        remove_nodes
        """
        # TODO(xiangsx): block do not support remove_edges
        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.')
        eids = utils.prepare_tensor(self, eids, 'u')
        if len(eids) == 0:
            # no edge to delete
            return
        assert self.number_of_edges(etype) > F.as_scalar(F.max(eids, dim=0)), \
            'The input eid {} is out of the range [0:{})'.format(
                F.as_scalar(F.max(eids, dim=0)), self.number_of_edges(etype))

        # edge_subgraph
        edges = {}
        u_type, e_type, v_type = self.to_canonical_etype(etype)
        for c_etype in self.canonical_etypes:
            # the target edge type
            if c_etype == (u_type, e_type, v_type):
                origin_eids = self.edges(form='eid', order='eid', etype=c_etype)
                edges[c_etype] = utils.compensate(eids, origin_eids)
            else:
                edges[c_etype] = self.edges(form='eid', order='eid', etype=c_etype)

625
        sub_g = self.edge_subgraph(edges, relabel_nodes=False, store_ids=store_ids)
626
627
628
629
        self._graph = sub_g._graph
        self._node_frames = sub_g._node_frames
        self._edge_frames = sub_g._edge_frames

630
    def remove_nodes(self, nids, ntype=None, store_ids=False):
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
        r"""Remove multiple nodes with the specified node type

        Edges that connect to the nodes will be removed as well. After removing
        nodes and edges, the rest nodes and edges will be re-indexed using
        consecutive integers from 0, with their relative order preserved.

        The features for the removed nodes/edges will be removed accordingly.

        Parameters
        ----------
        nids : int, tensor, numpy.ndarray, list
            Nodes to remove.
        ntype : str, optional
            The type of the nodes to remove. Can be omitted if there is
            only one node type in the graph.
646
647
648
649
        store_ids : bool, optional
            If True, it will store the raw IDs of the extracted nodes and edges in the ``ndata``
            and ``edata`` of the resulting graph under name ``dgl.NID`` and ``dgl.EID``,
            respectively.
650

651
652
653
654
655
656
657
658
        Notes
        -----

        This function discards the batch information. Please use
        :func:`dgl.DGLGraph.set_batch_num_nodes`
        and :func:`dgl.DGLGraph.set_batch_num_edges` on the transformed graph
        to maintain the information.

659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
        Examples
        --------

        >>> import dgl
        >>> import torch

        **Homogeneous Graphs or Heterogeneous Graphs with A Single Node Type**

        >>> g = dgl.graph((torch.tensor([0, 0, 2]), torch.tensor([0, 1, 2])))
        >>> g.ndata['hv'] = torch.arange(3).float().reshape(-1, 1)
        >>> g.edata['he'] = torch.arange(3).float().reshape(-1, 1)
        >>> g.remove_nodes(torch.tensor([0, 1]))
        >>> g
        Graph(num_nodes=1, num_edges=1,
            ndata_schemes={'hv': Scheme(shape=(1,), dtype=torch.float32)}
            edata_schemes={'he': Scheme(shape=(1,), dtype=torch.float32)})
        >>> g.ndata['hv']
        tensor([[2.]])
        >>> g.edata['he']
        tensor([[2.]])

        **Heterogeneous Graphs with Multiple Node Types**

        >>> g = dgl.heterograph({
683
684
685
686
687
        ...     ('user', 'plays', 'game'): (torch.tensor([0, 1, 1, 2]),
        ...                                 torch.tensor([0, 0, 1, 1])),
        ...     ('developer', 'develops', 'game'): (torch.tensor([0, 1]),
        ...                                         torch.tensor([0, 1]))
        ...     })
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
        >>> g.remove_nodes(torch.tensor([0, 1]))
        DGLError: Node type name must be specified
        if there are more than one node types.
        >>> g.remove_nodes(torch.tensor([0, 1]), ntype='game')
        >>> g.num_nodes('user')
        3
        >>> g.num_nodes('game')
        0
        >>> g.num_edges('plays')
        0

        See Also
        --------
        add_nodes
        add_edges
        remove_edges
        """
        # TODO(xiangsx): block do not support remove_nodes
        if ntype is None:
            if self._graph.number_of_ntypes() != 1:
                raise DGLError('Node type name must be specified if there are more than one ' \
                               'node types.')
Minjie Wang's avatar
Minjie Wang committed
710

711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
        nids = utils.prepare_tensor(self, nids, 'u')
        if len(nids) == 0:
            # no node to delete
            return
        assert self.number_of_nodes(ntype) > F.as_scalar(F.max(nids, dim=0)), \
            'The input nids {} is out of the range [0:{})'.format(
                F.as_scalar(F.max(nids, dim=0)), self.number_of_nodes(ntype))

        ntid = self.get_ntype_id(ntype)
        nodes = {}
        for c_ntype in self.ntypes:
            if self.get_ntype_id(c_ntype) == ntid:
                original_nids = self.nodes(c_ntype)
                nodes[c_ntype] = utils.compensate(nids, original_nids)
            else:
                nodes[c_ntype] = self.nodes(c_ntype)

        # node_subgraph
729
        sub_g = self.subgraph(nodes, store_ids=store_ids)
730
731
732
733
734
735
736
        self._graph = sub_g._graph
        self._node_frames = sub_g._node_frames
        self._edge_frames = sub_g._edge_frames

    def _reset_cached_info(self):
        """Some info like batch_num_nodes may be stale after mutation
        Clean these cached info
Minjie Wang's avatar
Minjie Wang committed
737
        """
738
739
740
        self._batch_num_nodes = None
        self._batch_num_edges = None

Minjie Wang's avatar
Minjie Wang committed
741
742
743
744

    #################################################################
    # Metagraph query
    #################################################################
Da Zheng's avatar
Da Zheng committed
745

746
747
748
749
750
751
    @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
752
        can be used to get the type, data, and nodes that belong to SRC and DST sets:
753
754
755
756
757
758
759
760
761
762
763

        * :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

764
    @property
Minjie Wang's avatar
Minjie Wang committed
765
    def ntypes(self):
766
        """Return all the node type names in the graph.
Mufei Li's avatar
Mufei Li committed
767
768
769

        Returns
        -------
770
771
772
773
774
775
776
        list[str]
            All the node type names in a list.

        Notes
        -----
        DGL internally assigns an integer ID for each node type. The returned
        node type names are sorted according to their IDs.
Mufei Li's avatar
Mufei Li committed
777
778
779

        Examples
        --------
780
781
782
783
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch
Mufei Li's avatar
Mufei Li committed
784

785
786
787
788
789
        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'follows', 'game'): (torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])),
        ...     ('user', 'plays', 'game'): (torch.tensor([1, 3]), torch.tensor([2, 3]))
        ... })
Mufei Li's avatar
Mufei Li committed
790
        >>> g.ntypes
791
        ['game', 'user']
Mufei Li's avatar
Mufei Li committed
792
        """
793
        return self._ntypes
Da Zheng's avatar
Da Zheng committed
794

795
    @property
Minjie Wang's avatar
Minjie Wang committed
796
    def etypes(self):
797
        """Return all the edge type names in the graph.
Mufei Li's avatar
Mufei Li committed
798
799
800

        Returns
        -------
801
802
        list[str]
            All the edge type names in a list.
803
804
805

        Notes
        -----
806
807
808
809
810
811
812
813
814
815
816
817
        DGL internally assigns an integer ID for each edge type. The returned
        edge type names are sorted according to their IDs.

        The complete format to specify an relation is a string triplet ``(str, str, str)``
        for source node type, edge type and destination node type. DGL calls this
        format *canonical edge type*. An edge type can appear in multiple canonical edge types.
        For example, ``'interacts'`` can appear in two canonical edge types
        ``('drug', 'interacts', 'drug')`` and ``('protein', 'interacts', 'protein')``.

        See Also
        --------
        canonical_etypes
Mufei Li's avatar
Mufei Li committed
818
819
820

        Examples
        --------
821
822
823
824
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch
Mufei Li's avatar
Mufei Li committed
825

826
827
828
829
830
        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'follows', 'game'): (torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])),
        ...     ('user', 'plays', 'game'): (torch.tensor([1, 3]), torch.tensor([2, 3]))
        ... })
Mufei Li's avatar
Mufei Li committed
831
        >>> g.etypes
832
        ['follows', 'follows', 'plays']
Mufei Li's avatar
Mufei Li committed
833
        """
834
        return self._etypes
Da Zheng's avatar
Da Zheng committed
835

Minjie Wang's avatar
Minjie Wang committed
836
837
    @property
    def canonical_etypes(self):
838
        """Return all the canonical edge types in the graph.
Minjie Wang's avatar
Minjie Wang committed
839

840
841
        A canonical edge type is a string triplet ``(str, str, str)``
        for source node type, edge type and destination node type.
Mufei Li's avatar
Mufei Li committed
842
843
844

        Returns
        -------
845
846
847
848
849
850
851
852
853
854
855
        list[(str, str, str)]
            All the canonical edge type triplets in a list.

        Notes
        -----
        DGL internally assigns an integer ID for each edge type. The returned
        edge type names are sorted according to their IDs.

        See Also
        --------
        etypes
Mufei Li's avatar
Mufei Li committed
856
857
858

        Examples
        --------
859
860
861
862
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch
Mufei Li's avatar
Mufei Li committed
863

864
865
866
867
868
        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'follows', 'game'): (torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])),
        ...     ('user', 'plays', 'game'): (torch.tensor([1, 3]), torch.tensor([2, 3]))
        ... })
Mufei Li's avatar
Mufei Li committed
869
        >>> g.canonical_etypes
870
871
872
        [('user', 'follows', 'user'),
         ('user', 'follows', 'game'),
         ('user', 'plays', 'game')]
Minjie Wang's avatar
Minjie Wang committed
873
874
875
        """
        return self._canonical_etypes

876
    @property
877
    def srctypes(self):
878
879
880
881
882
883
884
885
        """Return all the source node type names in this graph.

        If the graph can further divide its node types into two subsets A and B where
        all the edeges are from nodes of types in A to nodes of types in B, we call
        this graph a *uni-bipartite* graph and the nodes in A being the *source*
        nodes and the ones in B being the *destination* nodes. If the graph is not
        uni-bipartite, the source and destination nodes are just the entire set of
        nodes in the graph.
886
887
888

        Returns
        -------
889
890
        list[str]
            All the source node type names in a list.
891

892
893
894
895
        See Also
        --------
        dsttypes
        is_unibipartite
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920

        Examples
        --------
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Query for a uni-bipartite graph.

        >>> g = dgl.heterograph({
        ...     ('user', 'plays', 'game'): (torch.tensor([0]), torch.tensor([1])),
        ...     ('developer', 'develops', 'game'): (torch.tensor([1]), torch.tensor([2]))
        ... })
        >>> g.srctypes
        ['developer', 'user']

        Query for a graph that is not uni-bipartite.

        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0]), torch.tensor([1])),
        ...     ('developer', 'develops', 'game'): (torch.tensor([1]), torch.tensor([2]))
        ... })
        >>> g.srctypes
        ['developer', 'game', 'user']
921
922
923
924
925
        """
        if self.is_unibipartite:
            return sorted(list(self._srctypes_invmap.keys()))
        else:
            return self.ntypes
926
927

    @property
928
    def dsttypes(self):
929
930
931
932
933
934
935
936
        """Return all the destination node type names in this graph.

        If the graph can further divide its node types into two subsets A and B where
        all the edeges are from nodes of types in A to nodes of types in B, we call
        this graph a *uni-bipartite* graph and the nodes in A being the *source*
        nodes and the ones in B being the *destination* nodes. If the graph is not
        uni-bipartite, the source and destination nodes are just the entire set of
        nodes in the graph.
937
938
939

        Returns
        -------
940
941
        list[str]
            All the destination node type names in a list.
942

943
944
945
946
        See Also
        --------
        srctypes
        is_unibipartite
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971

        Examples
        --------
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Query for a uni-bipartite graph.

        >>> g = dgl.heterograph({
        ...     ('user', 'plays', 'game'): (torch.tensor([0]), torch.tensor([1])),
        ...     ('developer', 'develops', 'game'): (torch.tensor([1]), torch.tensor([2]))
        ... })
        >>> g.dsttypes
        ['game']

        Query for a graph that is not uni-bipartite.

        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0]), torch.tensor([1])),
        ...     ('developer', 'develops', 'game'): (torch.tensor([1]), torch.tensor([2]))
        ... })
        >>> g.dsttypes
        ['developer', 'game', 'user']
972
973
974
975
976
        """
        if self.is_unibipartite:
            return sorted(list(self._dsttypes_invmap.keys()))
        else:
            return self.ntypes
977

Da Zheng's avatar
Da Zheng committed
978
    def metagraph(self):
979
        """Return the metagraph of the heterograph.
980

981
982
983
        The metagraph (or network schema) of a heterogeneous network specifies type constraints
        on the sets of nodes and edges between the nodes. For a formal definition, refer to
        `Yizhou et al. <https://www.kdd.org/exploration_files/V14-02-03-Sun.pdf>`_.
Minjie Wang's avatar
Minjie Wang committed
984
985
986
987

        Returns
        -------
        networkx.MultiDiGraph
988
            The metagraph.
Mufei Li's avatar
Mufei Li committed
989
990
991

        Examples
        --------
992
        The following example uses PyTorch backend.
Mufei Li's avatar
Mufei Li committed
993

994
995
        >>> import dgl
        >>> import torch
Mufei Li's avatar
Mufei Li committed
996

997
998
999
1000
1001
1002
        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'follows', 'game'): (torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])),
        ...     ('user', 'plays', 'game'): (torch.tensor([1, 3]), torch.tensor([2, 3]))
        ... })
        >>> meta_g = g.metagraph()
Mufei Li's avatar
Mufei Li committed
1003
1004
1005
        >>> meta_g.nodes()
        NodeView(('user', 'game'))
        >>> meta_g.edges()
1006
        OutMultiEdgeDataView([('user', 'user'), ('user', 'game'), ('user', 'game')])
Minjie Wang's avatar
Minjie Wang committed
1007
        """
1008
1009
1010
1011
1012
1013
        nx_graph = self._graph.metagraph.to_networkx()
        nx_metagraph = nx.MultiDiGraph()
        for u_v in nx_graph.edges:
            srctype, etype, dsttype = self.canonical_etypes[nx_graph.edges[u_v]['id']]
            nx_metagraph.add_edge(srctype, dsttype, etype)
        return nx_metagraph
Minjie Wang's avatar
Minjie Wang committed
1014
1015

    def to_canonical_etype(self, etype):
1016
        """Convert an edge type to the corresponding canonical edge type in the graph.
Minjie Wang's avatar
Minjie Wang committed
1017

1018
1019
1020
1021
1022
        A canonical edge type is a string triplet ``(str, str, str)``
        for source node type, edge type and destination node type.

        The function expects the given edge type name can uniquely identify a canonical edge
        type. DGL will raise error if this is not the case.
Minjie Wang's avatar
Minjie Wang committed
1023
1024
1025

        Parameters
        ----------
1026
        etype : str or (str, str, str)
1027
            If :attr:`etype` is an edge type (str), it returns the corresponding canonical edge
1028
1029
            type in the graph. If :attr:`etype` is already a canonical edge type,
            it directly returns the input unchanged.
Minjie Wang's avatar
Minjie Wang committed
1030
1031
1032

        Returns
        -------
1033
        (str, str, str)
1034
1035
            The canonical edge type corresponding to the edge type.

Mufei Li's avatar
Mufei Li committed
1036
1037
        Examples
        --------
1038
        The following example uses PyTorch backend.
Mufei Li's avatar
Mufei Li committed
1039

1040
1041
1042
1043
        >>> import dgl
        >>> import torch

        Create a heterograph.
Mufei Li's avatar
Mufei Li committed
1044

1045
1046
1047
1048
1049
        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): ([0, 1], [1, 2]),
        ...     ('user', 'plays', 'game'): ([0, 1, 1, 2], [0, 0, 1, 1]),
        ...     ('developer', 'follows', 'game'): ([0, 1], [0, 1])
        ... })
Mufei Li's avatar
Mufei Li committed
1050

1051
        Map an edge type to its corresponding canonical edge type.
Mufei Li's avatar
Mufei Li committed
1052
1053
1054
1055
1056

        >>> g.to_canonical_etype('plays')
        ('user', 'plays', 'game')
        >>> g.to_canonical_etype(('user', 'plays', 'game'))
        ('user', 'plays', 'game')
1057
1058
1059
1060

        See Also
        --------
        canonical_etypes
1061
        """
1062
1063
1064
1065
1066
        if etype is None:
            if len(self.etypes) != 1:
                raise DGLError('Edge type name must be specified if there are more than one '
                               'edge types.')
            etype = self.etypes[0]
Minjie Wang's avatar
Minjie Wang committed
1067
1068
        if isinstance(etype, tuple):
            return etype
1069
        else:
Minjie Wang's avatar
Minjie Wang committed
1070
1071
1072
1073
            ret = self._etype2canonical.get(etype, None)
            if ret is None:
                raise DGLError('Edge type "{}" does not exist.'.format(etype))
            if len(ret) == 0:
1074
1075
                raise DGLError('Edge type "%s" is ambiguous. Please use canonical edge type '
                               'in the form of (srctype, etype, dsttype)' % etype)
Minjie Wang's avatar
Minjie Wang committed
1076
1077
1078
            return ret

    def get_ntype_id(self, ntype):
1079
        """Return the ID of the given node type.
Minjie Wang's avatar
Minjie Wang committed
1080
1081
1082

        ntype can also be None. If so, there should be only one node type in the
        graph.
1083

Minjie Wang's avatar
Minjie Wang committed
1084
1085
1086
1087
        Parameters
        ----------
        ntype : str
            Node type
Da Zheng's avatar
Da Zheng committed
1088
1089
1090

        Returns
        -------
Minjie Wang's avatar
Minjie Wang committed
1091
1092
        int
        """
1093
        if self.is_unibipartite and ntype is not None:
1094
1095
1096
1097
1098
1099
1100
1101
            # 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
Minjie Wang's avatar
Minjie Wang committed
1102
        if ntype is None:
1103
            if self.is_unibipartite or len(self._srctypes_invmap) != 1:
Minjie Wang's avatar
Minjie Wang committed
1104
1105
1106
                raise DGLError('Node type name must be specified if there are more than one '
                               'node types.')
            return 0
1107
        ntid = self._srctypes_invmap.get(ntype, self._dsttypes_invmap.get(ntype, None))
Minjie Wang's avatar
Minjie Wang committed
1108
1109
1110
        if ntid is None:
            raise DGLError('Node type "{}" does not exist.'.format(ntype))
        return ntid
Da Zheng's avatar
Da Zheng committed
1111

1112
    def get_ntype_id_from_src(self, ntype):
1113
        """Internal function to return the ID of the given SRC node type.
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130

        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.')
1131
            return next(iter(self._srctypes_invmap.values()))
1132
1133
1134
1135
1136
1137
        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):
1138
        """Internal function to return the ID of the given DST node type.
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155

        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.')
1156
            return next(iter(self._dsttypes_invmap.values()))
1157
1158
1159
1160
1161
        ntid = self._dsttypes_invmap.get(ntype, None)
        if ntid is None:
            raise DGLError('DST node type "{}" does not exist.'.format(ntype))
        return ntid

Minjie Wang's avatar
Minjie Wang committed
1162
1163
    def get_etype_id(self, etype):
        """Return the id of the given edge type.
1164

Minjie Wang's avatar
Minjie Wang committed
1165
1166
1167
1168
1169
1170
1171
        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
Da Zheng's avatar
Da Zheng committed
1172

1173
1174
        Returns
        -------
Minjie Wang's avatar
Minjie Wang committed
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
        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
Da Zheng's avatar
Da Zheng committed
1186

1187
1188
1189
1190
1191
    #################################################################
    # Batching
    #################################################################
    @property
    def batch_size(self):
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
        """Return the number of graphs in the batched graph.

        Returns
        -------
        int
            The Number of graphs in the batch. If the graph is not a batched one,
            it will return 1.

        Examples
        --------

        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Query for homogeneous graphs.

        >>> g1 = dgl.graph((torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])))
        >>> g1.batch_size
        1
        >>> g2 = dgl.graph((torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])))
        >>> bg = dgl.batch([g1, g2])
        >>> bg.batch_size
        2

        Query for heterogeneous graphs.

        >>> hg1 = dgl.heterograph({
        ...       ('user', 'plays', 'game') : (torch.tensor([0, 1]), torch.tensor([0, 0]))})
        >>> hg1.batch_size
        1
        >>> hg2 = dgl.heterograph({
        ...       ('user', 'plays', 'game') : (torch.tensor([0, 0]), torch.tensor([1, 0]))})
        >>> bg = dgl.batch([hg1, hg2])
        >>> bg.batch_size
        2
        """
1230
1231
1232
        return len(self.batch_num_nodes(self.ntypes[0]))

    def batch_num_nodes(self, ntype=None):
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
        """Return the number of nodes for each graph in the batch with the specified node type.

        Parameters
        ----------
        ntype : str, optional
            The node type for query. If the graph has multiple node types, one must
            specify the argument. Otherwise, it can be omitted. If the graph is not a batched
            one, it will return a list of length 1 that holds the number of nodes in the graph.

        Returns
        -------
        Tensor
            The number of nodes with the specified type for each graph in the batch. The i-th
            element of it is the number of nodes with the specified type for the i-th graph.

        Examples
        --------

        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Query for homogeneous graphs.

        >>> g1 = dgl.graph((torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])))
        >>> g1.batch_num_nodes()
        tensor([4])
        >>> g2 = dgl.graph((torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])))
        >>> bg = dgl.batch([g1, g2])
        >>> bg.batch_num_nodes()
        tensor([4, 3])

        Query for heterogeneous graphs.

        >>> hg1 = dgl.heterograph({
        ...       ('user', 'plays', 'game') : (torch.tensor([0, 1]), torch.tensor([0, 0]))})
        >>> hg2 = dgl.heterograph({
        ...       ('user', 'plays', 'game') : (torch.tensor([0, 0]), torch.tensor([1, 0]))})
        >>> bg = dgl.batch([hg1, hg2])
        >>> bg.batch_num_nodes('user')
        tensor([2, 1])
        """
        if ntype is not None and ntype not in self.ntypes:
            raise DGLError('Expect ntype in {}, got {}'.format(self.ntypes, ntype))

1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
        if self._batch_num_nodes is None:
            self._batch_num_nodes = {}
            for ty in self.ntypes:
                bnn = F.copy_to(F.tensor([self.number_of_nodes(ty)], F.int64), self.device)
                self._batch_num_nodes[ty] = bnn
        if ntype is None:
            if len(self.ntypes) != 1:
                raise DGLError('Node type name must be specified if there are more than one '
                               'node types.')
            ntype = self.ntypes[0]
        return self._batch_num_nodes[ntype]

    def set_batch_num_nodes(self, val):
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
        """Manually set the number of nodes for each graph in the batch with the specified node
        type.

        Parameters
        ----------
        val : Tensor or Mapping[str, Tensor]
            The dictionary storing number of nodes for each graph in the batch for all node types.
            If the graph has only one node type, ``val`` can also be a single array indicating the
            number of nodes per graph in the batch.

        Notes
        -----
        This API is always used together with ``set_batch_num_edges`` to specify batching
        information of a graph, it also do not check the correspondance between the graph structure
        and batching information and user must guarantee there will be no cross-graph edges in the
        batch.

        Examples
        --------

        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Create a homogeneous graph.
1318

1319
1320
1321
        >>> g = dgl.graph(([0, 1, 2, 3, 4, 5], [1, 2, 0, 4, 5, 3]))

        Manually set batch information
1322

1323
1324
        >>> g.set_batch_num_nodes(torch.tensor([3, 3]))
        >>> g.set_batch_num_edges(torch.tensor([3, 3]))
1325
1326

        Unbatch the graph.
1327

1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
        >>> dgl.unbatch(g)
        [Graph(num_nodes=3, num_edges=3,
              ndata_schemes={}
              edata_schemes={}), Graph(num_nodes=3, num_edges=3,
              ndata_schemes={}
              edata_schemes={})]

        Create a heterogeneous graph.

        >>> hg = dgl.heterograph({
        ...      ('user', 'plays', 'game') : ([0, 1, 2, 3, 4, 5], [0, 1, 1, 3, 3, 2]),
        ...      ('developer', 'develops', 'game') : ([0, 1, 2, 3], [1, 0, 3, 2])})

        Manually set batch information.

        >>> hg.set_batch_num_nodes({
        ...     'user': torch.tensor([3, 3]),
        ...     'game': torch.tensor([2, 2]),
        ...     'developer': torch.tensor([2, 2])})
        >>> hg.set_batch_num_edges({
        ...     ('user', 'plays', 'game'): torch.tensor([3, 3]),
        ...     ('developer', 'develops', 'game'): torch.tensor([2, 2])})

        Unbatch the graph.

        >>> g1, g2 = dgl.unbatch(hg)
        >>> g1
        Graph(num_nodes={'developer': 2, 'game': 2, 'user': 3},
              num_edges={('developer', 'develops', 'game'): 2, ('user', 'plays', 'game'): 3},
              metagraph=[('developer', 'game', 'develops'), ('user', 'game', 'plays')])
        >>> g2
        Graph(num_nodes={'developer': 2, 'game': 2, 'user': 3},
              num_edges={('developer', 'develops', 'game'): 2, ('user', 'plays', 'game'): 3},
              metagraph=[('developer', 'game', 'develops'), ('user', 'game', 'plays')])

        See Also
        --------
        set_batch_num_edges
        batch
        unbatch
        """
1369
1370
1371
1372
1373
1374
1375
        if not isinstance(val, Mapping):
            if len(self.ntypes) != 1:
                raise DGLError('Must provide a dictionary when there are multiple node types.')
            val = {self.ntypes[0] : val}
        self._batch_num_nodes = val

    def batch_num_edges(self, etype=None):
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
        """Return the number of edges for each graph in the batch with the specified edge type.

        Parameters
        ----------
        etype : str or tuple of str, optional
            The edge type for query, which can be an edge type (str) or a canonical edge type
            (3-tuple of str). When an edge type appears in multiple canonical edge types, one
            must use a canonical edge type. If the graph has multiple edge types, one must
            specify the argument. Otherwise, it can be omitted.

        Returns
        -------
        Tensor
            The number of edges with the specified type for each graph in the batch. The i-th
            element of it is the number of edges with the specified type for the i-th graph.
            If the graph is not a batched one, it will return a list of length 1 that holds
            the number of edges in the graph.

        Examples
        --------

        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Query for homogeneous graphs.

        >>> g1 = dgl.graph((torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])))
        >>> g1.batch_num_edges()
        tensor([3])
        >>> g2 = dgl.graph((torch.tensor([0, 0, 0, 1]), torch.tensor([0, 1, 2, 0])))
        >>> bg = dgl.batch([g1, g2])
        >>> bg.batch_num_edges()
        tensor([3, 4])

        Query for heterogeneous graphs.

        >>> hg1 = dgl.heterograph({
        ...       ('user', 'plays', 'game') : (torch.tensor([0, 1]), torch.tensor([0, 0]))})
        >>> hg2 = dgl.heterograph({
        ...       ('user', 'plays', 'game') : (torch.tensor([0, 0]), torch.tensor([1, 0]))})
        >>> bg = dgl.batch([hg1, hg2])
        >>> bg.batch_num_edges('plays')
        tensor([2, 2])
        """
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
        if self._batch_num_edges is None:
            self._batch_num_edges = {}
            for ty in self.canonical_etypes:
                bne = F.copy_to(F.tensor([self.number_of_edges(ty)], F.int64), self.device)
                self._batch_num_edges[ty] = bne
        if etype is None:
            if len(self.etypes) != 1:
                raise DGLError('Edge type name must be specified if there are more than one '
                               'edge types.')
            etype = self.canonical_etypes[0]
1432
1433
        else:
            etype = self.to_canonical_etype(etype)
1434
1435
1436
        return self._batch_num_edges[etype]

    def set_batch_num_edges(self, val):
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
        """Manually set the number of edges for each graph in the batch with the specified edge
        type.

        Parameters
        ----------
        val : Tensor or Mapping[str, Tensor]
            The dictionary storing number of edges for each graph in the batch for all edge types.
            If the graph has only one edge type, ``val`` can also be a single array indicating the
            number of edges per graph in the batch.

        Notes
        -----
1449
        This API is always used together with ``set_batch_num_nodes`` to specify batching
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
        information of a graph, it also do not check the correspondance between the graph structure
        and batching information and user must guarantee there will be no cross-graph edges in the
        batch.

        Examples
        --------

        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Create a homogeneous graph.
1463

1464
1465
1466
        >>> g = dgl.graph(([0, 1, 2, 3, 4, 5], [1, 2, 0, 4, 5, 3]))

        Manually set batch information
1467

1468
1469
        >>> g.set_batch_num_nodes(torch.tensor([3, 3]))
        >>> g.set_batch_num_edges(torch.tensor([3, 3]))
1470
1471

        Unbatch the graph.
1472

1473
        >>> dgl.unbatch(g)
1474
1475
1476
1477
1478
        [Graph(num_nodes=3, num_edges=3,
              ndata_schemes={}
              edata_schemes={}), Graph(num_nodes=3, num_edges=3,
              ndata_schemes={}
              edata_schemes={})]
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513

        Create a heterogeneous graph.

        >>> hg = dgl.heterograph({
        ...      ('user', 'plays', 'game') : ([0, 1, 2, 3, 4, 5], [0, 1, 1, 3, 3, 2]),
        ...      ('developer', 'develops', 'game') : ([0, 1, 2, 3], [1, 0, 3, 2])})

        Manually set batch information.

        >>> hg.set_batch_num_nodes({
        ...     'user': torch.tensor([3, 3]),
        ...     'game': torch.tensor([2, 2]),
        ...     'developer': torch.tensor([2, 2])})
        >>> hg.set_batch_num_edges(
        ...     {('user', 'plays', 'game'): torch.tensor([3, 3]),
        ...     ('developer', 'develops', 'game'): torch.tensor([2, 2])})

        Unbatch the graph.

        >>> g1, g2 = dgl.unbatch(hg)
        >>> g1
        Graph(num_nodes={'developer': 2, 'game': 2, 'user': 3},
              num_edges={('developer', 'develops', 'game'): 2, ('user', 'plays', 'game'): 3},
              metagraph=[('developer', 'game', 'develops'), ('user', 'game', 'plays')])
        >>> g2
        Graph(num_nodes={'developer': 2, 'game': 2, 'user': 3},
              num_edges={('developer', 'develops', 'game'): 2, ('user', 'plays', 'game'): 3},
              metagraph=[('developer', 'game', 'develops'), ('user', 'game', 'plays')])

        See Also
        --------
        set_batch_num_nodes
        batch
        unbatch
        """
1514
1515
1516
1517
1518
1519
        if not isinstance(val, Mapping):
            if len(self.etypes) != 1:
                raise DGLError('Must provide a dictionary when there are multiple edge types.')
            val = {self.canonical_etypes[0] : val}
        self._batch_num_edges = val

Minjie Wang's avatar
Minjie Wang committed
1520
1521
1522
    #################################################################
    # View
    #################################################################
Da Zheng's avatar
Da Zheng committed
1523

1524
    @property
Minjie Wang's avatar
Minjie Wang committed
1525
    def nodes(self):
1526
1527
1528
1529
1530
1531
        """Return a node view

        One can use it for:

        1. Getting the node IDs for a single node type.
        2. Setting/getting features for all nodes of a single node type.
Da Zheng's avatar
Da Zheng committed
1532

Minjie Wang's avatar
Minjie Wang committed
1533
1534
        Examples
        --------
Mufei Li's avatar
Mufei Li committed
1535
1536
        The following example uses PyTorch backend.

1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
        >>> import dgl
        >>> import torch

        Create a homogeneous graph and a heterogeneous graph of two node types.

        >>> g = dgl.graph((torch.tensor([0, 1]), torch.tensor([1, 2])))
        >>> hg = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'plays', 'game'): (torch.tensor([3, 4]), torch.tensor([5, 6]))
        ... })

        Get the node IDs of the homogeneous graph.

        >>> g.nodes()
        tensor([0, 1, 2])

        Get the node IDs of the heterogeneous graph. With multiple node types introduced,
        one needs to specify the node type for query.

        >>> hg.nodes('user')
        tensor([0, 1, 2, 3, 4])
Mufei Li's avatar
Mufei Li committed
1558

1559
1560
1561
1562
1563
1564
1565
        Set and get a feature 'h' for all nodes of a single type in the heterogeneous graph.

        >>> hg.nodes['user'].data['h'] = torch.ones(5, 1)
        >>> hg.nodes['user'].data['h']
        tensor([[1.], [1.], [1.], [1.], [1.]])

        To set node features for a graph with a single node type, use :func:`DGLGraph.ndata`.
Mufei Li's avatar
Mufei Li committed
1566
1567
1568
1569

        See Also
        --------
        ndata
1570
        """
1571
        # Todo (Mufei) Replace the syntax g.nodes[...].ndata[...] with g.nodes[...][...]
1572
1573
1574
1575
        return HeteroNodeView(self, self.get_ntype_id)

    @property
    def srcnodes(self):
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
        """Return a node view for source nodes

        If the graph is a uni-bipartite graph (see :func:`is_unibipartite` for reference),
        this is :func:`nodes` restricted to source node types. Otherwise, it is an alias
        for :func:`nodes`.

        One can use it for:

        1. Getting the node IDs for a single node type.
        2. Setting/getting features for all nodes of a single node type.
1586
1587
1588
1589
1590

        Examples
        --------
        The following example uses PyTorch backend.

1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
        >>> import dgl
        >>> import torch

        Create a uni-bipartite graph.

        >>> g = dgl.heterograph({
        ...     ('user', 'plays', 'game'): (torch.tensor([0]), torch.tensor([1])),
        ...     ('developer', 'develops', 'game'): (torch.tensor([1]), torch.tensor([2]))
        ... })

        Get the node IDs for source node types.

        >>> g.srcnodes('user')
        tensor([0])
        >>> g.srcnodes('developer')
        tensor([0, 1])

        Set/get features for source node types.

        >>> g.srcnodes['user'].data['h'] = torch.ones(1, 1)
        >>> g.srcnodes['user'].data['h']
        tensor([[1.]])
1613

1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
        Create a graph that is not uni-bipartite.

        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0]), torch.tensor([1])),
        ...     ('developer', 'develops', 'game'): (torch.tensor([1]), torch.tensor([2]))
        ... })

        :func:`dgl.DGLGraph.srcnodes` falls back to :func:`dgl.DGLGraph.nodes` and one can
        get the node IDs for both source and destination node types.

        >>> g.srcnodes('game')
        tensor([0, 1, 2])

        One can also set/get features for destination node types in this case.

        >>> g.srcnodes['game'].data['h'] = torch.ones(3, 1)
        >>> g.srcnodes['game'].data['h']
        tensor([[1.],
                [1.],
                [1.]])
1634
1635
1636
1637
1638
1639
1640
1641
1642

        See Also
        --------
        srcdata
        """
        return HeteroNodeView(self, self.get_ntype_id_from_src)

    @property
    def dstnodes(self):
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
        """Return a node view for destination nodes

        If the graph is a uni-bipartite graph (see :func:`is_unibipartite` for reference),
        this is :func:`nodes` restricted to destination node types. Otherwise, it is an alias
        for :func:`nodes`.

        One can use it for:

        1. Getting the node IDs for a single node type.
        2. Setting/getting features for all nodes of a single node type.
1653
1654
1655
1656
1657

        Examples
        --------
        The following example uses PyTorch backend.

1658
1659
1660
1661
        >>> import dgl
        >>> import torch

        Create a uni-bipartite graph.
1662

1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
        >>> g = dgl.heterograph({
        ...     ('user', 'plays', 'game'): (torch.tensor([0]), torch.tensor([1])),
        ...     ('developer', 'develops', 'game'): (torch.tensor([1]), torch.tensor([2]))
        ... })

        Get the node IDs for destination node types.

        >>> g.dstnodes('game')
        tensor([0, 1, 2])

        Set/get features for destination node types.

        >>> g.dstnodes['game'].data['h'] = torch.ones(3, 1)
        >>> g.dstnodes['game'].data['h']
        tensor([[1.],
                [1.],
                [1.]])

        Create a graph that is not uni-bipartite.

        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0]), torch.tensor([1])),
        ...     ('developer', 'develops', 'game'): (torch.tensor([1]), torch.tensor([2]))
        ... })

        :func:`dgl.DGLGraph.dstnodes` falls back to :func:`dgl.DGLGraph.nodes` and one can
        get the node IDs for both source and destination node types.

        >>> g.dstnodes('developer')
        tensor([0, 1])

        One can also set/get features for source node types in this case.

        >>> g.dstnodes['developer'].data['h'] = torch.ones(2, 1)
        >>> g.dstnodes['developer'].data['h']
        tensor([[1.],
                [1.]])
1700
1701
1702
1703
1704
1705

        See Also
        --------
        dstdata
        """
        return HeteroNodeView(self, self.get_ntype_id_from_dst)
Da Zheng's avatar
Da Zheng committed
1706

1707
    @property
Minjie Wang's avatar
Minjie Wang committed
1708
    def ndata(self):
1709
1710
1711
1712
1713
        """Return a node data view for setting/getting node features

        Let ``g`` be a DGLGraph. If ``g`` is a graph of a single node type, ``g.ndata[feat]``
        returns the node feature associated with the name ``feat``. One can also set a node
        feature associated with the name ``feat`` by setting ``g.ndata[feat]`` to a tensor.
Da Zheng's avatar
Da Zheng committed
1714

1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
        If ``g`` is a graph of multiple node types, ``g.ndata[feat]`` returns a
        dict[str, Tensor] mapping node types to the node features associated with the name
        ``feat`` for the corresponding type. One can also set a node feature associated
        with the name ``feat`` for some node type(s) by setting ``g.ndata[feat]`` to a
        dictionary as described.

        Notes
        -----
        For setting features, the device of the features must be the same as the device
        of the graph.
Minjie Wang's avatar
Minjie Wang committed
1725
1726
1727

        Examples
        --------
Mufei Li's avatar
Mufei Li committed
1728
1729
        The following example uses PyTorch backend.

1730
1731
        >>> import dgl
        >>> import torch
Mufei Li's avatar
Mufei Li committed
1732

1733
        Set and get feature 'h' for a graph of a single node type.
1734

1735
1736
        >>> g = dgl.graph((torch.tensor([0, 1]), torch.tensor([1, 2])))
        >>> g.ndata['h'] = torch.ones(3, 1)
1737
        >>> g.ndata['h']
1738
1739
1740
        tensor([[1.],
                [1.],
                [1.]])
1741

1742
        Set and get feature 'h' for a graph of multiple node types.
1743

1744
1745
1746
1747
1748
        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([1, 2]), torch.tensor([3, 4])),
        ...     ('player', 'plays', 'game'): (torch.tensor([2, 2]), torch.tensor([1, 1]))
        ... })
        >>> g.ndata['h'] = {'game': torch.zeros(2, 1), 'player': torch.ones(3, 1)}
1749
        >>> g.ndata['h']
1750
1751
1752
        {'game': tensor([[0.], [0.]]),
         'player': tensor([[1.], [1.], [1.]])}
        >>> g.ndata['h'] = {'game': torch.ones(2, 1)}
1753
        >>> g.ndata['h']
1754
1755
        {'game': tensor([[1.], [1.]]),
         'player': tensor([[1.], [1.], [1.]])}
1756

Mufei Li's avatar
Mufei Li committed
1757
1758
1759
        See Also
        --------
        nodes
Da Zheng's avatar
Da Zheng committed
1760
        """
1761
1762
1763
1764
1765
1766
1767
1768
1769
        if len(self.ntypes) == 1:
            ntid = self.get_ntype_id(None)
            ntype = self.ntypes[0]
            return HeteroNodeDataView(self, ntype, ntid, ALL)
        else:
            ntids = [self.get_ntype_id(ntype) for ntype in self.ntypes]
            ntypes = self.ntypes
            return HeteroNodeDataView(self, ntypes, ntids, ALL)

1770
1771
    @property
    def srcdata(self):
1772
        """Return a node data view for setting/getting source node features.
1773

1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
        Let ``g`` be a DGLGraph. If ``g`` is a graph of a single source node type,
        ``g.srcdata[feat]`` returns the source node feature associated with the name ``feat``.
        One can also set a source node feature associated with the name ``feat`` by
        setting ``g.srcdata[feat]`` to a tensor.

        If ``g`` is a graph of multiple source node types, ``g.srcdata[feat]`` returns a
        dict[str, Tensor] mapping source node types to the node features associated with
        the name ``feat`` for the corresponding type. One can also set a node feature
        associated with the name ``feat`` for some source node type(s) by setting
        ``g.srcdata[feat]`` to a dictionary as described.

        Notes
        -----
        For setting features, the device of the features must be the same as the device
        of the graph.
1789
1790
1791
1792
1793

        Examples
        --------
        The following example uses PyTorch backend.

1794
1795
        >>> import dgl
        >>> import torch
1796

1797
        Set and get feature 'h' for a graph of a single source node type.
1798
1799

        >>> g = dgl.heterograph({
1800
1801
1802
1803
1804
        ...     ('user', 'plays', 'game'): (torch.tensor([0, 1]), torch.tensor([1, 2]))})
        >>> g.srcdata['h'] = torch.ones(2, 1)
        >>> g.srcdata['h']
        tensor([[1.],
                [1.]])
1805

1806
        Set and get feature 'h' for a graph of multiple source node types.
1807
1808

        >>> g = dgl.heterograph({
1809
1810
1811
1812
        ...     ('user', 'plays', 'game'): (torch.tensor([1, 2]), torch.tensor([3, 4])),
        ...     ('player', 'plays', 'game'): (torch.tensor([2, 2]), torch.tensor([1, 1]))
        ... })
        >>> g.srcdata['h'] = {'user': torch.zeros(3, 1), 'player': torch.ones(3, 1)}
1813
        >>> g.srcdata['h']
1814
1815
1816
        {'player': tensor([[1.], [1.], [1.]]),
         'user': tensor([[0.], [0.], [0.]])}
        >>> g.srcdata['h'] = {'user': torch.ones(3, 1)}
1817
        >>> g.srcdata['h']
1818
1819
        {'player': tensor([[1.], [1.], [1.]]),
         'user': tensor([[1.], [1.], [1.]])}
1820
1821
1822
1823

        See Also
        --------
        nodes
1824
1825
        ndata
        srcnodes
1826
        """
1827
1828
1829
1830
1831
1832
1833
1834
        if len(self.srctypes) == 1:
            ntype = self.srctypes[0]
            ntid = self.get_ntype_id_from_src(ntype)
            return HeteroNodeDataView(self, ntype, ntid, ALL)
        else:
            ntypes = self.srctypes
            ntids = [self.get_ntype_id_from_src(ntype) for ntype in ntypes]
            return HeteroNodeDataView(self, ntypes, ntids, ALL)
1835
1836
1837

    @property
    def dstdata(self):
1838
1839
1840
1841
1842
1843
        """Return a node data view for setting/getting destination node features.

        Let ``g`` be a DGLGraph. If ``g`` is a graph of a single destination node type,
        ``g.dstdata[feat]`` returns the destination node feature associated with the name
        ``feat``. One can also set a destination node feature associated with the name
        ``feat`` by setting ``g.dstdata[feat]`` to a tensor.
1844

1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
        If ``g`` is a graph of multiple destination node types, ``g.dstdata[feat]`` returns a
        dict[str, Tensor] mapping destination node types to the node features associated with
        the name ``feat`` for the corresponding type. One can also set a node feature
        associated with the name ``feat`` for some destination node type(s) by setting
        ``g.dstdata[feat]`` to a dictionary as described.

        Notes
        -----
        For setting features, the device of the features must be the same as the device
        of the graph.
1855
1856
1857
1858
1859

        Examples
        --------
        The following example uses PyTorch backend.

1860
1861
        >>> import dgl
        >>> import torch
1862

1863
        Set and get feature 'h' for a graph of a single destination node type.
1864
1865

        >>> g = dgl.heterograph({
1866
1867
1868
1869
1870
1871
        ...     ('user', 'plays', 'game'): (torch.tensor([0, 1]), torch.tensor([1, 2]))})
        >>> g.dstdata['h'] = torch.ones(3, 1)
        >>> g.dstdata['h']
        tensor([[1.],
                [1.],
                [1.]])
1872

1873
        Set and get feature 'h' for a graph of multiple destination node types.
1874
1875

        >>> g = dgl.heterograph({
1876
1877
1878
1879
        ...     ('user', 'plays', 'game'): (torch.tensor([1, 2]), torch.tensor([1, 2])),
        ...     ('user', 'watches', 'movie'): (torch.tensor([2, 2]), torch.tensor([1, 1]))
        ... })
        >>> g.dstdata['h'] = {'game': torch.zeros(3, 1), 'movie': torch.ones(2, 1)}
1880
        >>> g.dstdata['h']
1881
1882
1883
        {'game': tensor([[0.], [0.], [0.]]),
         'movie': tensor([[1.], [1.]])}
        >>> g.dstdata['h'] = {'game': torch.ones(3, 1)}
1884
        >>> g.dstdata['h']
1885
1886
        {'game': tensor([[1.], [1.], [1.]]),
         'movie': tensor([[1.], [1.]])}
1887
1888
1889
1890

        See Also
        --------
        nodes
1891
1892
        ndata
        dstnodes
1893
        """
1894
1895
1896
1897
1898
1899
1900
1901
        if len(self.dsttypes) == 1:
            ntype = self.dsttypes[0]
            ntid = self.get_ntype_id_from_dst(ntype)
            return HeteroNodeDataView(self, ntype, ntid, ALL)
        else:
            ntypes = self.dsttypes
            ntids = [self.get_ntype_id_from_dst(ntype) for ntype in ntypes]
            return HeteroNodeDataView(self, ntypes, ntids, ALL)
1902

1903
    @property
Minjie Wang's avatar
Minjie Wang committed
1904
    def edges(self):
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
        """Return an edge view

        One can use it for:

        1. Getting the edges for a single edge type. In this case, it can take the
           following optional arguments:

            - form : str, optional
                  The return form, which can be one of the following:

                  - ``'uv'`` (default): The returned result is a 2-tuple of 1D tensors
                    :math:`(U, V)`, representing the source and destination nodes of all edges.
                    For each :math:`i`, :math:`(U[i], V[i])` forms an edge.
                  - ``'eid'``: The returned result is a 1D tensor :math:`EID`, representing
                    the IDs of all edges.
                  - ``'all'``: The returned result is a 3-tuple of 1D tensors :math:`(U, V, EID)`,
                    representing the source nodes, destination nodes and IDs of all edges.
                    For each :math:`i`, :math:`(U[i], V[i])` forms an edge with ID :math:`EID[i]`.
            - order : str, optional
                  The order of the returned edges, which can be one of the following:

                  - ``'eid'`` (default): The edges are sorted by their IDs.
                  - ``'srcdst'``: The edges are sorted first by their source node IDs and then
                    by their destination node IDs to break ties.
            - etype : str or tuple of str, optional
                  The edge type for query, which can be an edge type (str) or a canonical edge
                  type (3-tuple of str). When an edge type appears in multiple canonical edge
                  types, one must use a canonical edge type. If the graph has multiple edge
                  types, one must specify the argument. Otherwise, it can be omitted.
        2. Setting/getting features for all edges of a single edge type. To set/get a feature
           ``feat`` for edges of type ``etype`` in a graph ``g``, one can use
           ``g.edges[etype].data[feat]``.
1937

Minjie Wang's avatar
Minjie Wang committed
1938
1939
        Examples
        --------
Mufei Li's avatar
Mufei Li committed
1940
1941
        The following example uses PyTorch backend.

1942
1943
        >>> import dgl
        >>> import torch
Mufei Li's avatar
Mufei Li committed
1944

1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
        **Get the Edges for a Single Edge Type**

        Create a graph with a single edge type.

        >>> g = dgl.graph((torch.tensor([1, 0, 0]), torch.tensor([1, 1, 0])))
        >>> g.edges()
        (tensor([1, 0, 0]), tensor([1, 1, 0]))

        Specify a different value for :attr:`form` and :attr:`order`.

        >>> g.edges(form='all', order='srcdst')
        (tensor([0, 0, 1]), tensor([0, 1, 1]), tensor([2, 1, 0]))

        For a graph of multiple edge types, it is required to specify the edge type in query.

        >>> hg = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'plays', 'game'): (torch.tensor([3, 4]), torch.tensor([5, 6]))
        ... })
        >>> hg.edges(etype='plays')
        (tensor([3, 4]), tensor([5, 6]))

        **Set/get Features for All Edges of a Single Edge Type**

        Create a heterogeneous graph of two edge types.

        >>> hg = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'plays', 'game'): (torch.tensor([3, 4]), torch.tensor([5, 6]))
        ... })

        Set and get a feature 'h' for all edges of a single type in the heterogeneous graph.

        >>> hg.edges['follows'].data['h'] = torch.ones(2, 1)
        >>> hg.edges['follows'].data['h']
        tensor([[1.], [1.]])

        To set edge features for a graph with a single edge type, use :func:`DGLGraph.edata`.
Mufei Li's avatar
Mufei Li committed
1983
1984
1985
1986

        See Also
        --------
        edata
1987
        """
1988
        # TODO(Mufei): Replace the syntax g.edges[...].edata[...] with g.edges[...][...]
Minjie Wang's avatar
Minjie Wang committed
1989
        return HeteroEdgeView(self)
1990
1991

    @property
Minjie Wang's avatar
Minjie Wang committed
1992
    def edata(self):
1993
1994
1995
1996
1997
        """Return an edge data view for setting/getting edge features.

        Let ``g`` be a DGLGraph. If ``g`` is a graph of a single edge type, ``g.edata[feat]``
        returns the edge feature associated with the name ``feat``. One can also set an
        edge feature associated with the name ``feat`` by setting ``g.edata[feat]`` to a tensor.
1998

1999
2000
2001
2002
2003
        If ``g`` is a graph of multiple edge types, ``g.edata[feat]`` returns a
        dict[str, Tensor] mapping canonical edge types to the edge features associated with
        the name ``feat`` for the corresponding type. One can also set an edge feature
        associated with the name ``feat`` for some edge type(s) by setting
        ``g.edata[feat]`` to a dictionary as described.
2004

2005
2006
2007
2008
        Notes
        -----
        For setting features, the device of the features must be the same as the device
        of the graph.
Minjie Wang's avatar
Minjie Wang committed
2009
2010
2011

        Examples
        --------
Mufei Li's avatar
Mufei Li committed
2012
2013
        The following example uses PyTorch backend.

2014
2015
        >>> import dgl
        >>> import torch
Mufei Li's avatar
Mufei Li committed
2016

2017
        Set and get feature 'h' for a graph of a single edge type.
2018

2019
2020
        >>> g = dgl.graph((torch.tensor([0, 1]), torch.tensor([1, 2])))
        >>> g.edata['h'] = torch.ones(2, 1)
2021
        >>> g.edata['h']
2022
2023
2024
2025
2026
2027
2028
2029
2030
2031
2032
2033
        tensor([[1.],
                [1.]])

        Set and get feature 'h' for a graph of multiple edge types.

        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([1, 2]), torch.tensor([3, 4])),
        ...     ('user', 'plays', 'user'): (torch.tensor([2, 2]), torch.tensor([1, 1])),
        ...     ('player', 'plays', 'game'): (torch.tensor([2, 2]), torch.tensor([1, 1]))
        ... })
        >>> g.edata['h'] = {('user', 'follows', 'user'): torch.zeros(2, 1),
        ...                 ('user', 'plays', 'user'): torch.ones(2, 1)}
2034
        >>> g.edata['h']
2035
2036
2037
        {('user', 'follows', 'user'): tensor([[0.], [0.]]),
         ('user', 'plays', 'user'): tensor([[1.], [1.]])}
        >>> g.edata['h'] = {('user', 'follows', 'user'): torch.ones(2, 1)}
2038
        >>> g.edata['h']
2039
2040
        {('user', 'follows', 'user'): tensor([[1.], [1.]]),
         ('user', 'plays', 'user'): tensor([[1.], [1.]])}
2041

Mufei Li's avatar
Mufei Li committed
2042
2043
2044
        See Also
        --------
        edges
2045
        """
2046
2047
2048
2049
        if len(self.canonical_etypes) == 1:
            return HeteroEdgeDataView(self, None, ALL)
        else:
            return HeteroEdgeDataView(self, self.canonical_etypes, ALL)
Minjie Wang's avatar
Minjie Wang committed
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061

    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.

2062
        You can get a relation slice with ``self[srctype, etype, dsttype]``, where
Minjie Wang's avatar
Minjie Wang committed
2063
2064
2065
2066
2067
2068
2069
2070
2071
2072
2073
2074
        ``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.

2075
        If there are multiple canonical edge types found, then the source/edge/destination
Minjie Wang's avatar
Minjie Wang committed
2076
2077
2078
2079
        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
2080
        common features of the original source/destination types.  Therefore they are not
Minjie Wang's avatar
Minjie Wang committed
2081
        shared with the original graph.  Edge type is similar.
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
2104
2105
2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136

        Parameters
        ----------
        key : str or tuple
            Either a string representing the edge type name, or a tuple in the form of
            ``(srctype, etype, dsttype)`` where ``srctype``, ``etype``, ``dsttype`` can be either
            strings representing type names or a full slice object (`:`).

        Returns
        -------
        DGLGraph
            The relation slice.

        Notes
        -----
        This function returns a new graph.  Changing the content of this graph does not reflect
        onto the original graph.

        If the graph combines multiple node types or edge types together, it will have the
        mapping of node/edge types and IDs from the new graph to the original graph.
        The mappings have the name ``dgl.NTYPE``, ``dgl.NID``, ``dgl.ETYPE`` and ``dgl.EID``,
        similar to the function :func:`dgl.to_homogenenous`.

        Examples
        --------
        >>> g = dgl.heterograph({
        ...     ('A1', 'AB1', 'B'): ([0, 1, 2], [1, 2, 3]),
        ...     ('A1', 'AB2', 'B'): ([1, 2, 3], [3, 4, 5]),
        ...     ('A2', 'AB2', 'B'): ([1, 3, 5], [2, 4, 6])})
        >>> new_g = g['A1', :, 'B']         # combines all edge types between A1 and B
        >>> new_g
        Graph(num_nodes={'A1': 4, 'B': 7},
              num_edges={('A1', 'AB1+AB2', 'B'): 6},
              metagraph=[('A1', 'B', 'AB1+AB2')])
        >>> new_g.edges()
        (tensor([0, 1, 2, 1, 2, 3]), tensor([1, 2, 3, 3, 4, 5]))
        >>> new_g2 = g[:, 'AB2', 'B']        # combines all node types that are source of AB2
        >>> new_g2
        Graph(num_nodes={'A1+A2': 10, 'B': 7},
              num_edges={('A1+A2', 'AB2+AB2', 'B'): 6},
              metagraph=[('A1+A2', 'B', 'AB2+AB2')])
        >>> new_g2.edges()
        (tensor([1, 2, 3, 5, 7, 9]), tensor([3, 4, 5, 2, 4, 6]))

        If a combination of multiple node types and edge types occur, one can find
        the mapping to the original node type and IDs like the following:

        >>> new_g1.edges['AB1+AB2'].data[dgl.EID]
        tensor([0, 1, 2, 0, 1, 2])
        >>> new_g1.edges['AB1+AB2'].data[dgl.ETYPE]
        tensor([0, 0, 0, 1, 1, 1])
        >>> new_g2.nodes['A1+A2'].data[dgl.NID]
        tensor([0, 1, 2, 3, 0, 1, 2, 3, 4, 5])
        >>> new_g2.nodes['A1+A2'].data[dgl.NTYPE]
        tensor([0, 0, 0, 0, 1, 1, 1, 1, 1, 1])
Minjie Wang's avatar
Minjie Wang committed
2137
2138
2139
2140
2141
        """
        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'])."

2142
        orig_key = key
Minjie Wang's avatar
Minjie Wang committed
2143
2144
2145
2146
2147
2148
2149
        if not isinstance(key, tuple):
            key = (SLICE_FULL, key, SLICE_FULL)

        if len(key) != 3:
            raise DGLError(err_msg)

        etypes = self._find_etypes(key)
2150
2151
2152
2153

        if len(etypes) == 0:
            raise DGLError('Invalid key "{}". Must be one of the edge types.'.format(orig_key))

Minjie Wang's avatar
Minjie Wang committed
2154
2155
2156
        if len(etypes) == 1:
            # no ambiguity: return the unitgraph itself
            srctype, etype, dsttype = self._canonical_etypes[etypes[0]]
2157
            stid = self.get_ntype_id_from_src(srctype)
Minjie Wang's avatar
Minjie Wang committed
2158
            etid = self.get_etype_id((srctype, etype, dsttype))
2159
            dtid = self.get_ntype_id_from_dst(dsttype)
Minjie Wang's avatar
Minjie Wang committed
2160
2161
2162
2163
2164
2165
            new_g = self._graph.get_relation_graph(etid)

            if stid == dtid:
                new_ntypes = [srctype]
                new_nframes = [self._node_frames[stid]]
            else:
2166
                new_ntypes = ([srctype], [dsttype])
Minjie Wang's avatar
Minjie Wang committed
2167
2168
2169
                new_nframes = [self._node_frames[stid], self._node_frames[dtid]]
            new_etypes = [etype]
            new_eframes = [self._edge_frames[etid]]
2170

2171
            return self.__class__(new_g, new_ntypes, new_etypes, new_nframes, new_eframes)
Minjie Wang's avatar
Minjie Wang committed
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
        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
2193
            new_hg = self.__class__(new_g, new_ntypes, new_etypes, new_nframes, new_eframes)
Minjie Wang's avatar
Minjie Wang committed
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
2207
2208
2209
2210
2211

            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):
2212
        """Alias of :meth:`num_nodes`"""
2213
2214
2215
        return self.num_nodes(ntype)

    def num_nodes(self, ntype=None):
2216
        """Return the number of nodes in the graph.
Da Zheng's avatar
Da Zheng committed
2217
2218
2219

        Parameters
        ----------
Minjie Wang's avatar
Minjie Wang committed
2220
        ntype : str, optional
2221
            The node type name. If given, it returns the number of nodes of the
2222
            type. If not given (default), it returns the total number of nodes of all types.
2223
2224
2225
2226

        Returns
        -------
        int
2227
            The number of nodes.
Da Zheng's avatar
Da Zheng committed
2228
2229
2230

        Examples
        --------
Mufei Li's avatar
Mufei Li committed
2231

2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
2251
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Create a graph with two node types -- 'user' and 'game'.

        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'plays', 'game'): (torch.tensor([3, 4]), torch.tensor([5, 6]))
        ... })

        Query for the number of nodes.

        >>> g.num_nodes('user')
        5
        >>> g.num_nodes('game')
        7
        >>> g.num_nodes()
        12
Da Zheng's avatar
Da Zheng committed
2252
        """
2253
2254
2255
2256
        if ntype is None:
            return sum([self._graph.number_of_nodes(ntid) for ntid in range(len(self.ntypes))])
        else:
            return self._graph.number_of_nodes(self.get_ntype_id(ntype))
2257

2258
    def number_of_src_nodes(self, ntype=None):
2259
        """Alias of :meth:`num_src_nodes`"""
2260
        return self.num_src_nodes(ntype)
2261

2262
    def num_src_nodes(self, ntype=None):
2263
2264
2265
2266
2267
2268
2269
2270
        """Return the number of source nodes in the graph.

        If the graph can further divide its node types into two subsets A and B where
        all the edeges are from nodes of types in A to nodes of types in B, we call
        this graph a *uni-bipartite* graph and the nodes in A being the *source*
        nodes and the ones in B being the *destination* nodes. If the graph is not
        uni-bipartite, the source and destination nodes are just the entire set of
        nodes in the graph.
2271
2272
2273
2274

        Parameters
        ----------
        ntype : str, optional
2275
2276
            The source node type name. If given, it returns the number of nodes for
            the source node type. If not given (default), it returns the number of
2277
            nodes summed over all source node types.
2278
2279
2280
2281
2282
2283

        Returns
        -------
        int
            The number of nodes

2284
2285
2286
2287
2288
        See Also
        --------
        num_dst_nodes
        is_unibipartite

2289
2290
        Examples
        --------
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Create a homogeneous graph for query.

        >>> g = dgl.graph((torch.tensor([0, 1]), torch.tensor([1, 2])))
        >>> g.num_src_nodes()
        3

        Create a heterogeneous graph with two source node types -- 'developer' and 'user'.

        >>> g = dgl.heterograph({
        ...     ('developer', 'develops', 'game'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'plays', 'game'): (torch.tensor([3, 4]), torch.tensor([5, 6]))
        ... })

        Query for the number of nodes.

        >>> g.num_src_nodes('developer')
2312
        2
2313
2314
2315
2316
        >>> g.num_src_nodes('user')
        5
        >>> g.num_src_nodes()
        7
2317
        """
2318
2319
2320
2321
2322
        if ntype is None:
            return sum([self._graph.number_of_nodes(self.get_ntype_id_from_src(nty))
                        for nty in self.srctypes])
        else:
            return self._graph.number_of_nodes(self.get_ntype_id_from_src(ntype))
2323
2324

    def number_of_dst_nodes(self, ntype=None):
2325
2326
        """Alias of :func:`num_dst_nodes`"""
        return self.num_dst_nodes(ntype)
2327

2328
    def num_dst_nodes(self, ntype=None):
2329
2330
2331
2332
2333
2334
2335
2336
        """Return the number of destination nodes in the graph.

        If the graph can further divide its node types into two subsets A and B where
        all the edeges are from nodes of types in A to nodes of types in B, we call
        this graph a *uni-bipartite* graph and the nodes in A being the *source*
        nodes and the ones in B being the *destination* nodes. If the graph is not
        uni-bipartite, the source and destination nodes are just the entire set of
        nodes in the graph.
2337
2338
2339
2340

        Parameters
        ----------
        ntype : str, optional
2341
2342
2343
            The destination node type name. If given, it returns the number of nodes of
            the destination node type. If not given (default), it returns the number of
            nodes summed over all the destination node types.
2344
2345
2346
2347
2348
2349

        Returns
        -------
        int
            The number of nodes

2350
2351
2352
2353
2354
        See Also
        --------
        num_src_nodes
        is_unibipartite

2355
2356
        Examples
        --------
2357
2358
2359
2360
2361
2362
2363
2364
2365
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Create a homogeneous graph for query.

        >>> g = dgl.graph((torch.tensor([0, 1]), torch.tensor([1, 2])))
        >>> g.num_dst_nodes()
2366
        3
2367
2368
2369
2370
2371
2372
2373
2374
2375
2376
2377
2378
2379
2380
2381
2382

        Create a heterogeneous graph with two destination node types -- 'user' and 'game'.

        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'plays', 'game'): (torch.tensor([3, 4]), torch.tensor([5, 6]))
        ... })

        Query for the number of nodes.

        >>> g.num_dst_nodes('user')
        5
        >>> g.num_dst_nodes('game')
        7
        >>> g.num_dst_nodes()
        12
2383
        """
2384
2385
2386
2387
2388
        if ntype is None:
            return sum([self._graph.number_of_nodes(self.get_ntype_id_from_dst(nty))
                        for nty in self.dsttypes])
        else:
            return self._graph.number_of_nodes(self.get_ntype_id_from_dst(ntype))
2389

Minjie Wang's avatar
Minjie Wang committed
2390
    def number_of_edges(self, etype=None):
2391
2392
2393
2394
        """Alias of :func:`num_edges`"""
        return self.num_edges(etype)

    def num_edges(self, etype=None):
2395
        """Return the number of edges in the graph.
2396
2397
2398

        Parameters
        ----------
2399
2400
2401
2402
2403
2404
2405
2406
2407
        etype : str or (str, str, str), optional
            The type name of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            If not provided, return the total number of edges regardless of the types
            in the graph.
Da Zheng's avatar
Da Zheng committed
2408
2409
2410
2411

        Returns
        -------
        int
2412
            The number of edges.
Da Zheng's avatar
Da Zheng committed
2413

2414
2415
        Examples
        --------
Mufei Li's avatar
Mufei Li committed
2416

2417
2418
2419
2420
2421
2422
2423
2424
2425
2426
2427
2428
2429
2430
2431
2432
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Create a graph with three canonical edge types.

        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'follows', 'game'): (torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])),
        ...     ('user', 'plays', 'game'): (torch.tensor([1, 3]), torch.tensor([2, 3]))
        ... })

        Query for the number of edges.

        >>> g.num_edges('plays')
Mufei Li's avatar
Mufei Li committed
2433
        2
2434
2435
2436
2437
2438
2439
        >>> g.num_edges()
        7

        Use a canonical edge type instead when there is ambiguity for an edge type.

        >>> g.num_edges(('user', 'follows', 'user'))
Mufei Li's avatar
Mufei Li committed
2440
        2
2441
2442
        >>> g.num_edges(('user', 'follows', 'game'))
        3
2443
        """
2444
2445
2446
2447
2448
        if etype is None:
            return sum([self._graph.number_of_edges(etid)
                        for etid in range(len(self.canonical_etypes))])
        else:
            return self._graph.number_of_edges(self.get_etype_id(etype))
Minjie Wang's avatar
Minjie Wang committed
2449
2450
2451

    @property
    def is_multigraph(self):
2452
        """Return whether the graph is a multigraph with parallel edges.
Mufei Li's avatar
Mufei Li committed
2453

2454
2455
2456
2457
        A multigraph has more than one edges between the same pair of nodes, called
        *parallel edges*.  For heterogeneous graphs, parallel edge further requires
        the canonical edge type to be the same (see :meth:`canonical_etypes` for the
        definition).
2458

Mufei Li's avatar
Mufei Li committed
2459
2460
2461
        Returns
        -------
        bool
2462
            True if the graph is a multigraph.
2463
2464
2465

        Notes
        -----
2466
        Checking whether the graph is a multigraph could be expensive for a large one.
2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478
2479
2480
2481
2482
2483
2484
2485
2486
2487
2488
2489
2490
2491
2492
2493
2494
2495
2496
2497
2498

        Examples
        --------

        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Check for homogeneous graphs.

        >>> g = dgl.graph((torch.tensor([0, 1]), torch.tensor([1, 3])))
        >>> g.is_multigraph
        False
        >>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([1, 3, 3])))
        >>> g.is_multigraph
        True

        Check for heterogeneous graphs.

        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'follows', 'game'): (torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3]))
        ... })
        >>> g.is_multigraph
        False
        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1, 1]), torch.tensor([1, 2, 2])),
        ...     ('user', 'follows', 'game'): (torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3]))
        ... })
        >>> g.is_multigraph
        True
Mufei Li's avatar
Mufei Li committed
2499
        """
2500
        return self._graph.is_multigraph()
Minjie Wang's avatar
Minjie Wang committed
2501

2502
2503
    @property
    def is_homogeneous(self):
2504
        """Return whether the graph is a homogeneous graph.
2505
2506
2507
2508
2509
2510

        A homogeneous graph only has one node type and one edge type.

        Returns
        -------
        bool
2511
            True if the graph is a homogeneous graph.
2512
2513
2514
2515
2516
2517
2518
2519
2520
2521
2522
2523
2524
2525
2526
2527
2528
2529
2530
2531
2532
2533
2534
2535
2536
2537

        Examples
        --------

        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Create a homogeneous graph for check.

        >>> g = dgl.graph((torch.tensor([0, 0, 1, 1]), torch.tensor([1, 0, 2, 3])))
        >>> g.is_homogeneous
        True

        Create a heterogeneous graph for check.

        If the graph has multiple edge types, one need to specify the edge type.

        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'game'): (torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3]))})
        >>> g.is_homogeneous
        False
        """
        return len(self.ntypes) == 1 and len(self.etypes) == 1

Minjie Wang's avatar
Minjie Wang committed
2538
2539
    @property
    def is_readonly(self):
2540
        """**DEPRECATED**: DGLGraph will always be mutable.
Mufei Li's avatar
Mufei Li committed
2541
2542
2543
2544
2545
2546

        Returns
        -------
        bool
            True if the graph is readonly, False otherwise.
        """
2547
2548
2549
2550
        dgl_warning('DGLGraph.is_readonly is deprecated in v0.5.\n'
                    'DGLGraph now always supports mutable operations like add_nodes'
                    ' and add_edges.')
        return False
Da Zheng's avatar
Da Zheng committed
2551

2552
2553
    @property
    def idtype(self):
2554
2555
        """The data type for storing the structure-related graph information
        such as node and edge IDs.
2556
2557
2558

        Returns
        -------
2559
2560
2561
2562
2563
2564
2565
2566
2567
2568
2569
2570
2571
2572
2573
2574
2575
2576
2577
        Framework-specific device object
            For example, this can be ``torch.int32`` or ``torch.int64`` for PyTorch.

        Examples
        --------

        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        >>> src_ids = torch.tensor([0, 0, 1])
        >>> dst_ids = torch.tensor([1, 2, 2])
        >>> g = dgl.graph((src_ids, dst_ids))
        >>> g.idtype
        torch.int64
        >>> g = dgl.graph((src_ids, dst_ids), idtype=torch.int32)
        >>> g.idtype
        torch.int32
2578
2579
2580
2581
2582

        See Also
        --------
        long
        int
2583
2584
2585
2586
2587
2588
2589
2590
2591
2592
2593
2594
2595
2596
        """
        return getattr(F, self._graph.dtype)

    @property
    def _idtype_str(self):
        """The dtype of graph index

        Returns
        -------
        backend dtype object
            th.int32/th.int64 or tf.int32/tf.int64 etc.
        """
        return self._graph.dtype

2597
    def __contains__(self, vid):
2598
        """**DEPRECATED**: please directly call :func:`has_nodes`."""
2599
2600
2601
2602
2603
        dgl_warning('DGLGraph.__contains__ is deprecated.'
                    ' Please directly call has_nodes.')
        return self.has_nodes(vid)

    def has_nodes(self, vid, ntype=None):
2604
        """Return whether the graph contains the given nodes.
Da Zheng's avatar
Da Zheng committed
2605
2606
2607

        Parameters
        ----------
2608
        vid : node ID(s)
2609
2610
2611
2612
2613
2614
            The nodes IDs. The allowed nodes ID formats are:

            * ``int``: The ID of a single node.
            * Int Tensor: Each element is a node ID. The tensor must have the same device type
              and ID data type as the graph's.
            * iterable[int]: Each element is a node ID.
2615

Minjie Wang's avatar
Minjie Wang committed
2616
        ntype : str, optional
2617
2618
            The node type name. Can be omitted if there is
            only one type of nodes in the graph.
Da Zheng's avatar
Da Zheng committed
2619
2620
2621

        Returns
        -------
2622
        bool or bool Tensor
2623
2624
            A tensor of bool flags where each element is True if the node is in the graph.
            If the input is a single node, return one bool value.
Da Zheng's avatar
Da Zheng committed
2625
2626
2627

        Examples
        --------
2628
2629
2630
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642

        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Create a graph with two node types -- 'user' and 'game'.

        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'plays', 'game'): (torch.tensor([3, 4]), torch.tensor([0, 1]))
        ... })

        Query for the nodes.

2643
        >>> g.has_nodes(0, 'user')
Da Zheng's avatar
Da Zheng committed
2644
        True
2645
        >>> g.has_nodes(3, 'game')
Da Zheng's avatar
Da Zheng committed
2646
        False
2647
2648
        >>> g.has_nodes(torch.tensor([3, 0, 1]), 'game')
        tensor([False,  True,  True])
Da Zheng's avatar
Da Zheng committed
2649
        """
2650
2651
2652
        vid_tensor = utils.prepare_tensor(self, vid, "vid")
        if len(vid_tensor) > 0 and F.as_scalar(F.min(vid_tensor, 0)) < 0 < len(vid_tensor):
            raise DGLError('All IDs must be non-negative integers.')
2653
        ret = self._graph.has_nodes(
2654
            self.get_ntype_id(ntype), vid_tensor)
2655
2656
2657
2658
        if isinstance(vid, numbers.Integral):
            return bool(F.as_scalar(ret))
        else:
            return F.astype(ret, F.bool)
Da Zheng's avatar
Da Zheng committed
2659

2660
    def has_node(self, vid, ntype=None):
2661
        """Whether the graph has a particular node of a given type.
Da Zheng's avatar
Da Zheng committed
2662

2663
        **DEPRECATED**: see :func:`~DGLGraph.has_nodes`
Da Zheng's avatar
Da Zheng committed
2664
        """
2665
2666
        dgl_warning("DGLGraph.has_node is deprecated. Please use DGLGraph.has_nodes")
        return self.has_nodes(vid, ntype)
Da Zheng's avatar
Da Zheng committed
2667

2668
    def has_edges_between(self, u, v, etype=None):
2669
        """Return whether the graph contains the given edges.
Da Zheng's avatar
Da Zheng committed
2670
2671
2672

        Parameters
        ----------
2673
2674
2675
2676
2677
2678
2679
2680
2681
2682
2683
2684
2685
2686
2687
2688
2689
2690
2691
2692
2693
2694
2695
2696
2697
        u : node IDs
            The source node IDs of the edges. The allowed formats are:

            * ``int``: A single node.
            * Int Tensor: Each element is a node ID. The tensor must have the same device type
              and ID data type as the graph's.
            * iterable[int]: Each element is a node ID.

        v : node IDs
            The destination node IDs of the edges. The allowed formats are:

            * ``int``: A single node.
            * Int Tensor: Each element is a node ID. The tensor must have the same device type
              and ID data type as the graph's.
            * iterable[int]: Each element is a node ID.

        etype : str or (str, str, str), optional
            The type names of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            Can be omitted if the graph has only one type of edges.

Da Zheng's avatar
Da Zheng committed
2698
2699
2700

        Returns
        -------
2701
        bool or bool Tensor
2702
2703
            A tensor of bool flags where each element is True if the node is in the graph.
            If the input is a single node, return one bool value.
Da Zheng's avatar
Da Zheng committed
2704
2705
2706

        Examples
        --------
Mufei Li's avatar
Mufei Li committed
2707

2708
2709
2710
2711
2712
2713
2714
2715
2716
2717
2718
2719
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Create a homogeneous graph.

        >>> g = dgl.graph((torch.tensor([0, 0, 1, 1]), torch.tensor([1, 0, 2, 3])))

        Query for the edges.

        >>> g.has_edges_between(1, 2)
Da Zheng's avatar
Da Zheng committed
2720
        True
2721
2722
2723
2724
2725
2726
2727
2728
2729
2730
2731
2732
2733
2734
2735
2736
2737
2738
2739
2740
2741
        >>> g.has_edges_between(torch.tensor([1, 2]), torch.tensor([2, 3]))
        tensor([ True, False])

        If the graph has multiple edge types, one need to specify the edge type.

        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'follows', 'game'): (torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])),
        ...     ('user', 'plays', 'game'): (torch.tensor([1, 3]), torch.tensor([2, 3]))
        ... })
        >>> g.has_edges_between(torch.tensor([1, 2]), torch.tensor([2, 3]), 'plays')
        tensor([ True, False])

        Use a canonical edge type instead when there is ambiguity for an edge type.

        >>> g.has_edges_between(torch.tensor([1, 2]), torch.tensor([2, 3]),
        ...                     ('user', 'follows', 'user'))
        tensor([ True, False])
        >>> g.has_edges_between(torch.tensor([1, 2]), torch.tensor([2, 3]),
        ...                     ('user', 'follows', 'game'))
        tensor([True, True])
Da Zheng's avatar
Da Zheng committed
2742
        """
2743
2744
2745
2746
2747
2748
2749
        srctype, _, dsttype = self.to_canonical_etype(etype)
        u_tensor = utils.prepare_tensor(self, u, 'u')
        if F.as_scalar(F.sum(self.has_nodes(u_tensor, ntype=srctype), dim=0)) != len(u_tensor):
            raise DGLError('u contains invalid node IDs')
        v_tensor = utils.prepare_tensor(self, v, 'v')
        if F.as_scalar(F.sum(self.has_nodes(v_tensor, ntype=dsttype), dim=0)) != len(v_tensor):
            raise DGLError('v contains invalid node IDs')
2750
2751
        ret = self._graph.has_edges_between(
            self.get_etype_id(etype),
2752
            u_tensor, v_tensor)
2753
2754
2755
2756
        if isinstance(u, numbers.Integral) and isinstance(v, numbers.Integral):
            return bool(F.as_scalar(ret))
        else:
            return F.astype(ret, F.bool)
Da Zheng's avatar
Da Zheng committed
2757

2758
    def has_edge_between(self, u, v, etype=None):
Mufei Li's avatar
Mufei Li committed
2759
        """Whether the graph has edges of type ``etype``.
Da Zheng's avatar
Da Zheng committed
2760

2761
        **DEPRECATED**: please use :func:`~DGLGraph.has_edge_between`.
Da Zheng's avatar
Da Zheng committed
2762
        """
2763
2764
2765
        dgl_warning("DGLGraph.has_edge_between is deprecated. "
                    "Please use DGLGraph.has_edges_between")
        return self.has_edges_between(u, v, etype)
Da Zheng's avatar
Da Zheng committed
2766

Minjie Wang's avatar
Minjie Wang committed
2767
    def predecessors(self, v, etype=None):
2768
        """Return the predecessor(s) of a particular node with the specified edge type.
Da Zheng's avatar
Da Zheng committed
2769

2770
2771
        Node ``u`` is a predecessor of node ``v`` if there is an edge ``(u, v)`` with type
        ``etype`` in the graph.
Da Zheng's avatar
Da Zheng committed
2772
2773
2774
2775

        Parameters
        ----------
        v : int
2776
2777
            The node ID. If the graph has multiple edge types, the ID is for the destination
            type corresponding to the edge type.
2778
2779
2780
2781
2782
2783
2784
2785
2786
        etype : str or (str, str, str), optional
            The type names of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            Can be omitted if the graph has only one type of edges.

Da Zheng's avatar
Da Zheng committed
2787
2788
2789

        Returns
        -------
2790
2791
        Tensor
            The predecessors of :attr:`v` with the specified edge type.
Da Zheng's avatar
Da Zheng committed
2792
2793
2794
2795
2796

        Examples
        --------
        The following example uses PyTorch backend.

2797
2798
2799
2800
2801
2802
2803
2804
2805
2806
2807
2808
2809
2810
2811
2812
2813
2814
2815
        >>> import dgl
        >>> import torch

        Create a homogeneous graph.

        >>> g = dgl.graph((torch.tensor([0, 0, 1, 1]), torch.tensor([1, 1, 2, 3])))

        Query for node 1.

        >>> g.predecessors(1)
        tensor([0, 0])

        For a graph of multiple edge types, it is required to specify the edge type in query.

        >>> hg = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'plays', 'game'): (torch.tensor([3, 4]), torch.tensor([5, 6]))
        ... })
        >>> hg.predecessors(1, etype='follows')
Mufei Li's avatar
Mufei Li committed
2816
        tensor([0])
Da Zheng's avatar
Da Zheng committed
2817
2818
2819
2820
2821

        See Also
        --------
        successors
        """
2822
2823
        if not self.has_nodes(v, self.to_canonical_etype(etype)[-1]):
            raise DGLError('Non-existing node ID {}'.format(v))
2824
        return self._graph.predecessors(self.get_etype_id(etype), v)
Da Zheng's avatar
Da Zheng committed
2825

Minjie Wang's avatar
Minjie Wang committed
2826
    def successors(self, v, etype=None):
2827
        """Return the successor(s) of a particular node with the specified edge type.
Da Zheng's avatar
Da Zheng committed
2828

2829
2830
        Node ``u`` is a successor of node ``v`` if there is an edge ``(v, u)`` with type
        ``etype`` in the graph.
Da Zheng's avatar
Da Zheng committed
2831
2832
2833
2834

        Parameters
        ----------
        v : int
2835
2836
            The node ID. If the graph has multiple edge types, the ID is for the source
            type corresponding to the edge type.
2837
2838
2839
2840
2841
2842
2843
2844
        etype : str or (str, str, str), optional
            The type names of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            Can be omitted if the graph has only one type of edges.
Da Zheng's avatar
Da Zheng committed
2845
2846
2847

        Returns
        -------
2848
2849
        Tensor
            The successors of :attr:`v` with the specified edge type.
Da Zheng's avatar
Da Zheng committed
2850
2851
2852
2853
2854

        Examples
        --------
        The following example uses PyTorch backend.

2855
2856
2857
2858
2859
2860
2861
2862
2863
2864
2865
2866
2867
2868
2869
2870
2871
2872
2873
2874
        >>> import dgl
        >>> import torch

        Create a homogeneous graph.

        >>> g = dgl.graph((torch.tensor([0, 0, 1, 1]), torch.tensor([1, 1, 2, 3])))

        Query for node 1.

        >>> g.successors(1)
        tensor([2, 3])

        For a graph of multiple edge types, it is required to specify the edge type in query.

        >>> hg = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'plays', 'game'): (torch.tensor([3, 4]), torch.tensor([5, 6]))
        ... })
        >>> hg.successors(1, etype='follows')
        tensor([2])
Da Zheng's avatar
Da Zheng committed
2875
2876
2877
2878
2879

        See Also
        --------
        predecessors
        """
2880
2881
        if not self.has_nodes(v, self.to_canonical_etype(etype)[0]):
            raise DGLError('Non-existing node ID {}'.format(v))
2882
        return self._graph.successors(self.get_etype_id(etype), v)
Da Zheng's avatar
Da Zheng committed
2883

2884
    def edge_id(self, u, v, force_multi=None, return_uv=False, etype=None):
Da Zheng's avatar
Da Zheng committed
2885
        """Return the edge ID, or an array of edge IDs, between source node
Mufei Li's avatar
Mufei Li committed
2886
        `u` and destination node `v`, with the specified edge type
Da Zheng's avatar
Da Zheng committed
2887

2888
2889
2890
2891
2892
2893
2894
        **DEPRECATED**: See edge_ids
        """
        dgl_warning("DGLGraph.edge_id is deprecated. Please use DGLGraph.edge_ids.")
        return self.edge_ids(u, v, force_multi=force_multi,
                             return_uv=return_uv, etype=etype)

    def edge_ids(self, u, v, force_multi=None, return_uv=False, etype=None):
2895
        """Return the edge ID(s) given the two endpoints of the edge(s).
2896

Da Zheng's avatar
Da Zheng committed
2897
2898
        Parameters
        ----------
2899
2900
2901
2902
2903
2904
2905
2906
2907
2908
2909
2910
2911
2912
2913
        u : node IDs
            The source node IDs of the edges. The allowed formats are:

            * ``int``: A single node.
            * Int Tensor: Each element is a node ID. The tensor must have the same device type
              and ID data type as the graph's.
            * iterable[int]: Each element is a node ID.

        v : node IDs
            The destination node IDs of the edges. The allowed formats are:

            * ``int``: A single node.
            * Int Tensor: Each element is a node ID. The tensor must have the same device type
              and ID data type as the graph's.
            * iterable[int]: Each element is a node ID.
Minjie Wang's avatar
Minjie Wang committed
2914
        force_multi : bool, optional
2915
            **DEPRECATED**, use :attr:`return_uv` instead. Whether to allow the graph to be a
2916
2917
2918
2919
2920
2921
            multigraph, i.e. there can be multiple edges from one node to another.
        return_uv : bool, optional
            Whether to return the source and destination node IDs along with the edges. If
            False (default), it assumes that the graph is a simple graph and there is only
            one edge from one node to another. If True, there can be multiple edges found
            from one node to another.
2922
2923
2924
2925
2926
2927
2928
2929
        etype : str or (str, str, str), optional
            The type names of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            Can be omitted if the graph has only one type of edges.
Da Zheng's avatar
Da Zheng committed
2930
2931
2932

        Returns
        -------
2933
        Tensor, or (Tensor, Tensor, Tensor)
Mufei Li's avatar
Mufei Li committed
2934

2935
2936
            * If ``return_uv=False``, it returns the edge IDs in a tensor, where the i-th
              element is the ID of the edge ``(u[i], v[i])``.
2937
2938
            * If ``return_uv=True``, it returns a tuple of three 1D tensors ``(eu, ev, e)``.
              ``e[i]`` is the ID of an edge from ``eu[i]`` to ``ev[i]``. It returns all edges
2939
              (including parallel edges) from ``eu[i]`` to ``ev[i]`` in this case.
Da Zheng's avatar
Da Zheng committed
2940
2941
2942

        Notes
        -----
2943
2944
        If the graph is a simple graph, ``return_uv=False``, and there are no edges
        between some pairs of node(s), it will raise an error.
Da Zheng's avatar
Da Zheng committed
2945

2946
2947
        If the graph is a multigraph, ``return_uv=False``, and there are multiple edges
        between some pairs of node(s), it returns an arbitrary one from them.
2948

Da Zheng's avatar
Da Zheng committed
2949
2950
2951
2952
        Examples
        --------
        The following example uses PyTorch backend.

2953
2954
        >>> import dgl
        >>> import torch
Mufei Li's avatar
Mufei Li committed
2955

2956
        Create a homogeneous graph.
Mufei Li's avatar
Mufei Li committed
2957

2958
        >>> g = dgl.graph((torch.tensor([0, 0, 1, 1, 1]), torch.tensor([1, 0, 2, 3, 2])))
Mufei Li's avatar
Mufei Li committed
2959

2960
2961
2962
2963
2964
2965
2966
2967
2968
2969
2970
2971
2972
2973
2974
2975
2976
2977
2978
2979
2980
2981
2982
2983
2984
2985
2986
2987
2988
2989
        Query for the edges.

        >>> g.edge_ids(0, 0)
        1
        >>> g.edge_ids(torch.tensor([1, 0]), torch.tensor([3, 1]))
        tensor([3, 0])

        Get all edges for pairs of nodes.

        >>> g.edge_ids(torch.tensor([1, 0]), torch.tensor([3, 1]), return_uv=True)
        (tensor([1, 0]), tensor([3, 1]), tensor([3, 0]))

        If the graph has multiple edge types, one need to specify the edge type.

        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'follows', 'game'): (torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])),
        ...     ('user', 'plays', 'game'): (torch.tensor([1, 3]), torch.tensor([2, 3]))
        ... })
        >>> g.edge_ids(torch.tensor([1]), torch.tensor([2]), etype='plays')
        tensor([0])

        Use a canonical edge type instead when there is ambiguity for an edge type.

        >>> g.edge_ids(torch.tensor([0, 1]), torch.tensor([1, 2]),
        ...            etype=('user', 'follows', 'user'))
        tensor([0, 1])
        >>> g.edge_ids(torch.tensor([1, 2]), torch.tensor([2, 3]),
        ...            etype=('user', 'follows', 'game'))
        tensor([1, 2])
Da Zheng's avatar
Da Zheng committed
2990
        """
2991
        is_int = isinstance(u, numbers.Integral) and isinstance(v, numbers.Integral)
2992
        srctype, _, dsttype = self.to_canonical_etype(etype)
2993
        u = utils.prepare_tensor(self, u, 'u')
2994
2995
        if F.as_scalar(F.sum(self.has_nodes(u, ntype=srctype), dim=0)) != len(u):
            raise DGLError('u contains invalid node IDs')
2996
        v = utils.prepare_tensor(self, v, 'v')
2997
2998
        if F.as_scalar(F.sum(self.has_nodes(v, ntype=dsttype), dim=0)) != len(v):
            raise DGLError('v contains invalid node IDs')
2999
3000
3001
3002
3003
3004
        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:
3005
            return self._graph.edge_ids_all(self.get_etype_id(etype), u, v)
3006
        else:
3007
3008
3009
3010
3011
3012
3013
3014
3015
            eid = self._graph.edge_ids_one(self.get_etype_id(etype), u, v)
            is_neg_one = F.equal(eid, -1)
            if F.as_scalar(F.sum(is_neg_one, 0)):
                # Raise error since some (u, v) pair is not a valid edge.
                idx = F.nonzero_1d(is_neg_one)
                raise DGLError("Error: (%d, %d) does not form a valid edge." % (
                    F.as_scalar(F.gather_row(u, idx)),
                    F.as_scalar(F.gather_row(v, idx))))
            return F.as_scalar(eid) if is_int else eid
Da Zheng's avatar
Da Zheng committed
3016

Minjie Wang's avatar
Minjie Wang committed
3017
    def find_edges(self, eid, etype=None):
3018
        """Return the source and destination node ID(s) given the edge ID(s).
Da Zheng's avatar
Da Zheng committed
3019
3020
3021

        Parameters
        ----------
3022
        eid : edge ID(s)
3023
3024
3025
3026
3027
3028
3029
3030
3031
3032
3033
3034
3035
3036
3037
            The edge IDs. The allowed formats are:

            * ``int``: A single ID.
            * Int Tensor: Each element is an ID. The tensor must have the same device type
              and ID data type as the graph's.
            * iterable[int]: Each element is an ID.

        etype : str or (str, str, str), optional
            The type names of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            Can be omitted if the graph has only one type of edges.
Da Zheng's avatar
Da Zheng committed
3038
3039
3040

        Returns
        -------
3041
        Tensor
3042
3043
            The source node IDs of the edges. The i-th element is the source node ID of
            the i-th edge.
3044
        Tensor
3045
3046
            The destination node IDs of the edges. The i-th element is the destination node
            ID of the i-th edge.
Da Zheng's avatar
Da Zheng committed
3047
3048
3049
3050
3051

        Examples
        --------
        The following example uses PyTorch backend.

3052
3053
3054
3055
3056
3057
3058
3059
3060
3061
3062
3063
3064
3065
3066
3067
3068
3069
3070
3071
        >>> import dgl
        >>> import torch

        Create a homogeneous graph.

        >>> g = dgl.graph((torch.tensor([0, 0, 1, 1]), torch.tensor([1, 0, 2, 3])))

        Find edges of IDs 0 and 2.

        >>> g.find_edges(torch.tensor([0, 2]))
        (tensor([0, 1]), tensor([1, 2]))

        For a graph of multiple edge types, it is required to specify the edge type in query.

        >>> hg = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'plays', 'game'): (torch.tensor([3, 4]), torch.tensor([5, 6]))
        ... })
        >>> hg.find_edges(torch.tensor([1, 0]), 'plays')
        (tensor([4, 3]), tensor([6, 5]))
Da Zheng's avatar
Da Zheng committed
3072
        """
3073
        eid = utils.prepare_tensor(self, eid, 'eid')
3074
3075
3076
3077
3078
3079
3080
3081
        if len(eid) > 0:
            min_eid = F.as_scalar(F.min(eid, 0))
            if min_eid < 0:
                raise DGLError('Invalid edge ID {:d}'.format(min_eid))
            max_eid = F.as_scalar(F.max(eid, 0))
            if max_eid >= self.num_edges(etype):
                raise DGLError('Invalid edge ID {:d}'.format(max_eid))

3082
        if len(eid) == 0:
3083
3084
            empty = F.copy_to(F.tensor([], self.idtype), self.device)
            return empty, empty
Minjie Wang's avatar
Minjie Wang committed
3085
        src, dst, _ = self._graph.find_edges(self.get_etype_id(etype), eid)
3086
        return src, dst
Da Zheng's avatar
Da Zheng committed
3087

Minjie Wang's avatar
Minjie Wang committed
3088
    def in_edges(self, v, form='uv', etype=None):
3089
        """Return the incoming edges of the given nodes.
Da Zheng's avatar
Da Zheng committed
3090
3091
3092

        Parameters
        ----------
3093
3094
        v : node ID(s)
            The node IDs. The allowed formats are:
3095

3096
3097
3098
3099
            * ``int``: A single node.
            * Int Tensor: Each element is a node ID. The tensor must have the same device type
              and ID data type as the graph's.
            * iterable[int]: Each element is a node ID.
Da Zheng's avatar
Da Zheng committed
3100
        form : str, optional
3101
            The result format, which can be one of the following:
3102
3103
3104
3105
3106
3107
3108
3109
3110

            - ``'eid'``: The returned result is a 1D tensor :math:`EID`, representing
              the IDs of all edges.
            - ``'uv'`` (default): The returned result is a 2-tuple of 1D tensors :math:`(U, V)`,
              representing the source and destination nodes of all edges. For each :math:`i`,
              :math:`(U[i], V[i])` forms an edge.
            - ``'all'``: The returned result is a 3-tuple of 1D tensors :math:`(U, V, EID)`,
              representing the source nodes, destination nodes and IDs of all edges.
              For each :math:`i`, :math:`(U[i], V[i])` forms an edge with ID :math:`EID[i]`.
3111
3112
3113
3114
3115
3116
3117
3118
        etype : str or (str, str, str), optional
            The type names of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            Can be omitted if the graph has only one type of edges.
Da Zheng's avatar
Da Zheng committed
3119
3120
3121

        Returns
        -------
3122
3123
3124
        Tensor or (Tensor, Tensor) or (Tensor, Tensor, Tensor)
            All incoming edges of the nodes with the specified type. For a description of the
            returned result, see the description of :attr:`form`.
Da Zheng's avatar
Da Zheng committed
3125
3126
3127
3128
3129

        Examples
        --------
        The following example uses PyTorch backend.

3130
3131
3132
3133
3134
3135
3136
3137
3138
3139
3140
3141
3142
3143
3144
3145
3146
3147
3148
3149
3150
3151
3152
3153
3154
3155
3156
3157
3158
3159
        >>> import dgl
        >>> import torch

        Create a homogeneous graph.

        >>> g = dgl.graph((torch.tensor([0, 0, 1, 1]), torch.tensor([1, 0, 2, 3])))

        Query for the nodes 1 and 0.

        >>> g.in_edges(torch.tensor([1, 0]))
        (tensor([0, 0]), tensor([1, 0]))

        Specify a different value for :attr:`form`.

        >>> g.in_edges(torch.tensor([1, 0]), form='all')
        (tensor([0, 0]), tensor([1, 0]), tensor([0, 1]))

        For a graph of multiple edge types, it is required to specify the edge type in query.

        >>> hg = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'plays', 'game'): (torch.tensor([3, 4]), torch.tensor([5, 6]))
        ... })
        >>> hg.in_edges(torch.tensor([1, 0]), etype='follows')
        (tensor([0]), tensor([1]))

        See Also
        --------
        edges
        out_edges
Da Zheng's avatar
Da Zheng committed
3160
        """
3161
        v = utils.prepare_tensor(self, v, 'v')
Minjie Wang's avatar
Minjie Wang committed
3162
        src, dst, eid = self._graph.in_edges(self.get_etype_id(etype), v)
3163
        if form == 'all':
3164
            return src, dst, eid
3165
        elif form == 'uv':
3166
            return src, dst
3167
        elif form == 'eid':
3168
            return eid
3169
        else:
3170
            raise DGLError('Invalid form: {}. Must be "all", "uv" or "eid".'.format(form))
3171

Mufei Li's avatar
Mufei Li committed
3172
    def out_edges(self, u, form='uv', etype=None):
3173
        """Return the outgoing edges of the given nodes.
Da Zheng's avatar
Da Zheng committed
3174
3175
3176

        Parameters
        ----------
3177
3178
        u : node ID(s)
            The node IDs. The allowed formats are:
3179

3180
3181
3182
3183
            * ``int``: A single node.
            * Int Tensor: Each element is a node ID. The tensor must have the same device type
              and ID data type as the graph's.
            * iterable[int]: Each element is a node ID.
Da Zheng's avatar
Da Zheng committed
3184
        form : str, optional
3185
3186
3187
3188
3189
3190
3191
3192
3193
3194
            The return form, which can be one of the following:

            - ``'eid'``: The returned result is a 1D tensor :math:`EID`, representing
              the IDs of all edges.
            - ``'uv'`` (default): The returned result is a 2-tuple of 1D tensors :math:`(U, V)`,
              representing the source and destination nodes of all edges. For each :math:`i`,
              :math:`(U[i], V[i])` forms an edge.
            - ``'all'``: The returned result is a 3-tuple of 1D tensors :math:`(U, V, EID)`,
              representing the source nodes, destination nodes and IDs of all edges.
              For each :math:`i`, :math:`(U[i], V[i])` forms an edge with ID :math:`EID[i]`.
3195
3196
3197
3198
3199
3200
3201
3202
        etype : str or (str, str, str), optional
            The type names of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            Can be omitted if the graph has only one type of edges.
3203
3204
3205
3206
3207
3208
3209
3210
3211
3212
3213
3214
3215
3216
3217
3218
3219
3220
3221
3222
3223
3224
3225
3226
3227
3228
3229

        Returns
        -------
        Tensor or (Tensor, Tensor) or (Tensor, Tensor, Tensor)
            All outgoing edges of the nodes with the specified type. For a description of the
            returned result, see the description of :attr:`form`.

        Examples
        --------
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Create a homogeneous graph.

        >>> g = dgl.graph((torch.tensor([0, 0, 1, 1]), torch.tensor([1, 0, 2, 3])))

        Query for the nodes 1 and 2.

        >>> g.out_edges(torch.tensor([1, 2]))
        (tensor([1, 1]), tensor([2, 3]))

        Specify a different value for :attr:`form`.

        >>> g.out_edges(torch.tensor([1, 2]), form='all')
        (tensor([1, 1]), tensor([2, 3]), tensor([2, 3]))
Da Zheng's avatar
Da Zheng committed
3230

3231
        For a graph of multiple edge types, it is required to specify the edge type in query.
Mufei Li's avatar
Mufei Li committed
3232

3233
3234
3235
3236
3237
3238
        >>> hg = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'plays', 'game'): (torch.tensor([3, 4]), torch.tensor([5, 6]))
        ... })
        >>> hg.out_edges(torch.tensor([1, 2]), etype='follows')
        (tensor([1]), tensor([2]))
Da Zheng's avatar
Da Zheng committed
3239

3240
        See Also
Da Zheng's avatar
Da Zheng committed
3241
        --------
3242
3243
        edges
        in_edges
Da Zheng's avatar
Da Zheng committed
3244
        """
3245
        u = utils.prepare_tensor(self, u, 'u')
3246
3247
3248
        srctype, _, _ = self.to_canonical_etype(etype)
        if F.as_scalar(F.sum(self.has_nodes(u, ntype=srctype), dim=0)) != len(u):
            raise DGLError('u contains invalid node IDs')
Mufei Li's avatar
Mufei Li committed
3249
        src, dst, eid = self._graph.out_edges(self.get_etype_id(etype), u)
3250
        if form == 'all':
3251
            return src, dst, eid
3252
        elif form == 'uv':
3253
            return src, dst
3254
        elif form == 'eid':
3255
            return eid
3256
        else:
3257
            raise DGLError('Invalid form: {}. Must be "all", "uv" or "eid".'.format(form))
3258

3259
3260
    def all_edges(self, form='uv', order='eid', etype=None):
        """Return all edges with the specified edge type.
Da Zheng's avatar
Da Zheng committed
3261
3262
3263
3264

        Parameters
        ----------
        form : str, optional
3265
3266
3267
3268
3269
3270
3271
3272
3273
3274
3275
3276
3277
3278
3279
3280
            The return form, which can be one of the following:

            - ``'eid'``: The returned result is a 1D tensor :math:`EID`, representing
              the IDs of all edges.
            - ``'uv'`` (default): The returned result is a 2-tuple of 1D tensors :math:`(U, V)`,
              representing the source and destination nodes of all edges. For each :math:`i`,
              :math:`(U[i], V[i])` forms an edge.
            - ``'all'``: The returned result is a 3-tuple of 1D tensors :math:`(U, V, EID)`,
              representing the source nodes, destination nodes and IDs of all edges.
              For each :math:`i`, :math:`(U[i], V[i])` forms an edge with ID :math:`EID[i]`.
        order : str, optional
            The order of the returned edges, which can be one of the following:

            - ``'srcdst'``: The edges are sorted first by their source node IDs and then
              by their destination node IDs to break ties.
            - ``'eid'`` (default): The edges are sorted by their IDs.
Minjie Wang's avatar
Minjie Wang committed
3281
        etype : str or tuple of str, optional
3282
3283
3284
3285
            The edge type for query, which can be an edge type (str) or a canonical edge type
            (3-tuple of str). When an edge type appears in multiple canonical edge types, one
            must use a canonical edge type. If the graph has multiple edge types, one must
            specify the argument. Otherwise, it can be omitted.
Da Zheng's avatar
Da Zheng committed
3286
3287
3288

        Returns
        -------
3289
3290
3291
        Tensor or (Tensor, Tensor) or (Tensor, Tensor, Tensor)
            All edges of the specified edge type. For a description of the returned result,
            see the description of :attr:`form`.
Da Zheng's avatar
Da Zheng committed
3292
3293
3294
3295
3296

        Examples
        --------
        The following example uses PyTorch backend.

3297
3298
3299
3300
3301
3302
3303
3304
3305
3306
3307
3308
3309
3310
        >>> import dgl
        >>> import torch

        Create a homogeneous graph.

        >>> g = dgl.graph((torch.tensor([0, 0, 1, 1]), torch.tensor([1, 0, 2, 3])))

        Query for edges.

        >>> g.all_edges()
        (tensor([0, 0, 1, 1]), tensor([1, 0, 2, 3]))

        Specify a different value for :attr:`form` and :attr:`order`.

Mufei Li's avatar
Mufei Li committed
3311
        >>> g.all_edges(form='all', order='srcdst')
3312
3313
3314
3315
3316
3317
3318
3319
3320
3321
3322
3323
3324
3325
3326
3327
        (tensor([0, 0, 1, 1]), tensor([0, 1, 2, 3]), tensor([1, 0, 2, 3]))

        For a graph of multiple edge types, it is required to specify the edge type in query.

        >>> hg = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'plays', 'game'): (torch.tensor([3, 4]), torch.tensor([5, 6]))
        ... })
        >>> hg.all_edges(etype='plays')
        (tensor([3, 4]), tensor([5, 6]))

        See Also
        --------
        edges
        in_edges
        out_edges
Da Zheng's avatar
Da Zheng committed
3328
        """
Minjie Wang's avatar
Minjie Wang committed
3329
        src, dst, eid = self._graph.edges(self.get_etype_id(etype), order)
3330
        if form == 'all':
3331
            return src, dst, eid
3332
        elif form == 'uv':
3333
            return src, dst
3334
        elif form == 'eid':
3335
            return eid
3336
        else:
3337
            raise DGLError('Invalid form: {}. Must be "all", "uv" or "eid".'.format(form))
3338

Minjie Wang's avatar
Minjie Wang committed
3339
    def in_degree(self, v, etype=None):
Mufei Li's avatar
Mufei Li committed
3340
        """Return the in-degree of node ``v`` with edges of type ``etype``.
Da Zheng's avatar
Da Zheng committed
3341

3342
        **DEPRECATED**: Please use in_degrees
Da Zheng's avatar
Da Zheng committed
3343
        """
3344
3345
        dgl_warning("DGLGraph.in_degree is deprecated. Please use DGLGraph.in_degrees")
        return self.in_degrees(v, etype)
Da Zheng's avatar
Da Zheng committed
3346

Minjie Wang's avatar
Minjie Wang committed
3347
    def in_degrees(self, v=ALL, etype=None):
3348
3349
3350
        """Return the in-degree(s) of the given nodes.

        It computes the in-degree(s) w.r.t. to the edges of the given edge type.
Da Zheng's avatar
Da Zheng committed
3351
3352
3353

        Parameters
        ----------
3354
3355
        v : node IDs
            The node IDs. The allowed formats are:
3356

3357
3358
3359
3360
            * ``int``: A single node.
            * Int Tensor: Each element is a node ID. The tensor must have the same device type
              and ID data type as the graph's.
            * iterable[int]: Each element is a node ID.
3361

3362
3363
3364
3365
3366
3367
3368
3369
3370
            If not given, return the in-degrees of all the nodes.
        etype : str or (str, str, str), optional
            The type name of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            Can be omitted if the graph has only one type of edges.
Da Zheng's avatar
Da Zheng committed
3371
3372
3373

        Returns
        -------
3374
3375
3376
        int or Tensor
            The in-degree(s) of the node(s) in a Tensor. The i-th element is the in-degree
            of the i-th input node. If :attr:`v` is an ``int``, return an ``int`` too.
Da Zheng's avatar
Da Zheng committed
3377
3378
3379
3380
3381

        Examples
        --------
        The following example uses PyTorch backend.

3382
3383
        >>> import dgl
        >>> import torch
Mufei Li's avatar
Mufei Li committed
3384

3385
        Create a homogeneous graph.
Mufei Li's avatar
Mufei Li committed
3386

3387
        >>> g = dgl.graph((torch.tensor([0, 0, 1, 1]), torch.tensor([1, 1, 2, 3])))
Mufei Li's avatar
Mufei Li committed
3388

3389
3390
3391
3392
3393
3394
3395
3396
3397
3398
3399
3400
3401
3402
3403
3404
3405
3406
3407
3408
3409
3410
        Query for all nodes.

        >>> g.in_degrees()
        tensor([0, 2, 1, 1])

        Query for nodes 1 and 2.

        >>> g.in_degrees(torch.tensor([1, 2]))
        tensor([2, 1])

        For a graph of multiple edge types, it is required to specify the edge type in query.

        >>> hg = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'plays', 'game'): (torch.tensor([3, 4]), torch.tensor([5, 6]))
        ... })
        >>> hg.in_degrees(torch.tensor([1, 0]), etype='follows')
        tensor([1, 0])

        See Also
        --------
        out_degrees
Da Zheng's avatar
Da Zheng committed
3411
        """
3412
        dsttype = self.to_canonical_etype(etype)[2]
Minjie Wang's avatar
Minjie Wang committed
3413
        etid = self.get_etype_id(etype)
3414
        if is_all(v):
3415
            v = self.dstnodes(dsttype)
3416
3417
        v_tensor = utils.prepare_tensor(self, v, 'v')
        deg = self._graph.in_degrees(etid, v_tensor)
3418
3419
        if isinstance(v, numbers.Integral):
            return F.as_scalar(deg)
3420
        else:
3421
            return deg
Da Zheng's avatar
Da Zheng committed
3422

Mufei Li's avatar
Mufei Li committed
3423
3424
    def out_degree(self, u, etype=None):
        """Return the out-degree of node `u` with edges of type ``etype``.
Da Zheng's avatar
Da Zheng committed
3425

3426
        DEPRECATED: please use DGL.out_degrees
3427
        """
3428
3429
        dgl_warning("DGLGraph.out_degree is deprecated. Please use DGLGraph.out_degrees")
        return self.out_degrees(u, etype)
3430

Mufei Li's avatar
Mufei Li committed
3431
    def out_degrees(self, u=ALL, etype=None):
3432
3433
3434
        """Return the out-degree(s) of the given nodes.

        It computes the out-degree(s) w.r.t. to the edges of the given edge type.
3435
3436
3437

        Parameters
        ----------
3438
3439
        u : node IDs
            The node IDs. The allowed formats are:
3440

3441
3442
3443
3444
            * ``int``: A single node.
            * Int Tensor: Each element is a node ID. The tensor must have the same device type
              and ID data type as the graph's.
            * iterable[int]: Each element is a node ID.
3445

3446
3447
3448
3449
3450
3451
3452
3453
3454
            If not given, return the in-degrees of all the nodes.
        etype : str or (str, str, str), optional
            The type names of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            Can be omitted if the graph has only one type of edges.
3455
3456
3457

        Returns
        -------
3458
3459
3460
        int or Tensor
            The out-degree(s) of the node(s) in a Tensor. The i-th element is the out-degree
            of the i-th input node. If :attr:`v` is an ``int``, return an ``int`` too.
3461
3462
3463
3464
3465

        Examples
        --------
        The following example uses PyTorch backend.

3466
3467
        >>> import dgl
        >>> import torch
Mufei Li's avatar
Mufei Li committed
3468

3469
        Create a homogeneous graph.
Mufei Li's avatar
Mufei Li committed
3470

3471
        >>> g = dgl.graph((torch.tensor([0, 0, 1, 1]), torch.tensor([1, 1, 2, 3])))
Mufei Li's avatar
Mufei Li committed
3472

3473
3474
3475
3476
3477
3478
3479
3480
3481
3482
3483
3484
3485
3486
3487
3488
3489
3490
        Query for all nodes.

        >>> g.out_degrees()
        tensor([2, 2, 0, 0])

        Query for nodes 1 and 2.

        >>> g.out_degrees(torch.tensor([1, 2]))
        tensor([2, 0])

        For a graph of multiple edge types, it is required to specify the edge type in query.

        >>> hg = dgl.heterograph({
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2])),
        ...     ('user', 'plays', 'game'): (torch.tensor([3, 4]), torch.tensor([5, 6]))
        ... })
        >>> hg.out_degrees(torch.tensor([1, 0]), etype='follows')
        tensor([1, 1])
3491
3492
3493

        See Also
        --------
3494
        in_degrees
3495
        """
3496
        srctype = self.to_canonical_etype(etype)[0]
Minjie Wang's avatar
Minjie Wang committed
3497
        etid = self.get_etype_id(etype)
Mufei Li's avatar
Mufei Li committed
3498
        if is_all(u):
3499
            u = self.srcnodes(srctype)
3500
3501
3502
        u_tensor = utils.prepare_tensor(self, u, 'u')
        if F.as_scalar(F.sum(self.has_nodes(u_tensor, ntype=srctype), dim=0)) != len(u_tensor):
            raise DGLError('u contains invalid node IDs')
3503
3504
3505
        deg = self._graph.out_degrees(etid, utils.prepare_tensor(self, u, 'u'))
        if isinstance(u, numbers.Integral):
            return F.as_scalar(deg)
3506
        else:
3507
            return deg
Minjie Wang's avatar
Minjie Wang committed
3508

3509
    def adjacency_matrix(self, transpose=True, ctx=F.cpu(), scipy_fmt=None, etype=None):
3510
        """Alias of :meth:`adj`"""
3511
3512
        return self.adj(transpose, ctx, scipy_fmt, etype)

3513
    def adj(self, transpose=True, ctx=F.cpu(), scipy_fmt=None, etype=None):
Minjie Wang's avatar
Minjie Wang committed
3514
        """Return the adjacency matrix of edges of the given edge type.
3515

Minjie Wang's avatar
Minjie Wang committed
3516
3517
        By default, a row of returned adjacency matrix represents the
        destination of an edge and the column represents the source.
3518

Minjie Wang's avatar
Minjie Wang committed
3519
3520
        When transpose is True, a row represents the source and a column
        represents a destination.
3521
3522
3523

        Parameters
        ----------
Mufei Li's avatar
Mufei Li committed
3524
        transpose : bool, optional
3525
            A flag to transpose the returned adjacency matrix. (Default: True)
Mufei Li's avatar
Mufei Li committed
3526
3527
3528
        ctx : context, optional
            The context of returned adjacency matrix. (Default: cpu)
        scipy_fmt : str, optional
Minjie Wang's avatar
Minjie Wang committed
3529
            If specified, return a scipy sparse matrix in the given format.
Mufei Li's avatar
Mufei Li committed
3530
            Otherwise, return a backend dependent sparse tensor. (Default: None)
3531
3532
3533
3534
3535
3536
3537
3538
3539
        etype : str or (str, str, str), optional
            The type names of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            Can be omitted if the graph has only one type of edges.

3540

Minjie Wang's avatar
Minjie Wang committed
3541
3542
3543
3544
        Returns
        -------
        SparseTensor or scipy.sparse.spmatrix
            Adjacency matrix.
Mufei Li's avatar
Mufei Li committed
3545
3546
3547
3548

        Examples
        --------

3549
3550
3551
3552
3553
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

Mufei Li's avatar
Mufei Li committed
3554
3555
        Instantiate a heterogeneous graph.

3556
3557
3558
3559
        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): ([0, 1], [0, 1]),
        ...     ('developer', 'develops', 'game'): ([0, 1], [0, 2])
        ... })
Mufei Li's avatar
Mufei Li committed
3560
3561
3562

        Get a backend dependent sparse tensor. Here we use PyTorch for example.

3563
        >>> g.adj(etype='develops')
3564
3565
        tensor(indices=tensor([[0, 1],
                               [0, 2]]),
Mufei Li's avatar
Mufei Li committed
3566
               values=tensor([1., 1.]),
3567
               size=(2, 3), nnz=2, layout=torch.sparse_coo)
Mufei Li's avatar
Mufei Li committed
3568
3569
3570

        Get a scipy coo sparse matrix.

3571
        >>> g.adj(scipy_fmt='coo', etype='develops')
3572
3573
        <2x3 sparse matrix of type '<class 'numpy.int64'>'
	    with 2 stored elements in COOrdinate format>
3574
        """
Minjie Wang's avatar
Minjie Wang committed
3575
3576
3577
3578
3579
        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)
3580

3581

3582
    def adjacency_matrix_scipy(self, transpose=True, fmt='csr', return_edge_ids=None):
3583
3584
3585
3586
3587
3588
3589
3590
        """DEPRECATED: please use ``dgl.adjacency_matrix(transpose, scipy_fmt=fmt)``.
        """
        dgl_warning('DGLGraph.adjacency_matrix_scipy is deprecated. '
                    'Please replace it with:\n\n\t'
                    'DGLGraph.adjacency_matrix(transpose, scipy_fmt="{}").\n'.format(fmt))

        return self.adjacency_matrix(transpose=transpose, scipy_fmt=fmt)

3591
    def inc(self, typestr, ctx=F.cpu(), etype=None):
Minjie Wang's avatar
Minjie Wang committed
3592
3593
        """Return the incidence matrix representation of edges with the given
        edge type.
3594

Mufei Li's avatar
Mufei Li committed
3595
        An incidence matrix is an n-by-m sparse matrix, where n is
Minjie Wang's avatar
Minjie Wang committed
3596
3597
3598
        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.
3599

Mufei Li's avatar
Mufei Li committed
3600
        There are three types of incidence matrices :math:`I`:
Da Zheng's avatar
Da Zheng committed
3601

Minjie Wang's avatar
Minjie Wang committed
3602
        * ``in``:
Da Zheng's avatar
Da Zheng committed
3603

Minjie Wang's avatar
Minjie Wang committed
3604
3605
3606
            - :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.
Da Zheng's avatar
Da Zheng committed
3607

Minjie Wang's avatar
Minjie Wang committed
3608
        * ``out``:
Da Zheng's avatar
Da Zheng committed
3609

Minjie Wang's avatar
Minjie Wang committed
3610
3611
3612
3613
3614
3615
3616
3617
3618
            - :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).
Da Zheng's avatar
Da Zheng committed
3619
3620
3621

        Parameters
        ----------
Minjie Wang's avatar
Minjie Wang committed
3622
3623
        typestr : str
            Can be either ``in``, ``out`` or ``both``
Mufei Li's avatar
Mufei Li committed
3624
3625
        ctx : context, optional
            The context of returned incidence matrix. (Default: cpu)
3626
3627
3628
3629
3630
3631
3632
3633
        etype : str or (str, str, str), optional
            The type names of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            Can be omitted if the graph has only one type of edges.
Minjie Wang's avatar
Minjie Wang committed
3634
3635
3636

        Returns
        -------
Mufei Li's avatar
Mufei Li committed
3637
        Framework SparseTensor
Minjie Wang's avatar
Minjie Wang committed
3638
            The incidence matrix.
Mufei Li's avatar
Mufei Li committed
3639
3640
3641
3642

        Examples
        --------

3643
3644
3645
3646
3647
3648
        The following example uses PyTorch backend.

        >>> import dgl

        >>> g = dgl.graph(([0, 1], [0, 2]))
        >>> g.inc('in')
Mufei Li's avatar
Mufei Li committed
3649
3650
3651
3652
        tensor(indices=tensor([[0, 2],
                               [0, 1]]),
               values=tensor([1., 1.]),
               size=(3, 2), nnz=2, layout=torch.sparse_coo)
3653
        >>> g.inc('out')
Mufei Li's avatar
Mufei Li committed
3654
3655
3656
3657
        tensor(indices=tensor([[0, 1],
                               [0, 1]]),
               values=tensor([1., 1.]),
               size=(3, 2), nnz=2, layout=torch.sparse_coo)
3658
        >>> g.inc('both')
Mufei Li's avatar
Mufei Li committed
3659
3660
3661
3662
        tensor(indices=tensor([[1, 2],
                               [1, 1]]),
               values=tensor([-1.,  1.]),
               size=(3, 2), nnz=2, layout=torch.sparse_coo)
Da Zheng's avatar
Da Zheng committed
3663
        """
Minjie Wang's avatar
Minjie Wang committed
3664
3665
3666
        etid = self.get_etype_id(etype)
        return self._graph.incidence_matrix(etid, typestr, ctx)[0]

3667
3668
    incidence_matrix = inc

Minjie Wang's avatar
Minjie Wang committed
3669
3670
3671
3672
3673
    #################################################################
    # Features
    #################################################################

    def node_attr_schemes(self, ntype=None):
Mufei Li's avatar
Mufei Li committed
3674
        """Return the node feature schemes for the specified type.
Da Zheng's avatar
Da Zheng committed
3675

3676
        The scheme of a feature describes the shape and data type of it.
Da Zheng's avatar
Da Zheng committed
3677
3678
3679

        Parameters
        ----------
Minjie Wang's avatar
Minjie Wang committed
3680
        ntype : str, optional
3681
3682
            The node type name. Can be omitted if there is only one type of nodes
            in the graph.
Da Zheng's avatar
Da Zheng committed
3683
3684
3685

        Returns
        -------
3686
3687
        dict[str, Scheme]
            A dictionary mapping a feature name to its associated feature scheme.
3688
3689
3690

        Examples
        --------
3691
3692
3693
3694
3695
3696
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Query for a homogeneous graph.
3697

3698
3699
3700
3701
3702
3703
3704
3705
3706
3707
3708
3709
3710
        >>> g = dgl.graph((torch.tensor([0, 1]), torch.tensor([1, 2])))
        >>> g.ndata['h1'] = torch.randn(3, 1)
        >>> g.ndata['h2'] = torch.randn(3, 2)
        >>> g.node_attr_schemes()
        {'h1': Scheme(shape=(1,), dtype=torch.float32),
         'h2': Scheme(shape=(2,), dtype=torch.float32)}

        Query for a heterogeneous graph of multiple node types.

        >>> g = dgl.heterograph({('user', 'plays', 'game'):
        ...                      (torch.tensor([1, 2]), torch.tensor([3, 4]))})
        >>> g.nodes['user'].data['h1'] = torch.randn(3, 1)
        >>> g.nodes['user'].data['h2'] = torch.randn(3, 2)
3711
        >>> g.node_attr_schemes('user')
3712
3713
        {'h1': Scheme(shape=(1,), dtype=torch.float32),
         'h2': Scheme(shape=(2,), dtype=torch.float32)}
Mufei Li's avatar
Mufei Li committed
3714
3715
3716
3717

        See Also
        --------
        edge_attr_schemes
Da Zheng's avatar
Da Zheng committed
3718
        """
Minjie Wang's avatar
Minjie Wang committed
3719
        return self._node_frames[self.get_ntype_id(ntype)].schemes
Da Zheng's avatar
Da Zheng committed
3720

Minjie Wang's avatar
Minjie Wang committed
3721
    def edge_attr_schemes(self, etype=None):
Mufei Li's avatar
Mufei Li committed
3722
        """Return the edge feature schemes for the specified type.
Da Zheng's avatar
Da Zheng committed
3723

3724
        The scheme of a feature describes the shape and data type of it.
Da Zheng's avatar
Da Zheng committed
3725
3726
3727

        Parameters
        ----------
3728
3729
3730
3731
3732
3733
3734
3735
3736
        etype : str or (str, str, str), optional
            The type names of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            Can be omitted if the graph has only one type of edges.

Da Zheng's avatar
Da Zheng committed
3737
3738
3739

        Returns
        -------
3740
3741
        dict[str, Scheme]
            A dictionary mapping a feature name to its associated feature scheme.
Da Zheng's avatar
Da Zheng committed
3742

3743
3744
        Examples
        --------
3745
3746
3747
3748
3749
3750
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Query for a homogeneous graph.
Da Zheng's avatar
Da Zheng committed
3751

3752
3753
3754
3755
3756
3757
3758
3759
3760
3761
3762
3763
3764
3765
3766
3767
3768
3769
        >>> g = dgl.graph((torch.tensor([0, 1]), torch.tensor([1, 2])))
        >>> g.edata['h1'] = torch.randn(2, 1)
        >>> g.edata['h2'] = torch.randn(2, 2)
        >>> g.edge_attr_schemes()
        {'h1': Scheme(shape=(1,), dtype=torch.float32),
         'h2': Scheme(shape=(2,), dtype=torch.float32)}

        Query for a heterogeneous graph of multiple edge types.

        >>> g = dgl.heterograph({('user', 'plays', 'game'):
        ...                      (torch.tensor([1, 2]), torch.tensor([3, 4])),
        ...                      ('user', 'follows', 'user'):
        ...                      (torch.tensor([3, 4]), torch.tensor([5, 6]))})
        >>> g.edges['plays'].data['h1'] = torch.randn(2, 1)
        >>> g.edges['plays'].data['h2'] = torch.randn(2, 2)
        >>> g.edge_attr_schemes('plays')
        {'h1': Scheme(shape=(1,), dtype=torch.float32),
         'h2': Scheme(shape=(2,), dtype=torch.float32)}
Mufei Li's avatar
Mufei Li committed
3770
3771
3772
3773

        See Also
        --------
        node_attr_schemes
Da Zheng's avatar
Da Zheng committed
3774
        """
Minjie Wang's avatar
Minjie Wang committed
3775
        return self._edge_frames[self.get_etype_id(etype)].schemes
Da Zheng's avatar
Da Zheng committed
3776

Minjie Wang's avatar
Minjie Wang committed
3777
    def set_n_initializer(self, initializer, field=None, ntype=None):
3778
        """Set the initializer for node features.
Da Zheng's avatar
Da Zheng committed
3779

3780
3781
3782
        When only part of the nodes have a feature (e.g. new nodes are added,
        features are set for a subset of nodes), the initializer initializes
        features for the rest nodes.
Minjie Wang's avatar
Minjie Wang committed
3783
3784
3785
3786

        Parameters
        ----------
        initializer : callable
3787
3788
3789
3790
3791
3792
3793
3794
3795
3796
3797
            A function of signature ``func(shape, dtype, ctx, id_range) -> Tensor``.
            The tensor will be the initialized features. The arguments are:

            - ``shape``: The shape of the tensor to return, which is a tuple of int.
              The first dimension is the number of nodes for feature initialization.
            - ``dtype``: The data type of the tensor to return, which is a
              framework-specific data type object.
            - ``ctx``: The device of the tensor to return, which is a framework-specific
              device object.
            - ``id_range``: The start and end ID of the nodes for feature initialization,
              which is a slice.
Minjie Wang's avatar
Minjie Wang committed
3798
        field : str, optional
3799
3800
            The name of the feature that the initializer applies. If not given, the
            initializer applies to all features.
Minjie Wang's avatar
Minjie Wang committed
3801
        ntype : str, optional
3802
            The type name of the nodes. Can be omitted if the graph has only one type of nodes.
Da Zheng's avatar
Da Zheng committed
3803

3804
        Notes
Minjie Wang's avatar
Minjie Wang committed
3805
        -----
3806
3807
        Without setting a node feature initializer, zero tensors are generated
        for nodes without a feature.
Da Zheng's avatar
Da Zheng committed
3808

3809
        Examples
Mufei Li's avatar
Mufei Li committed
3810
        --------
3811
3812
3813
3814
3815
3816
3817
3818
3819
3820
3821
3822
3823
3824
3825
3826
3827
3828
3829
3830
3831
3832
3833
3834
3835
3836
3837
3838
3839
3840
3841
3842
3843
3844
3845
3846
3847
3848
3849
3850
3851
3852
3853
3854
3855
3856
3857
3858
3859
3860

        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Define a function for initializer.

        >>> def init_feats(shape, dtype, device, id_range):
        ...     return torch.ones(shape, dtype=dtype, device=device)

        An example for a homogeneous graph.

        >>> g = dgl.graph((torch.tensor([0]), torch.tensor([1])))
        >>> g.ndata['h1'] = torch.zeros(2, 2)
        >>> g.ndata['h2'] = torch.ones(2, 1)
        >>> # Apply the initializer to feature 'h2' only.
        >>> g.set_n_initializer(init_feats, field='h2')
        >>> g.add_nodes(1)
        >>> print(g.ndata['h1'])
        tensor([[0., 0.],
                [0., 0.],
                [0., 0.]])
        >>> print(g.ndata['h2'])
        tensor([[1.], [1.], [1.]])

        An example for a heterogeneous graph of multiple node types.

        >>> g = dgl.heterograph({
        ...     ('user', 'plays', 'game'): (torch.tensor([0, 1, 1, 2]),
        ...                                 torch.tensor([0, 0, 1, 1])),
        ...     ('developer', 'develops', 'game'): (torch.tensor([0, 1]),
        ...                                         torch.tensor([0, 1]))
        ...     })
        >>> g.nodes['user'].data['h'] = torch.zeros(3, 2)
        >>> g.nodes['game'].data['w'] = torch.ones(2, 2)
        >>> g.set_n_initializer(init_feats, ntype='game')
        >>> g.add_nodes(1, ntype='user')
        >>> # Initializer not set for 'user', use zero tensors by default
        >>> g.nodes['user'].data['h']
        tensor([[0., 0.],
                [0., 0.],
                [0., 0.],
                [0., 0.]])
        >>> # Initializer set for 'game'
        >>> g.add_nodes(1, ntype='game')
        >>> g.nodes['game'].data['w']
        tensor([[1., 1.],
                [1., 1.],
                [1., 1.]])
Da Zheng's avatar
Da Zheng committed
3861
        """
Minjie Wang's avatar
Minjie Wang committed
3862
3863
        ntid = self.get_ntype_id(ntype)
        self._node_frames[ntid].set_initializer(initializer, field)
Da Zheng's avatar
Da Zheng committed
3864

Minjie Wang's avatar
Minjie Wang committed
3865
    def set_e_initializer(self, initializer, field=None, etype=None):
3866
        """Set the initializer for edge features.
Da Zheng's avatar
Da Zheng committed
3867

3868
3869
3870
        When only part of the edges have a feature (e.g. new edges are added,
        features are set for a subset of edges), the initializer initializes
        features for the rest edges.
Minjie Wang's avatar
Minjie Wang committed
3871
3872
3873
3874

        Parameters
        ----------
        initializer : callable
3875
3876
3877
3878
3879
3880
3881
3882
3883
3884
3885
            A function of signature ``func(shape, dtype, ctx, id_range) -> Tensor``.
            The tensor will be the initialized features. The arguments are:

            - ``shape``: The shape of the tensor to return, which is a tuple of int.
              The first dimension is the number of edges for feature initialization.
            - ``dtype``: The data type of the tensor to return, which is a
              framework-specific data type object.
            - ``ctx``: The device of the tensor to return, which is a framework-specific
              device object.
            - ``id_range``: The start and end ID of the edges for feature initialization,
              which is a slice.
Minjie Wang's avatar
Minjie Wang committed
3886
        field : str, optional
3887
3888
            The name of the feature that the initializer applies. If not given, the
            initializer applies to all features.
3889
3890
3891
3892
3893
3894
3895
3896
3897
        etype : str or (str, str, str), optional
            The type names of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            Can be omitted if the graph has only one type of edges.

Minjie Wang's avatar
Minjie Wang committed
3898

3899
        Notes
Minjie Wang's avatar
Minjie Wang committed
3900
        -----
3901
3902
        Without setting an edge feature initializer, zero tensors are generated
        for edges without a feature.
Mufei Li's avatar
Mufei Li committed
3903

3904
        Examples
Mufei Li's avatar
Mufei Li committed
3905
        --------
3906
3907
3908
3909
3910
3911
3912
3913
3914
3915
3916
3917
3918
3919
3920
3921
3922
3923
3924
3925
3926
3927
3928
3929
3930
3931
3932
3933
3934
3935
3936
3937
3938
3939
3940
3941
3942
3943
3944
3945
3946
3947
3948
3949
3950
3951
3952
3953

        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Define a function for initializer.

        >>> def init_feats(shape, dtype, device, id_range):
        ...     return torch.ones(shape, dtype=dtype, device=device)

        An example for a homogeneous graph.

        >>> g = dgl.graph((torch.tensor([0]), torch.tensor([1])))
        >>> g.edata['h1'] = torch.zeros(1, 2)
        >>> g.edata['h2'] = torch.ones(1, 1)
        >>> # Apply the initializer to feature 'h2' only.
        >>> g.set_e_initializer(init_feats, field='h2')
        >>> g.add_edges(torch.tensor([1]), torch.tensor([1]))
        >>> print(g.edata['h1'])
        tensor([[0., 0.],
                [0., 0.]])
        >>> print(g.edata['h2'])
        tensor([[1.], [1.]])

        An example for a heterogeneous graph of multiple edge types.

        >>> g = dgl.heterograph({
        ...     ('user', 'plays', 'game'): (torch.tensor([0, 1]),
        ...                                 torch.tensor([0, 0])),
        ...     ('developer', 'develops', 'game'): (torch.tensor([0, 1]),
        ...                                         torch.tensor([0, 1]))
        ...     })
        >>> g.edges['plays'].data['h'] = torch.zeros(2, 2)
        >>> g.edges['develops'].data['w'] = torch.ones(2, 2)
        >>> g.set_e_initializer(init_feats, etype='plays')
        >>> # Initializer not set for 'develops', use zero tensors by default
        >>> g.add_edges(torch.tensor([1]), torch.tensor([1]), etype='develops')
        >>> g.edges['develops'].data['w']
        tensor([[1., 1.],
                [1., 1.],
                [0., 0.]])
        >>> # Initializer set for 'plays'
        >>> g.add_edges(torch.tensor([1]), torch.tensor([1]), etype='plays')
        >>> g.edges['plays'].data['h']
        tensor([[0., 0.],
                [0., 0.],
                [1., 1.]])
Da Zheng's avatar
Da Zheng committed
3954
        """
Minjie Wang's avatar
Minjie Wang committed
3955
3956
        etid = self.get_etype_id(etype)
        self._edge_frames[etid].set_initializer(initializer, field)
Da Zheng's avatar
Da Zheng committed
3957

3958
    def _set_n_repr(self, ntid, u, data):
Minjie Wang's avatar
Minjie Wang committed
3959
        """Internal API to set node features.
Da Zheng's avatar
Da Zheng committed
3960
3961
3962
3963
3964
3965

        `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).

3966
        All updates will be done out of place to work with autograd.
Da Zheng's avatar
Da Zheng committed
3967
3968
3969

        Parameters
        ----------
Minjie Wang's avatar
Minjie Wang committed
3970
3971
        ntid : int
            Node type id.
Da Zheng's avatar
Da Zheng committed
3972
3973
        u : node, container or tensor
            The node(s).
Minjie Wang's avatar
Minjie Wang committed
3974
3975
        data : dict of tensor
            Node representation.
Da Zheng's avatar
Da Zheng committed
3976
        """
3977
        if is_all(u):
Minjie Wang's avatar
Minjie Wang committed
3978
            num_nodes = self._graph.number_of_nodes(ntid)
3979
        else:
3980
            u = utils.prepare_tensor(self, u, 'u')
3981
3982
3983
3984
3985
3986
            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))
3987
3988
3989
3990
            if F.context(val) != self.device:
                raise DGLError('Cannot assign node feature "{}" on device {} to a graph on'
                               ' device {}. Call DGLGraph.to() to copy the graph to the'
                               ' same device.'.format(key, F.context(val), self.device))
3991
3992

        if is_all(u):
3993
            self._node_frames[ntid].update(data)
3994
        else:
3995
            self._node_frames[ntid].update_row(u, data)
Da Zheng's avatar
Da Zheng committed
3996

Minjie Wang's avatar
Minjie Wang committed
3997
    def _get_n_repr(self, ntid, u):
Da Zheng's avatar
Da Zheng committed
3998
3999
4000
4001
4002
4003
        """Get node(s) representation of a single node type.

        The returned feature tensor batches multiple node features on the first dimension.

        Parameters
        ----------
Minjie Wang's avatar
Minjie Wang committed
4004
4005
        ntid : int
            Node type id.
Da Zheng's avatar
Da Zheng committed
4006
4007
4008
4009
4010
4011
4012
4013
        u : node, container or tensor
            The node(s).

        Returns
        -------
        dict
            Representation dict from feature name to feature tensor.
        """
4014
        if is_all(u):
4015
            return self._node_frames[ntid]
4016
        else:
4017
            u = utils.prepare_tensor(self, u, 'u')
4018
            return self._node_frames[ntid].subframe(u)
Da Zheng's avatar
Da Zheng committed
4019

Minjie Wang's avatar
Minjie Wang committed
4020
4021
    def _pop_n_repr(self, ntid, key):
        """Internal API to get and remove the specified node feature.
Da Zheng's avatar
Da Zheng committed
4022
4023
4024

        Parameters
        ----------
Minjie Wang's avatar
Minjie Wang committed
4025
4026
        ntid : int
            Node type id.
Da Zheng's avatar
Da Zheng committed
4027
4028
4029
4030
4031
4032
4033
4034
        key : str
            The attribute name.

        Returns
        -------
        Tensor
            The popped representation
        """
Minjie Wang's avatar
Minjie Wang committed
4035
        return self._node_frames[ntid].pop(key)
Da Zheng's avatar
Da Zheng committed
4036

4037
    def _set_e_repr(self, etid, edges, data):
Minjie Wang's avatar
Minjie Wang committed
4038
        """Internal API to set edge(s) features.
Da Zheng's avatar
Da Zheng committed
4039
4040
4041
4042
4043

        `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.

4044
        All update will be done out of place to work with autograd.
Da Zheng's avatar
Da Zheng committed
4045
4046
4047

        Parameters
        ----------
Minjie Wang's avatar
Minjie Wang committed
4048
4049
        etid : int
            Edge type id.
Da Zheng's avatar
Da Zheng committed
4050
4051
4052
4053
4054
4055
4056
4057
        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.
Minjie Wang's avatar
Minjie Wang committed
4058
4059
        data : tensor or dict of tensor
            Edge representation.
Da Zheng's avatar
Da Zheng committed
4060
        """
4061
        # parse argument
4062
4063
        if not is_all(edges):
            eid = utils.parse_edges_arg_to_eid(self, edges, etid, 'edges')
4064
4065
4066
4067
4068
4069

        # sanity check
        if not utils.is_dict_like(data):
            raise DGLError('Expect dictionary type for feature data.'
                           ' Got "%s" instead.' % type(data))

4070
        if is_all(edges):
Minjie Wang's avatar
Minjie Wang committed
4071
            num_edges = self._graph.number_of_edges(etid)
4072
4073
4074
4075
4076
4077
4078
        else:
            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))
4079
4080
4081
4082
4083
            if F.context(val) != self.device:
                raise DGLError('Cannot assign edge feature "{}" on device {} to a graph on'
                               ' device {}. Call DGLGraph.to() to copy the graph to the'
                               ' same device.'.format(key, F.context(val), self.device))

4084
        # set
4085
4086
        if is_all(edges):
            self._edge_frames[etid].update(data)
4087
        else:
4088
            self._edge_frames[etid].update_row(eid, data)
Da Zheng's avatar
Da Zheng committed
4089

Minjie Wang's avatar
Minjie Wang committed
4090
4091
    def _get_e_repr(self, etid, edges):
        """Internal API to get edge features.
Da Zheng's avatar
Da Zheng committed
4092
4093
4094

        Parameters
        ----------
Minjie Wang's avatar
Minjie Wang committed
4095
4096
        etid : int
            Edge type id.
Da Zheng's avatar
Da Zheng committed
4097
4098
4099
4100
4101
4102
4103
4104
4105
        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
        """
4106
4107
        # parse argument
        if is_all(edges):
Minjie Wang's avatar
Minjie Wang committed
4108
            return dict(self._edge_frames[etid])
4109
        else:
4110
4111
            eid = utils.parse_edges_arg_to_eid(self, edges, etid, 'edges')
            return self._edge_frames[etid].subframe(eid)
Da Zheng's avatar
Da Zheng committed
4112

Minjie Wang's avatar
Minjie Wang committed
4113
    def _pop_e_repr(self, etid, key):
Da Zheng's avatar
Da Zheng committed
4114
4115
4116
4117
        """Get and remove the specified edge repr of a single edge type.

        Parameters
        ----------
Minjie Wang's avatar
Minjie Wang committed
4118
4119
        etid : int
            Edge type id.
Da Zheng's avatar
Da Zheng committed
4120
4121
4122
4123
4124
4125
4126
4127
        key : str
          The attribute name.

        Returns
        -------
        Tensor
            The popped representation
        """
Minjie Wang's avatar
Minjie Wang committed
4128
        self._edge_frames[etid].pop(key)
Da Zheng's avatar
Da Zheng committed
4129

Minjie Wang's avatar
Minjie Wang committed
4130
4131
4132
4133
4134
    #################################################################
    # Message passing
    #################################################################

    def apply_nodes(self, func, v=ALL, ntype=None, inplace=False):
4135
        """Update the features of the specified nodes by the provided function.
Da Zheng's avatar
Da Zheng committed
4136
4137
4138

        Parameters
        ----------
4139
4140
4141
4142
4143
4144
4145
4146
4147
4148
4149
4150
        func : callable
            The function to update node features. It must be
            a :ref:`apiudf`.
        v : node IDs
            The node IDs. The allowed formats are:

            * ``int``: A single node.
            * Int Tensor: Each element is a node ID. The tensor must have the same device type
              and ID data type as the graph's.
            * iterable[int]: Each element is a node ID.

            If not given (default), use all the nodes in the graph.
Minjie Wang's avatar
Minjie Wang committed
4151
        ntype : str, optional
4152
4153
            The node type name. Can be omitted if there is
            only one type of nodes in the graph.
Minjie Wang's avatar
Minjie Wang committed
4154
        inplace : bool, optional
4155
            **DEPRECATED**.
Da Zheng's avatar
Da Zheng committed
4156
4157
4158

        Examples
        --------
4159
4160
4161
4162
4163
4164
4165
4166
4167
4168
4169
4170
4171
4172
4173
4174
4175
4176
4177
4178

        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        **Homogeneous graph**

        >>> g = dgl.graph(([0, 1, 2, 3], [1, 2, 3, 4]))
        >>> g.ndata['h'] = torch.ones(5, 2)
        >>> g.apply_nodes(lambda nodes: {'x' : nodes.data['h'] * 2})
        >>> g.ndata['x']
        tensor([[2., 2.],
                [2., 2.],
                [2., 2.],
                [2., 2.],
                [2., 2.]])

        **Heterogeneous graph**

4179
        >>> g = dgl.heterograph({('user', 'follows', 'user'): ([0, 1], [1, 2])})
Minjie Wang's avatar
Minjie Wang committed
4180
4181
4182
        >>> 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']
Da Zheng's avatar
Da Zheng committed
4183
4184
4185
        tensor([[2., 2., 2., 2., 2.],
                [2., 2., 2., 2., 2.],
                [2., 2., 2., 2., 2.]])
Mufei Li's avatar
Mufei Li committed
4186
4187
4188
4189

        See Also
        --------
        apply_edges
Da Zheng's avatar
Da Zheng committed
4190
        """
4191
4192
        if inplace:
            raise DGLError('The `inplace` option is removed in v0.5.')
Minjie Wang's avatar
Minjie Wang committed
4193
        ntid = self.get_ntype_id(ntype)
4194
        ntype = self.ntypes[ntid]
Minjie Wang's avatar
Minjie Wang committed
4195
        if is_all(v):
4196
            v_id = self.nodes(ntype)
Minjie Wang's avatar
Minjie Wang committed
4197
        else:
4198
4199
            v_id = utils.prepare_tensor(self, v, 'v')
        ndata = core.invoke_node_udf(self, v_id, ntype, func, orig_nid=v_id)
4200
        self._set_n_repr(ntid, v, ndata)
Minjie Wang's avatar
Minjie Wang committed
4201
4202

    def apply_edges(self, func, edges=ALL, etype=None, inplace=False):
4203
        """Update the features of the specified edges by the provided function.
Da Zheng's avatar
Da Zheng committed
4204
4205
4206

        Parameters
        ----------
4207
4208
4209
4210
4211
4212
4213
4214
4215
4216
4217
4218
4219
4220
4221
4222
4223
4224
4225
4226
4227
4228
4229
4230
4231
4232
        func : dgl.function.BuiltinFunction or callable
            The function to generate new edge features. It must be either
            a :ref:`api-built-in` or a :ref:`apiudf`.
        edges : edges
            The edges to update features on. The allowed input formats are:

            * ``int``: A single edge ID.
            * Int Tensor: Each element is an edge ID.  The tensor must have the same device type
              and ID data type as the graph's.
            * iterable[int]: Each element is an edge ID.
            * (Tensor, Tensor): The node-tensors format where the i-th elements
              of the two tensors specify an edge.
            * (iterable[int], iterable[int]): Similar to the node-tensors format but
              stores edge endpoints in python iterables.

            Default value specifies all the edges in the graph.

        etype : str or (str, str, str), optional
            The type name of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            Can be omitted if the graph has only one type of edges.

Da Zheng's avatar
Da Zheng committed
4233
        inplace: bool, optional
4234
4235
4236
4237
4238
4239
4240
            **DEPRECATED**.

        Notes
        -----
        DGL recommends using DGL's bulit-in function for the :attr:`func` argument,
        because DGL will invoke efficient kernels that avoids copying node features to
        edge features in this case.
Da Zheng's avatar
Da Zheng committed
4241
4242
4243

        Examples
        --------
4244
4245
4246
4247
4248
4249
4250
4251
4252
4253
4254
4255
4256
4257
4258
4259
4260
4261
4262
4263
4264
4265
4266
4267
4268
4269
4270
4271
4272

        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        **Homogeneous graph**

        >>> g = dgl.graph(([0, 1, 2, 3], [1, 2, 3, 4]))
        >>> g.ndata['h'] = torch.ones(5, 2)
        >>> g.apply_edges(lambda edges: {'x' : edges.src['h'] + edges.dst['h']})
        >>> g.edata['x']
        tensor([[2., 2.],
                [2., 2.],
                [2., 2.],
                [2., 2.]])

        Use built-in function

        >>> import dgl.function as fn
        >>> g.apply_edges(fn.u_add_v('h', 'h', 'x'))
        >>> g.edata['x']
        tensor([[2., 2.],
                [2., 2.],
                [2., 2.],
                [2., 2.]])

        **Heterogeneous graph**

4273
        >>> g = dgl.heterograph({('user', 'plays', 'game'): ([0, 1, 1, 2], [0, 0, 2, 1])})
Minjie Wang's avatar
Minjie Wang committed
4274
4275
4276
        >>> 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']
Da Zheng's avatar
Da Zheng committed
4277
        tensor([[2., 2., 2., 2., 2.],
4278
                [2., 2., 2., 2., 2.],
Da Zheng's avatar
Da Zheng committed
4279
4280
                [2., 2., 2., 2., 2.],
                [2., 2., 2., 2., 2.]])
Mufei Li's avatar
Mufei Li committed
4281
4282
4283
4284

        See Also
        --------
        apply_nodes
Da Zheng's avatar
Da Zheng committed
4285
        """
4286
4287
        if inplace:
            raise DGLError('The `inplace` option is removed in v0.5.')
Minjie Wang's avatar
Minjie Wang committed
4288
        etid = self.get_etype_id(etype)
4289
4290
        etype = self.canonical_etypes[etid]
        g = self if etype is None else self[etype]
Minjie Wang's avatar
Minjie Wang committed
4291
        if is_all(edges):
4292
            eid = ALL
Minjie Wang's avatar
Minjie Wang committed
4293
        else:
4294
4295
4296
            eid = utils.parse_edges_arg_to_eid(self, edges, etid, 'edges')
        if core.is_builtin(func):
            if not is_all(eid):
4297
                g = g.edge_subgraph(eid, relabel_nodes=False)
4298
4299
4300
4301
            edata = core.invoke_gsddmm(g, func)
        else:
            edata = core.invoke_edge_udf(g, eid, etype, func)
        self._set_e_repr(etid, eid, edata)
Minjie Wang's avatar
Minjie Wang committed
4302

4303
4304
4305
4306
4307
4308
4309
    def send_and_recv(self,
                      edges,
                      message_func,
                      reduce_func,
                      apply_node_func=None,
                      etype=None,
                      inplace=False):
4310
4311
        """Send messages along the specified edges and reduce them on
        the destination nodes to update their features.
4312

Da Zheng's avatar
Da Zheng committed
4313
4314
        Parameters
        ----------
4315
4316
4317
4318
4319
4320
4321
4322
4323
4324
4325
4326
4327
4328
4329
4330
4331
4332
        edges : edges
            The edges to send and receive messages on. The allowed input formats are:

            * ``int``: A single edge ID.
            * Int Tensor: Each element is an edge ID.  The tensor must have the same device type
              and ID data type as the graph's.
            * iterable[int]: Each element is an edge ID.
            * (Tensor, Tensor): The node-tensors format where the i-th elements
              of the two tensors specify an edge.
            * (iterable[int], iterable[int]): Similar to the node-tensors format but
              stores edge endpoints in python iterables.

        message_func : dgl.function.BuiltinFunction or callable
            The message function to generate messages along the edges.
            It must be either a :ref:`api-built-in` or a :ref:`apiudf`.
        reduce_func : dgl.function.BuiltinFunction or callable
            The reduce function to aggregate the messages.
            It must be either a :ref:`api-built-in` or a :ref:`apiudf`.
4333
        apply_node_func : callable, optional
4334
4335
4336
4337
4338
4339
4340
4341
4342
4343
4344
            An optional apply function to further update the node features
            after the message reduction. It must be a :ref:`apiudf`.
        etype : str or (str, str, str), optional
            The type name of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            Can be omitted if the graph has only one type of edges.

Da Zheng's avatar
Da Zheng committed
4345
        inplace: bool, optional
4346
4347
4348
4349
4350
4351
4352
4353
            **DEPRECATED**.

        Notes
        -----
        DGL recommends using DGL's bulit-in function for the :attr:`message_func`
        and the :attr:`reduce_func` arguments,
        because DGL will invoke efficient kernels that avoids copying node features to
        edge features in this case.
Mufei Li's avatar
Mufei Li committed
4354
4355
4356
4357
4358
4359
4360
4361

        Examples
        --------

        >>> import dgl
        >>> import dgl.function as fn
        >>> import torch

4362
4363
4364
4365
4366
4367
4368
4369
4370
4371
4372
4373
4374
4375
4376
4377
4378
4379
4380
4381
4382
4383
4384
        **Homogeneous graph**

        >>> g = dgl.graph(([0, 1, 2, 3], [1, 2, 3, 4]))
        >>> g.ndata['x'] = torch.ones(5, 2)
        >>> # Specify edges using (Tensor, Tensor).
        >>> g.send_and_recv(([1, 2], [2, 3]), fn.copy_u('x', 'm'), fn.sum('m', 'h'))
        >>> g.ndata['h']
        tensor([[0., 0.],
                [0., 0.],
                [1., 1.],
                [1., 1.],
                [0., 0.]])
        >>> # Specify edges using IDs.
        >>> g.send_and_recv([0, 2, 3], fn.copy_u('x', 'm'), fn.sum('m', 'h'))
        >>> g.ndata['h']
        tensor([[0., 0.],
                [1., 1.],
                [0., 0.],
                [1., 1.],
                [1., 1.]])

        **Heterogeneous graph**

4385
4386
4387
4388
        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): ([0, 1], [1, 2]),
        ...     ('user', 'plays', 'game'): ([0, 1, 1, 2], [0, 0, 1, 1])
        ... })
Mufei Li's avatar
Mufei Li committed
4389
4390
        >>> g.nodes['user'].data['h'] = torch.tensor([[0.], [1.], [2.]])
        >>> g.send_and_recv(g['follows'].edges(), fn.copy_src('h', 'm'),
4391
        ...                 fn.sum('m', 'h'), etype='follows')
Mufei Li's avatar
Mufei Li committed
4392
4393
4394
4395
        >>> g.nodes['user'].data['h']
        tensor([[0.],
                [0.],
                [1.]])
4396
4397
4398
4399
4400
4401
4402
4403
4404
4405
4406
4407
4408
4409
4410
4411
4412
4413
4414
4415
4416
4417
4418

        **``send_and_recv`` using user-defined functions**

        >>> import torch as th
        >>> g = dgl.graph(([0, 1], [1, 2]))
        >>> g.ndata['x'] = th.tensor([[1.], [2.], [3.]])

        >>> # Define the function for sending node features as messages.
        >>> def send_source(edges):
        ...     return {'m': edges.src['x']}
        >>> # Sum the messages received and use this to replace the original node feature.
        >>> def simple_reduce(nodes):
        ...     return {'x': nodes.mailbox['m'].sum(1)}

        Send and receive messages.

        >>> g.send_and_recv(g.edges())
        >>> g.ndata['x']
        tensor([[1.],
                [1.],
                [2.]])

        Note that the feature of node 0 remains the same as it has no incoming edges.
Da Zheng's avatar
Da Zheng committed
4419
        """
4420
4421
4422
        if inplace:
            raise DGLError('The `inplace` option is removed in v0.5.')
        # edge type
Minjie Wang's avatar
Minjie Wang committed
4423
        etid = self.get_etype_id(etype)
4424
4425
4426
4427
4428
4429
        _, dtid = self._graph.metagraph.find_edge(etid)
        etype = self.canonical_etypes[etid]
        # edge IDs
        eid = utils.parse_edges_arg_to_eid(self, edges, etid, 'edges')
        if len(eid) == 0:
            # no computation
4430
            return
4431
4432
        u, v = self.find_edges(eid, etype=etype)
        # call message passing onsubgraph
4433
        g = self if etype is None else self[etype]
4434
4435
4436
        compute_graph, _, dstnodes, _ = _create_compute_graph(g, u, v, eid)
        ndata = core.message_passing(
            compute_graph, message_func, reduce_func, apply_node_func)
4437
        self._set_n_repr(dtid, dstnodes, ndata)
Minjie Wang's avatar
Minjie Wang committed
4438

Da Zheng's avatar
Da Zheng committed
4439
4440
    def pull(self,
             v,
Minjie Wang's avatar
Minjie Wang committed
4441
4442
             message_func,
             reduce_func,
4443
             apply_node_func=None,
Minjie Wang's avatar
Minjie Wang committed
4444
             etype=None,
Da Zheng's avatar
Da Zheng committed
4445
             inplace=False):
4446
4447
        """Pull messages from the specified node(s)' predecessors along the
        specified edge type, aggregate them to update the node features.
4448

Da Zheng's avatar
Da Zheng committed
4449
4450
        Parameters
        ----------
4451
4452
4453
4454
4455
4456
4457
4458
4459
4460
4461
4462
4463
4464
        v : node IDs
            The node IDs. The allowed formats are:

            * ``int``: A single node.
            * Int Tensor: Each element is a node ID. The tensor must have the same device type
              and ID data type as the graph's.
            * iterable[int]: Each element is a node ID.

        message_func : dgl.function.BuiltinFunction or callable
            The message function to generate messages along the edges.
            It must be either a :ref:`api-built-in` or a :ref:`apiudf`.
        reduce_func : dgl.function.BuiltinFunction or callable
            The reduce function to aggregate the messages.
            It must be either a :ref:`api-built-in` or a :ref:`apiudf`.
4465
        apply_node_func : callable, optional
4466
4467
4468
4469
4470
4471
4472
4473
4474
4475
4476
            An optional apply function to further update the node features
            after the message reduction. It must be a :ref:`apiudf`.
        etype : str or (str, str, str), optional
            The type name of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            Can be omitted if the graph has only one type of edges.

Minjie Wang's avatar
Minjie Wang committed
4477
        inplace: bool, optional
4478
4479
4480
4481
4482
4483
4484
4485
4486
4487
4488
4489
            **DEPRECATED**.

        Notes
        -----
        * If some of the given nodes :attr:`v` has no in-edges, DGL does not invoke
          message and reduce functions for these nodes and fill their aggregated messages
          with zero. Users can control the filled values via :meth:`set_n_initializer`.
          DGL still invokes :attr:`apply_node_func` if provided.
        * DGL recommends using DGL's bulit-in function for the :attr:`message_func`
          and the :attr:`reduce_func` arguments,
          because DGL will invoke efficient kernels that avoids copying node features to
          edge features in this case.
Mufei Li's avatar
Mufei Li committed
4490
4491
4492
4493
4494
4495
4496
4497

        Examples
        --------

        >>> import dgl
        >>> import dgl.function as fn
        >>> import torch

4498
4499
4500
4501
4502
4503
4504
4505
4506
4507
4508
4509
4510
        **Homogeneous graph**

        >>> g = dgl.graph(([0, 1, 2, 3], [1, 2, 3, 4]))
        >>> g.ndata['x'] = torch.ones(5, 2)
        >>> g.pull([0, 3, 4], fn.copy_u('x', 'm'), fn.sum('m', 'h'))
        >>> g.ndata['h']
        tensor([[0., 0.],
                [0., 0.],
                [0., 0.],
                [1., 1.],
                [1., 1.]])

        **Heterogeneous graph**
Mufei Li's avatar
Mufei Li committed
4511

4512
4513
4514
4515
        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): ([0, 1], [1, 2]),
        ...     ('user', 'plays', 'game'): ([0, 2], [0, 1])
        ... })
Mufei Li's avatar
Mufei Li committed
4516
4517
4518
4519
4520
4521
4522
4523
4524
        >>> 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.]])
Da Zheng's avatar
Da Zheng committed
4525
        """
4526
4527
        if inplace:
            raise DGLError('The `inplace` option is removed in v0.5.')
4528
        v = utils.prepare_tensor(self, v, 'v')
4529
        if len(v) == 0:
4530
            # no computation
4531
            return
4532
4533
4534
4535
4536
4537
        etid = self.get_etype_id(etype)
        _, dtid = self._graph.metagraph.find_edge(etid)
        etype = self.canonical_etypes[etid]
        g = self if etype is None else self[etype]
        # call message passing on subgraph
        src, dst, eid = g.in_edges(v, form='all')
4538
4539
4540
4541
        compute_graph, _, dstnodes, _ = _create_compute_graph(g, src, dst, eid, v)
        ndata = core.message_passing(
            compute_graph, message_func, reduce_func, apply_node_func)
        self._set_n_repr(dtid, dstnodes, ndata)
Minjie Wang's avatar
Minjie Wang committed
4542

Da Zheng's avatar
Da Zheng committed
4543
4544
    def push(self,
             u,
Minjie Wang's avatar
Minjie Wang committed
4545
4546
             message_func,
             reduce_func,
4547
             apply_node_func=None,
Minjie Wang's avatar
Minjie Wang committed
4548
             etype=None,
Da Zheng's avatar
Da Zheng committed
4549
             inplace=False):
4550
4551
        """Send message from the specified node(s) to their successors
        along the specified edge type and update their node features.
4552

Da Zheng's avatar
Da Zheng committed
4553
4554
        Parameters
        ----------
4555
4556
4557
4558
4559
4560
4561
4562
4563
4564
4565
4566
4567
4568
        v : node IDs
            The node IDs. The allowed formats are:

            * ``int``: A single node.
            * Int Tensor: Each element is a node ID. The tensor must have the same device type
              and ID data type as the graph's.
            * iterable[int]: Each element is a node ID.

        message_func : dgl.function.BuiltinFunction or callable
            The message function to generate messages along the edges.
            It must be either a :ref:`api-built-in` or a :ref:`apiudf`.
        reduce_func : dgl.function.BuiltinFunction or callable
            The reduce function to aggregate the messages.
            It must be either a :ref:`api-built-in` or a :ref:`apiudf`.
4569
        apply_node_func : callable, optional
4570
4571
4572
4573
4574
4575
4576
4577
4578
4579
4580
            An optional apply function to further update the node features
            after the message reduction. It must be a :ref:`apiudf`.
        etype : str or (str, str, str), optional
            The type name of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            Can be omitted if the graph has only one type of edges.

Da Zheng's avatar
Da Zheng committed
4581
        inplace: bool, optional
4582
4583
4584
4585
4586
4587
4588
4589
            **DEPRECATED**.

        Notes
        -----
        DGL recommends using DGL's bulit-in function for the :attr:`message_func`
        and the :attr:`reduce_func` arguments,
        because DGL will invoke efficient kernels that avoids copying node features to
        edge features in this case.
Mufei Li's avatar
Mufei Li committed
4590
4591
4592
4593
4594
4595
4596
4597

        Examples
        --------

        >>> import dgl
        >>> import dgl.function as fn
        >>> import torch

4598
4599
4600
4601
4602
4603
4604
4605
4606
4607
4608
4609
4610
        **Homogeneous graph**

        >>> g = dgl.graph(([0, 1, 2, 3], [1, 2, 3, 4]))
        >>> g.ndata['x'] = torch.ones(5, 2)
        >>> g.push([0, 1], fn.copy_u('x', 'm'), fn.sum('m', 'h'))
        >>> g.ndata['h']
        tensor([[0., 0.],
                [1., 1.],
                [1., 1.],
                [0., 0.],
                [0., 0.]])

        **Heterogeneous graph**
Mufei Li's avatar
Mufei Li committed
4611

4612
        >>> g = dgl.heterograph({('user', 'follows', 'user'): ([0, 0], [1, 2])})
Mufei Li's avatar
Mufei Li committed
4613
4614
4615
4616
4617
4618
4619
4620
4621
        >>> 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.]])
Da Zheng's avatar
Da Zheng committed
4622
        """
4623
4624
4625
4626
        if inplace:
            raise DGLError('The `inplace` option is removed in v0.5.')
        edges = self.out_edges(u, form='eid', etype=etype)
        self.send_and_recv(edges, message_func, reduce_func, apply_node_func, etype=etype)
Da Zheng's avatar
Da Zheng committed
4627
4628

    def update_all(self,
Minjie Wang's avatar
Minjie Wang committed
4629
4630
4631
4632
                   message_func,
                   reduce_func,
                   apply_node_func=None,
                   etype=None):
4633
4634
        """Send messages along all the edges of the specified type
        and update all the nodes of the corresponding destination type.
4635

Da Zheng's avatar
Da Zheng committed
4636
4637
        Parameters
        ----------
4638
4639
4640
4641
4642
4643
        message_func : dgl.function.BuiltinFunction or callable
            The message function to generate messages along the edges.
            It must be either a :ref:`api-built-in` or a :ref:`apiudf`.
        reduce_func : dgl.function.BuiltinFunction or callable
            The reduce function to aggregate the messages.
            It must be either a :ref:`api-built-in` or a :ref:`apiudf`.
4644
        apply_node_func : callable, optional
4645
4646
4647
4648
4649
4650
4651
4652
4653
4654
4655
4656
4657
4658
4659
4660
4661
4662
4663
4664
4665
            An optional apply function to further update the node features
            after the message reduction. It must be a :ref:`apiudf`.
        etype : str or (str, str, str), optional
            The type name of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            Can be omitted if the graph has only one type of edges.

        Notes
        -----
        * If some of the nodes in the graph has no in-edges, DGL does not invoke
          message and reduce functions for these nodes and fill their aggregated messages
          with zero. Users can control the filled values via :meth:`set_n_initializer`.
          DGL still invokes :attr:`apply_node_func` if provided.
        * DGL recommends using DGL's bulit-in function for the :attr:`message_func`
          and the :attr:`reduce_func` arguments,
          because DGL will invoke efficient kernels that avoids copying node features to
          edge features in this case.
Mufei Li's avatar
Mufei Li committed
4666
4667
4668
4669
4670

        Examples
        --------
        >>> import dgl
        >>> import dgl.function as fn
4671
        >>> import torch
Mufei Li's avatar
Mufei Li committed
4672

4673
4674
4675
4676
4677
4678
4679
4680
4681
4682
4683
4684
4685
        **Homogeneous graph**

        >>> g = dgl.graph(([0, 1, 2, 3], [1, 2, 3, 4]))
        >>> g.ndata['x'] = torch.ones(5, 2)
        >>> g.update_all(fn.copy_u('x', 'm'), fn.sum('m', 'h'))
        >>> g.ndata['h']
        tensor([[0., 0.],
                [1., 1.],
                [1., 1.],
                [1., 1.],
                [1., 1.]])

        **Heterogeneous graph**
Mufei Li's avatar
Mufei Li committed
4686

4687
        >>> g = dgl.heterograph({('user', 'follows', 'user'): ([0, 1, 2], [1, 2, 2])})
Mufei Li's avatar
Mufei Li committed
4688
4689
4690
4691
4692
4693
4694
4695
4696

        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.]])
Da Zheng's avatar
Da Zheng committed
4697
        """
Minjie Wang's avatar
Minjie Wang committed
4698
        etid = self.get_etype_id(etype)
4699
4700
4701
4702
4703
        etype = self.canonical_etypes[etid]
        _, dtid = self._graph.metagraph.find_edge(etid)
        g = self if etype is None else self[etype]
        ndata = core.message_passing(g, message_func, reduce_func, apply_node_func)
        self._set_n_repr(dtid, ALL, ndata)
4704

4705
4706
4707
    #################################################################
    # Message passing on heterograph
    #################################################################
Da Zheng's avatar
Da Zheng committed
4708

Mufei Li's avatar
Mufei Li committed
4709
    def multi_update_all(self, etype_dict, cross_reducer, apply_node_func=None):
4710
4711
4712
        r"""Send messages along all the edges, reduce them by first type-wisely
        then across different types, and then update the node features of all
        the nodes.
Da Zheng's avatar
Da Zheng committed
4713
4714
4715

        Parameters
        ----------
Mufei Li's avatar
Mufei Li committed
4716
        etype_dict : dict
4717
4718
4719
4720
4721
4722
4723
4724
4725
4726
4727
4728
4729
4730
4731
4732
4733
            Arguments for edge-type-wise message passing. The keys are edge types
            while the values are message passing arguments.

            The allowed key formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            The value must be a tuple ``(message_func, reduce_func, [apply_node_func])``, where

            * message_func : dgl.function.BuiltinFunction or callable
                The message function to generate messages along the edges.
                It must be either a :ref:`api-built-in` or a :ref:`apiudf`.
            * reduce_func : dgl.function.BuiltinFunction or callable
                The reduce function to aggregate the messages.
                It must be either a :ref:`api-built-in` or a :ref:`apiudf`.
Mufei Li's avatar
Mufei Li committed
4734
            * apply_node_func : callable, optional
4735
4736
4737
                An optional apply function to further update the node features
                after the message reduction. It must be a :ref:`apiudf`.

4738
4739
4740
4741
4742
        cross_reducer : str or callable function
            Cross type reducer. One of ``"sum"``, ``"min"``, ``"max"``, ``"mean"``, ``"stack"``
            or a callable function. If a callable function is provided, the input argument must be
            a single list of tensors containing aggregation results from each edge type, and the
            output of function must be a single tensor.
4743
4744
4745
4746
4747
4748
4749
4750
4751
4752
4753
4754
        apply_node_func : callable, optional
            An optional apply function after the messages are reduced both
            type-wisely and across different types.
            It must be a :ref:`apiudf`.

        Notes
        -----
        DGL recommends using DGL's bulit-in function for the message_func
        and the reduce_func in the type-wise message passing arguments,
        because DGL will invoke efficient kernels that avoids copying node features to
        edge features in this case.

Mufei Li's avatar
Mufei Li committed
4755
4756
4757
4758
4759
4760
4761
4762
4763

        Examples
        --------
        >>> import dgl
        >>> import dgl.function as fn
        >>> import torch

        Instantiate a heterograph.

4764
4765
4766
4767
        >>> g = dgl.heterograph({
        ...     ('user', 'follows', 'user'): ([0, 1], [1, 1]),
        ...     ('game', 'attracts', 'user'): ([0], [1])
        ... })
Mufei Li's avatar
Mufei Li committed
4768
4769
4770
4771
4772
4773
        >>> g.nodes['user'].data['h'] = torch.tensor([[1.], [2.]])
        >>> g.nodes['game'].data['h'] = torch.tensor([[1.]])

        Update all.

        >>> g.multi_update_all(
4774
4775
4776
        ...     {'follows': (fn.copy_src('h', 'm'), fn.sum('m', 'h')),
        ...      'attracts': (fn.copy_src('h', 'm'), fn.sum('m', 'h'))},
        ... "sum")
Mufei Li's avatar
Mufei Li committed
4777
4778
4779
4780
        >>> g.nodes['user'].data['h']
        tensor([[0.],
                [4.]])
        """
Minjie Wang's avatar
Minjie Wang committed
4781
        all_out = defaultdict(list)
4782
        merge_order = defaultdict(list)
4783
4784
4785
4786
4787
4788
4789
4790
        for etype, args in etype_dict.items():
            etid = self.get_etype_id(etype)
            _, dtid = self._graph.metagraph.find_edge(etid)
            args = pad_tuple(args, 3)
            if args is None:
                raise DGLError('Invalid arguments for edge type "{}". Should be '
                               '(msg_func, reduce_func, [apply_node_func])'.format(etype))
            mfunc, rfunc, afunc = args
4791
4792
            g = self if etype is None else self[etype]
            all_out[dtid].append(core.message_passing(g, mfunc, rfunc, afunc))
4793
            merge_order[dtid].append(etid)  # use edge type id as merge order hint
Minjie Wang's avatar
Minjie Wang committed
4794
4795
        for dtid, frames in all_out.items():
            # merge by cross_reducer
4796
            self._node_frames[dtid].update(
4797
                reduce_dict_data(frames, cross_reducer, merge_order[dtid]))
Minjie Wang's avatar
Minjie Wang committed
4798
            # apply
Mufei Li's avatar
Mufei Li committed
4799
            if apply_node_func is not None:
4800
4801
4802
4803
4804
                self.apply_nodes(apply_node_func, ALL, self.ntypes[dtid])

    #################################################################
    # Message propagation
    #################################################################
Minjie Wang's avatar
Minjie Wang committed
4805
4806
4807
4808
4809
4810
4811

    def prop_nodes(self,
                   nodes_generator,
                   message_func,
                   reduce_func,
                   apply_node_func=None,
                   etype=None):
Mufei Li's avatar
Mufei Li committed
4812
        """Propagate messages using graph traversal by sequentially triggering
Minjie Wang's avatar
Minjie Wang committed
4813
4814
4815
4816
4817
4818
        :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.
Da Zheng's avatar
Da Zheng committed
4819
4820
4821

        Parameters
        ----------
4822
4823
4824
4825
4826
4827
4828
4829
4830
4831
        nodes_generator : iterable[node IDs]
            The generator of node frontiers. Each frontier is a set of node IDs
            stored in Tensor or python iterables.
            It specifies which nodes perform :func:`pull` at each step.
        message_func : dgl.function.BuiltinFunction or callable
            The message function to generate messages along the edges.
            It must be either a :ref:`api-built-in` or a :ref:`apiudf`.
        reduce_func : dgl.function.BuiltinFunction or callable
            The reduce function to aggregate the messages.
            It must be either a :ref:`api-built-in` or a :ref:`apiudf`.
Minjie Wang's avatar
Minjie Wang committed
4832
        apply_node_func : callable, optional
4833
4834
4835
4836
4837
4838
4839
4840
4841
4842
            An optional apply function to further update the node features
            after the message reduction. It must be a :ref:`apiudf`.
        etype : str or (str, str, str), optional
            The type name of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            Can be omitted if the graph has only one type of edges.
Mufei Li's avatar
Mufei Li committed
4843
4844
4845
4846
4847
4848
4849
4850
4851

        Examples
        --------
        >>> import torch
        >>> import dgl
        >>> import dgl.function as fn

        Instantiate a heterogrph and perform multiple rounds of message passing.

4852
        >>> g = dgl.heterograph({('user', 'follows', 'user'): ([0, 1, 2, 3], [2, 3, 4, 4])})
Mufei Li's avatar
Mufei Li committed
4853
4854
        >>> g.nodes['user'].data['h'] = torch.tensor([[1.], [2.], [3.], [4.], [5.]])
        >>> g['follows'].prop_nodes([[2, 3], [4]], fn.copy_src('h', 'm'),
4855
        ...                         fn.sum('m', 'h'), etype='follows')
Mufei Li's avatar
Mufei Li committed
4856
4857
4858
4859
4860
        tensor([[1.],
                [2.],
                [1.],
                [2.],
                [3.]])
Da Zheng's avatar
Da Zheng committed
4861

Minjie Wang's avatar
Minjie Wang committed
4862
4863
4864
        See Also
        --------
        prop_edges
Da Zheng's avatar
Da Zheng committed
4865
        """
Minjie Wang's avatar
Minjie Wang committed
4866
4867
        for node_frontier in nodes_generator:
            self.pull(node_frontier, message_func, reduce_func, apply_node_func, etype=etype)
Da Zheng's avatar
Da Zheng committed
4868

Minjie Wang's avatar
Minjie Wang committed
4869
4870
4871
4872
4873
4874
    def prop_edges(self,
                   edges_generator,
                   message_func,
                   reduce_func,
                   apply_node_func=None,
                   etype=None):
Mufei Li's avatar
Mufei Li committed
4875
        """Propagate messages using graph traversal by sequentially triggering
Minjie Wang's avatar
Minjie Wang committed
4876
        :func:`send_and_recv()` on edges.
Da Zheng's avatar
Da Zheng committed
4877

Minjie Wang's avatar
Minjie Wang committed
4878
4879
4880
        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.
Da Zheng's avatar
Da Zheng committed
4881

Mufei Li's avatar
Mufei Li committed
4882
        Edges in the same frontier will be triggered together, and edges in
Minjie Wang's avatar
Minjie Wang committed
4883
        different frontiers will be triggered according to the generating order.
Da Zheng's avatar
Da Zheng committed
4884
4885
4886

        Parameters
        ----------
Minjie Wang's avatar
Minjie Wang committed
4887
4888
        edges_generator : generator
            The generator of edge frontiers.
4889
4890
4891
4892
4893
4894
        message_func : dgl.function.BuiltinFunction or callable
            The message function to generate messages along the edges.
            It must be either a :ref:`api-built-in` or a :ref:`apiudf`.
        reduce_func : dgl.function.BuiltinFunction or callable
            The reduce function to aggregate the messages.
            It must be either a :ref:`api-built-in` or a :ref:`apiudf`.
Minjie Wang's avatar
Minjie Wang committed
4895
        apply_node_func : callable, optional
4896
4897
4898
4899
4900
4901
4902
4903
4904
4905
            An optional apply function to further update the node features
            after the message reduction. It must be a :ref:`apiudf`.
        etype : str or (str, str, str), optional
            The type name of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            Can be omitted if the graph has only one type of edges.
Mufei Li's avatar
Mufei Li committed
4906
4907
4908
4909
4910
4911
4912
4913
4914

        Examples
        --------
        >>> import torch
        >>> import dgl
        >>> import dgl.function as fn

        Instantiate a heterogrph and perform multiple rounds of message passing.

4915
        >>> g = dgl.heterograph({('user', 'follows', 'user'): ([0, 1, 2, 3], [2, 3, 4, 4])})
Mufei Li's avatar
Mufei Li committed
4916
4917
        >>> 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'),
4918
        ...                         fn.sum('m', 'h'), etype='follows')
Mufei Li's avatar
Mufei Li committed
4919
4920
4921
4922
4923
4924
        >>> g.nodes['user'].data['h']
        tensor([[1.],
                [2.],
                [1.],
                [2.],
                [3.]])
Da Zheng's avatar
Da Zheng committed
4925

Minjie Wang's avatar
Minjie Wang committed
4926
4927
4928
        See Also
        --------
        prop_nodes
Da Zheng's avatar
Da Zheng committed
4929
        """
Minjie Wang's avatar
Minjie Wang committed
4930
4931
4932
        for edge_frontier in edges_generator:
            self.send_and_recv(edge_frontier, message_func, reduce_func,
                               apply_node_func, etype=etype)
Da Zheng's avatar
Da Zheng committed
4933

Minjie Wang's avatar
Minjie Wang committed
4934
4935
4936
    #################################################################
    # Misc
    #################################################################
Da Zheng's avatar
Da Zheng committed
4937

Minjie Wang's avatar
Minjie Wang committed
4938
    def filter_nodes(self, predicate, nodes=ALL, ntype=None):
4939
        """Return the IDs of the nodes with the given node type that satisfy
Da Zheng's avatar
Da Zheng committed
4940
4941
4942
4943
4944
        the given predicate.

        Parameters
        ----------
        predicate : callable
4945
4946
4947
            A function of signature ``func(nodes) -> Tensor``.
            ``nodes`` are :class:`dgl.NodeBatch` objects.
            Its output tensor should be a 1D boolean tensor with
Da Zheng's avatar
Da Zheng committed
4948
4949
            each element indicating whether the corresponding node in
            the batch satisfies the predicate.
4950
4951
4952
4953
4954
4955
4956
4957
4958
        nodes : node ID(s), optional
            The node(s) for query. The allowed formats are:

            - Tensor: A 1D tensor that contains the node(s) for query, whose data type
              and device should be the same as the :py:attr:`idtype` and device of the graph.
            - iterable[int] : Similar to the tensor, but stores node IDs in a sequence
              (e.g. list, tuple, numpy.ndarray).

            By default, it considers all nodes.
Minjie Wang's avatar
Minjie Wang committed
4959
        ntype : str, optional
4960
4961
            The node type for query. If the graph has multiple node types, one must
            specify the argument. Otherwise, it can be omitted.
Da Zheng's avatar
Da Zheng committed
4962
4963
4964

        Returns
        -------
4965
        Tensor
4966
            A 1D tensor that contains the ID(s) of the node(s) that satisfy the predicate.
Mufei Li's avatar
Mufei Li committed
4967
4968
4969

        Examples
        --------
4970
4971
4972

        The following example uses PyTorch backend.

Mufei Li's avatar
Mufei Li committed
4973
        >>> import dgl
4974
4975
4976
4977
4978
4979
4980
4981
4982
4983
4984
4985
4986
4987
4988
4989
4990
4991
4992
4993
4994
4995
4996
4997
4998
4999
5000
5001
5002
        >>> import torch

        Define a predicate function.

        >>> def nodes_with_feature_one(nodes):
        ...     # Whether a node has feature 1
        ...     return (nodes.data['h'] == 1.).squeeze(1)

        Filter nodes for a homogeneous graph.

        >>> g = dgl.graph((torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])))
        >>> g.ndata['h'] = torch.tensor([[0.], [1.], [1.], [0.]])
        >>> print(g.filter_nodes(nodes_with_feature_one))
        tensor([1, 2])

        Filter on nodes with IDs 0 and 1

        >>> print(g.filter_nodes(nodes_with_feature_one, nodes=torch.tensor([0, 1])))
        tensor([1])

        Filter nodes for a heterogeneous graph.

        >>> g = dgl.heterograph({
        ...     ('user', 'plays', 'game'): (torch.tensor([0, 1, 1, 2]),
        ...                                 torch.tensor([0, 0, 1, 1]))})
        >>> g.nodes['user'].data['h'] = torch.tensor([[0.], [1.], [1.]])
        >>> g.nodes['game'].data['h'] = torch.tensor([[0.], [1.]])
        >>> # Filter for 'user' nodes
        >>> print(g.filter_nodes(nodes_with_feature_one, ntype='user'))
Mufei Li's avatar
Mufei Li committed
5003
        tensor([1, 2])
Da Zheng's avatar
Da Zheng committed
5004
        """
5005
5006
5007
5008
5009
5010
        if is_all(nodes):
            nodes = self.nodes(ntype)
        v = utils.prepare_tensor(self, nodes, 'nodes')
        if F.as_scalar(F.sum(self.has_nodes(v, ntype=ntype), dim=0)) != len(v):
            raise DGLError('v contains invalid node IDs')

5011
5012
5013
5014
5015
5016
5017
        with self.local_scope():
            self.apply_nodes(lambda nbatch: {'_mask' : predicate(nbatch)}, nodes, ntype)
            ntype = self.ntypes[0] if ntype is None else ntype
            mask = self.nodes[ntype].data['_mask']
            if is_all(nodes):
                return F.nonzero_1d(mask)
            else:
5018
                return F.boolean_mask(v, F.gather_row(mask, v))
Minjie Wang's avatar
Minjie Wang committed
5019
5020

    def filter_edges(self, predicate, edges=ALL, etype=None):
5021
        """Return the IDs of the edges with the given edge type that satisfy
Da Zheng's avatar
Da Zheng committed
5022
5023
5024
5025
5026
        the given predicate.

        Parameters
        ----------
        predicate : callable
5027
5028
5029
            A function of signature ``func(edges) -> Tensor``.
            ``edges`` are :class:`dgl.EdgeBatch` objects.
            Its output tensor should be a 1D boolean tensor with
Da Zheng's avatar
Da Zheng committed
5030
5031
            each element indicating whether the corresponding edge in
            the batch satisfies the predicate.
5032
5033
        edges : edges
            The edges to send and receive messages on. The allowed input formats are:
5034

5035
5036
5037
5038
5039
5040
5041
5042
5043
5044
5045
5046
5047
5048
5049
5050
5051
5052
            * ``int``: A single edge ID.
            * Int Tensor: Each element is an edge ID.  The tensor must have the same device type
              and ID data type as the graph's.
            * iterable[int]: Each element is an edge ID.
            * (Tensor, Tensor): The node-tensors format where the i-th elements
              of the two tensors specify an edge.
            * (iterable[int], iterable[int]): Similar to the node-tensors format but
              stores edge endpoints in python iterables.

            By default, it considers all the edges.
        etype : str or (str, str, str), optional
            The type name of the edges. The allowed type name formats are:

            * ``(str, str, str)`` for source node type, edge type and destination node type.
            * or one ``str`` edge type name if the name can uniquely identify a
              triplet format in the graph.

            Can be omitted if the graph has only one type of edges.
Da Zheng's avatar
Da Zheng committed
5053
5054
5055

        Returns
        -------
5056
        Tensor
5057
            A 1D tensor that contains the ID(s) of the edge(s) that satisfy the predicate.
Mufei Li's avatar
Mufei Li committed
5058
5059
5060

        Examples
        --------
5061
5062
5063

        The following example uses PyTorch backend.

Mufei Li's avatar
Mufei Li committed
5064
        >>> import dgl
5065
5066
5067
5068
5069
5070
5071
5072
5073
5074
5075
5076
5077
5078
5079
5080
5081
5082
5083
5084
5085
5086
5087
5088
5089
5090
5091
5092
5093
        >>> import torch

        Define a predicate function.

        >>> def edges_with_feature_one(edges):
        ...     # Whether an edge has feature 1
        ...     return (edges.data['h'] == 1.).squeeze(1)

        Filter edges for a homogeneous graph.

        >>> g = dgl.graph((torch.tensor([0, 1, 2]), torch.tensor([1, 2, 3])))
        >>> g.edata['h'] = torch.tensor([[0.], [1.], [1.]])
        >>> print(g.filter_edges(edges_with_feature_one))
        tensor([1, 2])

        Filter on edges with IDs 0 and 1

        >>> print(g.filter_edges(edges_with_feature_one, edges=torch.tensor([0, 1])))
        tensor([1])

        Filter edges for a heterogeneous graph.

        >>> g = dgl.heterograph({
        ...     ('user', 'plays', 'game'): (torch.tensor([0, 1, 1, 2]),
        ...                                 torch.tensor([0, 0, 1, 1])),
        ...     ('user', 'follows', 'user'): (torch.tensor([0, 1]), torch.tensor([1, 2]))})
        >>> g.edges['plays'].data['h'] = torch.tensor([[0.], [1.], [1.], [0.]])
        >>> # Filter for 'plays' nodes
        >>> print(g.filter_edges(edges_with_feature_one, etype='plays'))
Mufei Li's avatar
Mufei Li committed
5094
        tensor([1, 2])
Da Zheng's avatar
Da Zheng committed
5095
        """
5096
5097
5098
5099
5100
5101
5102
5103
5104
5105
5106
5107
5108
5109
5110
5111
5112
5113
5114
5115
5116
5117
        if is_all(edges):
            pass
        elif isinstance(edges, tuple):
            u, v = edges
            srctype, _, dsttype = self.to_canonical_etype(etype)
            u = utils.prepare_tensor(self, u, 'u')
            if F.as_scalar(F.sum(self.has_nodes(u, ntype=srctype), dim=0)) != len(u):
                raise DGLError('edges[0] contains invalid node IDs')
            v = utils.prepare_tensor(self, v, 'v')
            if F.as_scalar(F.sum(self.has_nodes(v, ntype=dsttype), dim=0)) != len(v):
                raise DGLError('edges[1] contains invalid node IDs')
        elif isinstance(edges, Iterable) or F.is_tensor(edges):
            edges = utils.prepare_tensor(self, edges, 'edges')
            min_eid = F.as_scalar(F.min(edges, 0))
            if len(edges) > 0 > min_eid:
                raise DGLError('Invalid edge ID {:d}'.format(min_eid))
            max_eid = F.as_scalar(F.max(edges, 0))
            if len(edges) > 0 and max_eid >= self.num_edges(etype):
                raise DGLError('Invalid edge ID {:d}'.format(max_eid))
        else:
            raise ValueError('Unsupported type of edges:', type(edges))

5118
5119
5120
5121
5122
5123
5124
5125
5126
5127
5128
        with self.local_scope():
            self.apply_edges(lambda ebatch: {'_mask' : predicate(ebatch)}, edges, etype)
            etype = self.canonical_etypes[0] if etype is None else etype
            mask = self.edges[etype].data['_mask']
            if is_all(edges):
                return F.nonzero_1d(mask)
            else:
                if isinstance(edges, tuple):
                    e = self.edge_ids(edges[0], edges[1], etype=etype)
                else:
                    e = utils.prepare_tensor(self, edges, 'edges')
5129
                return F.boolean_mask(e, F.gather_row(mask, e))
Minjie Wang's avatar
Minjie Wang committed
5130

5131
5132
    @property
    def device(self):
5133
5134
5135
5136
5137
5138
5139
        """Get the device of the graph.

        Returns
        -------
        device context
            The device of the graph, which should be a framework-specific device object
            (e.g., ``torch.device``).
5140
5141
5142
5143
5144

        Examples
        --------
        The following example uses PyTorch backend.

5145
5146
5147
5148
5149
5150
        >>> import dgl
        >>> import torch

        Create a homogeneous graph for demonstration.

        >>> g = dgl.graph((torch.tensor([0, 1]), torch.tensor([1, 2])))
5151
5152
5153
        >>> print(g.device)
        device(type='cpu')

5154
        The case of heterogeneous graphs is the same.
5155
5156
5157
        """
        return F.to_backend_ctx(self._graph.ctx)

5158
5159
    def to(self, device, **kwargs):  # pylint: disable=invalid-name
        """Move ndata, edata and graph structure to the targeted device (cpu/gpu).
Da Zheng's avatar
Da Zheng committed
5160

5161
5162
5163
        If the graph is already on the specified device, the function directly returns it.
        Otherwise, it returns a cloned graph on the specified device.

Da Zheng's avatar
Da Zheng committed
5164
5165
        Parameters
        ----------
5166
        device : Framework-specific device context object
5167
            The context to move data to (e.g., ``torch.device``).
5168
5169
        kwargs : Key-word arguments.
            Key-word arguments fed to the framework copy function.
Minjie Wang's avatar
Minjie Wang committed
5170

5171
5172
        Returns
        -------
5173
5174
        DGLGraph
            The graph on the specified device.
5175

Minjie Wang's avatar
Minjie Wang committed
5176
5177
5178
5179
        Examples
        --------
        The following example uses PyTorch backend.

5180
        >>> import dgl
Minjie Wang's avatar
Minjie Wang committed
5181
        >>> import torch
5182
5183
5184
5185

        >>> g = dgl.graph((torch.tensor([1, 0]), torch.tensor([1, 2])))
        >>> g.ndata['h'] = torch.ones(3, 1)
        >>> g.edata['h'] = torch.zeros(2, 2)
5186
5187
5188
        >>> g1 = g.to(torch.device('cuda:0'))
        >>> print(g1.device)
        device(type='cuda', index=0)
5189
5190
5191
5192
5193
5194
5195
        >>> print(g1.ndata['h'].device)
        device(type='cuda', index=0)
        >>> print(g1.nodes().device)
        device(type='cuda', index=0)

        The original graph is still on CPU.

5196
5197
        >>> print(g.device)
        device(type='cpu')
5198
5199
5200
5201
5202
5203
        >>> print(g.ndata['h'].device)
        device(type='cpu')
        >>> print(g.nodes().device)
        device(type='cpu')

        The case of heterogeneous graphs is the same.
Da Zheng's avatar
Da Zheng committed
5204
        """
5205
        if device is None or self.device == device:
5206
            return self
5207
5208
5209
5210
5211
5212
5213
5214

        ret = copy.copy(self)

        # 1. Copy graph structure
        ret._graph = self._graph.copy_to(utils.to_dgl_context(device))

        # 2. Copy features
        # TODO(minjie): handle initializer
5215
5216
        new_nframes = []
        for nframe in self._node_frames:
5217
            new_nframes.append(nframe.to(device, **kwargs))
5218
5219
        ret._node_frames = new_nframes

5220
5221
        new_eframes = []
        for eframe in self._edge_frames:
5222
            new_eframes.append(eframe.to(device, **kwargs))
5223
5224
5225
5226
5227
5228
5229
5230
5231
5232
5233
5234
5235
5236
5237
5238
5239
5240
5241
5242
5243
5244
5245
5246
5247
5248
5249
5250
5251
5252
5253
5254
5255
5256
5257
5258
5259
5260
5261
5262
5263
5264
5265
5266
5267
5268
5269
5270
5271
5272
5273
5274
5275
5276
5277
5278
5279
5280
        ret._edge_frames = new_eframes

        # 2. Copy misc info
        if self._batch_num_nodes is not None:
            new_bnn = {k : F.copy_to(num, device, **kwargs)
                       for k, num in self._batch_num_nodes.items()}
            ret._batch_num_nodes = new_bnn
        if self._batch_num_edges is not None:
            new_bne = {k : F.copy_to(num, device, **kwargs)
                       for k, num in self._batch_num_edges.items()}
            ret._batch_num_edges = new_bne

        return ret

    def cpu(self):
        """Return a new copy of this graph on CPU.

        Returns
        -------
        DGLHeteroGraph
            Graph on CPU.

        See Also
        --------
        to
        """
        return self.to(F.cpu())

    def clone(self):
        """Return a heterograph object that is a clone of current graph.

        Returns
        -------
        DGLHeteroGraph
            The graph object that is a clone of current graph.
        """
        # XXX(minjie): Do a shallow copy first to clone some internal metagraph information.
        #   Not a beautiful solution though.
        ret = copy.copy(self)

        # Clone the graph structure
        meta_edges = []
        for s_ntype, _, d_ntype in self.canonical_etypes:
            meta_edges.append((self.get_ntype_id(s_ntype), self.get_ntype_id(d_ntype)))

        metagraph = graph_index.from_edge_list(meta_edges, True)
        # rebuild graph idx
        num_nodes_per_type = [self.number_of_nodes(c_ntype) for c_ntype in self.ntypes]
        relation_graphs = [self._graph.get_relation_graph(self.get_etype_id(c_etype))
                           for c_etype in self.canonical_etypes]
        ret._graph = heterograph_index.create_heterograph_from_relations(
            metagraph, relation_graphs, utils.toindex(num_nodes_per_type, "int64"))

        # Clone the frames
        ret._node_frames = [fr.clone() for fr in self._node_frames]
        ret._edge_frames = [fr.clone() for fr in self._edge_frames]

        return ret
Da Zheng's avatar
Da Zheng committed
5281

Minjie Wang's avatar
Minjie Wang committed
5282
    def local_var(self):
5283
        """Return a graph object for usage in a local function scope.
Minjie Wang's avatar
Minjie Wang committed
5284
5285
5286

        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,
5287
        thus making it easier to use in a function scope (e.g. forward computation of a model).
Minjie Wang's avatar
Minjie Wang committed
5288
5289
5290
5291

        If set, the local graph object will use same initializers for node features and
        edge features.

Mufei Li's avatar
Mufei Li committed
5292
5293
        Returns
        -------
5294
5295
        DGLGraph
            The graph object for a local variable.
Mufei Li's avatar
Mufei Li committed
5296
5297
5298

        Notes
        -----
5299
5300
        Inplace operations do reflect to the original graph. This function also has little
        overhead when the number of feature tensors in this graph is small.
Mufei Li's avatar
Mufei Li committed
5301

Minjie Wang's avatar
Minjie Wang committed
5302
5303
        Examples
        --------
5304

Minjie Wang's avatar
Minjie Wang committed
5305
5306
        The following example uses PyTorch backend.

5307
5308
5309
5310
        >>> import dgl
        >>> import torch

        Create a function for computation on graphs.
Minjie Wang's avatar
Minjie Wang committed
5311
5312

        >>> def foo(g):
5313
5314
5315
5316
5317
5318
5319
5320
5321
5322
        ...     g = g.local_var()
        ...     g.edata['h'] = torch.ones((g.num_edges(), 3))
        ...     g.edata['h2'] = torch.ones((g.num_edges(), 3))
        ...     return g.edata['h']

        ``local_var`` avoids changing the graph features when exiting the function.

        >>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([0, 0, 2])))
        >>> g.edata['h'] = torch.zeros((g.num_edges(), 3))
        >>> newh = foo(g)
Mufei Li's avatar
Mufei Li committed
5323
        >>> print(g.edata['h'])  # still get tensor of all zeros
5324
5325
5326
5327
5328
        tensor([[0., 0., 0.],
                [0., 0., 0.],
                [0., 0., 0.]])
        >>> 'h2' in g.edata      # new feature set in the function scope is not found
        False
Minjie Wang's avatar
Minjie Wang committed
5329

5330
        In-place operations will still reflect to the original graph.
Minjie Wang's avatar
Minjie Wang committed
5331
5332

        >>> def foo(g):
5333
5334
5335
5336
5337
5338
5339
5340
5341
5342
5343
5344
        ...     g = g.local_var()
        ...     # in-place operation
        ...     g.edata['h'] += 1
        ...     return g.edata['h']

        >>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([0, 0, 2])))
        >>> g.edata['h'] = torch.zeros((g.num_edges(), 1))
        >>> newh = foo(g)
        >>> print(g.edata['h'])  # the result changes
        tensor([[1.],
                [1.],
                [1.]])
Minjie Wang's avatar
Minjie Wang committed
5345
5346
5347

        See Also
        --------
5348
        local_scope
Minjie Wang's avatar
Minjie Wang committed
5349
        """
5350
        ret = copy.copy(self)
5351
5352
        ret._node_frames = [fr.clone() for fr in self._node_frames]
        ret._edge_frames = [fr.clone() for fr in self._edge_frames]
5353
        return ret
Minjie Wang's avatar
Minjie Wang committed
5354
5355
5356

    @contextmanager
    def local_scope(self):
5357
        """Enter a local scope context for the graph.
Minjie Wang's avatar
Minjie Wang committed
5358
5359

        By entering a local scope, any out-place mutation to the feature data will
5360
5361
        not reflect to the original graph, thus making it easier to use in a function scope
        (e.g. forward computation of a model).
Minjie Wang's avatar
Minjie Wang committed
5362
5363
5364
5365

        If set, the local scope will use same initializers for node features and
        edge features.

5366
5367
5368
5369
5370
        Notes
        -----
        Inplace operations do reflect to the original graph. This function also has little
        overhead when the number of feature tensors in this graph is small.

Minjie Wang's avatar
Minjie Wang committed
5371
5372
        Examples
        --------
5373

Minjie Wang's avatar
Minjie Wang committed
5374
5375
        The following example uses PyTorch backend.

5376
5377
5378
5379
        >>> import dgl
        >>> import torch

        Create a function for computation on graphs.
Minjie Wang's avatar
Minjie Wang committed
5380
5381

        >>> def foo(g):
5382
5383
5384
5385
5386
5387
5388
5389
5390
5391
        ...     with g.local_scope():
        ...         g.edata['h'] = torch.ones((g.num_edges(), 3))
        ...         g.edata['h2'] = torch.ones((g.num_edges(), 3))
        ...         return g.edata['h']

        ``local_scope`` avoids changing the graph features when exiting the function.

        >>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([0, 0, 2])))
        >>> g.edata['h'] = torch.zeros((g.num_edges(), 3))
        >>> newh = foo(g)
Mufei Li's avatar
Mufei Li committed
5392
        >>> print(g.edata['h'])  # still get tensor of all zeros
5393
5394
5395
5396
5397
        tensor([[0., 0., 0.],
                [0., 0., 0.],
                [0., 0., 0.]])
        >>> 'h2' in g.edata      # new feature set in the function scope is not found
        False
Minjie Wang's avatar
Minjie Wang committed
5398

5399
        In-place operations will still reflect to the original graph.
Minjie Wang's avatar
Minjie Wang committed
5400
5401

        >>> def foo(g):
5402
5403
5404
5405
5406
5407
5408
5409
5410
5411
5412
5413
        ...     with g.local_scope():
        ...         # in-place operation
        ...         g.edata['h'] += 1
        ...         return g.edata['h']

        >>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([0, 0, 2])))
        >>> g.edata['h'] = torch.zeros((g.num_edges(), 1))
        >>> newh = foo(g)
        >>> print(g.edata['h'])  # the result changes
        tensor([[1.],
                [1.],
                [1.]])
Minjie Wang's avatar
Minjie Wang committed
5414
5415
5416
5417
5418
5419
5420

        See Also
        --------
        local_var
        """
        old_nframes = self._node_frames
        old_eframes = self._edge_frames
5421
5422
        self._node_frames = [fr.clone() for fr in self._node_frames]
        self._edge_frames = [fr.clone() for fr in self._edge_frames]
Minjie Wang's avatar
Minjie Wang committed
5423
5424
5425
5426
        yield
        self._node_frames = old_nframes
        self._edge_frames = old_eframes

5427
5428
5429
5430
5431
5432
5433
5434
    def formats(self, formats=None):
        r"""Get a cloned graph with the specified sparse format(s) or query
        for the usage status of sparse formats

        The API copies both the graph structure and the features.

        If the input graph has multiple edge types, they will have the same
        sparse format.
5435

5436
5437
        Parameters
        ----------
5438
5439
5440
5441
        formats : str or list of str or None

            * If formats is None, return the usage status of sparse formats
            * Otherwise, it can be ``'coo'``/``'csr'``/``'csc'`` or a sublist of
5442
              them, specifying the sparse formats to use.
5443

5444
5445
        Returns
        -------
5446
5447
5448
5449
5450
5451
        dict or DGLGraph

            * If formats is None, the result will be a dict recording the usage
              status of sparse formats.
            * Otherwise, a DGLGraph will be returned, which is a clone of the
              original graph with the specified sparse format(s) ``formats``.
5452

5453
5454
5455
        Examples
        --------

5456
        The following example uses PyTorch backend.
5457

5458
5459
        >>> import dgl
        >>> import torch
5460

5461
5462
        **Homographs or Heterographs with A Single Edge Type**

5463
        >>> g = dgl.graph(([0, 0, 1], [2, 3, 2]))
5464
5465
5466
5467
5468
5469
5470
5471
5472
5473
5474
5475
5476
5477
5478
        >>> g.ndata['h'] = torch.ones(4, 1)
        >>> # Check status of format usage
        >>> g.formats()
        {'created': ['coo'], 'not created': ['csr', 'csc']}
        >>> # Get a clone of the graph with 'csr' format
        >>> csr_g = g.formats('csr')
        >>> # Only allowed formats will be displayed in the status query
        >>> csr_g.formats()
        {'created': ['csr'], 'not created': []}
        >>> # Features are copied as well
        >>> csr_g.ndata['h']
        tensor([[1.],
                [1.],
                [1.],
                [1.]])
5479

5480
        **Heterographs with Multiple Edge Types**
5481

5482
        >>> g = dgl.heterograph({
5483
5484
5485
5486
5487
        ...     ('user', 'plays', 'game'): (torch.tensor([0, 1, 1, 2]),
        ...                                 torch.tensor([0, 0, 1, 1])),
        ...     ('developer', 'develops', 'game'): (torch.tensor([0, 1]),
        ...                                         torch.tensor([0, 1]))
        ...     })
5488
5489
5490
5491
5492
5493
5494
        >>> g.formats()
        {'created': ['coo'], 'not created': ['csr', 'csc']}
        >>> # Get a clone of the graph with 'csr' format
        >>> csr_g = g.formats('csr')
        >>> # Only allowed formats will be displayed in the status query
        >>> csr_g.formats()
        {'created': ['csr'], 'not created': []}
5495
        """
5496
5497
5498
5499
5500
5501
5502
5503
        if formats is None:
            # Return the format information
            return self._graph.formats()
        else:
            # Convert the graph to use another format
            ret = copy.copy(self)
            ret._graph = self._graph.formats(formats)
            return ret
5504

5505
    def create_formats_(self):
5506
        r"""Create all sparse matrices allowed for the graph.
5507

5508
5509
5510
        By default, we create sparse matrices for a graph only when necessary.
        In some cases we may want to create them immediately (e.g. in a
        multi-process data loader), which can be achieved via this API.
5511
5512
5513
5514

        Examples
        --------

5515
        The following example uses PyTorch backend.
5516

5517
5518
        >>> import dgl
        >>> import torch
5519

5520
        **Homographs or Heterographs with A Single Edge Type**
5521

5522
        >>> g = dgl.graph(([0, 0, 1], [2, 3, 2]))
5523
5524
        >>> g.format()
        {'created': ['coo'], 'not created': ['csr', 'csc']}
5525
        >>> g.create_formats_()
5526
5527
        >>> g.format()
        {'created': ['coo', 'csr', 'csc'], 'not created': []}
5528

5529
        **Heterographs with Multiple Edge Types**
5530
5531

        >>> g = dgl.heterograph({
5532
5533
5534
5535
5536
        ...     ('user', 'plays', 'game'): (torch.tensor([0, 1, 1, 2]),
        ...                                 torch.tensor([0, 0, 1, 1])),
        ...     ('developer', 'develops', 'game'): (torch.tensor([0, 1]),
        ...                                         torch.tensor([0, 1]))
        ...     })
5537
5538
        >>> g.format()
        {'created': ['coo'], 'not created': ['csr', 'csc']}
5539
        >>> g.create_formats_()
5540
5541
        >>> g.format()
        {'created': ['coo', 'csr', 'csc'], 'not created': []}
5542
        """
5543
        return self._graph.create_formats_()
5544

5545
5546
    def astype(self, idtype):
        """Cast this graph to use another ID type.
5547

5548
        Features are copied (shallow copy) to the new graph.
5549
5550
5551

        Parameters
        ----------
5552
5553
        idtype : Data type object.
            New ID type. Can only be int32 or int64.
5554
5555
5556

        Returns
        -------
5557
5558
        DGLHeteroGraph
            Graph in the new ID type.
5559
        """
5560
5561
        if idtype is None:
            return self
5562
        utils.check_valid_idtype(idtype)
5563
5564
5565
5566
5567
5568
        if self.idtype == idtype:
            return self
        bits = 32 if idtype == F.int32 else 64
        ret = copy.copy(self)
        ret._graph = self._graph.asbits(bits)
        return ret
5569

5570
5571
5572
5573
5574
5575
5576
5577
5578
5579
5580
5581
    # TODO: Formats should not be specified, just saving all the materialized formats
    def shared_memory(self, name, formats=('coo', 'csr', 'csc')):
        """Return a copy of this graph in shared memory, without node data or edge data.

        It moves the graph index to shared memory and returns a DGLHeterograph object which
        has the same graph structure, node types and edge types but does not contain node data
        or edge data.

        Parameters
        ----------
        name : str
            The name of the shared memory.
5582
        formats : str or a list of str (optional)
5583
5584
5585
5586
5587
5588
5589
5590
5591
            Desired formats to be materialized.

        Returns
        -------
        HeteroGraph
            The graph in shared memory
        """
        assert len(name) > 0, "The name of shared memory cannot be empty"
        assert len(formats) > 0
5592
5593
        if isinstance(formats, str):
            formats = [formats]
5594
        for fmt in formats:
5595
            assert fmt in ("coo", "csr", "csc"), '{} is not coo, csr or csc'.format(fmt)
5596
5597
5598
        gidx = self._graph.shared_memory(name, self.ntypes, self.etypes, formats)
        return DGLHeteroGraph(gidx, self.ntypes, self.etypes)

5599
    def long(self):
5600
        """Cast the graph to one with idtype int64
5601

5602
5603
        If the graph already has idtype int64, the function directly returns it. Otherwise,
        it returns a cloned graph of idtype int64 with features copied (shallow copy).
5604
5605
5606

        Returns
        -------
5607
5608
        DGLGraph
            The graph of idtype int64.
5609
5610
5611
5612

        Examples
        --------

5613
5614
5615
5616
5617
5618
5619
5620
5621
5622
5623
5624
5625
5626
5627
5628
5629
5630
5631
5632
5633
5634
5635
5636
5637
5638
5639
5640
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Create a graph of idtype int32.

        >>> # (0, 1), (0, 2), (1, 2)
        >>> g = dgl.graph((torch.tensor([0, 0, 1]).int(), torch.tensor([1, 2, 2]).int()))
        >>> g.ndata['feat'] = torch.ones(3, 1)
        >>> g.idtype
        torch.int32

        Cast the graph to one of idtype int64.

        >>> # A cloned graph with an idtype of int64
        >>> g_long = g.long()
        >>> g_long.idtype
        torch.int64
        >>> # The idtype of the original graph does not change.
        >>> g.idtype
        torch.int32
        >>> g_long.edges()
        (tensor([0, 0, 1]), tensor([1, 2, 2]))
        >>> g_long.ndata
        {'feat': tensor([[1.],
                         [1.],
                         [1.]])}
5641
5642
5643
5644

        See Also
        --------
        int
5645
        idtype
5646
        """
5647
        return self.astype(F.int64)
5648
5649

    def int(self):
5650
5651
5652
5653
        """Cast the graph to one with idtype int32

        If the graph already has idtype int32, the function directly returns it. Otherwise,
        it returns a cloned graph of idtype int32 with features copied (shallow copy).
5654
5655
5656

        Returns
        -------
5657
5658
        DGLGraph
            The graph of idtype int32.
5659
5660
5661
5662

        Examples
        --------

5663
5664
5665
5666
5667
5668
5669
5670
5671
5672
5673
5674
5675
5676
5677
5678
5679
5680
5681
5682
5683
5684
5685
5686
5687
5688
5689
5690
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        Create a graph of idtype int64.

        >>> # (0, 1), (0, 2), (1, 2)
        >>> g = dgl.graph((torch.tensor([0, 0, 1]), torch.tensor([1, 2, 2])))
        >>> g.ndata['feat'] = torch.ones(3, 1)
        >>> g.idtype
        torch.int64

        Cast the graph to one of idtype int32.

        >>> # A cloned graph with an idtype of int32
        >>> g_int = g.int()
        >>> g_int.idtype
        torch.int32
        >>> # The idtype of the original graph does not change.
        >>> g.idtype
        torch.int64
        >>> g_int.edges()
        (tensor([0, 0, 1], dtype=torch.int32), tensor([1, 2, 2], dtype=torch.int32))
        >>> g_int.ndata
        {'feat': tensor([[1.],
                         [1.],
                         [1.]])}
5691
5692
5693
5694

        See Also
        --------
        long
5695
        idtype
5696
        """
5697
        return self.astype(F.int32)
5698

5699
5700
5701
5702
5703
5704
5705
5706
5707
5708
5709
5710
5711
5712
5713
5714
5715
5716
5717
5718
5719
5720
5721
5722
5723
5724
5725
5726
5727
5728
5729
5730
5731
5732
5733
5734
5735
5736
5737
5738
5739
5740
5741
5742
5743
5744
5745
5746
5747
5748
5749
5750
5751
5752
5753
5754
5755
5756
5757
5758
5759
5760
5761
5762
5763
5764
5765
5766
5767
5768
5769
5770
5771
5772
5773
5774
5775
5776
5777
5778
5779
5780
5781
5782
5783
5784
5785
5786
5787
5788
5789
5790
5791
5792
5793
5794
5795
5796
5797
5798
5799
5800
5801
5802
5803
5804
    #################################################################
    # DEPRECATED: from the old DGLGraph
    #################################################################

    def from_networkx(self, nx_graph, node_attrs=None, edge_attrs=None):
        """DEPRECATED: please use

            ``dgl.from_networkx(nx_graph, node_attrs, edge_attrs)``

        which will return a new graph created from the networkx graph.
        """
        raise DGLError('DGLGraph.from_networkx is deprecated. Please call the following\n\n'
                       '\t dgl.from_networkx(nx_graph, node_attrs, edge_attrs)\n\n'
                       ', which creates a new DGLGraph from the networkx graph.')

    def from_scipy_sparse_matrix(self, spmat, multigraph=None):
        """DEPRECATED: please use

            ``dgl.from_scipy(spmat)``

        which will return a new graph created from the scipy matrix.
        """
        raise DGLError('DGLGraph.from_scipy_sparse_matrix is deprecated. '
                       'Please call the following\n\n'
                       '\t dgl.from_scipy(spmat)\n\n'
                       ', which creates a new DGLGraph from the scipy matrix.')

    def register_apply_node_func(self, func):
        """Deprecated: please directly call :func:`apply_nodes` with ``func``
        as argument.
        """
        raise DGLError('DGLGraph.register_apply_node_func is deprecated.'
                       ' Please directly call apply_nodes with func as the argument.')

    def register_apply_edge_func(self, func):
        """Deprecated: please directly call :func:`apply_edges` with ``func``
        as argument.
        """
        raise DGLError('DGLGraph.register_apply_edge_func is deprecated.'
                       ' Please directly call apply_edges with func as the argument.')

    def register_message_func(self, func):
        """Deprecated: please directly call :func:`update_all` with ``func``
        as argument.
        """
        raise DGLError('DGLGraph.register_message_func is deprecated.'
                       ' Please directly call update_all with func as the argument.')

    def register_reduce_func(self, func):
        """Deprecated: please directly call :func:`update_all` with ``func``
        as argument.
        """
        raise DGLError('DGLGraph.register_reduce_func is deprecated.'
                       ' Please directly call update_all with func as the argument.')

    def group_apply_edges(self, group_by, func, edges=ALL, etype=None, inplace=False):
        """**DEPRECATED**: The API is removed in 0.5."""
        raise DGLError('DGLGraph.group_apply_edges is removed in 0.5.')

    def send(self, edges, message_func, etype=None):
        """Send messages along the given edges with the same edge type.

        DEPRECATE: please use send_and_recv, update_all.
        """
        raise DGLError('DGLGraph.send is deprecated. As a replacement, use DGLGraph.apply_edges\n'
                       ' API to compute messages as edge data. Then use DGLGraph.send_and_recv\n'
                       ' and set the message function as dgl.function.copy_e to conduct message\n'
                       ' aggregation.')

    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`.

        DEPRECATE: please use send_and_recv, update_all.
        """
        raise DGLError('DGLGraph.recv is deprecated. As a replacement, use DGLGraph.apply_edges\n'
                       ' API to compute messages as edge data. Then use DGLGraph.send_and_recv\n'
                       ' and set the message function as dgl.function.copy_e to conduct message\n'
                       ' aggregation.')

    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.

        DEPRECATE: please use multi_send_and_recv, multi_update_all.
        """
        raise DGLError('DGLGraph.multi_recv is deprecated. As a replacement,\n'
                       ' use DGLGraph.apply_edges API to compute messages as edge data.\n'
                       ' Then use DGLGraph.multi_send_and_recv and set the message function\n'
                       ' as dgl.function.copy_e to conduct message aggregation.')

    def multi_send_and_recv(self, etype_dict, cross_reducer, apply_node_func=None, inplace=False):
        r"""**DEPRECATED**: The API is removed in v0.5."""
        raise DGLError('DGLGraph.multi_pull is removed in v0.5. As a replacement,\n'
                       ' use DGLGraph.edge_subgraph to extract the subgraph first \n'
                       ' and then call DGLGraph.multi_update_all.')

    def multi_pull(self, v, etype_dict, cross_reducer, apply_node_func=None, inplace=False):
        r"""**DEPRECATED**: The API is removed in v0.5."""
        raise DGLError('DGLGraph.multi_pull is removed in v0.5. As a replacement,\n'
                       ' use DGLGraph.edge_subgraph to extract the subgraph first \n'
                       ' and then call DGLGraph.multi_update_all.')

    def readonly(self, readonly_state=True):
        """Deprecated: DGLGraph will always be mutable."""
        dgl_warning('DGLGraph.readonly is deprecated in v0.5.\n'
                    'DGLGraph now always supports mutable operations like add_nodes'
                    ' and add_edges.')
5805

Minjie Wang's avatar
Minjie Wang committed
5806
5807
5808
5809
5810
5811
5812
5813
5814
5815
5816
5817
5818
5819
5820
5821
5822
5823
5824
5825
5826
5827
5828
5829
5830
5831
5832
5833
5834
############################################################
# 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()))
5835
5836
5837
    if (len(etypes) == 1 and len(ntypes) == 1):
        return [(ntypes[0], etypes[0], ntypes[0])]
    src, dst, eid = metagraph.edges(order="eid")
Minjie Wang's avatar
Minjie Wang committed
5838
5839
5840
    rst = [(ntypes[sid], etypes[eid], ntypes[did]) for sid, did, eid in zip(src, dst, eid)]
    return rst

5841
5842
5843
5844
5845
5846
5847
5848
5849
5850
5851
5852
5853
5854
5855
5856
5857
5858
5859
5860
5861
5862
5863
5864
5865
5866
5867
5868
5869
5870
5871
5872
5873
5874
5875
5876
5877
5878
5879
5880
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.
    """
5881
5882
    ret = _CAPI_DGLFindSrcDstNtypes(metagraph)
    if ret is None:
5883
        return None
5884
5885
    else:
        src, dst = ret
5886
5887
        srctypes = {ntypes[tid] : tid for tid in src}
        dsttypes = {ntypes[tid] : tid for tid in dst}
5888
        return srctypes, dsttypes
5889

Minjie Wang's avatar
Minjie Wang committed
5890
5891
5892
5893
5894
5895
5896
5897
5898
5899
5900
5901
5902
5903
5904
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))

5905
5906
def reduce_dict_data(frames, reducer, order=None):
    """Merge tensor dictionaries into one. Resolve conflict fields using reducer.
Minjie Wang's avatar
Minjie Wang committed
5907
5908
5909

    Parameters
    ----------
5910
5911
    frames : list[dict[str, Tensor]]
        Input tensor dictionaries
5912
5913
5914
5915
5916
    reducer : str or callable function
        One of "sum", "max", "min", "mean", "stack" or a callable function.
        If a callable function is provided, the input arguments must be a single list
        of tensors containing aggregation results from each edge type, and the
        output of function must be a single tensor.
5917
5918
5919
5920
5921
5922
    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.
Minjie Wang's avatar
Minjie Wang committed
5923
5924
5925

    Returns
    -------
5926
    dict[str, Tensor]
Minjie Wang's avatar
Minjie Wang committed
5927
        Merged frame
Da Zheng's avatar
Da Zheng committed
5928
    """
5929
5930
5931
    if len(frames) == 1 and reducer != 'stack':
        # Directly return the only one input. Stack reducer requires
        # modifying tensor shape.
Minjie Wang's avatar
Minjie Wang committed
5932
        return frames[0]
5933
5934
5935
    if callable(reducer):
        merger = reducer
    elif reducer == 'stack':
5936
5937
5938
5939
5940
        # 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]
Minjie Wang's avatar
Minjie Wang committed
5941
5942
5943
5944
5945
5946
5947
5948
        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):
5949
            return redfn(F.stack(flist, 0), 0) if len(flist) > 1 else flist[0]
Minjie Wang's avatar
Minjie Wang committed
5950
5951
5952
    keys = set()
    for frm in frames:
        keys.update(frm.keys())
5953
    ret = {}
Minjie Wang's avatar
Minjie Wang committed
5954
5955
5956
5957
5958
    for k in keys:
        flist = []
        for frm in frames:
            if k in frm:
                flist.append(frm[k])
5959
        ret[k] = merger(flist)
Minjie Wang's avatar
Minjie Wang committed
5960
5961
    return ret

5962
def combine_frames(frames, ids, col_names=None):
Minjie Wang's avatar
Minjie Wang committed
5963
5964
5965
5966
    """Merge the frames into one frame, taking the common columns.

    Return None if there is no common columns.

Da Zheng's avatar
Da Zheng committed
5967
5968
    Parameters
    ----------
5969
    frames : List[Frame]
Minjie Wang's avatar
Minjie Wang committed
5970
5971
5972
        List of frames
    ids : List[int]
        List of frame IDs
5973
5974
    col_names : List[str], optional
        Column names to consider. If not given, it considers all columns.
Minjie Wang's avatar
Minjie Wang committed
5975
5976
5977

    Returns
    -------
5978
    Frame
Minjie Wang's avatar
Minjie Wang committed
5979
        The resulting frame
Da Zheng's avatar
Da Zheng committed
5980
    """
Minjie Wang's avatar
Minjie Wang committed
5981
    # find common columns and check if their schemes match
5982
5983
5984
5985
    if col_names is None:
        schemes = {key: scheme for key, scheme in frames[ids[0]].schemes.items()}
    else:
        schemes = {key: frames[ids[0]].schemes[key] for key in col_names}
Minjie Wang's avatar
Minjie Wang committed
5986
5987
    for frame_id in ids:
        frame = frames[frame_id]
5988
5989
5990
5991
5992
5993
5994
5995
        if frame.num_rows != 0:
            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]
Minjie Wang's avatar
Minjie Wang committed
5996
5997
5998
5999
6000
6001
6002

    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}
6003
    return Frame(cols)
Minjie Wang's avatar
Minjie Wang committed
6004
6005
6006
6007
6008
6009
6010
6011
6012
6013
6014
6015
6016
6017
6018
6019
6020
6021
6022
6023
6024

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)

6025
6026
6027
6028
6029
6030
6031
6032
6033
6034
6035
6036
6037
6038
6039
6040
6041
6042
6043
6044
6045
6046
6047
6048
6049
6050
class DGLBlock(DGLHeteroGraph):
    """Subclass that signifies the graph is a block created from
    :func:`dgl.to_block`.
    """
    # (BarclayII) I'm making a subclass because I don't want to make another version of
    # serialization that contains the is_block flag.
    is_block = True

    def __repr__(self):
        if len(self.srctypes) == 1 and len(self.dsttypes) == 1 and len(self.etypes) == 1:
            ret = 'Block(num_src_nodes={srcnode}, num_dst_nodes={dstnode}, num_edges={edge})'
            return ret.format(
                srcnode=self.number_of_src_nodes(),
                dstnode=self.number_of_dst_nodes(),
                edge=self.number_of_edges())
        else:
            ret = ('Block(num_src_nodes={srcnode},\n'
                   '      num_dst_nodes={dstnode},\n'
                   '      num_edges={edge},\n'
                   '      metagraph={meta})')
            nsrcnode_dict = {ntype : self.number_of_src_nodes(ntype)
                             for ntype in self.srctypes}
            ndstnode_dict = {ntype : self.number_of_dst_nodes(ntype)
                             for ntype in self.dsttypes}
            nedge_dict = {etype : self.number_of_edges(etype)
                          for etype in self.canonical_etypes}
6051
            meta = str(self.metagraph().edges(keys=True))
6052
6053
6054
            return ret.format(
                srcnode=nsrcnode_dict, dstnode=ndstnode_dict, edge=nedge_dict, meta=meta)

6055
6056
6057
6058
6059
6060
6061
6062
6063
6064
6065
6066
6067
6068
6069
6070
6071
6072
6073
6074
6075
6076
6077
6078
6079
6080
6081
6082
6083
6084
6085
6086
6087
6088
6089
6090
6091
6092
6093
6094
6095
6096
6097
6098
6099
6100
6101
6102
6103
6104
6105
6106
6107
6108
6109
6110
6111
6112
6113
6114
6115
6116
6117

def _create_compute_graph(graph, u, v, eid, recv_nodes=None):
    """Create a computation graph from the given edges.

    The compute graph is a uni-directional bipartite graph with only
    one edge type. Similar to subgraph extraction, it stores the original node IDs
    in the srcdata[NID] and dstdata[NID] and extracts features accordingly.
    Edges are not relabeled.

    This function is typically used during message passing to generate
    a graph that contains only the active set of edges.

    Parameters
    ----------
    graph : DGLGraph
        The input graph.
    u : Tensor
        Src nodes.
    v : Tensor
        Dst nodes.
    eid : Tensor
        Edge IDs.
    recv_nodes : Tensor
        Nodes that receive messages. If None, it is equal to unique(v).
        Otherwise, it must be a superset of v and can contain nodes
        that have no incoming edges.

    Returns
    -------
    DGLGraph
        A computation graph.
    """
    if len(u) == 0:
        # The computation graph has no edge and will not trigger message
        # passing. However, because of the apply node phase, we still construct
        # an empty graph to continue.
        unique_src = new_u = new_v = u
        assert recv_nodes is not None
        unique_dst, _ = utils.relabel(recv_nodes)
    else:
        # relabel u and v to starting from 0
        unique_src, src_map = utils.relabel(u)
        if recv_nodes is None:
            unique_dst, dst_map = utils.relabel(v)
        else:
            unique_dst, dst_map = utils.relabel(recv_nodes)
        new_u = F.gather_row(src_map, u)
        new_v = F.gather_row(dst_map, v)

    srctype, etype, dsttype = graph.canonical_etypes[0]
    # create graph
    hgidx = heterograph_index.create_unitgraph_from_coo(
        2, len(unique_src), len(unique_dst), new_u, new_v, ['coo', 'csr', 'csc'])
    # create frame
    srcframe = graph._node_frames[graph.get_ntype_id(srctype)].subframe(unique_src)
    srcframe[NID] = unique_src
    dstframe = graph._node_frames[graph.get_ntype_id(dsttype)].subframe(unique_dst)
    dstframe[NID] = unique_dst
    eframe = graph._edge_frames[0].subframe(eid)
    eframe[EID] = eid

    return DGLHeteroGraph(hgidx, ([srctype], [dsttype]), [etype],
                          node_frames=[srcframe, dstframe],
6118
                          edge_frames=[eframe]), unique_src, unique_dst, eid
6119

6120
_init_api("dgl.heterograph")