python tests passing

Author: liaopeiyuan

csrc/cpu/compat.h

@@ -0,0 +1,5 @@
+#ifdef VERSION_GE_1_3
+#define DATA_PTR data_ptr
+#else
+#define DATA_PTR data
+#endif

csrc/cpu/radius_cpu.cpp

@@ -0,0 +1,67 @@
+#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;
+}

csrc/cpu/radius_cpu.h

@@ -0,0 +1,46 @@
+#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);
+}
+*/

csrc/cpu/utils/cloud.h

@@ -0,0 +1,78 @@
+// 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; }
+
+
+};