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 <cstddef>
#include <unordered_map>
#include <string>
#include <vector>
#include <array>
 
namespace kaori {
 
template<typename Trie_>
inline void fill_library(const std::vector<const char*>& options, std::unordered_map<std::string, BarcodeIndex>& exact, Trie_& trie, bool reverse) {
    std::size_t len = trie.length();
    auto nopt = options.size();
 
    for (decltype(nopt) i = 0; i < nopt; ++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] = STATUS_UNMATCHED;
            }
        }
    }
 
    trie.optimize();
    return;
}
class SimpleBarcodeSearch {
public:
    struct Options {
        int max_mismatches = 0;
 
        bool reverse = false;
 
        DuplicateAction duplicates = DuplicateAction::ERROR;
    };
 
public:
    SimpleBarcodeSearch() = default;
 
    SimpleBarcodeSearch(const BarcodePool& barcode_pool, const Options& options) : 
        my_trie(barcode_pool.length(), options.duplicates), 
        my_max_mm(options.max_mismatches) 
    {
        fill_library(barcode_pool.pool(), my_exact, my_trie, options.reverse);
    }
 
private:
    AnyMismatches my_trie;
    int my_max_mm;
    std::unordered_map<std::string, BarcodeIndex> my_exact;
 
    struct CacheEntry {
        CacheEntry() = default;
        CacheEntry(BarcodeIndex index, int mismatches) : index(index), mismatches(mismatches) {}
        BarcodeIndex index;
        int mismatches;
    };
    std::unordered_map<std::string, CacheEntry> my_cache;
 
public:
    struct State {
        BarcodeIndex index = 0;
 
        int mismatches = 0;
 
        std::unordered_map<std::string, CacheEntry> cache;
    };
 
    State initialize() const {
        return State();
    }
 
    void reduce(State& state) {
        my_cache.merge(state.cache);
        state.cache.clear();
    }
 
public:
    void search(const std::string& search_seq, State& state) const {
        search(search_seq, state, my_max_mm);
        return;
    }
 
    void search(const std::string& search_seq, State& state, int allowed_mismatches) const {
        auto it = my_exact.find(search_seq);
        if (it != my_exact.end()) {
            state.index = it->second;
            state.mismatches = 0;
            return;
        }
 
        auto set_from_cache = [&](const CacheEntry& cached) -> void {
            if (cached.mismatches > allowed_mismatches) {
                // technically cached.mismatches is only a lower bound if index == UNMATCHED,
                // but if it's already UNMATCHED, then the result will be UNMATCHED either way.
                state.index = STATUS_UNMATCHED;
            } else {
                state.index = cached.index;
            }
            state.mismatches = cached.mismatches;
        };
 
        auto cIt = my_cache.find(search_seq);
        if (cIt != my_cache.end()) {
            set_from_cache(cIt->second);
            return;
        }
 
        auto lIt = state.cache.find(search_seq);
        if (lIt != state.cache.end()) {
            set_from_cache(lIt->second);
            return;
        }
 
        auto missed = my_trie.search(search_seq.c_str(), allowed_mismatches);
        if (is_barcode_index_ok(missed.index)) {
            // No need to check against allowed_mismatches, as we explicitly searched for that in the trie.
            state.index = missed.index;
            state.mismatches = missed.mismatches;
            state.cache[search_seq] = CacheEntry(missed.index, missed.mismatches);
            return;
        }
 
        // The trie search breaks early when it hits allowed_mismatches, but
        // allowed_mismatches might be different across calls to search(). 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 number of mismatches that was specified in the constructor.
        //
        // Of course, if the search failed because of ambiguity, then it would
        // have failed even if we were searching with maximum mismatches;
        // so we happily cache that.
        if (allowed_mismatches == my_max_mm || missed.index == STATUS_AMBIGUOUS) {
            state.cache[search_seq] = CacheEntry(missed.index, missed.mismatches);
        }
 
        state.index = missed.index;
        state.mismatches = missed.mismatches;
    }
};
 
