BarcodeSearch.hpp

Go to the documentation of this file.

#ifndef KAORI_VARIABLE_LIBRARY_HPP
#define KAORI_VARIABLE_LIBRARY_HPP
 
#include "BarcodePool.hpp"
#include "MismatchTrie.hpp"
#include "utils.hpp"
#include <unordered_map>
#include <string>
#include <vector>
#include <array>
 
namespace kaori {
 
template<class Trie>
void fill_library(
    const std::vector<const char*>& options, 
    std::unordered_map<std::string, int>& exact,
    Trie& trie,
    bool reverse
) {
    size_t len = trie.get_length();
 
    for (size_t i = 0; i < options.size(); ++i) {
        auto ptr = options[i];
 
        std::string current;
        if (!reverse) {
            current = std::string(ptr, ptr + len);
        } else {
            current.reserve(len);
            for (size_t j = 0; j < len; ++j) {
                current += complement_base<true, true>(ptr[len - j - 1]);
            }
        }
 
        // Note that this must be called, even if the sequence is duplicated;
        // otherwise the trie's internal counter will not be properly incremented.
        auto status = trie.add(current.c_str());
 
        if (!status.has_ambiguous) {
            if (!status.is_duplicate || status.duplicate_replaced) {
                exact[current] = i;
            } else if (status.duplicate_cleared) {
                exact[current] = -1;
            }
        }
    }
 
    trie.optimize();
    return;
}
 
template<class Methods, class Cache, class Trie, class Result, class Mismatch>
void matcher_in_the_rye(const std::string& x, const Cache& cache, const Trie& trie, Result& res, const Mismatch& mismatches, const Mismatch& max_mismatches) {
    // Seeing if it's any of the caches; otherwise searching the trie.
    auto cit = cache.find(x);
    if (cit == cache.end()) {
        auto lit = res.cache.find(x);
        if (lit != res.cache.end()) {
            Methods::update(res, lit->second, mismatches);
 
        } else {
            auto missed = trie.search(x.c_str(), mismatches);
 
            // The trie search breaks early when it hits the mismatch cap,
            // but the cap might be different across calls. If we break
            // early and report a miss, the miss will be cached and
            // returned in cases where there is a higher cap (and thus
            // might actually be a hit). As such, we should only store a
            // miss in the cache when the requested number of mismatches is
            // equal to the maximum value specified in the constructor.
            if (Methods::index(missed) >= 0 || mismatches == max_mismatches) {
                res.cache[x] = missed;
            }
 
            // No need to pass the requested number of mismatches,
            // as we explicitly searched for that in the trie.
            Methods::update(res, missed);
        }
    } else {
        Methods::update(res, cit->second, mismatches);
    }
    return;
}
class SimpleBarcodeSearch {
public:
    struct Options {
        int max_mismatches = 0;
 
        bool reverse = false;
 
        DuplicateAction duplicates = DuplicateAction::ERROR;
    };
 
public:
    SimpleBarcodeSearch() {}
 
    SimpleBarcodeSearch(const BarcodePool& barcode_pool, const Options& options) : 
        trie(barcode_pool.length, options.duplicates), 
        max_mm(options.max_mismatches) 
    {
        fill_library(barcode_pool.pool, exact, trie, options.reverse);
        return;
    }
 
public:
    struct State {
        int index = 0;
 
        int mismatches = 0;
        
        std::unordered_map<std::string, std::pair<int, int> > cache;
    };
 
    State initialize() const {
        return State();
    }
 
    void reduce(State& state) {
        cache.merge(state.cache);
        state.cache.clear();
    }
 
private:
    struct Methods {
        static int index(const std::pair<int, int>& val) {
            return val.first;
        }
 
        static void update(State& state, const std::pair<int, int>& val) {
            state.index = val.first;
            state.mismatches = val.second;
            return;
        }
 
        static void update(State& state, const std::pair<int, int>& val, int mismatches) {
            state.index = (val.second > mismatches ? -1 : val.first);
            state.mismatches = val.second;
            return;
        }
    };
 
public:
    void search(const std::string& search_seq, State& state) const {
        search(search_seq, state, max_mm);
        return;
    }
 
