diff --git a/cpp/include/seqwin/graph.hpp b/cpp/include/seqwin/graph.hpp index 3142795..a61c799 100644 --- a/cpp/include/seqwin/graph.hpp +++ b/cpp/include/seqwin/graph.hpp @@ -22,8 +22,8 @@ struct Kmer { /** * @brief Minimizer graph node for one unique minimizer hash. * - * The `[start, stop)` range is a half-open interval into the parallel - * `Graph.kmers` and `Graph.idx` arrays for minimizers with this hash. + * The `[start, stop)` range is a half-open interval into `Graph.kmers` + * for minimizers with this hash. */ struct Node { /** Hash value of the minimizers represented by this node. */ @@ -58,8 +58,6 @@ struct Edge { struct Graph { /** Minimizer occurrences in all assemblies, grouped and sorted by hash. */ NoInitArray kmers; - /** Parallel to `kmers` and stores each minimizer's original generation index, ordered by genomic position. */ - NoInitArray idx; /** Sorted by hash. */ NoInitArray nodes; /** Sorted by hash. */ diff --git a/cpp/include/seqwin/helpers.hpp b/cpp/include/seqwin/helpers.hpp index f028415..77e6392 100644 --- a/cpp/include/seqwin/helpers.hpp +++ b/cpp/include/seqwin/helpers.hpp @@ -14,15 +14,14 @@ class ThreadPool; /** * @brief Thread-local minimizer graph node. * - * `start` and `count` define a range in `ThreadGraph.idx`. - * Values in `ThreadGraph.idx` index into `ThreadGraph.kmers`. + * `start` and `count` define a range in `ThreadGraph.kmers`. */ struct ThreadNode { /** Hash value of the minimizers represented by this node. */ std::uint64_t hash; - /** Start of entries in `ThreadGraph.idx` for this node. */ + /** Start of entries in `ThreadGraph.kmers` for this node. */ std::size_t start; - /** Number of entries in `ThreadGraph.idx` for this node. */ + /** 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; @@ -36,17 +35,8 @@ struct ThreadNode { * @brief Partial minimizer graph built by one worker thread. */ struct ThreadGraph { - /** - * Minimizer occurrences generated by this worker, stored in genomic position order. - * Uses `std::vector` for `push_back()` support (exact kmer count is not known). - */ - std::vector kmers; - /** - * Indices into `kmers`, grouped by minimizer hash. - * The order is stable within each hash group, preserving the original genomic - * position order of minimizers with the same hash. - */ - NoInitArray idx; + /** Minimizer occurrences generated by this worker, grouped by hash. */ + NoInitArray kmers; /** Unsorted. */ NoInitArray nodes; /** Unsorted. */ @@ -88,7 +78,6 @@ Graph merge_thread_graphs( Graph filter_kmers( const Kmer* kmers, - const std::size_t* idx, const Node* nodes, std::size_t n_nodes, std::vector used_hashes diff --git a/cpp/src/bindings/python_bindings.cpp b/cpp/src/bindings/python_bindings.cpp index 559163b..362e2a3 100644 --- a/cpp/src/bindings/python_bindings.cpp +++ b/cpp/src/bindings/python_bindings.cpp @@ -65,7 +65,6 @@ PYBIND11_MODULE(_core, m) { } auto kmers = array_to_numpy(std::move(graph.kmers)); - auto idx = array_to_numpy(std::move(graph.idx)); auto nodes = array_to_numpy(std::move(graph.nodes)); auto edges = array_to_numpy(std::move(graph.edges)); auto record_offsets = array_to_numpy(std::move(graph.record_offsets)); @@ -79,7 +78,7 @@ PYBIND11_MODULE(_core, m) { ids_by_assembly.append(ids_tuple); } - return py::make_tuple(kmers, idx, nodes, edges, record_offsets, ids_by_assembly); + return py::make_tuple(kmers, nodes, edges, record_offsets, ids_by_assembly); }, py::arg("assembly_paths"), py::arg("kmerlen"), @@ -90,16 +89,13 @@ PYBIND11_MODULE(_core, m) { m.def("_filter_kmers_native", [](py::array_t kmers, - py::array_t idx, py::array_t nodes, const std::vector& used_hashes ) { auto kmers_buf = kmers.request(); - auto idx_buf = idx.request(); auto nodes_buf = nodes.request(); const auto* kmers_ptr = static_cast(kmers_buf.ptr); - const auto* idx_ptr = static_cast(idx_buf.ptr); const auto* nodes_ptr = static_cast(nodes_buf.ptr); const auto n_nodes = static_cast(nodes_buf.shape[0]); @@ -108,7 +104,6 @@ PYBIND11_MODULE(_core, m) { py::gil_scoped_release release; graph = seqwin::filter_kmers( kmers_ptr, - idx_ptr, nodes_ptr, n_nodes, used_hashes @@ -116,13 +111,11 @@ PYBIND11_MODULE(_core, m) { } auto kmers_new = array_to_numpy(std::move(graph.kmers)); - auto idx_new = array_to_numpy(std::move(graph.idx)); auto nodes_new = array_to_numpy(std::move(graph.nodes)); - return py::make_tuple(kmers_new, idx_new, nodes_new); + return py::make_tuple(kmers_new, nodes_new); }, py::arg("kmers").noconvert(), - py::arg("idx").noconvert(), py::arg("nodes").noconvert(), py::arg("used_hashes") ); diff --git a/cpp/src/seqwin/graph.cpp b/cpp/src/seqwin/graph.cpp index 5aa3e38..63c22c3 100644 --- a/cpp/src/seqwin/graph.cpp +++ b/cpp/src/seqwin/graph.cpp @@ -19,6 +19,11 @@ namespace { constexpr std::size_t map_reserve_divisor = 100; +struct RawKmer { + std::uint64_t hash; + Kmer kmer; +}; + struct NodeState { std::size_t count = 0; std::size_t start = 0; @@ -60,17 +65,15 @@ ThreadGraph build_worker( ), windowsize ); + std::vector raw_kmers; + raw_kmers.reserve(n_kmers_est); ThreadGraph graph; - graph.kmers.reserve(n_kmers_est); 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; - std::vector hashes; // Used to build ThreadGraph.idx - hashes.reserve(n_kmers_est); - // Reserving for unordered_map will actually allocate physical memory std::unordered_map node_map; std::unordered_map edge_map; @@ -99,11 +102,13 @@ ThreadGraph build_worker( const auto mins = btllib::minimize_sequence(record.sequence, kmerlen, windowsize); for (const auto& m : mins) { - graph.kmers.push_back(Kmer{ - static_cast(m.pos), - static_cast(record_idx) + raw_kmers.push_back(RawKmer{ + m.out_hash, + Kmer{ + static_cast(m.pos), + static_cast(record_idx) + } }); - hashes.push_back(m.out_hash); // Add this minimizer to an existing node, or create a new node auto [node_it, node_inserted] = node_map.try_emplace(m.out_hash); @@ -151,8 +156,8 @@ ThreadGraph build_worker( } std::unordered_map().swap(edge_map); - // Build ThreadGraph.idx (grouped by minimizer hash) - graph.idx = NoInitArray(graph.n_kmers); + // 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; @@ -160,12 +165,11 @@ ThreadGraph build_worker( state.cursor = cursor; cursor += state.count; } - for (std::size_t i = 0; i < graph.n_kmers; ++i) { - const auto hash = hashes[i]; - auto node_it = node_map.find(hash); - graph.idx[node_it->second.cursor++] = i; + 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(hashes); + std::vector().swap(raw_kmers); graph.nodes = NoInitArray(node_map.size()); std::size_t node_i = 0; diff --git a/cpp/src/seqwin/helpers.cpp b/cpp/src/seqwin/helpers.cpp index 0f6642c..ca433b6 100644 --- a/cpp/src/seqwin/helpers.cpp +++ b/cpp/src/seqwin/helpers.cpp @@ -17,9 +17,9 @@ namespace seqwin { namespace { /** - * @brief Describes a contiguous local `idx` segment and its output position. + * @brief Describes a contiguous thread-local k-mer segment and its output position. */ -struct IdxSegment { +struct KmerSegment { std::size_t thread_id; std::size_t local_start; std::size_t out_start; @@ -144,13 +144,13 @@ static void lsd_radix_sort( * * 1. Sort nodes by hash. * 2. Merge nodes with the same hash into final graph nodes. - * 3. Return local `idx` segment metadata to materialize `kmers`. + * 3. Return local k-mer segment metadata to materialize final `kmers`. * * @param nodes Thread-local node array to sort and merge. * @param pool Thread pool used for radix sorting. - * @return Final merged nodes and segment metadata for rebuilding `kmers` and `idx`. + * @return Final merged nodes and segment metadata for rebuilding `kmers`. */ -static std::pair, std::vector> merge_nodes( +static std::pair, std::vector> merge_nodes( NoInitArray& nodes, ThreadPool& pool ) { @@ -173,9 +173,9 @@ static std::pair, std::vector> merge_nodes( } NoInitArray merged(unique_count); - // Aggregate nodes and keep idx ranges - std::vector idx_segments; - idx_segments.reserve(n_nodes); + // Aggregate nodes and keep k-mer ranges + std::vector kmer_segments; + kmer_segments.reserve(n_nodes); std::size_t n_kmers = 0; std::size_t write_i = 0; @@ -192,7 +192,7 @@ static std::pair, std::vector> merge_nodes( const auto count = nodes[i].count; if (count != 0) { - idx_segments.push_back(IdxSegment{ + kmer_segments.push_back(KmerSegment{ nodes[i].thread_id, nodes[i].start, n_kmers, @@ -206,29 +206,27 @@ static std::pair, std::vector> merge_nodes( const auto stop = n_kmers; merged[write_i++] = Node{hash, start, stop, n_tar, n_neg}; } - return {std::move(merged), std::move(idx_segments)}; + return {std::move(merged), std::move(kmer_segments)}; } static NoInitArray merge_kmers( const std::vector& graphs, std::size_t total_kmers, - const std::vector& idx_segments, - const std::vector& thread_record_offsets, + const std::vector& kmer_segments, + const std::vector& thread_record_offsets, ThreadPool& pool ) { NoInitArray kmers(total_kmers); - pool.parallel_for(idx_segments.size(), [&](std::size_t start, std::size_t end, std::size_t) { + pool.parallel_for(kmer_segments.size(), [&](std::size_t start, std::size_t end, std::size_t) { for (std::size_t s = start; s < end; ++s) { - const auto& segment = idx_segments[s]; - const auto& local_idx = graphs[segment.thread_id].idx; + const auto& segment = kmer_segments[s]; const auto& local_kmers = graphs[segment.thread_id].kmers; const auto offset = thread_record_offsets[segment.thread_id]; for (std::size_t k = 0; k < segment.count; ++k) { - const auto local_kmer_i = local_idx[segment.local_start + k]; - auto kmer = local_kmers[local_kmer_i]; - kmer.record_idx += static_cast(offset); + auto kmer = local_kmers[segment.local_start + k]; + kmer.record_idx += offset; kmers[segment.out_start + k] = kmer; } } @@ -236,29 +234,6 @@ static NoInitArray merge_kmers( return kmers; } -static NoInitArray merge_idx( - const std::vector& graphs, - std::size_t total_kmers, - const std::vector& idx_segments, - const std::vector& kmer_offsets, - ThreadPool& pool -) { - NoInitArray idx(total_kmers); - - pool.parallel_for(idx_segments.size(), [&](std::size_t start, std::size_t end, std::size_t) { - for (std::size_t s = start; s < end; ++s) { - const auto& segment = idx_segments[s]; - const auto offset = kmer_offsets[segment.thread_id]; - const auto& local_idx = graphs[segment.thread_id].idx; - - for (std::size_t k = 0; k < segment.count; ++k) { - idx[segment.out_start + k] = local_idx[segment.local_start + k] + offset; - } - } - }); - return idx; -} - static void merge_edges(NoInitArray& edges, ThreadPool& pool) { const std::size_t n_edges = edges.size(); @@ -335,33 +310,22 @@ Graph merge_thread_graphs( } merge_edges(graph.edges, pool); // Sort only - auto [nodes, idx_segments] = merge_nodes(graph.nodes, pool); + auto [nodes, kmer_segments] = merge_nodes(graph.nodes, pool); graph.nodes.reset(); - std::vector thread_record_offsets{0}; + std::vector thread_record_offsets{0}; auto kmers = merge_kmers( graphs, graph.n_kmers, - idx_segments, + kmer_segments, thread_record_offsets, pool ); - std::vector().swap(graph.kmers); - - std::vector kmer_offsets{0}; - auto idx = merge_idx( - graphs, - graph.n_kmers, - idx_segments, - kmer_offsets, - pool - ); - graph.idx.reset(); - std::vector().swap(idx_segments); + graph.kmers.reset(); + std::vector().swap(kmer_segments); return { std::move(kmers), - std::move(idx), std::move(nodes), std::move(graph.edges), std::move(graph.record_offsets), @@ -372,7 +336,7 @@ Graph merge_thread_graphs( log_python(" - Merging from " + std::to_string(graphs.size()) + " threads..."); // Merge record offsets - std::vector thread_record_offsets(graphs.size()); + std::vector thread_record_offsets(graphs.size()); std::vector record_offsets; record_offsets.reserve(n_assemblies + 1); record_offsets.push_back(0); @@ -382,7 +346,7 @@ Graph merge_thread_graphs( auto& local_offsets = graphs[t].record_offsets; const auto base = total_records; - thread_record_offsets[t] = base; + thread_record_offsets[t] = static_cast(base); total_records += local_offsets.back(); for (std::size_t i = 1; i < local_offsets.size(); ++i) { @@ -399,26 +363,19 @@ Graph merge_thread_graphs( merge_edges(edges, pool); auto thread_nodes = concat_nodes(graphs, pool); - auto [nodes, idx_segments] = merge_nodes(thread_nodes, pool); + auto [nodes, kmer_segments] = merge_nodes(thread_nodes, pool); thread_nodes.reset(); - std::vector kmer_offsets(graphs.size()); std::size_t total_kmers = 0; - for (std::size_t t = 0; t < graphs.size(); ++t) { - kmer_offsets[t] = total_kmers; - total_kmers += graphs[t].n_kmers; - } - - auto kmers = merge_kmers(graphs, total_kmers, idx_segments, thread_record_offsets, pool); - for (auto& graph : graphs) { - std::vector().swap(graph.kmers); + for (const auto& graph : graphs) { + total_kmers += graph.n_kmers; } - auto idx = merge_idx(graphs, total_kmers, idx_segments, kmer_offsets, pool); + auto kmers = merge_kmers(graphs, total_kmers, kmer_segments, thread_record_offsets, pool); for (auto& graph : graphs) { - graph.idx.reset(); + graph.kmers.reset(); } - std::vector().swap(idx_segments); + std::vector().swap(kmer_segments); std::vector> ids_by_assembly; ids_by_assembly.reserve(n_assemblies); @@ -432,7 +389,6 @@ Graph merge_thread_graphs( return { std::move(kmers), - std::move(idx), std::move(nodes), std::move(edges), std::move(record_offsets), @@ -442,7 +398,6 @@ Graph merge_thread_graphs( Graph filter_kmers( const Kmer* kmers, - const std::size_t* idx, const Node* nodes, std::size_t n_nodes, std::vector used_hashes @@ -478,7 +433,6 @@ Graph filter_kmers( graph.nodes = NoInitArray(used_node_indices.size()); graph.kmers = NoInitArray(n_kmers); - graph.idx = NoInitArray(n_kmers); std::size_t new_start = 0; for (std::size_t out_node_i = 0; out_node_i < used_node_indices.size(); ++out_node_i) { @@ -498,7 +452,6 @@ Graph filter_kmers( const auto out_i = new_start + k; const auto in_i = old_start + k; graph.kmers[out_i] = kmers[in_i]; - graph.idx[out_i] = idx[in_i]; } new_start += size; diff --git a/src/seqwin/_version.py b/src/seqwin/_version.py index b703f5c..9bdd4d2 100644 --- a/src/seqwin/_version.py +++ b/src/seqwin/_version.py @@ -1 +1 @@ -__version__ = '0.4.1' \ No newline at end of file +__version__ = '0.5.0' \ No newline at end of file diff --git a/src/seqwin/config.py b/src/seqwin/config.py index dcdf19d..3a0fbf1 100644 --- a/src/seqwin/config.py +++ b/src/seqwin/config.py @@ -30,7 +30,7 @@ - BLASTCONFIG (BlastConfig) - EDGE_W (str) - NODE_P (str) -- CONSEC_KMER_TH (int) +- CONSEC_KMER_MUL (float) - LEN_TH_MUL (float) - NO_BLAST_DIV (float) """ @@ -255,7 +255,7 @@ class WorkingDir: config (str): The `Config` instance saved as JSON. ['config.json'] assemblies_dir (str): Directory for downloaded assemblies. ['assemblies'] assemblies_csv (str): The `Assemblies` instance (`assemblies.py`) saved as CSV. ['assemblies.csv'] - graph (str): NumPy arrays representing the k-mer graph (`kmers`, `idx`, `nodes` and `edges`). ['graph.npz'] + graph (str): NumPy arrays representing the k-mer graph (`kmers`, `nodes` and `edges`). ['graph.npz'] mash (str): Mash sketch of all assemblies (.msh will be added by Mash). ['sketches'] blast_dir (str): Directory for the BLAST database. ['blastdb'] blast_log (str): Console output of the makeblastdb command, inside `blast_dir`. ['makeblastdb.log'] @@ -347,6 +347,5 @@ def config_logger(file: Path, level: int) -> None: EDGE_W: str = 'w' # Key for edge weight, used in networkx graphs. ['w'] NODE_P: str = 'p' # Key for node penalty, used in networkx graphs. ['p'] -CONSEC_KMER_TH: int = 2 # Max index difference between consecutive k-mers. [2] -LEN_TH_MUL: float = 1.5 # Deprecated. Multiplier for candidate sequence length threshold. [1.5] +CONSEC_KMER_MUL: float = 1.5 # Multiplier of windowsize for max position difference between consecutive k-mers. [1.5] NO_BLAST_DIV: float = 0.5 # Deprecated. Assumed divergence when there is no BLAST hit. [0.5] diff --git a/src/seqwin/core.py b/src/seqwin/core.py index ff5b65c..2899c35 100644 --- a/src/seqwin/core.py +++ b/src/seqwin/core.py @@ -140,7 +140,6 @@ def run(self) -> None: graph_path, allow_pickle=False, kmers=kmers.kmers, - idx=kmers.idx, nodes=kmers.nodes, edges=kmers.edges, record_offsets=kmers.record_offsets, diff --git a/src/seqwin/graph/__init__.py b/src/seqwin/graph/__init__.py index 01b0888..859ac5e 100644 --- a/src/seqwin/graph/__init__.py +++ b/src/seqwin/graph/__init__.py @@ -66,7 +66,6 @@ def build( n_cpu: int = 1 ) -> tuple[ NDArray[np.void], - NDArray[np.uintp], NDArray[np.void], NDArray[np.void], NDArray[np.uintp], @@ -77,7 +76,7 @@ def build( Example usage: ```python >>> from seqwin.graph import build - >>> kmers, idx, nodes, edges, record_offsets, record_ids = build( + >>> kmers, nodes, edges, record_offsets, record_ids = build( >>> assembly_paths = ..., >>> kmerlen = 21, >>> windowsize = 200, @@ -87,14 +86,12 @@ def build( ``` - `assembly_paths` and `is_targets` are parallel lists. - `kmers` stores minimizer occurrences in all assemblies, grouped and sorted by hash. - - `idx` is parallel to `kmers` and stores each minimizer's original generation index, ordered by genomic position. - `nodes` and `edges` are sorted by hash. The `[start, stop)` range in each node identifies minimizers with this hash. ```python >>> kmer_group = kmers[node['start']:node['stop']] >>> group_hash = node['hash'] - >>> original_indices = idx[node['start']:node['stop']] # strictly increasing ``` Use `record_offsets` to recover the original assembly and record index of each minimizer. @@ -121,8 +118,7 @@ def build( Dtype: `KMER_DTYPE` - 'pos' (uint32): 0-based position of the minimizer within its FASTA record. - 'record_idx' (uint32): 0-based global index of the FASTA record. - 2. NDArray[np.uintp]: The original indices assigned when minimizers are generated (ordered by genomic positions). - 3. NDArray[np.void]: A 1-D NumPy structured array of minimizer nodes. + 2. NDArray[np.void]: A 1-D NumPy structured array of minimizer nodes. Dtype: `NODE_DTYPE` - 'hash' (uint64): Hash value of the minimizers represented by this node. - 'start' (uintp): Start of the half-open range for this node's minimizer entries. @@ -130,13 +126,13 @@ def build( - 'n_tar' (uint32): Number of target assemblies containing this minimizer hash. - 'n_neg' (uint32): Number of non-target assemblies containing this minimizer hash. - 'penalty' (float64): Node penalty score used for downstream graph filtering. - 4. NDArray[np.void]: A 1-D NumPy structured array of weighted, undirected edges. + 3. NDArray[np.void]: A 1-D NumPy structured array of weighted, undirected edges. Dtype: `EDGE_DTYPE` - 'first' (uint64): Smaller endpoint hash of the undirected edge. - 'second' (uint64): Larger endpoint hash of the undirected edge. - 'weight' (uintp): Number of assemblies where the endpoints are adjacent. - 5. NDArray[np.uintp]: Cumulative global FASTA record offsets by assembly. - 6. list[tuple[str, ...]]: FASTA record IDs of each assembly. + 4. NDArray[np.uintp]: Cumulative global FASTA record offsets by assembly. + 5. list[tuple[str, ...]]: FASTA record IDs of each assembly. """ return _build_native( list(str(p) for p in assembly_paths), @@ -149,28 +145,24 @@ def build( def _filter_kmers( kmers: NDArray[np.void], - idx: NDArray[np.uintp], nodes: NDArray[np.void], used_hashes: frozenset[np.uint64] ) -> tuple[ NDArray[np.void], - NDArray[np.uintp], NDArray[np.void] ]: """ - 1. Remove k-mers (`kmers`, `idx` and `nodes`) not included in `used_hashes`. + 1. Remove k-mers (`kmers` and `nodes`) not included in `used_hashes`. 2. Update 'start' and 'stop' in nodes. Args: kmers (NDArray): See `KmerGraph.kmers`. - idx (NDArray): See `KmerGraph.idx`. nodes (NDArray): See `KmerGraph.nodes`. used_hashes (frozenset[np.uint64]): K-mers and nodes with these hash values are kept. Returns: tuple: A tuple containing 1. NDArray: See `KmerGraph.kmers`. - 2. NDArray: See `KmerGraph.idx`. - 3. NDArray: See `KmerGraph.nodes`. + 2. NDArray: See `KmerGraph.nodes`. """ - return _filter_kmers_native(kmers, idx, nodes, used_hashes) + return _filter_kmers_native(kmers, nodes, used_hashes) diff --git a/src/seqwin/kmers.py b/src/seqwin/kmers.py index 9bd8be6..c5fb706 100644 --- a/src/seqwin/kmers.py +++ b/src/seqwin/kmers.py @@ -49,9 +49,7 @@ class KmerGraph(object): 2. Extract low-penalty subgraphs from the k-mer graph with `self.filter()`. Attributes: - kmers (NDArray[np.void]): A 1-D NumPy structured array of k-mers from all assemblies, grouped and sorted by k-mer hashes. - idx (NDArray[np.uintp] | None): The original indices assigned when k-mers are generated (ordered by genomic positions). - Parallel to `kmers`. + kmers (NDArray[np.void]): A 1-D NumPy structured array of k-mers from all assemblies, grouped by node range. nodes (NDArray[np.void]): A 1-D NumPy structured array of k-mer nodes. For each node, `kmers[node['start']:node['stop']]` is the k-mer group with `node['hash']`. edges (NDArray[np.void]): A 1-D NumPy structured array of weighted, undirected edges. @@ -62,10 +60,9 @@ class KmerGraph(object): Generated with `self.filter()`. """ __slots__ = ( - 'kmers', 'idx', 'nodes', 'edges', 'record_offsets', 'graph', 'subgraphs', '_filtered_flag' + 'kmers', 'nodes', 'edges', 'record_offsets', 'graph', 'subgraphs', '_filtered_flag' ) kmers: NDArray[np.void] - idx: NDArray[np.uintp] nodes: NDArray[np.void] edges: NDArray[np.void] record_offsets: NDArray[np.uintp] @@ -88,7 +85,7 @@ def __init__(self, assemblies: Assemblies, kmerlen: int, windowsize: int, n_cpu: logger.info(f'Building minimizer graph from {n_assemblies} assemblies...') tik = time() - kmers, idx, nodes, edges, record_offsets, record_ids = build( + kmers, nodes, edges, record_offsets, record_ids = build( assemblies.path, kmerlen, windowsize, @@ -111,7 +108,6 @@ def __init__(self, assemblies: Assemblies, kmerlen: int, windowsize: int, n_cpu: print_time_delta(time()-tik) self.kmers = kmers - self.idx = idx self.nodes = nodes self.edges = edges self.record_offsets = record_offsets @@ -135,7 +131,6 @@ def filter( rng (random.Random): See `RunState` in `config.py`. """ kmers = self.kmers - idx = self.idx nodes = self.nodes edges = self.edges filtered_flag = self._filtered_flag @@ -160,12 +155,11 @@ def filter( # remove unused k-mers logger.info(' - Removing k-mers not included in any of the subgraphs...') - kmers, idx, nodes = _filter_kmers(kmers, idx, nodes, used_hashes) + kmers, nodes = _filter_kmers(kmers, nodes, used_hashes) logger.info(f' - {len(kmers)} k-mers left') print_time_delta(time()-tik) self.kmers = kmers - self.idx = idx self.nodes = nodes self.edges = edges self.graph = graph diff --git a/src/seqwin/markers.py b/src/seqwin/markers.py index d0b8e63..04ac46b 100644 --- a/src/seqwin/markers.py +++ b/src/seqwin/markers.py @@ -51,7 +51,7 @@ from .ncbi import blast from .graph import OrderedKmers from .utils import print_time_delta, log_and_raise, file_to_write, mp_wrapper -from .config import Config, RunState, HAS_BLAST, WORKINGDIR, BLASTCONFIG, CONSEC_KMER_TH, LEN_TH_MUL +from .config import Config, RunState, HAS_BLAST, WORKINGDIR, BLASTCONFIG, CONSEC_KMER_MUL # Set ConnectedKmers.is_bad as True if any of these warnings is present _BAD_WARNINGS = frozenset(( @@ -130,7 +130,7 @@ class ConnectedKmers(object): warnings: set is_bad: bool - def __init__(self, graph: nx.Graph, kmers: pd.DataFrame, kmerlen: int) -> None: + def __init__(self, graph: nx.Graph, kmers: pd.DataFrame, kmerlen: int, windowsize: int) -> None: """Given a subgraph of the k-mer graph, 1. Determine the boundary of the subgraph in each assembly. 2. Determine the representative k-mer order. @@ -142,6 +142,7 @@ def __init__(self, graph: nx.Graph, kmers: pd.DataFrame, kmerlen: int) -> None: It's a subset of `KmerGraph.kmers`, with index inherited. K-mers with adjacent indices are also adjacent in the assembly sequence. kmerlen (int): See `Config` in `config.py`. + windowsize (int): See `Config` in `config.py`. """ warnings = set() # passed to methods to add warnings in-place @@ -149,7 +150,7 @@ def __init__(self, graph: nx.Graph, kmers: pd.DataFrame, kmerlen: int) -> None: # kmers['strand'] = kmers['strand'].astype(str) # determine the boundary of the subgraph in each assembly - loc = ConnectedKmers.__get_loc(kmers, kmerlen) + loc = ConnectedKmers.__get_loc(kmers, kmerlen, windowsize) # determine the representative k-mer order and the number of targets having this order rep_order, n_rep = ConnectedKmers.__get_rep_order(loc, warnings) @@ -189,7 +190,7 @@ def __init__(self, graph: nx.Graph, kmers: pd.DataFrame, kmerlen: int) -> None: self.is_bad = is_bad @staticmethod - def __get_loc(kmers: pd.DataFrame, kmerlen: int) -> pd.DataFrame: + def __get_loc(kmers: pd.DataFrame, kmerlen: int, windowsize: int) -> pd.DataFrame: """Determine the location / boundary of the subgraph in each assembly. 1. Find consecutive k-mers for each assembly. 2. Determine the boundaries (start & stop) of each group of consecutive k-mers. @@ -198,22 +199,22 @@ def __get_loc(kmers: pd.DataFrame, kmerlen: int) -> pd.DataFrame: Args: kmers (pd.DataFrame): See `ConnectedKmers.__init__()`. kmerlen (int): See `Config` in `config.py`. + windowsize (int): See `Config` in `config.py`. Returns: pd.DataFrame: See `ConnectedKmers.loc`. """ - # sort by k-mer index - # essentially the same as sorting by k-mer position ['assembly_idx', 'record_idx', 'pos'] - kmers.sort_index(inplace=True) + # sort by genomic position + kmers.sort_values( + by=['assembly_idx', 'record_idx', 'pos'], + inplace=True, ignore_index=True + ) # since a subgraph might be repetitive in an assembly, we cannot simply groupby ['assembly_idx', 'record_idx'] - # we need to find runs of consecutive k-mers, by their indices (index is sorted) - # if two k-mers have consecutive index, then they are also adjacent on the assembly (if on the same sequence record) - # check KmerGraph.kmers for more info + # we need to find runs of consecutive k-mers, by their genomic positions (sorted) # consec_gp: consecutive group - # the definition of "consecutive" can be adjusted by changing CONSEC_KMER_TH - # we are doing this because the subgraph might occur more than one time in the same assembly - kmers['consec_gp'] = (kmers.index.diff() > CONSEC_KMER_TH).cumsum() + # the definition of "consecutive" can be adjusted by changing CONSEC_KMER_MUL + kmers['consec_gp'] = (kmers['pos'].diff() > CONSEC_KMER_MUL * windowsize).cumsum() # find the start and stop of each consecutive group # df.groupby preserves the order of rows within each group @@ -356,22 +357,22 @@ def _create_ck( graph: nx.Graph, nodes: tuple[np.uint64], kmers: tuple, - idx: tuple, record_offsets: NDArray[np.uintp], n_tar: int, - kmerlen: int + kmerlen: int, + windowsize: int ) -> ConnectedKmers: """Create a ConnectedKmers instance by taking the outputs of `_get_create_ck_args()`. """ - # add hash and idx to each group + # add hash to each group kmers_df = list() - for h, kmer_group, idx_group in zip(nodes, kmers, idx): - df = pd.DataFrame(kmer_group, index=idx_group) + for h, kmer_group in zip(nodes, kmers): + df = pd.DataFrame(kmer_group) df['hash'] = h kmers_df.append(df) # recover assembly_idx - kmers_df = pd.concat(kmers_df, ignore_index=False) + kmers_df = pd.concat(kmers_df, ignore_index=True) record_idx = kmers_df['record_idx'].to_numpy() assembly_idx = np.searchsorted( record_offsets, @@ -382,24 +383,24 @@ def _create_ck( kmers_df['assembly_idx'] = assembly_idx kmers_df['is_target'] = assembly_idx < n_tar - return ConnectedKmers(graph, kmers_df, kmerlen) + return ConnectedKmers(graph, kmers_df, kmerlen, windowsize) def _get_create_ck_args( - kg: KmerGraph, kmerlen: int, n_tar: int + kg: KmerGraph, n_tar: int, kmerlen: int, windowsize: int ) -> Generator[tuple[nx.Graph, pd.DataFrame, int], None, None]: """Generate input arguments for `_create_ck()`. Args: kg (KmerGraph): See `KmerGraph` in `kmers.py`. - kmerlen (int): See `Config` in `config.py`. n_tar (int): See `RunState` in `config.py`. + kmerlen (int): See `Config` in `config.py`. + windowsize (int): See `Config` in `config.py`. Yields: tuple: Input arguments of `_create_ck()`. """ kmers = kg.kmers - idx = kg.idx nodes = kg.nodes graph = kg.graph subgraphs = kg.subgraphs @@ -407,11 +408,9 @@ def _get_create_ck_args( # create a dict of hash -> k-mer group kmer_groups = dict() - idx_groups = dict() for node in nodes: h, start, stop = node['hash'], node['start'], node['stop'] kmer_groups[h] = kmers[start:stop] - idx_groups[h] = idx[start:stop] # yield function args for sg in subgraphs: @@ -422,11 +421,8 @@ def _get_create_ck_args( arg_kmers = tuple( kmer_groups.pop(h) for h in arg_nodes ) - arg_idx = tuple( - idx_groups.pop(h) for h in arg_nodes - ) - yield arg_graph, arg_nodes, arg_kmers, arg_idx, record_offsets, n_tar, kmerlen + yield arg_graph, arg_nodes, arg_kmers, record_offsets, n_tar, kmerlen, windowsize def _fetch_cks_seq( @@ -476,7 +472,13 @@ def _fetch_cks_seq( def _get_cks( - kmers: KmerGraph, assemblies: Assemblies, kmerlen: int, min_len: int, n_tar: int, n_cpu: int + kmers: KmerGraph, + n_tar: int, + kmerlen: int, + windowsize: int, + min_len: int, + assemblies: Assemblies, + n_cpu: int ) -> tuple[list[ConnectedKmers], list[str]]: """ 1. Create a ConnectedKmers instance for each low-penalty subgraph of the k-mer graph (`KmerGraph.subgraphs`). @@ -485,10 +487,11 @@ def _get_cks( Args: kmers (KmerGraph): See `KmerGraph` in `kmers.py`. - assemblies (Assemblies): See `Assemblies` in `assemblies.py`. + n_tar (int): See `RunState` in `config.py`. kmerlen (int): See `Config` in `config.py`. + windowsize (int): See `Config` in `config.py`. min_len (int): See `Config` in `config.py`. - n_tar (int): See `RunState` in `config.py`. + assemblies (Assemblies): See `Assemblies` in `assemblies.py`. n_cpu (int): See `Config` in `config.py`. Returns: @@ -503,7 +506,7 @@ def _get_cks( logger.info(' - Processing each subgraph...') all_cks: list[ConnectedKmers] = mp_wrapper( _create_ck, - _get_create_ck_args(kmers, kmerlen, n_tar), + _get_create_ck_args(kmers, n_tar, kmerlen, windowsize), n_cpu=n_cpu, n_jobs=len(kmers.subgraphs) ) @@ -740,6 +743,7 @@ def get_markers( """ overwrite = config.overwrite kmerlen = config.kmerlen + windowsize = config.windowsize min_len = config.min_len run_blast = config.run_blast blast_neg_only = config.blast_neg_only @@ -751,7 +755,7 @@ def get_markers( # extract marker from each low-penalty subgraph # all_cks: ConnectedKmers (ck) instances; all_reps: representative sequences - all_cks, all_reps = _get_cks(kmers, assemblies, kmerlen, min_len, n_tar, n_cpu) + all_cks, all_reps = _get_cks(kmers, n_tar, kmerlen, windowsize, min_len, assemblies, n_cpu) # evaluate each marker with BLAST if run_blast and HAS_BLAST: diff --git a/tests/smoke/fixtures/expected/graph.npz b/tests/smoke/fixtures/expected/graph.npz index c1973e8..9a7f105 100644 Binary files a/tests/smoke/fixtures/expected/graph.npz and b/tests/smoke/fixtures/expected/graph.npz differ diff --git a/tests/smoke/test_graph.py b/tests/smoke/test_graph.py index 5060567..a88bc07 100644 --- a/tests/smoke/test_graph.py +++ b/tests/smoke/test_graph.py @@ -13,19 +13,16 @@ def _sorted_edges(edges: np.ndarray) -> np.ndarray: return edge_values[idx] -def _assert_idx_invariants(kmers: np.ndarray, idx: np.ndarray, nodes: np.ndarray) -> None: - assert idx.dtype == np.uintp - assert len(idx) == len(kmers) - assert np.array_equal(np.sort(idx), np.arange(len(kmers), dtype=np.uintp)) + +def _assert_node_ranges(kmers: np.ndarray, nodes: np.ndarray) -> None: + total = 0 for node in nodes: start = int(node['start']) stop = int(node['stop']) - assert stop > start - block = kmers[start:stop] - assert len(block) == (stop - start) - node_idx = idx[start:stop] - assert len(node_idx) == len(block) - assert not np.array_equal(node_idx, np.arange(start, stop, dtype=np.uintp)) + assert 0 <= start <= stop <= len(kmers) + assert len(kmers[start:stop]) == (stop - start) + total += stop - start + assert total == len(kmers) def test_dtype_layouts() -> None: @@ -59,21 +56,21 @@ def test_build_threading_equivalence(targets_dir, non_targets_dir) -> None: ] is_targets = [True, True, False, False] - kmers_1, idx_1, nodes_1, edges_1, record_offsets_1, record_ids_1 = build( + kmers_1, nodes_1, edges_1, record_offsets_1, record_ids_1 = build( assembly_paths, kmerlen=7, windowsize=10, is_targets=is_targets, n_cpu=1, ) - kmers_2, idx_2, nodes_2, edges_2, record_offsets_2, record_ids_2 = build( + kmers_2, nodes_2, edges_2, record_offsets_2, record_ids_2 = build( assembly_paths, kmerlen=7, windowsize=10, is_targets=is_targets, n_cpu=2, ) - kmers_many, idx_many, nodes_many, edges_many, record_offsets_many, record_ids_many = build( + kmers_many, nodes_many, edges_many, record_offsets_many, record_ids_many = build( assembly_paths, kmerlen=7, windowsize=10, @@ -84,12 +81,6 @@ def test_build_threading_equivalence(targets_dir, non_targets_dir) -> None: assert kmers_1.dtype == KMER_DTYPE assert kmers_1.dtype.itemsize == 8 assert kmers_1.dtype.names == ('pos', 'record_idx') - for node in nodes_1: - start = int(node['start']) - stop = int(node['stop']) - assert stop > start - assert len(kmers_1[start:stop]) == (stop - start) - assert np.array_equal(record_offsets_1, np.array([0, 1, 2, 3, 4], dtype=np.uintp)) assert np.array_equal(np.unique(kmers_1['record_idx']), np.arange(4, dtype=np.uint32)) assert np.all(nodes_1['n_tar'] + nodes_1['n_neg'] > 0) @@ -102,14 +93,12 @@ def test_build_threading_equivalence(targets_dir, non_targets_dir) -> None: assert np.array_equal(edge_values_1[:, 2], edges_1["weight"]) assert nodes_1.dtype == NODE_DTYPE - _assert_idx_invariants(kmers_1, idx_1, nodes_1) - _assert_idx_invariants(kmers_2, idx_2, nodes_2) - _assert_idx_invariants(kmers_many, idx_many, nodes_many) + _assert_node_ranges(kmers_1, nodes_1) + _assert_node_ranges(kmers_2, nodes_2) + _assert_node_ranges(kmers_many, nodes_many) assert np.array_equal(kmers_1, kmers_2) assert np.array_equal(kmers_1, kmers_many) - assert np.array_equal(idx_1, idx_2) - assert np.array_equal(idx_1, idx_many) assert np.array_equal(nodes_1, nodes_2) assert np.array_equal(nodes_1, nodes_many) @@ -134,7 +123,7 @@ def write_fasta(path: Path, n_records: int) -> None: write_fasta(path, n_records) assembly_paths.append(path) - kmers, _, _, _, record_offsets, record_ids = build( + kmers, _, _, record_offsets, record_ids = build( assembly_paths, kmerlen=7, windowsize=10, @@ -147,34 +136,6 @@ def write_fasta(path: Path, n_records: int) -> None: assert np.array_equal(np.unique(kmers['record_idx']), np.arange(7, dtype=np.uint32)) -def test_create_ck_reconstructs_assembly_and_local_record_indices(monkeypatch) -> None: - captured = {} - - class FakeConnectedKmers: - def __init__(self, graph, kmers_df, kmerlen): - captured['graph'] = graph - captured['kmers_df'] = kmers_df.copy() - captured['kmerlen'] = kmerlen - - monkeypatch.setattr('seqwin.markers.ConnectedKmers', FakeConnectedKmers) - - graph = nx.Graph() - kmers = ( - np.array([(5, 0), (6, 2), (7, 5)], dtype=KMER_DTYPE), - ) - idx = (np.array([10, 20, 30], dtype=np.uintp),) - record_offsets = np.array([0, 2, 5, 6], dtype=np.uintp) - - _create_ck(graph, (np.uint64(123),), kmers, idx, record_offsets, n_tar=2, kmerlen=7) - - df = captured['kmers_df'].sort_index() - assert df['assembly_idx'].tolist() == [0, 1, 2] - assert df['record_idx'].tolist() == [0, 0, 0] - assert df['is_target'].tolist() == [True, True, False] - assert df['hash'].tolist() == [123, 123, 123] - assert captured['kmerlen'] == 7 - - def test_filter_kmers() -> None: kmers = np.array([ (10, 0), @@ -184,14 +145,13 @@ def test_filter_kmers() -> None: (31, 2), (32, 2), ], dtype=KMER_DTYPE) - idx = np.array([100, 101, 200, 300, 301, 302], dtype=np.uintp) nodes = np.array([ (10, 0, 2, 1, 0, 0.1), (20, 2, 3, 1, 0, 0.2), (30, 3, 6, 1, 1, 0.3), ], dtype=NODE_DTYPE) - kmers_new, idx_new, nodes_new = _filter_kmers(kmers, idx, nodes, {30, 10}) + kmers_new, nodes_new = _filter_kmers(kmers, nodes, {30, 10}) assert np.array_equal(nodes_new['hash'], np.array([10, 30], dtype=np.uint64)) assert np.array_equal(nodes_new['start'], np.array([0, 2], dtype=np.uintp)) @@ -204,7 +164,5 @@ def test_filter_kmers() -> None: (31, 2), (32, 2), ], dtype=KMER_DTYPE) - expected_idx = np.array([100, 101, 300, 301, 302], dtype=np.uintp) - assert np.array_equal(kmers_new, expected_kmers) - assert np.array_equal(idx_new, expected_idx) + _assert_node_ranges(kmers_new, nodes_new) diff --git a/tests/smoke/test_outputs.py b/tests/smoke/test_outputs.py index 9f680c1..5df9f61 100644 --- a/tests/smoke/test_outputs.py +++ b/tests/smoke/test_outputs.py @@ -25,7 +25,7 @@ def _assert_graph_matches_expected(actual_path: Path, expected_path: Path) -> No actual_keys = set(actual.files) expected_keys = set(expected.files) - required_keys = {'kmers', 'idx', 'nodes', 'edges', 'record_offsets'} + required_keys = {'kmers', 'nodes', 'edges', 'record_offsets'} assert actual_keys == expected_keys, f'graph arrays mismatch: actual={sorted(actual_keys)}, expected={sorted(expected_keys)}' assert actual_keys == required_keys, f'graph arrays must be exactly {sorted(required_keys)}, got {sorted(actual_keys)}' @@ -36,7 +36,6 @@ def _assert_graph_matches_expected(actual_path: Path, expected_path: Path) -> No assert actual_array.dtype == expected_array.dtype, f'{name} dtype mismatch: actual={actual_array.dtype}, expected={expected_array.dtype}' assert actual_array.shape == expected_array.shape, f'{name} shape mismatch: actual={actual_array.shape}, expected={expected_array.shape}' - np.testing.assert_array_equal(actual['idx'], expected['idx'], err_msg='idx array values mismatch') np.testing.assert_array_equal(actual['kmers'], expected['kmers'], err_msg='kmers array values mismatch') np.testing.assert_array_equal(actual['edges'], expected['edges'], err_msg='edges array values mismatch')