python tests passing

csrc/cpu/compat.h

+#ifdef VERSION_GE_1_3
+#define DATA_PTR data_ptr
+#else
+#define DATA_PTR data
+#endif

csrc/cpu/radius_cpu.cpp

+#include "radius_cpu.h"
+
+#include "utils.h"
+
+torch::Tensor radius_cpu(torch::Tensor query, torch::Tensor support, 
+             torch::Tensor ptr_x, torch::Tensor ptr_y, 
+			 float radius, int max_num){
+
+	CHECK_CPU(query);
+	CHECK_CPU(support);
+
+	/*
+	x = torch.cat([x, 2 * r * batch_x.view(-1, 1).to(x.dtype)], dim=-1)
+    y = torch.cat([y, 2 * r * batch_y.view(-1, 1).to(y.dtype)], dim=-1)
+	auto batch_x = ptr_x.clone();
+	auto batch_y = ptr_y.clone();
+
+	batch_x._mul(2*radius);
+	batch_y._mul(2*radius);
+
+	auto query = torch::cat({query,batch_x},-1);
+	auto support = torch::cat({support,batch_y},-1);
+	*/
+
+	torch::Tensor out;
+	std::vector<long> neighbors_indices;
+	auto options = torch::TensorOptions().dtype(torch::kLong).device(torch::kCPU);
+	int max_count = 0;
+
+	AT_DISPATCH_ALL_TYPES(query.scalar_type(), "radius_cpu", [&] {
+
+	auto data_q = query.DATA_PTR<scalar_t>();
+	auto data_s = support.DATA_PTR<scalar_t>();
+	std::vector<scalar_t> queries_stl = std::vector<scalar_t>(data_q,
+								   data_q + query.size(0)*query.size(1));
+	std::vector<scalar_t> supports_stl = std::vector<scalar_t>(data_s,
+								   data_s + support.size(0)*support.size(1));
+
+	int dim = torch::size(query, 1);
+
+	max_count = nanoflann_neighbors<scalar_t>(queries_stl, supports_stl ,neighbors_indices, radius, dim, max_num);
+
+	});
+
+	long* neighbors_indices_ptr = neighbors_indices.data();
+	out = torch::from_blob(neighbors_indices_ptr, {neighbors_indices.size()/2, 2}, options=options);
+
+	return out.t().clone();
+}
+
+void get_size_batch(const vector<long>& batch, vector<long>& res){
+
+	res.resize(batch[batch.size()-1]-batch[0]+1, 0);
+	long ind = batch[0];
+	long incr = 1;
+	for(int i=1; i < batch.size(); i++){
+
+		if(batch[i] == ind)
+			incr++;
+		else{
+			res[ind-batch[0]] = incr;
+			incr =1;
+			ind = batch[i];
+		}
+	}
+	res[ind-batch[0]] = incr;
+}
\ No newline at end of file

csrc/cpu/radius_cpu.h

