L2Normalized.hpp

Go to the documentation of this file.

#ifndef KNNCOLLE_L2_NORMALIZED_HP
#define KNNCOLLE_L2_NORMALIZED_HP
 
#include <vector>
#include <cmath>
#include <memory>
#include <limits>
#include <cstddef>
#include <type_traits>
#include <cstring>
#include <string>
#include <filesystem>
#include <algorithm>
 
#include "Searcher.hpp"
#include "Prebuilt.hpp"
#include "Builder.hpp"
#include "Matrix.hpp"
#include "NumericType.hpp"
#include "utils.hpp"
 
namespace knncolle {
 
inline static constexpr const char* l2normalized_prebuilt_save_name = "knncolle::L2Normalized";
 
template<typename Normalized_>
std::function<void(const std::filesystem::path&)>& custom_save_for_l2normalized_normalized() {
    static std::function<void(const std::filesystem::path&)> fun;
    return fun;
}
 
template<typename Data_, typename Normalized_>
void l2norm(const Data_* ptr, std::size_t ndim, Normalized_* buffer) {
    if constexpr(std::is_same<Data_, Normalized_>::value) {
        // If it's the same type, we don't need to do an up-front cast of 'ptr[d]' to Normalized_.
        Normalized_ l2 = 0;
        for (std::size_t d = 0; d < ndim; ++d) {
            l2 += ptr[d] * ptr[d];
        }
 
        if (l2 > 0) {
            l2 = std::sqrt(l2);
            for (std::size_t d = 0; d < ndim; ++d) {
                buffer[d] = ptr[d] / l2;
            }
        } else {
            std::fill_n(buffer, ndim, 0);
        }    
 
    } else {
        // Otherwise, we do the cast first to avoid any surprises, e.g., from the squared value overflowing an integer 'Data_' type.
        std::copy_n(ptr, ndim, buffer);
 
        Normalized_ l2 = 0;
        for (std::size_t d = 0; d < ndim; ++d) {
            l2 += buffer[d] * buffer[d];
        }
 
        if (l2 > 0) {
            l2 = std::sqrt(l2);
            for (std::size_t d = 0; d < ndim; ++d) {
                buffer[d] /= l2;
            }
        }    
    }
}
 
template<typename Index_, typename Data_, typename Distance_, typename Normalized_, class Searcher_>
class L2NormalizedSearcher final : public Searcher<Index_, Data_, Distance_> {
public:
    L2NormalizedSearcher(std::unique_ptr<Searcher_> searcher, std::size_t num_dimensions) : 
        my_searcher(std::move(searcher)),
        buffer(sanisizer::cast<I<decltype(buffer.size())> >(num_dimensions))
    {}
 
private:
    // No way around this; the L2-normalized values must be floating-point,
    // so the internal searcher must accept floats.
    static_assert(std::is_floating_point<Normalized_>::value);
 
    std::unique_ptr<Searcher_> my_searcher;
    std::vector<Normalized_> buffer;
 
public:
    void search(Index_ i, Index_ k, std::vector<Index_>* output_indices, std::vector<Distance_>* output_distances) {
        my_searcher->search(i, k, output_indices, output_distances);
    }
 
    void search(const Data_* ptr, Index_ k, std::vector<Index_>* output_indices, std::vector<Distance_>* output_distances) {
        auto normalized = buffer.data();
        l2norm(ptr, buffer.size(), normalized);
        my_searcher->search(normalized, k, output_indices, output_distances);
    }
 
public:
    bool can_search_all() const {
        return my_searcher->can_search_all();
    }
 
    Index_ search_all(Index_ i, Distance_ threshold, std::vector<Index_>* output_indices, std::vector<Distance_>* output_distances) {
        return my_searcher->search_all(i, threshold, output_indices, output_distances);
    }
 
    Index_ search_all(const Data_* ptr, Distance_ threshold, std::vector<Index_>* output_indices, std::vector<Distance_>* output_distances) {
        auto normalized = buffer.data();
        l2norm(ptr, buffer.size(), normalized);
        return my_searcher->search_all(normalized, threshold, output_indices, output_distances);
    }
};
 
template<typename Index_, typename Data_, typename Distance_>
Prebuilt<Index_, Data_, Distance_>* load_prebuilt_raw(const std::filesystem::path&);
 
template<typename Index_, typename Data_, typename Distance_, typename Normalized_>
class L2NormalizedPrebuilt final : public Prebuilt<Index_, Data_, Distance_> {
public:
    L2NormalizedPrebuilt(std::unique_ptr<Prebuilt<Index_, Normalized_, Distance_> > prebuilt) : my_prebuilt(std::move(prebuilt)) {}
 
private:
    std::unique_ptr<Prebuilt<Index_, Normalized_, Distance_> > my_prebuilt;
 
public:
    Index_ num_observations() const {
        return my_prebuilt->num_observations();
    }
 
    std::size_t num_dimensions() const {
        return my_prebuilt->num_dimensions();
    }
 
public:
    std::unique_ptr<Searcher<Index_, Data_, Distance_> > initialize() const {
        return initialize_known();
    }
 
