Construction and search of Kd-Tree in c++

Question

I am not very fluent in C++ and want to construct a Kd-Tree in C++ based on this paper Efficient Locally Weighted Polynomial Regression Predictions, to use it to conduct some fast implementations of local linear regression . However, the speed for creating the tree and search( especially) is very slow. I have tried a stack implementation, and it just improves the speed by a little. Below are the relevant part of my code.

KDtreeiteration.h

#pragma once
#include <memory> 
#include <vector>
#include <Eigen/Dense>
#include <chrono>
#include <iostream>

using all_point_t = std::vector<Eigen::VectorXd>; 

template<typename T, typename U> 
std::pair<T, U> operator+(const std::pair<T,U>&, const std::pair<T,U>&);

template<typename T> 
std::ostream& operator<<(std::ostream&, const std::vector<T>&); 

class kdnode{ 
            public: 
                int n_below; 
                int split_d; 
                double split_v;
                Eigen::MatrixXd XtX; 
                Eigen::VectorXd XtY; 
                std::vector<double> dim_max; 
                std::vector<double> dim_min;
                std::shared_ptr<kdnode> right_child; 
                std::shared_ptr<kdnode> left_child; 
                double sumY; 

                kdnode(); // constructor 
                kdnode(kdnode&&) ;  //move 
                ~kdnode();  // destructor    
};

class kdtree{
    public:     
        kdtree(); 
        ~kdtree();  
        double weight_sf;
        int tracker;
        std::shared_ptr<kdnode> root; 
        std::shared_ptr<kdnode> leaf;
        std::shared_ptr<kdnode> build_tree(all_point_t::iterator, all_point_t::iterator, int, double, int, size_t, std::vector<double>, std::vector<double>);
        std::shared_ptr<kdnode> build_exacttree(all_point_t::iterator, all_point_t::iterator, int, double, int, size_t, std::vector<double>, std::vector<double>);
        explicit kdtree(all_point_t, int, int); 
        std::pair<Eigen::MatrixXd, Eigen::VectorXd> get_XtXXtY(const Eigen::VectorXd& query, std::vector<double> dim_max, std::vector<double> dim_min, std::shared_ptr<kdnode>& root, const Eigen::VectorXd& h, int kcode);
        std::pair<Eigen::MatrixXd, Eigen::VectorXd> getapprox_XtXXtY(const Eigen::VectorXd& query, std::vector<double> dim_max, std::vector<double> dim_min, std::shared_ptr<kdnode>& root, double epsilon, const Eigen::VectorXd& h, int kcode);
        std::pair<Eigen::MatrixXd, Eigen::VectorXd> find_XtXXtY(const Eigen::VectorXd& query, int method, double epsilon, const Eigen::VectorXd& h, int kcode);
        // test functions; 
        void test_XtX(Eigen::MatrixXd);
        void test_XtY(Eigen::MatrixXd);
        void test_XtXXtY(Eigen::MatrixXd);
};


class Timer
{
    public:
        Timer();
        void reset();
        double elapsed() const;

    private:
        typedef std::chrono::high_resolution_clock clock_;
        typedef std::chrono::duration<double, std::ratio<1> > second_;
        std::chrono::time_point<clock_> beg_;
};

// functions 
all_point_t convert_to_vector(const Eigen::MatrixXd& XY_mat);
all_point_t convert_to_query(const Eigen::MatrixXd& XY_mat);
all_point_t convert_to_queryX(const Eigen::MatrixXd& X_mat);
double eval_kernel(int kcode, const double& z);
std::pair<double, double> calculate_weight(int kcode, const Eigen::VectorXd& query, const std::vector<double>& dim_max, const std::vector<double>& dim_min , const Eigen::VectorXd& h); 
Eigen::MatrixXd form_ll_XtX(const Eigen::MatrixXd& XtX, const Eigen::VectorXd& query ); 
Eigen::VectorXd form_ll_XtY(const Eigen::VectorXd& XtY , const Eigen::VectorXd& query);
Eigen::VectorXd calculate_mx(const Eigen::MatrixXd& XtX, const Eigen::VectorXd& XtY);