+#pragma once
+
+#include <torch/extension.h>
+#include <ATen/ATen.h>
+#include "utils/neighbors.h"
+#include "utils/neighbors.cpp"
+#include <iostream>
+#include "compat.h"
+
+torch::Tensor radius_cpu(torch::Tensor query,
+			 torch::Tensor support,torch::Tensor ptr_x,
+                     torch::Tensor ptr_y, 
+			 float radius, int max_num);
+/*
+using namespace pybind11::literals;
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+	m.def("radius_search",
+	      &radius_search,
+	      "compute the radius search of a point cloud using nanoflann"
+	      "-query : a pytorch tensor of size N1 x d,. used to query the nearest neighbors"
+	      "- support : a pytorch tensor of size N2 x d. used to build the tree"
+	      "-  radius : float number, size of the ball for the radius search."
+	      "- max_num : int number, indicate the maximum of neaghbors allowed(if -1 then all the possible neighbors will be computed). "
+	      " - mode : int number that indicate which format for the neighborhood"
+	      " mode=0 mean a matrix of neighbors(-1 for shadow neighbors)"
+	      "mode=1 means a matrix of edges of size Num_edge x 2"
+	      "return a tensor of size N1 x M where M is either max_num or the maximum number of neighbors found if mode = 0, if mode=1 return a tensor of size Num_edge x 2.",
+	      "query"_a, "support"_a, "radius"_a, "dim"_a, "max_num"_a=-1, "mode"_a=0);
+	m.def("batch_radius_search",
+	      &batch_radius_search,
+		"compute the radius search of a point cloud for each batch using nanoflann"
+		"-query : a pytorch tensor (float) of size N1 x d,. used to query the nearest neighbors"
+		"- support : a pytorch tensor(float) of size N2 x d. used to build the tree"
+		"- query_batch : a pytorch tensor(long) contains indices of the batch of the query size N1"
+	      "NB : the batch must be sorted"
+		"- support_batch: a pytorch tensor(long) contains indices of the batch of the support size N2"
+	      "NB: the batch must be sorted"
+		"-radius: float number, size of the ball for the radius search."
+		"- max_num : int number, indicate the maximum of neaghbors allowed(if -1 then all the possible neighbors wrt the radius will be computed)."
+		"- mode : int number that indicate which format for the neighborhood"
+		"mode=0 mean a matrix of neighbors(N2 for shadow neighbors)"
+		"mode=1 means a matrix of edges of size Num_edge x 2"
+	      "return a tensor of size N1 x M where M is either max_num or the maximum number of neighbors found if mode = 0, if mode=1 return a tensor of size Num_edge x 2.",
+	      "query"_a, "support"_a, "query_batch"_a, "support_batch"_a, "radius"_a, "dim"_a, "max_num"_a=-1, "mode"_a=0);
+}
+*/
\ No newline at end of file

csrc/cpu/utils/cloud.h

+// Author: Peiyuan Liao (alexander_liao@outlook.com)
+//
+
+
+# pragma once
+
+#include <ATen/ATen.h>
+#include <cmath>
+#include <vector>
+#include <unordered_map>
+#include <map>
+#include <algorithm>
+#include <numeric>
+#include <iostream>
+#include <iomanip>
+
+#include <time.h>
+
+
+template<typename scalar_t>
+struct PointCloud
+{
+	std::vector<std::vector<scalar_t>> pts;
+
+	void set(std::vector<scalar_t> new_pts, int dim){
+
+		// pts = std::vector<Point>((Point*)new_pts, (Point*)new_pts+new_pts.size()/3);
+		std::vector<std::vector<scalar_t>> temp(new_pts.size()/dim);
+		for(unsigned int i=0; i < new_pts.size(); i++){
+			if(i%dim == 0){
+
+				//Point point;
+				std::vector<scalar_t> point(dim);
+				//std::vector<scalar_t> vect(sizeof(scalar_t)*dim, 0)
+				//point.pt = temp;
+				for (unsigned int j = 0; j < dim; j++) {
+					point[j]=new_pts[i+j];
+					//point.pt[j] = new_pts[i+j];
+				}
+				temp[i/dim] = point;
+			}
+		}
+
+		pts = temp;
+	}
+	void set_batch(std::vector<scalar_t> new_pts, int begin, int size, int dim){
+		std::vector<std::vector<scalar_t>> temp(size);
+		for(int i=0; i < size; i++){
+			//std::vector<scalar_t> temp(sizeof(scalar_t)*dim, 0);
+			//point.pt = temp;
+			std::vector<scalar_t> point(dim);
+			for (unsigned int j = 0; j < dim; j++) {
+					point[j] = new_pts[dim*(begin+i)+j];
+			}
+
+			temp[i] = point;
+
+		}
+		pts = temp;
+	}
+
+	// Must return the number of data points
+	inline size_t kdtree_get_point_count() const { return pts.size(); }
+
+	// Returns the dim'th component of the idx'th point in the class:
+	inline scalar_t kdtree_get_pt(const size_t idx, const size_t dim) const
+	{
+		return pts[idx][dim];
+	}
+
+	// Optional bounding-box computation: return false to default to a standard bbox computation loop.
+	//   Return true if the BBOX was already computed by the class and returned in "bb" so it can be avoided to redo it again.
+	//   Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 for point clouds)
+	template <class BBOX>
+	bool kdtree_get_bbox(BBOX& /* bb */) const { return false; }
+
+
+};