template<int 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() = default;
 
    SegmentedBarcodeSearch(
        const BarcodePool& barcode_pool, 
        std::array<SeqLength, num_segments_> segments, 
        const Options& options
    ) : 
        my_trie(
            [&]{
                if (options.reverse) {
                    std::reverse(segments.begin(), segments.end());
                }
                return segments;
            }(),
            options.duplicates
        ), 
        my_max_mm(
            [&]{
                auto copy = options.max_mismatches;
                if (options.reverse) {
                    std::reverse(copy.begin(), copy.end());
                }
                return copy;
            }()
        )
    {
        if (barcode_pool.length() != my_trie.length()) {
            throw std::runtime_error("variable sequences should have the same length as the sum of segment lengths");
        }
        fill_library(barcode_pool.pool(), my_exact, my_trie, options.reverse);
    }
 
private:
    SegmentedMismatches<num_segments_> my_trie;
    std::array<int, num_segments_> my_max_mm;
    std::unordered_map<std::string, BarcodeIndex> my_exact;
 
    struct CacheEntry {
        CacheEntry() = default;
        CacheEntry(BarcodeIndex index, int mismatches, std::array<int, num_segments_> per_segment) :
            index(index), mismatches(mismatches), per_segment(per_segment) {}
        BarcodeIndex index;
        int mismatches;
        std::array<int, num_segments_> per_segment;
    };
    std::unordered_map<std::string, CacheEntry> my_cache;
 
public:
    struct State {
        BarcodeIndex index = 0;
 
        int mismatches = 0;
 
        std::array<int, num_segments_> per_segment;
        
        State() : per_segment() {}
 
        std::unordered_map<std::string, CacheEntry> cache;
    };
 
    State initialize() const {
        return State();
    }
 
    void reduce(State& state) {
        my_cache.merge(state.cache);
        state.cache.clear();
    }
 
public:
    void search(const std::string& search_seq, State& state) const {
        search(search_seq, state, my_max_mm);
        return;
    }
 
    void search(const std::string& search_seq, State& state, std::array<int, num_segments_> allowed_mismatches) const {
        auto it = my_exact.find(search_seq);
        if (it != my_exact.end()) {
            state.index = it->second;
            state.mismatches = 0;
            std::fill_n(state.per_segment.begin(), num_segments_, 0);
            return;
        }
 
        auto set_from_cache = [&](const CacheEntry& cached) -> void {
            state.mismatches = cached.mismatches;
            state.per_segment = cached.per_segment;
            for (int s = 0; s < num_segments_; ++s) {
                if (cached.per_segment[s] > allowed_mismatches[s]) {
                    // technically cached.mismatches is only a lower bound if index == UNMATCHED,
                    // but if it's already UNMATCHED, then the result will be UNMATCHED either way.
                    state.index = STATUS_UNMATCHED;
                    return;
                }
            }
            state.index = cached.index;
        };
 
        auto cIt = my_cache.find(search_seq);
        if (cIt != my_cache.end()) {
            set_from_cache(cIt->second);
            return;
        }
 
        auto lIt = state.cache.find(search_seq);
        if (lIt != state.cache.end()) {
            set_from_cache(lIt->second);
            return;
        }
 
        auto missed = my_trie.search(search_seq.c_str(), allowed_mismatches);
        if (is_barcode_index_ok(missed.index)) {
            // No need to check against allowed_mismatches, as we explicitly searched for that in the trie.
            state.index = missed.index;
            state.mismatches = missed.mismatches;
            state.per_segment = missed.per_segment;
            state.cache[search_seq] = CacheEntry(missed.index, missed.mismatches, missed.per_segment);
            return;
        }
 
        // The trie search breaks early when it hits allowed_mismatches, but
        // the allowed_mismatches might be different across calls to search().
        // 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 number of mismatches that was specified in the constructor.
        //
        // Of course, if the search failed because of ambiguity, then it would
        // have failed even if we were searching with maximum mismatches;
        // so we happily cache that.
        if (allowed_mismatches == my_max_mm || missed.index == STATUS_AMBIGUOUS) {
            state.cache[search_seq] = CacheEntry(missed.index, missed.mismatches, missed.per_segment);
        }
 
        state.index = missed.index;
        state.mismatches = missed.mismatches;
        state.per_segment = missed.per_segment;
    }
};
 
}
 
#endif