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& 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..ac5f895 100644 --- a/cpp/include/seqwin/helpers.hpp +++ b/cpp/include/seqwin/helpers.hpp @@ -3,6 +3,8 @@ #include #include #include +#include +#include #include #include "seqwin/graph.hpp" @@ -47,10 +49,22 @@ struct ThreadGraph { std::vector> 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 KmerMaps = std::vector>; + /** * @brief Emit a message through Python's logging module. * @@ -68,12 +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 merge_thread_graphs( std::vector& graphs, std::size_t n_assemblies, - ThreadPool& pool + ThreadPool& pool, + bool low_memory ); 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& 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..e0f2b40 100644 --- a/cpp/src/seqwin/graph.cpp +++ b/cpp/src/seqwin/graph.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include #include @@ -48,6 +49,13 @@ struct EdgeState { std::size_t last_seen_assembly = std::numeric_limits::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& assembly_paths, std::size_t kmerlen, @@ -55,7 +63,8 @@ ThreadGraph build_worker( const std::vector& 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 +75,16 @@ ThreadGraph build_worker( windowsize ); std::vector 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 node_map; @@ -102,13 +114,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(m.pos), - static_cast(record_idx) - } - }); + if (!low_memory) { + raw_kmers.push_back(RawKmer{ + m.out_hash, + Kmer{ + static_cast(m.pos), + static_cast(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 +170,24 @@ ThreadGraph build_worker( } std::unordered_map().swap(edge_map); - // Build ThreadGraph.kmers (grouped by hash) - graph.kmers = NoInitArray(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(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().swap(raw_kmers); } - std::vector().swap(raw_kmers); + graph.n_nodes = node_map.size(); graph.nodes = NoInitArray(node_map.size()); std::size_t node_i = 0; for (const auto& [hash, state] : node_map) { @@ -187,6 +204,73 @@ 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 recompute_kmers( + const std::vector& assembly_paths, + std::size_t kmerlen, + std::size_t windowsize, + const std::vector& graphs, + const std::vector& record_offsets, + KmerMaps& kmer_maps, + ThreadPool& pool +) { + std::size_t total_kmers = 0; + for (const auto& graph : graphs) { + total_kmers += graph.n_kmers; + } + NoInitArray 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 = kmer_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::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(m.pos), + static_cast(record_idx) + }; + } + } + } + } + }); + + return kmers; +} + } // namespace Graph build( @@ -194,7 +278,8 @@ Graph build( std::size_t kmerlen, std::size_t windowsize, const std::vector& 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 +301,41 @@ 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); + auto [graph, kmer_maps] = merge_thread_graphs( + graphs, + n_assemblies, + pool, + low_memory + ); + if (low_memory) { + log_python(" - Recomputing minimizers for low-memory mode..."); + graph.kmers = recompute_kmers( + assembly_paths, + kmerlen, + windowsize, + graphs, + graph.record_offsets, + kmer_maps, + pool + ); + } + return std::move(graph); } } // namespace seqwin diff --git a/cpp/src/seqwin/helpers.cpp b/cpp/src/seqwin/helpers.cpp index ca433b6..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 nodes; + std::vector kmer_segments; + KmerMaps kmer_maps; +}; + template NoInitArray concat(std::vector& 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, std::vector> merge_nodes( +static MergedNodes merge_nodes( NoInitArray& nodes, - ThreadPool& pool + const std::vector& 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, std::vector> merge_nodes( } ++unique_count; } - NoInitArray merged(unique_count); + merged.nodes = NoInitArray(unique_count); - // Aggregate nodes and keep k-mer ranges - std::vector 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,31 +209,35 @@ static std::pair, std::vector> 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; } + 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(kmer_segments)}; + return merged; } static NoInitArray merge_kmers( const std::vector& graphs, - std::size_t total_kmers, const std::vector& kmer_segments, const std::vector& thread_record_offsets, ThreadPool& pool ) { + std::size_t total_kmers = 0; + for (const auto& graph : graphs) { + total_kmers += graph.n_kmers; + } NoInitArray kmers(total_kmers); pool.parallel_for(kmer_segments.size(), [&](std::size_t start, std::size_t end, std::size_t) { @@ -298,10 +320,11 @@ void log_python(const std::string& message, const std::string& level) } } -Graph merge_thread_graphs( +std::pair merge_thread_graphs( std::vector& graphs, std::size_t n_assemblies, - ThreadPool& pool + ThreadPool& pool, + bool low_memory ) { if (graphs.size() == 1) { auto& graph = graphs[0]; @@ -310,26 +333,31 @@ Graph merge_thread_graphs( } merge_edges(graph.edges, pool); // Sort only - auto [nodes, kmer_segments] = merge_nodes(graph.nodes, pool); + + auto merged = merge_nodes(graph.nodes, graphs, pool, low_memory); graph.nodes.reset(); - std::vector thread_record_offsets{0}; - auto kmers = merge_kmers( - graphs, - graph.n_kmers, - kmer_segments, - thread_record_offsets, - pool - ); - graph.kmers.reset(); - std::vector().swap(kmer_segments); + NoInitArray kmers; + if (!low_memory) { + kmers = merge_kmers( + graphs, + merged.kmer_segments, + std::vector{0}, + pool + ); + graph.kmers.reset(); + std::vector().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) }; } @@ -363,19 +391,17 @@ Graph merge_thread_graphs( merge_edges(edges, pool); auto thread_nodes = concat_nodes(graphs, pool); - auto [nodes, kmer_segments] = merge_nodes(thread_nodes, pool); + auto merged = merge_nodes(thread_nodes, graphs, pool, low_memory); thread_nodes.reset(); - std::size_t total_kmers = 0; - for (const auto& graph : graphs) { - total_kmers += graph.n_kmers; - } - - auto kmers = merge_kmers(graphs, total_kmers, kmer_segments, thread_record_offsets, pool); - for (auto& graph : graphs) { - graph.kmers.reset(); + NoInitArray kmers; + if (!low_memory) { + kmers = merge_kmers(graphs, merged.kmer_segments, thread_record_offsets, pool); + for (auto& graph : graphs) { + graph.kmers.reset(); + } + std::vector().swap(merged.kmer_segments); } - std::vector().swap(kmer_segments); std::vector> ids_by_assembly; ids_by_assembly.reserve(n_assemblies); @@ -385,14 +411,15 @@ Graph merge_thread_graphs( } } - std::vector().swap(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) }; } 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) ) 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)