csrc/cpu/utils/neighbors.cpp

+
+// Taken from https://github.com/HuguesTHOMAS/KPConv
+
+#include "neighbors.h"
+
+template<typename scalar_t>
+int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
+			vector<long>& neighbors_indices, float radius, int dim, int max_num){
+
+	// Initiate variables
+	// ******************
+
+	// square radius
+
+	const scalar_t search_radius = static_cast<scalar_t>(radius*radius);
+
+	// indices
+	int i0 = 0;
+
+	// Counting vector
+	int max_count = 1;
+	float d2;
+
+	// Nanoflann related variables
+	// ***************************
+
+	// CLoud variable
+	PointCloud<scalar_t> pcd;
+	pcd.set(supports, dim);
+	//Cloud query
+	PointCloud<scalar_t> pcd_query;
+	pcd_query.set(queries, dim);
+
+	// Tree parameters
+	nanoflann::KDTreeSingleIndexAdaptorParams tree_params(10 /* max leaf */);
+
+	// KDTree type definition
+	typedef nanoflann::KDTreeSingleIndexAdaptor< nanoflann::L2_Adaptor<scalar_t, PointCloud<scalar_t> > , PointCloud<scalar_t>> my_kd_tree_t;
+
+	// Pointer to trees
+	my_kd_tree_t* index;
+	index = new my_kd_tree_t(dim, pcd, tree_params);
+	index->buildIndex();
+	// Search neigbors indices
+	// ***********************
+
+	// Search params
+	nanoflann::SearchParams search_params;
+	search_params.sorted = true;
+	std::vector< std::vector<std::pair<size_t, scalar_t> > > list_matches(pcd_query.pts.size());
+
+	float eps = 0.00001;
+
+	for (auto& p0 : pcd_query.pts){
+
+		// Find neighbors
+		scalar_t query_pt[dim];
+		std::copy(p0.begin(), p0.end(), query_pt); 
+
+		//for(int i=0; i < p0.size(); i++)
+		//std::cout << query_pt[i] << '\n';
+
+		list_matches[i0].reserve(max_count);
+		std::vector<std::pair<size_t, scalar_t> > ret_matches;
+		//nanoflann::RadiusResultSet<float,size_t> resultSet(r2, indices_dists);
+
+		const size_t nMatches = index->radiusSearch(&query_pt[0], search_radius+eps, ret_matches, search_params);
+		list_matches[i0] = ret_matches;
+		if(max_count < nMatches) max_count = nMatches;
+		i0++;
+
+
+		// Get worst (furthest) point, without sorting:
+		
+		// cout << "\n neighbors: " << nMatches << "\n";
+
+		// Get worst (furthest) point, without sorting:
+		// for(int i=0; i < ret_matches.size(); i++)
+		// std::cout << ret_matches.at(i) << '\n';
+
+	}
+	// Reserve the memory
+	if(max_num > 0) {
+		max_count = max_num;
+	}
+	
+	int size = 0; // total number of edges
+	for (auto& inds : list_matches){
+		if(inds.size() <= max_count)
+			size += inds.size();
+		else
+			size += max_count;
+	}
+	
+	neighbors_indices.resize(size*2);
+	int i1 = 0; // index of the query points
+	int u = 0; // curent index of the neighbors_indices
+	for (auto& inds : list_matches){
+		for (int j = 0; j < max_count; j++){
+			if(j < inds.size()){
+				neighbors_indices[u] = inds[j].first;
+				neighbors_indices[u + 1] = i1;
+				u += 2;
+			}
+		}
+		i1++;
+	}
+
+	return max_count;
+
+
+
+
+}
+template<typename scalar_t>
+int batch_nanoflann_neighbors (vector<scalar_t>& queries,
+                               vector<scalar_t>& supports,
+                               vector<long>& q_batches,
+                               vector<long>& s_batches,
+                               vector<long>& neighbors_indices,
+                               float radius, int dim, int max_num,
+			       int mode){
+
+
+// Initiate variables
+// ******************
+// indices
+	int i0 = 0;
+
+// Square radius
+	float r2 = radius * radius;
+
+	// Counting vector
+	int max_count = 0;
+	float d2;
+
+
+	// batch index
+	long b = 0;
+	long sum_qb = 0;
+	long sum_sb = 0;
+
+	// Nanoflann related variables
+	// ***************************
+
+	// CLoud variable
+	PointCloud<scalar_t> current_cloud;
+	PointCloud<scalar_t> query_pcd;
+	query_pcd.set(queries, dim);
+	vector<vector<pair<size_t, scalar_t> > > all_inds_dists(query_pcd.pts.size());
+
+	// Tree parameters
+	nanoflann::KDTreeSingleIndexAdaptorParams tree_params(10 /* max leaf */);
+
+	// KDTree type definition
+	typedef nanoflann::KDTreeSingleIndexAdaptor< nanoflann::L2_Adaptor<scalar_t, PointCloud<scalar_t> > , PointCloud<scalar_t>> my_kd_tree_t;
+
+// Pointer to trees
+	my_kd_tree_t* index;
+    // Build KDTree for the first batch element
+	current_cloud.set_batch(supports, sum_sb, s_batches[b], dim);
+	index = new my_kd_tree_t(dim, current_cloud, tree_params);
+	index->buildIndex();
+// Search neigbors indices
+// ***********************
+// Search params
+	nanoflann::SearchParams search_params;
+	search_params.sorted = true;
+
+	for (auto& p0 : query_pcd.pts){
+// Check if we changed batch
+
+		scalar_t query_pt[dim];
+		std::copy(p0.begin(), p0.end(), query_pt); 
+
+		/*
+		std::cout << "\n ========== \n";
+		for(int i=0; i < dim; i++)
+			std::cout << query_pt[i] << '\n';
+		std::cout << "\n ========== \n";
+		*/
+	
+		if (i0 == sum_qb + q_batches[b]){
+			sum_qb += q_batches[b];
+			sum_sb += s_batches[b];
+			b++;
+
+// Change the points
+			current_cloud.pts.clear();
+			current_cloud.set_batch(supports, sum_sb, s_batches[b], dim);
+// Build KDTree of the current element of the batch
+			delete index;
+			index = new my_kd_tree_t(dim, current_cloud, tree_params);
+			index->buildIndex();
+		}
+// Initial guess of neighbors size
+		all_inds_dists[i0].reserve(max_count);
+// Find neighbors
+		size_t nMatches = index->radiusSearch(query_pt, r2, all_inds_dists[i0], search_params);
+// Update max count
+
+		std::vector<std::pair<size_t, float> > indices_dists;
+		nanoflann::RadiusResultSet<float,size_t> resultSet(r2, indices_dists);
+
+		index->findNeighbors(resultSet, query_pt, search_params);
+
+		if (nMatches > max_count)
+			max_count = nMatches;
+// Increment query idx
+		i0++;
+	}
+	// how many neighbors do we keep
+	if(max_num > 0) {
+		max_count = max_num;
+	}
+// Reserve the memory
+	if(mode == 0){
+		neighbors_indices.resize(query_pcd.pts.size() * max_count);
+		i0 = 0;
+		sum_sb = 0;
+		sum_qb = 0;
+		b = 0;
+
+		for (auto& inds_dists : all_inds_dists){// Check if we changed batch
+
+			if (i0 == sum_qb + q_batches[b]){
+				sum_qb += q_batches[b];
+				sum_sb += s_batches[b];
+				b++;
+			}
+
+			for (int j = 0; j < max_count; j++){
+				if (j < inds_dists.size())
+					neighbors_indices[i0 * max_count + j] = inds_dists[j].first + sum_sb;
+				else
+					neighbors_indices[i0 * max_count + j] = supports.size();
+
+			}
+
+			i0++;
+		}
+		delete index;
+	}
+	else if(mode == 1){
+		int size = 0; // total number of edges
+		for (auto& inds_dists : all_inds_dists){
+			if(inds_dists.size() <= max_count)
+				size += inds_dists.size();
+			else
+				size += max_count;
+		}
+		neighbors_indices.resize(size * 2);
+		i0 = 0;
+		sum_sb = 0;
+		sum_qb = 0;
+		b = 0;
+		int u = 0;
+		for (auto& inds_dists : all_inds_dists){
+			if (i0 == sum_qb + q_batches[b]){
+				sum_qb += q_batches[b];
+				sum_sb += s_batches[b];
+				b++;
+			}
+			for (int j = 0; j < max_count; j++){
+				if (j < inds_dists.size()){
+					neighbors_indices[u] = inds_dists[j].first + sum_sb;
+					neighbors_indices[u + 1] = i0;
+					u += 2;
+				}
+			}
+			i0++;
+		}
+	}
+	return max_count;
+}

