MismatchTrie.hpp

Go to the documentation of this file.

#ifndef KAORI_MISMATCH_TRIE_HPP
#define KAORI_MISMATCH_TRIE_HPP
 
#include <array>
#include <vector>
#include <stdexcept>
#include <numeric>
#include "utils.hpp"
 
namespace kaori {
 
class MismatchTrie {
public:
    static constexpr int STATUS_MISSING = -1;
 
    static constexpr int STATUS_AMBIGUOUS = -2;
 
protected:
    static constexpr int NUM_BASES = 4;
public:
    MismatchTrie() {}
 
    MismatchTrie(size_t barcode_length, DuplicateAction duplicates) : 
        length(barcode_length), 
        pointers(NUM_BASES, STATUS_MISSING), 
        duplicates(duplicates), 
        counter(0) 
    {}
 
public:
    struct AddStatus {
        bool has_ambiguous = false;
 
        bool is_duplicate = false;
 
        bool duplicate_replaced = false;
 
        bool duplicate_cleared = false;
    };
 
private:
    void next(int shift, int& position) {
        auto& current = pointers[position + shift];
        if (current < 0) {
            current = pointers.size();
            position = current;
            pointers.resize(position + NUM_BASES, STATUS_MISSING);
        } else {
            position = current;
        }
    }
 
    void end(int shift, int position, AddStatus& status) {
        auto& current = pointers[position + shift];
        if (current >= 0) {
            status.is_duplicate = true;
            switch(duplicates) {
                case DuplicateAction::FIRST:
                    break;
                case DuplicateAction::LAST:
                    status.duplicate_replaced = true;
                    current = counter;
                    break;
                case DuplicateAction::NONE:
                    status.duplicate_cleared = true;
                    current = STATUS_AMBIGUOUS;
                    break;
                case DuplicateAction::ERROR:
                    throw std::runtime_error("duplicate sequences detected (" + 
                        std::to_string(current + 1) + ", " + 
                        std::to_string(counter + 1) + ") when constructing the trie");
            }
 
        } else if (current == STATUS_MISSING) {
            current = counter;
        } else if (current == STATUS_AMBIGUOUS) {
            status.is_duplicate = true;
        }
    }
 
    void recursive_add(size_t i, int position, const char* barcode_seq, AddStatus& status) {
        // Processing a stretch of non-ambiguous codes, where possible.
        // This reduces the recursion depth among the (hopefully fewer) ambiguous codes.
        while (1) {
            auto shift = base_shift<true>(barcode_seq[i]);
            if (shift == NON_STANDARD_BASE) {
               break;
            }
 
            if ((++i) == length) {
                end(shift, position, status);
                return;
            } else {
                next(shift, position);
            }
        } 
 
        // Processing the ambiguous codes.
        status.has_ambiguous = true;
 
        auto process = [&](char base) -> void {
            auto shift = base_shift(base);
            if (i + 1 == length) {
                end(shift, position, status);
            } else {
                auto curpos = position;
                next(shift, curpos);
                recursive_add(i + 1, curpos, barcode_seq, status);
            }
        };
 
        switch(barcode_seq[i]) {
            case 'R': case 'r':
                process('A'); process('G'); break;
            case 'Y': case 'y':
                process('C'); process('T'); break;
            case 'S': case 's':
                process('C'); process('G'); break;
            case 'W': case 'w':
                process('A'); process('T'); break;
            case 'K': case 'k':
                process('G'); process('T'); break;
            case 'M': case 'm':
                process('A'); process('C'); break;
            case 'B': case 'b':
                process('C'); process('G'); process('T'); break;
            case 'D': case 'd':
                process('A'); process('G'); process('T'); break;
            case 'H': case 'h':
                process('A'); process('C'); process('T'); break;
            case 'V': case 'v':
                process('A'); process('C'); process('G'); break;
            case 'N': case 'n':
                process('A'); process('C'); process('G'); process('T'); break;
            default:
                throw std::runtime_error("unknown base '" + std::string(1, barcode_seq[i]) + "' detected when constructing the trie");
        }
    }
 
public:
    AddStatus add(const char* barcode_seq) {
        AddStatus status;
        recursive_add(0, 0, barcode_seq, status);
        ++counter;
        return status;
    }
 
public:
    size_t get_length() const {
        return length;
    }
 
    int size() const {
        return counter;
    }
 
protected:
    size_t length;
    std::vector<int> pointers;
 
    static constexpr int NON_STANDARD_BASE = -1;
 
    template<bool allow_unknown = false>
    static int base_shift(char base) {
        int shift = 0;
        switch (base) {
            case 'A': case 'a':
                break;
            case 'C': case 'c':
                shift = 1;
                break;
            case 'G': case 'g':
                shift = 2;
                break;
            case 'T': case 't':
                shift = 3;
                break;
            default:
                if constexpr(allow_unknown) {
                    shift = NON_STANDARD_BASE; 
                } else {
                    throw std::runtime_error("unknown base '" + std::string(1, base) + "' detected when constructing the trie");
                }
        }
        return shift;
    }
protected:
    DuplicateAction duplicates;
 
