Optimize memory usage during graph construction

cpp/include/seqwin/build.hpp

@@ -11,12 +11,9 @@ namespace seqwin {
 /**
  * @brief Build a minimizer graph from assembly FASTA files.
  *
- * `assembly_paths` and `is_targets` are parallel lists.
- *
  * @param assembly_paths Paths to input assemblies in FASTA format (plain or gzipped).
  * @param kmerlen K-mer length for minimizer sketch.
  * @param windowsize Window size for minimizer sketch.
- * @param is_targets Whether each assembly is a target assembly.
  * @param n_cpu Number of worker threads to use.
  * @param low_memory Recompute minimizers in a second pass to reduce peak memory.
  * @return Minimizer graph.
@@ -26,7 +23,6 @@ Graph build(
     const std::vector<std::string>& assembly_paths,
     std::size_t kmerlen,
     std::size_t windowsize,
-    const std::vector<bool>& is_targets,
     std::size_t n_cpu = 1,
     bool low_memory = false
 );

cpp/include/seqwin/filter.hpp

@@ -8,6 +8,17 @@
 
 namespace seqwin {
 
+void get_penalty(
+    const Kmer* kmers,
+    Node* nodes,
+    std::size_t n_nodes,
+    const std::size_t* record_offsets,
+    std::size_t n_record_offsets,
+    const bool* is_targets,
+    std::size_t n_assemblies,
+    std::size_t n_cpu = 1
+);
+
 Graph filter_kmers(
     const Kmer* kmers,
     const Node* nodes,

cpp/include/seqwin/graph.hpp

@@ -32,11 +32,11 @@ struct Node {
     std::size_t start;
     /** End of the half-open range for this node's minimizer entries. */
     std::size_t stop;
-    /** Number of target assemblies containing this minimizer hash. */
-    std::uint32_t n_tar;
-    /** Number of non-target assemblies containing this minimizer hash. */
-    std::uint32_t n_neg;
-    /** Placeholder for node penalty score, for downstream graph filtering. */
+    /** Node scoring placeholder. */
+    std::uint32_t n_tar = 0;
+    /** Node scoring placeholder. */
+    std::uint32_t n_neg = 0;
+    /** Node scoring placeholder. */
     double penalty = 0.0;
 };
 
