From df6075ba02d04b60a0eb9d79687126cb712b2742 Mon Sep 17 00:00:00 2001
From: mxwang66 <54554363+mxwang66@users.noreply.github.com>
Date: Fri, 5 Jun 2026 18:25:10 -0500
Subject: [PATCH 1/4] Add low-memory graph build mode

---
 cpp/include/seqwin/graph.hpp         |   4 +-
 cpp/include/seqwin/helpers.hpp       |  17 ++-
 cpp/src/bindings/python_bindings.cpp |   9 +-
 cpp/src/seqwin/graph.cpp             | 155 +++++++++++++++++++++------
 cpp/src/seqwin/helpers.cpp           | 130 +++++++++++++++++-----
 src/seqwin/graph/__init__.py         |   6 +-
 6 files changed, 256 insertions(+), 65 deletions(-)

diff --git a/cpp/include/seqwin/graph.hpp b/cpp/include/seqwin/graph.hpp
index a61c799..422b204 100644
--- a/cpp/include/seqwin/graph.hpp
+++ b/cpp/include/seqwin/graph.hpp
@@ -78,6 +78,7 @@ struct Graph {
  * @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.
  * @throws `std::runtime_error` If input sizes are inconsistent or counts exceed supported ranges.
  */
@@ -86,7 +87,8 @@ Graph build(
     std::size_t kmerlen,
     std::size_t windowsize,
     const std::vector<bool>& is_targets,
-    std::size_t n_cpu
+    std::size_t n_cpu = 1,
+    bool low_memory = false
 );
 
 } // namespace seqwin
diff --git a/cpp/include/seqwin/helpers.hpp b/cpp/include/seqwin/helpers.hpp
index 77e6392..03592be 100644
--- a/cpp/include/seqwin/helpers.hpp
+++ b/cpp/include/seqwin/helpers.hpp
@@ -3,6 +3,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <string>
+#include <unordered_map>
 #include <vector>
 
 #include "seqwin/graph.hpp"
@@ -47,10 +48,22 @@ struct ThreadGraph {
     std::vector<std::vector<std::string>> ids_by_assembly;
     /** Total number of minimizers generated by this worker. */
     std::size_t n_kmers = 0;
+    /** Number of thread-local nodes materialized by this worker. */
+    std::size_t n_nodes = 0;
     /** Global index of the first assembly assigned to this worker. */
     std::size_t start_assembly = 0;
+    /** Global index one past the last assembly assigned to this worker. */
+    std::size_t end_assembly = 0;
 };
 