    // To be called in the middle steps of the recursive search (i.e., for all but the last position).
    template<class SearchResult_>
    void replace_best_with_chosen(SearchResult_& best, int& best_index, int best_score, const SearchResult_& chosen, int chosen_index, int chosen_score) const {
        if (chosen_index >= 0) {
            if (chosen_score < best_score) {
                best = chosen;
            } else if (chosen_score == best_score) { 
                if (chosen_index != best_index) { // protect against multiple occurrences of IUPAC code-containing barcodes.
                    if (duplicates == DuplicateAction::FIRST) {
                        if (chosen_index < best_index) {
                            best_index = chosen_index;
                        }
                    } else if (duplicates == DuplicateAction::LAST) {
                        if (chosen_index > best_index) {
                            best_index = chosen_index;
                        }
                    } else {
                        best_index = STATUS_AMBIGUOUS; 
                    }
                }
            }
 
        } else if (chosen_index == STATUS_AMBIGUOUS) {
            if (chosen_score < best_score) {
                best = chosen;
            } else if (chosen_score == best_score) {
                // Ambiguity is infectious. Each ambiguous status indicates that there
                // are already 2+ barcodes on this score, so it doesn't matter how
                // many other unambiguous barcodes are here; we're already ambiguous.
                best_index = STATUS_AMBIGUOUS;
            }
        }
    }
 
    // To be called in the last step of the recursive search.
    void scan_final_position_with_mismatch(int node, int refshift, int& current_index, int current_mismatches, int& mismatch_cap) const {
        bool found = false;
        for (int s = 0; s < NUM_BASES; ++s) {
            if (s == refshift) { 
                continue;
            }
 
            int candidate = pointers[node + s];
            if (candidate >= 0) {
                if (found) { 
                    if (candidate != current_index) { // protect against multiple occurrences of IUPAC-containg barcodes.
                        if (duplicates == DuplicateAction::FIRST) {
                            if (current_index > candidate) {
                                current_index = candidate;
                            }
                        } else if (duplicates == DuplicateAction::LAST) {
                            if (current_index < candidate) {
                                current_index = candidate;
                            }
                        } else {
                            current_index = STATUS_AMBIGUOUS; // ambiguous, so we quit early.
                            break;
                        }
                    }
                } else {
                    current_index = candidate;
                    mismatch_cap = current_mismatches;
                    found = true;
                }
 
            } else if (candidate == STATUS_AMBIGUOUS) {
                // If an ambiguity is present on a base at the last position, 
                // and we're accepting a mismatch on the last position, then 
                // we already have at least two known barcodes that match the 
                // input sequence. The behavior of the other bases is irrelevant;
                // even if they refer to non-ambiguous barcodes, that just adds to the
                // set of 2+ barcodes that the input sequence already matches.
                // So, we have no choice but to fail the match due to ambiguity.
                current_index = STATUS_AMBIGUOUS;
                mismatch_cap = current_mismatches;
                break;
            }
        }
    }
private:
    int counter;
 
public:
    void optimize() {
        int maxed = 0;
        if (!is_optimal(0, 0, maxed)) {
            std::vector<int> replacement;
            replacement.reserve(pointers.size());
            optimize(0, 0, replacement);
            pointers.swap(replacement);
        }
    }
 
private:
    // Optimization involves reorganizing the nodes so that the pointers are
    // always increasing. This promotes memory locality of similar sequences
    // in a depth-first search (which is what search() does anyway).
    bool is_optimal(int node, size_t pos, int& maxed) const {
        ++pos;
        if (pos < length) {
            for (int s = 0; s < NUM_BASES; ++s) {
                auto v = pointers[node + s];
                if (v < 0) {
                    continue;
                }
 
                if (v < maxed) {
                    return false;
                }
 
                maxed = v;
                if (!is_optimal(v, pos, maxed)) {
                    return false;
                }
            }
        }
        return true;
    }
 
    void optimize(int node, size_t pos, std::vector<int>& trie) const {
        auto it = pointers.begin() + node;
        size_t new_node = trie.size();
        trie.insert(trie.end(), it, it + NUM_BASES);
 
        ++pos;
        if (pos < length) {
            for (int s = 0; s < NUM_BASES; ++s) {
                auto& v = trie[new_node + s];
                if (v < 0) {
                    continue;
                }
 
                auto original = v;
                v = trie.size();
                optimize(original, pos, trie);
            }
        }
    }
};
 
class AnyMismatches : public MismatchTrie {
public:
    AnyMismatches() {}
 
