chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,267 @@
|
||||
"""Views of DGLGraph."""
|
||||
from __future__ import absolute_import
|
||||
|
||||
from collections import defaultdict, namedtuple
|
||||
from collections.abc import MutableMapping
|
||||
|
||||
from . import backend as F
|
||||
from .base import ALL, DGLError
|
||||
from .frame import LazyFeature
|
||||
|
||||
NodeSpace = namedtuple("NodeSpace", ["data"])
|
||||
EdgeSpace = namedtuple("EdgeSpace", ["data"])
|
||||
|
||||
|
||||
class HeteroNodeView(object):
|
||||
"""A NodeView class to act as G.nodes for a DGLGraph."""
|
||||
|
||||
__slots__ = ["_graph", "_typeid_getter"]
|
||||
|
||||
def __init__(self, graph, typeid_getter):
|
||||
self._graph = graph
|
||||
self._typeid_getter = typeid_getter
|
||||
|
||||
def __getitem__(self, key):
|
||||
if isinstance(key, slice):
|
||||
# slice
|
||||
if not (
|
||||
key.start is None and key.stop is None and key.step is None
|
||||
):
|
||||
raise DGLError('Currently only full slice ":" is supported')
|
||||
nodes = ALL
|
||||
ntype = None
|
||||
elif isinstance(key, tuple):
|
||||
nodes, ntype = key
|
||||
elif key is None or isinstance(key, str):
|
||||
nodes = ALL
|
||||
ntype = key
|
||||
else:
|
||||
nodes = key
|
||||
ntype = None
|
||||
ntid = self._typeid_getter(ntype)
|
||||
return NodeSpace(
|
||||
data=HeteroNodeDataView(self._graph, ntype, ntid, nodes)
|
||||
)
|
||||
|
||||
def __call__(self, ntype=None):
|
||||
"""Return the nodes."""
|
||||
ntid = self._typeid_getter(ntype)
|
||||
ret = F.arange(
|
||||
0,
|
||||
self._graph._graph.num_nodes(ntid),
|
||||
dtype=self._graph.idtype,
|
||||
ctx=self._graph.device,
|
||||
)
|
||||
return ret
|
||||
|
||||
|
||||
class HeteroNodeDataView(MutableMapping):
|
||||
"""The data view class when G.ndata[ntype] is called."""
|
||||
|
||||
__slots__ = ["_graph", "_ntype", "_ntid", "_nodes"]
|
||||
|
||||
def __init__(self, graph, ntype, ntid, nodes):
|
||||
self._graph = graph
|
||||
self._ntype = ntype
|
||||
self._ntid = ntid
|
||||
self._nodes = nodes
|
||||
|
||||
def __getitem__(self, key):
|
||||
if isinstance(self._ntype, list):
|
||||
ret = {}
|
||||
for (i, ntype) in enumerate(self._ntype):
|
||||
value = self._graph._get_n_repr(self._ntid[i], self._nodes).get(
|
||||
key, None
|
||||
)
|
||||
if value is not None:
|
||||
ret[ntype] = value
|
||||
return ret
|
||||
else:
|
||||
return self._graph._get_n_repr(self._ntid, self._nodes)[key]
|
||||
|
||||
def __setitem__(self, key, val):
|
||||
if isinstance(val, LazyFeature):
|
||||
self._graph._node_frames[self._ntid][key] = val
|
||||
elif isinstance(self._ntype, list):
|
||||
assert isinstance(val, dict), (
|
||||
"Current HeteroNodeDataView has multiple node types, "
|
||||
"please passing the node type and the corresponding data through a dict."
|
||||
)
|
||||
|
||||
for (ntype, data) in val.items():
|
||||
ntid = self._graph.get_ntype_id(ntype)
|
||||
self._graph._set_n_repr(ntid, self._nodes, {key: data})
|
||||
else:
|
||||
assert isinstance(val, dict) is False, (
|
||||
"The HeteroNodeDataView has only one node type. "
|
||||
"please pass a tensor directly"
|
||||
)
|
||||
self._graph._set_n_repr(self._ntid, self._nodes, {key: val})
|
||||
|
||||
def __delitem__(self, key):
|
||||
if isinstance(self._ntype, list):
|
||||
for ntid in self._ntid:
|
||||
if self._graph._get_n_repr(ntid, ALL).get(key, None) is None:
|
||||
continue
|
||||
self._graph._pop_n_repr(ntid, key)
|
||||
else:
|
||||
self._graph._pop_n_repr(self._ntid, key)
|
||||
|
||||
def _transpose(self, as_dict=False):
|
||||
if isinstance(self._ntype, list):
|
||||
ret = defaultdict(dict)
|
||||
for (i, ntype) in enumerate(self._ntype):
|
||||
data = self._graph._get_n_repr(self._ntid[i], self._nodes)
|
||||
for key in self._graph._node_frames[self._ntid[i]]:
|
||||
ret[key][ntype] = data[key]
|
||||
else:
|
||||
ret = self._graph._get_n_repr(self._ntid, self._nodes)
|
||||
if as_dict:
|
||||
ret = {
|
||||
key: ret[key]
|
||||
for key in self._graph._node_frames[self._ntid]
|
||||
}
|
||||
return ret
|
||||
|
||||
def __len__(self):
|
||||
return len(self._transpose())
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self._transpose())
|
||||
|
||||
def keys(self):
|
||||
return self._transpose().keys()
|
||||
|
||||
def values(self):
|
||||
return self._transpose().values()
|
||||
|
||||
def __repr__(self):
|
||||
return repr(self._transpose(as_dict=True))
|
||||
|
||||
|
||||
class HeteroEdgeView(object):
|
||||
"""A EdgeView class to act as G.edges for a DGLGraph."""
|
||||
|
||||
__slots__ = ["_graph"]
|
||||
|
||||
def __init__(self, graph):
|
||||
self._graph = graph
|
||||
|
||||
def __getitem__(self, key):
|
||||
if isinstance(key, slice):
|
||||
# slice
|
||||
if not (
|
||||
key.start is None and key.stop is None and key.step is None
|
||||
):
|
||||
raise DGLError('Currently only full slice ":" is supported')
|
||||
edges = ALL
|
||||
etype = None
|
||||
elif key is None:
|
||||
edges = ALL
|
||||
etype = None
|
||||
elif isinstance(key, tuple):
|
||||
if len(key) == 3:
|
||||
edges = ALL
|
||||
etype = key
|
||||
else:
|
||||
edges = key
|
||||
etype = None
|
||||
elif isinstance(key, str):
|
||||
edges = ALL
|
||||
etype = key
|
||||
else:
|
||||
edges = key
|
||||
etype = None
|
||||
return EdgeSpace(data=HeteroEdgeDataView(self._graph, etype, edges))
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
"""Return all the edges."""
|
||||
return self._graph.all_edges(*args, **kwargs)
|
||||
|
||||
|
||||
class HeteroEdgeDataView(MutableMapping):
|
||||
"""The data view class when G.edata[etype] is called."""
|
||||
|
||||
__slots__ = ["_graph", "_etype", "_etid", "_edges"]
|
||||
|
||||
def __init__(self, graph, etype, edges):
|
||||
self._graph = graph
|
||||
self._etype = etype
|
||||
self._etid = (
|
||||
[self._graph.get_etype_id(t) for t in etype]
|
||||
if isinstance(etype, list)
|
||||
else self._graph.get_etype_id(etype)
|
||||
)
|
||||
self._edges = edges
|
||||
|
||||
def __getitem__(self, key):
|
||||
if isinstance(self._etype, list):
|
||||
ret = {}
|
||||
for (i, etype) in enumerate(self._etype):
|
||||
value = self._graph._get_e_repr(self._etid[i], self._edges).get(
|
||||
key, None
|
||||
)
|
||||
if value is not None:
|
||||
ret[etype] = value
|
||||
return ret
|
||||
else:
|
||||
return self._graph._get_e_repr(self._etid, self._edges)[key]
|
||||
|
||||
def __setitem__(self, key, val):
|
||||
if isinstance(val, LazyFeature):
|
||||
self._graph._edge_frames[self._etid][key] = val
|
||||
elif isinstance(self._etype, list):
|
||||
assert isinstance(val, dict), (
|
||||
"Current HeteroEdgeDataView has multiple edge types, "
|
||||
"please pass the edge type and the corresponding data through a dict."
|
||||
)
|
||||
|
||||
for (etype, data) in val.items():
|
||||
etid = self._graph.get_etype_id(etype)
|
||||
self._graph._set_e_repr(etid, self._edges, {key: data})
|
||||
else:
|
||||
assert isinstance(val, dict) is False, (
|
||||
"The HeteroEdgeDataView has only one edge type. "
|
||||
"please pass a tensor directly"
|
||||
)
|
||||
self._graph._set_e_repr(self._etid, self._edges, {key: val})
|
||||
|
||||
def __delitem__(self, key):
|
||||
if isinstance(self._etype, list):
|
||||
for etid in self._etid:
|
||||
if self._graph._get_e_repr(etid, ALL).get(key, None) is None:
|
||||
continue
|
||||
self._graph._pop_e_repr(etid, key)
|
||||
else:
|
||||
self._graph._pop_e_repr(self._etid, key)
|
||||
|
||||
def _transpose(self, as_dict=False):
|
||||
if isinstance(self._etype, list):
|
||||
ret = defaultdict(dict)
|
||||
for (i, etype) in enumerate(self._etype):
|
||||
data = self._graph._get_e_repr(self._etid[i], self._edges)
|
||||
for key in self._graph._edge_frames[self._etid[i]]:
|
||||
ret[key][etype] = data[key]
|
||||
else:
|
||||
ret = self._graph._get_e_repr(self._etid, self._edges)
|
||||
if as_dict:
|
||||
ret = {
|
||||
key: ret[key]
|
||||
for key in self._graph._edge_frames[self._etid]
|
||||
}
|
||||
return ret
|
||||
|
||||
def __len__(self):
|
||||
return len(self._transpose())
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self._transpose())
|
||||
|
||||
def keys(self):
|
||||
return self._transpose().keys()
|
||||
|
||||
def values(self):
|
||||
return self._transpose().values()
|
||||
|
||||
def __repr__(self):
|
||||
return repr(self._transpose(as_dict=True))
|
||||
Reference in New Issue
Block a user