+/**
+ * @brief Per-thread maps from a minimizer hash to its start position in `Graph.kmers`.
+ *
+ * Each thread/hash pair owns a disjoint segment inside `Graph.kmers`.
+ * When building `Graph.kmers` in the low memory mode, these starts are mutated as cursors.
+ */
+using NodeMaps = std::vector<std::unordered_map<std::uint64_t, std::size_t>>;
+
 /**
  * @brief Emit a message through Python's logging module.
  *
@@ -73,7 +86,9 @@ void log_python(
 Graph merge_thread_graphs(
     std::vector<ThreadGraph>& graphs,
     std::size_t n_assemblies,
-    ThreadPool& pool
+    ThreadPool& pool,
+    bool low_memory,
+    NodeMaps* node_maps
 );
 
 Graph filter_kmers(
diff --git a/cpp/src/bindings/python_bindings.cpp b/cpp/src/bindings/python_bindings.cpp
index 362e2a3..6d45028 100644
--- a/cpp/src/bindings/python_bindings.cpp
+++ b/cpp/src/bindings/python_bindings.cpp
@@ -50,7 +50,8 @@ PYBIND11_MODULE(_core, m) {
            std::size_t kmerlen,
            std::size_t windowsize,
            const std::vector<bool>& is_targets,
-           std::size_t n_cpu
+           std::size_t n_cpu,
+           bool low_memory
         ) {
             seqwin::Graph graph;
             {
@@ -60,7 +61,8 @@ PYBIND11_MODULE(_core, m) {
                     kmerlen,
                     windowsize,
                     is_targets,
-                    n_cpu
+                    n_cpu,
+                    low_memory
                 );
             }
 
@@ -84,7 +86,8 @@ PYBIND11_MODULE(_core, m) {
         py::arg("kmerlen"),
         py::arg("windowsize"),
         py::arg("is_targets"),
-        py::arg("n_cpu") = 1
+        py::arg("n_cpu") = 1,
+        py::arg("low_memory") = false
     );
 
     m.def("_filter_kmers_native",
diff --git a/cpp/src/seqwin/graph.cpp b/cpp/src/seqwin/graph.cpp
index 63c22c3..8f3321f 100644
--- a/cpp/src/seqwin/graph.cpp
+++ b/cpp/src/seqwin/graph.cpp
@@ -5,6 +5,7 @@
 #include <exception>
 #include <limits>
 #include <stdexcept>
+#include <string>
 #include <unordered_map>
 #include <utility>
 #include <vector>
@@ -55,7 +56,8 @@ ThreadGraph build_worker(
     const std::vector<bool>& is_targets,
     std::size_t start_assembly,
     std::size_t end_assembly,
-    std::size_t thread_id
+    std::size_t thread_id,
+    bool low_memory
 ) {
     // Estimate total minimizer count in all assemblies
     const auto n_kmers_est = seqwin::est_kmer_number(
@@ -66,13 +68,16 @@ ThreadGraph build_worker(
         windowsize
     );
     std::vector<RawKmer> raw_kmers;
-    raw_kmers.reserve(n_kmers_est);
+    if (!low_memory) {
+        raw_kmers.reserve(n_kmers_est);
+    }
 
     ThreadGraph graph;
     graph.record_offsets.reserve(end_assembly - start_assembly + 1);
     graph.record_offsets.push_back(0);
     graph.ids_by_assembly.reserve(end_assembly - start_assembly);
     graph.start_assembly = start_assembly;
+    graph.end_assembly = end_assembly;
 
     // Reserving for unordered_map will actually allocate physical memory
     std::unordered_map<std::uint64_t, NodeState> node_map;
@@ -102,13 +107,15 @@ ThreadGraph build_worker(
             const auto mins = btllib::minimize_sequence(record.sequence, kmerlen, windowsize);
 
             for (const auto& m : mins) {
-                raw_kmers.push_back(RawKmer{
-                    m.out_hash,
-                    Kmer{
-                        static_cast<std::uint32_t>(m.pos),
-                        static_cast<std::uint32_t>(record_idx)
-                    }
-                });
+                if (!low_memory) {
+                    raw_kmers.push_back(RawKmer{
+                        m.out_hash,
+                        Kmer{
+                            static_cast<std::uint32_t>(m.pos),
+                            static_cast<std::uint32_t>(record_idx)
+                        }
+                    });
+                }
 
                 // Add this minimizer to an existing node, or create a new node
                 auto [node_it, node_inserted] = node_map.try_emplace(m.out_hash);
@@ -156,21 +163,24 @@ ThreadGraph build_worker(
     }
     std::unordered_map<EdgeKey, EdgeState, EdgeKeyHash>().swap(edge_map);
 
-    // Build ThreadGraph.kmers (grouped by hash)
-    graph.kmers = NoInitArray<Kmer>(graph.n_kmers);
-    std::size_t cursor = 0;
-    for (auto& [hash, state] : node_map) {
-        (void)hash;
-        state.start = cursor;
-        state.cursor = cursor;
-        cursor += state.count;
-    }
-    for (const auto& rk : raw_kmers) {
-        auto node_it = node_map.find(rk.hash);
-        graph.kmers[node_it->second.cursor++] = rk.kmer;
+    if (!low_memory) {
+        // Build ThreadGraph.kmers (grouped by hash)
+        graph.kmers = NoInitArray<Kmer>(graph.n_kmers);
+        std::size_t cursor = 0;
+        for (auto& [hash, state] : node_map) {
+            (void)hash;
+            state.start = cursor;
+            state.cursor = cursor;
+            cursor += state.count;
+        }
+        for (const auto& rk : raw_kmers) {
+            auto node_it = node_map.find(rk.hash);
+            graph.kmers[node_it->second.cursor++] = rk.kmer;
+        }
+        std::vector<RawKmer>().swap(raw_kmers);
     }
-    std::vector<RawKmer>().swap(raw_kmers);
 
+    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) {
@@ -187,6 +197,67 @@ ThreadGraph build_worker(
     return graph;
 }
 
+NoInitArray<Kmer> recompute_kmers(
+    const std::vector<std::string>& assembly_paths,
+    std::size_t kmerlen,
+    std::size_t windowsize,
+    const std::vector<ThreadGraph>& graphs,
+    const std::vector<std::size_t>& record_offsets,
+    NodeMaps& node_maps,
+    ThreadPool& pool
+) {
+    std::size_t total_kmers = 0;
+    for (const auto& graph : graphs) {
+        total_kmers += graph.n_kmers;
+    }
+    NoInitArray<Kmer> kmers(total_kmers);
+
+    pool.parallel_for(graphs.size(), [&](std::size_t start, std::size_t end, std::size_t) {
+        for (std::size_t thread_id = start; thread_id < end; ++thread_id) {
+            const auto& graph = graphs[thread_id];
+            auto& hash_to_cursor = node_maps[thread_id];
+
+            for (std::size_t assembly_i = graph.start_assembly;
+                 assembly_i < graph.end_assembly;
+                 ++assembly_i) {
+                auto records = seqwin::read_fasta(assembly_paths[assembly_i]);
+
+                for (std::size_t record_i = 0; record_i < records.size(); ++record_i) {
+                    const auto record_idx = record_offsets[assembly_i] + record_i;
+                    auto& record = records[record_i];
+                    if (record.sequence.size() > std::numeric_limits<std::uint32_t>::max()) {
+                        throw std::runtime_error(
+                            "Sequence length exceeds uint32 range for record " +
+                            record.id + " in assembly " + assembly_paths[assembly_i]);
+                    }
+
+                    const auto mins = btllib::minimize_sequence(
+                        record.sequence,
+                        kmerlen,
+                        windowsize
+                    );
+
+                    for (const auto& m : mins) {
+                        auto cursor_it = hash_to_cursor.find(m.out_hash);
+                        if (cursor_it == hash_to_cursor.end()) {
+                            throw std::runtime_error(
+                                "Low-memory minimizer recomputation found an unknown minimizer hash"
+                            );
+                        }
+                        auto& cursor = cursor_it->second;
+                        kmers[cursor++] = Kmer{
+                            static_cast<std::uint32_t>(m.pos),
+                            static_cast<std::uint32_t>(record_idx)
+                        };
+                    }
+                }
+            }
+        }
+    });
+
+    return kmers;
+}
+
 } // namespace
 
 Graph build(
@@ -194,7 +265,8 @@ Graph build(
     std::size_t kmerlen,
     std::size_t windowsize,
     const std::vector<bool>& is_targets,
-    std::size_t n_cpu
+    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");
@@ -216,22 +288,43 @@ Graph build(
     const std::size_t rem = n_assemblies % n_workers;
 
     pool.parallel_for(n_workers, [&](std::size_t start, std::size_t end, std::size_t) {
-        for (std::size_t i = start; i < end; ++i) {
-            std::size_t chunk_start = i * base + std::min(i, rem);
-            std::size_t chunk_end = chunk_start + base + (i < rem ? 1 : 0);
-            graphs[i] = build_worker(
+        for (std::size_t thread_id = start; thread_id < end; ++thread_id) {
+            std::size_t start_assembly = thread_id * base + std::min(thread_id, rem);
+            std::size_t end_assembly = start_assembly + base + (thread_id < rem ? 1 : 0);
+            graphs[thread_id] = build_worker(
                 assembly_paths,
                 kmerlen,
                 windowsize,
                 is_targets,
-                chunk_start,
-                chunk_end,
-                i
+                start_assembly,
+                end_assembly,
+                thread_id,
+                low_memory
             );
         }
     });
 
-    return merge_thread_graphs(graphs, n_assemblies, pool);
+    NodeMaps node_maps;
+    auto graph = merge_thread_graphs(
+        graphs,
+        n_assemblies,
+        pool,
+        low_memory,
+        low_memory ? &node_maps : nullptr
+    );
+    if (low_memory) {
+        log_python(" - Recomputing minimizers for low-memory mode...");
+        graph.kmers = recompute_kmers(
+            assembly_paths,
+            kmerlen,
+            windowsize,
+            graphs,
+            graph.record_offsets,
+            node_maps,
+            pool
+        );
+    }
+    return graph;
 }
 
 } // namespace seqwin
diff --git a/cpp/src/seqwin/helpers.cpp b/cpp/src/seqwin/helpers.cpp
index ca433b6..80edf29 100644
--- a/cpp/src/seqwin/helpers.cpp
+++ b/cpp/src/seqwin/helpers.cpp
@@ -209,13 +209,74 @@ static std::pair<NoInitArray<Node>, std::vector<KmerSegment>> merge_nodes(
     return {std::move(merged), std::move(kmer_segments)};
 }
 
+static std::pair<NoInitArray<Node>, NodeMaps> merge_nodes_lm(
+    NoInitArray<ThreadNode>& nodes,
+    const std::vector<ThreadGraph>& graphs,
+    ThreadPool& pool
+) {
+    const std::size_t n_nodes = nodes.size();
+    if (n_nodes == 0) {
+        NodeMaps node_maps(graphs.size());
+        return {NoInitArray<Node>{}, std::move(node_maps)};
+    }
+
+    lsd_radix_sort(nodes, pool, &ThreadNode::hash);
+
+    // Determine final node count
+    std::size_t unique_count = 0;
+    std::size_t i = 0;
+    while (i < n_nodes) {
+        const auto hash = nodes[i].hash;
+        while (i < n_nodes && nodes[i].hash == hash) {
+            ++i;
+        }
+        ++unique_count;
+    }
+    NoInitArray<Node> merged(unique_count);
+
+    // Aggregate nodes and build node_maps
+    NodeMaps node_maps(graphs.size());
+    for (i = 0; i < graphs.size(); ++i) {
+        node_maps[i].reserve(graphs[i].n_nodes);
+    }
+
+    std::size_t n_kmers = 0;
+    std::size_t write_i = 0;
+    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 (count != 0) {
+                node_maps[nodes[i].thread_id][hash] = n_kmers;
+                n_kmers += count;
+            }
+            ++i;
+        }
+
+        const auto stop = n_kmers;
+        merged[write_i++] = Node{hash, start, stop, n_tar, n_neg};
+    }
+    return {std::move(merged), std::move(node_maps)};
+}
+
 static NoInitArray<Kmer> merge_kmers(
     const std::vector<ThreadGraph>& graphs,
-    std::size_t total_kmers,
     const std::vector<KmerSegment>& kmer_segments,
     const std::vector<std::uint32_t>& thread_record_offsets,
     ThreadPool& pool
 ) {
+    std::size_t total_kmers = 0;
+    for (const auto& graph : graphs) {
+        total_kmers += graph.n_kmers;
+    }
     NoInitArray<Kmer> kmers(total_kmers);
 
     pool.parallel_for(kmer_segments.size(), [&](std::size_t start, std::size_t end, std::size_t) {
@@ -301,8 +362,14 @@ void log_python(const std::string& message, const std::string& level)
 Graph merge_thread_graphs(
     std::vector<ThreadGraph>& graphs,
     std::size_t n_assemblies,
-    ThreadPool& pool
+    ThreadPool& pool,
+    bool low_memory,
+    NodeMaps* node_maps
 ) {
+    if (low_memory && node_maps == nullptr) {
+        throw std::runtime_error("low_memory merge requires node_maps");
+    }
+
     if (graphs.size() == 1) {
         auto& graph = graphs[0];
         if (graph.record_offsets.back() > std::numeric_limits<std::uint32_t>::max()) {
@@ -310,19 +377,24 @@ Graph merge_thread_graphs(
         }
 
         merge_edges(graph.edges, pool); // Sort only
-        auto [nodes, kmer_segments] = merge_nodes(graph.nodes, pool);
-        graph.nodes.reset();
-
-        std::vector<std::uint32_t> thread_record_offsets{0};
-        auto kmers = merge_kmers(
-            graphs,
-            graph.n_kmers,
-            kmer_segments,
-            thread_record_offsets,
-            pool
-        );
-        graph.kmers.reset();
-        std::vector<KmerSegment>().swap(kmer_segments);
+
+        NoInitArray<Node> nodes;
+        NoInitArray<Kmer> kmers;
+        if (low_memory) {
+            auto [nodes_tmp, node_maps_tmp] = merge_nodes_lm(graph.nodes, graphs, pool);
+            graph.nodes.reset();
+            nodes = std::move(nodes_tmp);
+            *node_maps = std::move(node_maps_tmp);
+        } else {
+            auto [nodes_tmp, kmer_segments] = merge_nodes(graph.nodes, pool);
+            graph.nodes.reset();
+            nodes = std::move(nodes_tmp);
+
+            std::vector<std::uint32_t> thread_record_offsets{0};
+            kmers = merge_kmers(graphs, kmer_segments, thread_record_offsets, pool);
+            graph.kmers.reset();
+            std::vector<KmerSegment>().swap(kmer_segments);
+        }
 
         return {
             std::move(kmers),
@@ -363,19 +435,23 @@ Graph merge_thread_graphs(
     merge_edges(edges, pool);
 
     auto thread_nodes = concat_nodes(graphs, pool);
-    auto [nodes, kmer_segments] = merge_nodes(thread_nodes, pool);
-    thread_nodes.reset();
-
-    std::size_t total_kmers = 0;
-    for (const auto& graph : graphs) {
-        total_kmers += graph.n_kmers;
-    }
+    NoInitArray<Node> nodes;
+    NoInitArray<Kmer> kmers;
+    if (low_memory) {
+        auto [nodes_tmp, node_maps_tmp] = merge_nodes_lm(thread_nodes, graphs, pool);
+        nodes = std::move(nodes_tmp);
+        *node_maps = std::move(node_maps_tmp);
+    } else {
+        auto [nodes_tmp, kmer_segments] = merge_nodes(thread_nodes, pool);
+        thread_nodes.reset();
+        nodes = std::move(nodes_tmp);
 
-    auto kmers = merge_kmers(graphs, total_kmers, kmer_segments, thread_record_offsets, pool);
-    for (auto& graph : graphs) {
-        graph.kmers.reset();
+        kmers = merge_kmers(graphs, kmer_segments, thread_record_offsets, pool);
+        for (auto& graph : graphs) {
+            graph.kmers.reset();
+        }
+        std::vector<KmerSegment>().swap(kmer_segments);
     }
-    std::vector<KmerSegment>().swap(kmer_segments);
 
     std::vector<std::vector<std::string>> ids_by_assembly;
     ids_by_assembly.reserve(n_assemblies);
@@ -385,8 +461,6 @@ Graph merge_thread_graphs(
         }
     }
 
-    std::vector<ThreadGraph>().swap(graphs);
-
     return {
         std::move(kmers),
         std::move(nodes),
diff --git a/src/seqwin/graph/__init__.py b/src/seqwin/graph/__init__.py
index 859ac5e..12e1db8 100644
--- a/src/seqwin/graph/__init__.py
+++ b/src/seqwin/graph/__init__.py
@@ -63,6 +63,7 @@ def build(
     kmerlen: int,
     windowsize: int,
     is_targets: Iterable[bool],
+    low_memory: bool = False,
     n_cpu: int = 1
 ) -> tuple[
     NDArray[np.void],
@@ -82,6 +83,7 @@ def build(
     >>>     windowsize = 200,
     >>>     is_targets = ...,
     >>>     n_cpu = 4,
+    >>>     low_memory = False
     >>> )
     ```
     - `assembly_paths` and `is_targets` are parallel lists.
@@ -111,6 +113,7 @@ def build(
         windowsize (int): Window size for minimizer sketch.
         is_targets (Iterable[bool]): Whether each assembly is a target assembly.
         n_cpu (int, optional): Number of worker threads to use. [1]
+        low_memory (bool, optional): Recompute minimizers in a second pass to reduce peak memory. [False]
 
     Returns:
         tuple: A tuple containing
@@ -139,7 +142,8 @@ def build(
         int(kmerlen),
         int(windowsize),
         list(bool(t) for t in is_targets),
-        int(n_cpu)
+        int(n_cpu),
+        bool(low_memory)
     )
 
 

From e62be5dd9133d90693e74224e46893713ad65322 Mon Sep 17 00:00:00 2001
From: mxwang66 <54554363+mxwang66@users.noreply.github.com>
Date: Fri, 5 Jun 2026 18:25:33 -0500
Subject: [PATCH 2/4] Add tests for low-memory mode

---
 tests/smoke/test_graph.py | 43 +++++++++++++++++++++++++++++++++++++--
 1 file changed, 41 insertions(+), 2 deletions(-)

diff --git a/tests/smoke/test_graph.py b/tests/smoke/test_graph.py
index a88bc07..27eec8c 100644
--- a/tests/smoke/test_graph.py
+++ b/tests/smoke/test_graph.py
@@ -1,10 +1,8 @@
 from pathlib import Path
 
-import networkx as nx
 import numpy as np
 
 from seqwin.graph import KMER_DTYPE, NODE_DTYPE, EDGE_DTYPE, build, _filter_kmers
-from seqwin.markers import _create_ck
 
 
 def _sorted_edges(edges: np.ndarray) -> np.ndarray:
@@ -13,6 +11,17 @@ def _sorted_edges(edges: np.ndarray) -> np.ndarray:
     return edge_values[idx]
 
 
+def _assert_graph_outputs_equal(standard, low_memory) -> None:
+    kmers_std, nodes_std, edges_std, offsets_std, ids_std = standard
+    kmers_lm, nodes_lm, edges_lm, offsets_lm, ids_lm = low_memory
+
+    assert np.array_equal(kmers_std, kmers_lm)
+    assert np.array_equal(nodes_std, nodes_lm)
+    assert np.array_equal(_sorted_edges(edges_std), _sorted_edges(edges_lm))
+    assert np.array_equal(offsets_std, offsets_lm)
+    assert ids_std == ids_lm
+    _assert_node_ranges(kmers_lm, nodes_lm)
+
 
 def _assert_node_ranges(kmers: np.ndarray, nodes: np.ndarray) -> None:
     total = 0
@@ -166,3 +175,33 @@ def test_filter_kmers() -> None:
     ], dtype=KMER_DTYPE)
     assert np.array_equal(kmers_new, expected_kmers)
     _assert_node_ranges(kmers_new, nodes_new)
+
+
+def test_low_memory_build_matches_standard(targets_dir, non_targets_dir) -> None:
+    assembly_paths = [
+        targets_dir / 'target-1.fasta',
+        targets_dir / 'target-2.fasta',
+        non_targets_dir / 'non-target-1.fasta',
+        non_targets_dir / 'non-target-2.fasta',
+    ]
+    is_targets = [True, True, False, False]
+
+    for n_cpu in (1, 2, 99):
+        standard = build(
+            assembly_paths,
+            kmerlen=7,
+            windowsize=10,
+            is_targets=is_targets,
+            n_cpu=n_cpu,
+            low_memory=False,
+        )
+        low_memory = build(
+            assembly_paths,
+            kmerlen=7,
+            windowsize=10,
+            is_targets=is_targets,
+            n_cpu=n_cpu,
+            low_memory=True,
+        )
+
+        _assert_graph_outputs_equal(standard, low_memory)

From 9b0de132e569553ce4ae633b80b204866e142cc3 Mon Sep 17 00:00:00 2001
From: mxwang66 <54554363+mxwang66@users.noreply.github.com>
Date: Fri, 5 Jun 2026 22:07:00 -0500
Subject: [PATCH 3/4] Clean up low-memory graph merge

---
 cpp/include/seqwin/helpers.hpp |  11 +-
 cpp/src/seqwin/graph.cpp       |  14 ++-
 cpp/src/seqwin/helpers.cpp     | 185 ++++++++++++---------------------
 3 files changed, 81 insertions(+), 129 deletions(-)

diff --git a/cpp/include/seqwin/helpers.hpp b/cpp/include/seqwin/helpers.hpp
index 03592be..ac5f895 100644
--- a/cpp/include/seqwin/helpers.hpp
+++ b/cpp/include/seqwin/helpers.hpp
@@ -4,6 +4,7 @@
 #include <cstdint>
 #include <string>
 #include <unordered_map>
+#include <utility>
 #include <vector>
 
 #include "seqwin/graph.hpp"
@@ -62,7 +63,7 @@ struct ThreadGraph {
  * Each thread/hash pair owns a disjoint segment inside `Graph.kmers`.
  * When building `Graph.kmers` in the low memory mode, these starts are mutated as cursors.
  */
