include/lshkit/mplsh.h

Go to the documentation of this file.

/* 
    Copyright (C) 2008 Wei Dong <wdong@princeton.edu>. All Rights Reserved.
  
    This file is part of LSHKIT.
  
    LSHKIT is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    LSHKIT is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with LSHKIT.  If not, see <http://www.gnu.org/licenses/>.
*/

#ifndef __LSHKIT_PROBE__
#define __LSHKIT_PROBE__

#include <lshkit/common.h>
#include <lshkit/lsh.h>
#include <lshkit/composite.h>
#include <lshkit/metric.h>
#include <lshkit/lsh-index.h>
#include <lshkit/mplsh-model.h>
#include <lshkit/topk.h>

namespace lshkit
{

struct Probe
{
    unsigned mask;
    unsigned shift;
    float score;
    unsigned reserve;
    bool operator < (const Probe &p) const { return score < p.score; }
    Probe operator + (const Probe &m) const
    {
        Probe ret;
        ret.mask = mask | m.mask;
        ret.shift = shift | m.shift;
        ret.score = score + m.score;
        return ret;
    }
    bool conflict (const Probe &m)
    {
        return (mask & m.mask) != 0;
    }
    static const unsigned MAX_M = 20;
    static const unsigned MAX_T = 200;
}; 

typedef std::vector<Probe> ProbeSequence;

void GenProbeSequenceTemplate (ProbeSequence &seq, unsigned M, unsigned T);

class ProbeSequenceTemplates: public std::vector<ProbeSequence>
{
public:
    ProbeSequenceTemplates(unsigned max_M, unsigned max_T)
        : std::vector<ProbeSequence>(max_M + 1)
    {
        for (unsigned i = 1; i <= max_M; ++i)
        {
            GenProbeSequenceTemplate(at(i), i, max_T);
        }
    }
};

extern ProbeSequenceTemplates __probeSequenceTemplates;

class MultiProbeLsh: public RepeatHash<GaussianLsh> 
{
    unsigned H_;
public:
    typedef RepeatHash<GaussianLsh> Super;
    typedef Super::Domain Domain;

    struct Parameter : public Super::Parameter {

        unsigned range;

        template<class Archive>
        void serialize(Archive & ar, const unsigned int version)
        {
            ar & range;
            ar & repeat;
            ar & dim;
            ar & W;
        }
    };

    MultiProbeLsh () {}

    template <typename RNG>
    void reset(const Parameter &param, RNG &rng)
    {
        H_ = param.range;
        Super::reset(param, rng);
    }

    template <typename RNG>
    MultiProbeLsh(const Parameter &param, RNG &rng)
    {
        H_ = param.range;
        Super::reset(param, rng);
    }

    unsigned getRange () const
    {
        return H_;
    }

    unsigned operator () (Domain obj) const
    {
        return Super::operator ()(obj) % H_;
    }

    template<class Archive>
    void serialize(Archive & ar, const unsigned int version)
    {
        Super::serialize(ar, version);
        ar & H_;
    }

    void genProbeSequence (Domain obj, std::vector<unsigned> &seq, unsigned T) const;
};


template <typename KEY>
class MultiProbeLshIndex: public LshIndex<MultiProbeLsh, KEY>
{
public:
    typedef LshIndex<MultiProbeLsh, KEY> Super;
    typedef typename Super::Parameter Parameter;

private:

    Parameter param_;
    MultiProbeLshRecallTable recall_;

public: 
    typedef typename Super::Domain Domain;
    typedef KEY Key;

    MultiProbeLshIndex() {
    } 


    template <typename Engine>
    void init (const Parameter &param, Engine &engine, unsigned L) {
        Super::init(param, engine, L);
        param_ = param;
        // we are going to normalize the distance by window size, so here we pass W = 1.0.
        // We tune adaptive probing for KNN distance range [0.0001W, 20W].
        recall_.reset(MultiProbeLshModel(Super::lshs_.size(), 1.0, param_.repeat, Probe::MAX_T), 200, 0.0001, 20.0);
    }

    void load (std::istream &ar) 
    {
        Super::load(ar);
        param_.serialize(ar, 0);
        recall_.load(ar);
        verify(ar);
    }

    void save (std::ostream &ar)
    {
        Super::save(ar);
        param_.serialize(ar, 0);
        recall_.save(ar);
        verify(ar);
    }


    template <typename SCANNER>
    void query (Domain obj, unsigned T, SCANNER &scanner)
    {
        std::vector<unsigned> seq;
        for (unsigned i = 0; i < Super::lshs_.size(); ++i) {
            Super::lshs_[i].genProbeSequence(obj, seq, T);
            for (unsigned j = 0; j < seq.size(); ++j) {
                typename Super::Bin &bin = Super::tables_[i][seq[j]];
                BOOST_FOREACH(Key key, bin) {
                    scanner(key);
                }
            }
        }
    }


    template <typename SCANNER>
    void query_recall (Domain obj, float recall, SCANNER &scanner) const
    {
        unsigned K = scanner.topk().getK();
        if (K == 0) throw std::logic_error("CANNOT ACCEPT R-NN QUERY");
        if (scanner.topk().size() < K) throw std::logic_error("ERROR");
        unsigned L = Super::lshs_.size();
        std::vector<std::vector<unsigned> > seqs(L);
        for (unsigned i = 0; i < L; ++i) {
            Super::lshs_[i].genProbeSequence(obj, seqs[i], Probe::MAX_T);
        }

        for (unsigned j = 0; j < Probe::MAX_T; ++j) {
            if (j >= seqs[0].size()) break;
            for (unsigned i = 0; i < L; ++i) {
                BOOST_FOREACH(Key key, Super::tables_[i][seqs[i][j]]) {
                    scanner(key);
                }
            }
            float r = 0.0;
            for (unsigned i = 0; i < K; ++i) {
                r += recall_.lookup(scanner.topk()[i].dist / param_.W, j + 1);
            }
            r /= K;
            if (r >= recall) break;
        }
    }
};

}


#endif

---