csrc/cpu/utils/neighbors.h

+
+
+#include "cloud.h"
+#include "nanoflann.hpp"
+#include <set>
+#include <cstdint>
+
+using namespace std;
+
+
+template<typename scalar_t>
+int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
+			vector<long>& neighbors_indices, float radius, int dim, int max_num, int mode);
+
+
+template<typename scalar_t>
+int batch_nanoflann_neighbors(vector<scalar_t>& queries,
+                                vector<scalar_t>& supports,
+                                vector<long>& q_batches,
+                                vector<long>& s_batches,
+                                vector<long>& neighbors_indices,
+			        float radius, int dim, int max_num, int mode);

csrc/radius.cpp

@@ -4,6 +4,7 @@
 #ifdef WITH_CUDA
 #include "cuda/radius_cuda.h"
 #endif
+#include "cpu/radius_cpu.h"
 
 #ifdef _WIN32
 PyMODINIT_FUNC PyInit__radius(void) { return NULL; }
@@ -18,7 +19,7 @@ torch::Tensor radius(torch::Tensor x, torch::Tensor y, torch::Tensor ptr_x,
     AT_ERROR("Not compiled with CUDA support");
 #endif
   } else {
-    AT_ERROR("No CPU version supported");
+    return radius_cpu(x, y, ptr_x, ptr_y, r, max_num_neighbors);
   }
 }
 