-using NodeMaps = std::vector<std::unordered_map<std::uint64_t, std::size_t>>;
+using KmerMaps = std::vector<std::unordered_map<std::uint64_t, std::size_t>>;
 
 /**
  * @brief Emit a message through Python's logging module.
@@ -81,14 +82,14 @@ void log_python(
  * @param graphs Thread-local graphs produced during parallel graph construction.
  * @param n_assemblies Total number of input assemblies.
  * @param pool Thread pool used for parallel merge steps.
- * @return Fully merged minimizer graph.
+ * @param low_memory If True, return k-mer position metadata to build `kmers` in the second pass.
+ * @return Merged minimizer graph and k-mer maps. The second item is empty in standard mode.
  */
-Graph merge_thread_graphs(
+std::pair<Graph, KmerMaps> merge_thread_graphs(
     std::vector<ThreadGraph>& graphs,
     std::size_t n_assemblies,
     ThreadPool& pool,
-    bool low_memory,
-    NodeMaps* node_maps
+    bool low_memory
 );
 
 Graph filter_kmers(
diff --git a/cpp/src/seqwin/graph.cpp b/cpp/src/seqwin/graph.cpp
index 8f3321f..389dcc9 100644
--- a/cpp/src/seqwin/graph.cpp
+++ b/cpp/src/seqwin/graph.cpp
@@ -203,7 +203,7 @@ NoInitArray<Kmer> recompute_kmers(
     std::size_t windowsize,
     const std::vector<ThreadGraph>& graphs,
     const std::vector<std::size_t>& record_offsets,
-    NodeMaps& node_maps,
+    KmerMaps& kmer_maps,
     ThreadPool& pool
 ) {
     std::size_t total_kmers = 0;
@@ -215,7 +215,7 @@ NoInitArray<Kmer> recompute_kmers(
     pool.parallel_for(graphs.size(), [&](std::size_t start, std::size_t end, std::size_t) {
         for (std::size_t thread_id = start; thread_id < end; ++thread_id) {
             const auto& graph = graphs[thread_id];
-            auto& hash_to_cursor = node_maps[thread_id];
+            auto& hash_to_cursor = kmer_maps[thread_id];
 
             for (std::size_t assembly_i = graph.start_assembly;
                  assembly_i < graph.end_assembly;
@@ -304,13 +304,11 @@ Graph build(
         }
     });
 
-    NodeMaps node_maps;
-    auto graph = merge_thread_graphs(
+    auto [graph, kmer_maps] = merge_thread_graphs(
         graphs,
         n_assemblies,
         pool,
-        low_memory,
-        low_memory ? &node_maps : nullptr
+        low_memory
     );
     if (low_memory) {
         log_python(" - Recomputing minimizers for low-memory mode...");
@@ -320,11 +318,11 @@ Graph build(
             windowsize,
             graphs,
             graph.record_offsets,
-            node_maps,
+            kmer_maps,
             pool
         );
     }
-    return graph;
+    return std::move(graph);
 }
 
 } // namespace seqwin
diff --git a/cpp/src/seqwin/helpers.cpp b/cpp/src/seqwin/helpers.cpp
index 80edf29..cc0a28c 100644
--- a/cpp/src/seqwin/helpers.cpp
+++ b/cpp/src/seqwin/helpers.cpp
@@ -26,6 +26,12 @@ struct KmerSegment {
     std::size_t count;
 };
 
+struct MergedNodes {
+    NoInitArray<Node> nodes;
+    std::vector<KmerSegment> kmer_segments;
+    KmerMaps kmer_maps;
+};
+
 template <typename T, typename MemberPtr>
 NoInitArray<T> concat(std::vector<ThreadGraph>& graphs, MemberPtr member, ThreadPool& pool)
 {
@@ -144,19 +150,26 @@ static void lsd_radix_sort(
  *
  * 1. Sort nodes by hash.
  * 2. Merge nodes with the same hash into final graph nodes.
- * 3. Return local k-mer segment metadata to materialize final `kmers`.
+ * 3. Track k-mer output positions to build the final `kmers`.
  *
  * @param nodes Thread-local node array to sort and merge.
+ * @param graphs Thread-local graphs used to size k-mer maps.
  * @param pool Thread pool used for radix sorting.
- * @return Final merged nodes and segment metadata for rebuilding `kmers`.
+ * @param low_memory Whether to emit k-mer maps or k-mer segments.
+ * @return Final merged nodes and mode-specific k-mer position metadata.
  */
-static std::pair<NoInitArray<Node>, std::vector<KmerSegment>> merge_nodes(
+static MergedNodes merge_nodes(
     NoInitArray<ThreadNode>& nodes,
-    ThreadPool& pool
+    const std::vector<ThreadGraph>& graphs,
+    ThreadPool& pool,
+    bool low_memory
 ) {
+    MergedNodes merged;
+    merged.kmer_maps = KmerMaps(graphs.size());
+
     const std::size_t n_nodes = nodes.size();
     if (n_nodes == 0) {
-        return {};
+        return merged;
     }
 
     lsd_radix_sort(nodes, pool, &ThreadNode::hash);
@@ -171,11 +184,16 @@ static std::pair<NoInitArray<Node>, std::vector<KmerSegment>> merge_nodes(
         }
         ++unique_count;
     }
-    NoInitArray<Node> merged(unique_count);
+    merged.nodes = NoInitArray<Node>(unique_count);
 
-    // Aggregate nodes and keep k-mer ranges
-    std::vector<KmerSegment> kmer_segments;
-    kmer_segments.reserve(n_nodes);
+    // Aggregate nodes and track final kmer output positions
+    if (low_memory) {
+        for (std::size_t i = 0; i < graphs.size(); ++i) {
+            merged.kmer_maps[i].reserve(graphs[i].n_nodes);
+        }
+    } else {
+        merged.kmer_segments.reserve(n_nodes);
+    }
 
     std::size_t n_kmers = 0;
     std::size_t write_i = 0;
@@ -191,80 +209,23 @@ static std::pair<NoInitArray<Node>, std::vector<KmerSegment>> merge_nodes(
             n_neg += nodes[i].n_neg;
             const auto count = nodes[i].count;
 
-            if (count != 0) {
-                kmer_segments.push_back(KmerSegment{
+            if (low_memory) {
+                merged.kmer_maps[nodes[i].thread_id][hash] = n_kmers;
+            } else {
+                merged.kmer_segments.push_back(KmerSegment{
                     nodes[i].thread_id,
                     nodes[i].start,
                     n_kmers,
                     count
                 });
-                n_kmers += count;
-            }
-            ++i;
-        }
-
-        const auto stop = n_kmers;
-        merged[write_i++] = Node{hash, start, stop, n_tar, n_neg};
-    }
-    return {std::move(merged), std::move(kmer_segments)};
-}
-
-static std::pair<NoInitArray<Node>, NodeMaps> merge_nodes_lm(
-    NoInitArray<ThreadNode>& nodes,
-    const std::vector<ThreadGraph>& graphs,
-    ThreadPool& pool
-) {
-    const std::size_t n_nodes = nodes.size();
-    if (n_nodes == 0) {
-        NodeMaps node_maps(graphs.size());
-        return {NoInitArray<Node>{}, std::move(node_maps)};
-    }
-
-    lsd_radix_sort(nodes, pool, &ThreadNode::hash);
-
-    // Determine final node count
-    std::size_t unique_count = 0;
-    std::size_t i = 0;
-    while (i < n_nodes) {
-        const auto hash = nodes[i].hash;
-        while (i < n_nodes && nodes[i].hash == hash) {
-            ++i;
-        }
-        ++unique_count;
-    }
-    NoInitArray<Node> merged(unique_count);
-
-    // Aggregate nodes and build node_maps
-    NodeMaps node_maps(graphs.size());
-    for (i = 0; i < graphs.size(); ++i) {
-        node_maps[i].reserve(graphs[i].n_nodes);
-    }
-
-    std::size_t n_kmers = 0;
-    std::size_t write_i = 0;
-    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 (count != 0) {
-                node_maps[nodes[i].thread_id][hash] = n_kmers;
-                n_kmers += count;
             }
+            n_kmers += count;
             ++i;
         }
 
-        const auto stop = n_kmers;
-        merged[write_i++] = Node{hash, start, stop, n_tar, n_neg};
+        merged.nodes[write_i++] = Node{hash, start, n_kmers, n_tar, n_neg};
     }
-    return {std::move(merged), std::move(node_maps)};
+    return merged;
 }
 
 static NoInitArray<Kmer> merge_kmers(
@@ -359,17 +320,12 @@ void log_python(const std::string& message, const std::string& level)
     }
 }
 
-Graph merge_thread_graphs(
+std::pair<Graph, KmerMaps> merge_thread_graphs(
     std::vector<ThreadGraph>& graphs,
     std::size_t n_assemblies,
     ThreadPool& pool,
-    bool low_memory,
-    NodeMaps* node_maps
+    bool low_memory
 ) {
-    if (low_memory && node_maps == nullptr) {
-        throw std::runtime_error("low_memory merge requires node_maps");
-    }
-
     if (graphs.size() == 1) {
         auto& graph = graphs[0];
         if (graph.record_offsets.back() > std::numeric_limits<std::uint32_t>::max()) {
@@ -378,30 +334,30 @@ Graph merge_thread_graphs(
 
         merge_edges(graph.edges, pool); // Sort only
 
-        NoInitArray<Node> nodes;
+        auto merged = merge_nodes(graph.nodes, graphs, pool, low_memory);
+        graph.nodes.reset();
+
         NoInitArray<Kmer> kmers;
-        if (low_memory) {
-            auto [nodes_tmp, node_maps_tmp] = merge_nodes_lm(graph.nodes, graphs, pool);
-            graph.nodes.reset();
-            nodes = std::move(nodes_tmp);
-            *node_maps = std::move(node_maps_tmp);
-        } else {
-            auto [nodes_tmp, kmer_segments] = merge_nodes(graph.nodes, pool);
-            graph.nodes.reset();
-            nodes = std::move(nodes_tmp);
-
-            std::vector<std::uint32_t> thread_record_offsets{0};
-            kmers = merge_kmers(graphs, kmer_segments, thread_record_offsets, pool);
+        if (!low_memory) {
+            kmers = merge_kmers(
+                graphs,
+                merged.kmer_segments,
+                std::vector<std::uint32_t>{0},
+                pool
+            );
             graph.kmers.reset();
-            std::vector<KmerSegment>().swap(kmer_segments);
+            std::vector<KmerSegment>().swap(merged.kmer_segments);
         }
 
         return {
-            std::move(kmers),
-            std::move(nodes),
-            std::move(graph.edges),
-            std::move(graph.record_offsets),
-            std::move(graph.ids_by_assembly)
+            Graph{
+                std::move(kmers),
+                std::move(merged.nodes),
+                std::move(graph.edges),
+                std::move(graph.record_offsets),
+                std::move(graph.ids_by_assembly)
+            },
+            std::move(merged.kmer_maps)
         };
     }
 
@@ -435,22 +391,16 @@ Graph merge_thread_graphs(
     merge_edges(edges, pool);
 
     auto thread_nodes = concat_nodes(graphs, pool);
-    NoInitArray<Node> nodes;
-    NoInitArray<Kmer> kmers;
-    if (low_memory) {
-        auto [nodes_tmp, node_maps_tmp] = merge_nodes_lm(thread_nodes, graphs, pool);
-        nodes = std::move(nodes_tmp);
-        *node_maps = std::move(node_maps_tmp);
-    } else {
-        auto [nodes_tmp, kmer_segments] = merge_nodes(thread_nodes, pool);
-        thread_nodes.reset();
-        nodes = std::move(nodes_tmp);
+    auto merged = merge_nodes(thread_nodes, graphs, pool, low_memory);
+    thread_nodes.reset();
 
-        kmers = merge_kmers(graphs, kmer_segments, thread_record_offsets, pool);
+    NoInitArray<Kmer> kmers;
+    if (!low_memory) {
+        kmers = merge_kmers(graphs, merged.kmer_segments, thread_record_offsets, pool);
         for (auto& graph : graphs) {
             graph.kmers.reset();
         }
-        std::vector<KmerSegment>().swap(kmer_segments);
+        std::vector<KmerSegment>().swap(merged.kmer_segments);
     }
 
     std::vector<std::vector<std::string>> ids_by_assembly;
@@ -462,11 +412,14 @@ Graph merge_thread_graphs(
     }
 
     return {
-        std::move(kmers),
-        std::move(nodes),
-        std::move(edges),
-        std::move(record_offsets),
-        std::move(ids_by_assembly)
+        Graph{
+            std::move(kmers),
+            std::move(merged.nodes),
+            std::move(edges),
+            std::move(record_offsets),
+            std::move(ids_by_assembly)
+        },
+        std::move(merged.kmer_maps)
     };
 }
 

From 44e05d8d07e086361295a43d1e6f8276a5683c2a Mon Sep 17 00:00:00 2001
From: mxwang66 <54554363+mxwang66@users.noreply.github.com>
Date: Fri, 5 Jun 2026 22:28:55 -0500
Subject: [PATCH 4/4] Add docstrings

---
 cpp/src/seqwin/graph.cpp | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/cpp/src/seqwin/graph.cpp b/cpp/src/seqwin/graph.cpp
index 389dcc9..e0f2b40 100644
--- a/cpp/src/seqwin/graph.cpp
+++ b/cpp/src/seqwin/graph.cpp
@@ -49,6 +49,13 @@ struct EdgeState {
     std::size_t last_seen_assembly = std::numeric_limits<std::size_t>::max();
 };
 
+/**
+ * @brief Builds the thread-local portion of a minimizer graph.
+ *
+ * In standard mode, materializes thread-local `kmers` (grouped by hash).
+ * In low-memory mode, skips k-mer materialization and records only the node metadata
+ * needed for second-pass k-mer recomputation.
+ */
 ThreadGraph build_worker(
     const std::vector<std::string>& assembly_paths,
     std::size_t kmerlen,
@@ -197,6 +204,12 @@ ThreadGraph build_worker(
     return graph;
 }
 
+/**
+ * @brief Recomputes final `kmers` for a merged graph in low-memory mode.
+ *
+ * Replays each worker's assemblies, recomputes minimizers, and writes them directly
+ * into their final `Graph.kmers` positions using the per-thread k-mer maps.
+ */
 NoInitArray<Kmer> recompute_kmers(
     const std::vector<std::string>& assembly_paths,
     std::size_t kmerlen,