From 7c7e1390258fc313d5596a57da32401b1d57cf09 Mon Sep 17 00:00:00 2001
From: Harun Mustafa <bilgisayaradami@gmail.com>
Date: Tue, 8 Oct 2024 11:41:18 +0200
Subject: [PATCH] test

---
 metagraph/src/cli/query.cpp                   |   5 +-
 .../src/graph/alignment/annotation_buffer.cpp |   1 +
 metagraph/src/graph/annotated_dbg.cpp         | 101 +++++++++++----
 metagraph/src/graph/annotated_dbg.hpp         |   7 +-
 .../graph/representation/canonical_dbg.cpp    |  16 ---
 .../graph/representation/hash/dbg_sshash.cpp  | 118 +++++++++++++++---
 .../graph/representation/hash/dbg_sshash.hpp  |   5 +-
 .../tests/annotation/test_annotated_dbg.cpp   |  33 ++---
 .../annotation/test_annotated_dbg_helpers.cpp |   2 +
 .../tests/graph/all/test_dbg_helpers.cpp      |  79 +++++++++++-
 .../tests/graph/all/test_dbg_helpers.hpp      |   8 ++
 metagraph/tests/graph/test_canonical_dbg.cpp  |   3 +-
 metagraph/tests/graph/test_dbg_canonical.cpp  |   3 +-
 13 files changed, 303 insertions(+), 78 deletions(-)