    auto initialize_known() const {
        typedef I<decltype(*(my_prebuilt->initialize_known()))> KnownSearcher;
        return std::make_unique<L2NormalizedSearcher<Index_, Data_, Distance_, Normalized_, KnownSearcher> >(my_prebuilt->initialize_known(), my_prebuilt->num_dimensions());
    }
 
public:
    void save(const std::filesystem::path& dir) const {
        quick_save(dir / "ALGORITHM", l2normalized_prebuilt_save_name, std::strlen(l2normalized_prebuilt_save_name));
        auto norm_type = get_numeric_type<Normalized_>();
        quick_save(dir / "NORMALIZED", &norm_type, 1);
 
        auto& cust = custom_save_for_l2normalized_normalized<Normalized_>();
        if (cust) {
            cust(dir);
        }
 
        const auto indexdir = dir / "INDEX";
        std::filesystem::create_directory(indexdir);
        my_prebuilt->save(indexdir);
    }
 
    L2NormalizedPrebuilt(const std::filesystem::path& dir) : my_prebuilt(load_prebuilt_raw<Index_, Normalized_, Distance_>(dir / "INDEX")) {}
};
 
template<typename Index_, typename Data_, typename Normalized_, typename Matrix_>
class L2NormalizedMatrix;
 
template<typename Index_, typename Data_, typename Normalized_, class Extractor_ = MatrixExtractor<Data_> >
class L2NormalizedMatrixExtractor final : public MatrixExtractor<Normalized_> {
public:
    L2NormalizedMatrixExtractor(std::unique_ptr<Extractor_> extractor, std::size_t dim) : 
        my_extractor(std::move(extractor)),
        buffer(sanisizer::cast<I<decltype(buffer.size())> >(dim))
    {}
 
private:
    std::unique_ptr<Extractor_> my_extractor;
    std::vector<Normalized_> buffer;
 
public:
    const Normalized_* next() {
        auto raw = my_extractor->next();
        auto normalized = buffer.data();
        l2norm(raw, buffer.size(), normalized);
        return normalized;
    }
};
template<typename Index_, typename Data_, typename Normalized_, typename Matrix_ = Matrix<Index_, Data_> >
class L2NormalizedMatrix final : public Matrix<Index_, Normalized_> {
public:
    L2NormalizedMatrix(const Matrix_& matrix) : my_matrix(matrix) {}
private:
    static_assert(std::is_same<decltype(std::declval<Matrix_>().num_observations()), Index_>::value);
    static_assert(std::is_same<typename std::remove_pointer<decltype(std::declval<Matrix_>().new_extractor()->next())>::type, const Data_>::value);
 
    const Matrix_& my_matrix;
 
public:
    std::size_t num_dimensions() const {
        return my_matrix.num_dimensions();
    }
 
    Index_ num_observations() const {
        return my_matrix.num_observations();
    }
 
    auto new_known_extractor() const {
        typedef I<decltype(*(my_matrix.new_known_extractor()))> KnownExtractor;
        return std::make_unique<L2NormalizedMatrixExtractor<Index_, Data_, Normalized_, KnownExtractor> >(my_matrix.new_known_extractor(), num_dimensions());
    }
 
    std::unique_ptr<MatrixExtractor<Normalized_> > new_extractor() const {
        return new_known_extractor();
    }
};
 
template<typename Index_, typename Data_, typename Distance_, typename Normalized_, class Matrix_ = Matrix<Index_, Data_> >
class L2NormalizedBuilder final : public Builder<Index_, Data_, Distance_, Matrix_> {
public:
    typedef L2NormalizedMatrix<Index_, Data_, Normalized_, Matrix_> NormalizedMatrix;
 
    typedef typename std::conditional<
        std::is_base_of<Matrix_, NormalizedMatrix>::value,
        Matrix_,
        NormalizedMatrix
    >::type BuilderMatrix;
 
public:
    L2NormalizedBuilder(std::shared_ptr<const Builder<Index_, Normalized_, Distance_, BuilderMatrix> > builder) : my_builder(std::move(builder)) {}
 
private:
    std::shared_ptr<const Builder<Index_, Normalized_, Distance_, BuilderMatrix> > my_builder;
 
public:
    Prebuilt<Index_, Data_, Distance_>* build_raw(const Matrix_& data) const {
        return build_known_raw(data);
    }
public:
    auto build_known_raw(const Matrix_& data) const {
        NormalizedMatrix normalized(data);
        return new L2NormalizedPrebuilt<Index_, Data_, Distance_, Normalized_>(my_builder->build_unique(normalized));
    }
 
    auto build_known_unique(const Matrix_& data) const {
        return std::unique_ptr<I<decltype(*(build_known_raw(data)))> >(build_known_raw(data));
    }
 
    auto build_known_shared(const Matrix_& data) const {
        return std::shared_ptr<I<decltype(*(build_known_raw(data)))> >(build_known_raw(data));
    }
};
 
}
 
#endif