neighbors.cpp

#include "cloud.h"
#include "nanoflann.hpp"
#include <set>
#include <cstdint>
#include <thread>
#include <iostream>

typedef struct thread_struct {
	void* kd_tree;
	void* matches;
	void* queries;
	size_t* max_count;
	std::mutex* ct_m;
	std::mutex* tree_m;
	size_t start;
	size_t end;
	double search_radius;
	bool small;
	bool option;
	size_t k;
} thread_args;

template<typename scalar_t>
void thread_routine(thread_args* targs) {
	typedef nanoflann::KDTreeSingleIndexAdaptor< nanoflann::L2_Adaptor<scalar_t, PointCloud<scalar_t> > , PointCloud<scalar_t>> my_kd_tree_t;
	typedef std::vector< std::vector<std::pair<size_t, scalar_t> > > kd_pair;
	my_kd_tree_t* index = (my_kd_tree_t*) targs->kd_tree;
	kd_pair* matches = (kd_pair*)targs->matches;
	PointCloud<scalar_t>* pcd_query = (PointCloud<scalar_t>*)targs->queries;
	size_t* max_count = targs->max_count;
	std::mutex* ct_m = targs->ct_m;
	std::mutex* tree_m = targs->tree_m;
	double eps;
	if (targs->small) {
		eps = 0.000001;
	}
	else {
		eps = 0;
	}
	double search_radius = (double) targs->search_radius;
	size_t start = targs->start;
	size_t end = targs->end;
	auto k = targs->k;
	for (size_t i = start; i < end; i++) {

		std::vector<scalar_t> p0 = *(((*pcd_query).pts)[i]);

		scalar_t* query_pt = new scalar_t[p0.size()];
		std::copy(p0.begin(), p0.end(), query_pt);
		(*matches)[i].reserve(*max_count);
		std::vector<std::pair<size_t, scalar_t> > ret_matches;
		std::vector<size_t>* knn_ret_matches = new std::vector<size_t>(k);
		std::vector<scalar_t>* knn_dist_matches = new std::vector<scalar_t>(k);

		tree_m->lock();

		size_t nMatches;
		if (targs->option){
			nMatches = index->radiusSearch(query_pt, (scalar_t)(search_radius+eps), ret_matches, nanoflann::SearchParams());
		}
		else {
			nMatches = index->knnSearch(query_pt, k, &(*knn_ret_matches)[0],&(* knn_dist_matches)[0]);
			auto temp = new std::vector<std::pair<size_t, scalar_t> >((*knn_dist_matches).size());
			for (size_t j = 0; j < (*knn_ret_matches).size(); j++){
				(*temp)[j] = std::make_pair( (*knn_ret_matches)[j],(*knn_dist_matches)[j] );
			}
			ret_matches = *temp;
		}
		tree_m->unlock();

		(*matches)[i] = ret_matches;
		
		ct_m->lock();
		if(*max_count < nMatches) {
			*max_count = nMatches;
		}
		ct_m->unlock();
	
	}

}

