From c9bb0a21f5e7094d978a2b63467928c3aa92be37 Mon Sep 17 00:00:00 2001
From: Harun Mustafa <bilgisayaradami@gmail.com>
Date: Wed, 21 Oct 2020 18:35:10 +0200
Subject: [PATCH 1/2] Differential assembly support for canonical and primary
 graphs

DBGMode used in build_anno_graph
---
 README.md                                     |   4 +-
 .../binary_matrix/base/binary_matrix.cpp      |  11 +
 .../binary_matrix/base/binary_matrix.hpp      |   7 +
 metagraph/src/cli/assemble.cpp                | 290 ++++++--
 metagraph/src/cli/config/config.cpp           |  42 +-
 metagraph/src/cli/config/config.hpp           |   8 +-
 .../src/cli/load/load_annotated_graph.cpp     |   3 +-
 metagraph/src/common/vectors/bitmap.cpp       |  49 ++
 metagraph/src/common/vectors/bitmap.hpp       |  36 +
 .../src/graph/annotated_graph_algorithm.cpp   | 639 +++++++++++++-----
 .../src/graph/annotated_graph_algorithm.hpp   |  95 ++-
 .../src/graph/representation/masked_graph.hpp |   1 +
 .../annotation/test_annotated_dbg_helpers.cpp |  37 +-
 .../annotation/test_annotated_dbg_helpers.hpp |  15 +-
 .../test_annotated_graph_algorithm.cpp        | 395 ++++++-----
 .../tests/annotation/test_matrix_helpers.cpp  |  43 ++
 16 files changed, 1148 insertions(+), 527 deletions(-)

diff --git a/README.md b/README.md
index 96d1a131d3..29a7676d4d 100644
--- a/README.md
+++ b/README.md
@@ -206,9 +206,7 @@ Requires `M*V/8 + Size(BRWT)` bytes of RAM, where `M` is the number of rows in t
 ./metagraph assemble -v <GRAPH_DIR>/graph.dbg \
                         --unitigs \
                         -a <GRAPH_DIR>/annotation.column.annodbg \
