distances.hpp

Go to the documentation of this file.

#ifndef KNNCOLLE_HNSW_DISTANCES_HPP
#define KNNCOLLE_HNSW_DISTANCES_HPP
 
#include <cmath>
#include <functional>
#include <cstddef>
 
namespace knncolle_hnsw {
 
template<typename HnswData_ = float>
struct DistanceConfig {
    std::function<hnswlib::SpaceInterface<HnswData_>*(std::size_t)> create;
 
    std::function<HnswData_(HnswData_)> normalize;
};
 
template<typename HnswData_ = float>
class ManhattanDistance : public hnswlib::SpaceInterface<HnswData_> {
private:
    std::size_t my_data_size;
    std::size_t my_dim;
 
public:
    ManhattanDistance(std::size_t dim) : my_data_size(dim * sizeof(HnswData_)), my_dim(dim) {}
 
public:
    std::size_t get_data_size() {
        return my_data_size;
    }
 
    hnswlib::DISTFUNC<HnswData_> get_dist_func() {
        return L1;
    }
 
    void * get_dist_func_param() {
        return &my_dim;
    }
 
private:
    static HnswData_ L1(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
        const HnswData_* pVect1 = static_cast<const HnswData_*>(pVect1v);
        const HnswData_* pVect2 = static_cast<const HnswData_*>(pVect2v);
        std::size_t qty = *((size_t *) qty_ptr);
        HnswData_ res = 0;
        for (; qty > 0; --qty, ++pVect1, ++pVect2) {
            res += std::abs(*pVect1 - *pVect2);
        }
        return res;
    }
};
 
template<typename HnswData_ = float>
class SquaredEuclideanDistance : public hnswlib::SpaceInterface<HnswData_> {
private:
    std::size_t my_data_size;
    std::size_t my_dim;
 
public:
    SquaredEuclideanDistance(std::size_t dim) : my_data_size(dim * sizeof(HnswData_)), my_dim(dim) {}
 
public:
    std::size_t get_data_size() {
        return my_data_size;
    }
 
    hnswlib::DISTFUNC<HnswData_> get_dist_func() {
        return L2;
    }
 
    void * get_dist_func_param() {
        return &my_dim;
    }
 
private:
    static HnswData_ L2(const void *pVect1v, const void *pVect2v, const void *qty_ptr) {
        const HnswData_* pVect1 = static_cast<const HnswData_*>(pVect1v);
        const HnswData_* pVect2 = static_cast<const HnswData_*>(pVect2v);
        std::size_t qty = *((size_t *) qty_ptr);
        HnswData_ res = 0;
        for (; qty > 0; --qty, ++pVect1, ++pVect2) {
            auto delta = *pVect1 - *pVect2;
            res += delta * delta;
        }
        return res;
    }
};
 
template<typename HnswData_ = float>
DistanceConfig<HnswData_> makeEuclideanDistanceConfig() {
    DistanceConfig<HnswData_> output;
    output.create = [](std::size_t dim) -> hnswlib::SpaceInterface<HnswData_>* {
        if constexpr(std::is_same<HnswData_, float>::value) {
            return static_cast<hnswlib::SpaceInterface<HnswData_>*>(new hnswlib::L2Space(dim));
        } else {
            return static_cast<hnswlib::SpaceInterface<HnswData_>*>(new SquaredEuclideanDistance<HnswData_>(dim));
        }
    };
    output.normalize = [](HnswData_ x) -> HnswData_ {
        return std::sqrt(x);
    };
    return output;
}
 
template<typename HnswData_ = float>
DistanceConfig<HnswData_> makeManhattanDistanceConfig() {
    DistanceConfig<HnswData_> output;
    output.create = [](std::size_t dim) -> hnswlib::SpaceInterface<HnswData_>* {
        return static_cast<hnswlib::SpaceInterface<HnswData_>*>(new knncolle_hnsw::ManhattanDistance<HnswData_>(dim));
    };
    return output;
}
 
}
 
#endif