// R function
Eigen::VectorXd loclinear_i(const Eigen::MatrixXd& XY_mat, int method, int kcode, 
                        double epsilon, const Eigen::VectorXd& h, int N_min);

// for bandwidth selection
std::pair<Eigen::VectorXd, double> calculate_mx_Xinv(int kcode, const Eigen::MatrixXd &XtX, const Eigen::MatrixXd &XtY);
Eigen::VectorXd h_select_i(const Eigen::MatrixXd& XY_mat, int method, int kcode, double epsilon, 
                         const Eigen::MatrixXd& bw, int N_min);
void pertube_XtX(Eigen::MatrixXd& XtX);
double max_weight(int kcode, const Eigen::VectorXd&h);
Eigen::VectorXd predict_i(const Eigen::MatrixXd& XY_mat, const Eigen::MatrixXd& Xpred_mat, int method, int kcode, 
                        double epsilon, const Eigen::VectorXd& h,  int N_min);

Implementation for creation of tree

std::shared_ptr<kdnode> kdtree::build_tree(all_point_t::iterator start, all_point_t::iterator end, 
                                           int split_d, double split_v, int N_min, size_t len,
                                           std::vector<double> dim_max_, std::vector<double> dim_min_){
        
    std::shared_ptr<kdnode> newnode = std::make_shared<kdnode>();
    if(end == start) {   
        return newnode; 
    }

    newnode->n_below = len;    
    newnode->dim_max = dim_max_;
    newnode->dim_min = dim_min_;

    if (end-start <= N_min) {            
        Eigen::MatrixXd XtX_(dim_max_.size() + 1 , dim_max_.size() + 1);
        XtX_.setZero(); 
        Eigen::VectorXd XtY_(dim_max_.size() + 1);
        XtY_.setZero();
        for (auto k =0; k <dim_max_.size(); k++ ){
            dim_max_[k] = (*start)(k+1);
            dim_min_[k] = (*start)(k+1);
        }                  
        for (auto i = start; i != end; i ++){
            Eigen::VectorXd XY = *i;    
            Eigen::VectorXd Y = XY.tail<1>(); 
            Eigen::VectorXd X = XY.head(XY.size()-1); 
            XtY_ += X*Y; 
            XtX_ += X*X.transpose();  
            for (auto j = 0; j < dim_max_.size(); j++) {
                if(X(j+1) > dim_max_[j]){
                    dim_max_[j] = X(j+1); 
                }
                if(X(j+1) < dim_min_[j]){
                    dim_min_[j] = X(j+1); 
                } 
            }
        } 
         newnode->dim_max = dim_max_;
         newnode->dim_min = dim_min_;   
        // std::cout << "dim_min_" << dim_min_ <<'\n'; 
        // std::cout << "dim_max_" << dim_max_ <<'\n';
         newnode->XtX = XtX_; 
         newnode->XtY = XtY_;  

         return newnode; 
    }

    else {  
        size_t l_len = len/2  ;          // left length
        size_t r_len = len - l_len;     // right length
        auto middle = start + len/2;   // middle iterator 
        int max = 0; 
        int dim = 0;
        for(int i = 0; i < newnode->dim_max.size(); i++){   
            double var = newnode->dim_max[i] - newnode->dim_min[i]; 
            if(var > max){
                max = var; 
                dim = i; 
            }
        }

        newnode -> split_d = dim; 
        int vector_dim = dim + 1;  

        std::nth_element(start, middle, end, [vector_dim](const Eigen::VectorXd& a, const Eigen::VectorXd & b) {
            return a(vector_dim) < b(vector_dim);    
        });           

        newnode->split_v = (*middle)[vector_dim];   

        dim_max_[dim] = newnode->split_v; 
        dim_min_[dim] = newnode->split_v;

        newnode-> left_child = build_tree(start, middle, newnode->split_d, newnode->split_v, N_min, l_len, dim_max_, newnode->dim_min);
        newnode-> right_child = build_tree(middle, end, newnode->split_d, newnode->split_v, N_min, r_len, newnode->dim_max, dim_min_);
        
        if ((newnode->left_child) && (newnode->right_child)){ 
            newnode->XtX = newnode->left_child->XtX + newnode ->right_child->XtX;  // sumY = the sum of the bottom 2 nodes
            newnode->XtY = newnode->left_child->XtY + newnode ->right_child->XtY;
        } 
        else if (newnode->left_child) {
            newnode->XtY = newnode->left_child->XtY;
            newnode->XtX = newnode->left_child->XtX;
        }
        else if (newnode->right_child) {
            newnode->XtX = newnode->right_child->XtX; 
            newnode->XtY = newnode->right_child->XtY; 
        }   
    }
    return newnode;    
}