diff --git a/metagraph/src/cli/query.cpp b/metagraph/src/cli/query.cpp
index cb7d2e35c9..fcb41c61b5 100644
--- a/metagraph/src/cli/query.cpp
+++ b/metagraph/src/cli/query.cpp
@@ -955,8 +955,9 @@ construct_query_graph(const AnnotatedDBG &anno_graph,
     #pragma omp parallel for num_threads(num_threads)
     for (size_t i = 0; i < contigs.size(); ++i) {
         contigs[i].second.reserve(contigs[i].first.length() - graph_init->get_k() + 1);
-        full_dbg.map_to_nodes(contigs[i].first,
-                              [&](node_index node) { contigs[i].second.push_back(node); });
+        call_annotated_nodes_offsets(full_dbg, contigs[i].first, [&](node_index node, int64_t) {
+            contigs[i].second.push_back(node);
+        });
     }
     logger->trace("[Query graph construction] Contigs mapped to the full graph in {} sec",
                   timer.elapsed());
diff --git a/metagraph/src/graph/alignment/annotation_buffer.cpp b/metagraph/src/graph/alignment/annotation_buffer.cpp
index 4020f312a7..16d501385b 100644
--- a/metagraph/src/graph/alignment/annotation_buffer.cpp
+++ b/metagraph/src/graph/alignment/annotation_buffer.cpp
@@ -48,6 +48,7 @@ void AnnotationBuffer::fetch_queued_annotations() {
 
     for (const auto &path : queued_paths_) {
         std::vector<node_index> base_path;
+        std::vector<int64_t> base_path_offsets;
         if (base_graph->get_mode() == DeBruijnGraph::CANONICAL) {
             // TODO: avoid this call of spell_path
             std::string query = spell_path(graph_, path);
diff --git a/metagraph/src/graph/annotated_dbg.cpp b/metagraph/src/graph/annotated_dbg.cpp
index 430365365a..6fc6ea257c 100644
--- a/metagraph/src/graph/annotated_dbg.cpp
+++ b/metagraph/src/graph/annotated_dbg.cpp
@@ -6,6 +6,7 @@
 #include "annotation/representation/row_compressed/annotate_row_compressed.hpp"
 #include "annotation/int_matrix/base/int_matrix.hpp"
 #include "graph/representation/canonical_dbg.hpp"
+#include "graph/representation/hash/dbg_sshash.hpp"
 #include "common/aligned_vector.hpp"
 #include "common/vectors/vector_algorithm.hpp"
 #include "common/vector_map.hpp"
@@ -23,6 +24,21 @@ using Column = mtg::annot::matrix::BinaryMatrix::Column;
 typedef AnnotatedDBG::Label Label;
 typedef std::pair<Label, size_t> StringCountPair;
 
+void call_annotated_nodes_offsets(const SequenceGraph &graph,
+                                  std::string_view sequence,
+                                  const std::function<void(SequenceGraph::node_index, int64_t)> &callback) {
+    if (const auto *sshash = dynamic_cast<const DBGSSHash*>(&graph)) {
+        if (sshash->is_monochromatic()) {
+            sshash->map_to_contigs_with_rc(sequence, [&](SequenceGraph::node_index i, int64_t offset, bool) {
+                callback(i, offset);
+            });
+            return;
+        }
+    }
+
+    graph.map_to_nodes(sequence, [&](SequenceGraph::node_index i) { callback(i, 0); });
+}
+
 
 AnnotatedSequenceGraph
 ::AnnotatedSequenceGraph(std::shared_ptr<SequenceGraph> graph,
@@ -46,10 +62,9 @@ ::annotate_sequence(std::string_view sequence,
 
     std::vector<row_index> indices;
     indices.reserve(sequence.size());
-
-    graph_->map_to_nodes(sequence, [&](node_index i) {
+    call_annotated_nodes_offsets(*graph_, sequence, [&](node_index i, int64_t) {
         if (i > 0)
-            indices.push_back(graph_to_anno_index(i));
+            indices.emplace_back(graph_to_anno_index(i));
     });
 
     if (!indices.size())
@@ -80,9 +95,9 @@ ::annotate_sequences(const std::vector<std::pair<std::string, std::vector<Label>
         if (!indices.capacity())
             indices.reserve(data[t].first.size());
 
-        graph_->map_to_nodes(data[t].first, [&](node_index i) {
+        call_annotated_nodes_offsets(*graph_, data[t].first, [&](node_index i, int64_t) {
             if (i > 0)
-                indices.push_back(graph_to_anno_index(i));
+                indices.emplace_back(graph_to_anno_index(i));
         });
     }
 
@@ -117,10 +132,10 @@ void AnnotatedDBG::add_kmer_counts(std::string_view sequence,
     indices.reserve(sequence.size() - dbg_.get_k() + 1);
     size_t end = 0;
 
-    graph_->map_to_nodes(sequence, [&](node_index i) {
+    call_annotated_nodes_offsets(dbg_, sequence, [&](node_index i, int64_t) {
         // only insert indexes for matched k-mers and shift counts accordingly
         if (i > 0) {
-            indices.push_back(graph_to_anno_index(i));
+            indices.emplace_back(graph_to_anno_index(i));
             kmer_counts[indices.size() - 1] = kmer_counts[end++];
         }
     });
@@ -139,15 +154,22 @@ void AnnotatedDBG::add_kmer_coord(std::string_view sequence,
     if (sequence.size() < dbg_.get_k())
         return;
 
-    std::vector<row_index> indices = map_to_nodes(dbg_, sequence);
+    std::vector<row_index> indices;
+    std::vector<int64_t> offsets;
+    call_annotated_nodes_offsets(dbg_, sequence, [&](node_index i, int64_t offset) {
+        indices.emplace_back(i);
+        offsets.emplace_back(offset);
+    });
 
     std::lock_guard<std::mutex> lock(mutex_);
 
+    auto it = offsets.begin();
     for (node_index i : indices) {
         // only insert coordinates for matched k-mers and increment the coordinates
         if (i > 0)
-            annotator_->add_label_coord(graph_to_anno_index(i), labels, coord);
+            annotator_->add_label_coord(graph_to_anno_index(i), labels, coord - *it);
         coord++;
+        ++it;
     }
 }
 
@@ -156,10 +178,17 @@ void AnnotatedDBG::add_kmer_coords(
     assert(check_compatibility());
 
     std::vector<std::vector<row_index>> ids;
+    std::vector<std::vector<int64_t>> offsets;
     ids.reserve(data.size());
     for (const auto &[sequence, labels, _] : data) {
-        if (sequence.size() >= dbg_.get_k())
-            ids.push_back(map_to_nodes(dbg_, sequence));
+        if (sequence.size() >= dbg_.get_k()) {
+            auto &id = ids.emplace_back();
+            auto &offset = offsets.emplace_back();
+            call_annotated_nodes_offsets(dbg_, sequence, [&](node_index i, int64_t o) {
+                id.emplace_back(i);
+                offset.emplace_back(o);
+            });
+        }
     }
 
     std::lock_guard<std::mutex> lock(mutex_);
@@ -168,11 +197,13 @@ void AnnotatedDBG::add_kmer_coords(
         const auto &labels = std::get<1>(data[t]);
         uint64_t coord = std::get<2>(data[t]);
 
+        auto it = offsets[t].begin();
         for (node_index i : ids[t]) {
             // only insert coordinates for matched k-mers and increment the coordinates
             if (i > 0)
-                annotator_->add_label_coord(graph_to_anno_index(i), labels, coord);
+                annotator_->add_label_coord(graph_to_anno_index(i), labels, coord - *it);
             coord++;
+            ++it;
         }
     }
 }
@@ -200,10 +231,10 @@ void AnnotatedDBG::annotate_kmer_coords(
             coords[last].reserve(sequence.size() - dbg_.get_k() + 1);
         }
 
-        graph_->map_to_nodes(sequence, [&](node_index i) {
+        call_annotated_nodes_offsets(*graph_, sequence, [&](node_index i, int64_t o) {
             if (i > 0) {
                 ids[last].push_back(graph_to_anno_index(i));
-                coords[last].emplace_back(graph_to_anno_index(i), coord);
+                coords[last].emplace_back(graph_to_anno_index(i), coord - o);
             }
             coord++;
         });
@@ -238,7 +269,7 @@ std::vector<Label> AnnotatedDBG::get_labels(std::string_view sequence,
     size_t num_present_kmers = 0;
     size_t num_missing_kmers = 0;
 
-    graph_->map_to_nodes(sequence, [&](node_index i) {
+    call_annotated_nodes_offsets(*graph_, sequence, [&](node_index i, int64_t) {
         if (i > 0) {
             index_counts[graph_to_anno_index(i)]++;
             num_present_kmers++;
@@ -312,7 +343,7 @@ AnnotatedDBG::get_top_labels(std::string_view sequence,
 
     size_t num_present_kmers = 0;
 
-    graph_->map_to_nodes(sequence, [&](node_index i) {
+    call_annotated_nodes_offsets(*graph_, sequence, [&](node_index i, int64_t) {
         if (i > 0) {
             index_counts[graph_to_anno_index(i)]++;
             num_present_kmers++;
@@ -342,7 +373,13 @@ AnnotatedDBG::get_kmer_counts(std::string_view sequence,
                               size_t num_top_labels,
                               double discovery_fraction,
                               double presence_fraction) const {
-    std::vector<node_index> nodes = map_to_nodes(dbg_, sequence);
+    std::vector<node_index> nodes;
+    if (sequence.size() >= dbg_.get_k()) {
+        nodes.reserve(sequence.size() - dbg_.get_k() + 1);
+        call_annotated_nodes_offsets(*graph_, sequence, [&](node_index i, int64_t) {
+            nodes.emplace_back(i);
+        });
+    }
     return get_kmer_counts(nodes, num_top_labels, discovery_fraction, presence_fraction);
 }
 
@@ -454,19 +491,30 @@ AnnotatedDBG::get_kmer_coordinates(std::string_view sequence,
                                    size_t num_top_labels,
                                    double discovery_fraction,
                                    double presence_fraction) const {
-    std::vector<node_index> nodes = map_to_nodes(dbg_, sequence);
-    return get_kmer_coordinates(nodes, num_top_labels, discovery_fraction, presence_fraction);
+    std::vector<node_index> nodes;
+    std::vector<int64_t> offsets;
+    if (sequence.size() >= dbg_.get_k()) {
+        nodes.reserve(sequence.size() - dbg_.get_k() + 1);
+        offsets.reserve(sequence.size() - dbg_.get_k() + 1);
+        call_annotated_nodes_offsets(*graph_, sequence, [&](node_index i, int64_t o) {
+            nodes.emplace_back(i);
+            offsets.emplace_back(o);
+        });
+    }
+    return get_kmer_coordinates(nodes, num_top_labels, discovery_fraction, presence_fraction, offsets);
 }
 
 std::vector<std::tuple<Label, size_t, std::vector<SmallVector<uint64_t>>>>
 AnnotatedDBG::get_kmer_coordinates(const std::vector<node_index> &nodes,
                                    size_t num_top_labels,
                                    double discovery_fraction,
-                                   double presence_fraction) const {
+                                   double presence_fraction,
+                                   const std::vector<int64_t> &offsets) const {
     assert(discovery_fraction >= 0.);
     assert(discovery_fraction <= 1.);
     assert(presence_fraction >= 0.);
     assert(presence_fraction <= 1.);
+    assert(offsets.empty() || offsets.size() == nodes.size());
     assert(check_compatibility());
 
     if (!nodes.size())
@@ -528,8 +576,15 @@ AnnotatedDBG::get_kmer_coordinates(const std::vector<node_index> &nodes,
     for (size_t i = 0; i < rows_tuples.size(); ++i) {
         // set the non-empty tuples
         for (auto &[j, tuple] : rows_tuples[i]) {
-            if (col_counts[j])
-                std::get<2>(result[col_counts[j] - 1])[kmer_positions[i]] = std::move(tuple);
+            if (col_counts[j]) {
+                auto &coords = std::get<2>(result[col_counts[j] - 1])[kmer_positions[i]];
+                coords = std::move(tuple);
+                if (offsets.size()) {
+                    for (auto &c : coords) {
+                        c += offsets[kmer_positions[i]];
+                    }
+                }
+            }
         }
     }
 
@@ -574,7 +629,7 @@ AnnotatedDBG::get_top_label_signatures(std::string_view sequence,
     kmer_positions.reserve(num_kmers);
 
     size_t j = 0;
-    graph_->map_to_nodes(sequence, [&](node_index i) {
+    call_annotated_nodes_offsets(*graph_, sequence, [&](node_index i, int64_t o) {
         if (i > 0) {
             kmer_positions.push_back(j);
             row_indices.push_back(graph_to_anno_index(i));
diff --git a/metagraph/src/graph/annotated_dbg.hpp b/metagraph/src/graph/annotated_dbg.hpp
index 4660a907d1..927a7b3cc2 100644
--- a/metagraph/src/graph/annotated_dbg.hpp
+++ b/metagraph/src/graph/annotated_dbg.hpp
@@ -149,7 +149,8 @@ class AnnotatedDBG : public AnnotatedSequenceGraph {
     get_kmer_coordinates(const std::vector<node_index> &nodes,
                          size_t num_top_labels,
                          double discovery_fraction,
-                         double presence_fraction) const;
+                         double presence_fraction,
+                         const std::vector<int64_t> &offsets = {}) const;
 
     std::vector<std::pair<Label, sdsl::bit_vector>>
     get_top_label_signatures(std::string_view sequence,
@@ -165,6 +166,10 @@ class AnnotatedDBG : public AnnotatedSequenceGraph {
     DeBruijnGraph &dbg_;
 };
 
+void call_annotated_nodes_offsets(const SequenceGraph &graph,
+                                  std::string_view sequence,
+                                  const std::function<void(SequenceGraph::node_index, int64_t)> &callback);
+
 } // namespace graph
 } // namespace mtg
 
diff --git a/metagraph/src/graph/representation/canonical_dbg.cpp b/metagraph/src/graph/representation/canonical_dbg.cpp
index 1da6e4ae53..a944cd958c 100644
--- a/metagraph/src/graph/representation/canonical_dbg.cpp
+++ b/metagraph/src/graph/representation/canonical_dbg.cpp
@@ -156,22 +156,6 @@ ::map_to_nodes_sequentially(std::string_view sequence,
 void CanonicalDBG::map_to_nodes(std::string_view sequence,
                                 const std::function<void(node_index)> &callback,
                                 const std::function<bool()> &terminate) const {
-    if (const auto sshash = std::dynamic_pointer_cast<const DBGSSHash>(graph_)) {
-        if (sshash->is_monochromatic()) {
-            sshash->map_to_contigs_with_rc<true>(
-                sequence,
-                [&](node_index node, int64_t, bool orientation) {
-                    if (orientation)
-                        node = reverse_complement(node);
-
-                    callback(get_base_node(node));
-                },
-                terminate
-            );
-            return;
-        }
-    }
-
     map_to_nodes_sequentially(sequence, [&](node_index i) {
         callback(get_base_node(i));
     }, terminate);
diff --git a/metagraph/src/graph/representation/hash/dbg_sshash.cpp b/metagraph/src/graph/representation/hash/dbg_sshash.cpp
index 591587b879..8c1a20869a 100644
--- a/metagraph/src/graph/representation/hash/dbg_sshash.cpp
+++ b/metagraph/src/graph/representation/hash/dbg_sshash.cpp
@@ -2,10 +2,15 @@
 
 #include <type_traits>
 
+#include <progress_bar.hpp>
+
 #include "common/seq_tools/reverse_complement.hpp"
 #include "common/threads/threading.hpp"
 #include "common/logger.hpp"
+#include "common/utils/string_utils.hpp"
+#include "common/algorithms.hpp"
 #include "kmer/kmer_extractor.hpp"
+#include "seq_io/sequence_io.hpp"
 
 
 namespace mtg {
@@ -39,8 +44,7 @@ size_t DBGSSHash::dict_size() const {
     return std::visit([](const auto &d) { return d.size(); }, dict_);
 }
 
-DBGSSHash::DBGSSHash(size_t k, Mode mode)
-      : k_(k), num_nodes_(0), mode_(mode), is_monochromatic_(false) {
+DBGSSHash::DBGSSHash(size_t k, Mode mode) : k_(k), num_nodes_(0), mode_(mode) {
     size_t odd_k = (k_ | 1);
 
     if (odd_k * bits_per_char <= 64) {
@@ -64,7 +68,6 @@ DBGSSHash::DBGSSHash(const std::string &input_filename,
     if (k <= 1)
         throw std::domain_error("k must be at least 2");
 
-    is_monochromatic_ = is_monochromatic;
     sshash::build_configuration build_config;
     build_config.k = k;
 
@@ -94,6 +97,60 @@ DBGSSHash::DBGSSHash(const std::string &input_filename,
         std::cout.copyfmt(orig_state);
 
     num_nodes_ = dict_size();
+
+    if (is_monochromatic) {
+        common::logger->trace("Marking colour boundaries");
+        sdsl::bit_vector id(num_nodes_);
+        {
+            size_t pos = 0;
+            ProgressBar progress_bar(num_nodes(),
+                                     "Parsing sequences",
+                                     std::cerr, !common::get_verbose());
+            seq_io::read_fasta_file_critical(input_filename, [&](auto *read_stream) {
+                id[pos] = true;
+                auto split_string = utils::split_string(std::string(read_stream->comment.s, read_stream->comment.l), " ");
+                size_t old_pos = pos;
+                size_t nkmers = read_stream->seq.l - k_ + 1;
+
+                for (const auto &part : split_string) {
+                    if (part[0] == 'C') {
+                        auto c_parts = utils::split_string(part, ":");
+                        assert(c_parts.size());
+                        id[pos] = true;
+                        pos += atoll(c_parts.back().c_str());
+                    }
+                }
+
+                if (pos == old_pos)
+                    pos += nkmers;
+
+                assert(pos == old_pos + nkmers);
+
+                progress_bar += nkmers;
+            });
+        }
+
+        size_t update_batch = 10000;
+        ProgressBar progress_bar(num_nodes_,
+                                 "Marking unitig breakpoints",
+                                 std::cerr, !common::get_verbose());
+        std::atomic_thread_fence(std::memory_order_release);
+        #pragma omp parallel for num_threads(get_num_threads())
+        for (size_t p = 0; p < num_nodes_; ++p) {
+            node_index node = sshash_to_graph_index(p);
+            if (!has_single_incoming(node) || !has_single_outgoing(node))
+                set_bit(id.data(), p, std::memory_order_relaxed);
+
+            if (p && (p % update_batch) == 0)
+                progress_bar += update_batch;
+        }
+        std::atomic_thread_fence(std::memory_order_acquire);
+
+        progress_bar += num_nodes_ % update_batch;
+
+        monotig_id_ = bit_vector_rrr<>(std::move(id));
+        common::logger->trace("Marked {} boundaries", monotig_id_.num_set_bits());
+    }
 }
 
 std::string DBGSSHash::file_extension() const {
@@ -115,12 +172,23 @@ ::map_to_nodes_with_rc_advanced(std::string_view sequence,
 
     std::visit([&](const auto &dict) {
         using kmer_t = get_kmer_t<decltype(dict)>;
+
+        std::vector<char> seq_encoded;
+        seq_encoded.reserve(sequence.size());
+        for (size_t i = 0; i < sequence.size(); ++i) {
+            char enc = kmer_t::canonicalize_basepair_forward_map[static_cast<uint8_t>(sequence[i])];
+            seq_encoded.emplace_back(enc == '\0');
+        }
+
+        auto invalid = utils::drag_and_mark_segments(seq_encoded, 1, k_);
+
         kmer_t uint_kmer = sshash::util::string_to_uint_kmer<kmer_t>(sequence.data(), k_ - 1);
         uint_kmer.pad_char();
         for (size_t i = k_ - 1; i < sequence.size() && !terminate(); ++i) {
             uint_kmer.drop_char();
             uint_kmer.kth_char_or(k_ - 1, kmer_t::char_to_uint(sequence[i]));
-            callback(dict.lookup_advanced_uint(uint_kmer, with_rc));
+            callback(!invalid[i] ? dict.lookup_advanced_uint(uint_kmer, with_rc)
+                                 : sshash::lookup_result());
         }
     }, dict_);
 }
@@ -167,9 +235,33 @@ void DBGSSHash::map_to_contigs_with_rc(std::string_view sequence,
         return;
     }
 
+    node_index last_node = npos;
+    uint64_t coloured_kmer_id = std::numeric_limits<uint64_t>::max();
+    uint64_t next_breakpoint = std::numeric_limits<uint64_t>::max();
     map_to_nodes_with_rc_advanced(sequence, [&](sshash::lookup_result res) {
         node_index node = sshash_to_graph_index(res.kmer_id);
         if (node != npos) {
+            if (is_monochromatic()) {
+                assert(monotig_id_.size() == num_nodes());
+                assert(monotig_id_[res.kmer_id] || res.kmer_id_in_contig);
+                if (last_node != npos && node == last_node + 1) {
+                    if (res.kmer_id == next_breakpoint) {
+                        coloured_kmer_id = res.kmer_id;
+                        next_breakpoint = coloured_kmer_id + 1 < monotig_id_.size()
+                            ? monotig_id_.next1(coloured_kmer_id + 1)
+                            : monotig_id_.size();
+                    }
+                } else {
+                    coloured_kmer_id = monotig_id_.prev1(res.kmer_id);
+                    next_breakpoint = coloured_kmer_id + 1 < monotig_id_.size()
+                        ? monotig_id_.next1(coloured_kmer_id + 1)
+                        : monotig_id_.size();
+                    assert(coloured_kmer_id >= res.kmer_id - res.kmer_id_in_contig);
+                }
+                res.kmer_id_in_contig = res.kmer_id - coloured_kmer_id;
+                last_node = node;
+            }
+
             node -= res.kmer_id_in_contig;
         } else {
             res.kmer_id_in_contig = 0;
@@ -191,15 +283,7 @@ void DBGSSHash::map_to_contigs_with_rc<false>(std::string_view,
 void DBGSSHash::map_to_nodes(std::string_view sequence,
                              const std::function<void(node_index)>& callback,
                              const std::function<bool()>& terminate) const {
-    if (is_monochromatic()) {
-        if (mode_ != CANONICAL) {
-            map_to_contigs_with_rc<false>(
-                sequence, [&](node_index node, int64_t, bool) { callback(node); }, terminate);
-        } else {
-            map_to_contigs_with_rc<true>(
-                sequence, [&](node_index node, int64_t, bool) { callback(node); }, terminate);
-        }
-    } else if (mode_ != CANONICAL) {
+    if (mode_ != CANONICAL) {
         map_to_nodes_sequentially(sequence, callback, terminate);
     } else {
         map_to_nodes_with_rc<true>(
@@ -414,10 +498,11 @@ void DBGSSHash::serialize(std::ostream& out) const {
     saver.visit(num_nodes_);
     saver.visit(k_);
     saver.visit(mode_);
-    saver.visit(is_monochromatic_);
 
     if (num_nodes())
         std::visit([&](const auto &d) { saver.visit(d); }, dict_);
+
+    monotig_id_.serialize(out);
 }
 
 void DBGSSHash::serialize(const std::string& filename) const {
@@ -431,19 +516,18 @@ bool DBGSSHash::load(std::istream &in) {
     size_t num_nodes;
     size_t k;
     Mode mode;
-    bool is_monochromatic;
     loader.visit(num_nodes);
     loader.visit(k);
     loader.visit(mode);
-    loader.visit(is_monochromatic);
 
     *this = DBGSSHash(k, mode);
     num_nodes_ = num_nodes;
-    is_monochromatic_ = is_monochromatic;
 
     if (num_nodes_)
         std::visit([&](auto &d) { d.visit(loader); }, dict_);
 
+    monotig_id_.load(in);
+
     return true;
 }
 
diff --git a/metagraph/src/graph/representation/hash/dbg_sshash.hpp b/metagraph/src/graph/representation/hash/dbg_sshash.hpp
index a44458bd8b..d57ec3d9ac 100644
--- a/metagraph/src/graph/representation/hash/dbg_sshash.hpp
+++ b/metagraph/src/graph/representation/hash/dbg_sshash.hpp
@@ -8,6 +8,7 @@
 #include <sdsl/uint256_t.hpp>
 
 #include "graph/representation/base/sequence_graph.hpp"
+#include "common/vectors/bit_vector_adaptive.hpp"
 
 namespace mtg::graph {
 
@@ -78,7 +79,7 @@ class DBGSSHash : public DeBruijnGraph {
 
     Mode get_mode() const override final { return mode_; }
 
-    bool is_monochromatic() const { return is_monochromatic_; }
+    bool is_monochromatic() const override { return monotig_id_.size(); }
 
     node_index traverse(node_index node, char next_char) const override;
     node_index traverse_back(node_index node, char prev_char) const override;
@@ -124,7 +125,7 @@ class DBGSSHash : public DeBruijnGraph {
     size_t k_;
     size_t num_nodes_;
     Mode mode_;
-    bool is_monochromatic_;
+    bit_vector_rrr<> monotig_id_;
 
     void map_to_nodes_with_rc_advanced(
             std::string_view sequence,
diff --git a/metagraph/tests/annotation/test_annotated_dbg.cpp b/metagraph/tests/annotation/test_annotated_dbg.cpp
index 278f01f1f0..d695445632 100644
--- a/metagraph/tests/annotation/test_annotated_dbg.cpp
+++ b/metagraph/tests/annotation/test_annotated_dbg.cpp
@@ -4,15 +4,12 @@
 #include "gtest/gtest.h"
 
 #include "../test_helpers.hpp"
+#include "../graph/all/test_dbg_helpers.hpp"
 
 #include "common/threads/threading.hpp"
 #include "common/vectors/bit_vector_dyn.hpp"
 #include "common/vectors/vector_algorithm.hpp"
 #include "annotation/representation/column_compressed/annotate_column_compressed.hpp"
-#include "graph/representation/bitmap/dbg_bitmap.hpp"
-#include "graph/representation/hash/dbg_hash_string.hpp"
-#include "graph/representation/hash/dbg_hash_ordered.hpp"
-#include "graph/representation/hash/dbg_hash_fast.hpp"
 
 #define protected public
 #define private public
@@ -31,16 +28,16 @@ void check_labels(const AnnotatedDBG &anno_graph,
                   const std::vector<std::string> labels_present,
                   const std::vector<std::string> labels_not_present) {
     std::set<SequenceGraph::node_index> indices;
-    anno_graph.get_graph().map_to_nodes(
-        sequence,
-        [&](const auto &index) {
-            ASSERT_NE(SequenceGraph::npos, index);
-            indices.insert(index);
-            EXPECT_EQ(labels_present.size(), anno_graph.get_labels(index).size());
-            EXPECT_EQ(convert_to_set(labels_present),
-                      convert_to_set(anno_graph.get_labels(index)));
-        }
-    );
+    call_annotated_nodes_offsets(anno_graph.get_graph(), sequence, [&](auto index, int64_t) {
+        ASSERT_NE(SequenceGraph::npos, index);
+        indices.insert(index);
+        EXPECT_EQ(labels_present.size(), anno_graph.get_labels(index).size());
+        EXPECT_EQ(convert_to_set(labels_present),
+                  convert_to_set(anno_graph.get_labels(index)));
+    });
+
+    EXPECT_EQ(convert_to_set(labels_present),
+              convert_to_set(anno_graph.get_labels(sequence)));
 
     for (const auto &label : labels_present) {
         std::set<SequenceGraph::node_index> cur_indices;
@@ -987,6 +984,8 @@ typedef ::testing::Types<std::pair<DBGBitmap, annot::ColumnCompressed<>>,
                          std::pair<DBGHashOrdered, annot::ColumnCompressed<>>,
                          std::pair<DBGHashFast, annot::ColumnCompressed<>>,
                          std::pair<DBGSuccinct, annot::ColumnCompressed<>>,
+                         std::pair<DBGSSHash, annot::ColumnCompressed<>>,
+                         std::pair<DBGSSHashMonochromatic, annot::ColumnCompressed<>>,
                          std::pair<DBGBitmap, annot::RowFlatAnnotator>,
                          std::pair<DBGHashString, annot::RowFlatAnnotator>,
                          std::pair<DBGHashOrdered, annot::RowFlatAnnotator>,
@@ -1016,6 +1015,8 @@ class AnnotatedDBGNoNTest : public ::testing::Test {};
 typedef ::testing::Types<std::pair<DBGBitmap, annot::ColumnCompressed<>>,
                          std::pair<DBGHashOrdered, annot::ColumnCompressed<>>,
                          std::pair<DBGHashFast, annot::ColumnCompressed<>>,
+                         std::pair<DBGSSHash, annot::ColumnCompressed<>>,
+                         std::pair<DBGSSHashMonochromatic, annot::ColumnCompressed<>>,
                          std::pair<DBGBitmap, annot::RowFlatAnnotator>,
                          std::pair<DBGHashOrdered, annot::RowFlatAnnotator>,
                          std::pair<DBGHashFast, annot::RowFlatAnnotator>,
@@ -1040,6 +1041,10 @@ TYPED_TEST(AnnotatedDBGWithNTest, check_labels) {
             k + 1, sequences, { "First", "Second" , "Third" }
         );
 
+        if (k > 3) {
+            ASSERT_FALSE(anno_graph->get_graph().find(std::string(k - 3, 'T') + std::string(3, 'G')));
+        }
+
         check_labels(*anno_graph, std::string(100, 'A'), { "First", "Third" }, { "Second" });
         check_labels(*anno_graph, std::string(100, 'T'), { "Second" }, { "First", "Third" });
         check_labels(*anno_graph, std::string(k, 'A') + std::string(k, 'C'), { "First" }, { "Second", "Third" });
diff --git a/metagraph/tests/annotation/test_annotated_dbg_helpers.cpp b/metagraph/tests/annotation/test_annotated_dbg_helpers.cpp
index 2c6b1f5735..8b7a8071a9 100644
--- a/metagraph/tests/annotation/test_annotated_dbg_helpers.cpp
+++ b/metagraph/tests/annotation/test_annotated_dbg_helpers.cpp
@@ -235,6 +235,8 @@ template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGBitmap, ColumnCompres
 template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGHashOrdered, ColumnCompressed<>>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&, DeBruijnGraph::Mode, bool);
 template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGHashFast, ColumnCompressed<>>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&, DeBruijnGraph::Mode, bool);
 template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGHashString, ColumnCompressed<>>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&, DeBruijnGraph::Mode, bool);
+template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGSSHash, ColumnCompressed<>>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&, DeBruijnGraph::Mode, bool);
+template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGSSHashMonochromatic, ColumnCompressed<>>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&, DeBruijnGraph::Mode, bool);
 
 template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGSuccinct, RowFlatAnnotator>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&, DeBruijnGraph::Mode, bool);
 template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGBitmap, RowFlatAnnotator>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&, DeBruijnGraph::Mode, bool);
diff --git a/metagraph/tests/graph/all/test_dbg_helpers.cpp b/metagraph/tests/graph/all/test_dbg_helpers.cpp
index 82c6878024..dae809d2d5 100644
--- a/metagraph/tests/graph/all/test_dbg_helpers.cpp
+++ b/metagraph/tests/graph/all/test_dbg_helpers.cpp
@@ -1,6 +1,10 @@
 #include "test_dbg_helpers.hpp"
 
+#include "../../annotation/test_annotated_dbg_helpers.hpp"
+#include "annotation/representation/column_compressed/annotate_column_compressed.hpp"
+
 #include "gtest/gtest.h"
+#include "graph/annotated_dbg.hpp"
 #include "graph/representation/canonical_dbg.hpp"
 #include "graph/representation/succinct/boss.hpp"
 #include "graph/representation/succinct/boss_construct.hpp"
@@ -43,6 +47,9 @@ template<> size_t max_test_k<DBGHashString>() {
 template<> size_t max_test_k<DBGSSHash>() {
     return 255 / DBGSSHash::kmer_t<uint64_t>::bits_per_char;
 }
+template<> size_t max_test_k<DBGSSHashMonochromatic>() {
+    return 255 / DBGSSHash::kmer_t<uint64_t>::bits_per_char;
+}
 
 template <class Graph>
 std::vector<std::string>
@@ -146,6 +153,7 @@ void writeFastaFile(const std::vector<std::string>& sequences, const std::string
 
     fastaFile.close();
 }
+
 template <>
 std::shared_ptr<DeBruijnGraph>
 build_graph<DBGSSHash>(uint64_t k,
@@ -155,7 +163,7 @@ build_graph<DBGSSHash>(uint64_t k,
         return std::make_shared<DBGSSHash>(k, mode);
 
     // use DBGHashString to get contigs for SSHash
-    auto string_graph = build_graph<DBGHashString>(k, sequences, mode);
+    auto string_graph = build_graph<DBGHashFast>(k, sequences, mode);
 
     std::vector<std::string> contigs;
     size_t num_kmers = 0;
@@ -182,6 +190,70 @@ build_graph<DBGSSHash>(uint64_t k,
     return graph;
 }
 
+template <>
+std::shared_ptr<DeBruijnGraph>
+build_graph<DBGSSHashMonochromatic>(uint64_t k,
+                                    std::vector<std::string> sequences,
+                                    DeBruijnGraph::Mode mode) {
+    if (sequences.empty())
+        return std::make_shared<DBGSSHash>(k, mode);
+
+    std::vector<std::string> labels;
+    labels.reserve(sequences.size());
+    char end_char = 'A' + sequences.size();
+    for (char c = 'A'; c < end_char; ++c) {
+        labels.emplace_back(1, c);
+    }
+
+    // use DBGHashString to get contigs for SSHash
+    auto string_anno_graph = build_anno_graph<DBGHashFast, annot::ColumnCompressed<>>(k, sequences, labels, mode);
+    const auto &string_graph = string_anno_graph->get_graph();
+
+    std::vector<std::string> contigs;
+    size_t num_kmers = 0;
+    size_t num_chars = 0;
+    string_graph.call_sequences([&](const std::string &contig, const auto &path) {
+        num_kmers += path.size();
+        auto begin = contig.begin();
+        auto end = contig.begin() + k;
+
+        auto sigs = string_anno_graph->get_top_label_signatures(contig, labels.size());
+        for (size_t i = 1; i + k <= contig.size(); ++i) {
+            for (size_t j = 0; j < sigs.size(); ++j) {
+                if (sigs[j].second[i] != sigs[j].second[i - 1]) {
+                    // we're at a breakpoint
+                    contigs.emplace_back(begin, end);
+                    num_chars += end - begin;
+                    begin = end - k + 1;
+                    assert(begin == contig.begin() + i);
+                    break;
+                }
+            }
+            ++end;
+        }
+
+        assert(end == contig.end());
+
+        contigs.emplace_back(begin, end);
+        num_chars += end - begin;
+    }, 1, mode != DeBruijnGraph::BASIC);
+
+    if (contigs.empty())
+        return std::make_shared<DBGSSHash>(k, mode);
+
+    std::string dump_path = "../tests/data/sshash_sequences/contigs.fa";
+    writeFastaFile(contigs, dump_path);
+
+    std::shared_ptr<DBGSSHash> graph;
+    graph = std::make_shared<DBGSSHash>(dump_path, k, mode, num_chars, true);
+
+    if (mode == DeBruijnGraph::PRIMARY)
+        return std::make_shared<CanonicalDBG>(
+                std::static_pointer_cast<DeBruijnGraph>(graph), 2 /* cache size */);
+
+    return graph;
+}
+
 template <>
 std::shared_ptr<DeBruijnGraph>
 build_graph<DBGSuccinct>(uint64_t k,
@@ -343,6 +415,10 @@ template
 std::shared_ptr<DeBruijnGraph>
 build_graph_batch<DBGSSHash>(uint64_t, std::vector<std::string>, DeBruijnGraph::Mode);
 
+template
+std::shared_ptr<DeBruijnGraph>
+build_graph_batch<DBGSSHashMonochromatic>(uint64_t, std::vector<std::string>, DeBruijnGraph::Mode);
+
 template <>
 std::shared_ptr<DeBruijnGraph>
 build_graph_batch<DBGBitmap>(uint64_t k,
@@ -549,6 +625,7 @@ template bool check_graph<DBGHashOrdered>(const std::string &, DeBruijnGraph::Mo
 template bool check_graph<DBGHashFast>(const std::string &, DeBruijnGraph::Mode, bool);
 template bool check_graph<DBGHashString>(const std::string &, DeBruijnGraph::Mode, bool);
 template bool check_graph<DBGSSHash>(const std::string &, DeBruijnGraph::Mode, bool);
+template bool check_graph<DBGSSHashMonochromatic>(const std::string &, DeBruijnGraph::Mode, bool);
 
 } // namespace test
 } // namespace mtg
diff --git a/metagraph/tests/graph/all/test_dbg_helpers.hpp b/metagraph/tests/graph/all/test_dbg_helpers.hpp
index bf394da089..d770d6134b 100644
--- a/metagraph/tests/graph/all/test_dbg_helpers.hpp
+++ b/metagraph/tests/graph/all/test_dbg_helpers.hpp
@@ -53,6 +53,13 @@ class DBGSuccinctCached : public DBGSuccinct {
           : DBGSuccinct(std::forward<Args>(args)...) {}
 };
 
+class DBGSSHashMonochromatic : public DBGSSHash {
+  public:
+    template <typename... Args>
+    DBGSSHashMonochromatic(Args&&... args)
+          : DBGSSHash(std::forward<Args>(args)...) {}
+};
+
 template <class Graph>
 std::shared_ptr<DeBruijnGraph>
 build_graph(uint64_t k,
@@ -93,6 +100,7 @@ typedef ::testing::Types<DBGBitmap,
                          DBGHashOrdered,
                          DBGHashFast,
                          DBGSSHash,
+                         DBGSSHashMonochromatic,
                          DBGSuccinct,
                          DBGSuccinctIndexed<1>,
                          DBGSuccinctIndexed<2>,
diff --git a/metagraph/tests/graph/test_canonical_dbg.cpp b/metagraph/tests/graph/test_canonical_dbg.cpp
index e0b3b0b8d0..8fc5b5643c 100644
--- a/metagraph/tests/graph/test_canonical_dbg.cpp
+++ b/metagraph/tests/graph/test_canonical_dbg.cpp
@@ -38,7 +38,8 @@ typedef ::testing::Types<DBGBitmap,
                          DBGSuccinct,
                          DBGSuccinctBloom<4, 1>,
                          DBGSuccinctCached,
-                         DBGSSHash> CanonicalGraphTypes;
+                         DBGSSHash,
+                         DBGSSHashMonochromatic> CanonicalGraphTypes;
 
 TYPED_TEST_SUITE(CanonicalDBGTest, CanonicalGraphTypes);
 
diff --git a/metagraph/tests/graph/test_dbg_canonical.cpp b/metagraph/tests/graph/test_dbg_canonical.cpp
index ef88f0caa9..e7c46e2d4e 100644
--- a/metagraph/tests/graph/test_dbg_canonical.cpp
+++ b/metagraph/tests/graph/test_dbg_canonical.cpp
@@ -26,7 +26,8 @@ typedef ::testing::Types<DBGBitmap,
                          DBGSuccinct,
                          DBGSuccinctBloom<4, 1>,
                          DBGSuccinctCached,
-                         DBGSSHash> CanonicalGraphTypes;
+                         DBGSSHash,
+                         DBGSSHashMonochromatic> CanonicalGraphTypes;
 TYPED_TEST_SUITE(DeBruijnGraphCanonicalTest, CanonicalGraphTypes);