wehub-resource-sync
106 lines
4.2 KiB
Python
106 lines
4.2 KiB
Python
"""Sampling utilities"""
|
|
from collections.abc import Mapping
|
|
|
|
import numpy as np
|
|
|
|
from .. import backend as F, transforms, utils
|
|
from ..base import EID
|
|
|
|
from ..utils import recursive_apply, recursive_apply_pair
|
|
|
|
|
|
def _locate_eids_to_exclude(frontier_parent_eids, exclude_eids):
|
|
"""Find the edges whose IDs in parent graph appeared in exclude_eids.
|
|
|
|
Note that both arguments are numpy arrays or numpy dicts.
|
|
"""
|
|
if not isinstance(frontier_parent_eids, Mapping):
|
|
return np.isin(frontier_parent_eids, exclude_eids).nonzero()[0]
|
|
result = {}
|
|
for k, v in frontier_parent_eids.items():
|
|
if k in exclude_eids:
|
|
result[k] = np.isin(v, exclude_eids[k]).nonzero()[0]
|
|
return recursive_apply(result, F.zerocopy_from_numpy)
|
|
|
|
|
|
class EidExcluder(object):
|
|
"""Class that finds the edges whose IDs in parent graph appeared in exclude_eids.
|
|
|
|
The edge IDs can be both CPU and GPU tensors.
|
|
"""
|
|
|
|
def __init__(self, exclude_eids):
|
|
device = None
|
|
if isinstance(exclude_eids, Mapping):
|
|
for _, v in exclude_eids.items():
|
|
if device is None:
|
|
device = F.context(v)
|
|
break
|
|
else:
|
|
device = F.context(exclude_eids)
|
|
self._exclude_eids = None
|
|
self._filter = None
|
|
|
|
if device == F.cpu():
|
|
# TODO(nv-dlasalle): Once Filter is implemented for the CPU, we
|
|
# should just use that irregardless of the device.
|
|
self._exclude_eids = (
|
|
recursive_apply(exclude_eids, F.zerocopy_to_numpy)
|
|
if exclude_eids is not None
|
|
else None
|
|
)
|
|
else:
|
|
self._filter = recursive_apply(exclude_eids, utils.Filter)
|
|
|
|
def _find_indices(self, parent_eids):
|
|
"""Find the set of edge indices to remove."""
|
|
if self._exclude_eids is not None:
|
|
parent_eids_np = recursive_apply(parent_eids, F.zerocopy_to_numpy)
|
|
return _locate_eids_to_exclude(parent_eids_np, self._exclude_eids)
|
|
else:
|
|
assert self._filter is not None
|
|
func = lambda x, y: x.find_included_indices(y)
|
|
return recursive_apply_pair(self._filter, parent_eids, func)
|
|
|
|
def __call__(self, frontier, weights=None):
|
|
parent_eids = frontier.edata[EID]
|
|
located_eids = self._find_indices(parent_eids)
|
|
|
|
if not isinstance(located_eids, Mapping):
|
|
# (BarclayII) If frontier already has a EID field and located_eids is empty,
|
|
# the returned graph will keep EID intact. Otherwise, EID will change
|
|
# to the mapping from the new graph to the old frontier.
|
|
# So we need to test if located_eids is empty, and do the remapping ourselves.
|
|
if len(located_eids) > 0:
|
|
frontier = transforms.remove_edges(
|
|
frontier, located_eids, store_ids=True
|
|
)
|
|
if (
|
|
weights is not None
|
|
and weights[0].shape[0] == frontier.num_edges()
|
|
):
|
|
weights[0] = F.gather_row(weights[0], frontier.edata[EID])
|
|
frontier.edata[EID] = F.gather_row(
|
|
parent_eids, frontier.edata[EID]
|
|
)
|
|
else:
|
|
# (BarclayII) remove_edges only accepts removing one type of edges,
|
|
# so I need to keep track of the edge IDs left one by one.
|
|
new_eids = parent_eids.copy()
|
|
for i, (k, v) in enumerate(located_eids.items()):
|
|
if len(v) > 0:
|
|
frontier = transforms.remove_edges(
|
|
frontier, v, etype=k, store_ids=True
|
|
)
|
|
new_eids[k] = F.gather_row(
|
|
parent_eids[k], frontier.edges[k].data[EID]
|
|
)
|
|
if weights is not None and weights[i].shape[
|
|
0
|
|
] == frontier.num_edges(k):
|
|
weights[i] = F.gather_row(
|
|
weights[i], frontier.edges[k].data[EID]
|
|
)
|
|
frontier.edata[EID] = new_eids
|
|
return frontier if weights is None else (frontier, weights)
|