IndexIVF_HNSW.h

#ifndef IVF_HNSW_LIB_IVF_HNSW_H
#define IVF_HNSW_LIB_IVF_HNSW_H

#include <iostream>
#include <fstream>
#include <cstdio>
#include <unordered_map>

#include <faiss/index_io.h>
#include <faiss/Heap.h>
#include <faiss/ProductQuantizer.h>
#include <faiss/VectorTransform.h>
#include <faiss/FaissAssert.h>
#include <faiss/utils.h>

#include <hnswlib/hnswalg.h>
#include "utils.h"

namespace ivfhnsw {
    struct IndexIVF_HNSW
    {
        typedef uint32_t idx_t;     

        size_t d;               
        size_t nc;              
        size_t code_size;       

        hnswlib::HierarchicalNSW *quantizer; 

        faiss::ProductQuantizer *pq;         
        faiss::ProductQuantizer *norm_pq;    
        faiss::LinearTransform *opq_matrix;  
        bool do_opq;                         

        size_t nprobe;        
        size_t max_codes;     

        std::vector<std::vector<idx_t> > ids;           
        std::vector<std::vector<uint8_t> > codes;       
        std::vector<std::vector<uint8_t> > norm_codes;  

    protected:
        std::vector<float> norms;           
        std::vector<float> centroid_norms;  

    public:
        explicit IndexIVF_HNSW(size_t dim, size_t ncentroids, size_t bytes_per_code, size_t nbits_per_idx);
        virtual ~IndexIVF_HNSW();

        void build_quantizer(const char *path_data, const char *path_info, const char *path_edges,
                             size_t M=16, size_t efConstruction = 500);

        void assign (size_t n, const float *x, idx_t *labels, size_t k = 1);

        virtual void search(size_t k, const float *x, float *distances, long *labels);

        virtual void add_batch(size_t n, const float *x, const idx_t *xids, const idx_t *precomputed_idx = nullptr);

        virtual void train_pq(size_t n, const float *x);

        virtual void write(const char *path);

        virtual void read(const char *path);

        void compute_centroid_norms();

        void rotate_quantizer();

    protected:
        std::vector<float> precomputed_table;

        float pq_L2sqr(const uint8_t *code);

    private:
        void reconstruct(size_t n, float *x, const float *decoded_residuals, const idx_t *keys);
        void compute_residuals(size_t n, const float *x, float *residuals, const idx_t *keys);
    };
}
#endif //IVF_HNSW_LIB_INDEX_HNSW_H