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