torch_cluster/radius.py

@@ -87,9 +87,13 @@ def radius(x: torch.Tensor, y: torch.Tensor, r: float,
         assert x.size(0) == batch_x.size(0)
         assert y.size(0) == batch_y.size(0)
 
+        
         x = torch.cat([x, 2 * r * batch_x.view(-1, 1).to(x.dtype)], dim=-1)
         y = torch.cat([y, 2 * r * batch_y.view(-1, 1).to(y.dtype)], dim=-1)
 
+        return torch.ops.torch_cluster.radius(x, y, x, y, r,
+                                              max_num_neighbors)
+        """
         tree = scipy.spatial.cKDTree(x.detach().numpy())
         col = tree.query_ball_point(y.detach().numpy(), r)
         col = [torch.tensor(c, dtype=torch.long) for c in col]
@@ -98,7 +102,7 @@ def radius(x: torch.Tensor, y: torch.Tensor, r: float,
         row, col = torch.cat(row, dim=0), torch.cat(col, dim=0)
         mask = col < int(tree.n)
         return torch.stack([row[mask], col[mask]], dim=0)
-
+        """
 
 def radius_graph(x: torch.Tensor, r: float,
                  batch: Optional[torch.Tensor] = None, loop: bool = False,
@@ -136,7 +140,11 @@ def radius_graph(x: torch.Tensor, r: float,
     assert flow in ['source_to_target', 'target_to_source']
     row, col = radius(x, x, r, batch, batch,
                       max_num_neighbors if loop else max_num_neighbors + 1)
+    if x.is_cuda:
         row, col = (col, row) if flow == 'source_to_target' else (row, col)
+    else:
+        row, col = (col, row) if flow == 'target_to_source' else (row, col)
+        
     if not loop:
         mask = row != col
         row, col = row[mask], col[mask]