KDTree.h

/*
 * (C) Copyright 2013 ECMWF.
 *
 * This software is licensed under the terms of the Apache Licence Version 2.0
 * which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
 * In applying this licence, ECMWF does not waive the privileges and immunities
 * granted to it by virtue of its status as an intergovernmental organisation
 * nor does it submit to any jurisdiction.
 */

#pragma once

#include "atlas/util/Geometry.h"
#include "atlas/util/ObjectHandle.h"
#include "atlas/util/detail/KDTree.h"

namespace atlas {
namespace util {

//------------------------------------------------------------------------------------------------------


template <typename PayloadT, typename PointT = Point3>
class KDTree : public ObjectHandle<detail::KDTreeBase<PayloadT, PointT>> {
public:
    using Handle         = ObjectHandle<detail::KDTreeBase<PayloadT, PointT>>;
    using Implementation = typename Handle::Implementation;
    using Point          = typename Implementation::Point;
    using Payload        = typename Implementation::Payload;
    using PayloadList    = typename Implementation::PayloadList;
    using Value          = typename Implementation::Value;
    using ValueList      = typename Implementation::ValueList;

    using Handle::get;
    using Handle::Handle;

public:
    //--------------------------------------------------------------------------------------
    // Constructors

    KDTree( const KDTree& handle ) : Handle( handle ) {}

    KDTree() : Handle( new detail::KDTreeMemory<Payload, Point>() ) {}

    KDTree( const Geometry& geometry ) : Handle( new detail::KDTreeMemory<Payload, Point>( geometry ) ) {}

    template <typename Tree>
    KDTree( const std::shared_ptr<Tree>& kdtree ) :
        Handle( new detail::KDTree_eckit<Tree, Payload, Point>( kdtree ) ) {}

    template <typename Tree>
    KDTree( const std::shared_ptr<Tree>& kdtree, const Geometry& geometry ) :
        Handle( new detail::KDTree_eckit<Tree, Payload, Point>( kdtree, geometry ) ) {}

    //--------------------------------------------------------------------------------------
    // Methods to build the KDTree

    void reserve( idx_t size ) { get()->reserve( size ); }

    template <typename Point>
    void insert( const Point& p, const Payload& payload ) {
        get()->insert( p, payload );
    }

    void insert( const Value& value ) { get()->insert( value ); }

    void build() { get()->build(); }

    template <typename Longitudes, typename Latitudes, typename Payloads>
    void build( const Longitudes& longitudes, const Latitudes& latitudes, const Payloads& payloads ) {
        get()->build( longitudes, latitudes, payloads );
    }

    template <typename LongitudesIterator, typename LatitudesIterator, typename PayloadsIterator>
    void build( const LongitudesIterator& longitudes_begin, const LongitudesIterator& longitudes_end,
                const LatitudesIterator& latitudes_begin, const LatitudesIterator& latitudes_end,
                const PayloadsIterator& payloads_begin, const PayloadsIterator& payloads_end ) {
        get()->build( longitudes_begin, longitudes_end, latitudes_begin, latitudes_end, payloads_begin, payloads_end );
    }

    template <typename Points, typename Payloads>
    void build( const Points& points, const Payloads& payloads ) {
        get()->build( points, payloads );
    }

    template <typename PointIterator, typename PayloadsIterator>
    void build( const PointIterator& points_begin, const PointIterator& points_end,
                const PayloadsIterator& payloads_begin, const PayloadsIterator& payloads_end ) {
        get()->build( points_begin, points_end, payloads_begin, payloads_end );
    }

    void build( const std::vector<Value>& values ) { get()->build( values ); }

    //--------------------------------------------------------------------------------------
    // Methods to access the KDTree

    bool empty() const { return get()->empty(); }

    size_t size() const { return get()->size(); }

    size_t footprint() const { return get()->footprint(); }

    template <typename Point>
    ValueList closestPoints( const Point& p, size_t k ) const {
        return get()->closestPoints( p, k );
    }

    template <typename Point>
    Value closestPoint( const Point& p ) const {
        return get()->closestPoint( p );
    }

    template <typename Point>
    ValueList closestPointsWithinRadius( const Point& p, double radius ) const {
        return get()->closestPointsWithinRadius( p, radius );
    }

    const Geometry& geometry() const { return get()->geometry(); }
};

//------------------------------------------------------------------------------------------------------

using IndexKDTree2D = KDTree<idx_t, Point2>;  // 2D search: implementation is using 2D points only
using IndexKDTree3D = KDTree<idx_t, Point3>;  // 3D search: lonlat (2D) to xyz (3D) conversion is done internally
using IndexKDTree   = IndexKDTree3D;

// ------------------------------------------------------------------
// C wrapper interfaces to C++ routines

extern "C" {
IndexKDTree::Implementation* atlas__IndexKDTree__new();
IndexKDTree::Implementation* atlas__IndexKDTree__new_geometry( const Geometry::Implementation* geometry );
void atlas__IndexKDTree__delete( IndexKDTree::Implementation* This );
void atlas__IndexKDTree__reserve( IndexKDTree::Implementation* This, const idx_t size );
void atlas__IndexKDTree__insert( IndexKDTree::Implementation* This, const double lon, const double lat,
                                 const idx_t index );
void atlas__IndexKDTree__build( IndexKDTree::Implementation* This );
void atlas__IndexKDTree__closestPoints( const IndexKDTree::Implementation* This, const double plon, const double plat,
                                        const size_t k, double*& lon, double*& lat, idx_t*& indices,
                                        double*& distances );
void atlas__IndexKDTree__closestPoint( const IndexKDTree::Implementation* This, const double plon, const double plat,
                                       double& lon, double& lat, idx_t& index, double& distance );
void atlas__IndexKDTree__closestPointsWithinRadius( const IndexKDTree::Implementation* This, const double plon,
                                                    const double plat, const double radius, size_t& k, double*& lon,
                                                    double*& lat, idx_t*& indices, double*& distances );
const Geometry::Implementation* atlas__IndexKDTree__geometry( const IndexKDTree::Implementation* This );
int atlas__IndexKDTree__empty( const IndexKDTree::Implementation* This );
idx_t atlas__IndexKDTree__size( const IndexKDTree::Implementation* This );
}

//------------------------------------------------------------------------------------------------------

}  // namespace util
}  // namespace atlas