chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,409 @@
|
||||
"""Data utilities."""
|
||||
|
||||
from collections import namedtuple
|
||||
|
||||
import networkx as nx
|
||||
import scipy as sp
|
||||
|
||||
from .. import backend as F
|
||||
from ..base import DGLError
|
||||
from . import checks
|
||||
|
||||
|
||||
def elist2tensor(elist, idtype):
|
||||
"""Function to convert an edge list to edge tensors.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
elist : iterable of int pairs
|
||||
List of (src, dst) node ID pairs.
|
||||
idtype : int32, int64, optional
|
||||
Integer ID type. Must be int32 or int64.
|
||||
|
||||
Returns
|
||||
-------
|
||||
(Tensor, Tensor)
|
||||
Edge tensors.
|
||||
"""
|
||||
if len(elist) == 0:
|
||||
u, v = [], []
|
||||
else:
|
||||
u, v = zip(*elist)
|
||||
u = list(u)
|
||||
v = list(v)
|
||||
return F.tensor(u, idtype), F.tensor(v, idtype)
|
||||
|
||||
|
||||
def scipy2tensor(spmat, idtype):
|
||||
"""Function to convert a scipy matrix to a sparse adjacency matrix tuple.
|
||||
|
||||
Note that the data array of the scipy matrix is discarded.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
spmat : scipy.sparse.spmatrix
|
||||
SciPy sparse matrix.
|
||||
idtype : int32, int64, optional
|
||||
Integer ID type. Must be int32 or int64.
|
||||
|
||||
Returns
|
||||
-------
|
||||
(str, tuple[Tensor])
|
||||
A tuple containing the format as well as the list of tensors representing
|
||||
the sparse matrix.
|
||||
"""
|
||||
if spmat.format in ["csr", "csc"]:
|
||||
indptr = F.tensor(spmat.indptr, idtype)
|
||||
indices = F.tensor(spmat.indices, idtype)
|
||||
data = F.tensor([], idtype)
|
||||
return SparseAdjTuple(spmat.format, (indptr, indices, data))
|
||||
else:
|
||||
spmat = spmat.tocoo()
|
||||
row = F.tensor(spmat.row, idtype)
|
||||
col = F.tensor(spmat.col, idtype)
|
||||
return SparseAdjTuple("coo", (row, col))
|
||||
|
||||
|
||||
def networkx2tensor(nx_graph, idtype, edge_id_attr_name=None):
|
||||
"""Function to convert a networkx graph to edge tensors.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
nx_graph : nx.Graph
|
||||
NetworkX graph.
|
||||
idtype : int32, int64, optional
|
||||
Integer ID type. Must be int32 or int64.
|
||||
edge_id_attr_name : str, optional
|
||||
Key name for edge ids in the NetworkX graph. If not found, we
|
||||
will consider the graph not to have pre-specified edge ids. (Default: None)
|
||||
|
||||
Returns
|
||||
-------
|
||||
(Tensor, Tensor)
|
||||
Edge tensors.
|
||||
"""
|
||||
if not nx_graph.is_directed():
|
||||
nx_graph = nx_graph.to_directed()
|
||||
|
||||
# Relabel nodes using consecutive integers
|
||||
nx_graph = nx.convert_node_labels_to_integers(nx_graph, ordering="sorted")
|
||||
has_edge_id = edge_id_attr_name is not None
|
||||
|
||||
if has_edge_id:
|
||||
num_edges = nx_graph.number_of_edges()
|
||||
src = [0] * num_edges
|
||||
dst = [0] * num_edges
|
||||
for u, v, attr in nx_graph.edges(data=True):
|
||||
eid = int(attr[edge_id_attr_name])
|
||||
if eid < 0 or eid >= nx_graph.number_of_edges():
|
||||
raise DGLError(
|
||||
"Expect edge IDs to be a non-negative integer smaller than {:d}, "
|
||||
"got {:d}".format(num_edges, eid)
|
||||
)
|
||||
src[eid] = u
|
||||
dst[eid] = v
|
||||
else:
|
||||
src = []
|
||||
dst = []
|
||||
for e in nx_graph.edges:
|
||||
src.append(e[0])
|
||||
dst.append(e[1])
|
||||
src = F.tensor(src, idtype)
|
||||
dst = F.tensor(dst, idtype)
|
||||
return src, dst
|
||||
|
||||
|
||||
SparseAdjTuple = namedtuple("SparseAdjTuple", ["format", "arrays"])
|
||||
|
||||
|
||||
def graphdata2tensors(
|
||||
data, idtype=None, bipartite=False, infer_node_count=True, **kwargs
|
||||
):
|
||||
"""Function to convert various types of data to edge tensors and infer
|
||||
the number of nodes.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : graph data
|
||||
Various kinds of graph data. Possible data types are:
|
||||
|
||||
- ``(row, col)``
|
||||
- ``('coo', (row, col))``
|
||||
- ``('csr', (indptr, indices, edge_ids))``
|
||||
- ``('csc', (indptr, indices, edge_ids))``
|
||||
- SciPy sparse matrix
|
||||
- NetworkX graph
|
||||
idtype : int32, int64, optional
|
||||
Integer ID type. If None, try infer from the data and if fail use
|
||||
int64.
|
||||
bipartite : bool, optional
|
||||
Whether infer number of nodes of a bipartite graph --
|
||||
num_src and num_dst can be different.
|
||||
infer_node_count : bool, optional
|
||||
Whether infer number of nodes at all. If False, num_src and num_dst
|
||||
are returned as None.
|
||||
kwargs
|
||||
|
||||
- edge_id_attr_name : The name (str) of the edge attribute that stores the edge
|
||||
IDs in the NetworkX graph.
|
||||
- top_map : The dictionary mapping the original IDs of the source nodes to the
|
||||
new ones.
|
||||
- bottom_map : The dictionary mapping the original IDs of the destination nodes
|
||||
to the new ones.
|
||||
|
||||
Returns
|
||||
-------
|
||||
data : SparseAdjTuple
|
||||
A tuple with the sparse matrix format and the adjacency matrix tensors.
|
||||
num_src : int
|
||||
Number of source nodes.
|
||||
num_dst : int
|
||||
Number of destination nodes.
|
||||
"""
|
||||
# Convert tuple to SparseAdjTuple
|
||||
if isinstance(data, tuple):
|
||||
if not isinstance(data[0], str):
|
||||
# (row, col) format, convert to ('coo', (row, col))
|
||||
data = ("coo", data)
|
||||
data = SparseAdjTuple(*data)
|
||||
|
||||
if idtype is None and not (
|
||||
isinstance(data, SparseAdjTuple) and F.is_tensor(data.arrays[0])
|
||||
):
|
||||
# preferred default idtype is int64
|
||||
# if data is tensor and idtype is None, infer the idtype from tensor
|
||||
idtype = F.int64
|
||||
checks.check_valid_idtype(idtype)
|
||||
|
||||
if isinstance(data, SparseAdjTuple) and (
|
||||
not all(F.is_tensor(a) for a in data.arrays)
|
||||
):
|
||||
# (Iterable, Iterable) type data, convert it to (Tensor, Tensor)
|
||||
if len(data.arrays[0]) == 0:
|
||||
# force idtype for empty list
|
||||
data = SparseAdjTuple(
|
||||
data.format, tuple(F.tensor(a, idtype) for a in data.arrays)
|
||||
)
|
||||
else:
|
||||
# convert the iterable to tensor and keep its native data type so we can check
|
||||
# its validity later
|
||||
data = SparseAdjTuple(
|
||||
data.format, tuple(F.tensor(a) for a in data.arrays)
|
||||
)
|
||||
|
||||
num_src, num_dst = None, None
|
||||
if isinstance(data, SparseAdjTuple):
|
||||
if idtype is not None:
|
||||
data = SparseAdjTuple(
|
||||
data.format, tuple(F.astype(a, idtype) for a in data.arrays)
|
||||
)
|
||||
if infer_node_count:
|
||||
num_src, num_dst = infer_num_nodes(data, bipartite=bipartite)
|
||||
elif isinstance(data, list):
|
||||
src, dst = elist2tensor(data, idtype)
|
||||
data = SparseAdjTuple("coo", (src, dst))
|
||||
if infer_node_count:
|
||||
num_src, num_dst = infer_num_nodes(data, bipartite=bipartite)
|
||||
elif isinstance(data, sp.sparse.spmatrix):
|
||||
# We can get scipy matrix's number of rows and columns easily.
|
||||
if infer_node_count:
|
||||
num_src, num_dst = infer_num_nodes(data, bipartite=bipartite)
|
||||
data = scipy2tensor(data, idtype)
|
||||
elif isinstance(data, nx.Graph):
|
||||
# We can get networkx graph's number of sources and destinations easily.
|
||||
if infer_node_count:
|
||||
num_src, num_dst = infer_num_nodes(data, bipartite=bipartite)
|
||||
edge_id_attr_name = kwargs.get("edge_id_attr_name", None)
|
||||
if bipartite:
|
||||
top_map = kwargs.get("top_map")
|
||||
bottom_map = kwargs.get("bottom_map")
|
||||
src, dst = networkxbipartite2tensors(
|
||||
data,
|
||||
idtype,
|
||||
top_map=top_map,
|
||||
bottom_map=bottom_map,
|
||||
edge_id_attr_name=edge_id_attr_name,
|
||||
)
|
||||
else:
|
||||
src, dst = networkx2tensor(
|
||||
data, idtype, edge_id_attr_name=edge_id_attr_name
|
||||
)
|
||||
data = SparseAdjTuple("coo", (src, dst))
|
||||
else:
|
||||
raise DGLError("Unsupported graph data type:", type(data))
|
||||
|
||||
return data, num_src, num_dst
|
||||
|
||||
|
||||
def networkxbipartite2tensors(
|
||||
nx_graph, idtype, top_map, bottom_map, edge_id_attr_name=None
|
||||
):
|
||||
"""Function to convert a networkx bipartite to edge tensors.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
nx_graph : nx.Graph
|
||||
NetworkX graph. It must follow the bipartite graph convention of networkx.
|
||||
Each node has an attribute ``bipartite`` with values 0 and 1 indicating
|
||||
which set it belongs to.
|
||||
top_map : dict
|
||||
The dictionary mapping the original node labels to the node IDs for the source type.
|
||||
bottom_map : dict
|
||||
The dictionary mapping the original node labels to the node IDs for the destination type.
|
||||
idtype : int32, int64, optional
|
||||
Integer ID type. Must be int32 or int64.
|
||||
edge_id_attr_name : str, optional
|
||||
Key name for edge ids in the NetworkX graph. If not found, we
|
||||
will consider the graph not to have pre-specified edge ids. (Default: None)
|
||||
|
||||
Returns
|
||||
-------
|
||||
(Tensor, Tensor)
|
||||
Edge tensors.
|
||||
"""
|
||||
has_edge_id = edge_id_attr_name is not None
|
||||
|
||||
if has_edge_id:
|
||||
num_edges = nx_graph.number_of_edges()
|
||||
src = [0] * num_edges
|
||||
dst = [0] * num_edges
|
||||
for u, v, attr in nx_graph.edges(data=True):
|
||||
if u not in top_map:
|
||||
raise DGLError(
|
||||
"Expect the node {} to have attribute bipartite=0 "
|
||||
"with edge {}".format(u, (u, v))
|
||||
)
|
||||
if v not in bottom_map:
|
||||
raise DGLError(
|
||||
"Expect the node {} to have attribute bipartite=1 "
|
||||
"with edge {}".format(v, (u, v))
|
||||
)
|
||||
eid = int(attr[edge_id_attr_name])
|
||||
if eid < 0 or eid >= nx_graph.number_of_edges():
|
||||
raise DGLError(
|
||||
"Expect edge IDs to be a non-negative integer smaller than {:d}, "
|
||||
"got {:d}".format(num_edges, eid)
|
||||
)
|
||||
src[eid] = top_map[u]
|
||||
dst[eid] = bottom_map[v]
|
||||
else:
|
||||
src = []
|
||||
dst = []
|
||||
for e in nx_graph.edges:
|
||||
u, v = e[0], e[1]
|
||||
if u not in top_map:
|
||||
raise DGLError(
|
||||
"Expect the node {} to have attribute bipartite=0 "
|
||||
"with edge {}".format(u, (u, v))
|
||||
)
|
||||
if v not in bottom_map:
|
||||
raise DGLError(
|
||||
"Expect the node {} to have attribute bipartite=1 "
|
||||
"with edge {}".format(v, (u, v))
|
||||
)
|
||||
src.append(top_map[u])
|
||||
dst.append(bottom_map[v])
|
||||
src = F.tensor(src, dtype=idtype)
|
||||
dst = F.tensor(dst, dtype=idtype)
|
||||
return src, dst
|
||||
|
||||
|
||||
def infer_num_nodes(data, bipartite=False):
|
||||
"""Function for inferring the number of nodes.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : graph data
|
||||
Supported types are:
|
||||
|
||||
* SparseTuple ``(sparse_fmt, arrays)`` where ``arrays`` can be either ``(src, dst)`` or
|
||||
``(indptr, indices, data)``.
|
||||
* SciPy matrix.
|
||||
* NetworkX graph.
|
||||
bipartite : bool, optional
|
||||
Whether infer number of nodes of a bipartite graph --
|
||||
num_src and num_dst can be different.
|
||||
|
||||
Returns
|
||||
-------
|
||||
num_src : int
|
||||
Number of source nodes.
|
||||
num_dst : int
|
||||
Number of destination nodes.
|
||||
|
||||
or
|
||||
|
||||
None
|
||||
If the inference failed.
|
||||
"""
|
||||
if isinstance(data, tuple) and len(data) == 2:
|
||||
if not isinstance(data[0], str):
|
||||
raise TypeError(
|
||||
"Expected sparse format as a str, but got %s" % type(data[0])
|
||||
)
|
||||
|
||||
if data[0] == "coo":
|
||||
# ('coo', (src, dst)) format
|
||||
u, v = data[1]
|
||||
nsrc = F.as_scalar(F.max(u, dim=0)) + 1 if len(u) > 0 else 0
|
||||
ndst = F.as_scalar(F.max(v, dim=0)) + 1 if len(v) > 0 else 0
|
||||
elif data[0] == "csr":
|
||||
# ('csr', (indptr, indices, eids)) format
|
||||
indptr, indices, _ = data[1]
|
||||
nsrc = F.shape(indptr)[0] - 1
|
||||
ndst = (
|
||||
F.as_scalar(F.max(indices, dim=0)) + 1
|
||||
if len(indices) > 0
|
||||
else 0
|
||||
)
|
||||
elif data[0] == "csc":
|
||||
# ('csc', (indptr, indices, eids)) format
|
||||
indptr, indices, _ = data[1]
|
||||
ndst = F.shape(indptr)[0] - 1
|
||||
nsrc = (
|
||||
F.as_scalar(F.max(indices, dim=0)) + 1
|
||||
if len(indices) > 0
|
||||
else 0
|
||||
)
|
||||
else:
|
||||
raise ValueError("unknown format %s" % data[0])
|
||||
elif isinstance(data, sp.sparse.spmatrix):
|
||||
nsrc, ndst = data.shape[0], data.shape[1]
|
||||
elif isinstance(data, nx.Graph):
|
||||
if data.number_of_nodes() == 0:
|
||||
nsrc = ndst = 0
|
||||
elif not bipartite:
|
||||
nsrc = ndst = data.number_of_nodes()
|
||||
else:
|
||||
nsrc = len(
|
||||
{n for n, d in data.nodes(data=True) if d["bipartite"] == 0}
|
||||
)
|
||||
ndst = data.number_of_nodes() - nsrc
|
||||
else:
|
||||
return None
|
||||
if not bipartite:
|
||||
nsrc = ndst = max(nsrc, ndst)
|
||||
return nsrc, ndst
|
||||
|
||||
|
||||
def to_device(data, device):
|
||||
"""Transfer the tensor or dictionary of tensors to the given device.
|
||||
|
||||
Nothing will happen if the device of the original tensor is the same as target device.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : Tensor or dict[str, Tensor]
|
||||
The data.
|
||||
device : device
|
||||
The target device.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Tensor or dict[str, Tensor]
|
||||
The output data.
|
||||
"""
|
||||
if isinstance(data, dict):
|
||||
return {k: F.copy_to(v, device) for k, v in data.items()}
|
||||
else:
|
||||
return F.copy_to(data, device)
|
||||
Reference in New Issue
Block a user