@@ -56,7 +56,10 @@ struct Edge {
  * @brief Container for the minimizer graph returned by `build()`.
  */
 struct Graph {
-    /** Minimizer occurrences in all assemblies, grouped and sorted by hash. */
+    /**
+     * Minimizer occurrences in all assemblies, grouped and sorted by hash.
+     * `record_idx` is nondecreasing within each node range.
+     */
     NoInitArray<Kmer> kmers;
     /** Sorted by hash. */
     NoInitArray<Node> nodes;

cpp/src/bindings/python_bindings.cpp

@@ -1,5 +1,6 @@
 #include <cstdint>
 #include <memory>
+#include <stdexcept>
 #include <string>
 #include <type_traits>
 #include <utility>
@@ -50,7 +51,6 @@ PYBIND11_MODULE(_core, m) {
         [](const std::vector<std::string>& assembly_paths,
            std::size_t kmerlen,
            std::size_t windowsize,
-           const std::vector<bool>& is_targets,
            std::size_t n_cpu,
            bool low_memory
         ) {
@@ -61,7 +61,6 @@ PYBIND11_MODULE(_core, m) {
                     assembly_paths,
                     kmerlen,
                     windowsize,
-                    is_targets,
                     n_cpu,
                     low_memory
                 );
@@ -86,11 +85,55 @@ PYBIND11_MODULE(_core, m) {
         py::arg("assembly_paths"),
         py::arg("kmerlen"),
         py::arg("windowsize"),
-        py::arg("is_targets"),
         py::arg("n_cpu") = 1,
         py::arg("low_memory") = false
     );
 
+    m.def("_get_penalty_native",
+        [](py::array_t<seqwin::Kmer, py::array::c_style> kmers,
+           py::array_t<seqwin::Node, py::array::c_style> nodes,
+           py::array_t<std::size_t, py::array::c_style> record_offsets,
+           py::array_t<bool, py::array::c_style> is_targets,
+           std::size_t n_cpu
+        ) {
+            auto kmers_buf = kmers.request();
+            auto nodes_buf = nodes.request();
+            auto record_offsets_buf = record_offsets.request();
+            auto is_targets_buf = is_targets.request();
+
+            if (nodes_buf.readonly) {
+                throw std::invalid_argument("nodes must be writable");
+            }
+
+            const auto* kmers_ptr = static_cast<const seqwin::Kmer*>(kmers_buf.ptr);
+            auto* nodes_ptr = static_cast<seqwin::Node*>(nodes_buf.ptr);
+            const auto* record_offsets_ptr = static_cast<const std::size_t*>(record_offsets_buf.ptr);
+            const auto* is_targets_ptr = static_cast<const bool*>(is_targets_buf.ptr);
+            const auto n_nodes = static_cast<std::size_t>(nodes_buf.shape[0]);
+            const auto n_record_offsets = static_cast<std::size_t>(record_offsets_buf.shape[0]);
+            const auto n_assemblies = static_cast<std::size_t>(is_targets_buf.shape[0]);
+
+            {
+                py::gil_scoped_release release;
+                seqwin::get_penalty(
+                    kmers_ptr,
+                    nodes_ptr,
+                    n_nodes,
+                    record_offsets_ptr,
+                    n_record_offsets,
+                    is_targets_ptr,
+                    n_assemblies,
+                    n_cpu
+                );
+            }
+        },
+        py::arg("kmers").noconvert(),
+        py::arg("nodes").noconvert(),
+        py::arg("record_offsets").noconvert(),
+        py::arg("is_targets").noconvert(),
+        py::arg("n_cpu") = 1
+    );
+
     m.def("_filter_kmers_native",
         [](py::array_t<seqwin::Kmer, py::array::c_style> kmers,
            py::array_t<seqwin::Node, py::array::c_style> nodes,

cpp/src/seqwin/build.cpp

@@ -61,14 +61,6 @@ struct RawKmer {
     Kmer kmer;
 };
 
-struct NodeState {
-    std::size_t count = 0;
-    std::size_t cursor = 0;
-    std::uint32_t n_tar = 0;
-    std::uint32_t n_neg = 0;
-    std::size_t last_seen_assembly = std::numeric_limits<std::size_t>::max();
-};
-
 using EdgeKey = std::pair<std::uint64_t, std::uint64_t>;
 
 struct EdgeKeyHash {
@@ -80,8 +72,8 @@ struct EdgeKeyHash {
 };
 
 struct EdgeState {
-    std::size_t weight = 0;
-    std::size_t last_seen_assembly = std::numeric_limits<std::size_t>::max();
+    std::uint32_t weight = 0;
+    std::uint32_t last_seen_assembly = std::numeric_limits<std::uint32_t>::max();
 };
 
 /**
@@ -95,7 +87,6 @@ ThreadGraph build_worker(
     const std::vector<std::string>& assembly_paths,
     std::size_t kmerlen,
     std::size_t windowsize,
-    const std::vector<bool>& is_targets,
     std::size_t start_assembly,
     std::size_t end_assembly,
     std::size_t thread_id,
@@ -121,15 +112,14 @@ ThreadGraph build_worker(
     graph.end_assembly = end_assembly;
 
     // Reserving for unordered_map will actually allocate physical memory
-    std::unordered_map<std::uint64_t, NodeState> node_map;
+    std::unordered_map<std::uint64_t, std::size_t> node_map;
     std::unordered_map<EdgeKey, EdgeState, EdgeKeyHash> edge_map;
     const auto n_map_entries_est = n_minimizers_est / map_reserve_divisor;
     node_map.reserve(n_map_entries_est);
     edge_map.reserve(n_map_entries_est);
 
     for (std::size_t assembly_i = start_assembly; assembly_i < end_assembly; ++assembly_i) {
-        const bool is_target = is_targets[assembly_i];
-
+        const auto assembly_i_u32 = static_cast<std::uint32_t>(assembly_i);
         auto records = read_fasta(assembly_paths[assembly_i]);
         std::vector<std::string> record_ids;
         record_ids.reserve(records.size());
@@ -159,16 +149,8 @@ ThreadGraph build_worker(
                 }
 
                 // Add this minimizer to an existing node, or create a new node
-                auto [node_it, node_inserted] = node_map.try_emplace(m.out_hash);
-                ++(node_it->second.count);
-                if (node_inserted || node_it->second.last_seen_assembly != assembly_i) {
-                    if (is_target) {
-                        ++(node_it->second.n_tar);
-                    } else {
-                        ++(node_it->second.n_neg);
-                    }
-                    node_it->second.last_seen_assembly = assembly_i;
-                }
+                auto node_it = node_map.try_emplace(m.out_hash, 0).first;
+                ++node_it->second;
                 ++graph.n_kmers;
             }
 
@@ -186,9 +168,9 @@ ThreadGraph build_worker(
                 }
                 const EdgeKey key{u, v};
                 auto edge_it = edge_map.try_emplace(key).first;
-                if (edge_it->second.last_seen_assembly != assembly_i) {
-                    ++(edge_it->second.weight);
-                    edge_it->second.last_seen_assembly = assembly_i;
+                if (edge_it->second.last_seen_assembly != assembly_i_u32) {
+                    ++edge_it->second.weight;
+                    edge_it->second.last_seen_assembly = assembly_i_u32;
                 }
             }
         }
@@ -204,37 +186,57 @@ ThreadGraph build_worker(
     }
     std::unordered_map<EdgeKey, EdgeState, EdgeKeyHash>().swap(edge_map);
 
+    graph.n_nodes = node_map.size();
+    graph.nodes = NoInitArray<ThreadNode>(graph.n_nodes);
+    std::size_t node_i = 0;
+
     if (!low_memory) {
         // Build ThreadGraph.kmers (grouped by hash)
         graph.kmers = NoInitArray<Kmer>(graph.n_kmers);
+
+        // node_map values are counts while assigning node ranges,
+        // then cursors while scattering raw k-mers into those ranges
         std::size_t cursor = 0;
-        for (auto& [hash, state] : node_map) {
-            (void)hash;
-            state.cursor = cursor;
-            cursor += state.count;
+        for (auto& [hash, node_val] : node_map) {
+            const std::size_t count = node_val;
+            const std::size_t start = cursor;
+
+            graph.nodes[node_i++] = ThreadNode{
+                hash,
+                start,
+                count,
+                thread_id
+            };
+
+            node_val = start;
+            cursor += count;
         }
+        if (cursor != graph.n_kmers) {
+            throw std::logic_error("Node k-mer counts do not sum to graph.n_kmers");
+        }
+
         for (const auto& rk : raw_kmers) {
             auto node_it = node_map.find(rk.hash);
-            graph.kmers[node_it->second.cursor++] = rk.kmer;
+            if (node_it == node_map.end()) {
+                throw std::logic_error(
+                    "Standard-mode k-mer scattering found an unknown minimizer hash"
+                );
+            }
+            graph.kmers[node_it->second++] = rk.kmer;
         }
         std::vector<RawKmer>().swap(raw_kmers);
+    } else {
+        // In low-memory mode node_map values remain counts
+        for (const auto& [hash, count] : node_map) {
+            graph.nodes[node_i++] = ThreadNode{
+                hash,
+                0,
+                count,
+                thread_id
+            };
+        }
     }
-
-    graph.n_nodes = node_map.size();
-    graph.nodes = NoInitArray<ThreadNode>(node_map.size());
-    std::size_t node_i = 0;
-    for (const auto& [hash, state] : node_map) {
-        const std::size_t start = low_memory ? 0 : state.cursor - state.count;
-
-        graph.nodes[node_i++] = ThreadNode{
-            hash,
-            start,
-            state.count,
-            state.n_tar,
-            state.n_neg,
-            thread_id
-        };
-    }
+    std::unordered_map<std::uint64_t, std::size_t>().swap(node_map);
 
     return graph;
 }
@@ -313,13 +315,9 @@ Graph build(
     const std::vector<std::string>& assembly_paths,
     std::size_t kmerlen,
     std::size_t windowsize,
-    const std::vector<bool>& is_targets,
     std::size_t n_cpu,
     bool low_memory
 ) {
-    if (assembly_paths.size() != is_targets.size()) {
-        throw std::runtime_error("assembly_paths and is_targets must have the same length");
-    }
     if (assembly_paths.size() > std::numeric_limits<std::uint32_t>::max()) {
         throw std::runtime_error("Number of input assemblies exceeds uint32 range");
     }
@@ -344,7 +342,6 @@ Graph build(
                 assembly_paths,
                 kmerlen,
                 windowsize,
-                is_targets,
                 start_assembly,
                 end_assembly,
                 thread_id,

cpp/src/seqwin/build_internals.cpp

@@ -198,13 +198,9 @@ static MergedNodes merge_nodes(
     i = 0;
     while (i < n_nodes) {
         const auto hash = nodes[i].hash;
-        std::uint32_t n_tar = 0;
-        std::uint32_t n_neg = 0;
         const auto start = n_kmers;
 
         while (i < n_nodes && nodes[i].hash == hash) {
-            n_tar += nodes[i].n_tar;
-            n_neg += nodes[i].n_neg;
             const auto count = nodes[i].count;
 
             if (low_memory) {
@@ -221,7 +217,7 @@ static MergedNodes merge_nodes(
             ++i;
         }
 
-        merged.nodes[write_i++] = Node{hash, start, n_kmers, n_tar, n_neg};
+        merged.nodes[write_i++] = Node{hash, start, n_kmers};
     }
     return merged;
 }

cpp/src/seqwin/build_internals.hpp

@@ -25,10 +25,6 @@ struct ThreadNode {
     std::size_t start;
     /** Number of entries in `ThreadGraph.kmers` for this node. */
     std::size_t count;
-    /** Number of target assemblies containing this minimizer hash. */
-    std::uint32_t n_tar;
-    /** Number of non-target assemblies containing this minimizer hash. */
-    std::uint32_t n_neg;
     /** Worker thread that produced this node. Used by the merging step. */
     std::size_t thread_id;
 };