    AnyMismatches(size_t barcode_length, DuplicateAction duplicates) : MismatchTrie(barcode_length, duplicates) {}
 
public:
    std::pair<int, int> search(const char* search_seq, int max_mismatches) const {
        return search(search_seq, 0, 0, 0, max_mismatches);
    }
 
private:
    std::pair<int, int> search(const char* seq, size_t pos, int node, int mismatches, int& max_mismatches) const {
        int shift = base_shift<true>(seq[pos]);
        int current = (shift >= 0 ? pointers[node + shift] : STATUS_MISSING);
 
        // At the end: we prepare to return the actual values. We also refine
        // the max number of mismatches so that we don't search for things with
        // more mismatches than the best hit that was already encountered.
        if (pos + 1 == length) {
            if (current >= 0 || current == STATUS_AMBIGUOUS) {
                max_mismatches = mismatches; // this assignment should always decrease max_mismatches, otherwise the search would have terminated earlier.
                return std::make_pair(current, mismatches);
            }
 
            int alt = STATUS_MISSING;
            ++mismatches;
            if (mismatches <= max_mismatches) {
                scan_final_position_with_mismatch(node, shift, alt, mismatches, max_mismatches);
            }
 
            return std::make_pair(alt, mismatches);
 
        } else {
            ++pos;
 
            std::pair<int, int> best(STATUS_MISSING, max_mismatches + 1);
            if (current >= 0) {
                best = search(seq, pos, current, mismatches, max_mismatches);
            }
 
            ++mismatches;
            if (mismatches <= max_mismatches) {
                for (int s = 0; s < NUM_BASES; ++s) {
                    if (shift == s) { 
                        continue;
                    } 
 
                    int alt = pointers[node + s];
                    if (alt < 0) {
                        continue;
                    }
 
                    if (mismatches <= max_mismatches) { // check again, just in case max_mismatches changed.
                        auto chosen = search(seq, pos, alt, mismatches, max_mismatches);
                        replace_best_with_chosen(best, best.first, best.second, chosen, chosen.first, chosen.second);
                    }
                }
            }
 
            return best;
        }
    }
};
 
template<size_t num_segments>
class SegmentedMismatches : public MismatchTrie {
public:
    SegmentedMismatches() {}
 
    SegmentedMismatches(std::array<int, num_segments> segments, DuplicateAction duplicates) : 
        MismatchTrie(std::accumulate(segments.begin(), segments.end(), 0), duplicates), 
        boundaries(segments)
    {
        for (size_t i = 1; i < num_segments; ++i) {
            boundaries[i] += boundaries[i-1];
        }
    }
 
public:
    struct Result {
        Result() : per_segment() {}
        int index = 0;
 
        int total = 0;
 
        std::array<int, num_segments> per_segment;
    };
 
    Result search(const char* search_seq, const std::array<int, num_segments>& max_mismatches) const {
        int total_mismatches = std::accumulate(max_mismatches.begin(), max_mismatches.end(), 0);
        return search(search_seq, 0, 0, Result(), max_mismatches, total_mismatches);
    }
 
private:
    Result search(
        const char* seq, 
        size_t pos, 
        size_t segment_id,
        Result state,
        const std::array<int, num_segments>& segment_mismatches, 
        int& total_mismatches
    ) const {
        // Note that, during recursion, state.index does double duty 
        // as the index of the node on the trie.
        int node = state.index;
 
        int shift = base_shift<true>(seq[pos]);
        int current = (shift >= 0 ? pointers[node + shift] : STATUS_MISSING);
 
        // At the end: we prepare to return the actual values. We also refine
        // the max number of mismatches so that we don't search for things with
        // more mismatches than the best hit that was already encountered.
        if (pos + 1 == length) {
            if (current >= 0 || current == STATUS_AMBIGUOUS) {
                total_mismatches = state.total; // this assignment should always decrease total_mismatches, otherwise the search would have terminated earlier.
                state.index = current;
                return state;
            }
 
            state.index = STATUS_MISSING;
            ++state.total;
            auto& current_segment_mm = state.per_segment[segment_id];
            ++current_segment_mm;
 
            if (state.total <= total_mismatches && current_segment_mm <= segment_mismatches[segment_id]) {
                scan_final_position_with_mismatch(node, shift, state.index, state.total, total_mismatches);
            }
 
            return state;
 
        } else {
            auto next_pos = pos + 1;
            auto next_segment_id = segment_id;
            if (static_cast<int>(next_pos) == boundaries[segment_id]) { // TODO: boundaries should probably be size_t's, thus avoiding the need for this cast.
                ++next_segment_id;
            }
 
            Result best;
            best.index = STATUS_MISSING;
            best.total = total_mismatches + 1;
 
            if (current >= 0) {
                state.index = current;
                best = search(seq, next_pos, next_segment_id, state, segment_mismatches, total_mismatches);
            }
 
            ++state.total;
            auto& current_segment_mm = state.per_segment[segment_id];
            ++current_segment_mm;
 
            if (state.total <= total_mismatches && current_segment_mm <= segment_mismatches[segment_id]) {
                for (int s = 0; s < NUM_BASES; ++s) {
                    if (shift == s) { 
                        continue;
                    } 
                    
                    int alt = pointers[node + s];
                    if (alt < 0) {
                        continue;
                    }
 
                    if (state.total <= total_mismatches) { // check again, just in case total_mismatches changed.
                        state.index = alt;
                        auto chosen = search(seq, next_pos, next_segment_id, state, segment_mismatches, total_mismatches);
                        replace_best_with_chosen(best, best.index, best.total, chosen, chosen.index, chosen.total);
                    }
                }
            }
 
            return best;
        }
    }
private:
    std::array<int, num_segments> boundaries;
};
 
}
 
#endif