Skip to content

Commit

Permalink
[Misc] Provide a "return 1s instead of edge IDs" option in scipy adja…
Browse files Browse the repository at this point in the history
…cency matrix (dmlc#730)

* return 1 option in scipy adjacency matrix

* lint

* use dgl warning

* i'm an idiot

* lint x2

* rename
  • Loading branch information
BarclayII authored and yzh119 committed Jul 29, 2019
1 parent 7ad663c commit 86f28d6
Show file tree
Hide file tree
Showing 3 changed files with 26 additions and 14 deletions.
4 changes: 2 additions & 2 deletions python/dgl/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,8 @@ def is_all(arg):
"""Return true if the argument is a special symbol for all nodes or edges."""
return isinstance(arg, str) and arg == ALL

def dgl_warning(msg):
def dgl_warning(msg, warn_type=UserWarning):
"""Print out warning messages."""
warnings.warn(msg)
warnings.warn(msg, warn_type)

_init_internal_api()
10 changes: 6 additions & 4 deletions python/dgl/graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -3040,7 +3040,7 @@ def edge_subgraph(self, edges, preserve_nodes=False):
sgi = self._graph.edge_subgraph(induced_edges, preserve_nodes=preserve_nodes)
return subgraph.DGLSubGraph(self, sgi)

def adjacency_matrix_scipy(self, transpose=False, fmt='csr'):
def adjacency_matrix_scipy(self, transpose=False, fmt='csr', return_edge_ids=None):
"""Return the scipy adjacency matrix representation of this graph.
By default, a row of returned adjacency matrix represents the destination
Expand All @@ -3049,22 +3049,24 @@ def adjacency_matrix_scipy(self, transpose=False, fmt='csr'):
When transpose is True, a row represents the source and a column represents
a destination.
The elements in the adajency matrix are edge ids.
Parameters
----------
transpose : bool, optional (default=False)
A flag to transpose the returned adjacency matrix.
fmt : str, optional (default='csr')
Indicates the format of returned adjacency matrix.
return_edge_ids : bool, optional (default=True)
If True, the elements in the adjacency matrix are edge ids.
Note that one of the element is 0. Proceed with caution.
If False, the elements will be always 1.
Returns
-------
scipy.sparse.spmatrix
The scipy representation of adjacency matrix.
"""
return self._graph.adjacency_matrix_scipy(transpose, fmt)
return self._graph.adjacency_matrix_scipy(transpose, fmt, return_edge_ids)

def adjacency_matrix(self, transpose=False, ctx=F.cpu()):
"""Return the adjacency matrix representation of this graph.
Expand Down
26 changes: 18 additions & 8 deletions python/dgl/graph_index.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@

from ._ffi.object import register_object, ObjectBase
from ._ffi.function import _init_api
from .base import DGLError
from .base import DGLError, dgl_warning
from . import backend as F
from . import utils

Expand Down Expand Up @@ -577,7 +577,7 @@ def edge_subgraph(self, e, preserve_nodes=False):
return SubgraphIndex(gidx, self, induced_nodes, e)

@utils.cached_member(cache='_cache', prefix='scipy_adj')
def adjacency_matrix_scipy(self, transpose, fmt):
def adjacency_matrix_scipy(self, 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
Expand All @@ -586,14 +586,14 @@ def adjacency_matrix_scipy(self, transpose, fmt):
When transpose is True, a row represents the source and a column represents
a destination.
The elements in the adajency matrix are edge ids.
Parameters
----------
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
-------
Expand All @@ -603,20 +603,30 @@ def adjacency_matrix_scipy(self, transpose, fmt):
if not isinstance(transpose, bool):
raise DGLError('Expect bool value for "transpose" arg,'
' but got %s.' % (type(transpose)))

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

rst = _CAPI_DGLGraphGetAdj(self, transpose, fmt)
if fmt == "csr":
indptr = utils.toindex(rst(0)).tonumpy()
indices = utils.toindex(rst(1)).tonumpy()
shuffle = utils.toindex(rst(2)).tonumpy()
data = utils.toindex(rst(2)).tonumpy() if return_edge_ids else np.ones_like(indices)
n = self.number_of_nodes()
return scipy.sparse.csr_matrix((shuffle, indices, indptr), shape=(n, n))
return scipy.sparse.csr_matrix((data, indices, indptr), shape=(n, n))
elif fmt == 'coo':
idx = utils.toindex(rst(0)).tonumpy()
n = self.number_of_nodes()
m = self.number_of_edges()
row, col = np.reshape(idx, (2, m))
shuffle = np.arange(0, m)
return scipy.sparse.coo_matrix((shuffle, (row, col)), shape=(n, n))
data = np.arange(0, m) if return_edge_ids else np.ones_like(row)
return scipy.sparse.coo_matrix((data, (row, col)), shape=(n, n))
else:
raise Exception("unknown format")

Expand Down

0 comments on commit 86f28d6

Please sign in to comment.