RowDiffChar annotation format for SSHash graphs

metagraph/integration_tests/test_transform_anno.py

@@ -7,7 +7,7 @@
 import glob
 import os
 import gzip
-from base import TestingBase, METAGRAPH, TEST_DATA_DIR, graph_file_extension, anno_file_extension
+from base import TestingBase, METAGRAPH, TEST_DATA_DIR, graph_file_extension, anno_file_extension, DNA_MODE
 from helpers import get_test_class_name
 
 """Test operations on annotation columns"""
@@ -546,5 +546,71 @@ def test_round_trip_column_to_brwt_to_flat_to_disk(self):
             self._assert_query_identical(column_path, converted)
 
 
+@parameterized_class(('graph_mode',),
+    # protein graphs have no reverse complement, so only basic mode
+    input_values=[('basic',)] + ([('canonical',), ('primary',)] if DNA_MODE else []),
+    class_name_func=get_test_class_name
+)
+class TestRowDiffCharSSHash(TestingBase):
+    """row_diff_brwt -> row_diff_char_brwt on an SSHash graph must not change
+    any query result: only the successor encoding differs, the diffs are the same."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.tempdir = TemporaryDirectory()
+        cls.graph_path = cls.tempdir.name + '/graph.sshashdbg'
+        # SSHash is built from contigs without repeated k-mers, as in the unit tests
+        # (building it from the raw reads makes it write minimizer files without bound)
+        contig_mode = 'canonical' if cls.graph_mode == 'primary' else cls.graph_mode
+        cls._run_command(f'{METAGRAPH} build -p {NUM_THREADS} --mode {contig_mode} -k 21 '
+                         f'-o {cls.tempdir.name}/contigs {TEST_DATA_DIR}/transcripts_100.fa',
+                         "Build command")
+        primary_flag = ' --primary-kmers' if cls.graph_mode == 'primary' else ''
+        cls._run_command(f'{METAGRAPH} transform -p {NUM_THREADS} --to-fasta{primary_flag} '
+                         f'-o {cls.tempdir.name}/contigs {cls.tempdir.name}/contigs.dbg',
+                         "Transform command")
+        input_fa = cls.tempdir.name + '/contigs.fasta.gz'
+        build_mode = cls.graph_mode
+        cls._run_command(f'{METAGRAPH} build -p {NUM_THREADS} --graph sshash --mode {build_mode} '
+                         f'-k 21 -o {cls.graph_path} {input_fa}', "Build command")
+
+        cls._annotate_graph(TEST_DATA_DIR + '/transcripts_100.fa',
+                            cls.graph_path,
+                            cls.tempdir.name + '/annotation',
+                            'row_diff_brwt')
+        cls.brwt_path = cls.tempdir.name + '/annotation.row_diff_brwt.annodbg'
+
+        out_base = cls.tempdir.name + '/annotation_char'
+        command = (f'{METAGRAPH} transform_anno --anno-type row_diff_char_brwt '
+                   f'-i {cls.graph_path} -o {out_base} -p {NUM_THREADS} {cls.brwt_path}')
+        cls._run_command(command, "Transform_anno to row_diff_char_brwt")
+        cls.char_path = out_base + '.row_diff_char_brwt.annodbg'
+
+    def test_stats_match(self):
+        src = self._get_stats(f'-a {self.brwt_path}')
+        dst = self._get_stats(f'-a {self.char_path}')
+        self.assertEqual(src['returncode'], 0)
+        self.assertEqual(dst['returncode'], 0)
+        self.assertEqual('100', dst['labels'])
+        self.assertEqual(src['labels'], dst['labels'])
+        self.assertEqual(src['objects'], dst['objects'])
+        self.assertEqual(src['density'], dst['density'])
+        self.assertEqual('row_diff_char_brwt', dst['representation'])
+
+    @parameterized.expand(itertools.product(['labels', 'matches', 'signature'], [1, NUM_THREADS]))
+    def test_query_identical(self, query_mode, num_threads):
+        def query(anno_path):
+            cmd = (f'{METAGRAPH} query -i {self.graph_path} -a {anno_path} '
+                   f'--query-mode {query_mode} -p {num_threads} '
+                   f'--min-kmers-fraction-label 0 {TEST_DATA_DIR}/transcripts_1000.fa')
+            res = subprocess.run(cmd.split(), stdout=PIPE, stderr=PIPE)
+            self.assertEqual(res.returncode, 0, msg=res.stderr.decode())
+            # with several threads, results are printed in no fixed order
+            return sorted(res.stdout.splitlines())
+        expected = query(self.brwt_path)
+        self.assertGreater(len(expected), 0)
+        self.assertEqual(expected, query(self.char_path))
+
+
 if __name__ == '__main__':
     unittest.main()

metagraph/src/annotation/annotation_converters.cpp

@@ -300,6 +300,84 @@ std::pair<std::string, std::string> get_anchors_and_fork_fnames(const std::strin
     return std::make_pair(anchors_file, fork_succ_file);
 }
 
+std::pair<bit_vector_small, sdsl::int_vector<>>
+compute_rd_succ_char(const graph::DeBruijnGraph &graph,
+                     const bit_vector &fork_succ) {
+    using node_index = graph::DeBruijnGraph::node_index;
+
+    const auto &alphabet = graph.alphabet();
+    // skip sentinel '$' if present at position 0
+    size_t offset = (alphabet[0] == '$') ? 1 : 0;
+    size_t alpha_size = alphabet.size() - offset;
+
+    // bits per character: ceil(log2(alpha_size))
+    uint8_t bits_per_char = 0;
+    for (size_t s = alpha_size - 1; s > 0; s >>= 1)
+        ++bits_per_char;
+
+    // build char-to-index map for the non-sentinel alphabet
+    std::array<uint8_t, 256> char_to_idx{};
+    for (size_t i = 0; i < alpha_size; ++i) {
+        char_to_idx[static_cast<unsigned char>(alphabet[i + offset])] = i;
+    }
+
+    // Nodes are processed in blocks in parallel. The block size is a multiple of 64,
+    // so no two threads write to the same word of |needs_char_bv|. Each block keeps
+    // its chars packed until all blocks are concatenated in node order.
+    const uint64_t num_rows = graph.max_index() + 1;
+    constexpr uint64_t kBlockSize = 64 << 14;
+    const uint64_t num_blocks = (num_rows + kBlockSize - 1) / kBlockSize;
+    sdsl::bit_vector needs_char_bv(num_rows, 0);
+    std::vector<sdsl::int_vector<>> block_chars(num_blocks);
+
+    #pragma omp parallel for num_threads(get_num_threads()) schedule(dynamic)
+    for (uint64_t b = 0; b < num_blocks; ++b) {
+        std::vector<uint8_t> chars;
+        const uint64_t end = std::min(num_rows, (b + 1) * kBlockSize);
+        for (node_index node = std::max<uint64_t>(b * kBlockSize, 1); node < end; ++node) {
+            // If the graph can trivially determine the successor, no need to store
+            if (!graph.in_graph(node) || graph.has_no_outgoing(node)
+                    || graph.trivial_outgoing(node) != graph::DeBruijnGraph::npos)
+                continue;
+
+            // Need to store the character for this node
+            node_index succ = row_diff_successor(graph, node, fork_succ);
+            int idx = -1;
+            graph.call_outgoing_kmers(node, [&](node_index adj, char c) {
+                if (adj == succ)
+                    idx = char_to_idx[static_cast<unsigned char>(c)];
+            });
+            if (idx < 0) {
+                logger->error("No outgoing edge of node {} leads to its row-diff "
+                              "successor {}", node, succ);
+                std::exit(1);
+            }
+            needs_char_bv[node] = true;
+            chars.push_back(idx);
+        }
+        block_chars[b] = sdsl::int_vector<>(chars.size(), 0, bits_per_char);
+        std::copy(chars.begin(), chars.end(), block_chars[b].begin());
+    }
+
+    size_t count = 0;
+    for (const auto &chars : block_chars) {
+        count += chars.size();
+    }
+    sdsl::int_vector<> rd_succ_char(count, 0, bits_per_char);
+    auto it = rd_succ_char.begin();
+    for (auto &chars : block_chars) {
+        it = std::copy(chars.begin(), chars.end(), it);
+        chars = sdsl::int_vector<>();
+    }
+
+    logger->trace("Computed rd_succ_char: {} nodes total, {} need stored char ({:.1f}%), "
+                  "{} bits/char",
+                  graph.max_index(), count,
+                  100.0 * count / std::max<size_t>(graph.max_index(), 1),
+                  bits_per_char);
+    return { bit_vector_small(std::move(needs_char_bv)), std::move(rd_succ_char) };
+}
+
 std::tuple<std::vector<std::unique_ptr<bit_vector>>, LabelEncoder<std::string>, uint64_t>
 load_row_diff_columns(const std::vector<std::string> &files) {
     std::vector<std::unique_ptr<bit_vector>> columns;

metagraph/src/annotation/annotation_converters.hpp

@@ -139,6 +139,13 @@ void convert_row_diff_to_col_compressed(const std::vector<std::string> &files,
 
 std::pair<std::string, std::string> get_anchors_and_fork_fnames(const std::string &fbase);
 
+// Compute sparse successor character data for CharSucc mode.
+// Returns (needs_char bitvector, compact char array). Only nodes where the graph
+// can't trivially determine the successor have stored characters.
+std::pair<bit_vector_small, sdsl::int_vector<>>
+compute_rd_succ_char(const graph::DeBruijnGraph &graph,
+                     const bit_vector &fork_succ);
+
 template <class Annotator>
 StaticBinRelAnnotator<matrix::TupleCSCMatrix<typename Annotator::binary_matrix_type>, std::string>
 load_coords(Annotator&& anno, const std::vector<std::string> &files);

metagraph/src/annotation/binary_matrix/row_diff/row_diff.cpp

@@ -38,6 +38,36 @@ node_index row_diff_successor(const graph::DeBruijnGraph &graph,
     }
 }
 
+node_index row_diff_successor_char(
+        const graph::DeBruijnGraph &graph,
+        node_index node,
+        const bit_vector_small &rd_succ_needs_char,
+        const sdsl::int_vector<> &rd_succ_char) {
+    // Fast path: if the graph can trivially resolve the successor (e.g.,
+    // succ_is_next_ in SSHash), use it directly without any char lookup.
+    node_index trivial = graph.trivial_outgoing(node);
+    if (trivial != graph::DeBruijnGraph::npos)
+        return trivial;
+
+    // Slow path: look up stored character for this node
+    assert(node < rd_succ_needs_char.size());
+    assert(rd_succ_needs_char[node] && "node must have a stored char");
+    uint64_t rank = rd_succ_needs_char.rank1(node);
+    assert(rank > 0 && rank <= rd_succ_char.size());
+    uint8_t idx = rd_succ_char[rank - 1];
+    const auto &alphabet = graph.alphabet();
+    size_t offset = (alphabet[0] == '$') ? 1 : 0;
+    assert(idx + offset < alphabet.size());
+    char c = alphabet[idx + offset];
+    node_index succ = graph.traverse(node, c);
+    if (succ == graph::DeBruijnGraph::npos) {
+        // called in parallel regions, where an exception would terminate anyway
+        common::logger->error("Corrupt rd_succ_char: traverse({}, '{}') returned npos", node, c);
+        std::exit(1);
+    }
+    return succ;
+}
+
 
 namespace matrix {
 
@@ -124,7 +154,16 @@ IRowDiff::get_rd_ids(std::span<const BinaryMatrix::Row> row_ids, size_t num_thre
             if (anchor_[row])
                 break;
 
-            node = row_diff_successor(*graph_, node, fork_succ_);
+            switch (rd_mode_) {
+                case RowDiffMode::CharSucc:
+                    node = row_diff_successor_char(*graph_, node,
+                                                   rd_succ_needs_char_,
+                                                   rd_succ_char_);
+                    break;
+                case RowDiffMode::Standard:
+                    node = row_diff_successor(*graph_, node, fork_succ_);
+                    break;
+            }
         }
     }
 

metagraph/src/annotation/binary_matrix/row_diff/row_diff.hpp

@@ -6,6 +6,8 @@
 #include <string>
 #include <vector>
 
+#include <sdsl/int_vector.hpp>
+
 #include "annotation/binary_matrix/base/binary_matrix.hpp"
 #include "annotation/binary_matrix/column_sparse/column_major.hpp"
 #include "common/batch_accumulator.hpp"
@@ -30,8 +32,20 @@ graph::DeBruijnGraph::node_index row_diff_successor(const graph::DeBruijnGraph &
                                                     graph::DeBruijnGraph::node_index node,
                                                     const bit_vector &rd_succ);
 
+// CharSucc mode: look up successor from sparse char storage or trivial graph path
+graph::DeBruijnGraph::node_index row_diff_successor_char(
+        const graph::DeBruijnGraph &graph,
+        graph::DeBruijnGraph::node_index node,
+        const bit_vector_small &rd_succ_needs_char,
+        const sdsl::int_vector<> &rd_succ_char);
+
 namespace matrix {
 
+enum class RowDiffMode : uint8_t {
+    Standard = 0,  // original fork_succ bitvector
+    CharSucc = 1,  // successor chars, stored only where trivial_outgoing() fails
+};
+
 const std::string kRowDiffAnchorExt = ".anchors";
 const std::string kRowDiffForkSuccExt = ".rd_succ";
 
@@ -53,6 +67,19 @@ class IRowDiff {
 
     const fork_succ_bv_type& fork_succ() const { return fork_succ_; }
 
+    const bit_vector_small& rd_succ_needs_char() const { return rd_succ_needs_char_; }
+    const sdsl::int_vector<>& rd_succ_char() const { return rd_succ_char_; }
+
+    RowDiffMode mode() const { return rd_mode_; }
+
+    // Switch to CharSucc mode: the successor chars replace fork_succ
+    void set_successor_chars(bit_vector_small&& needs_char, sdsl::int_vector<>&& succ_char) {
+        rd_succ_needs_char_ = std::move(needs_char);
+        rd_succ_char_ = std::move(succ_char);
+        fork_succ_ = fork_succ_bv_type();
+        rd_mode_ = RowDiffMode::CharSucc;
+    }
+
   protected:
     template <std::derived_from<BinaryMatrix> Matrix>
     std::vector<BinaryMatrix::Row>
@@ -74,8 +101,11 @@ class IRowDiff {
                    H decode_diffs, Callback call_row, size_t num_threads) const;
 
     const graph::DeBruijnGraph *graph_ = nullptr;
+    RowDiffMode rd_mode_ = RowDiffMode::Standard;
     anchor_bv_type anchor_;
     fork_succ_bv_type fork_succ_;
+    bit_vector_small rd_succ_needs_char_;
+    sdsl::int_vector<> rd_succ_char_;
 };
 
 /**
@@ -151,6 +181,23 @@ class RowDiff : public IRowDiff, public BinaryMatrix {
 };
 
 
+/**
+ * RowDiffChar: RowDiff variant that stores successor characters only for nodes
+ * whose successor trivial_outgoing() cannot resolve, instead of the 1-bit
+ * fork_succ bitvector. Enables single-lookup traversal
+ * via graph.traverse(node, char) instead of enumerating all outgoing neighbors.
+ */
+template <class BaseMatrix>
+class RowDiffChar : public RowDiff<BaseMatrix> {
+  public:
+    template <typename... Args>
+    RowDiffChar(const graph::DeBruijnGraph *graph = nullptr, Args&&... args)
+        : RowDiff<BaseMatrix>(graph, std::forward<Args>(args)...) {
+        this->rd_mode_ = RowDiffMode::CharSucc;
+    }
+};
+
+
 template <std::derived_from<BinaryMatrix> Matrix>
 std::vector<BinaryMatrix::Row>
 IRowDiff::get_column_from_rows(const Matrix &matrix, BinaryMatrix::Column column) const {
@@ -296,7 +343,21 @@ template <class BaseMatrix>
 bool RowDiff<BaseMatrix>::load(std::istream &f) {
     auto pos = f.tellg();
     std::string version(4, '\0');
-    if (f.read(version.data(), 4) && version == "v2.0") {
+    if (f.read(version.data(), 4) && version == "v3.0") {
+        // CharSucc mode: anchor + needs_char(bit_vector_small) + rd_succ_char(int_vector<>) + diffs
+        rd_mode_ = RowDiffMode::CharSucc;
+        if constexpr(!std::is_same_v<BaseMatrix, ColumnMajor>) {
+            auto anchor_start = f.tellg();
+            if (!anchor_.load(f) || !rd_succ_needs_char_.load(f))
+                return false;
+            rd_succ_char_.load(f);
+            if (!f.good())
+                return false;
+            // anchor_ / successor chars are accessed randomly, hint the kernel.
+            utils::madvise_random_range(f, anchor_start, f.tellg() - anchor_start);
+        }
+    } else if (version == "v2.0") {
+        rd_mode_ = RowDiffMode::Standard;
         if constexpr(!std::is_same_v<BaseMatrix, ColumnMajor>) {
             auto anchor_start = f.tellg();
             if (!anchor_.load(f) || !fork_succ_.load(f))
@@ -307,6 +368,7 @@ bool RowDiff<BaseMatrix>::load(std::istream &f) {
     } else {
         // backward compatibility
         f.seekg(pos);
+        rd_mode_ = RowDiffMode::Standard;
         if constexpr(!std::is_same_v<BaseMatrix, ColumnMajor>) {
             auto anchor_start = f.tellg();
             if (!anchor_.load(f))
@@ -325,10 +387,23 @@ bool RowDiff<BaseMatrix>::load(std::istream &f) {
 
 template <class BaseMatrix>
 void RowDiff<BaseMatrix>::serialize(std::ostream &f) const {
-    f.write("v2.0", 4);
-    if constexpr(!std::is_same_v<BaseMatrix, ColumnMajor>) {
-        anchor_.serialize(f);
-        fork_succ_.serialize(f);
+    switch (rd_mode_) {
+        case RowDiffMode::CharSucc:
+            f.write("v3.0", 4);
+            if constexpr(!std::is_same_v<BaseMatrix, ColumnMajor>) {
+                anchor_.serialize(f);
+                rd_succ_needs_char_.serialize(f);
+                rd_succ_char_.serialize(f);
+            }
+            break;
+        case RowDiffMode::Standard:
+        default:
+            f.write("v2.0", 4);
+            if constexpr(!std::is_same_v<BaseMatrix, ColumnMajor>) {
+                anchor_.serialize(f);
+                fork_succ_.serialize(f);
+            }
+            break;
     }
     diffs_.serialize(f);
 }

metagraph/src/annotation/representation/annotation_matrix/annotation_matrix.cpp

@@ -198,6 +198,8 @@ template class StaticBinRelAnnotator<RowDiff<RowSparse>, std::string>;
 
 template class StaticBinRelAnnotator<RowDiff<RowDisk>, std::string>;
 
+template class StaticBinRelAnnotator<RowDiffChar<BRWT>, std::string>;
+
 template class StaticBinRelAnnotator<IntRowDiff<IntRowDisk>, std::string>;
 
 template class StaticBinRelAnnotator<TupleRowDiff<CoordRowDisk>, std::string>;