Function _locate_eids_to_exclude

python/dgl/sampling/utils.py:12–23

Find the edges whose IDs in parent graph appeared in exclude_eids. Note that both arguments are numpy arrays or numpy dicts.

(frontier_parent_eids, exclude_eids)

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):