neighbor.cc 18.1 KB
Newer Older
1
/*!
2
 *  Copyright (c) 2020-2021 by Contributors
3
4
5
6
7
8
9
 * \file graph/sampling/neighbor.cc
 * \brief Definition of neighborhood-based sampler APIs.
 */

#include <dgl/runtime/container.h>
#include <dgl/packed_func_ext.h>
#include <dgl/array.h>
10
#include <dgl/aten/macro.h>
11
12
13
14
15
16
17
18
19
20
#include <dgl/sampling/neighbor.h>
#include "../../../c_api_common.h"
#include "../../unit_graph.h"

using namespace dgl::runtime;
using namespace dgl::aten;

namespace dgl {
namespace sampling {

21
22
23
24
25
26
27
28
29
30
31
32
33
HeteroSubgraph ExcludeCertainEdges(
    const HeteroSubgraph& sg,
    const std::vector<IdArray>& exclude_edges) {

    HeteroGraphPtr hg_view = HeteroGraphRef(sg.graph).sptr();
    std::vector<IdArray> remain_induced_edges(hg_view->NumEdgeTypes());
    std::vector<IdArray> remain_edges(hg_view->NumEdgeTypes());

    for (dgl_type_t etype = 0; etype < hg_view->NumEdgeTypes(); ++etype) {
      IdArray edge_ids = Range(0,
                               sg.induced_edges[etype]->shape[0],
                               sg.induced_edges[etype]->dtype.bits,
                               sg.induced_edges[etype]->ctx);
34
      if (exclude_edges[etype].GetSize() == 0 || edge_ids.GetSize() == 0) {
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
        remain_edges[etype] = edge_ids;
        remain_induced_edges[etype] = sg.induced_edges[etype];
        continue;
      }
      ATEN_ID_TYPE_SWITCH(hg_view->DataType(), IdType, {
        IdType* idx_data = edge_ids.Ptr<IdType>();
        IdType* induced_edges_data = sg.induced_edges[etype].Ptr<IdType>();
        const IdType exclude_edges_len = exclude_edges[etype]->shape[0];
        std::sort(exclude_edges[etype].Ptr<IdType>(),
                  exclude_edges[etype].Ptr<IdType>() + exclude_edges_len);
        const IdType* exclude_edges_data = exclude_edges[etype].Ptr<IdType>();
        IdType outId = 0;
        for (IdType i = 0; i != sg.induced_edges[etype]->shape[0]; ++i) {
          if (!std::binary_search(exclude_edges_data,
                                  exclude_edges_data + exclude_edges_len,
                                  induced_edges_data[i])) {
            induced_edges_data[outId] = induced_edges_data[i];
            idx_data[outId] = idx_data[i];
            ++outId;
          }
        }
        remain_edges[etype] = aten::IndexSelect(edge_ids, 0, outId);
        remain_induced_edges[etype] = aten::IndexSelect(sg.induced_edges[etype], 0, outId);
      });
    }
    HeteroSubgraph subg = hg_view->EdgeSubgraph(remain_edges, true);
    subg.induced_edges = std::move(remain_induced_edges);
    return subg;
}

65
66
67
68
69
HeteroSubgraph SampleNeighbors(
    const HeteroGraphPtr hg,
    const std::vector<IdArray>& nodes,
    const std::vector<int64_t>& fanouts,
    EdgeDir dir,
70
    const std::vector<NDArray>& prob_or_mask,
71
    const std::vector<IdArray>& exclude_edges,
72
73
74
75
76
77
78
    bool replace) {

  // sanity check
  CHECK_EQ(nodes.size(), hg->NumVertexTypes())
    << "Number of node ID tensors must match the number of node types.";
  CHECK_EQ(fanouts.size(), hg->NumEdgeTypes())
    << "Number of fanout values must match the number of edge types.";
79
  CHECK_EQ(prob_or_mask.size(), hg->NumEdgeTypes())
80
    << "Number of probability tensors must match the number of edge types.";
81

82
  DGLContext ctx = aten::GetContextOf(nodes);
83

84
85
86
87
88
89
90
91
  std::vector<HeteroGraphPtr> subrels(hg->NumEdgeTypes());
  std::vector<IdArray> induced_edges(hg->NumEdgeTypes());
  for (dgl_type_t etype = 0; etype < hg->NumEdgeTypes(); ++etype) {
    auto pair = hg->meta_graph()->FindEdge(etype);
    const dgl_type_t src_vtype = pair.first;
    const dgl_type_t dst_vtype = pair.second;
    const IdArray nodes_ntype = nodes[(dir == EdgeDir::kOut)? src_vtype : dst_vtype];
    const int64_t num_nodes = nodes_ntype->shape[0];
92

93
94
    if (num_nodes == 0 || fanouts[etype] == 0) {
      // Nothing to sample for this etype, create a placeholder relation graph
95
96
97
98
      subrels[etype] = UnitGraph::Empty(
        hg->GetRelationGraph(etype)->NumVertexTypes(),
        hg->NumVertices(src_vtype),
        hg->NumVertices(dst_vtype),
99
100
        hg->DataType(), ctx);
      induced_edges[etype] = aten::NullArray(hg->DataType(), ctx);
101
    } else {
102
      COOMatrix sampled_coo;
103
      // sample from one relation graph
104
      auto req_fmt = (dir == EdgeDir::kOut)? CSR_CODE : CSC_CODE;
105
106
      auto avail_fmt = hg->SelectFormat(etype, req_fmt);
      switch (avail_fmt) {
107
        case SparseFormat::kCOO:
108
109
110
          if (dir == EdgeDir::kIn) {
            sampled_coo = aten::COOTranspose(aten::COORowWiseSampling(
              aten::COOTranspose(hg->GetCOOMatrix(etype)),
111
              nodes_ntype, fanouts[etype], prob_or_mask[etype], replace));
112
113
          } else {
            sampled_coo = aten::COORowWiseSampling(
114
              hg->GetCOOMatrix(etype), nodes_ntype, fanouts[etype], prob_or_mask[etype], replace);
115
116
          }
          break;
117
        case SparseFormat::kCSR:
118
119
          CHECK(dir == EdgeDir::kOut) << "Cannot sample out edges on CSC matrix.";
          sampled_coo = aten::CSRRowWiseSampling(
120
            hg->GetCSRMatrix(etype), nodes_ntype, fanouts[etype], prob_or_mask[etype], replace);
121
          break;
122
        case SparseFormat::kCSC:
123
124
          CHECK(dir == EdgeDir::kIn) << "Cannot sample in edges on CSR matrix.";
          sampled_coo = aten::CSRRowWiseSampling(
125
            hg->GetCSCMatrix(etype), nodes_ntype, fanouts[etype], prob_or_mask[etype], replace);
126
127
128
129
130
          sampled_coo = aten::COOTranspose(sampled_coo);
          break;
        default:
          LOG(FATAL) << "Unsupported sparse format.";
      }
131

132
      subrels[etype] = UnitGraph::CreateFromCOO(
133
134
        hg->GetRelationGraph(etype)->NumVertexTypes(), sampled_coo.num_rows, sampled_coo.num_cols,
        sampled_coo.row, sampled_coo.col);
135
      induced_edges[etype] = sampled_coo.data;
136
137
138
139
    }
  }

  HeteroSubgraph ret;
140
  ret.graph = CreateHeteroGraph(hg->meta_graph(), subrels, hg->NumVerticesPerType());
141
142
  ret.induced_vertices.resize(hg->NumVertexTypes());
  ret.induced_edges = std::move(induced_edges);
143
144
145
  if (!exclude_edges.empty()) {
    return ExcludeCertainEdges(ret, exclude_edges);
  }
146
147
148
  return ret;
}

149
150
151
HeteroSubgraph SampleNeighborsEType(
    const HeteroGraphPtr hg,
    const IdArray nodes,
152
    const std::vector<int64_t>& eid2etype_offset,
153
    const std::vector<int64_t>& fanouts,
154
    EdgeDir dir,
155
    const std::vector<FloatArray>& prob,
156
    bool replace,
157
    bool rowwise_etype_sorted) {
158
159
160
161
162
163
164
165
166
167
168
169

  CHECK_EQ(1, hg->NumVertexTypes())
    << "SampleNeighborsEType only work with homogeneous graph";
  CHECK_EQ(1, hg->NumEdgeTypes())
    << "SampleNeighborsEType only work with homogeneous graph";

  std::vector<HeteroGraphPtr> subrels(1);
  std::vector<IdArray> induced_edges(1);
  const int64_t num_nodes = nodes->shape[0];
  dgl_type_t etype = 0;
  const dgl_type_t src_vtype = 0;
  const dgl_type_t dst_vtype = 0;
170
171
172
173
174
175
176
177
178
179
180

  bool same_fanout = true;
  int64_t fanout_value = fanouts[0];
  for (auto fanout : fanouts) {
    if (fanout != fanout_value) {
      same_fanout = false;
      break;
    }
  }

  if (num_nodes == 0 || (same_fanout && fanout_value == 0)) {
181
182
183
184
185
186
    subrels[etype] = UnitGraph::Empty(1,
      hg->NumVertices(src_vtype),
      hg->NumVertices(dst_vtype),
      hg->DataType(), hg->Context());
    induced_edges[etype] = aten::NullArray();
  } else {
187
    COOMatrix sampled_coo;
188
189
190
191
192
193
194
195
196
    // sample from graph
    // the edge type is stored in etypes
    auto req_fmt = (dir == EdgeDir::kOut)? CSR_CODE : CSC_CODE;
    auto avail_fmt = hg->SelectFormat(etype, req_fmt);
    switch (avail_fmt) {
      case SparseFormat::kCOO:
        if (dir == EdgeDir::kIn) {
          sampled_coo = aten::COOTranspose(aten::COORowWisePerEtypeSampling(
            aten::COOTranspose(hg->GetCOOMatrix(etype)),
197
            nodes, eid2etype_offset, fanouts, prob, replace));
198
199
        } else {
          sampled_coo = aten::COORowWisePerEtypeSampling(
200
            hg->GetCOOMatrix(etype), nodes, eid2etype_offset, fanouts, prob, replace);
201
202
203
204
205
        }
        break;
      case SparseFormat::kCSR:
        CHECK(dir == EdgeDir::kOut) << "Cannot sample out edges on CSC matrix.";
        sampled_coo = aten::CSRRowWisePerEtypeSampling(
206
207
            hg->GetCSRMatrix(etype), nodes, eid2etype_offset,
            fanouts, prob, replace, rowwise_etype_sorted);
208
209
210
211
          break;
      case SparseFormat::kCSC:
        CHECK(dir == EdgeDir::kIn) << "Cannot sample in edges on CSR matrix.";
        sampled_coo = aten::CSRRowWisePerEtypeSampling(
212
213
            hg->GetCSCMatrix(etype), nodes, eid2etype_offset,
            fanouts, prob, replace, rowwise_etype_sorted);
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
        sampled_coo = aten::COOTranspose(sampled_coo);
        break;
      default:
        LOG(FATAL) << "Unsupported sparse format.";
    }

    subrels[etype] = UnitGraph::CreateFromCOO(
      1, sampled_coo.num_rows, sampled_coo.num_cols,
      sampled_coo.row, sampled_coo.col);
    induced_edges[etype] = sampled_coo.data;
  }

  HeteroSubgraph ret;
  ret.graph = CreateHeteroGraph(hg->meta_graph(), subrels, hg->NumVerticesPerType());
  ret.induced_vertices.resize(hg->NumVertexTypes());
  ret.induced_edges = std::move(induced_edges);
  return ret;
}

233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
HeteroSubgraph SampleNeighborsTopk(
    const HeteroGraphPtr hg,
    const std::vector<IdArray>& nodes,
    const std::vector<int64_t>& k,
    EdgeDir dir,
    const std::vector<FloatArray>& weight,
    bool ascending) {
  // sanity check
  CHECK_EQ(nodes.size(), hg->NumVertexTypes())
    << "Number of node ID tensors must match the number of node types.";
  CHECK_EQ(k.size(), hg->NumEdgeTypes())
    << "Number of k values must match the number of edge types.";
  CHECK_EQ(weight.size(), hg->NumEdgeTypes())
    << "Number of weight tensors must match the number of edge types.";

  std::vector<HeteroGraphPtr> subrels(hg->NumEdgeTypes());
  std::vector<IdArray> induced_edges(hg->NumEdgeTypes());
  for (dgl_type_t etype = 0; etype < hg->NumEdgeTypes(); ++etype) {
    auto pair = hg->meta_graph()->FindEdge(etype);
    const dgl_type_t src_vtype = pair.first;
    const dgl_type_t dst_vtype = pair.second;
    const IdArray nodes_ntype = nodes[(dir == EdgeDir::kOut)? src_vtype : dst_vtype];
    const int64_t num_nodes = nodes_ntype->shape[0];
256
257
    if (num_nodes == 0 || k[etype] == 0) {
      // Nothing to sample for this etype, create a placeholder relation graph
258
259
260
261
262
      subrels[etype] = UnitGraph::Empty(
        hg->GetRelationGraph(etype)->NumVertexTypes(),
        hg->NumVertices(src_vtype),
        hg->NumVertices(dst_vtype),
        hg->DataType(), hg->Context());
263
      induced_edges[etype] = aten::NullArray();
264
265
    } else {
      // sample from one relation graph
266
      auto req_fmt = (dir == EdgeDir::kOut)? CSR_CODE : CSC_CODE;
267
268
269
      auto avail_fmt = hg->SelectFormat(etype, req_fmt);
      COOMatrix sampled_coo;
      switch (avail_fmt) {
270
        case SparseFormat::kCOO:
271
272
273
274
275
276
277
278
279
          if (dir == EdgeDir::kIn) {
            sampled_coo = aten::COOTranspose(aten::COORowWiseTopk(
              aten::COOTranspose(hg->GetCOOMatrix(etype)),
              nodes_ntype, k[etype], weight[etype], ascending));
          } else {
            sampled_coo = aten::COORowWiseTopk(
              hg->GetCOOMatrix(etype), nodes_ntype, k[etype], weight[etype], ascending);
          }
          break;
280
        case SparseFormat::kCSR:
281
282
283
284
          CHECK(dir == EdgeDir::kOut) << "Cannot sample out edges on CSC matrix.";
          sampled_coo = aten::CSRRowWiseTopk(
            hg->GetCSRMatrix(etype), nodes_ntype, k[etype], weight[etype], ascending);
          break;
285
        case SparseFormat::kCSC:
286
287
288
289
290
291
292
293
294
          CHECK(dir == EdgeDir::kIn) << "Cannot sample in edges on CSR matrix.";
          sampled_coo = aten::CSRRowWiseTopk(
            hg->GetCSCMatrix(etype), nodes_ntype, k[etype], weight[etype], ascending);
          sampled_coo = aten::COOTranspose(sampled_coo);
          break;
        default:
          LOG(FATAL) << "Unsupported sparse format.";
      }
      subrels[etype] = UnitGraph::CreateFromCOO(
295
296
        hg->GetRelationGraph(etype)->NumVertexTypes(), sampled_coo.num_rows, sampled_coo.num_cols,
        sampled_coo.row, sampled_coo.col);
297
      induced_edges[etype] = sampled_coo.data;
298
299
300
301
    }
  }

  HeteroSubgraph ret;
302
  ret.graph = CreateHeteroGraph(hg->meta_graph(), subrels, hg->NumVerticesPerType());
303
304
305
306
307
  ret.induced_vertices.resize(hg->NumVertexTypes());
  ret.induced_edges = std::move(induced_edges);
  return ret;
}

308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
HeteroSubgraph SampleNeighborsBiased(
    const HeteroGraphPtr hg,
    const IdArray& nodes,
    const int64_t fanout,
    const NDArray& bias,
    const NDArray& tag_offset,
    const EdgeDir dir,
    const bool replace
) {
  CHECK_EQ(hg->NumEdgeTypes(), 1) << "Only homogeneous or bipartite graphs are supported";
  auto pair = hg->meta_graph()->FindEdge(0);
  const dgl_type_t src_vtype = pair.first;
  const dgl_type_t dst_vtype = pair.second;
  const dgl_type_t nodes_ntype = (dir == EdgeDir::kOut) ? src_vtype : dst_vtype;

  // sanity check
  CHECK_EQ(tag_offset->ndim, 2) << "The shape of tag_offset should be [num_nodes, num_tags + 1]";
  CHECK_EQ(tag_offset->shape[0], hg->NumVertices(nodes_ntype))
    << "The shape of tag_offset should be [num_nodes, num_tags + 1]";
  CHECK_EQ(tag_offset->shape[1], bias->shape[0] + 1)
    << "The sizes of tag_offset and bias are inconsistent";

  const int64_t num_nodes = nodes->shape[0];
  HeteroGraphPtr subrel;
  IdArray induced_edges;
  const dgl_type_t etype = 0;
  if (num_nodes == 0 || fanout == 0) {
      // Nothing to sample for this etype, create a placeholder relation graph
      subrel = UnitGraph::Empty(
        hg->GetRelationGraph(etype)->NumVertexTypes(),
        hg->NumVertices(src_vtype),
        hg->NumVertices(dst_vtype),
        hg->DataType(), hg->Context());
      induced_edges = aten::NullArray();
    } else {
      // sample from one relation graph
      const auto req_fmt = (dir == EdgeDir::kOut)? CSR_CODE : CSC_CODE;
      const auto created_fmt = hg->GetCreatedFormats();
      COOMatrix sampled_coo;

      switch (req_fmt) {
        case CSR_CODE:
          CHECK(created_fmt & CSR_CODE) << "A sorted CSR Matrix is required.";
          sampled_coo = aten::CSRRowWiseSamplingBiased(
            hg->GetCSRMatrix(etype), nodes, fanout, tag_offset, bias, replace);
          break;
        case CSC_CODE:
          CHECK(created_fmt & CSC_CODE) << "A sorted CSC Matrix is required.";
          sampled_coo = aten::CSRRowWiseSamplingBiased(
            hg->GetCSCMatrix(etype), nodes, fanout, tag_offset, bias, replace);
          sampled_coo = aten::COOTranspose(sampled_coo);
          break;
        default:
          LOG(FATAL) << "Unsupported sparse format.";
      }
      subrel = UnitGraph::CreateFromCOO(
        hg->GetRelationGraph(etype)->NumVertexTypes(), sampled_coo.num_rows, sampled_coo.num_cols,
        sampled_coo.row, sampled_coo.col);
      induced_edges = sampled_coo.data;
    }

  HeteroSubgraph ret;
  ret.graph = CreateHeteroGraph(hg->meta_graph(), {subrel}, hg->NumVerticesPerType());
  ret.induced_vertices.resize(hg->NumVertexTypes());
  ret.induced_edges = {induced_edges};
  return ret;
}

376
377
378
379
DGL_REGISTER_GLOBAL("sampling.neighbor._CAPI_DGLSampleNeighborsEType")
.set_body([] (DGLArgs args, DGLRetValue *rv) {
    HeteroGraphRef hg = args[0];
    IdArray nodes = args[1];
380
    const std::vector<int64_t>& eid2etype_offset = ListValueToVector<int64_t>(args[2]);
381
    IdArray fanout = args[3];
382
    const std::string dir_str = args[4];
383
    const auto& prob = ListValueToVector<FloatArray>(args[5]);
384
    const bool replace = args[6];
385
    const bool rowwise_etype_sorted = args[7];
386
387
388
389

    CHECK(dir_str == "in" || dir_str == "out")
      << "Invalid edge direction. Must be \"in\" or \"out\".";
    EdgeDir dir = (dir_str == "in")? EdgeDir::kIn : EdgeDir::kOut;
390
391
    CHECK_INT64(fanout, "fanout");
    std::vector<int64_t> fanout_vec = fanout.ToVector<int64_t>();
392
393
394

    std::shared_ptr<HeteroSubgraph> subg(new HeteroSubgraph);
    *subg = sampling::SampleNeighborsEType(
395
        hg.sptr(), nodes, eid2etype_offset, fanout_vec, dir, prob, replace, rowwise_etype_sorted);
396
397
398
    *rv = HeteroSubgraphRef(subg);
  });

399
400
401
402
DGL_REGISTER_GLOBAL("sampling.neighbor._CAPI_DGLSampleNeighbors")
.set_body([] (DGLArgs args, DGLRetValue *rv) {
    HeteroGraphRef hg = args[0];
    const auto& nodes = ListValueToVector<IdArray>(args[1]);
403
404
    IdArray fanouts_array = args[2];
    const auto& fanouts = fanouts_array.ToVector<int64_t>();
405
    const std::string dir_str = args[3];
406
    const auto& prob_or_mask = ListValueToVector<NDArray>(args[4]);
407
408
    const auto& exclude_edges = ListValueToVector<IdArray>(args[5]);
    const bool replace = args[6];
409
410
411
412
413
414
415

    CHECK(dir_str == "in" || dir_str == "out")
      << "Invalid edge direction. Must be \"in\" or \"out\".";
    EdgeDir dir = (dir_str == "in")? EdgeDir::kIn : EdgeDir::kOut;

    std::shared_ptr<HeteroSubgraph> subg(new HeteroSubgraph);
    *subg = sampling::SampleNeighbors(
416
        hg.sptr(), nodes, fanouts, dir, prob_or_mask, exclude_edges, replace);
417
418
419
420
421
422
423
424

    *rv = HeteroSubgraphRef(subg);
  });

DGL_REGISTER_GLOBAL("sampling.neighbor._CAPI_DGLSampleNeighborsTopk")
.set_body([] (DGLArgs args, DGLRetValue *rv) {
    HeteroGraphRef hg = args[0];
    const auto& nodes = ListValueToVector<IdArray>(args[1]);
425
426
    IdArray k_array = args[2];
    const auto& k = k_array.ToVector<int64_t>();
427
428
429
430
    const std::string dir_str = args[3];
    const auto& weight = ListValueToVector<FloatArray>(args[4]);
    const bool ascending = args[5];

431
432
433
    CHECK(dir_str == "in" || dir_str == "out")
      << "Invalid edge direction. Must be \"in\" or \"out\".";
      EdgeDir dir = (dir_str == "in")? EdgeDir::kIn : EdgeDir::kOut;
434
435
436
437
438
439
440
441

    std::shared_ptr<HeteroSubgraph> subg(new HeteroSubgraph);
    *subg = sampling::SampleNeighborsTopk(
        hg.sptr(), nodes, k, dir, weight, ascending);

    *rv = HeteroGraphRef(subg);
  });

442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
DGL_REGISTER_GLOBAL("sampling.neighbor._CAPI_DGLSampleNeighborsBiased")
.set_body([] (DGLArgs args, DGLRetValue *rv) {
    HeteroGraphRef hg = args[0];
    const IdArray nodes = args[1];
    const int64_t fanout = args[2];
    const NDArray bias = args[3];
    const NDArray tag_offset = args[4];
    const std::string dir_str = args[5];
    const bool replace = args[6];

    CHECK(dir_str == "in" || dir_str == "out")
      << "Invalid edge direction. Must be \"in\" or \"out\".";
      EdgeDir dir = (dir_str == "in")? EdgeDir::kIn : EdgeDir::kOut;

    std::shared_ptr<HeteroSubgraph> subg(new HeteroSubgraph);
    *subg = sampling::SampleNeighborsBiased(
        hg.sptr(), nodes, fanout, bias, tag_offset, dir, replace);

    *rv = HeteroGraphRef(subg);
  });

463
464
}  // namespace sampling
}  // namespace dgl