diff --git a/cpp/include/seqwin/graph.hpp b/cpp/include/seqwin/graph.hpp index c5c0934..30b396a 100644 --- a/cpp/include/seqwin/graph.hpp +++ b/cpp/include/seqwin/graph.hpp @@ -15,10 +15,8 @@ namespace seqwin { struct Kmer { /** 0-based position of the minimizer within its FASTA record. */ std::uint32_t pos; - /** 0-based index of the FASTA record within the assembly. */ - std::uint16_t record_idx; - /** Assembly index assigned to the source assembly. */ - std::uint16_t assembly_idx; + /** 0-based global index of the FASTA record. */ + std::uint32_t record_idx; }; /** @@ -70,6 +68,8 @@ struct Graph { NoInitArray nodes; /** Sorted by hash. */ NoInitArray edges; + /** Cumulative global FASTA record offsets by assembly. */ + std::vector record_offsets; /** FASTA record IDs of each assembly. */ std::vector> ids_by_assembly; }; @@ -77,22 +77,20 @@ struct Graph { /** * @brief Build a minimizer graph from assembly FASTA files. * - * `assembly_paths`, `assembly_indices`, and `is_targets` are parallel lists. + * `assembly_paths` and `is_targets` are parallel lists. * * @param assembly_paths Paths to input assemblies in FASTA format (plain or gzipped). * @param kmerlen K-mer length for minimizer sketch. * @param windowsize Window size for minimizer sketch. - * @param assembly_indices Assembly indices, expected to be `0..N-1`. * @param is_targets Whether each assembly is a target assembly. * @param n_cpu Number of worker threads to use. * @return Minimizer graph. - * @throws `std::runtime_error` If input sizes are inconsistent or indices exceed supported ranges. + * @throws `std::runtime_error` If input sizes are inconsistent or counts exceed supported ranges. */ Graph build( const std::vector& assembly_paths, std::size_t kmerlen, std::size_t windowsize, - const std::vector& assembly_indices, const std::vector& is_targets, std::size_t n_cpu ); diff --git a/cpp/include/seqwin/helpers.hpp b/cpp/include/seqwin/helpers.hpp index f5e1993..c04c4ca 100644 --- a/cpp/include/seqwin/helpers.hpp +++ b/cpp/include/seqwin/helpers.hpp @@ -30,6 +30,8 @@ struct ThreadGraph { NoInitArray nodes; /** Unsorted. */ NoInitArray edges; + /** Thread-local cumulative FASTA record offsets by assembly. */ + std::vector record_offsets; /** FASTA record IDs of each assembly. */ std::vector> ids_by_assembly; /** Total number of minimizers generated by this worker. */ diff --git a/cpp/src/bindings/python_bindings.cpp b/cpp/src/bindings/python_bindings.cpp index e0ff226..0f5ff38 100644 --- a/cpp/src/bindings/python_bindings.cpp +++ b/cpp/src/bindings/python_bindings.cpp @@ -1,6 +1,8 @@ #include #include #include +#include +#include #include #include @@ -12,8 +14,32 @@ namespace py = pybind11; +namespace { + +template +auto array_to_numpy(Array&& values) { + using Owner = typename std::decay::type; + using T = typename std::remove_cv< + typename std::remove_pointer().data())>::type + >::type; + + auto* owner = new Owner(std::forward(values)); + auto capsule = py::capsule(owner, [](void* ptr) { + delete static_cast(ptr); + }); + + return py::array_t( + {static_cast(owner->size())}, + {static_cast(sizeof(T))}, + owner->data(), + capsule + ); +} + +} // namespace + PYBIND11_MODULE(_core, m) { - PYBIND11_NUMPY_DTYPE(seqwin::Kmer, pos, record_idx, assembly_idx); + PYBIND11_NUMPY_DTYPE(seqwin::Kmer, pos, record_idx); PYBIND11_NUMPY_DTYPE(seqwin::Node, hash, start, stop, n_tar, n_neg, penalty); PYBIND11_NUMPY_DTYPE(seqwin::Edge, first, second, weight); @@ -23,7 +49,6 @@ PYBIND11_MODULE(_core, m) { [](const std::vector& assembly_paths, std::size_t kmerlen, std::size_t windowsize, - const std::vector& assembly_indices, const std::vector& is_targets, std::size_t n_cpu ) { @@ -34,45 +59,19 @@ PYBIND11_MODULE(_core, m) { assembly_paths, kmerlen, windowsize, - assembly_indices, is_targets, n_cpu ); } - auto owner = std::make_shared(std::move(graph)); - auto capsule = py::capsule( - new std::shared_ptr(owner), - [](void* ptr) { - delete static_cast*>(ptr); - } - ); - - auto kmers = py::array_t( - {static_cast(owner->kmers.size())}, - {static_cast(sizeof(seqwin::Kmer))}, - owner->kmers.data(), - capsule - ); - auto idx = py::array_t( - {static_cast(owner->idx.size())}, - {static_cast(sizeof(std::uint64_t))}, - owner->idx.data(), - capsule - ); - auto nodes = py::array_t( - {static_cast(owner->nodes.size())}, - {static_cast(sizeof(seqwin::Node))}, - owner->nodes.data(), - capsule - ); - auto edges = py::array_t( - {static_cast(owner->edges.size())}, - {static_cast(sizeof(seqwin::Edge))}, - owner->edges.data(), - capsule - ); + + 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)); + py::list ids_by_assembly; - for (const auto& ids : owner->ids_by_assembly) { + for (const auto& ids : graph.ids_by_assembly) { py::tuple ids_tuple(ids.size()); for (std::size_t i = 0; i < ids.size(); ++i) { ids_tuple[i] = ids[i]; @@ -80,12 +79,11 @@ PYBIND11_MODULE(_core, m) { ids_by_assembly.append(ids_tuple); } - return py::make_tuple(kmers, idx, nodes, edges, ids_by_assembly); + return py::make_tuple(kmers, idx, nodes, edges, record_offsets, ids_by_assembly); }, py::arg("assembly_paths"), py::arg("kmerlen"), py::arg("windowsize"), - py::arg("assembly_indices"), py::arg("is_targets"), py::arg("n_cpu") = 1 ); @@ -116,32 +114,10 @@ PYBIND11_MODULE(_core, m) { used_hashes ); } - auto owner = std::make_shared(std::move(graph)); - auto capsule = py::capsule( - new std::shared_ptr(owner), - [](void* ptr) { - delete static_cast*>(ptr); - } - ); - - auto kmers_new = py::array_t( - {static_cast(owner->kmers.size())}, - {static_cast(sizeof(seqwin::Kmer))}, - owner->kmers.data(), - capsule - ); - auto idx_new = py::array_t( - {static_cast(owner->idx.size())}, - {static_cast(sizeof(std::uint64_t))}, - owner->idx.data(), - capsule - ); - auto nodes_new = py::array_t( - {static_cast(owner->nodes.size())}, - {static_cast(sizeof(seqwin::Node))}, - owner->nodes.data(), - capsule - ); + + 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); }, diff --git a/cpp/src/seqwin/graph.cpp b/cpp/src/seqwin/graph.cpp index c4f87a1..caaed74 100644 --- a/cpp/src/seqwin/graph.cpp +++ b/cpp/src/seqwin/graph.cpp @@ -47,7 +47,6 @@ ThreadGraph build_worker( const std::vector& assembly_paths, std::size_t kmerlen, std::size_t windowsize, - const std::vector& assembly_indices, const std::vector& is_targets, std::size_t start_assembly, std::size_t end_assembly, @@ -64,7 +63,9 @@ ThreadGraph build_worker( ThreadGraph graph; graph.kmers.reserve(n_kmers_est); - graph.ids_by_assembly.resize(end_assembly - start_assembly); + 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 @@ -78,31 +79,29 @@ ThreadGraph build_worker( edge_map.reserve(n_map_entries_est); for (std::size_t assembly_i = start_assembly; assembly_i < end_assembly; ++assembly_i) { - const auto assembly_idx = assembly_indices[assembly_i]; const bool is_target = is_targets[assembly_i]; auto records = seqwin::read_fasta(assembly_paths[assembly_i]); - auto& record_ids = graph.ids_by_assembly[assembly_i - start_assembly]; - record_ids.resize(records.size()); - - for (std::size_t record_idx = 0; record_idx < records.size(); ++record_idx) { - if (record_idx > std::numeric_limits::max()) { - throw std::runtime_error("Number of FASTA records exceeds uint16 range"); + std::vector record_ids; + record_ids.reserve(records.size()); + + for (std::size_t record_i = 0; record_i < records.size(); ++record_i) { + const auto record_idx = graph.record_offsets.back() + 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]); } - auto& record = records[record_idx]; - record_ids[record_idx] = std::move(record.id); + record_ids.push_back(std::move(record.id)); // Generate minimizers for the current record const auto mins = btllib::minimize_sequence(record.sequence, kmerlen, windowsize); for (const auto& m : mins) { - if (m.pos > std::numeric_limits::max()) { - throw std::runtime_error("Minimizer position exceeds uint32 range"); - } graph.kmers.push_back(Kmer{ static_cast(m.pos), - static_cast(record_idx), - static_cast(assembly_idx) + static_cast(record_idx) }); hashes.push_back(m.out_hash); @@ -140,6 +139,8 @@ ThreadGraph build_worker( } } } + graph.record_offsets.push_back(graph.record_offsets.back() + records.size()); + graph.ids_by_assembly.push_back(std::move(record_ids)); } // Materialize edges first to reduce peak memory @@ -164,6 +165,7 @@ ThreadGraph build_worker( auto node_it = node_map.find(hash); graph.idx[node_it->second.cursor++] = static_cast(i); } + std::vector().swap(hashes); graph.nodes = NoInitArray(node_map.size()); std::size_t node_i = 0; @@ -188,19 +190,14 @@ Graph build( const std::vector& assembly_paths, std::size_t kmerlen, std::size_t windowsize, - const std::vector& assembly_indices, const std::vector& is_targets, std::size_t n_cpu ) { - if (assembly_paths.size() != assembly_indices.size() || - assembly_paths.size() != is_targets.size()) { - throw std::runtime_error( - "assembly_paths, assembly_indices, and is_targets must have the same length"); + if (assembly_paths.size() != is_targets.size()) { + throw std::runtime_error("assembly_paths and is_targets must have the same length"); } - for (std::size_t i = 0; i < assembly_indices.size(); ++i) { - if (assembly_indices[i] > std::numeric_limits::max()) { - throw std::runtime_error("Assembly index exceeds uint16 range"); - } + if (assembly_paths.size() > std::numeric_limits::max()) { + throw std::runtime_error("Number of input assemblies exceeds uint32 range"); } const auto n_assemblies = assembly_paths.size(); @@ -223,7 +220,6 @@ Graph build( assembly_paths, kmerlen, windowsize, - assembly_indices, is_targets, chunk_start, chunk_end, diff --git a/cpp/src/seqwin/helpers.cpp b/cpp/src/seqwin/helpers.cpp index 190a94b..df55df2 100644 --- a/cpp/src/seqwin/helpers.cpp +++ b/cpp/src/seqwin/helpers.cpp @@ -4,6 +4,8 @@ #include #include #include +#include +#include #include #include @@ -181,6 +183,7 @@ static NoInitArray merge_kmers( const std::vector& idx_segments, std::vector& graphs, std::size_t total_kmers, + const std::vector& thread_record_offsets, ThreadPool& pool ) { NoInitArray kmers(total_kmers); @@ -190,10 +193,13 @@ static NoInitArray merge_kmers( const auto& segment = idx_segments[s]; const auto& local_idx = graphs[segment.thread_id].idx; 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.length; ++k) { const auto local_kmer_i = local_idx[segment.local_start + k]; - kmers[segment.out_start + k] = local_kmers[static_cast(local_kmer_i)]; + auto kmer = local_kmers[static_cast(local_kmer_i)]; + kmer.record_idx += static_cast(offset); + kmers[segment.out_start + k] = kmer; } } }); @@ -333,14 +339,19 @@ Graph merge_thread_graphs( ) { if (graphs.size() == 1) { auto& graph = graphs[0]; + if (graph.record_offsets.back() > std::numeric_limits::max()) { + throw std::runtime_error("Total number of FASTA records exceeds uint32 range"); + } merge_edges(graph.edges, pool); auto idx_segments = merge_nodes(graph.nodes, pool); + std::vector thread_record_offsets{0}; auto kmers = merge_kmers( idx_segments, graphs, graph.n_kmers, + thread_record_offsets, pool ); std::vector().swap(graph.kmers); @@ -360,12 +371,36 @@ Graph merge_thread_graphs( std::move(idx), std::move(graph.nodes), std::move(graph.edges), + std::move(graph.record_offsets), std::move(graph.ids_by_assembly) }; } log_python(" - Merging from " + std::to_string(graphs.size()) + " threads..."); + // Merge record offsets + std::vector thread_record_offsets(graphs.size()); + std::vector record_offsets; + record_offsets.reserve(n_assemblies + 1); + record_offsets.push_back(0); + + std::uint64_t total_records = 0; + for (std::size_t t = 0; t < graphs.size(); ++t) { + auto& local_offsets = graphs[t].record_offsets; + + const auto base = total_records; + thread_record_offsets[t] = base; + total_records += local_offsets.back(); + + for (std::size_t i = 1; i < local_offsets.size(); ++i) { + record_offsets.push_back(base + local_offsets[i]); + } + std::vector().swap(local_offsets); + } + if (total_records > std::numeric_limits::max()) { + throw std::runtime_error("Total number of FASTA records exceeds uint32 range"); + } + // Merge edges and nodes first to reduce peak memory auto edges = concat_edges(graphs, pool); merge_edges(edges, pool); @@ -375,12 +410,12 @@ Graph merge_thread_graphs( std::vector kmer_offsets(graphs.size()); std::uint64_t total_kmers = 0; - for (std::size_t r = 0; r < graphs.size(); ++r) { - kmer_offsets[r] = total_kmers; - total_kmers += graphs[r].n_kmers; + 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(idx_segments, graphs, total_kmers, pool); + auto kmers = merge_kmers(idx_segments, graphs, total_kmers, thread_record_offsets, pool); for (auto& graph : graphs) { std::vector().swap(graph.kmers); } @@ -404,6 +439,7 @@ Graph merge_thread_graphs( std::move(idx), std::move(nodes), std::move(edges), + std::move(record_offsets), std::move(ids_by_assembly) }; } diff --git a/src/seqwin/core.py b/src/seqwin/core.py index cc8867d..ff5b65c 100644 --- a/src/seqwin/core.py +++ b/src/seqwin/core.py @@ -143,6 +143,7 @@ def run(self) -> None: idx=kmers.idx, nodes=kmers.nodes, edges=kmers.edges, + record_offsets=kmers.record_offsets, ) logger.info(f'Filtering is turned off. Raw minimizer graph is saved as {graph_path}') else: diff --git a/src/seqwin/graph/__init__.py b/src/seqwin/graph/__init__.py index fdca009..f7cb1f7 100644 --- a/src/seqwin/graph/__init__.py +++ b/src/seqwin/graph/__init__.py @@ -7,15 +7,8 @@ Usage: ------ ```python ->>> from pathlib import Path >>> from seqwin.graph import build ->>> kmers, idx, nodes, edges, record_ids = build( ->>> assembly_paths=[Path('example1.fa'), Path('example2.fa.gz')], ->>> kmerlen=21, ->>> windowsize=200, ->>> assembly_idx=[0, 1], ->>> is_target=[True, False], ->>> ) +>>> help(build) ``` Dependencies: @@ -46,8 +39,7 @@ KMER_DTYPE = np.dtype([ ('pos', np.uint32), - ('record_idx', np.uint16), - ('assembly_idx', np.uint16) + ('record_idx', np.uint32), ]) NODE_DTYPE = np.dtype([ @@ -70,7 +62,6 @@ def build( assembly_paths: Iterable[Path], kmerlen: int, windowsize: int, - assembly_indices: Iterable[int], is_targets: Iterable[bool], n_cpu: int = 1 ) -> tuple[ @@ -78,6 +69,7 @@ def build( NDArray[np.uint64], NDArray[np.void], NDArray[np.void], + NDArray[np.uint64], list[tuple[str, ...]] ]: """Build a Seqwin minimizer graph. @@ -85,16 +77,15 @@ def build( Example usage: ```python >>> from seqwin.graph import build - >>> kmers, idx, nodes, edges, record_ids = build( + >>> kmers, idx, nodes, edges, record_offsets, record_ids = build( >>> assembly_paths = ..., >>> kmerlen = 21, >>> windowsize = 200, - >>> assembly_indices = ..., >>> is_targets = ..., >>> n_cpu = 4, >>> ) ``` - - `assembly_paths`, `assembly_indices`, and `is_targets` are parallel lists. + - `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. @@ -106,11 +97,21 @@ def build( >>> original_indices = idx[node['start']:node['stop']] # strictly increasing ``` + Use `record_offsets` to recover the original assembly and record index of each minimizer. + ```python + >>> import numpy as np + >>> assembly_idx = np.searchsorted( + >>> record_offsets, + >>> kmers['record_idx'], + >>> side='right', + >>> ) - 1 + >>> record_idx = kmers['record_idx'] - record_offsets[assembly_idx] + ``` + Args: assembly_paths (Iterable[Path]): Paths to input assemblies in FASTA format (plain or gzipped). kmerlen (int): K-mer length for minimizer sketch. windowsize (int): Window size for minimizer sketch. - assembly_indices (Iterable[int]): Assembly indices. is_targets (Iterable[bool]): Whether each assembly is a target assembly. n_cpu (int, optional): Number of worker threads to use. [1] @@ -119,8 +120,7 @@ def build( 1. NDArray[np.void]: A 1-D NumPy structured array of minimizers from all assemblies. Dtype: `KMER_DTYPE` - 'pos' (uint32): 0-based position of the minimizer within its FASTA record. - - 'record_idx' (uint16): 0-based index of the FASTA record within the assembly. - - 'assembly_idx' (uint16): Assembly index assigned to the source assembly. + - 'record_idx' (uint32): 0-based global index of the FASTA record. 2. NDArray[np.uint64]: 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. Dtype: `NODE_DTYPE` @@ -135,13 +135,13 @@ def build( - 'first' (uint64): Smaller endpoint hash of the undirected edge. - 'second' (uint64): Larger endpoint hash of the undirected edge. - 'weight' (uint64): Number of assemblies where the endpoints are adjacent. - 5. list[tuple[str, ...]]: FASTA record IDs of each assembly. + 5. NDArray[np.uint64]: Cumulative global FASTA record offsets by assembly. + 6. list[tuple[str, ...]]: FASTA record IDs of each assembly. """ return _build_native( list(str(p) for p in assembly_paths), int(kmerlen), int(windowsize), - list(int(i) for i in assembly_indices), list(bool(t) for t in is_targets), int(n_cpu) ) diff --git a/src/seqwin/kmers.py b/src/seqwin/kmers.py index 1721819..b8fd2c6 100644 --- a/src/seqwin/kmers.py +++ b/src/seqwin/kmers.py @@ -56,17 +56,19 @@ class KmerGraph(object): 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. Edge weight is the number of assemblies where the two k-mers are adjacent. + record_offsets (NDArray[np.uint64]): Cumulative global FASTA record offsets by assembly. graph (nx.Graph): The graph instance built from filtered nodes and edges. subgraphs (tuple[frozenset[np.uint64], ...] | None): Low-penalty subgraphs. Each subgraph is a set of k-mer hash values. Generated with `self.filter()`. """ __slots__ = ( - 'kmers', 'idx', 'nodes', 'edges', 'graph', 'subgraphs', '_filtered_flag' + 'kmers', 'idx', 'nodes', 'edges', 'record_offsets', 'graph', 'subgraphs', '_filtered_flag' ) kmers: NDArray[np.void] idx: NDArray[np.uint64] nodes: NDArray[np.void] edges: NDArray[np.void] + record_offsets: NDArray[np.uint64] graph: nx.Graph subgraphs: tuple[frozenset[np.uint64], ...] | None _filtered_flag: bool # True if `self.filter()` is called @@ -86,11 +88,10 @@ 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_ids = build( + kmers, idx, nodes, edges, record_offsets, record_ids = build( assemblies.path, kmerlen, windowsize, - assemblies.index, assemblies.is_target, n_cpu=n_cpu, ) @@ -113,6 +114,7 @@ def __init__(self, assemblies: Assemblies, kmerlen: int, windowsize: int, n_cpu: self.idx = idx self.nodes = nodes self.edges = edges + self.record_offsets = record_offsets self.graph = None # create graph after filtering nodes and edges self.subgraphs = None self._filtered_flag = False diff --git a/src/seqwin/markers.py b/src/seqwin/markers.py index bf8b837..ac72e68 100644 --- a/src/seqwin/markers.py +++ b/src/seqwin/markers.py @@ -44,6 +44,7 @@ import numpy as np import pandas as pd import networkx as nx +from numpy.typing import NDArray from .assemblies import Assemblies from .kmers import KmerGraph @@ -58,8 +59,6 @@ 'dup', # has duplicate k-mers 'rev' # k-mer ordering is reversible )) -# Translation table for complement k-mer strands, deprecated. ['+' -> '-', '-' -> '+'] -_KMER_STRAND_COMP = str.maketrans('+-', '-+') @dataclass(slots=True, frozen=True) @@ -354,7 +353,13 @@ def __get_graph_order(graph: nx.Graph, rep_order: OrderedKmers, warnings: set) - def _create_ck( - graph: nx.Graph, nodes: tuple[np.uint64], kmers: tuple, idx: tuple, n_tar: int, kmerlen: int + graph: nx.Graph, + nodes: tuple[np.uint64], + kmers: tuple, + idx: tuple, + record_offsets: NDArray[np.uint64], + n_tar: int, + kmerlen: int ) -> ConnectedKmers: """Create a ConnectedKmers instance by taking the outputs of `_get_create_ck_args()`. """ @@ -363,10 +368,19 @@ def _create_ck( for h, kmer_group, idx_group in zip(nodes, kmers, idx): df = pd.DataFrame(kmer_group, index=idx_group) df['hash'] = h - df['is_target'] = df['assembly_idx'] < n_tar kmers_df.append(df) + # recover assembly_idx kmers_df = pd.concat(kmers_df, ignore_index=False) + record_idx = kmers_df['record_idx'].to_numpy() + assembly_idx = np.searchsorted( + record_offsets, + record_idx, + side='right', + ) - 1 + kmers_df['record_idx'] = record_idx - record_offsets[assembly_idx] + kmers_df['assembly_idx'] = assembly_idx + kmers_df['is_target'] = assembly_idx < n_tar return ConnectedKmers(graph, kmers_df, kmerlen) @@ -389,6 +403,7 @@ def _get_create_ck_args( nodes = kg.nodes graph = kg.graph subgraphs = kg.subgraphs + record_offsets = kg.record_offsets # create a dict of hash -> k-mer group kmer_groups = dict() @@ -411,7 +426,7 @@ def _get_create_ck_args( idx_groups.pop(h) for h in arg_nodes ) - yield arg_graph, arg_nodes, arg_kmers, arg_idx, n_tar, kmerlen + yield arg_graph, arg_nodes, arg_kmers, arg_idx, record_offsets, n_tar, kmerlen def _fetch_cks_seq( diff --git a/tests/smoke/fixtures/expected/graph.npz b/tests/smoke/fixtures/expected/graph.npz index 49cfa34..c1973e8 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 f77e076..3bd9403 100644 --- a/tests/smoke/test_graph.py +++ b/tests/smoke/test_graph.py @@ -1,6 +1,10 @@ +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: @@ -26,8 +30,8 @@ def _assert_idx_invariants(kmers: np.ndarray, idx: np.ndarray, nodes: np.ndarray def test_dtype_layouts() -> None: assert KMER_DTYPE.itemsize == 8 - assert KMER_DTYPE.names == ('pos', 'record_idx', 'assembly_idx') - assert KMER_DTYPE['assembly_idx'] == np.dtype(np.uint16) + assert KMER_DTYPE.names == ('pos', 'record_idx') + assert KMER_DTYPE['record_idx'] == np.dtype(np.uint32) assert NODE_DTYPE.names == ('hash', 'start', 'stop', 'n_tar', 'n_neg', 'penalty') assert NODE_DTYPE["n_tar"] == np.dtype(np.uint32) @@ -48,44 +52,41 @@ def test_build_threading_equivalence(targets_dir, non_targets_dir) -> None: non_targets_dir / 'non-target-1.fasta', non_targets_dir / 'non-target-2.fasta', ] - assembly_indices = [0, 1, 2, 3] is_targets = [True, True, False, False] - kmers_1, idx_1, nodes_1, edges_1, record_ids_1 = build( + kmers_1, idx_1, nodes_1, edges_1, record_offsets_1, record_ids_1 = build( assembly_paths, kmerlen=7, windowsize=10, - assembly_indices=assembly_indices, is_targets=is_targets, n_cpu=1, ) - kmers_2, idx_2, nodes_2, edges_2, record_ids_2 = build( + kmers_2, idx_2, nodes_2, edges_2, record_offsets_2, record_ids_2 = build( assembly_paths, kmerlen=7, windowsize=10, - assembly_indices=assembly_indices, is_targets=is_targets, n_cpu=2, ) - kmers_many, idx_many, nodes_many, edges_many, record_ids_many = build( + kmers_many, idx_many, nodes_many, edges_many, record_offsets_many, record_ids_many = build( assembly_paths, kmerlen=7, windowsize=10, - assembly_indices=assembly_indices, is_targets=is_targets, n_cpu=99, ) assert kmers_1.dtype == KMER_DTYPE assert kmers_1.dtype.itemsize == 8 - assert kmers_1.dtype.names == ('pos', 'record_idx', 'assembly_idx') + 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 set(np.unique(kmers_1['assembly_idx']).tolist()) == {0, 1, 2, 3} + assert np.array_equal(record_offsets_1, np.array([0, 1, 2, 3, 4], dtype=np.uint64)) + 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) assert edges_1.ndim == 1 @@ -110,19 +111,73 @@ def test_build_threading_equivalence(targets_dir, non_targets_dir) -> None: assert record_ids_1 == record_ids_2 assert record_ids_1 == record_ids_many assert len(record_ids_1) == len(assembly_paths) + assert np.array_equal(record_offsets_1, record_offsets_2) + assert np.array_equal(record_offsets_1, record_offsets_many) assert np.array_equal(_sorted_edges(edges_1), _sorted_edges(edges_2)) assert np.array_equal(_sorted_edges(edges_1), _sorted_edges(edges_many)) +def test_multi_thread_record_offsets_and_global_record_indices(tmp_path: Path) -> None: + def write_fasta(path: Path, n_records: int) -> None: + seq = 'ACGT' * 20 + path.write_text(''.join(f'>r{i}\n{seq}\n' for i in range(n_records))) + + assembly_paths = [] + for i, n_records in enumerate([2, 1, 3, 1]): + path = tmp_path / f'a{i}.fasta' + write_fasta(path, n_records) + assembly_paths.append(path) + + kmers, _, _, _, record_offsets, record_ids = build( + assembly_paths, + kmerlen=7, + windowsize=10, + is_targets=[True, True, False, False], + n_cpu=2, + ) + + assert [len(ids) for ids in record_ids] == [2, 1, 3, 1] + assert np.array_equal(record_offsets, np.array([0, 2, 3, 6, 7], dtype=np.uint64)) + 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.uint64),) + record_offsets = np.array([0, 2, 5, 6], dtype=np.uint64) + + _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, 0), - (11, 0, 0), - (20, 0, 1), - (30, 1, 1), - (31, 1, 1), - (32, 1, 1), + (10, 0), + (11, 0), + (20, 1), + (30, 2), + (31, 2), + (32, 2), ], dtype=KMER_DTYPE) idx = np.array([100, 101, 200, 300, 301, 302], dtype=np.uint64) nodes = np.array([ @@ -138,11 +193,11 @@ def test_filter_kmers() -> None: assert np.array_equal(nodes_new['stop'], np.array([2, 5], dtype=np.uint64)) expected_kmers = np.array([ - (10, 0, 0), - (11, 0, 0), - (30, 1, 1), - (31, 1, 1), - (32, 1, 1), + (10, 0), + (11, 0), + (30, 2), + (31, 2), + (32, 2), ], dtype=KMER_DTYPE) expected_idx = np.array([100, 101, 300, 301, 302], dtype=np.uint64) diff --git a/tests/smoke/test_outputs.py b/tests/smoke/test_outputs.py index c146b61..9f680c1 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'} + required_keys = {'kmers', 'idx', '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)}'