subgraph.py 4.81 KB
Newer Older
1
"""Class for subgraph data structure."""
Minjie Wang's avatar
Minjie Wang committed
2
3
from __future__ import absolute_import

Minjie Wang's avatar
Minjie Wang committed
4
5
6
from .frame import Frame, FrameRef
from .graph import DGLGraph
from . import utils
7
from .base import DGLError
8
from .graph_index import map_to_subgraph_nid
Minjie Wang's avatar
Minjie Wang committed
9
10

class DGLSubGraph(DGLGraph):
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
    """The subgraph class.

    There are two subgraph modes: shared and non-shared.

    For the "non-shared" mode, the user needs to explicitly call
    ``copy_from_parent`` to copy node/edge features from its parent graph.
    * If the user tries to get node/edge features before ``copy_from_parent``,
      s/he will get nothing.
    * If the subgraph already has its own node/edge features, ``copy_from_parent``
      will override them.
    * Any update on the subgraph's node/edge features will not be seen
      by the parent graph. As such, the memory consumption is of the order
      of the subgraph size.
    * To write the subgraph's node/edge features back to parent graph. There are two options:
      (1) Use ``copy_to_parent`` API to write node/edge features back.
      (2) [TODO] Use ``dgl.merge`` to merge multiple subgraphs back to one parent.

    The "shared" mode is currently not supported.

30
    The subgraph is read-only on structure; graph mutation is not allowed.
31
32
33
34
35

    Parameters
    ----------
    parent : DGLGraph
        The parent graph
36
37
    sgi : SubgraphIndex
        Internal subgraph data structure.
38
39
40
    shared : bool, optional
        Whether the subgraph shares node/edge features with the parent graph.
    """
41
42
43
    def __init__(self, parent, sgi, shared=False):
        super(DGLSubGraph, self).__init__(graph_data=sgi.graph,
                                          readonly=True)
Minjie Wang's avatar
Minjie Wang committed
44
45
        if shared:
            raise DGLError('Shared mode is not yet supported.')
Da Zheng's avatar
Da Zheng committed
46
        self._parent = parent
47
48
        self._parent_nid = sgi.induced_nodes
        self._parent_eid = sgi.induced_edges
49
50

    # override APIs
Minjie Wang's avatar
Minjie Wang committed
51
    def add_nodes(self, num, data=None):
52
        """Add nodes. Disabled because subgraph is read-only."""
Minjie Wang's avatar
Minjie Wang committed
53
        raise DGLError('Readonly graph. Mutation is not allowed.')
54

Minjie Wang's avatar
Minjie Wang committed
55
    def add_edge(self, u, v, data=None):
56
        """Add one edge. Disabled because subgraph is read-only."""
Minjie Wang's avatar
Minjie Wang committed
57
        raise DGLError('Readonly graph. Mutation is not allowed.')
58

Minjie Wang's avatar
Minjie Wang committed
59
    def add_edges(self, u, v, data=None):
60
        """Add many edges. Disabled because subgraph is read-only."""
Minjie Wang's avatar
Minjie Wang committed
61
        raise DGLError('Readonly graph. Mutation is not allowed.')
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76

    @property
    def parent_nid(self):
        """Get the parent node ids.

        The returned tensor can be used as a map from the node id
        in this subgraph to the node id in the parent graph.

        Returns
        -------
        Tensor
            The parent node id array.
        """
        return self._parent_nid.tousertensor()

Da Zheng's avatar
Da Zheng committed
77
78
79
80
81
82
83
84
85
    def _get_parent_eid(self):
        # The parent eid might be lazily evaluated and thus may not
        # be an index. Instead, it's a lambda function that returns
        # an index.
        if isinstance(self._parent_eid, utils.Index):
            return self._parent_eid
        else:
            return self._parent_eid()

86
87
88
89
90
91
92
93
94
95
96
97
    @property
    def parent_eid(self):
        """Get the parent edge ids.

        The returned tensor can be used as a map from the edge id
        in this subgraph to the edge id in the parent graph.

        Returns
        -------
        Tensor
            The parent edge id array.
        """
Da Zheng's avatar
Da Zheng committed
98
        return self._get_parent_eid().tousertensor()
Minjie Wang's avatar
Minjie Wang committed
99

Da Zheng's avatar
Da Zheng committed
100
    def copy_to_parent(self, inplace=False):
101
102
103
104
105
106
107
108
        """Write node/edge features to the parent graph.

        Parameters
        ----------
        inplace : bool
            If true, use inplace write (no gradient but faster)
        """
        self._parent._node_frame.update_rows(
Minjie Wang's avatar
Minjie Wang committed
109
            self._parent_nid, self._node_frame, inplace=inplace)
Da Zheng's avatar
Da Zheng committed
110
111
        if self._parent._edge_frame.num_rows != 0:
            self._parent._edge_frame.update_rows(
Minjie Wang's avatar
Minjie Wang committed
112
                self._get_parent_eid(), self._edge_frame, inplace=inplace)
Da Zheng's avatar
Da Zheng committed
113
114

    def copy_from_parent(self):
Minjie Wang's avatar
Minjie Wang committed
115
116
117
118
        """Copy node/edge features from the parent graph.

        All old features will be removed.
        """
Da Zheng's avatar
Da Zheng committed
119
        if self._parent._node_frame.num_rows != 0:
120
121
            self._node_frame = FrameRef(Frame(
                self._parent._node_frame[self._parent_nid]))
Da Zheng's avatar
Da Zheng committed
122
        if self._parent._edge_frame.num_rows != 0:
123
            self._edge_frame = FrameRef(Frame(
Da Zheng's avatar
Da Zheng committed
124
                self._parent._edge_frame[self._get_parent_eid()]))
125
126

    def map_to_subgraph_nid(self, parent_vids):
127
128
129
130
131
132
133
134
135
136
137
138
139
        """Map the node Ids in the parent graph to the node Ids in the subgraph.

        Parameters
        ----------
        parent_vids : list, tensor
            The node ID array in the parent graph.

        Returns
        -------
        tensor
            The node ID array in the subgraph.
        """
        return map_to_subgraph_nid(self._graph, utils.toindex(parent_vids)).tousertensor()