Add directed=False (induced subgraph) support for homogeneous CPU neighbor sampling

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added support for `directed=False` (induced subgraph) neighbor sampling for homogeneous graphs on CPU. Currently does not support `disjoint=True` or `distributed=True` ([#690](https://github.com/pyg-team/pyg-lib/pull/690))
 - Added Linux `aarch64` manylinux wheel and CUDA Docker image builds ([#683](https://github.com/pyg-team/pyg-lib/pull/683))
 
 ### Changed

pyg_lib/csrc/sampler/cpu/neighbor_kernel.cpp

@@ -159,6 +159,42 @@ class NeighborSampler {
     }
   }
 
+  std::tuple<at::Tensor, at::Tensor, c10::optional<at::Tensor>>
+  get_subgraph_edges(pyg::sampler::Mapper<node_t, scalar_t>& dst_mapper,
+                     std::vector<node_t>& out_global_dst_nodes,
+                     bool csc = false) {
+    std::vector<scalar_t> out_row, out_col, out_edge_id;
+    if (save_edge_ids) {
+      out_edge_id = std::vector<scalar_t>{};
+    }
+    for (size_t i = 0; i < out_global_dst_nodes.size(); ++i) {
+      const auto global_src_node = out_global_dst_nodes[i];
+      const auto row_start = rowptr_[to_scalar_t(global_src_node)];
+      const auto row_end = rowptr_[to_scalar_t(global_src_node) + 1];
+
+      for (auto edge_id = row_start; edge_id < row_end; ++edge_id) {
+        const auto v = to_node_t(col_[edge_id], global_src_node);
+        const auto local_idx = dst_mapper.map(v);
+        if (local_idx >= 0) {
+          out_row.push_back(i);
+          out_col.push_back(local_idx);
+          out_edge_id.push_back(edge_id);
+        }
+      }
+    }
+    const auto row = pyg::utils::from_vector(out_row);
+    const auto col = pyg::utils::from_vector(out_col);
+    c10::optional<at::Tensor> edge_id =
+        save_edge_ids
+            ? c10::optional<at::Tensor>(pyg::utils::from_vector(out_edge_id))
+            : c10::nullopt;
+    if (!csc) {
+      return std::make_tuple(row, col, edge_id);
+    } else {
+      return std::make_tuple(col, row, edge_id);
+    }
+  }
+
  private:
   inline scalar_t to_scalar_t(const scalar_t& node) { return node; }
   inline scalar_t to_scalar_t(const std::pair<scalar_t, scalar_t>& node) {
@@ -498,11 +534,17 @@ sample(const at::Tensor& rowptr,
     }
 
     out_node_id = pyg::utils::from_vector(sampled_nodes);
-    TORCH_CHECK(directed, "Undirected subgraphs not yet supported");
     if (directed) {
       std::tie(out_row, out_col, out_edge_id) = sampler.get_sampled_edges(csc);
     } else {
+      TORCH_CHECK(
+          !distributed,
+          "Induced subgraph sampling not yet supported in distributed mode");
       TORCH_CHECK(!disjoint, "Disjoint subgraphs not yet supported");
+      if constexpr (!disjoint) {
+        std::tie(out_row, out_col, out_edge_id) =
+            sampler.get_subgraph_edges(mapper, sampled_nodes, csc);
+      }
     }
 
     num_sampled_edges_per_hop = sampler.num_sampled_edges_per_hop;

test/csrc/sampler/test_neighbor.cpp

@@ -375,3 +375,163 @@ TEST(HeteroBiasedNeighborTest, BasicAssertions) {
   auto expected_edges = at::tensor({0, 2}, options);
   EXPECT_TRUE(at::equal(std::get<3>(out).value().at(rel_key), expected_edges));
 }
+
+TEST(NeighborSamplerTest, InducedSubgraphHasMoreEdgesThanDirected) {
+  const int64_t num_nodes = 5;
+  const auto options = at::TensorOptions().dtype(at::kLong);
+  at::Tensor rowptr, col;
+  std::tie(rowptr, col) = cycle_graph(num_nodes, options);
+
+  const auto seed = at::tensor({0}, options);
+  std::vector<int64_t> num_neighbors = {2, 2};
+
+  const auto directed_out = pyg::sampler::neighbor_sample(
+      /*rowptr=*/rowptr,
+      /*col=*/col,
+      /*seed=*/seed,
+      /*num_neighbors=*/num_neighbors,
+      /*node_time=*/c10::nullopt,
+      /*edge_time=*/c10::nullopt,
+      /*seed_time=*/c10::nullopt,
+      /*edge_weight=*/c10::nullopt,
+      /*csc=*/false,
+      /*replace=*/false,
+      /*directed=*/true,
+      /*disjoint=*/false,
+      /*temporal_strategy=*/"uniform",
+      /*return_edge_id=*/false);
+
+  const auto undirected_out = pyg::sampler::neighbor_sample(
+      /*rowptr=*/rowptr,
+      /*col=*/col,
+      /*seed=*/seed,
+      /*num_neighbors=*/num_neighbors,
+      /*node_time=*/c10::nullopt,
+      /*edge_time=*/c10::nullopt,
+      /*seed_time=*/c10::nullopt,
+      /*edge_weight=*/c10::nullopt,
+      /*csc=*/false,
+      /*replace=*/false,
+      /*directed=*/false,
+      /*disjoint=*/false,
+      /*temporal_strategy=*/"uniform",
+      /*return_edge_id=*/false);
+
+  const auto& d_row = std::get<0>(directed_out);
+  const auto& u_row = std::get<0>(undirected_out);
+  const auto& u_node = std::get<2>(undirected_out);
+
+  EXPECT_EQ(u_node.numel(), num_nodes);
+  EXPECT_EQ(u_row.numel(), 2 * num_nodes);
+  EXPECT_LT(d_row.numel(), u_row.numel());
+}
+
+TEST(NeighborSamplerTest, InducedSubgraphRejectsDisjoint) {
+  const int64_t num_nodes = 5;
+  const auto options = at::TensorOptions().dtype(at::kLong);
+  at::Tensor rowptr, col;
+  std::tie(rowptr, col) = cycle_graph(num_nodes, options);
+
+  const auto seed = at::tensor({0, 1}, options);
+  std::vector<int64_t> num_neighbors = {2};
+
+  EXPECT_THROW(pyg::sampler::neighbor_sample(
+                   /*rowptr=*/rowptr,
+                   /*col=*/col,
+                   /*seed=*/seed,
+                   /*num_neighbors=*/num_neighbors,
+                   /*node_time=*/c10::nullopt,
+                   /*edge_time=*/c10::nullopt,
+                   /*seed_time=*/c10::nullopt,
+                   /*edge_weight=*/c10::nullopt,
+                   /*csc=*/false,
+                   /*replace=*/false,
+                   /*directed=*/false,
+                   /*disjoint=*/true,
+                   /*temporal_strategy =*/"uniform",
+                   /*return_edge_id =*/false),
+               c10::Error);
+}
+
+TEST(NeighborSamplerTest, InducedSubgraphReturnsEdgeIds) {
+  const int64_t num_nodes = 5;
+  const auto options = at::TensorOptions().dtype(at::kLong);
+  at::Tensor rowptr, col;
+  std::tie(rowptr, col) = cycle_graph(num_nodes, options);
+  const auto seed = at::tensor({0}, options);
+  std::vector<int64_t> num_neighbors = {2, 2};
+
+  const auto out = pyg::sampler::neighbor_sample(
+      /*rowptr=*/rowptr,
+      /*col=*/col,
+      /*seed=*/seed,
+      /*num_neighbors=*/num_neighbors,
+      /*node_time=*/c10::nullopt,
+      /*edge_time=*/c10::nullopt,
+      /*seed_time=*/c10::nullopt,
+      /*edge_weight=*/c10::nullopt,
+      /*csc=*/false,
+      /*replace=*/false,
+      /*directed=*/false,
+      /*disjoint=*/false,
+      /*temporal_strategy=*/"uniform",
+      /*return_edge_id=*/true);
+
+  const auto& out_row = std::get<0>(out);
+  const auto& out_edge_id = std::get<3>(out);
+  ASSERT_TRUE(out_edge_id.has_value());
+  EXPECT_EQ(out_edge_id.value().numel(), out_row.numel());
+}
+
+TEST(NeighborSamplerTest, InducedSubgraphCscSwapsRowCol) {
+  const int64_t num_nodes = 5;
+  const auto options = at::TensorOptions().dtype(at::kLong);
+  at::Tensor rowptr, col;
+  std::tie(rowptr, col) = cycle_graph(num_nodes, options);
+  const auto seed = at::tensor({0}, options);
+  std::vector<int64_t> num_neighbors = {2, 2};
+
+  const auto csr_out = pyg::sampler::neighbor_sample(
+      /*rowptr=*/rowptr,
+      /*col=*/col,
+      /*seed=*/seed,
+      /*num_neighbors=*/num_neighbors,
+      /*node_time=*/c10::nullopt,
+      /*edge_time=*/c10::nullopt,
+      /*seed_time=*/c10::nullopt,
+      /*edge_weight=*/c10::nullopt,
+      /*csc=*/false,
+      /*replace=*/false,
+      /*directed=*/false,
+      /*disjoint=*/false,
+      /*temporal_strategy=*/"uniform",
+      /*return_edge_id=*/false);
+
+  const auto csc_out = pyg::sampler::neighbor_sample(
+      /*rowptr=*/rowptr,
+      /*col=*/col,
+      /*seed=*/seed,
+      /*num_neighbors=*/num_neighbors,
+      /*node_time=*/c10::nullopt,
+      /*edge_time=*/c10::nullopt,
+      /*seed_time=*/c10::nullopt,
+      /*edge_weight=*/c10::nullopt,
+      /*csc=*/true,
+      /*replace=*/false,
+      /*directed=*/false,
+      /*disjoint=*/false,
+      /*temporal_strategy=*/"uniform",
+      /*return_edge_id=*/false);
+
+  const auto& csr_row = std::get<0>(csr_out);
+  const auto& csr_col = std::get<1>(csr_out);
+  const auto& csc_row = std::get<0>(csc_out);
+  const auto& csc_col = std::get<1>(csc_out);
+
+  // Same number of edges either way.
+  EXPECT_EQ(csr_row.numel(), csc_row.numel());
+
+  // csc output should be the row/col swap of csr output.
+  EXPECT_TRUE(at::equal(csr_row, csc_col));
+  EXPECT_TRUE(at::equal(csr_col, csc_row));
+}