Boost RTree with spherical coordinate system

test_spherical.cpp

// This code snippet test rtree with spherical coordinate system.
// It will generate grid points of a cubic grid and store them in an rtree.
// And then query is done by find points within some (approximate) distance
// from query point. The query point can be specified by command line args.
//
// Commandline options:
//     no args: query point (0, 0, 0) with radius 2
//     3 args: query point (x, y, z) with radius 2
//     4 args: query point (x, y, z) with radius r

// Two compilation options:
// #define USE3D // use 2d (phi, theta) or 3d (phi, theta, r) points in rtree
                 // with 2d point, the query area is cone-like
                 // with 3d point, the query area is sphere-like
// #define USE_EQUATORIAL // use sphere_equatorial
// #define USE_CARTESIAN // use normal rtree for comparison, although
                         // the coordinate will still be shperical representation

#define _USE_MATH_DEFINES
#include <cmath>
#include <cstdlib>
#include <iostream>
#include <vector>
#include <tuple>
#include <unordered_set>
#include <boost/geometry.hpp>
#include <boost/geometry/geometries/point.hpp>
#include <boost/geometry/index/rtree.hpp>


using namespace std;
namespace bg = boost::geometry;
namespace bgm = bg::model;
namespace bgi = bg::index;

#ifdef USE_CARTESIAN
typedef bg::cs::cartesian scoord_t;
#else
#ifdef USE_EQUATORIAL
typedef bg::cs::spherical_equatorial<bg::radian> scoord_t;
#else
typedef bg::cs::spherical<bg::radian> scoord_t;
#endif
#endif

#ifdef USE3D
typedef bg::model::point<float, 3, scoord_t> spoint_t; // point in spherical coordinate
#else
typedef bg::model::point<float, 2, scoord_t> spoint_t; // point in spherical coordinate
#endif
typedef std::pair<spoint_t, float> svalue_t;
typedef bg::model::box<spoint_t> sbox_t;
typedef bgi::rstar<1024, 256> params_t;
typedef bgi::rtree<svalue_t, params_t> srtree_t; // rtree in spherical coordinate

spoint_t cartesian_to_spherical(float x, float y, float z, float &r)
{
    float phi = atan2(y, x);
    if (phi < 0)
        phi += 2 * M_PI;
    float rxy2 = x * x + y * y;
#ifdef USE_EQUATORIAL
    float theta = atan2(z, sqrt(rxy2));
#else
    float theta = M_PI / 2 - atan2(z, sqrt(rxy2));
#endif
    r = sqrt(rxy2 + z * z);
#ifdef USE3D
    return {phi, theta, r};
#else
    return {phi, theta};
#endif
}

template <class T>
inline void hash_combine(std::size_t &seed, T const &v)
{
    seed ^= std::hash<T>()(v) + 0x9e3779b9 + (seed << 6) + (seed >> 2);
}

struct spoint_hash
{
    std::size_t operator()(const spoint_t &p) const
    {
        size_t seed = 0;
        hash_combine(seed, bg::get<0>(p));
        hash_combine(seed, bg::get<1>(p));
#ifdef USE3D
        hash_combine(seed, bg::get<2>(p));
#endif
        return seed;
    }
};

struct spoint_eq
{
    bool operator()(spoint_t lhs, spoint_t rhs) const
    {
        return bg::get<0>(lhs) == bg::get<0>(rhs) && bg::get<1>(lhs) == bg::get<1>(rhs)
#ifdef USE3D
               && bg::get<2>(lhs) == bg::get<2>(rhs)
#endif
            ;
    }
};

int main(int argc, char *argv[])
{
    vector<svalue_t> points;
    float _; // dummy
    for (int i = -5; i <= 5; i++)
        for (int j = -5; j <= 5; j++)
            for (int k = -5; k <= 5; k++)
            {
                if (i == 0 && j == 0 && k == 0)
                    continue;
                points.push_back(std::make_pair(cartesian_to_spherical(i, j, k, _), 1));
            }

    srtree_t rtree(points.cbegin(), points.cend());
    cout << rtree.size() << endl;

    spoint_t q;
    float r, radius;
    if (argc == 4)
    {
        q = cartesian_to_spherical(atoi(argv[1]), atoi(argv[2]), atoi(argv[3]), r);
        radius = 2;
    }
    else if (argc == 5)
    {
        q = cartesian_to_spherical(atoi(argv[1]), atoi(argv[2]), atoi(argv[3]), r);
        radius = atof(argv[4]);
    }
    else
    {
        q = cartesian_to_spherical(0, 0, 0, r);
        radius = 2;
    }
    float radius_in_sufrace = radius / r;
#ifdef USE3D
    const spoint_t p1 = {bg::get<0>(q) - radius_in_sufrace, std::max(bg::get<1>(q) - radius_in_sufrace, (float)-M_PI), bg::get<2>(q) - radius};
    const spoint_t p2 = {bg::get<0>(q) + radius_in_sufrace, std::min(bg::get<1>(q) + radius_in_sufrace, (float)M_PI), bg::get<2>(q) + radius};
#else
    const spoint_t p1 = {bg::get<0>(q) - radius_in_sufrace, std::max(bg::get<1>(q) - radius_in_sufrace, (float)-M_PI)};
    const spoint_t p2 = {bg::get<0>(q) + radius_in_sufrace, std::min(bg::get<1>(q) + radius_in_sufrace, (float)M_PI)};
#endif
    const sbox_t bq(p1, p2);

    auto print_point = [](const std::string title, const spoint_t &p) {
#ifdef USE3D
        cout << title << bg::get<0>(p) << ", " << bg::get<1>(p) << ", " << bg::get<2>(p) << endl;
#else
        cout << title << bg::get<0>(p) << ", " << bg::get<1>(p) << endl;
#endif
    };
    print_point("query min: ", p1);
    print_point("query max: ", p2);

    cout << "==========" << endl;
    int aresult = 0;
    std::unordered_set<spoint_t, spoint_hash, spoint_eq> aset;
    for (auto pit = points.begin(); pit != points.end(); pit++)
    {
#ifdef USE3D
        if (bg::get<0>(pit->first) > bg::get<0>(p1) && bg::get<1>(pit->first) > bg::get<1>(p1) && bg::get<2>(pit->first) > bg::get<2>(p1) &&
            bg::get<0>(pit->first) < bg::get<0>(p2) && bg::get<1>(pit->first) < bg::get<1>(p2) && bg::get<2>(pit->first) < bg::get<2>(p2))
#else
        if (bg::get<0>(pit->first) > bg::get<0>(p1) && bg::get<1>(pit->first) > bg::get<1>(p1) &&
            bg::get<0>(pit->first) < bg::get<0>(p2) && bg::get<1>(pit->first) < bg::get<1>(p2))
#endif
        {
            aresult += 1;
            aset.insert(pit->first);
            // print_point("actual point: ", pit->first);
        }
    }
    int qresult = 0;
    std::unordered_set<spoint_t, spoint_hash, spoint_eq> qset;
    for (auto it = rtree.qbegin(bgi::within(bq)); it != rtree.qend(); ++it)
    {
        qresult += it->second;
        qset.insert(it->first);
        // print_point("query point: ", it->first);
    }

    for (auto sit = aset.begin(); sit != aset.end(); sit++)
        if (qset.find(*sit) == qset.end())
            print_point("missed point: ", *sit);
    for (auto sit = qset.begin(); sit != qset.end(); sit++)
        if (aset.find(*sit) == aset.end())
            print_point("extra point: ", *sit);

    cout << "==========" << endl;
    cout << aresult << " actual points" << endl;
    cout << qresult << " query points" << endl;
}