template<typename scalar_t>
size_t nanoflann_neighbors(std::vector<scalar_t>& queries, std::vector<scalar_t>& supports,
			std::vector<size_t>*& neighbors_indices, double radius, int dim, 
			int64_t max_num, int64_t n_threads, int64_t k, int option){

	const scalar_t search_radius = static_cast<scalar_t>(radius*radius);

	// Counting vector
	size_t* max_count = new size_t();
	*max_count = 1;

	size_t ssize = supports.size();
	// CLoud variable
	PointCloud<scalar_t> pcd;
	pcd.set(supports, dim);
	// Cloud query
	PointCloud<scalar_t>* pcd_query = new PointCloud<scalar_t>();
	(*pcd_query).set(queries, dim);

	// Tree parameters
	nanoflann::KDTreeSingleIndexAdaptorParams tree_params(15 /* max leaf */);

	// KDTree type definition
	typedef nanoflann::KDTreeSingleIndexAdaptor< nanoflann::L2_Adaptor<scalar_t, PointCloud<scalar_t> > , PointCloud<scalar_t>> my_kd_tree_t;
	typedef std::vector< std::vector<std::pair<size_t, scalar_t> > > kd_pair;

	// 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;
	kd_pair* list_matches = new kd_pair((*pcd_query).pts.size());

	// single threaded routine
	if (n_threads == 1){
		size_t i0 = 0;
		double eps;
		if (ssize < 10) {
			eps = 0.000001;
		}
		else {
			eps = 0;
		}

		for (auto& p : (*pcd_query).pts){
			auto p0 = *p;
			// Find neighbors
			scalar_t* query_pt = new scalar_t[dim];
			std::copy(p0.begin(), p0.end(), query_pt); 

			(*list_matches)[i0].reserve(*max_count);
			std::vector<std::pair<size_t, scalar_t> > ret_matches;
			std::vector<size_t>* knn_ret_matches = new std::vector<size_t>(k);
			std::vector<scalar_t>* knn_dist_matches = new std::vector<scalar_t>(k);

			size_t nMatches;
			if (!!(option)){
				nMatches = index->radiusSearch(query_pt, (scalar_t)(search_radius+eps), ret_matches, search_params);
			}
			else {
				nMatches = index->knnSearch(query_pt, (size_t)k, &(*knn_ret_matches)[0],&(* knn_dist_matches)[0]);
				auto temp = new std::vector<std::pair<size_t, scalar_t> >((*knn_dist_matches).size());
				for (size_t j = 0; j < (*knn_ret_matches).size(); j++){
					(*temp)[j] = std::make_pair( (*knn_ret_matches)[j],(*knn_dist_matches)[j] );
				}
				ret_matches = *temp;
			}
			(*list_matches)[i0] = ret_matches;
			if(*max_count < nMatches) *max_count = nMatches;
			i0++;

		}
	}
	else {// Multi-threaded routine
		std::mutex* mtx = new std::mutex();
		std::mutex* mtx_tree = new std::mutex();

		size_t n_queries = (*pcd_query).pts.size();
		size_t actual_threads = std::min((long long)n_threads, (long long)n_queries);

		std::vector<std::thread*> tid(actual_threads);

		size_t start, end;
		size_t length;
		if (n_queries) {
			length = 1;
		}
		else {
			auto res = std::lldiv((long long)n_queries, (long long)n_threads);
			length = (size_t)res.quot;
		}
		for (size_t t = 0; t < actual_threads; t++) {
			start = t*length;
			if (t == actual_threads-1) {
				end = n_queries;
			}
			else {
				end = (t+1)*length;
			}
			thread_args* targs = new thread_args();
			targs->kd_tree = index;
			targs->matches = list_matches;
			targs->max_count = max_count;
			targs->ct_m = mtx;
			targs->tree_m = mtx_tree;
			targs->search_radius = search_radius;
			targs->queries = pcd_query;
			targs->start = start;
			targs->end = end;
			if (ssize < 10) {
				targs->small = true;
			}
			else {
				targs->small = false;
			}
			targs->option = !!(option);
			targs->k = k;
			std::thread* temp = new std::thread(thread_routine<scalar_t>, targs);
			tid[t] = temp;
		}

		for (size_t t = 0; t < actual_threads; t++){
			tid[t]->join();
		}
	}

	// Reserve the memory
	if(max_num > 0) {
		*max_count = max_num;
	}
	
	size_t 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);
	size_t i1 = 0; // index of the query points
	size_t u = 0; // curent index of the neighbors_indices
	for (auto& inds : *list_matches){
		for (size_t 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>
size_t batch_nanoflann_neighbors (std::vector<scalar_t>& queries,
                               std::vector<scalar_t>& supports,
                               std::vector<long>& q_batches,
                               std::vector<long>& s_batches,
                               std::vector<size_t>*& neighbors_indices,
                               double radius, int dim, int64_t max_num, int64_t k, int option){


	// indices
	size_t i0 = 0;

	// Square radius
	const scalar_t r2 = static_cast<scalar_t>(radius*radius);

	// Counting vector
	size_t max_count = 0;

	// batch index
	size_t b = 0;
	size_t sum_qb = 0;
	size_t sum_sb = 0;

	double eps;
	if (supports.size() < 10){
		eps = 0.000001;
	}
	else {
		eps = 0;
	}
	// Nanoflann related variables

	// CLoud variable
	PointCloud<scalar_t> current_cloud;
	PointCloud<scalar_t> query_pcd;
	query_pcd.set(queries, dim);
	std::vector<std::vector<std::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& p : query_pcd.pts){
		auto p0 = *p;
		// Check if we changed batch

		scalar_t* query_pt = new scalar_t[dim];
		std::copy(p0.begin(), p0.end(), query_pt); 

		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;
		if (!!option) {
			nMatches = index->radiusSearch(query_pt, r2+eps, all_inds_dists[i0], search_params);
			// Update max count
		}
		else {
			std::vector<size_t>* knn_ret_matches = new std::vector<size_t>(k);
			std::vector<scalar_t>* knn_dist_matches = new std::vector<scalar_t>(k);
			nMatches = index->knnSearch(query_pt, (size_t)k, &(*knn_ret_matches)[0],&(*knn_dist_matches)[0]);
			auto temp = new std::vector<std::pair<size_t, scalar_t> >((*knn_dist_matches).size());
			for (size_t j = 0; j < (*knn_ret_matches).size(); j++){
				(*temp)[j] = std::make_pair( (*knn_ret_matches)[j],(*knn_dist_matches)[j] );
			}
			all_inds_dists[i0] = *temp;
		}
		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
	
	size_t 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;
	size_t 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 (size_t 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;
}