-                        --label-mask-in LABEL_1 \
-                        --label-mask-in LABEL_2 \
-                        --label-mask-out LABEL_3 \
+                        --label-mask-file diff_assembly_experiment_file.txt \
                         -o diff_assembled.fa
 ```
 
diff --git a/metagraph/src/annotation/binary_matrix/base/binary_matrix.cpp b/metagraph/src/annotation/binary_matrix/base/binary_matrix.cpp
index d06d312463..bfad5bfc2d 100644
--- a/metagraph/src/annotation/binary_matrix/base/binary_matrix.cpp
+++ b/metagraph/src/annotation/binary_matrix/base/binary_matrix.cpp
@@ -1,5 +1,7 @@
 #include "binary_matrix.hpp"
 
+#include "common/vectors/bitmap.hpp"
+#include "common/vectors/bit_vector_adaptive.hpp"
 #include "common/serialization.hpp"
 
 
@@ -32,6 +34,15 @@ BinaryMatrix::slice_rows(const std::vector<Row> &row_ids) const {
     return slice;
 }
 
+void BinaryMatrix::slice_columns(const std::vector<Column> &column_ids,
+                                 const ColumnCallback &callback) const {
+    size_t nrows = num_rows();
+    for (size_t k = 0; k < column_ids.size(); ++k) {
+        Column j = column_ids[k];
+        callback(j, bitmap_generator(get_column(j), nrows));
+    }
+}
+
 template <typename RowType>
 StreamRows<RowType>::StreamRows(const std::string &filename, size_t offset) {
     std::ifstream instream(filename, std::ios::binary);
diff --git a/metagraph/src/annotation/binary_matrix/base/binary_matrix.hpp b/metagraph/src/annotation/binary_matrix/base/binary_matrix.hpp
index 085a968192..c1ea903700 100644
--- a/metagraph/src/annotation/binary_matrix/base/binary_matrix.hpp
+++ b/metagraph/src/annotation/binary_matrix/base/binary_matrix.hpp
@@ -9,6 +9,8 @@
 #include "common/vector.hpp"
 
 
+class bitmap;
+
 namespace mtg {
 namespace annot {
 namespace binmat {
@@ -21,6 +23,7 @@ class BinaryMatrix {
     typedef Vector<Column> SetBitPositions;
     typedef std::function<void(const SetBitPositions &)> RowCallback;
     typedef std::function<void(Row, Column)> ValueCallback;
+    typedef std::function<void(Column, bitmap&&)> ColumnCallback;
 
     virtual ~BinaryMatrix() {}
 
@@ -32,9 +35,13 @@ class BinaryMatrix {
     virtual SetBitPositions get_row(Row row) const = 0;
     virtual std::vector<SetBitPositions> get_rows(const std::vector<Row> &rows) const;
     virtual std::vector<Row> get_column(Column column) const = 0;
+
     // get all selected rows appended with -1 and concatenated
     virtual std::vector<Column> slice_rows(const std::vector<Row> &rows) const;
 
+    virtual void slice_columns(const std::vector<Column> &columns,
+                               const ColumnCallback &callback) const;
+
     virtual bool load(std::istream &in) = 0;
     virtual void serialize(std::ostream &out) const = 0;
 
diff --git a/metagraph/src/cli/assemble.cpp b/metagraph/src/cli/assemble.cpp
index 541864ac50..75c29d02ce 100644
--- a/metagraph/src/cli/assemble.cpp
+++ b/metagraph/src/cli/assemble.cpp
@@ -2,8 +2,10 @@
 
 #include <fmt/format.h>
 
+#include "common/algorithms.hpp"
 #include "common/logger.hpp"
 #include "common/unix_tools.hpp"
+#include "common/threads/threading.hpp"
 #include "seq_io/sequence_io.hpp"
 #include "config/config.hpp"
 #include "load/load_graph.hpp"
@@ -19,91 +21,188 @@ using mtg::common::logger;
 using mtg::graph::DeBruijnGraph;
 using mtg::graph::MaskedDeBruijnGraph;
 using mtg::graph::AnnotatedDBG;
+using mtg::graph::DifferentialAssemblyConfig;
+
+
+void clean_label_set(const AnnotatedDBG &anno_graph,
+                     std::vector<std::string> &label_set) {
+    label_set.erase(std::remove_if(label_set.begin(), label_set.end(),
+        [&](const std::string &label) {
+            bool exists = anno_graph.label_exists(label);
+            if (!exists)
+                logger->trace("Removing label {}", label);
+
+            return !exists;
+        }
+    ), label_set.end());
+
+    std::sort(label_set.begin(), label_set.end());
+    auto end = std::unique(label_set.begin(), label_set.end());
+    for (auto it = end; it != label_set.end(); ++it) {
+        logger->trace("Removing duplicate label {}", *it);
+    }
+    label_set.erase(end, label_set.end());
+}
 
 
 std::unique_ptr<MaskedDeBruijnGraph>
-mask_graph(const AnnotatedDBG &anno_graph, Config *config) {
+mask_graph_from_labels(const AnnotatedDBG &anno_graph,
+                       const std::vector<std::string> &label_mask_in,
+                       const std::vector<std::string> &label_mask_out,
+                       const std::vector<std::string> &label_mask_in_post,
+                       const std::vector<std::string> &label_mask_out_post,
+                       const DifferentialAssemblyConfig &diff_config,
+                       size_t num_threads) {
     auto graph = std::dynamic_pointer_cast<const DeBruijnGraph>(anno_graph.get_graph_ptr());
 
     if (!graph.get())
         throw std::runtime_error("Masking only supported for DeBruijnGraph");
 
-    // Remove non-present labels
-    config->label_mask_in.erase(
-        std::remove_if(config->label_mask_in.begin(),
-                       config->label_mask_in.end(),
-                       [&](const auto &label) {
-                           bool exists = anno_graph.label_exists(label);
-                           if (!exists)
-                               logger->trace("Removing mask-in label {}", label);
-
-                           return !exists;
-                       }),
-        config->label_mask_in.end()
-    );
-
-    config->label_mask_out.erase(
-        std::remove_if(config->label_mask_out.begin(),
-                       config->label_mask_out.end(),
-                       [&](const auto &label) {
-                           bool exists = anno_graph.label_exists(label);
-                           if (!exists)
-                               logger->trace("Removing mask-out label {}", label);
-
-                           return !exists;
-                       }),
-        config->label_mask_out.end()
-    );
-
-    logger->trace("Masked in: {}", fmt::join(config->label_mask_in, " "));
-    logger->trace("Masked out: {}", fmt::join(config->label_mask_out, " "));
-
-    if (!config->filter_by_kmer) {
-        return std::make_unique<MaskedDeBruijnGraph>(
-            graph,
-            mask_nodes_by_unitig_labels(
-                anno_graph,
-                config->label_mask_in,
-                config->label_mask_out,
-                std::max(1u, get_num_threads()),
-                config->label_mask_in_fraction,
-                config->label_mask_out_fraction,
-                config->label_other_fraction
-            )
-        );
+    std::vector<const std::vector<std::string>*> label_sets {
+        &label_mask_in, &label_mask_out,
+        &label_mask_in_post, &label_mask_out_post
+    };
+
+    for (const auto *label_set : label_sets) {
+        for (const auto *other_label_set : label_sets) {
+            if (label_set == other_label_set)
+                continue;
+
+            if (utils::count_intersection(label_set->begin(), label_set->end(),
+                                          other_label_set->begin(), other_label_set->end()))
+                logger->warn("Overlapping label sets");
+        }
     }
 
-    return std::make_unique<MaskedDeBruijnGraph>(
-        graph,
-        mask_nodes_by_node_label(
-            anno_graph,
-            config->label_mask_in,
-            config->label_mask_out,
-            [config,&anno_graph](auto index,
-                                 auto get_num_in_labels,
-                                 auto get_num_out_labels) {
-                assert(index != DeBruijnGraph::npos);
+    logger->trace("Masked in: {}", fmt::join(label_mask_in, " "));
+    logger->trace("Masked in (post-processing): {}", fmt::join(label_mask_in_post, " "));
+    logger->trace("Masked out: {}", fmt::join(label_mask_out, " "));
+    logger->trace("Masked out (post-processing): {}", fmt::join(label_mask_out_post, " "));
+
+    return std::make_unique<MaskedDeBruijnGraph>(mask_nodes_by_label(
+        anno_graph,
+        label_mask_in, label_mask_out,
+        label_mask_in_post, label_mask_out_post,
+        diff_config, num_threads
+    ));
+}
 
-                size_t num_in_labels = get_num_in_labels();
+DifferentialAssemblyConfig parse_diff_config(const std::string &config_str,
+                                             bool canonical) {
+    DifferentialAssemblyConfig diff_config;
+    diff_config.add_complement = canonical;
 
-                if (num_in_labels < config->label_mask_in_fraction
-                                        * config->label_mask_in.size())
-                    return false;
+    auto vals = utils::split_string(config_str, ",");
+    auto it = vals.begin();
+    if (it != vals.end()) {
+        diff_config.label_mask_in_kmer_fraction = std::stof(*it);
+        ++it;
+    }
+    if (it != vals.end()) {
+        diff_config.label_mask_in_unitig_fraction = std::stof(*it);
+        ++it;
+    }
+    if (it != vals.end()) {
+        diff_config.label_mask_out_kmer_fraction = std::stof(*it);
+        ++it;
+    }
+    if (it != vals.end()) {
+        diff_config.label_mask_out_unitig_fraction = std::stof(*it);
+        ++it;
+    }
+    if (it != vals.end()) {
+        diff_config.label_mask_other_unitig_fraction = std::stof(*it);
+        ++it;
+    }
 
-                size_t num_out_labels = get_num_out_labels();
+    assert(it == vals.end());
 
-                if (num_out_labels < config->label_mask_out_fraction
-                                        * config->label_mask_out.size())
-                    return false;
+    logger->trace("Per-kmer mask in fraction: {}", diff_config.label_mask_in_kmer_fraction);
+    logger->trace("Per-unitig mask in fraction: {}", diff_config.label_mask_in_unitig_fraction);
+    logger->trace("Per-kmer mask out fraction: {}", diff_config.label_mask_out_kmer_fraction);
+    logger->trace("Per-unitig mask out fraction: {}", diff_config.label_mask_out_unitig_fraction);
+    logger->trace("Per-unitig other label fraction: {}", diff_config.label_mask_other_unitig_fraction);
+    logger->trace("Include reverse complements: {}", diff_config.add_complement);
 
-                size_t num_total_labels = anno_graph.get_labels(index).size();
+    return diff_config;
+}
 
-                return num_total_labels - num_in_labels - num_out_labels
-                            <= config->label_other_fraction * num_total_labels;
-            },
-            std::max(1u, get_num_threads())
-        )
-    );
+typedef std::function<void(const graph::MaskedDeBruijnGraph&,
+                           const std::string& /* header */)> CallMaskedGraphHeader;
+
+void call_masked_graphs(const AnnotatedDBG &anno_graph, Config *config,
+                        const CallMaskedGraphHeader &callback,
+                        size_t num_parallel_graphs_masked = 1,
+                        size_t num_threads_per_graph = 1) {
+    assert(!config->label_mask_file.empty());
+
+    std::ifstream fin(config->label_mask_file);
+    if (!fin.good()) {
+        throw std::iostream::failure("Failed to read label mask file");
+        exit(1);
+    }
+
+    ThreadPool thread_pool(num_parallel_graphs_masked);
+    std::vector<std::string> shared_foreground_labels;
+    std::vector<std::string> shared_background_labels;
+
+    std::string line;
+    while (std::getline(fin, line)) {
+        if (line.empty() || line[0] == '#')
+            continue;
+
+        if (line[0] == '@') {
+            logger->trace("Counting shared k-mers");
+
+            // shared in and out labels
+            auto line_split = utils::split_string(line, "\t", false);
+            if (line_split.size() <= 1 || line_split.size() > 3)
+                throw std::iostream::failure("Each line in mask file must have 2-3 fields.");
+
+            // sync all assembly jobs before clearing shared labels
+            thread_pool.join();
+
+            shared_foreground_labels = utils::split_string(line_split[1], ",");
+            shared_background_labels = utils::split_string(
+                line_split.size() == 3 ? line_split[2] : "",
+                ","
+            );
+
+            clean_label_set(anno_graph, shared_foreground_labels);
+            clean_label_set(anno_graph, shared_background_labels);
+
+            continue;
+        }
+
+        thread_pool.enqueue([&](std::string line) {
+            auto line_split = utils::split_string(line, "\t", false);
+            if (line_split.size() <= 2 || line_split.size() > 4)
+                throw std::iostream::failure("Each line in mask file must have 3-4 fields.");
+
+            auto diff_config = parse_diff_config(line_split[1], config->canonical);
+
+            if (config->enumerate_out_sequences)
+                line_split[0] += ".";
+
+            auto foreground_labels = utils::split_string(line_split[2], ",");
+            auto background_labels = utils::split_string(
+                line_split.size() == 4 ? line_split[3] : "",
+                ","
+            );
+
+            clean_label_set(anno_graph, foreground_labels);
+            clean_label_set(anno_graph, background_labels);
+
+            callback(*mask_graph_from_labels(anno_graph,
+                                             foreground_labels, background_labels,
+                                             shared_foreground_labels,
+                                             shared_background_labels,
+                                             diff_config, num_threads_per_graph),
+                     line_split[0]);
+        }, std::move(line));
+    }
+
+    thread_pool.join();
 }
 
 
@@ -122,15 +221,56 @@ int assemble(Config *config) {
 
     logger->trace("Graph loaded in {} sec", timer.elapsed());
 
-    std::unique_ptr<graph::AnnotatedDBG> anno_graph;
     if (config->infbase_annotators.size()) {
-        anno_graph = initialize_annotated_dbg(graph, *config);
+        assert(config->label_mask_file.size());
+        auto anno_graph = initialize_annotated_dbg(graph, *config);
+
+        logger->trace("Generating masked graphs...");
+
+        std::filesystem::remove(
+            utils::remove_suffix(config->outfbase, ".gz", ".fasta") + ".fasta.gz"
+        );
+
+        std::mutex file_open_mutex;
+        std::mutex write_mutex;
 
-        logger->trace("Masking graph...");
+        size_t num_parallel_assemblies = std::max(
+            1u, std::min(config->parallel_assemblies, get_num_threads())
+        );
+        size_t num_threads_per_traversal = std::max(
+            size_t(1), get_num_threads() / num_parallel_assemblies
+        );
 
-        graph = mask_graph(*anno_graph, config);
+        call_masked_graphs(*anno_graph, config,
+            [&](const graph::MaskedDeBruijnGraph &graph, const std::string &header) {
+                std::lock_guard<std::mutex> file_lock(file_open_mutex);
+                seq_io::FastaWriter writer(config->outfbase, header,
+                                           config->enumerate_out_sequences,
+                                           get_num_threads() > 1, /* async write */
+                                           "a" /* append mode */);
+
+                if (config->unitigs || config->min_tip_size > 1) {
+                    graph.call_unitigs([&](const auto &unitig, auto&&) {
+                                           std::lock_guard<std::mutex> lock(write_mutex);
+                                           writer.write(unitig);
+                                       },
+                                       num_threads_per_traversal,
+                                       config->min_tip_size,
+                                       config->kmers_in_single_form);
+                } else {
+                    graph.call_sequences([&](const auto &seq, auto&&) {
+                                             std::lock_guard<std::mutex> lock(write_mutex);
+                                             writer.write(seq);
+                                         },
+                                         num_threads_per_traversal,
+                                         config->kmers_in_single_form);
+                }
+            },
+            num_parallel_assemblies,
+            num_threads_per_traversal
+        );
 
-        logger->trace("Masked in {} sec", timer.elapsed());
+        return 0;
     }
 
     logger->trace("Extracting sequences from graph...");
@@ -179,6 +319,8 @@ int assemble(Config *config) {
             get_num_threads(),
             config->min_tip_size
         );
+
+        return 0;
     }
 
     seq_io::FastaWriter writer(config->outfbase, config->header,
diff --git a/metagraph/src/cli/config/config.cpp b/metagraph/src/cli/config/config.cpp
index 19dac3e31a..ec28012047 100644
--- a/metagraph/src/cli/config/config.cpp
+++ b/metagraph/src/cli/config/config.cpp
@@ -142,6 +142,8 @@ Config::Config(int argc, char *argv[]) {
             set_num_threads(atoi(get_value(i++)));
         } else if (!strcmp(argv[i], "--parallel-nodes")) {
             parallel_nodes = atoi(get_value(i++));
+        } else if (!strcmp(argv[i], "--parallel-assemblies")) {
+            parallel_assemblies = atoi(get_value(i++));
         } else if (!strcmp(argv[i], "--max-path-length")) {
             max_path_length = atoi(get_value(i++));
         } else if (!strcmp(argv[i], "--parts-total")) {
@@ -250,6 +252,8 @@ Config::Config(int argc, char *argv[]) {
             host_address = get_value(i++);
         }else if (!strcmp(argv[i], "--suffix")) {
             suffix = get_value(i++);
+        } else if (!strcmp(argv[i], "--label-mask-file")) {
+            label_mask_file = get_value(i++);
         } else if (!strcmp(argv[i], "--initialize-bloom")) {
             initialize_bloom = true;
         } else if (!strcmp(argv[i], "--bloom-fpp")) {
@@ -330,18 +334,6 @@ Config::Config(int argc, char *argv[]) {
             print_welcome_message();
             print_usage(argv[0], identity);
             exit(0);
-        } else if (!strcmp(argv[i], "--label-mask-in")) {
-            label_mask_in.emplace_back(get_value(i++));
-        } else if (!strcmp(argv[i], "--label-mask-out")) {
-            label_mask_out.emplace_back(get_value(i++));
-        } else if (!strcmp(argv[i], "--label-mask-in-fraction")) {
-            label_mask_in_fraction = std::stof(get_value(i++));
-        } else if (!strcmp(argv[i], "--label-mask-out-fraction")) {
-            label_mask_out_fraction = std::stof(get_value(i++));
-        } else if (!strcmp(argv[i], "--label-other-fraction")) {
-            label_other_fraction = std::stof(get_value(i++));
-        } else if (!strcmp(argv[i], "--filter-by-kmer")) {
-            filter_by_kmer = true;
         } else if (!strcmp(argv[i], "--disk-swap")) {
             tmp_dir = get_value(i++);
         } else if (!strcmp(argv[i], "--disk-cap-gb")) {
@@ -358,6 +350,9 @@ Config::Config(int argc, char *argv[]) {
     if (parallel_nodes == static_cast<unsigned int>(-1))
         parallel_nodes = get_num_threads();
 
+    if (parallel_assemblies == static_cast<unsigned int>(-1))
+        parallel_assemblies = get_num_threads();
+
     if (identity == TRANSFORM && to_fasta)
         identity = ASSEMBLE;
 
@@ -500,6 +495,12 @@ Config::Config(int argc, char *argv[]) {
         print_usage_and_exit = true;
     }
 
+    if ((identity == ASSEMBLE || identity == TRANSFORM)
+            && (infbase_annotators.size() && label_mask_file.empty())) {
+        std::cerr << "Error: annotator passed, but no label mask file provided" << std::endl;
+        print_usage_and_exit = true;
+    }
+
     if (identity == ANNOTATE_COORDINATES && outfbase.empty())
         outfbase = utils::remove_suffix(infbase, ".dbg",
                                                  ".orhashdbg",
@@ -926,13 +927,16 @@ void Config::print_usage(const std::string &prog_name, IdentityType identity) {
             fprintf(stderr, "\t   --header [STR] \theader for sequences in FASTA output []\n");
             fprintf(stderr, "\t-p --parallel [INT] \tuse multiple threads for computation [1]\n");
             fprintf(stderr, "\n");
-            fprintf(stderr, "\t-a --annotator [STR] \t\t\tannotator to load []\n");
-            fprintf(stderr, "\t   --label-mask-in [STR] \t\tlabel to include in masked graph\n");
-            fprintf(stderr, "\t   --label-mask-out [STR] \t\tlabel to exclude from masked graph\n");
-            fprintf(stderr, "\t   --label-mask-in-fraction [FLOAT] \tminimum fraction of mask-in labels among the set of masked labels [1.0]\n");
-            fprintf(stderr, "\t   --label-mask-out-fraction [FLOAT] \tmaximum fraction of mask-out labels among the set of masked labels [0.0]\n");
-            fprintf(stderr, "\t   --label-other-fraction [FLOAT] \tmaximum fraction of other labels allowed [1.0]\n");
-            fprintf(stderr, "\t   --filter-by-kmer \t\t\tmask out graph k-mers individually [off]\n");
+            fprintf(stderr, "\t-a --annotator [STR] \t\tannotator to load []\n");
+            fprintf(stderr, "\t   --parallel-assemblies [INT] \tnumber of assembly experiments to run in parallel [n_threads]\n");
+            fprintf(stderr, "\t   --label-mask-file [STR] \tfile describing labels to mask in and out and their relative fractions []\n");
+            fprintf(stderr, "\t                       \t\tA k-mer is an in-k-mer if it has at least in_kmer_frac in-labels.\n");
+            fprintf(stderr, "\t                       \t\tA k-mer is an out-k-mer if it has more than out_kmer_frac out-labels.\n");
+            fprintf(stderr, "\t                       \t\tA unitig is included if it has at least in_unitig_frac in-k-mers and at most out_unitig_frac out-k-mers.\n");
+            fprintf(stderr, "\t                       \t\tA unitig is excluded if more than other_unitig_frac of its k-mers are not in the in- or out-label sets.\n");
+            fprintf(stderr, "\t                       \t\texample: '<exp_label>\\t<in_kmer_frac>,<in_unitig_frac>,<out_kmer_frac>,<out_unitig_frac>,<other_unitig_frac>\\t<inlabel1>,<inlabel2>,...\\t<outlabel1>,<outlabel2>,...'\n");
+            fprintf(stderr, "\t                       \t\tIf exp_label == @, then the second and third fields are lists of in- and out-labels, respectively, that apply to all subsequent lines.\n");
+            fprintf(stderr, "\t                       \t\texample: '@\\t<inlabel1>,<inlabel2>,...\\t<outlabel1>,<outlabel2>,...'\n");
         } break;
         case STATS: {
             fprintf(stderr, "Usage: %s stats [options] GRAPH1 [[GRAPH2] ...]\n\n", prog_name.c_str());
diff --git a/metagraph/src/cli/config/config.hpp b/metagraph/src/cli/config/config.hpp
index eb4458265a..aabbbd078b 100644
--- a/metagraph/src/cli/config/config.hpp
+++ b/metagraph/src/cli/config/config.hpp
@@ -56,7 +56,6 @@ class Config {
     bool map_sequences = false;
     bool align_sequences = false;
     bool align_both_strands = false;
-    bool filter_by_kmer = false;
     bool output_json = false;
     bool optimize = false;
 
@@ -86,6 +85,7 @@ class Config {
     unsigned int bloom_max_num_hash_functions = 10;
     unsigned int num_columns_cached = 10;
     unsigned int max_hull_forks = 4;
+    unsigned int parallel_assemblies = -1; // by default, run |parallel| assemblies in parallel, one per thread
 
     unsigned long long int query_batch_size_in_bytes = 100'000'000;
     unsigned long long int num_rows_subsampled = 1'000'000;
@@ -115,9 +115,6 @@ class Config {
     size_t alignment_max_num_seeds_per_locus = std::numeric_limits<size_t>::max();
 
     double discovery_fraction = 0.7;
-    double label_mask_in_fraction = 1.0;
-    double label_mask_out_fraction = 0.0;
-    double label_other_fraction = 1.0;
     double min_count_quantile = 0.;
     double max_count_quantile = 1.;
     double bloom_fpp = 1.0;
@@ -129,8 +126,6 @@ class Config {
     std::vector<std::string> fnames;
     std::vector<std::string> anno_labels;
     std::vector<std::string> infbase_annotators;
-    std::vector<std::string> label_mask_in;
-    std::vector<std::string> label_mask_out;
     std::string outfbase;
     std::string infbase;
     std::string rename_instructions_file;
@@ -141,6 +136,7 @@ class Config {
     std::string fasta_anno_comment_delim = UNINITIALIZED_STR;
     std::string header = "";
     std::string host_address;
+    std::string label_mask_file;
     uint32_t max_path_length = 50;
 
     std::filesystem::path tmp_dir;
diff --git a/metagraph/src/cli/load/load_annotated_graph.cpp b/metagraph/src/cli/load/load_annotated_graph.cpp
index aa4d0c31bc..9a16f451fc 100644
--- a/metagraph/src/cli/load/load_annotated_graph.cpp
+++ b/metagraph/src/cli/load/load_annotated_graph.cpp
@@ -22,7 +22,7 @@ using mtg::common::logger;
 std::unique_ptr<AnnotatedDBG> initialize_annotated_dbg(std::shared_ptr<DeBruijnGraph> graph,
                                                        const Config &config) {
     // TODO: check and wrap into canonical only if the graph is primary
-    if (config.canonical && !graph->is_canonical_mode())
+    if (config.canonical && !graph->is_canonical_mode() && config.identity != Config::ASSEMBLE)
         graph = std::make_shared<CanonicalDBG>(graph, true);
 
     uint64_t max_index = graph->max_index();
@@ -85,6 +85,5 @@ std::unique_ptr<AnnotatedDBG> initialize_annotated_dbg(const Config &config) {
     return initialize_annotated_dbg(load_critical_dbg(config.infbase), config);
 }
 
-
 } // namespace cli
 } // namespace mtg
diff --git a/metagraph/src/common/vectors/bitmap.cpp b/metagraph/src/common/vectors/bitmap.cpp
index 507582c3af..2e998f7ac0 100644
--- a/metagraph/src/common/vectors/bitmap.cpp
+++ b/metagraph/src/common/vectors/bitmap.cpp
@@ -369,3 +369,52 @@ void bitmap_lazy::call_ones_in_range(uint64_t begin, uint64_t end,
             callback(i);
     }
 }
+
+
+////////////////////////////////////////////////////////////////
+//                         bitmap_generator                   //
+////////////////////////////////////////////////////////////////
+
+bitmap_generator::bitmap_generator(size_t size, bool value)
+      : bitmap_generator([value,size](const auto &index_callback) {
+                             if (value) {
+                                 for (size_t i = 0; i < size; ++i) {
+                                     index_callback(i);
+                                 }
+                             }
+                         }, size, size * value) {}
+
+bitmap_generator
+::bitmap_generator(std::function<void(const VoidCall<uint64_t>&)>&& generator,
+                   size_t size,
+                   size_t num_set_bits) noexcept
+      : size_(size), num_set_bits_(num_set_bits), generator_(std::move(generator)) {}
+
+bitmap_generator
+::bitmap_generator(bitmap&& base, std::function<uint64_t(uint64_t)>&& index_transformer,
+                   size_t size, size_t num_set_bits) noexcept
+      : size_(size == static_cast<size_t>(-1) ? base.size() : size),
+        num_set_bits_(num_set_bits == static_cast<size_t>(-1) ? base.num_set_bits() : num_set_bits) {
+    auto get_generator = [&](bitmap&& v, auto&& transformer) {
+        return [&](const auto &callback) {
+            v.call_ones([&](auto i) { callback(transformer(i)); });
+        };
+    };
+
+    generator_ = get_generator(std::move(base), std::move(index_transformer));
+}
+
+void bitmap_generator::add_to(sdsl::bit_vector *other) const {
+    assert(other);
+    assert(other->size() == size());
+    call_ones([other](auto i) { (*other)[i] = true; });
+}
+
+void bitmap_generator::call_ones_in_range(uint64_t begin, uint64_t end,
+                                          const VoidCall<uint64_t> &callback) const {
+    generator_([&](uint64_t i) {
+        assert(i < size_);
+        if (i >= begin && i < end)
+            callback(i);
+    });
+}
diff --git a/metagraph/src/common/vectors/bitmap.hpp b/metagraph/src/common/vectors/bitmap.hpp
index 3f7dfa98c6..2a0869b12a 100644
--- a/metagraph/src/common/vectors/bitmap.hpp
+++ b/metagraph/src/common/vectors/bitmap.hpp
@@ -187,4 +187,40 @@ class bitmap_lazy : public bitmap {
     size_t num_set_bits_;
 };
 
+class bitmap_generator : public bitmap {
+  public:
+    explicit bitmap_generator(size_t size = 0, bool value = false);
+
+    bitmap_generator(std::vector<uint64_t>&& set_bits, size_t size) noexcept
+          : size_(size), num_set_bits_(set_bits.size()),
+            generator_([s=std::move(set_bits)](const VoidCall<uint64_t> &callback) {
+                           std::for_each(s.begin(), s.end(), callback);
+                       }) {}
+
+    bitmap_generator(std::function<void(const VoidCall<uint64_t>&)>&& generator,
+                     size_t size,
+                     size_t num_set_bits = -1) noexcept;
+
+    bitmap_generator(bitmap&& base,
+                     std::function<uint64_t(uint64_t)>&& index_transformer
+                         = [](auto i) { return i; },
+                     size_t size = -1,
+                     size_t num_set_bits = -1) noexcept;
+
+    bool operator[](uint64_t) const { throw std::runtime_error("Not implemented"); }
+    uint64_t get_int(uint64_t, uint32_t) const { throw std::runtime_error("Not implemented"); }
+
+    uint64_t size() const { return size_; }
+    uint64_t num_set_bits() const { return num_set_bits_; }
+
+    void add_to(sdsl::bit_vector *other) const;
+    void call_ones_in_range(uint64_t begin, uint64_t end,
+                            const VoidCall<uint64_t> &callback) const;
+
+  private:
+    size_t size_;
+    size_t num_set_bits_;
+    std::function<void(const VoidCall<uint64_t>&)> generator_;
+};
+
 #endif // __BITMAP_HPP__
diff --git a/metagraph/src/graph/annotated_graph_algorithm.cpp b/metagraph/src/graph/annotated_graph_algorithm.cpp
index e888219f87..908565dd44 100644
--- a/metagraph/src/graph/annotated_graph_algorithm.cpp
+++ b/metagraph/src/graph/annotated_graph_algorithm.cpp
@@ -1,253 +1,520 @@
 #include "annotated_graph_algorithm.hpp"
 
-#include <unordered_set>
+#include <tsl/hopscotch_set.h>
 
+#include "common/logger.hpp"
+#include "common/seq_tools/reverse_complement.hpp"
 #include "common/vectors/vector_algorithm.hpp"
-#include "common/vector_map.hpp"
-#include "kmer/alphabets.hpp"
-#include "annotation/representation/column_compressed/annotate_column_compressed.hpp"
-#include "graph/alignment/aligner_helper.hpp"
+#include "common/threads/threading.hpp"
+#include "common/vectors/bitmap.hpp"
 #include "graph/representation/masked_graph.hpp"
+#include "graph/representation/succinct/dbg_succinct.hpp"
+#include "graph/representation/succinct/boss_construct.hpp"
 
 
 namespace mtg {
 namespace graph {
 
+using mtg::graph::boss::BOSS;
+using mtg::graph::boss::BOSSConstructor;
+using mtg::common::logger;
+
 typedef AnnotatedDBG::node_index node_index;
 typedef AnnotatedDBG::row_index row_index;
 typedef AnnotatedDBG::Annotator::Label Label;
-typedef Alignment<DeBruijnGraph::node_index> DBGAlignment;
 
+typedef std::function<size_t()> LabelCountCallback;
+
+constexpr bool MAKE_BOSS = true;
+
+
+/**
+ * Return an int_vector<>, bitmap pair, each of length anno_graph.get_graph().max_index().
+ * For an index i, the int_vector will contain a packed integer representing the
+ * number of labels in labels_in and labels_out which the k-mer of index i is
+ * annotated with. The least significant half of each integer represents the count
+ * from labels_in, while the most significant half represents the count from
+ * labels_out.
+ * The returned bitmap is a binarization of the int_vector
+ */
+std::pair<sdsl::int_vector<>, std::unique_ptr<bitmap>>
+construct_diff_label_count_vector(const AnnotatedDBG &anno_graph,
+                                  const std::vector<Label> &labels_in,
+                                  const std::vector<Label> &labels_out,
+                                  size_t num_threads);
+
+
+template <class GetKeptIntervals>
+void update_masked_graph_by_unitig(MaskedDeBruijnGraph &masked_graph,
+                                   const GetKeptIntervals &get_kept_intervals,
+                                   size_t num_threads);
+
+template <class KeepNode>
+void update_masked_graph_by_node(MaskedDeBruijnGraph &masked_graph,
+                                 const KeepNode &keep_node,
+                                 size_t num_threads);
+
+std::shared_ptr<MaskedDeBruijnGraph>
+make_initial_masked_graph(std::shared_ptr<const DeBruijnGraph> &graph_ptr,
+                          sdsl::int_vector<> &counts,
+                          std::unique_ptr<bitmap>&& mask,
+                          bool add_complement,
+                          bool make_boss,
+                          size_t num_threads);
+
+
+MaskedDeBruijnGraph mask_nodes_by_label(const AnnotatedDBG &anno_graph,
+                                        const std::vector<Label> &labels_in,
+                                        const std::vector<Label> &labels_out,
+                                        const std::vector<Label> &labels_in_post,
+                                        const std::vector<Label> &labels_out_post,
+                                        const DifferentialAssemblyConfig &config,
+                                        size_t num_threads) {
+    auto graph_ptr = std::dynamic_pointer_cast<const DeBruijnGraph>(
+        anno_graph.get_graph_ptr()
+    );
+
+    logger->trace("Generating initial mask");
+
+    // Construct initial masked graph from union of labels in labels_in
+    auto count_vector = construct_diff_label_count_vector(
+        anno_graph, labels_in, labels_out, num_threads
+    );
+    auto &[counts, union_mask] = count_vector;
+
+    // counts is a double-width, interleaved vector where the significant bits
+    // represent the out-label count and the least significant bits represent
+    // the in-label count
+    uint32_t width = counts.width() / 2;
+    uint64_t int_mask = (uint64_t(1) << width) - 1;
+
+    auto masked_graph = make_initial_masked_graph(graph_ptr,
+                                                  counts, std::move(union_mask),
+                                                  config.add_complement, MAKE_BOSS,
+                                                  num_threads);
+
+    // Filter unitigs from masked graph based on filtration criteria
+    logger->trace("Filtering out background");
+
+    tsl::hopscotch_set<std::string> masked_labels;
+    tsl::hopscotch_set<std::string> labels_in_post_set(labels_in_post.begin(),
+                                                       labels_in_post.end());
+    tsl::hopscotch_set<std::string> labels_out_post_set(labels_out_post.begin(),
+                                                        labels_out_post.end());
+
+    if (config.label_mask_other_unitig_fraction != 1.0) {
+        masked_labels.insert(labels_in.begin(), labels_in.end());
+        masked_labels.insert(labels_out.begin(), labels_out.end());
+    } else if (config.label_mask_in_unitig_fraction == 0.0
+            && config.label_mask_out_unitig_fraction == 1.0
+            && labels_in_post.empty() && labels_out_post.empty()) {
+        if (config.label_mask_in_kmer_fraction == 0.0
+                && config.label_mask_out_kmer_fraction == 1.0) {
+            logger->trace("Bypassing background filtration");
+            return std::move(*masked_graph);
+        }
+
+        logger->trace("Filtering by node");
+        update_masked_graph_by_node(*masked_graph, [&](node_index node) {
+            uint64_t count = count_vector.first[node];
+            uint64_t in_count = count & int_mask;
+            uint64_t out_count = count >> width;
+            uint64_t sum = in_count + out_count;
+            return in_count >= config.label_mask_in_kmer_fraction * sum
+                && out_count <= config.label_mask_out_kmer_fraction * sum;
+        }, num_threads);
+
+        return std::move(*masked_graph);
+    }
+
+    logger->trace("Filtering by unitig");
 
-std::unique_ptr<bitmap_vector>
-mask_nodes_by_unitig(const DeBruijnGraph &graph,
-                     const KeepUnitigPath &keep_unitig,
-                     size_t num_threads) {
-    sdsl::bit_vector unitig_mask(graph.max_index() + 1, false);
+    size_t num_labels = anno_graph.get_annotation().num_labels();
+    bool check_other = config.label_mask_other_unitig_fraction != 1.0;
+    counts.width(width);
 
-    const bool atomic = num_threads > 1;
+    update_masked_graph_by_unitig(*masked_graph,
+                                  [&](const auto &unitig, const auto &path)
+                                      -> std::vector<std::pair<size_t, size_t>> {
+        auto &counts = count_vector.first;
+        sdsl::bit_vector in_mask(path.size(), false);
+        sdsl::bit_vector out_mask(path.size(), false);
+        sdsl::bit_vector other_mask(check_other ? path.size() : 0, false);
+
+        size_t min_label_in_count = config.label_mask_in_kmer_fraction * labels_in.size();
+        size_t max_label_out_count = config.label_mask_out_kmer_fraction * labels_out.size();
+
+        size_t label_in_cutoff = std::ceil(config.label_mask_in_unitig_fraction * path.size());
+        size_t label_out_cutoff = std::floor(config.label_mask_out_unitig_fraction * path.size());
+        size_t other_cutoff = std::floor(config.label_mask_other_unitig_fraction * path.size());
+        size_t in_kmer_count = 0;
+        size_t out_kmer_count = 0;
+
+        for (size_t i = 0; i < path.size(); ++i) {
+            if (counts[2 * path[i]] >= min_label_in_count && !in_mask[i]) {
+                in_mask[i] = true;
+                ++in_kmer_count;
+            }
 
-    graph.call_unitigs([&](const std::string &unitig, const auto &path) {
-        if (keep_unitig(unitig, path)) {
-            for (DeBruijnGraph::node_index node : path) {
-                set_bit(unitig_mask.data(), node, atomic);
+            if (counts[2 * path[i] + 1] > max_label_out_count && !out_mask[i]) {
+                out_mask[i] = true;
+                ++out_kmer_count;
             }
         }
+
+        if (check_other
+                || (in_kmer_count < label_in_cutoff && labels_in_post.size())
+                || (out_kmer_count < label_out_cutoff && labels_out_post.size())) {
+            for (auto &[label, sig] : anno_graph.get_top_label_signatures(unitig, num_labels)) {
+                if (check_other && !masked_labels.count(label))
+                    bitmap_vector(std::move(sig)).add_to(&other_mask);
+
+                if (in_kmer_count < label_in_cutoff && labels_in_post_set.count(label)) {
+                    for (size_t i = 0; i < path.size(); ++i) {
+                        if (!sig[i])
+                            continue;
+
+                        if (!in_mask[i] && ++counts[2 * path[i]] >= min_label_in_count) {
+                            in_mask[i] = true;
+                            ++in_kmer_count;
+                        }
+                    }
+                }
+
+                if (out_kmer_count < label_out_cutoff && labels_out_post_set.count(label)) {
+                    for (size_t i = 0; i < path.size(); ++i) {
+                        if (!sig[i])
+                            continue;
+
+                        if (!out_mask[i] && ++counts[2 * path[i] + 1] >= max_label_out_count) {
+                            out_mask[i] = true;
+                            ++out_kmer_count;
+                        }
+                    }
+                }
+            }
+        }
+
+        size_t begin = next_bit(in_mask, 0);
+
+        if (begin == in_mask.size())
+            return {};
+
+        size_t end = prev_bit(in_mask, in_mask.size() - 1) + 1;
+        assert(end > begin);
+        out_kmer_count -= count_ones(out_mask, 0, begin) + count_ones(out_mask, end, out_mask.size());
+        size_t other_kmer_count = check_other ? count_ones(other_mask, begin, end) : 0;
+
+        if (out_kmer_count > label_out_cutoff)
+            return {};
+
+        if (other_kmer_count > other_cutoff)
+            return {};
+
+        return { std::make_pair(begin, end) };
+
     }, num_threads);
 
-    return std::make_unique<bitmap_vector>(std::move(unitig_mask));
+    return std::move(*masked_graph);
 }
 
-std::unique_ptr<bitmap_vector>
-mask_nodes_by_unitig_labels(const AnnotatedDBG &anno_graph,
-                            const std::vector<Label> &labels_in,
-                            const std::vector<Label> &labels_out,
-                            size_t num_threads,
-                            double label_mask_in_fraction,
-                            double label_mask_out_fraction,
-                            double label_other_fraction) {
-    const auto &dbg = anno_graph.get_graph();
-    const auto &annotation = anno_graph.get_annotation();
-    const auto &label_encoder = annotation.get_label_encoder();
-
-    assert(labels_in.size() + labels_out.size() <= annotation.num_labels());
-
-    const double label_in_factor = label_mask_in_fraction * labels_in.size();
-    const double label_out_factor = label_mask_out_fraction * labels_out.size();
-
-    std::unordered_set<uint64_t> labels_in_enc;
-    for (const auto &label_in : labels_in) {
-        labels_in_enc.emplace(label_encoder.encode(label_in));
-    }
 
-    std::unordered_set<uint64_t> labels_out_enc;
-    for (const auto &label_out : labels_out) {
-        labels_out_enc.emplace(label_encoder.encode(label_out));
-    }
+std::shared_ptr<MaskedDeBruijnGraph>
+make_initial_masked_graph(std::shared_ptr<const DeBruijnGraph> &graph_ptr,
+                          sdsl::int_vector<> &counts,
+                          std::unique_ptr<bitmap>&& mask,
+                          bool add_complement,
+                          bool make_boss,
+                          size_t num_threads) {
+    auto masked_graph = std::make_shared<MaskedDeBruijnGraph>(
+        graph_ptr,
+        std::move(mask),
+        graph_ptr->is_canonical_mode()
+    );
+    logger->trace("Constructed masked graph with {} nodes", masked_graph->num_nodes());
+    bool masked_canonical = add_complement || graph_ptr->is_canonical_mode();
+
+    if (make_boss || (add_complement && !graph_ptr->is_canonical_mode())) {
+        // we can't guarantee that the reverse complement is present, so
+        // construct a new subgraph
+        logger->trace("Constructing BOSS from labeled subgraph");
+        uint32_t width = counts.width() / 2;
+        std::vector<std::pair<std::string, sdsl::int_vector<>>> contigs;
+        std::mutex add_mutex;
+        BOSSConstructor constructor(
+            graph_ptr->get_k() - 1,
+            masked_canonical,
+            0 /* count width */, "" /* suffix */, 1 /* num_threads */,
+            masked_graph->num_nodes() * 32
+        );
+        masked_graph->call_sequences([&](const std::string &seq, const auto &path) {
+            sdsl::int_vector<> path_counts(path.size(), 0, width * 2);
+            auto it = path_counts.begin();
+            for (node_index i : path) {
+                *it = counts[i];
+                ++it;
+            }
 
-    return mask_nodes_by_unitig(dbg, [&](const auto &, const auto &path) {
-        VectorMap<row_index, size_t> index_counts;
-        for (node_index i : path) {
-            index_counts[anno_graph.graph_to_anno_index(i)]++;
-        }
+            std::lock_guard<std::mutex> lock(add_mutex);
+            contigs.emplace_back(seq, std::move(path_counts));
+            constructor.add_sequence(seq);
+        }, num_threads, masked_canonical);
+
+        auto dbg_succ = std::make_shared<DBGSuccinct>(new BOSS(&constructor),
+                                                      masked_canonical);
+
+        // instead of keeping multiple masks (one for valid nodes and another
+        // for the masked graph), transfer the valid node mask to the masked graph
+        dbg_succ->mask_dummy_kmers(num_threads, false);
+        std::unique_ptr<bit_vector> dummy_mask(dbg_succ->release_mask());
+        graph_ptr = dbg_succ;
+        sdsl::bit_vector new_indicator(graph_ptr->max_index() + 1, false);
+        dummy_mask->add_to(&new_indicator);
+        masked_graph = std::make_shared<MaskedDeBruijnGraph>(
+            graph_ptr,
+            std::make_unique<bitmap_vector>(std::move(new_indicator)),
+            masked_canonical
+        );
+
+        logger->trace("Reconstructing count vector");
+        counts = aligned_int_vector(graph_ptr->max_index() + 1, 0, width * 2, 16);
 
-        size_t other_count = 0;
-        size_t in_count = 0;
-        size_t out_count = 0;
-        const size_t out_count_cutoff = label_out_factor * path.size();
+        std::atomic_thread_fence(std::memory_order_release);
 
-        for (const auto &pair : annotation.count_labels(index_counts.values_container())) {
-            if (labels_in_enc.find(pair.first) != labels_in_enc.end()) {
-                in_count += pair.second;
+        std::mutex vector_backup_mutex;
+        constexpr std::memory_order memorder = std::memory_order_relaxed;
 
-            } else if (labels_out_enc.find(pair.first) != labels_out_enc.end()) {
-                // early cutoff
-                if ((out_count += pair.second) > out_count_cutoff)
-                    return false;
+        #pragma omp parallel for num_threads(num_threads) schedule(dynamic)
+        for (size_t l = 0; l < contigs.size(); ++l) {
+            auto &[seq, seq_counts] = contigs[l];
+            size_t j = 0;
+            graph_ptr->map_to_nodes_sequentially(seq, [&](node_index i) {
+                atomic_exchange(counts, i, contigs[l].second[j++],
+                                vector_backup_mutex, memorder);
+            });
+            assert(j == seq_counts.size());
+        }
 
-            } else {
-                other_count += pair.second;
+        if (masked_canonical) {
+            #pragma omp parallel for num_threads(num_threads) schedule(dynamic)
+            for (size_t l = 0; l < contigs.size(); ++l) {
+                auto &[seq, seq_counts] = contigs[l];
+                reverse_complement(seq.begin(), seq.end());
+
+                // counts stores counts for in labels and out labels interleaved,
+                // so to add values properly, it should be reshaped first
+                size_t j = seq_counts.size();
+                graph_ptr->map_to_nodes_sequentially(seq, [&](node_index i) {
+                    uint64_t old_val = atomic_exchange(counts, i, contigs[l].second[--j],
+                                                       vector_backup_mutex, memorder);
+                    std::ignore = old_val;
+                    assert(!old_val);
+                });
+                assert(!j);
             }
         }
 
-        return (in_count >= label_in_factor * path.size())
-            && (other_count <= label_other_fraction * (in_count + out_count + other_count));
-    }, num_threads);
+        std::atomic_thread_fence(std::memory_order_acquire);
+
+        logger->trace("Constructed BOSS with {} nodes", graph_ptr->num_nodes());
+    }
+
+    assert(counts.size() == graph_ptr->max_index() + 1);
+
+    return masked_graph;
 }
 
-sdsl::int_vector<>
+/**
+ * Return an int_vector<>, bitmap pair, each of length anno_graph.get_graph().max_index().
+ * For an index i, the int_vector will contain a packed integer representing the
+ * number of labels in labels_in and labels_out which the k-mer of index i is
+ * annotated with. The least significant half of each integer represents the count
+ * from labels_in, while the most significant half represents the count from
+ * labels_out.
+ * The returned bitmap is a binarization of the int_vector
+ */
+std::pair<sdsl::int_vector<>, std::unique_ptr<bitmap>>
 construct_diff_label_count_vector(const AnnotatedDBG &anno_graph,
                                   const std::vector<Label> &labels_in,
                                   const std::vector<Label> &labels_out,
                                   size_t num_threads) {
+    auto graph = std::dynamic_pointer_cast<const DeBruijnGraph>(
+        anno_graph.get_graph_ptr()
+    );
+
+    size_t width = sdsl::bits::hi(std::max(labels_in.size(), labels_out.size())) + 1;
+    sdsl::bit_vector indicator(graph->max_index() + 1, false);
+
     // at this stage, the width of counts is twice what it should be, since
     // the intention is to store the in label and out label counts interleaved
-    // in the beginning, it's the correct size, but double width
-    size_t width = sdsl::bits::hi(std::max(labels_in.size(), labels_out.size())) + 1;
-    auto counts = aligned_int_vector(anno_graph.get_graph().max_index() + 1, 0, width << 1);
+    sdsl::int_vector<> counts = aligned_int_vector(graph->max_index() + 1, 0, width * 2, 16);
+
+    const auto &label_encoder = anno_graph.get_annotation().get_label_encoder();
+    const auto &binmat = anno_graph.get_annotation().get_matrix();
+
+    std::vector<uint64_t> label_in_codes(labels_in.size());
+    std::vector<uint64_t> label_out_codes(labels_out.size());
+    for (size_t i = 0; i < labels_in.size(); ++i) {
+        label_in_codes[i] = label_encoder.encode(labels_in[i]);
+    }
+    for (size_t i = 0; i < labels_out.size(); ++i) {
+        label_out_codes[i] = label_encoder.encode(labels_out[i]);
+    }
 
-    std::mutex backup_mutex;
     constexpr std::memory_order memorder = std::memory_order_relaxed;
 
-    std::atomic_thread_fence(std::memory_order_release);
-    #pragma omp parallel for num_threads(num_threads) schedule(dynamic)
-    for (size_t j = 0; j < labels_in.size(); ++j) {
-        const std::string &label_in = labels_in[j];
-        anno_graph.call_annotated_nodes(label_in, [&](DeBruijnGraph::node_index i) {
-            atomic_fetch_and_add(counts, i, 1, backup_mutex, memorder);
+    std::mutex vector_backup_mutex;
+    #pragma omp parallel num_threads(num_threads)
+    #pragma omp single
+    {
+        std::atomic_thread_fence(std::memory_order_release);
+        bool async = num_threads > 1;
+        binmat.slice_columns(label_in_codes, [&](auto, const bitmap &rows) {
+            rows.call_ones([&](auto r) {
+                node_index i = AnnotatedDBG::anno_to_graph_index(r);
+                set_bit(indicator.data(), i, async, memorder);
+                atomic_fetch_and_add(counts, i, 1, vector_backup_mutex, memorder);
+            });
         });
-    }
-    std::atomic_thread_fence(std::memory_order_acquire);
 
-    // correct the width of counts, making it single-width
-    counts.width(width);
+        #pragma omp taskwait
+        std::atomic_thread_fence(std::memory_order_acquire);
 
-    std::atomic_thread_fence(std::memory_order_release);
-    #pragma omp parallel for num_threads(num_threads) schedule(dynamic)
-    for (size_t j = 0; j < labels_out.size(); ++j) {
-        const std::string &label_out = labels_out[j];
-        anno_graph.call_annotated_nodes(label_out, [&](DeBruijnGraph::node_index i) {
-            atomic_fetch_and_add(counts, (i << 1) + 1, 1, backup_mutex, memorder);
+        // correct the width of counts, making it single-width
+        counts.width(width);
+        std::atomic_thread_fence(std::memory_order_release);
+
+        binmat.slice_columns(label_out_codes, [&](auto, const bitmap &rows) {
+            rows.call_ones([&](auto r) {
+                node_index i = AnnotatedDBG::anno_to_graph_index(r);
+                set_bit(indicator.data(), i, async, memorder);
+                atomic_fetch_and_add(counts, i * 2 + 1, 1, vector_backup_mutex, memorder);
+            });
         });
+
+        #pragma omp taskwait
+        std::atomic_thread_fence(std::memory_order_acquire);
     }
-    std::atomic_thread_fence(std::memory_order_acquire);
+
+    std::unique_ptr<bitmap> union_mask = std::make_unique<bitmap_vector>(
+        std::move(indicator)
+    );
+
+#if 0
+    // MAKE_BOSS is set to true, so for now, comment this out
+    if (!MAKE_BOSS && graph->is_canonical_mode()) {
+        logger->trace("Adding reverse complements");
+
+        MaskedDeBruijnGraph masked_graph(graph, std::move(union_mask));
+
+        std::mutex count_mutex;
+        std::atomic_thread_fence(std::memory_order_release);
+        masked_graph.update_mask([&](const auto &callback) {
+            masked_graph.call_sequences([&](const std::string &seq, const auto &path) {
+                auto it = path.rbegin();
+                auto rev = seq;
+                reverse_complement(rev.begin(), rev.end());
+                graph->map_to_nodes_sequentially(rev, [&](node_index i) {
+                    if (i != *it) {
+                        assert(i != DeBruijnGraph::npos);
+                        assert(it != path.rend());
+                        callback(i, true);
+                        atomic_fetch_and_add(counts, i * 2,
+                                             atomic_fetch(counts, *it * 2, count_mutex),
+                                             count_mutex, memorder);
+                        atomic_fetch_and_add(counts, i * 2 + 1,
+                                             atomic_fetch(counts, *it * 2 + 1, count_mutex),
+                                             count_mutex, memorder);
+                    }
+                    ++it;
+                });
+                assert(it == path.rend());
+            }, num_threads);
+        }, true, num_threads > 1, memorder);
+        std::atomic_thread_fence(std::memory_order_acquire);
+
+        union_mask = std::unique_ptr<bitmap>(masked_graph.release_mask());
+    }
+#endif
 
     // set the width to be double again
-    counts.width(width << 1);
+    counts.width(width * 2);
 
-    return counts;
+    return std::make_pair(std::move(counts), std::move(union_mask));
 }
 
-std::function<uint64_t(SequenceGraph::node_index)>
-build_label_counter(const AnnotatedDBG &anno_graph,
-                    const std::vector<Label> &labels_to_check) {
-    return [&anno_graph, labels_to_check](SequenceGraph::node_index i) {
-        assert(i != SequenceGraph::npos);
-        return std::count_if(labels_to_check.begin(), labels_to_check.end(),
-            [&](const auto &label) { return anno_graph.has_label(i, label); }
-        );
-    };
-}
 
-std::unique_ptr<bitmap>
-mask_nodes_by_node_label(const AnnotatedDBG &anno_graph,
-                         const std::vector<Label> &labels_in,
-                         const std::vector<Label> &labels_out,
-                         const std::function<bool(DeBruijnGraph::node_index,
-                                                  const LabelCountCallback & /* get_num_labels_in */,
-                                                  const LabelCountCallback & /* get_num_labels_out */)> &is_node_in_mask,
-                         size_t num_threads,
-                         double min_frequency_for_frequent_label) {
-    if (!anno_graph.get_graph().num_nodes())
-        return std::make_unique<bitmap_lazy>([](uint64_t) { return false; },
-                                             anno_graph.get_graph().max_index() + 1,
-                                             0);
-
-    const auto *columns = dynamic_cast<const annot::ColumnCompressed<>*>(
-        &anno_graph.get_annotation()
-    );
+template <class GetKeptIntervals>
+void update_masked_graph_by_unitig(MaskedDeBruijnGraph &masked_graph,
+                                   const GetKeptIntervals &get_kept_intervals,
+                                   size_t num_threads) {
+    std::atomic<uint64_t> kept_unitigs(0);
+    std::atomic<uint64_t> total_unitigs(0);
+    std::atomic<uint64_t> num_kept_nodes(0);
+    bool parallel = num_threads > 1;
+    constexpr std::memory_order memorder = std::memory_order_relaxed;
 
-    //TODO: why do we have num_nodes + 1 here and not just num_nodes?
-    if (columns) {
-        size_t frequent_column_label_min_count
-            = (anno_graph.get_graph().num_nodes() + 1)
-                * min_frequency_for_frequent_label;
-
-        // Partition labels into frequent and infrequent sets
-        std::vector<Label> labels_in_infrequent,
-                           labels_in_frequent,
-                           labels_out_infrequent,
-                           labels_out_frequent;
-
-        for (const auto &label_in : labels_in) {
-            if (columns->get_column(label_in).num_set_bits()
-                    >= frequent_column_label_min_count) {
-                labels_in_frequent.push_back(label_in);
-            } else {
-                labels_in_infrequent.push_back(label_in);
+    std::atomic_thread_fence(std::memory_order_release);
+
+    // TODO: the hack below relies on the fact that the implementation of call_unitigs
+    // on masked DBGSuccinct makes a copy of the underlying bit vector, so we can
+    // safely modify the mask during the call. Find a way to avoid this.
+    auto &mask = const_cast<sdsl::bit_vector&>(
+        dynamic_cast<bitmap_vector*>(masked_graph.get_mask_ptr().get())->data()
+    );
+    masked_graph.call_unitigs([&](const std::string &unitig, const auto &path) {
+        total_unitigs.fetch_add(1, memorder);
+
+        size_t last = 0;
+        for (const auto &[begin, end] : get_kept_intervals(unitig, path)) {
+            kept_unitigs.fetch_add(1, memorder);
+            num_kept_nodes.fetch_add(end - begin, memorder);
+            for ( ; last < begin; ++last) {
+                unset_bit(mask.data(), path[last], parallel, memorder);
             }
+            last = end;
         }
 
-        for (const auto &label_out : labels_out) {
-            if (columns->get_column(label_out).num_set_bits()
-                    >= frequent_column_label_min_count) {
-                labels_out_frequent.push_back(label_out);
-            } else {
-                labels_out_infrequent.push_back(label_out);
-            }
+        for ( ; last < path.size(); ++last) {
+            unset_bit(mask.data(), path[last], parallel, memorder);
         }
 
-        // If at least one infrequent label exists, construct a count vector to
-        // reduce calls to the annotator
-        if (labels_in_infrequent.size() || labels_out_infrequent.size()) {
-            sdsl::int_vector<> counts = construct_diff_label_count_vector(
-                anno_graph, labels_in_infrequent, labels_out_infrequent, num_threads
-            );
-
-            // the width of counts is double, since it's both in and out counts interleaved
-            uint32_t width = counts.width() >> 1;
-            uint64_t int_mask = (uint64_t(1) << width) - 1;
-
-            // Flatten count vector to bitmap if all labels were infrequent
-            if (labels_in_frequent.empty() && labels_out_frequent.empty()) {
-                sdsl::bit_vector mask(anno_graph.get_graph().max_index() + 1, false);
-
-                call_nonzeros(counts, [&](DeBruijnGraph::node_index i, size_t count) {
-                    if (i != DeBruijnGraph::npos
-                            && is_node_in_mask(i, [&]() { return count & int_mask; },
-                                                  [&]() { return count >> width; }))
-                        mask[i] = true;
-                });
+    }, num_threads);
+    std::atomic_thread_fence(std::memory_order_acquire);
 
-                return std::make_unique<bitmap_vector>(std::move(mask));
-            }
+    logger->trace("Kept {} out of {} unitigs with average length {}",
+                  kept_unitigs, total_unitigs,
+                  static_cast<double>(num_kept_nodes + kept_unitigs * (masked_graph.get_k() - 1))
+                      / kept_unitigs);
+}
 
-            // If at least one of the labels was frequent, construct a lazy bitmap
-            // which references both the Annotator and the count vector
-            auto count_frequent_in_labels = build_label_counter(anno_graph,
-                                                                labels_in_frequent);
-            auto count_frequent_out_labels = build_label_counter(anno_graph,
-                                                                 labels_out_frequent);
-            return std::make_unique<bitmap_lazy>([=](uint64_t i) {
-                auto count = counts[i];
-                return i != DeBruijnGraph::npos
-                    && is_node_in_mask(i,
-                            [&]() { return (count & int_mask) + count_frequent_in_labels(i); },
-                            [&]() { return (count >> width) + count_frequent_out_labels(i); });
-            }, counts.size());
-        }
-    }
+template <class KeepNode>
+void update_masked_graph_by_node(MaskedDeBruijnGraph &masked_graph,
+                                 const KeepNode &keep_node,
+                                 size_t num_threads) {
+    // TODO: parallel
+    std::ignore = num_threads;
+    bool parallel = false;
+    constexpr std::memory_order memorder = std::memory_order_relaxed;
+
+    size_t kept_nodes = 0;
+    size_t total_nodes = masked_graph.num_nodes();
 
-    // If all of the labels were frequent, or if the Annotator is not ColumnCompressed,
-    // construct a lazy bitmap
-    auto count_frequent_in_labels = build_label_counter(anno_graph, labels_in);
-    auto count_frequent_out_labels = build_label_counter(anno_graph, labels_out);
+    auto &mask = const_cast<sdsl::bit_vector&>(
+        dynamic_cast<bitmap_vector*>(masked_graph.get_mask_ptr().get())->data()
+    );
+
+    std::atomic_thread_fence(std::memory_order_release);
+    masked_graph.call_nodes([&](node_index node) {
+        if (keep_node(node)) {
+            ++kept_nodes;
+        } else {
+            unset_bit(mask.data(), node, parallel, memorder);
+        }
+    });
+    std::atomic_thread_fence(std::memory_order_acquire);
 
-    return std::make_unique<bitmap_lazy>([=](DeBruijnGraph::node_index i) {
-        return i != DeBruijnGraph::npos
-            && is_node_in_mask(i, [&]() { return count_frequent_in_labels(i); },
-                                  [&]() { return count_frequent_out_labels(i); });
-    }, anno_graph.get_graph().max_index() + 1);
+    logger->trace("Kept {} out of {} nodes", kept_nodes, total_nodes);
 }
 
 } // namespace graph
diff --git a/metagraph/src/graph/annotated_graph_algorithm.hpp b/metagraph/src/graph/annotated_graph_algorithm.hpp
index 0f517192dc..bbaac91db5 100644
--- a/metagraph/src/graph/annotated_graph_algorithm.hpp
+++ b/metagraph/src/graph/annotated_graph_algorithm.hpp
@@ -2,71 +2,52 @@
 #define __ANNOTATED_GRAPH_ALGORITHM_HPP__
 
 #include <vector>
-#include <functional>
 
-#include "common/threads/threading.hpp"
-#include "common/vectors/bitmap.hpp"
 #include "graph/annotated_dbg.hpp"
-#include "graph/representation/masked_graph.hpp"
 
 
 namespace mtg {
 namespace graph {
 
-typedef std::function<size_t()> LabelCountCallback;
 
-template <typename NodeType>
-class Alignment;
-
-typedef std::function<bool(const std::string&,
-                           const std::vector<DeBruijnGraph::node_index>&)> KeepUnitigPath;
-
-// Given a DeBruijnGraph and a bool-returning string callback, return a bitmap of
-// length graph.max_index() + 1. An index is set to 1 if it is contained in a
-// unitig satisfying keep_unitig(unitig).
-std::unique_ptr<bitmap_vector>
-mask_nodes_by_unitig(const DeBruijnGraph &graph,
-                     const KeepUnitigPath &keep_unitig,
-                     size_t num_threads = 1);
-
-// Given an AnnotatedDBG and sets of foreground (in) and background (out) labels,
-// return a bitmap of length anno_graph.get_graph().max_index() + 1. An index i
-// is set to 1 if there is a unitig containing i such that
-// at least (label_mask_in_fraction * 100)% of the total possible number of in labels is present,
-// at most (label_mask_out_fraction * 100)% of the total possible number of out labels is present,
-// and (label_other_fraction * 100)% of the labels are neither in or out masked.
-std::unique_ptr<bitmap_vector>
-mask_nodes_by_unitig_labels(const AnnotatedDBG &anno_graph,
-                            const std::vector<AnnotatedDBG::Annotator::Label> &labels_in,
-                            const std::vector<AnnotatedDBG::Annotator::Label> &labels_out,
-                            size_t num_threads = 1,
-                            double label_mask_in_fraction = 1.0,
-                            double label_mask_out_fraction = 0.0,
-                            double label_other_fraction = 1.0);
-
-// Given an AnnotatedDBG and vectors of foreground (labels_in) and
-// background (labels_out) labels, construct a bitmap of length
-// anno_graph.get_graph().max_index() + 1. An index i is set to 1 if
-//
-// is_node_in_mask(i,
-//                 get # of labels of node i in labels_in,
-//                 get # of labels of node i in labels_out)
-//
-// where the getters for the numbers of overlapping labels are callbacks of
-// type LabelCountCallback.
-//
-// For labels which appear in more than
-// num_nodes() * lazy_evaluation_label_frequency_cutoff
-// nodes, precompute the numbers of overlapping labels.
-std::unique_ptr<bitmap>
-mask_nodes_by_node_label(const AnnotatedDBG &anno_graph,
-                         const std::vector<AnnotatedDBG::Annotator::Label> &labels_in,
-                         const std::vector<AnnotatedDBG::Annotator::Label> &labels_out,
-                         const std::function<bool(DeBruijnGraph::node_index,
-                                                  const LabelCountCallback & /* get_num_labels_in */,
-                                                  const LabelCountCallback & /* get_num_labels_out */)> &is_node_in_mask,
-                         size_t num_threads = 1,
-                         double min_frequency_for_frequent_label = 0.05);
+class MaskedDeBruijnGraph;
+
+/**
+ * A container for differential assembly parameters. These are:
+ * label_mask_in_kmer_fraction: minimum fraction of in-labels which should be
+ *                              present for a node to be kept
+ * label_mask_in_unitig_fraction: minimum fraction of nodes within a unitig which
+ *                                must satisfy the in-label criteria
+ * label_mask_out_kmer_fraction: maximum fraction of out-labels which can be
+ *                               present in a kept node
+ * label_mask_out_unitig_fraction: maximum fraction of nodes within a unitig which
+ *                                 can satisfy the out-label criteria
+ * label_mask_other_unitig_fraction: maximum fraction of nodes within a unitig which
+ *                                   can contain labels not in the in- and out-groups
+ * add_complement: also consider the reverse complements of nodes
+ */
+struct DifferentialAssemblyConfig {
+    double label_mask_in_unitig_fraction = 1.0;
+    double label_mask_in_kmer_fraction = 1.0;
+    double label_mask_out_unitig_fraction = 0.0;
+    double label_mask_out_kmer_fraction = 0.0;
+    double label_mask_other_unitig_fraction = 1.0;
+    bool add_complement = false;
+};
+
+/**
+ * Given an AnnotatedDBG and sets of foreground (in) and background (out) labels,
+ * return a MaskedDeBruijnGraph with the nodes of anno_graph masked according to
+ * the parameters specified by config.
+ */
+MaskedDeBruijnGraph
+mask_nodes_by_label(const AnnotatedDBG &anno_graph,
+                    const std::vector<typename AnnotatedDBG::Annotator::Label> &labels_in,
+                    const std::vector<typename AnnotatedDBG::Annotator::Label> &labels_out,
+                    const std::vector<typename AnnotatedDBG::Annotator::Label> &labels_in_post,
+                    const std::vector<typename AnnotatedDBG::Annotator::Label> &labels_out_post,
+                    const DifferentialAssemblyConfig &config,
+                    size_t num_threads = 1);
 
 } // namespace graph
 } // namespace mtg
diff --git a/metagraph/src/graph/representation/masked_graph.hpp b/metagraph/src/graph/representation/masked_graph.hpp
index 2024b1bd9f..0a62c30598 100644
--- a/metagraph/src/graph/representation/masked_graph.hpp
+++ b/metagraph/src/graph/representation/masked_graph.hpp
@@ -119,6 +119,7 @@ class MaskedDeBruijnGraph : public DeBruijnGraph {
     virtual void set_mask(bitmap *mask) { kmers_in_graph_.reset(mask); }
 
     virtual const bitmap& get_mask() const { return *kmers_in_graph_; }
+    virtual std::unique_ptr<bitmap>& get_mask_ptr() { return kmers_in_graph_; }
 
   private:
     std::shared_ptr<const DeBruijnGraph> graph_;
diff --git a/metagraph/tests/annotation/test_annotated_dbg_helpers.cpp b/metagraph/tests/annotation/test_annotated_dbg_helpers.cpp
index 7412d4f39d..645a3067b3 100644
--- a/metagraph/tests/annotation/test_annotated_dbg_helpers.cpp
+++ b/metagraph/tests/annotation/test_annotated_dbg_helpers.cpp
@@ -4,6 +4,7 @@
 
 #include "annotation/annotation_converters.hpp"
 #include "graph/annotated_graph_algorithm.hpp"
+#include "graph/representation/canonical_dbg.hpp"
 
 
 namespace mtg {
@@ -12,14 +13,16 @@ namespace test {
 using namespace mtg::graph;
 
 template <class Graph, class Annotation>
-std::unique_ptr<AnnotatedDBG>
-build_anno_graph(uint64_t k,
-                 const std::vector<std::string> &sequences,
-                 const std::vector<std::string> &labels) {
+std::unique_ptr<AnnotatedDBG> build_anno_graph(uint64_t k,
+                                               const std::vector<std::string> &sequences,
+                                               const std::vector<std::string> &labels,
+                                               DBGMode mode) {
     assert(sequences.size() == labels.size());
-    auto graph = build_graph_batch<Graph>(k, sequences);
+    auto graph = build_graph_batch<Graph>(k, sequences, mode);
+    auto canonical = std::dynamic_pointer_cast<CanonicalDBG>(graph);
+    assert(!canonical || mode == DBGMode::CANONICAL_WRAPPER);
 
-    uint64_t max_index = graph->max_index();
+    uint64_t max_index = canonical ? canonical->get_graph().max_index() : graph->max_index();
 
     auto anno_graph = std::make_unique<AnnotatedDBG>(
         graph,
@@ -45,17 +48,17 @@ build_anno_graph(uint64_t k,
     return anno_graph;
 }
 
-template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGSuccinct, annot::ColumnCompressed<>>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&);
-template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGBitmap, annot::ColumnCompressed<>>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&);
-template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGHashOrdered, annot::ColumnCompressed<>>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&);
-template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGHashFast, annot::ColumnCompressed<>>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&);
-template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGHashString, annot::ColumnCompressed<>>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&);
-
-template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGSuccinct, annot::RowFlatAnnotator>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&);
-template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGBitmap, annot::RowFlatAnnotator>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&);
-template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGHashOrdered, annot::RowFlatAnnotator>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&);
-template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGHashFast, annot::RowFlatAnnotator>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&);
-template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGHashString, annot::RowFlatAnnotator>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&);
+template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGSuccinct, annot::ColumnCompressed<>>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&, DBGMode);
+template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGBitmap, annot::ColumnCompressed<>>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&, DBGMode);
+template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGHashOrdered, annot::ColumnCompressed<>>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&, DBGMode);
+template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGHashFast, annot::ColumnCompressed<>>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&, DBGMode);
+template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGHashString, annot::ColumnCompressed<>>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&, DBGMode);
+
+template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGSuccinct, annot::RowFlatAnnotator>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&, DBGMode);
+template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGBitmap, annot::RowFlatAnnotator>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&, DBGMode);
+template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGHashOrdered, annot::RowFlatAnnotator>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&, DBGMode);
+template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGHashFast, annot::RowFlatAnnotator>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&, DBGMode);
+template std::unique_ptr<AnnotatedDBG> build_anno_graph<DBGHashString, annot::RowFlatAnnotator>(uint64_t, const std::vector<std::string> &, const std::vector<std::string>&, DBGMode);
 
 
 } // namespace test
diff --git a/metagraph/tests/annotation/test_annotated_dbg_helpers.hpp b/metagraph/tests/annotation/test_annotated_dbg_helpers.hpp
index 6e497b1655..7fd1259cb7 100644
--- a/metagraph/tests/annotation/test_annotated_dbg_helpers.hpp
+++ b/metagraph/tests/annotation/test_annotated_dbg_helpers.hpp
@@ -20,6 +20,8 @@
 #include "graph/representation/bitmap/dbg_bitmap.hpp"
 #include "graph/representation/masked_graph.hpp"
 
+#include "../graph/all/test_dbg_helpers.hpp"
+
 
 namespace mtg {
 namespace test {
@@ -29,7 +31,8 @@ using namespace mtg::graph;
 template <class Graph, class Annotation = annot::ColumnCompressed<>>
 std::unique_ptr<AnnotatedDBG> build_anno_graph(uint64_t k,
                                                const std::vector<std::string> &sequences,
-                                               const std::vector<std::string> &labels);
+                                               const std::vector<std::string> &labels,
+                                               DBGMode mode = NORMAL);
 
 typedef ::testing::Types<DBGBitmap,
                          DBGHashString,
@@ -66,6 +69,16 @@ typedef ::testing::Types<std::pair<DBGHashString, annot::ColumnCompressed<>>,
                          std::pair<DBGSuccinct, annot::RowFlatAnnotator>
                         > GraphWithNAnnotationPairTypes;
 
+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<DBGBitmap, annot::RowFlatAnnotator>,
+                         std::pair<DBGHashOrdered, annot::RowFlatAnnotator>,
+                         std::pair<DBGHashFast, annot::RowFlatAnnotator>,
+                         std::pair<DBGSuccinct, annot::RowFlatAnnotator>
+                        > GraphAnnotationCanonicalPairTypes;
+
 } // namespace test
 } // namespace mtg
 
diff --git a/metagraph/tests/annotation/test_annotated_graph_algorithm.cpp b/metagraph/tests/annotation/test_annotated_graph_algorithm.cpp
index e62c446fcb..b6ba470e38 100644
--- a/metagraph/tests/annotation/test_annotated_graph_algorithm.cpp
+++ b/metagraph/tests/annotation/test_annotated_graph_algorithm.cpp
@@ -14,194 +14,265 @@ using namespace mtg::graph;
 using namespace mtg::test;
 
 
-MaskedDeBruijnGraph build_masked_graph(const AnnotatedDBG &anno_graph,
-                                       const std::vector<std::string> &ingroup,
-                                       const std::vector<std::string> &outgroup,
-                                       double mask_in_label_fraction,
-                                       double mask_out_label_fraction,
-                                       double other_label_fraction,
-                                       double lazy_evaluation_density_cutoff,
-                                       size_t num_threads) {
-    size_t insize = ingroup.size();
-    size_t outsize = outgroup.size();
-    return MaskedDeBruijnGraph(
-        std::dynamic_pointer_cast<const DeBruijnGraph>(anno_graph.get_graph_ptr()),
-        graph::mask_nodes_by_node_label(
-            anno_graph,
-            ingroup,
-            outgroup,
-            [=,&anno_graph](auto index, auto get_num_in_labels, auto get_num_out_labels) {
-                assert(index != DeBruijnGraph::npos);
-
-                size_t num_in_labels = get_num_in_labels();
-                if (num_in_labels < mask_in_label_fraction * insize)
-                    return false;
-
-                size_t num_out_labels = get_num_out_labels();
-                if (num_out_labels < mask_out_label_fraction * outsize)
-                    return false;
-
-                size_t num_total_labels = anno_graph.get_labels(index).size();
-
-                return (num_total_labels - num_in_labels - num_out_labels)
-                    <= other_label_fraction * num_total_labels;
-            },
-            num_threads,
-            lazy_evaluation_density_cutoff
-        )
-    );
-}
-
-
 template <typename GraphAnnotationPair>
 class MaskedDeBruijnGraphAlgorithm : public ::testing::Test {};
-TYPED_TEST_SUITE(MaskedDeBruijnGraphAlgorithm, GraphAnnotationPairTypes);
 
-template <class Graph, class Annotation = annot::ColumnCompressed<>>
-void test_mask_indices(double density_cutoff, size_t num_threads) {
-    const std::vector<std::string> ingroup { "B", "C" };
-    const std::vector<std::string> outgroup { "A" };
-
-    for (size_t k = 3; k < 15; ++k) {
-        const std::vector<std::string> sequences {
-            std::string("T") + std::string(k - 1, 'A') + std::string(100, 'T'),
-            std::string("T") + std::string(k - 1, 'A') + "C",
-            std::string("T") + std::string(k - 1, 'A') + "C",
-            std::string("T") + std::string(k - 1, 'A') + "A",
-            std::string("T") + std::string(k - 1, 'A') + "G"
-        };
-        const std::vector<std::string> labels { "A", "B", "C", "D", "E" };
-
-        auto anno_graph = build_anno_graph<Graph, Annotation>(k, sequences, labels);
-
-        std::unordered_set<std::string> obs_labels, obs_kmers;
-        const std::unordered_set<std::string> ref_kmers {
-            std::string(k - 1, 'A') + "C"
-        };
-        const std::unordered_set<std::string> ref_labels {
-            "B", "C"
-        };
-
-        auto masked_dbg = build_masked_graph(*anno_graph,
-                                             ingroup, outgroup,
-                                             1.0, 0.0, 0.0,
-                                             density_cutoff,
-                                             num_threads);
-
-        // FYI: num_nodes() throws exception for masked graph with lazy node mask
-        // EXPECT_EQ(anno_graph->get_graph().num_nodes(), masked_dbg.num_nodes());
-        ASSERT_EQ(anno_graph->get_graph().max_index(), masked_dbg.max_index());
-
-        masked_dbg.call_kmers([&](auto i, const auto &kmer) {
-            auto cur_labels = anno_graph->get_labels(i);
-            obs_labels.insert(cur_labels.begin(), cur_labels.end());
-            obs_kmers.insert(kmer);
-        });
-
-        EXPECT_EQ(ref_labels, obs_labels) << k << " " << density_cutoff;
-        EXPECT_EQ(ref_kmers, obs_kmers) << k << " " << density_cutoff;
-    }
-}
+TYPED_TEST_SUITE(MaskedDeBruijnGraphAlgorithm, GraphAnnotationCanonicalPairTypes);
 
 TYPED_TEST(MaskedDeBruijnGraphAlgorithm, MaskIndicesByLabel) {
-    for (size_t num_threads : { 1, 4 }) {
-        for (double d = 0.0; d <= 1.0; d += 0.05) {
-            test_mask_indices<typename TypeParam::first_type,
-                              typename TypeParam::second_type>(d, num_threads);
+    typedef typename TypeParam::first_type Graph;
+    typedef typename TypeParam::second_type Annotation;
+    for (size_t num_threads = 1; num_threads < 5; num_threads += 3) {
+        const std::vector<std::string> ingroup { "B", "C" };
+        const std::vector<std::string> outgroup { "A" };
+
+        for (size_t k = 3; k < 15; ++k) {
+            const std::vector<std::string> sequences {
+                std::string("T") + std::string(k - 1, 'A') + std::string(100, 'T'),
+                std::string("T") + std::string(k - 1, 'A') + "C",
+                std::string("T") + std::string(k - 1, 'A') + "C",
+                std::string("T") + std::string(k - 1, 'A') + "A",
+                std::string("T") + std::string(k - 1, 'A') + "G"
+            };
+            const std::vector<std::string> labels { "A", "B", "C", "D", "E" };
+
+            auto anno_graph = build_anno_graph<Graph, Annotation>(k, sequences, labels);
+
+            std::unordered_set<std::string> obs_labels, obs_kmers;
+            const std::unordered_set<std::string> ref_kmers {
+                std::string(k - 1, 'A') + "C"
+            };
+            const std::unordered_set<std::string> ref_labels {
+                "B", "C"
+            };
+
+            DifferentialAssemblyConfig config {
+                .label_mask_in_unitig_fraction = 0.0,
+                .label_mask_in_kmer_fraction = 1.0,
+                .label_mask_out_unitig_fraction = 1.0,
+                .label_mask_out_kmer_fraction = 0.0,
+                .label_mask_other_unitig_fraction = 1.0
+            };
+
+            auto masked_dbg = mask_nodes_by_label(*anno_graph,
+                                                  ingroup, outgroup,
+                                                  {}, {},
+                                                  config, num_threads);
+
+            masked_dbg.call_kmers([&](auto, const std::string &kmer) {
+                auto cur_labels = anno_graph->get_labels(kmer, 0.0);
+                obs_labels.insert(cur_labels.begin(), cur_labels.end());
+                obs_kmers.insert(kmer);
+            });
+
+            EXPECT_EQ(ref_labels, obs_labels) << k;
+            EXPECT_EQ(ref_kmers, obs_kmers) << k;
         }
     }
 }
 
-
 template <class Graph, class Annotation = annot::ColumnCompressed<>>
-void
-test_mask_unitigs(double inlabel_fraction,
-                  double outlabel_fraction,
-                  double other_label_fraction,
-                  const std::unordered_set<std::string> &ref_kmers,
-                  size_t num_threads) {
-    const std::vector<std::string> ingroup { "B", "C" };
-    const std::vector<std::string> outgroup { "A" };
-    size_t k = 3;
-
-    {
-        /*
-           CGA                 GCC-CCT
-              \               /
-               GAA-AAT-ATG-TGC
-              /               \
-           GGA                 GCA-CAC
-        */
-        const std::vector<std::string> sequences {
-            "TGCCT",
-            "CGAATGCCT",
-            "GGAATGCAC",
-            "TTTTTTTTTTTTTT"
-        };
-        const std::vector<std::string> labels { "A", "B", "C", "D" };
-
-        auto anno_graph = build_anno_graph<Graph, Annotation>(k, sequences, labels);
-
-        std::unordered_set<std::string> obs_kmers;
-
-        MaskedDeBruijnGraph masked_dbg(
-            std::dynamic_pointer_cast<const DeBruijnGraph>(anno_graph->get_graph_ptr()),
-            mask_nodes_by_unitig_labels(
-                *anno_graph,
-                ingroup,
-                outgroup,
-                num_threads,
-                inlabel_fraction,
-                outlabel_fraction,
-                other_label_fraction
-            )
-        );
-
-        EXPECT_EQ(anno_graph->get_graph().max_index(), masked_dbg.max_index());
-
-        masked_dbg.call_kmers([&](auto, const auto &kmer) { obs_kmers.insert(kmer); });
-
-        EXPECT_EQ(ref_kmers, obs_kmers)
-            << k << " "
-            << inlabel_fraction << " "
-            << outlabel_fraction << " "
-            << other_label_fraction;
+void test_mask_unitigs(double inlabel_fraction,
+                       double outlabel_fraction,
+                       double other_label_fraction,
+                       const std::unordered_set<std::string> &ref_kmers) {
+    for (size_t num_threads = 1; num_threads < 5; num_threads += 3) {
+        const std::vector<std::string> ingroup { "B", "C" };
+        const std::vector<std::string> outgroup { "A" };
+        size_t k = 3;
+
+        {
+            /*  B                  AB  AB
+               CGA                 GCC-CCT
+                  \ BC  BC  BC ABC/
+                   GAA-AAT-ATG-TGC
+                 C/               \  C   C
+               GGA                 GCA-CAC
+            */
+            const std::vector<std::string> sequences {
+                    "TGCCT",
+                "CGAATGCCT",
+                "GGAATGCAC",
+                "TTTTTTTTTTTTTT"
+            };
+            const std::vector<std::string> labels { "A", "B", "C", "D" };
+
+            auto anno_graph = build_anno_graph<Graph, Annotation>(k, sequences, labels);
+
+            std::unordered_set<std::string> obs_kmers;
+
+            DifferentialAssemblyConfig config {
+                .label_mask_in_unitig_fraction = inlabel_fraction,
+                .label_mask_in_kmer_fraction = 1.0,
+                .label_mask_out_unitig_fraction = outlabel_fraction,
+                .label_mask_out_kmer_fraction = 0.0,
+                .label_mask_other_unitig_fraction = other_label_fraction
+            };
+
+            auto masked_dbg = mask_nodes_by_label(*anno_graph,
+                                                  ingroup, outgroup,
+                                                  {}, {},
+                                                  config, num_threads);
+
+            masked_dbg.call_kmers([&](auto, const std::string &kmer) { obs_kmers.insert(kmer); });
+
+            EXPECT_EQ(ref_kmers, obs_kmers)
+                << k << " "
+                << inlabel_fraction << " "
+                << outlabel_fraction << " "
+                << other_label_fraction;
+        }
     }
 }
 
 TYPED_TEST(MaskedDeBruijnGraphAlgorithm, MaskUnitigsByLabel) {
-    for (size_t num_threads : { 1, 4 }) {
+    for (double other_frac : std::vector<double>{ 0.0, 1.0 }) {
         std::unordered_set<std::string> ref_kmers;
+        test_mask_unitigs<typename TypeParam::first_type,
+                          typename TypeParam::second_type>(1.0, 0.0, other_frac, ref_kmers);
 
-        test_mask_unitigs<typename TypeParam::first_type, typename TypeParam::second_type>(
-            1.0, 0.0, 0.0, ref_kmers, num_threads
-        );
-
-        test_mask_unitigs<typename TypeParam::first_type, typename TypeParam::second_type>(
-            1.0, 0.24, 0.0, ref_kmers, num_threads
-        );
+        test_mask_unitigs<typename TypeParam::first_type,
+                          typename TypeParam::second_type>(1.0, 0.24, other_frac, ref_kmers);
 
         ref_kmers.insert("GAA");
         ref_kmers.insert("AAT");
         ref_kmers.insert("ATG");
         ref_kmers.insert("TGC");
 
-        test_mask_unitigs<typename TypeParam::first_type, typename TypeParam::second_type>(
-            1.0, 0.25, 0.0, ref_kmers, num_threads
-        );
+        test_mask_unitigs<typename TypeParam::first_type,
+                          typename TypeParam::second_type>(1.0, 0.25, other_frac, ref_kmers);
 
 
-        test_mask_unitigs<typename TypeParam::first_type, typename TypeParam::second_type>(
-            1.0, 0.50, 0.0, ref_kmers, num_threads
-        );
+        test_mask_unitigs<typename TypeParam::first_type,
+                          typename TypeParam::second_type>(1.0, 0.50, other_frac, ref_kmers);
 
-        test_mask_unitigs<typename TypeParam::first_type, typename TypeParam::second_type>(
-            1.0, 0.75, 0.0, ref_kmers, num_threads
-        );
+        test_mask_unitigs<typename TypeParam::first_type,
+                          typename TypeParam::second_type>(1.0, 0.75, other_frac, ref_kmers);
 
-        test_mask_unitigs<typename TypeParam::first_type, typename TypeParam::second_type>(
-            1.0, 1.0, 0.0, ref_kmers, num_threads
-        );
+        test_mask_unitigs<typename TypeParam::first_type,
+                          typename TypeParam::second_type>(1.0, 1.0, other_frac, ref_kmers);
+    }
+}
+
+
+template <class Graph, class Annotation = annot::ColumnCompressed<>>
+void test_mask_unitigs_canonical(double inlabel_fraction,
+                                 double outlabel_fraction,
+                                 double other_label_fraction,
+                                 const std::unordered_set<std::string> &ref_kmers,
+                                 bool add_complement,
+                                 DBGMode mode) {
+    for (size_t num_threads = 1; num_threads < 5; num_threads += 3) {
+        const std::vector<std::string> ingroup { "B", "C" };
+        const std::vector<std::string> outgroup { "A" };
+        size_t k = 5;
+
+        {
+            /*
+                B                AB    AB
+               CGAAT             ATGCC-TGCCT
+                    \ BC   ABC  /
+                     GAATG-AATGC
+                 C  /           \  C     C
+               GGAAT             ATGCA-TGCAC
+                                   C  /  C                B
+                                 GTGCA-TGCAT             ATTCG
+                                            \ABC    BC  /
+                                             GCATT-CATTC
+                                 AB    AB   /           \  C
+                                 AGGCA-GGCAT             ATTCC
+            */
+            const std::vector<std::string> sequences {
+                  "AATGCCT",     // "AGGCATT"
+                "CGAATGCCT",     // "AGGCATTCG"
+                "GGAATGCAC",     // "GTGCATTCC"
+                "TTTTTTTTTTTTTT" // "AAAAAAAAAAAAAA"
+            };
+            const std::vector<std::string> labels { "A", "B", "C", "D" };
+
+            // if add_complement, then the differential
+            auto anno_graph = build_anno_graph<Graph, Annotation>(
+                k, sequences, labels, mode
+            );
+
+            std::unordered_set<std::string> obs_kmers;
+
+            DifferentialAssemblyConfig config {
+                .label_mask_in_unitig_fraction = inlabel_fraction,
+                .label_mask_in_kmer_fraction = 1.0,
+                .label_mask_out_unitig_fraction = outlabel_fraction,
+                .label_mask_out_kmer_fraction = 0.0,
+                .label_mask_other_unitig_fraction = other_label_fraction,
+                .add_complement = add_complement
+            };
+
+            auto masked_dbg = mask_nodes_by_label(*anno_graph,
+                                                  ingroup, outgroup,
+                                                  {}, {},
+                                                  config, num_threads);
+
+            masked_dbg.call_kmers([&](auto, const std::string &kmer) { obs_kmers.insert(kmer); });
+
+            EXPECT_EQ(ref_kmers, obs_kmers)
+                << k << " "
+                << inlabel_fraction << " "
+                << outlabel_fraction << " "
+                << other_label_fraction << " "
+                << add_complement << " "
+                << mode << " "
+                << num_threads;
+        }
+    }
+}
+
+TYPED_TEST(MaskedDeBruijnGraphAlgorithm, MaskUnitigsByLabelCanonical) {
+    std::vector<std::pair<bool, DBGMode>> params {
+        { true, DBGMode::NORMAL },
+        { false, DBGMode::CANONICAL },
+        { true, DBGMode::CANONICAL },
+        { false, DBGMode::CANONICAL_WRAPPER },
+        { true, DBGMode::CANONICAL_WRAPPER },
+        { true, DBGMode::PRIMARY },
+    };
+    for (const auto &[add_complement, mode] : params) {
+        for (double other_frac : std::vector<double>{ 1.0, 0.0 }) {
+            std::unordered_set<std::string> ref_kmers;
+            test_mask_unitigs_canonical<typename TypeParam::first_type,
+                                        typename TypeParam::second_type>(
+                1.0, 0.0, other_frac, ref_kmers, add_complement, mode
+            );
+
+            test_mask_unitigs_canonical<typename TypeParam::first_type,
+                                        typename TypeParam::second_type>(
+                1.0, 0.25, other_frac, ref_kmers, add_complement, mode
+            );
+
+
+            test_mask_unitigs_canonical<typename TypeParam::first_type,
+                                        typename TypeParam::second_type>(
+                1.0, 0.49, other_frac, ref_kmers, add_complement, mode
+            );
+
+            ref_kmers.insert("GAATG");
+            ref_kmers.insert("AATGC");
+            ref_kmers.insert("GCATT");
+            ref_kmers.insert("CATTC");
+
+            test_mask_unitigs_canonical<typename TypeParam::first_type,
+                                        typename TypeParam::second_type>(
+                1.0, 0.50, other_frac, ref_kmers, add_complement, mode
+            );
+
+            test_mask_unitigs_canonical<typename TypeParam::first_type,
+                                        typename TypeParam::second_type>(
+                1.0, 0.75, other_frac, ref_kmers, add_complement, mode
+            );
+
+            test_mask_unitigs_canonical<typename TypeParam::first_type,
+                                        typename TypeParam::second_type>(
+                1.0, 1.0, other_frac, ref_kmers, add_complement, mode
+            );
+        }
     }
 }
 
diff --git a/metagraph/tests/annotation/test_matrix_helpers.cpp b/metagraph/tests/annotation/test_matrix_helpers.cpp
index a2d5830f55..bb06f5d825 100644
--- a/metagraph/tests/annotation/test_matrix_helpers.cpp
+++ b/metagraph/tests/annotation/test_matrix_helpers.cpp
@@ -198,6 +198,49 @@ void test_matrix(const TypeParam &matrix, const BitVectorPtrArray &columns) {
         }
     }
 
+    // check slice_columns
+    for (size_t m : { size_t(0),
+                      size_t(matrix.num_columns() / 2),
+                      size_t(matrix.num_columns()) }) {
+        std::vector<uint64_t> indices(m);
+        std::iota(indices.begin(), indices.end(), 0);
+
+        std::vector<std::vector<BinaryMatrix::Row>> column_map(m);
+        matrix.slice_columns(indices, [&](auto j, bitmap&& rows) {
+            rows.call_ones([&](auto i) { column_map[j].push_back(i); });
+        });
+        std::vector<uint64_t> slice;
+        for (size_t j = 0; j < column_map.size(); ++j) {
+            slice.insert(slice.end(), column_map[j].begin(), column_map[j].end());
+            slice.push_back(std::numeric_limits<BinaryMatrix::Row>::max());
+        }
+
+        ASSERT_GE(slice.size(), indices.size());
+
+        auto column_begin = slice.begin();
+
+        for (size_t j = 0; j < indices.size(); ++j) {
+            // every row in `slice` ends with `-1`
+            auto column_end = std::find(column_begin, slice.end(),
+                                        std::numeric_limits<BinaryMatrix::Row>::max());
+            std::vector<BinaryMatrix::Row> column_set_bits(column_begin, column_end);
+            column_begin = column_end + 1;
+
+            // make sure all returned indexes are unique
+            ASSERT_EQ(column_set_bits.size(), convert_to_set(column_set_bits).size());
+
+            for (auto i : column_set_bits) {
+                ASSERT_TRUE(i < matrix.num_rows());
+                EXPECT_TRUE((*columns[j])[i]);
+            }
+
+            auto set_bits = convert_to_set(column_set_bits);
+            for (size_t i = 0; i < matrix.num_rows(); ++i) {
+                EXPECT_EQ((*columns[j])[i], set_bits.count(i));
+            }
+        }
+    }
+
     // check get_row
     for (size_t i = 0, n_rows = matrix.num_rows(); i < n_rows; ++i) {
         auto row_set_bits = matrix.get_row(i);

From 420437bc93227f67b345f5057ab30a5837346198 Mon Sep 17 00:00:00 2001
From: Harun Mustafa <bilgisayaradami@gmail.com>
Date: Wed, 21 Oct 2020 19:57:08 +0200
Subject: [PATCH 2/2] Improved slice_columns in BRWT and Rainbow

---
 metagraph/benchmarks/benchmark_matrix.cpp     | 94 ++++++++++++++++++
 .../binary_matrix/multi_brwt/brwt.cpp         | 97 +++++++++++++++++++
 .../binary_matrix/multi_brwt/brwt.hpp         |  3 +
 .../binary_matrix/rainbowfish/rainbow.cpp     | 19 ++++
 .../binary_matrix/rainbowfish/rainbow.hpp     |  3 +
 5 files changed, 216 insertions(+)

diff --git a/metagraph/benchmarks/benchmark_matrix.cpp b/metagraph/benchmarks/benchmark_matrix.cpp
index ef43b6ba1f..657dd15a0b 100644
--- a/metagraph/benchmarks/benchmark_matrix.cpp
+++ b/metagraph/benchmarks/benchmark_matrix.cpp
@@ -115,4 +115,98 @@ BENCHMARK_TEMPLATE(BM_BRWTQueryRows, 3000000, 100, 30, 2, false, 0)
     ->Unit(benchmark::kMillisecond)
     ->DenseRange(0, 10, 1);
 
+
+template <size_t rows_arg = 300000,
+          size_t cols_arg = 100,
+          size_t unique_arg = 10,
+          size_t arity_arg = 2,
+          bool greedy_arg = true,
+          size_t relax_arg = 2>
+static void BM_BRWTQueryColumns(benchmark::State& state) {
+    DataGenerator generator;
+    generator.set_seed(42);
+
+    auto density_arg = std::vector<double>(unique_arg, state.range(0) / 100.);
+    auto generated_columns = generator.generate_random_columns(
+        rows_arg,
+        unique_arg,
+        get_densities(unique_arg, density_arg),
+        std::vector<uint32_t>(unique_arg, cols_arg / unique_arg)
+    );
+
+    std::unique_ptr<annot::binmat::BinaryMatrix> matrix = experiments::generate_brwt_from_rows(
+        std::move(generated_columns),
+        arity_arg,
+        greedy_arg,
+        relax_arg
+    );
+
+    std::vector<uint64_t> indexes;
+    call_ones(generator.generate_random_column(matrix->num_columns(), 1. / 10),
+        [&](uint64_t i) { indexes.push_back(i); }
+    );
+
+    for (auto _ : state) {
+        uint64_t j = 0;
+        #pragma omp parallel for num_threads(3)
+        for (size_t i = 0; i < indexes.size(); ++i) {
+            j += i;
+            for (auto t : matrix->get_column(indexes[i])) {
+                j += t;
+            }
+        }
+    }
+}
+
+BENCHMARK_TEMPLATE(BM_BRWTQueryColumns, 3000000, 100, 30, 2, false, 0)
+    ->Unit(benchmark::kMillisecond)
+    ->DenseRange(0, 10, 1);
+
+template <size_t rows_arg = 300000,
+          size_t cols_arg = 100,
+          size_t unique_arg = 10,
+          size_t arity_arg = 2,
+          bool greedy_arg = true,
+          size_t relax_arg = 2>
+static void BM_BRWTSliceColumns(benchmark::State& state) {
+    DataGenerator generator;
+    generator.set_seed(42);
+
+    auto density_arg = std::vector<double>(unique_arg, state.range(0) / 100.);
+    auto generated_columns = generator.generate_random_columns(
+        rows_arg,
+        unique_arg,
+        get_densities(unique_arg, density_arg),
+        std::vector<uint32_t>(unique_arg, cols_arg / unique_arg)
+    );
+
+    std::unique_ptr<annot::binmat::BinaryMatrix> matrix = experiments::generate_brwt_from_rows(
+        std::move(generated_columns),
+        arity_arg,
+        greedy_arg,
+        relax_arg
+    );
+
+    std::vector<uint64_t> indexes;
+    call_ones(generator.generate_random_column(matrix->num_columns(), 1. / 10),
+        [&](uint64_t i) { indexes.push_back(i); }
+    );
+
+    for (auto _ : state) {
+        uint64_t j = 0;
+        #pragma omp parallel num_threads(3)
+        #pragma omp single
+        {
+            matrix->slice_columns(indexes, [&](auto i, auto&& bitmap) {
+                j += i;
+                bitmap.call_ones([&](auto t) { j += t; });
+            });
+        }
+    }
+}
+
+BENCHMARK_TEMPLATE(BM_BRWTSliceColumns, 3000000, 100, 30, 2, false, 0)
+    ->Unit(benchmark::kMillisecond)
+    ->DenseRange(0, 10, 1);
+
 } // namespace
diff --git a/metagraph/src/annotation/binary_matrix/multi_brwt/brwt.cpp b/metagraph/src/annotation/binary_matrix/multi_brwt/brwt.cpp
index f71ae11ed9..9a7a524597 100644
--- a/metagraph/src/annotation/binary_matrix/multi_brwt/brwt.cpp
+++ b/metagraph/src/annotation/binary_matrix/multi_brwt/brwt.cpp
@@ -3,6 +3,10 @@
 #include <queue>
 #include <numeric>
 
+#include <omp.h>
+
+#include <tsl/hopscotch_map.h>
+
 #include "common/algorithms.hpp"
 #include "common/serialization.hpp"
 
@@ -189,6 +193,99 @@ std::vector<BRWT::Column> BRWT::slice_rows(const std::vector<Row> &row_ids) cons
     return slice;
 }
 
+void BRWT::slice_columns(const std::vector<Column> &column_ids,
+                         const ColumnCallback &callback) const {
+    if (column_ids.empty())
+        return;
+
+    auto num_nonzero_rows = nonzero_rows_->num_set_bits();
+
+    // check if the column is empty
+    if (!num_nonzero_rows)
+        return;
+
+    // check whether it is a leaf
+    if (!child_nodes_.size()) {
+        // return the index column
+        for (size_t k = 0; k < column_ids.size(); ++k) {
+            callback(column_ids[k], std::move(*nonzero_rows_->copy()));
+        }
+
+        return;
+    }
+
+    tsl::hopscotch_map<uint32_t, std::vector<Column>> child_columns_map;
+    for (size_t i = 0; i < column_ids.size(); ++i) {
+        assert(column_ids[i] < num_columns());
+        auto child_node = assignments_.group(column_ids[i]);
+        auto child_column = assignments_.rank(column_ids[i]);
+
+        auto it = child_columns_map.find(child_node);
+        if (it == child_columns_map.end())
+            it = child_columns_map.emplace(child_node, std::vector<Column>{}).first;
+
+        it.value().push_back(child_column);
+    }
+
+    auto process = [&](auto child_node, auto *child_columns_ptr) {
+        if (num_nonzero_rows == nonzero_rows_->size()) {
+            child_nodes_[child_node]->slice_columns(*child_columns_ptr,
+                [&](Column j, bitmap&& rows) {
+                    callback(assignments_.get(child_node, j), std::move(rows));
+                }
+            );
+        } else {
+            const BRWT *child_node_brwt = dynamic_cast<const BRWT*>(
+                child_nodes_[child_node].get()
+            );
+            if (child_node_brwt
+                    && child_columns_ptr->size() > 1
+                    && !child_node_brwt->child_nodes_.size()) {
+                // if there are multiple column ids corresponding to the same leaf
+                // node, then this branch avoids doing redundant select1 calls
+                const auto *nonzero_rows = child_node_brwt->nonzero_rows_.get();
+                size_t num_nonzero_rows = nonzero_rows->num_set_bits();
+                if (num_nonzero_rows) {
+                    std::vector<uint64_t> set_bits;
+                    set_bits.reserve(num_nonzero_rows);
+                    nonzero_rows->call_ones([&](auto i) {
+                        set_bits.push_back(nonzero_rows->select1(i + 1));
+                    });
+
+                    for (size_t k = 0; k < child_columns_ptr->size() - 1; ++k) {
+                        callback(assignments_.get(child_node, (*child_columns_ptr)[k]),
+                                 bitmap_generator(std::move(set_bits), num_rows()));
+                    }
+
+                    callback(assignments_.get(child_node, child_columns_ptr->back()),
+                             bitmap_generator(std::move(set_bits), num_rows()));
+                }
+            } else {
+                child_nodes_[child_node]->slice_columns(*child_columns_ptr,
+                    [&](Column j, bitmap&& rows) {
+                        size_t num_set_bits = rows.num_set_bits();
+                        callback(assignments_.get(child_node, j),
+                                 bitmap_generator(std::move(rows), [&](uint64_t i) {
+                                     return nonzero_rows_->select1(i + 1);
+                                 }, num_rows(), num_set_bits));
+                    }
+                );
+            }
+        }
+    };
+
+    for (auto it = ++child_columns_map.begin(); it != child_columns_map.end(); ++it) {
+        auto child_node = it->first;
+        auto *child_columns_ptr = &it->second;
+        #pragma omp task firstprivate(child_node, child_columns_ptr)
+        process(child_node, child_columns_ptr);
+    }
+
+    process(child_columns_map.begin()->first, &child_columns_map.begin()->second);
+
+    #pragma omp taskwait
+}
+
 std::vector<BRWT::Row> BRWT::get_column(Column column) const {
     assert(column < num_columns());
 
diff --git a/metagraph/src/annotation/binary_matrix/multi_brwt/brwt.hpp b/metagraph/src/annotation/binary_matrix/multi_brwt/brwt.hpp
index 829341cf76..641e0a5cf7 100644
--- a/metagraph/src/annotation/binary_matrix/multi_brwt/brwt.hpp
+++ b/metagraph/src/annotation/binary_matrix/multi_brwt/brwt.hpp
@@ -35,6 +35,9 @@ class BRWT : public BinaryMatrix {
     // get all selected rows appended with -1 and concatenated
     std::vector<Column> slice_rows(const std::vector<Row> &rows) const override;
 
+    void slice_columns(const std::vector<Column> &columns,
+                       const ColumnCallback &callback) const override;
+
     bool load(std::istream &in) override;
     void serialize(std::ostream &out) const override;
 
diff --git a/metagraph/src/annotation/binary_matrix/rainbowfish/rainbow.cpp b/metagraph/src/annotation/binary_matrix/rainbowfish/rainbow.cpp
index 20161658fa..8e16517f10 100644
--- a/metagraph/src/annotation/binary_matrix/rainbowfish/rainbow.cpp
+++ b/metagraph/src/annotation/binary_matrix/rainbowfish/rainbow.cpp
@@ -131,6 +131,25 @@ Rainbow<MatrixType>::get_column(Column column) const {
     return row_indices;
 }
 
+template <class MatrixType>
+void
+Rainbow<MatrixType>::slice_columns(const std::vector<Column> &columns,
+                                   const ColumnCallback &callback) const {
+    uint64_t nrows = num_rows();
+    sdsl::bit_vector code_column(reduced_matrix_.num_rows());
+    reduced_matrix_.slice_columns(columns, [&](Column j, bitmap&& rows) {
+        sdsl::util::set_to_value(code_column, false);
+        rows.add_to(&code_column);
+
+        callback(j, bitmap_generator([&](const auto &index_callback) {
+            for (uint64_t i = 0; i < nrows; ++i) {
+                if (code_column[get_code(i)])
+                    index_callback(i);
+            }
+        }, nrows));
+    });
+}
+
 template <class MatrixType>
 bool Rainbow<MatrixType>::load(std::istream &in) {
     try {
diff --git a/metagraph/src/annotation/binary_matrix/rainbowfish/rainbow.hpp b/metagraph/src/annotation/binary_matrix/rainbowfish/rainbow.hpp
index b5ef7c187b..fc7ced91aa 100644
--- a/metagraph/src/annotation/binary_matrix/rainbowfish/rainbow.hpp
+++ b/metagraph/src/annotation/binary_matrix/rainbowfish/rainbow.hpp
@@ -40,6 +40,9 @@ class Rainbow : public RainbowMatrix {
                                           size_t num_threads = 1) const override;
     std::vector<Row> get_column(Column column) const override;
 
+    void slice_columns(const std::vector<Column> &columns,
+                       const ColumnCallback &callback) const override;
+
     bool load(std::istream &in) override;
     void serialize(std::ostream &out) const override;