项目文件夹

文件
Quan (Andy) Gan a9dabcc769 [Feature] Random Walk for 0.5 (#1209)
* trying to refactor IndexSelect

* partial implementation

* add index select and assign for floats as well

* move to random choice source

* more updates

* fixes

* fixes

* more fixes

* adding python impl

* fixes

* unit test

* lint

* lint x2

* lint x3

* update metapath2vec

* debugging performance

* still debugging for performance

* tuning

* switching to succvec

* redo

* revert non-uniform sampler to use vector

* still not fast

* why does this crash with OpenMP???

* because there was a data race!!!

* add documentations and remove assign op

* lint

* lint x2

* lol what have i done

* lint x3

* fix and disable gpu testing

* bugfix

* generic random walk

* reorg the random walk source code

* Update randomwalks.h

* Update randomwalks_cpu.cc

* rename file

* move internal function to anonymous ns

* reorg & docstrings

* constant restart probability

* docstring fix

* more commit

* random walk with restart, tested

* some fixes

* switch to using NDArray for choice

* massive fix & docstring

* lint x?

* lint x??

* fix

* export symbols

* skip gpu test

* addresses comments

* replaces another VecToIdArray

* add randomwalks.h to include

* replace void * with template
2020-01-31 19:52:08 +08:00

86 行
3.3 KiB
C++

/*!
* Copyright (c) 2019 by Contributors
* \file dgl/samplinig/randomwalks.h
* \brief Random walk functions.
*/
#ifndef DGL_SAMPLING_RANDOMWALKS_H_
#define DGL_SAMPLING_RANDOMWALKS_H_
#include <dgl/base_heterograph.h>
#include <dgl/array.h>
#include <vector>
#include <utility>
namespace dgl {
namespace sampling {
/*!
* \brief Metapath-based random walk.
* \param hg The heterograph.
* \param seeds A 1D array of seed nodes, with the type the source type of the first
* edge type in the metapath.
* \param metapath A 1D array of edge types representing the metapath.
* \param prob A vector of 1D float arrays, indicating the transition probability of
* each edge by edge type. An empty float array assumes uniform transition.
* \return A pair of
* 1. One 2D array of shape (len(seeds), len(metapath) + 1) with node IDs. The
* paths that terminated early are padded with -1.
* 2. One 1D array of shape (len(metapath) + 1) with node type IDs.
*/
std::pair<IdArray, TypeArray> RandomWalk(
const HeteroGraphPtr hg,
const IdArray seeds,
const TypeArray metapath,
const std::vector<FloatArray> &prob);
/*!
* \brief Metapath-based random walk with restart probability.
* \param hg The heterograph.
* \param seeds A 1D array of seed nodes, with the type the source type of the first
* edge type in the metapath.
* \param metapath A 1D array of edge types representing the metapath.
* \param prob A vector of 1D float arrays, indicating the transition probability of
* each edge by edge type. An empty float array assumes uniform transition.
* \param restart_prob Restart probability
* \return A pair of
* 1. One 2D array of shape (len(seeds), len(metapath) + 1) with node IDs. The
* paths that terminated early are padded with -1.
* 2. One 1D array of shape (len(metapath) + 1) with node type IDs.
*/
std::pair<IdArray, TypeArray> RandomWalkWithRestart(
const HeteroGraphPtr hg,
const IdArray seeds,
const TypeArray metapath,
const std::vector<FloatArray> &prob,
double restart_prob);
/*!
* \brief Metapath-based random walk with stepwise restart probability. Useful
* for PinSAGE-like models.
* \param hg The heterograph.
* \param seeds A 1D array of seed nodes, with the type the source type of the first
* edge type in the metapath.
* \param metapath A 1D array of edge types representing the metapath.
* \param prob A vector of 1D float arrays, indicating the transition probability of
* each edge by edge type. An empty float array assumes uniform transition.
* \param restart_prob Restart probability array which has the same number of elements
* as \c metapath, indicating the probability to terminate after transition.
* \return A pair of
* 1. One 2D array of shape (len(seeds), len(metapath) + 1) with node IDs. The
* paths that terminated early are padded with -1.
* 2. One 1D array of shape (len(metapath) + 1) with node type IDs.
*/
std::pair<IdArray, TypeArray> RandomWalkWithStepwiseRestart(
const HeteroGraphPtr hg,
const IdArray seeds,
const TypeArray metapath,
const std::vector<FloatArray> &prob,
FloatArray restart_prob);
}; // namespace sampling
}; // namespace dgl
#endif // DGL_SAMPLING_RANDOMWALKS_H_