Randomize barcode downsampling

Author: Itolstoganov

src/common/barcode_index/barcode_index_builder.hpp

@@ -14,6 +14,8 @@
 #include "alignment/sequence_mapper_notifier.hpp"
 #include "alignment/sequence_mapper.hpp"
 
+#include <algorithm>
+#include <cmath>
 #include <random>
 #include <string>
 #include <unordered_set>
@@ -199,31 +201,29 @@ class FrameBarcodeIndexBuilder {
     template<class ReadType>
     void ConstructBarcodeIndex(io::ReadStreamList<ReadType> read_streams,
                                FrameBarcodeIndex<Graph> &barcode_index,
-                               const io::SequencingLibraryBase &lib,
                                bool is_tellseq);
 
-    void DownsampleBarcodeIndex(FrameBarcodeIndex<Graph> &downsampled_index, FrameBarcodeIndex<Graph> &original_index, double sampling_factor) {
+    void DownsampleBarcodeIndex(FrameBarcodeIndex<Graph> &downsampled_index, 
+                                FrameBarcodeIndex<Graph> &original_index, 
+                                double sampling_factor,
+                                int seed) {
         std::unordered_set<BarcodeId> barcodes;
-        std::unordered_set<BarcodeId> passed_barcodes;
-        BarcodeId min_barcode = std::numeric_limits<BarcodeId>::max();
-        BarcodeId max_barcode = std::numeric_limits<BarcodeId>::min();
         for (auto it = original_index.begin(); it != original_index.end(); ++it) {
             const auto &barcode_distribution = it->second.GetDistribution();
             for (const auto &entry: barcode_distribution) {
-                BarcodeId current_barcode = entry.first;
-                barcodes.insert(current_barcode);
-                min_barcode = std::min(min_barcode, current_barcode);
-                max_barcode = std::max(max_barcode, current_barcode);
+                barcodes.insert(entry.first);
             }
         }
         INFO("Number of encountered barcodes: " << barcodes.size());
-        INFO("Barcode id range: " << min_barcode << ", " << max_barcode);
-        double barcode_thr = static_cast<double>(max_barcode - min_barcode) * sampling_factor;
-        for (const auto &barcode: barcodes) {
-            if (math::le(static_cast<double>(barcode - min_barcode), barcode_thr)) {
-                passed_barcodes.insert(barcode);
-            }
-        }
+        size_t target = static_cast<size_t>(std::round(static_cast<double>(barcodes.size()) * sampling_factor));
+        std::unordered_set<BarcodeId> passed_barcodes;
+        passed_barcodes.reserve(target);
+        std::mt19937 rng(seed);
+        std::sample(barcodes.begin(), 
+                    barcodes.end(), 
+                    std::inserter(passed_barcodes, passed_barcodes.end()), 
+                    target, 
+                    rng);
         INFO("Passed barcodes: " << passed_barcodes.size());
 
         downsampled_index.InitialFillMap();
@@ -248,7 +248,6 @@ class FrameBarcodeIndexBuilder {
 template<class ReadType>
 void FrameBarcodeIndexBuilder::ConstructBarcodeIndex(io::ReadStreamList<ReadType> read_streams,
                                                      FrameBarcodeIndex<Graph> &barcode_index,
-                                                     const io::SequencingLibraryBase &lib,
                                                      bool is_tellseq) {
     {
         size_t starting_barcode = 0;

src/projects/splitter/barcode_index_construction.cpp

@@ -35,11 +35,11 @@ void ConstructBarcodeIndex(barcode_index::FrameBarcodeIndex<debruijn_graph::Grap
         FrameBarcodeIndexBuilder barcode_index_builder(graph, mapper, barcode_prefices, frame_size, nthreads);
         bool is_tellseq = lib.type() == io::LibraryType::TellSeqReads;
         if (!is_tellseq) {
-            barcode_index_builder.ConstructBarcodeIndex(io::paired_easy_readers(lib, false, 0), barcode_index, lib, is_tellseq);
+            barcode_index_builder.ConstructBarcodeIndex(io::paired_easy_readers(lib, false, 0), barcode_index, is_tellseq);
         }
-        if (is_tellseq) {
+        else {
             INFO("Constructing from tellseq lib");
-            barcode_index_builder.ConstructBarcodeIndex(io::tellseq_easy_readers(lib, false, 0), barcode_index, lib, is_tellseq);
+            barcode_index_builder.ConstructBarcodeIndex(io::tellseq_easy_readers(lib, false, 0), barcode_index, is_tellseq);
         }
         INFO("Barcode index construction finished.");
 
@@ -69,14 +69,15 @@ void DownsampleBarcodeIndex(const debruijn_graph::Graph &graph,
                             unsigned nthreads,
                             barcode_index::FrameBarcodeIndex<debruijn_graph::Graph> &barcode_index,
                             barcode_index::FrameBarcodeIndex<debruijn_graph::Graph> &downsampled_index,
-                            double sampling_fraction) {
+                            double sampling_fraction,
+                            int seed) {
     VERIFY_MSG(math::ls(sampling_fraction, 1.0), "Sampling fraction must be less than 1");
     const size_t mapping_k = 31;
     const std::vector<std::string> barcode_prefices = {"BC:Z:", "BX:Z:"};
     debruijn_graph::Graph empty_graph(mapping_k);
     alignment::BWAReadMapper<debruijn_graph::Graph> mapper(empty_graph);
     FrameBarcodeIndexBuilder barcode_index_builder(graph, mapper, barcode_prefices, barcode_index.GetFrameSize(), nthreads);
-    barcode_index_builder.DownsampleBarcodeIndex(downsampled_index, barcode_index, sampling_fraction);
+    barcode_index_builder.DownsampleBarcodeIndex(downsampled_index, barcode_index, sampling_fraction, seed);
 }
 
 }

src/projects/splitter/barcode_index_construction.hpp

@@ -36,5 +36,6 @@ void DownsampleBarcodeIndex(const debruijn_graph::Graph &graph,
                             unsigned nthreads,
                             barcode_index::FrameBarcodeIndex<debruijn_graph::Graph> &barcode_index,
                             barcode_index::FrameBarcodeIndex<debruijn_graph::Graph> &downsampled_index,
-                            double sampling_factor);
+                            double sampling_factor,
+                            int seed);
 }

src/projects/splitter/main.cpp

@@ -57,6 +57,7 @@ struct gcfg {
   size_t frame_size = 40000;
   size_t read_linkage_distance = 40000;
   double sampling_factor = 1.0;
+  int seed = 999;
 
   //graph construction
   double graph_score_threshold = 2.0;
@@ -89,6 +90,7 @@ static void process_cmdline(int argc, char** argv, gcfg& cfg) {
         (option("-l") & integer("value", cfg.libindex)) % "library index (0-based, default: 0)",
         (option("-t") & integer("value", cfg.nthreads)) % "# of threads to use",
         (option("--sampling-factor") & value("sampling-factor", cfg.sampling_factor)) % "Sampling factor for read downsampling",
+        (option("--seed") & value("seed", cfg.seed)) % "Seed for barcode downsampling",
         (option("--assembly-info") & value("assembly-info", assembly_info))
             % "Path to metaflye assembly_info.txt file (meta mode, metaFlye graphs only)",
         (with_prefix("-G",
@@ -162,11 +164,11 @@ struct TimeTracerRAII {
 gfa::GFAReader ReadGraph(const gcfg &cfg,
                          debruijn_graph::Graph &graph,
                          io::IdMapper<std::string> *id_mapper) {
-        gfa::GFAReader gfa(cfg.graph);
-        gfa.to_graph(graph, id_mapper);
-        INFO("GFA segments: " << gfa.num_edges() << ", links: " << gfa.num_links() << ", paths: "
-                              << gfa.num_paths());
-        return gfa;
+    gfa::GFAReader gfa(cfg.graph);
+    gfa.to_graph(graph, id_mapper);
+    INFO("GFA segments: " << gfa.num_edges() << ", links: " << gfa.num_links() << ", paths: "
+                            << gfa.num_paths());
+    return gfa;
 }
 
 std::unordered_set<debruijn_graph::EdgeId> ParseRepetitiveEdges(const debruijn_graph::Graph &graph,
@@ -401,7 +403,7 @@ int main(int argc, char** argv) {
     if (not math::eq(cfg.sampling_factor, 1.0)) {
         INFO("Downsampling the barcode index with factor " << cfg.sampling_factor);
         cont_index::DownsampleBarcodeIndex(graph, cfg.nthreads, barcode_index, downsampled_index,
-                                            cfg.sampling_factor);
+                                            cfg.sampling_factor, cfg.seed);
         barcode_extractor_ptr = std::make_shared<BarcodeExtractor>(downsampled_index, graph);
     }