Implementation of search

std::pair<Eigen::MatrixXd, Eigen::VectorXd> kdtree::get_XtXXtY(const Eigen::VectorXd& query, 
                                                               std::vector<double> dim_max, 
                                                               std::vector<double> dim_min, 
                                                               std::shared_ptr<kdnode>& root, 
                                                               const Eigen::VectorXd& h, 
                                                               int kcode){

    std::pair<double,double> weights;
    std::stack<std::shared_ptr<kdnode>> storage; 
    std::shared_ptr<kdnode> curr = root; 
    Eigen::MatrixXd XtX = Eigen::MatrixXd::Zero(curr->XtX.rows() , curr->XtX.cols()); 
    Eigen::VectorXd XtY = Eigen::MatrixXd::Zero(curr->XtY.rows() , curr->XtY.cols());
    weights = calculate_weight(kcode, query, dim_max, dim_min,h);
    double w_max = weights.first; 
    double w_min = weights.second;

    while (w_max != w_min || storage.empty() == false){
        while (w_max != w_min ){   // if condition fufilled
            storage.push(curr);
            curr = curr->left_child;  
            weights = calculate_weight(kcode, query, curr->dim_max, curr->dim_min, h);  // calculate max and min weight 
            w_max = weights.first;                      
            w_min = weights.second; 

            if(w_max == w_min ){ 
                XtX += w_max*curr->XtX;
                XtY += w_max*curr->XtY; 
            }           
        }

        curr = storage.top();
        storage.pop(); 
        curr = curr->right_child; 
        weights = calculate_weight(kcode, query, curr->dim_max, curr->dim_min, h);  // calculate max and min weight 
        w_max = weights.first;   
        w_min = weights.second; 

        if(w_max == w_min){ 
            XtX += w_max*curr->XtX;
            XtY += w_max*curr->XtY;
        }
    }   
    return std::make_pair(XtX,XtY); 
}

All of my constructors are default. Any ideas on how to make the code faster? Thank you!

Answer 1

Fix compiler warnings

The compiler will warn about some unused parameters and of comparisons between integers of different signedness. For the latter, the solution is simple; instead of:

for (auto i = 0; i < some.size(); i++) {

Write the type of k explicitly:

for (size_t i = 0; i < some.size(); i++) {

The unused parameters of build_tree() should be removed.

Move kdnode inside kdtree

Since kdnode is just an implementation detail of kdtree, move it into class kdtree, and to avoid needless repetition, I would rename it node:

class kdtree {
public:
    class node {
        ...
    };

    ...
    shared_ptr<node> root;
    ...
};

Don't repeat yourself

Avoid repeating code and typenames. For example:

std::shared_ptr<kdnode> newnode = std::make_shared<kdnode>();

Either try to specify the type only on the left hand side:

std::shared_ptr<kdnode> newnode{new kdnode};

This still requires specifying kdnode twice though. Or only on the right hand side:

auto newnode = std::make_shared<kdnode>();

If you see you are using dim_max_.size() a lot, create a variable for it:

auto size = dim_max_.size();
...
Eigen::MatrixXd XtX_(size + 1, size + 1);
...
for (size_t k = 0; k < size; k++) {
    ...
}

Use const references where appropriate

I see some std::vector's being passed by value, which is almost always wrong. Since they are input arguments, pass them by const reference. This avoids unnecessary copies of a potentially large arrays being made.

It looks like you are modifying dim_max_ and dim_min_ in kdtree::build_tree, but then you copy them into newnode->dim_max and nedwnode->dim_min after modification. Just copy first (which you already do), and then modify the copy instead.

kdtree::build_tree() always returns a node

The first line of kdtree::build_tree() always constructs a new kdnode. If end == start, you then return the shared pointer to this new node. However, later on you recursively call build_tree(), and check whether the return value is an empty shared pointer, so that makes me believe you actually intended to return an empty shared pointer. To do that, write the following:

std::shared_ptr<kdnode> kdtree::build_tree(...) {
   if (end == start)
       return {};

   auto newnode = std::make_shared<kdnode>();
   ...
}

Don't write unnecessary comments

The comments you wrote are not very good. They are almost all completely redundant. For example:

kdnode(); // constructor 
kdnode(kdnode&&); // move 
~kdnode(); // destructor    

Yes, those functions are the (move) constructors and the destructor. That should be obvious from the function name and the fact that there is no type in front of them.

// test functions;
void test_XtX(Eigen::MatrixXd);
...

I already see that those are functions, and they have test in the name, so I can already see that they are test functions.

// functions 
all_point_t convert_to_vector(const Eigen::MatrixXd& XY_mat);
...

That is the most pointless comment I have ever seen.

// R functions
Eigen::VectorXd loclinear_i(const Eigen::MatrixXd& XY_mat, int method, int kcode, 
                            double epsilon, const Eigen::VectorXd& h, int N_min);

What is an "R" function? The only thing that pops into my mind when I read that is the R programming language, but clearly this is just a C++ function.

// for bandwidth selection
...

Ok, here there might be some value. But do these functions themselves calculate bandwidth, or are they just generic helper functions? A function like max_weight() sounds like it could be used for many things, not just "bandwidth selection". Is this comment really helpful?

size_t l_len = len/2;         // left length
size_t r_len = len - l_len;   // right length
auto middle = start + len/2;  // middle iterator 

The comments just mirror the names of the variables. If you think l_len is not clear enough, then instead of adding a comment, change the variable names instead: size_t left_length = ..., size_t right_length = ..., and so on.

newnode->XtX = newnode->left_child->XtX + newnode->right_child->XtX;  // sumY = the sum of the bottom 2 nodes

This one is very confusing! What is sumY? I don't see that variable name anywhere! If you mean newnode->XtX is the sum of its child nodes, which are also at the bottom of the tree, then that does look redundant. And why only add this comment here? Is this line special, or the fact that you sum is special?

while (w_max != w_min) {  // if condition fufilled

This is also redundant. Of course a while-statement will run while its condition is fulfilled.

weights = calculate_weight(kcode, query, curr->dim_max, curr->dim_min, h);  // calculate max and min weight 

The function name suggests you are calculating a weight. The comment says you are calculating the max and min weight. Maybe you should change the name of the function instead? Also, the result is a std::pair, if possible use C++17's structured bindings to decompose it directly into two variables, min_weight and max_weight. Ideally this line should look something like:

auto [min_weight, max_weight] = calculate_weight_range(...);

And yes, make the minimum come before the maximum, as that is much more common, so is less surprising.

Write useful comments

What is missing is some comments with high-level descriptions of what you are doing, and why. For example, there are three different scenarios in kdtree::build_tree(): either end == start, end <= N_min or end > N_min. In the first case, there is nothing to do because the range is empty. So you could write:

if (end == start) {
    // The range is empty, so we do not have to create a new node.
    return {};
}

Do something similar for the other two conditions: in the case of end <= N_min, mention that you only need to create a signle node, and if end > N_min, that you bisect the range and recursively iterate over both halves.

In kdtree::get_XtXXtY(), explain how you are traversing the tree structure to get the answer you want.

Prefer std::unique_ptr over std::shared_ptr

You should use std::shared_ptr if there are multiple owners of an object, which possibly non-overlapping ownership lifetimes. However, in the case of a tree structure, there is a clear hierarchy of ownership, and a std::shared_ptr is overkill. Use a std::unique_ptr instead for kdnode::left_child, kdnode::right_child and kdtree::root.

In kdtree::get_XtXXtY(), you have to maintain a std::stack of pointers to nodes. This stack is not owning any nodes, it's just to contain references to existing nodes that are guaranteed to exist while this function is still running, so you can just store raw pointers:

std::stack<kdnode *> storage;
auto curr = root.get();
...
storage.push(curr.get());
...
curr = storage.pop();