    void search(const std::string& search_seq, State& state, int allowed_mismatches) const {
        auto it = exact.find(search_seq);
        if (it != exact.end()) {
            state.index = it->second;
            state.mismatches = 0;
        } else {
            matcher_in_the_rye<Methods>(search_seq, cache, trie, state, allowed_mismatches, max_mm);
        }
    }
 
private:
    std::unordered_map<std::string, int> exact;
    AnyMismatches trie;
    std::unordered_map<std::string, std::pair<int, int> > cache;
    int max_mm;
};
 
// Using some template recursion instead of a fixed-length loop.
// Maybe it compiles down to the same thing. 
template<size_t total, size_t position>
struct HasMore {
    static bool check(const std::array<int, total>& left, const std::array<int, total>& right) {
        return (HasMore<total, position + 1>::check(left, right) || left[position] > right[position]);
    }
};
 
template<size_t total>
struct HasMore<total, total> {
    static bool check(const std::array<int, total>&, const std::array<int, total>&) { return false; }
};
template<size_t num_segments>
class SegmentedBarcodeSearch {
public:
    struct Options {
        Options(int max_mismatch_per_segment = 0) {
            max_mismatches.fill(max_mismatch_per_segment);
        }
        
        std::array<int, num_segments> max_mismatches;
 
        bool reverse = false;
 
        DuplicateAction duplicates = DuplicateAction::ERROR;
    };
 
public:
    SegmentedBarcodeSearch() {}
 
    SegmentedBarcodeSearch(
        const BarcodePool& barcode_pool, 
        std::array<int, num_segments> segments, 
        const Options& options
    ) : 
        trie(
            (!options.reverse ? 
                segments :
                [&]{
                    auto copy = segments;
                    std::reverse(copy.begin(), copy.end());
                    return copy;
                }()
            ),
            options.duplicates
        ), 
        max_mm(
            (!options.reverse ?
                options.max_mismatches :
                [&]{
                    auto copy = options.max_mismatches;
                    std::reverse(copy.begin(), copy.end());
                    return copy;
                }()
            )
        )
    {
        if (barcode_pool.length != trie.get_length()) {
            throw std::runtime_error("variable sequences should have the same length as the sum of segment lengths");
        }
        fill_library(barcode_pool.pool, exact, trie, options.reverse);
        return;
    }
 
public:
    struct State {
        int index = 0;
 
        int mismatches = 0;
 
        std::array<int, num_segments> per_segment;
        
        State() : per_segment() {}
 
        std::unordered_map<std::string, typename SegmentedMismatches<num_segments>::Result> cache;
    };
 
    State initialize() const {
        return State();
    }
 
    void reduce(State& state) {
        cache.merge(state.cache);
        state.cache.clear();
    }
 
private:
    typedef typename SegmentedMismatches<num_segments>::Result SegmentedResult;
 
    struct Methods {
        static int index(const SegmentedResult& val) {
            return val.index;
        }
 
        static void update(State& state, const SegmentedResult& val) {
            state.index = val.index;
            state.mismatches = val.total;
            state.per_segment = val.per_segment;
            return;
        }
 
        static void update(State& state, const SegmentedResult& val, const std::array<int, num_segments>& mismatches) {
            state.index = (HasMore<num_segments, 0>::check(val.per_segment, mismatches) ? -1 : val.index);
            state.mismatches = val.total;
            state.per_segment = val.per_segment;
            return;
        }
    };
 
public:
    void search(const std::string& search_seq, State& state) const {
        search(search_seq, state, max_mm);
        return;
    }
 
    void search(const std::string& search_seq, State& state, std::array<int, num_segments> allowed_mismatches) const {
        auto it = exact.find(search_seq);
        if (it != exact.end()) {
            state.index = it->second;
            state.mismatches = 0;
            std::fill_n(state.per_segment.begin(), num_segments, 0);
        } else {
            matcher_in_the_rye<Methods>(search_seq, cache, trie, state, allowed_mismatches, max_mm);
        }
    }
 
private:
    std::unordered_map<std::string, int> exact;
    SegmentedMismatches<num_segments> trie;
    std::unordered_map<std::string, SegmentedResult> cache;
    std::array<int, num_segments> max_mm;
};
 
}
 
#endif