diff --git a/cpp/include/seqwin/build.hpp b/cpp/include/seqwin/build.hpp index f3218cd..2a23250 100644 --- a/cpp/include/seqwin/build.hpp +++ b/cpp/include/seqwin/build.hpp @@ -11,12 +11,9 @@ namespace seqwin { /** * @brief Build a minimizer graph from assembly FASTA files. * - * `assembly_paths` and `is_targets` are parallel lists. - * * @param assembly_paths Paths to input assemblies in FASTA format (plain or gzipped). * @param kmerlen K-mer length for minimizer sketch. * @param windowsize Window size for minimizer sketch. - * @param is_targets Whether each assembly is a target assembly. * @param n_cpu Number of worker threads to use. * @param low_memory Recompute minimizers in a second pass to reduce peak memory. * @return Minimizer graph. @@ -26,7 +23,6 @@ Graph build( const std::vector& assembly_paths, std::size_t kmerlen, std::size_t windowsize, - const std::vector& is_targets, std::size_t n_cpu = 1, bool low_memory = false ); diff --git a/cpp/include/seqwin/filter.hpp b/cpp/include/seqwin/filter.hpp index b2a4405..12b0b3c 100644 --- a/cpp/include/seqwin/filter.hpp +++ b/cpp/include/seqwin/filter.hpp @@ -8,6 +8,17 @@ namespace seqwin { +void get_penalty( + const Kmer* kmers, + Node* nodes, + std::size_t n_nodes, + const std::size_t* record_offsets, + std::size_t n_record_offsets, + const bool* is_targets, + std::size_t n_assemblies, + std::size_t n_cpu = 1 +); + Graph filter_kmers( const Kmer* kmers, const Node* nodes, diff --git a/cpp/include/seqwin/graph.hpp b/cpp/include/seqwin/graph.hpp index c0ec30f..86b1132 100644 --- a/cpp/include/seqwin/graph.hpp +++ b/cpp/include/seqwin/graph.hpp @@ -32,11 +32,11 @@ struct Node { std::size_t start; /** End of the half-open range for this node's minimizer entries. */ std::size_t stop; - /** Number of target assemblies containing this minimizer hash. */ - std::uint32_t n_tar; - /** Number of non-target assemblies containing this minimizer hash. */ - std::uint32_t n_neg; - /** Placeholder for node penalty score, for downstream graph filtering. */ + /** Node scoring placeholder. */ + std::uint32_t n_tar = 0; + /** Node scoring placeholder. */ + std::uint32_t n_neg = 0; + /** Node scoring placeholder. */ double penalty = 0.0; }; @@ -56,7 +56,10 @@ struct Edge { * @brief Container for the minimizer graph returned by `build()`. */ struct Graph { - /** Minimizer occurrences in all assemblies, grouped and sorted by hash. */ + /** + * Minimizer occurrences in all assemblies, grouped and sorted by hash. + * `record_idx` is nondecreasing within each node range. + */ NoInitArray kmers; /** Sorted by hash. */ NoInitArray nodes; diff --git a/cpp/src/bindings/python_bindings.cpp b/cpp/src/bindings/python_bindings.cpp index db443e6..2c956ee 100644 --- a/cpp/src/bindings/python_bindings.cpp +++ b/cpp/src/bindings/python_bindings.cpp @@ -1,5 +1,6 @@ #include #include +#include #include #include #include @@ -50,7 +51,6 @@ PYBIND11_MODULE(_core, m) { [](const std::vector& assembly_paths, std::size_t kmerlen, std::size_t windowsize, - const std::vector& is_targets, std::size_t n_cpu, bool low_memory ) { @@ -61,7 +61,6 @@ PYBIND11_MODULE(_core, m) { assembly_paths, kmerlen, windowsize, - is_targets, n_cpu, low_memory ); @@ -86,11 +85,55 @@ PYBIND11_MODULE(_core, m) { py::arg("assembly_paths"), py::arg("kmerlen"), py::arg("windowsize"), - py::arg("is_targets"), py::arg("n_cpu") = 1, py::arg("low_memory") = false ); + m.def("_get_penalty_native", + [](py::array_t kmers, + py::array_t nodes, + py::array_t record_offsets, + py::array_t is_targets, + std::size_t n_cpu + ) { + auto kmers_buf = kmers.request(); + auto nodes_buf = nodes.request(); + auto record_offsets_buf = record_offsets.request(); + auto is_targets_buf = is_targets.request(); + + if (nodes_buf.readonly) { + throw std::invalid_argument("nodes must be writable"); + } + + const auto* kmers_ptr = static_cast(kmers_buf.ptr); + auto* nodes_ptr = static_cast(nodes_buf.ptr); + const auto* record_offsets_ptr = static_cast(record_offsets_buf.ptr); + const auto* is_targets_ptr = static_cast(is_targets_buf.ptr); + const auto n_nodes = static_cast(nodes_buf.shape[0]); + const auto n_record_offsets = static_cast(record_offsets_buf.shape[0]); + const auto n_assemblies = static_cast(is_targets_buf.shape[0]); + + { + py::gil_scoped_release release; + seqwin::get_penalty( + kmers_ptr, + nodes_ptr, + n_nodes, + record_offsets_ptr, + n_record_offsets, + is_targets_ptr, + n_assemblies, + n_cpu + ); + } + }, + py::arg("kmers").noconvert(), + py::arg("nodes").noconvert(), + py::arg("record_offsets").noconvert(), + py::arg("is_targets").noconvert(), + py::arg("n_cpu") = 1 + ); + m.def("_filter_kmers_native", [](py::array_t kmers, py::array_t nodes, diff --git a/cpp/src/seqwin/build.cpp b/cpp/src/seqwin/build.cpp index 622d7c5..0515f85 100644 --- a/cpp/src/seqwin/build.cpp +++ b/cpp/src/seqwin/build.cpp @@ -61,14 +61,6 @@ struct RawKmer { Kmer kmer; }; -struct NodeState { - std::size_t count = 0; - std::size_t cursor = 0; - std::uint32_t n_tar = 0; - std::uint32_t n_neg = 0; - std::size_t last_seen_assembly = std::numeric_limits::max(); -}; - using EdgeKey = std::pair; struct EdgeKeyHash { @@ -80,8 +72,8 @@ struct EdgeKeyHash { }; struct EdgeState { - std::size_t weight = 0; - std::size_t last_seen_assembly = std::numeric_limits::max(); + std::uint32_t weight = 0; + std::uint32_t last_seen_assembly = std::numeric_limits::max(); }; /** @@ -95,7 +87,6 @@ ThreadGraph build_worker( const std::vector& assembly_paths, std::size_t kmerlen, std::size_t windowsize, - const std::vector& is_targets, std::size_t start_assembly, std::size_t end_assembly, std::size_t thread_id, @@ -121,15 +112,14 @@ ThreadGraph build_worker( graph.end_assembly = end_assembly; // Reserving for unordered_map will actually allocate physical memory - std::unordered_map node_map; + std::unordered_map node_map; std::unordered_map edge_map; const auto n_map_entries_est = n_minimizers_est / map_reserve_divisor; node_map.reserve(n_map_entries_est); edge_map.reserve(n_map_entries_est); for (std::size_t assembly_i = start_assembly; assembly_i < end_assembly; ++assembly_i) { - const bool is_target = is_targets[assembly_i]; - + const auto assembly_i_u32 = static_cast(assembly_i); auto records = read_fasta(assembly_paths[assembly_i]); std::vector record_ids; record_ids.reserve(records.size()); @@ -159,16 +149,8 @@ ThreadGraph build_worker( } // Add this minimizer to an existing node, or create a new node - auto [node_it, node_inserted] = node_map.try_emplace(m.out_hash); - ++(node_it->second.count); - if (node_inserted || node_it->second.last_seen_assembly != assembly_i) { - if (is_target) { - ++(node_it->second.n_tar); - } else { - ++(node_it->second.n_neg); - } - node_it->second.last_seen_assembly = assembly_i; - } + auto node_it = node_map.try_emplace(m.out_hash, 0).first; + ++node_it->second; ++graph.n_kmers; } @@ -186,9 +168,9 @@ ThreadGraph build_worker( } const EdgeKey key{u, v}; auto edge_it = edge_map.try_emplace(key).first; - if (edge_it->second.last_seen_assembly != assembly_i) { - ++(edge_it->second.weight); - edge_it->second.last_seen_assembly = assembly_i; + if (edge_it->second.last_seen_assembly != assembly_i_u32) { + ++edge_it->second.weight; + edge_it->second.last_seen_assembly = assembly_i_u32; } } } @@ -204,37 +186,57 @@ ThreadGraph build_worker( } std::unordered_map().swap(edge_map); + graph.n_nodes = node_map.size(); + graph.nodes = NoInitArray(graph.n_nodes); + std::size_t node_i = 0; + if (!low_memory) { // Build ThreadGraph.kmers (grouped by hash) graph.kmers = NoInitArray(graph.n_kmers); + + // node_map values are counts while assigning node ranges, + // then cursors while scattering raw k-mers into those ranges std::size_t cursor = 0; - for (auto& [hash, state] : node_map) { - (void)hash; - state.cursor = cursor; - cursor += state.count; + for (auto& [hash, node_val] : node_map) { + const std::size_t count = node_val; + const std::size_t start = cursor; + + graph.nodes[node_i++] = ThreadNode{ + hash, + start, + count, + thread_id + }; + + node_val = start; + cursor += count; } + if (cursor != graph.n_kmers) { + throw std::logic_error("Node k-mer counts do not sum to graph.n_kmers"); + } + for (const auto& rk : raw_kmers) { auto node_it = node_map.find(rk.hash); - graph.kmers[node_it->second.cursor++] = rk.kmer; + if (node_it == node_map.end()) { + throw std::logic_error( + "Standard-mode k-mer scattering found an unknown minimizer hash" + ); + } + graph.kmers[node_it->second++] = rk.kmer; } std::vector().swap(raw_kmers); + } else { + // In low-memory mode node_map values remain counts + for (const auto& [hash, count] : node_map) { + graph.nodes[node_i++] = ThreadNode{ + hash, + 0, + count, + thread_id + }; + } } - - graph.n_nodes = node_map.size(); - graph.nodes = NoInitArray(node_map.size()); - std::size_t node_i = 0; - for (const auto& [hash, state] : node_map) { - const std::size_t start = low_memory ? 0 : state.cursor - state.count; - - graph.nodes[node_i++] = ThreadNode{ - hash, - start, - state.count, - state.n_tar, - state.n_neg, - thread_id - }; - } + std::unordered_map().swap(node_map); return graph; } @@ -313,13 +315,9 @@ Graph build( const std::vector& assembly_paths, std::size_t kmerlen, std::size_t windowsize, - const std::vector& is_targets, std::size_t n_cpu, bool low_memory ) { - if (assembly_paths.size() != is_targets.size()) { - throw std::runtime_error("assembly_paths and is_targets must have the same length"); - } if (assembly_paths.size() > std::numeric_limits::max()) { throw std::runtime_error("Number of input assemblies exceeds uint32 range"); } @@ -344,7 +342,6 @@ Graph build( assembly_paths, kmerlen, windowsize, - is_targets, start_assembly, end_assembly, thread_id, diff --git a/cpp/src/seqwin/build_internals.cpp b/cpp/src/seqwin/build_internals.cpp index 27cf0e7..d5086c1 100644 --- a/cpp/src/seqwin/build_internals.cpp +++ b/cpp/src/seqwin/build_internals.cpp @@ -198,13 +198,9 @@ static MergedNodes merge_nodes( i = 0; while (i < n_nodes) { const auto hash = nodes[i].hash; - std::uint32_t n_tar = 0; - std::uint32_t n_neg = 0; const auto start = n_kmers; while (i < n_nodes && nodes[i].hash == hash) { - n_tar += nodes[i].n_tar; - n_neg += nodes[i].n_neg; const auto count = nodes[i].count; if (low_memory) { @@ -221,7 +217,7 @@ static MergedNodes merge_nodes( ++i; } - merged.nodes[write_i++] = Node{hash, start, n_kmers, n_tar, n_neg}; + merged.nodes[write_i++] = Node{hash, start, n_kmers}; } return merged; } diff --git a/cpp/src/seqwin/build_internals.hpp b/cpp/src/seqwin/build_internals.hpp index aea68ce..91232f9 100644 --- a/cpp/src/seqwin/build_internals.hpp +++ b/cpp/src/seqwin/build_internals.hpp @@ -25,10 +25,6 @@ struct ThreadNode { std::size_t start; /** Number of entries in `ThreadGraph.kmers` for this node. */ std::size_t count; - /** Number of target assemblies containing this minimizer hash. */ - std::uint32_t n_tar; - /** Number of non-target assemblies containing this minimizer hash. */ - std::uint32_t n_neg; /** Worker thread that produced this node. Used by the merging step. */ std::size_t thread_id; }; diff --git a/cpp/src/seqwin/filter.cpp b/cpp/src/seqwin/filter.cpp index fbbc0fa..aa7707a 100644 --- a/cpp/src/seqwin/filter.cpp +++ b/cpp/src/seqwin/filter.cpp @@ -1,12 +1,128 @@ #include "seqwin/filter.hpp" #include +#include #include #include +#include +#include #include +#include "utils/thread_pool.hpp" + namespace seqwin { +void get_penalty( + const Kmer* kmers, + Node* nodes, + std::size_t n_nodes, + const std::size_t* record_offsets, + std::size_t n_record_offsets, + const bool* is_targets, + std::size_t n_assemblies, + std::size_t n_cpu +) { + if (n_record_offsets != n_assemblies + 1) { + throw std::invalid_argument("len(record_offsets) must equal len(is_targets) + 1"); + } + if (n_record_offsets == 0 || record_offsets[0] != 0) { + throw std::invalid_argument("record_offsets must start with 0"); + } + if (n_assemblies > std::numeric_limits::max()) { + throw std::invalid_argument("Number of assemblies exceeds uint32 range"); + } + + std::size_t total_targets = 0; + std::size_t total_non_targets = 0; + for (std::size_t i = 0; i < n_assemblies; ++i) { + if (record_offsets[i + 1] < record_offsets[i]) { + throw std::invalid_argument("record_offsets must be nondecreasing"); + } + if (is_targets[i]) { + ++total_targets; + } else { + ++total_non_targets; + } + } + if (total_targets == 0) { + throw std::invalid_argument("is_targets must contain at least one target assembly"); + } + if (total_non_targets == 0) { + throw std::invalid_argument("is_targets must contain at least one non-target assembly"); + } + + std::size_t n_workers = std::max(1, n_cpu); + if (n_nodes > 0) { + n_workers = std::min(n_workers, n_nodes); + } + internal::ThreadPool pool(n_workers); + + pool.parallel_for(n_nodes, [&](std::size_t start, std::size_t end, std::size_t) { + for (std::size_t node_i = start; node_i < end; ++node_i) { + auto& node = nodes[node_i]; + if (node.start == node.stop) { + node.n_tar = 0; + node.n_neg = 0; + node.penalty = 1.0; + continue; + } + + std::uint32_t n_tar = 0; + std::uint32_t n_neg = 0; + + // Monotonic scan of record_idx and record_offsets + // Each node range has nondecreasing record_idx values, so one upper_bound() + // maps the first record to its assembly and the scan only advances forward + std::uint32_t previous_record_idx = kmers[node.start].record_idx; + if (static_cast(previous_record_idx) >= record_offsets[n_assemblies]) { + throw std::invalid_argument("record_idx is outside record_offsets range"); + } + const auto* offset_it = std::upper_bound( + record_offsets, + record_offsets + n_record_offsets, + static_cast(previous_record_idx) + ); + std::size_t assembly_idx = static_cast(offset_it - record_offsets - 1); + // Count each assembly once per node + std::size_t last_counted_assembly = std::numeric_limits::max(); + + for (std::size_t kmer_i = node.start; kmer_i < node.stop; ++kmer_i) { + const std::uint32_t record_idx_u32 = kmers[kmer_i].record_idx; + if (record_idx_u32 < previous_record_idx) { + throw std::invalid_argument("record_idx must be nondecreasing within each node range"); + } + previous_record_idx = record_idx_u32; + + const std::size_t record_idx = static_cast(record_idx_u32); + if (record_idx >= record_offsets[n_assemblies]) { + throw std::invalid_argument("record_idx is outside record_offsets range"); + } + // Duplicate record offsets are allowed for zero-record assemblies + while ( + assembly_idx + 1 < n_assemblies && + record_idx >= record_offsets[assembly_idx + 1] + ) { + ++assembly_idx; + } + if (assembly_idx != last_counted_assembly) { + if (is_targets[assembly_idx]) { + ++n_tar; + } else { + ++n_neg; + } + last_counted_assembly = assembly_idx; + } + } + + node.n_tar = n_tar; + node.n_neg = n_neg; + const double frac_tar = static_cast(n_tar) / static_cast(total_targets); + const double frac_neg = static_cast(n_neg) / static_cast(total_non_targets); + node.penalty = std::hypot(1.0 - frac_tar, frac_neg); + } + }); +} + Graph filter_kmers( const Kmer* kmers, const Node* nodes, diff --git a/src/seqwin/graph/__init__.py b/src/seqwin/graph/__init__.py index 12e1db8..b1aa47d 100644 --- a/src/seqwin/graph/__init__.py +++ b/src/seqwin/graph/__init__.py @@ -33,7 +33,7 @@ import numpy as np from numpy.typing import NDArray -from ._core import _build_native, _filter_kmers_native +from ._core import _build_native, _get_penalty_native, _filter_kmers_native from .utils import OrderedKmers @@ -62,7 +62,6 @@ def build( assembly_paths: Iterable[Path], kmerlen: int, windowsize: int, - is_targets: Iterable[bool], low_memory: bool = False, n_cpu: int = 1 ) -> tuple[ @@ -81,12 +80,10 @@ def build( >>> assembly_paths = ..., >>> kmerlen = 21, >>> windowsize = 200, - >>> is_targets = ..., >>> n_cpu = 4, >>> low_memory = False >>> ) ``` - - `assembly_paths` and `is_targets` are parallel lists. - `kmers` stores minimizer occurrences in all assemblies, grouped and sorted by hash. - `nodes` and `edges` are sorted by hash. @@ -111,7 +108,6 @@ def build( 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. - 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] @@ -126,9 +122,9 @@ def build( - '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. - 'stop' (uintp): End of the half-open range for this node's minimizer entries. - - '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. + - 'n_tar' (uint32): Node scoring placeholder initialized to 0. + - 'n_neg' (uint32): Node scoring placeholder initialized to 0. + - 'penalty' (float64): Node scoring placeholder initialized to 0.0. 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. @@ -141,12 +137,36 @@ def build( list(str(p) for p in assembly_paths), int(kmerlen), int(windowsize), - list(bool(t) for t in is_targets), int(n_cpu), bool(low_memory) ) +def _get_penalty( + kmers: NDArray[np.void], + nodes: NDArray[np.void], + record_offsets: NDArray[np.uintp], + is_targets: Iterable[bool], + n_cpu: int = 1 +) -> None: + """Populate node target counts and penalty scores in place. + + Args: + kmers (NDArray): See `KmerGraph.kmers`. + nodes (NDArray): See `KmerGraph.nodes`. + record_offsets (NDArray[np.uintp]): See `KmerGraph.record_offsets`. + is_targets (Iterable[bool]): Whether each assembly is a target assembly. + n_cpu (int, optional): Number of worker threads to use. [1] + """ + _get_penalty_native( + kmers, + nodes, + record_offsets, + np.asarray(is_targets, dtype=np.bool_, order='C'), + int(n_cpu) + ) + + def _filter_kmers( kmers: NDArray[np.void], nodes: NDArray[np.void], diff --git a/src/seqwin/kmers.py b/src/seqwin/kmers.py index 55e6594..6868f88 100644 --- a/src/seqwin/kmers.py +++ b/src/seqwin/kmers.py @@ -36,7 +36,7 @@ import networkx as nx from numpy.typing import NDArray -from .graph import build, _filter_kmers +from .graph import build, _get_penalty, _filter_kmers from .assemblies import Assemblies from .helpers import get_subgraphs from .utils import print_time_delta @@ -90,16 +90,16 @@ def __init__(self, assemblies: Assemblies, kmerlen: int, windowsize: int, n_cpu: assemblies.path, kmerlen, windowsize, - assemblies.is_target, n_cpu=n_cpu, low_memory=low_memory ) # calculate penalty for each node - n_tar = sum(assemblies.is_target) - n_neg = n_assemblies - n_tar - nodes['penalty'] = _frac_to_penalty( - nodes['n_tar'] / n_tar, - nodes['n_neg'] / n_neg + _get_penalty( + kmers, + nodes, + record_offsets, + assemblies.is_target, + n_cpu=n_cpu ) assemblies.record_ids = record_ids diff --git a/tests/smoke/test_graph.py b/tests/smoke/test_graph.py index 27eec8c..629fb6b 100644 --- a/tests/smoke/test_graph.py +++ b/tests/smoke/test_graph.py @@ -2,7 +2,7 @@ import numpy as np -from seqwin.graph import KMER_DTYPE, NODE_DTYPE, EDGE_DTYPE, build, _filter_kmers +from seqwin.graph import KMER_DTYPE, NODE_DTYPE, EDGE_DTYPE, build, _get_penalty, _filter_kmers def _sorted_edges(edges: np.ndarray) -> np.ndarray: @@ -30,6 +30,7 @@ def _assert_node_ranges(kmers: np.ndarray, nodes: np.ndarray) -> None: stop = int(node['stop']) assert 0 <= start <= stop <= len(kmers) assert len(kmers[start:stop]) == (stop - start) + assert np.all(kmers[start:stop]['record_idx'][:-1] <= kmers[start:stop]['record_idx'][1:]) total += stop - start assert total == len(kmers) @@ -63,27 +64,22 @@ 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', ] - is_targets = [True, True, False, False] - 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, 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, nodes_many, edges_many, record_offsets_many, record_ids_many = build( assembly_paths, kmerlen=7, windowsize=10, - is_targets=is_targets, n_cpu=99, ) @@ -92,7 +88,9 @@ def test_build_threading_equivalence(targets_dir, non_targets_dir) -> None: assert kmers_1.dtype.names == ('pos', 'record_idx') 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) + assert np.all(nodes_1['n_tar'] == 0) + assert np.all(nodes_1['n_neg'] == 0) + assert np.all(nodes_1['penalty'] == 0.0) assert edges_1.ndim == 1 assert edges_1.dtype == EDGE_DTYPE @@ -121,6 +119,20 @@ def test_build_threading_equivalence(targets_dir, non_targets_dir) -> None: assert np.array_equal(_sorted_edges(edges_1), _sorted_edges(edges_many)) +def test_build_rejects_is_targets_argument(targets_dir, non_targets_dir) -> None: + assembly_paths = [ + targets_dir / 'target-1.fasta', + non_targets_dir / 'non-target-1.fasta', + ] + with np.testing.assert_raises(TypeError): + build( + assembly_paths, + kmerlen=7, + windowsize=10, + is_targets=[True, False], + ) + + 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 @@ -136,7 +148,6 @@ def write_fasta(path: Path, n_records: int) -> None: assembly_paths, kmerlen=7, windowsize=10, - is_targets=[True, True, False, False], n_cpu=2, ) @@ -184,14 +195,11 @@ def test_low_memory_build_matches_standard(targets_dir, non_targets_dir) -> None 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, ) @@ -199,9 +207,87 @@ def test_low_memory_build_matches_standard(targets_dir, non_targets_dir) -> None assembly_paths, kmerlen=7, windowsize=10, - is_targets=is_targets, n_cpu=n_cpu, low_memory=True, ) _assert_graph_outputs_equal(standard, low_memory) + + +def _synthetic_penalty_inputs() -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + kmers = np.array([ + (0, 0), (1, 0), (2, 1), (3, 2), (4, 4), # node 10: asm 0, 1, 2 => tar 2 neg 1 + (5, 2), (6, 3), # node 20: asm 1 => neg 1 + (7, 5), (8, 6), # node 30: asm 3 => neg 1 + (9, 4), # node 40: asm 2 => tar 1 + ], dtype=KMER_DTYPE) + nodes = np.array([ + (10, 0, 5, 0, 0, 0.0), + (20, 5, 7, 0, 0, 0.0), + (30, 7, 9, 0, 0, 0.0), + (40, 9, 10, 0, 0, 0.0), + (50, 10, 10, 9, 9, 9.0), + (60, 5, 9, 0, 0, 0.0), + ], dtype=NODE_DTYPE) + record_offsets = np.array([0, 2, 4, 5, 7], dtype=np.uintp) + is_targets = np.array([True, False, True, False], dtype=np.bool_) + return kmers, nodes, record_offsets, is_targets + + +def test_get_penalty_exact_scoring() -> None: + kmers, nodes, record_offsets, is_targets = _synthetic_penalty_inputs() + + assert _get_penalty(kmers, nodes, record_offsets, is_targets, n_cpu=1) is None + + assert np.array_equal(nodes['n_tar'], np.array([2, 0, 0, 1, 0, 0], dtype=np.uint32)) + assert np.array_equal(nodes['n_neg'], np.array([1, 1, 1, 0, 0, 2], dtype=np.uint32)) + np.testing.assert_allclose( + nodes['penalty'], + np.array([0.5, np.hypot(1.0, 0.5), np.hypot(1.0, 0.5), 0.5, 1.0, np.sqrt(2.0)]) + ) + + +def test_get_penalty_parallel_equivalence() -> None: + kmers, nodes_1, record_offsets, is_targets = _synthetic_penalty_inputs() + nodes_2 = nodes_1.copy() + nodes_many = nodes_1.copy() + + _get_penalty(kmers, nodes_1, record_offsets, is_targets, n_cpu=1) + _get_penalty(kmers, nodes_2, record_offsets, is_targets, n_cpu=2) + _get_penalty(kmers, nodes_many, record_offsets, is_targets, n_cpu=99) + + assert np.array_equal(nodes_1, nodes_2) + assert np.array_equal(nodes_1, nodes_many) + + +def test_get_penalty_validation() -> None: + kmers, nodes, record_offsets, is_targets = _synthetic_penalty_inputs() + + read_only = nodes.copy() + read_only.flags.writeable = False + with np.testing.assert_raises(ValueError): + _get_penalty(kmers, read_only, record_offsets, is_targets) + with np.testing.assert_raises(ValueError): + _get_penalty(kmers, nodes.copy(), record_offsets[:-1], is_targets) + with np.testing.assert_raises(ValueError): + _get_penalty(kmers, nodes.copy(), np.array([1, 2, 4, 5, 7], dtype=np.uintp), is_targets) + with np.testing.assert_raises(ValueError): + _get_penalty(kmers, nodes.copy(), np.array([0, 3, 2, 5, 7], dtype=np.uintp), is_targets) + with np.testing.assert_raises(ValueError): + _get_penalty(kmers, nodes.copy(), record_offsets, [False, False, False, False]) + with np.testing.assert_raises(ValueError): + _get_penalty(kmers, nodes.copy(), record_offsets, [True, True, True, True]) + bad_nodes = nodes.copy() + bad_nodes[0]['stop'] = len(kmers) + 1 + with np.testing.assert_raises(ValueError): + _get_penalty(kmers, bad_nodes, record_offsets, is_targets) + bad_kmers = kmers.copy() + bad_kmers[0]['record_idx'] = 7 + with np.testing.assert_raises(ValueError): + _get_penalty(bad_kmers, nodes.copy(), record_offsets, is_targets) + descending = kmers.copy() + descending[3]['record_idx'] = 0 + with np.testing.assert_raises(ValueError): + _get_penalty(descending, nodes.copy(), record_offsets, is_targets) + with np.testing.assert_raises(ValueError): + _get_penalty(kmers, nodes.copy(), record_offsets, np.array([[True, False], [True, False]]))