diff --git a/RELEASE_NOTES.md b/RELEASE_NOTES.md index 7443e53..2d76ef4 100644 --- a/RELEASE_NOTES.md +++ b/RELEASE_NOTES.md @@ -5,10 +5,22 @@ ### Added - Add left_multiplication and right_multiplication to ThreadingInstructions (#99) +- OpenMP parallelism for Li-Stephens Viterbi and hets phases across targets (#127) +- NumPy batch API for zero-copy genotype transfer from Python to C++ (#127) ### Changed - Build wheels on macOS 14 for arm64 and macOS 15 for x86_64 (#108) +- Replace hash map traceback storage with deque arena in Viterbi (#127) +- Replace hash map state lookup with flat vectors in ThreadsLowMem (#127) +- Boolean array neighbor search in Matcher replaces red-black tree (#127) + +### Fixed + +- Memory leak: raw HMM pointer replaced with unique_ptr (#127) +- Sign-extension in pair_key/coord_id_key hash functions (#127) +- Bounds check after array access in ThreadsFastLS (#127) +- exit(1) calls replaced with exceptions for proper Python error propagation (#127) ## [0.2.1] - 2025-06-03 diff --git a/pyproject.toml b/pyproject.toml index 00935f6..7883b28 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,16 +16,9 @@ authors = [ requires-python = ">=3.9" dependencies = [ - "click", - "xarray", - "h5py", - "pandas", - "numba", "numpy", - "tszip", "arg-needle-lib==1.2.1", "cyvcf2", - "ray", "pgenlib", "tqdm" ] @@ -43,7 +36,24 @@ classifiers = [ [project.optional-dependencies] dev = [ - "pytest" + "pytest", + "h5py" +] +convert = [ + "tszip" +] +impute = [ + "pandas", + "scipy" +] +normalize = [ + "msprime" +] +all = [ + "tszip", + "msprime", + "pandas", + "scipy" ] [project.scripts] diff --git a/src/AlleleAges.cpp b/src/AlleleAges.cpp index ed7ca29..1af27fc 100644 --- a/src/AlleleAges.cpp +++ b/src/AlleleAges.cpp @@ -15,14 +15,11 @@ // along with this program. If not, see . #include "AlleleAges.hpp" +#include "GenotypeIterator.hpp" -#include -#include +#include +#include #include -#include -#include - -#include AgeEstimator::AgeEstimator(const ThreadingInstructions& instructions) { num_samples = instructions.num_samples; @@ -55,7 +52,12 @@ void AgeEstimator::process_site(const std::vector& genotypes) { } else { if (genotypes.at(i) == 1) { int target = threading_iterators.at(i).current_target; - path_lengths[i] = path_lengths[target] + 1; + // Self-referencing targets don't extend carrier chains + if (static_cast(target) != i) { + path_lengths[i] = path_lengths[target] + 1; + } else { + path_lengths[i] = 1; + } } else { path_lengths[i] = 0; } @@ -84,64 +86,94 @@ void AgeEstimator::process_site(const std::vector& genotypes) { size_t start_tmp = path_start; while (start_tmp > 0) { int next_sample = threading_iterators.at(start_tmp).current_target; + // Skip self-referencing targets to avoid infinite loops + if (next_sample == static_cast(start_tmp)) { + break; + } double this_tmrca = threading_iterators.at(start_tmp).current_tmrca; running_max = std::max(running_max, this_tmrca); tmrcas.at(next_sample) = running_max; start_tmp = next_sample; } - if (start_tmp != 0) { - throw std::runtime_error("Invalid threading instruction traversal."); - } for (int i = 0; i < num_samples; i++) { if (tmrcas.at(i) < 0) { int target = threading_iterators.at(i).current_target; double tmrca = threading_iterators.at(i).current_tmrca; - tmrcas.at(i) = std::max(tmrcas.at(target), tmrca); + if (target == i || target < 0 || tmrcas.at(target) < 0) { + // Self-ref or unfilled target: use own tmrca + tmrcas.at(i) = tmrca; + } else { + tmrcas.at(i) = std::max(tmrcas.at(target), tmrca); + } } } - // Create a sorted unique list of coalescence times as transform_reduce - // below must be done in order. For performance, boost's flat_set is - // faster than std::set or sorting a std::vector in this instance. - boost::container::flat_set unique_tmrcas(tmrcas.begin(), tmrcas.end()); - - // For each sample, check its tmrca with path_start and - // update the score for each tmrca bin accordingly - std::map scores; - for (double t : unique_tmrcas) { - scores[t] = std::transform_reduce( - tmrcas.begin(), - tmrcas.end(), - genotypes.begin(), - 0, - std::plus<>(), - [t](double tmrca, int genotype) { - bool is_carrier = genotype > 0; - if (is_carrier && tmrca <= t) { - return 1; - } - if (!is_carrier && tmrca > t) { - return 1; - } - return 0; - } - ); + // Sort samples by tmrca, then sweep to find the threshold that + // maximizes: carriers_at_or_below(t) + non_carriers_above(t). + struct TmrcaSample { + double tmrca; + int genotype; + }; + std::vector sorted_samples; + sorted_samples.reserve(num_samples); + int total_non_carriers = 0; + for (int i = 0; i < num_samples; i++) { + sorted_samples.push_back({tmrcas[i], genotypes[i]}); + if (genotypes[i] == 0) total_non_carriers++; } - - std::vector age_bin_boundaries; - for (auto const& imap: scores) - age_bin_boundaries.push_back(imap.first); - std::vector age_bins; - - for (size_t k = 0; k < age_bin_boundaries.size(); k++) { - int score = scores.at(age_bin_boundaries.at(k)); - if (score > max_score) { - max_score = score; - allele_age = (k == age_bin_boundaries.size() - 1) - ? age_bin_boundaries.at(k) + 1 - : (age_bin_boundaries.at(k) + age_bin_boundaries.at(k + 1)) / 2.; + // NaN-safe sort: put NaN values last + std::sort(sorted_samples.begin(), sorted_samples.end(), + [](const TmrcaSample& a, const TmrcaSample& b) { + if (std::isnan(a.tmrca)) return false; + if (std::isnan(b.tmrca)) return true; + return a.tmrca < b.tmrca; + }); + + // Initial score: threshold below all tmrcas → 0 carriers correct, + // all non-carriers correct (they're all above threshold). + int score = total_non_carriers; + int best_score = score; + double best_boundary = sorted_samples.front().tmrca; + double next_boundary = sorted_samples.front().tmrca; + + // Sweep through sorted samples in groups of equal tmrca + size_t i_sweep = 0; + size_t n_sorted = sorted_samples.size(); + while (i_sweep < n_sorted) { + double current_t = sorted_samples[i_sweep].tmrca; + // Stop at NaN values (they are sorted to the end) + if (std::isnan(current_t)) break; + // Process all samples at this tmrca + int carriers_at_t = 0; + int non_carriers_at_t = 0; + size_t group_end = i_sweep; + while (group_end < n_sorted && sorted_samples[group_end].tmrca == current_t) { + if (sorted_samples[group_end].genotype > 0) + carriers_at_t++; + else + non_carriers_at_t++; + group_end++; + } + // Moving threshold to include this group: + // carriers become correctly classified (+), non-carriers become incorrect (-) + score += carriers_at_t - non_carriers_at_t; + + if (score > best_score) { + best_score = score; + best_boundary = current_t; + next_boundary = (group_end < n_sorted) + ? sorted_samples[group_end].tmrca + : current_t; } + i_sweep = group_end; + } + + if (best_score > max_score) { + max_score = best_score; + allele_age = (best_boundary == next_boundary) + ? best_boundary + 1 + : (best_boundary + next_boundary) / 2.; } } @@ -152,3 +184,12 @@ void AgeEstimator::process_site(const std::vector& genotypes) { std::vector AgeEstimator::get_inferred_ages() const { return estimated_ages; } + +std::vector estimate_ages(const ThreadingInstructions& instructions) { + GenotypeIterator gt_it(instructions); + AgeEstimator estimator(instructions); + while (gt_it.has_next_genotype()) { + estimator.process_site(gt_it.next_genotype()); + } + return estimator.get_inferred_ages(); +} diff --git a/src/AlleleAges.hpp b/src/AlleleAges.hpp index 86736b0..858f5e5 100644 --- a/src/AlleleAges.hpp +++ b/src/AlleleAges.hpp @@ -37,4 +37,8 @@ class AgeEstimator { std::vector estimated_ages; }; +// Bulk function: creates GenotypeIterator + AgeEstimator internally, +// processes all sites in C++, returns estimated ages. +std::vector estimate_ages(const ThreadingInstructions& instructions); + #endif // THREADS_ARG_ALLELE_AGES_HPP \ No newline at end of file diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 55931b1..9431dfd 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -18,6 +18,18 @@ find_package(Boost REQUIRED) message(STATUS "Found Boost ${Boost_VERSION}") +# Find HDF5 C library (for .threads file I/O) +find_package(HDF5 REQUIRED COMPONENTS C) +message(STATUS "Found HDF5 ${HDF5_VERSION}") + +# Optional OpenMP for parallel Viterbi across targets +find_package(OpenMP) +if(OpenMP_CXX_FOUND) + message(STATUS "OpenMP found — parallel Viterbi enabled") +else() + message(STATUS "OpenMP not found — building single-threaded") +endif() + # Threads static library set(threads_arg_src Demography.cpp @@ -34,6 +46,8 @@ set(threads_arg_src AlleleAges.cpp GenotypeIterator.cpp VCFWriter.cpp + ForwardBackward.cpp + ThreadsIO.cpp ) set(threads_arg_hdr @@ -52,6 +66,8 @@ set(threads_arg_hdr AlleleAges.hpp GenotypeIterator.hpp VCFWriter.hpp + ForwardBackward.hpp + ThreadsIO.hpp ) add_library(threads_arg STATIC @@ -72,8 +88,17 @@ target_link_libraries(threads_arg PRIVATE Boost::headers project_warnings + $<$:OpenMP::OpenMP_CXX> + PUBLIC + HDF5::HDF5 ) +# Native-architecture tuning + LTO for hot numeric kernels +target_compile_options(threads_arg PRIVATE + $<$:-march=native -O3> +) +set_property(TARGET threads_arg PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE) + # Conditionally create python bindings if(PYTHON_BINDINGS) set_target_properties(threads_arg @@ -92,5 +117,7 @@ if(PYTHON_BINDINGS) project_warnings ) + set_property(TARGET threads_arg_python_bindings PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE) + install(TARGETS threads_arg_python_bindings LIBRARY DESTINATION .) endif() diff --git a/src/DataConsistency.cpp b/src/DataConsistency.cpp index 2294690..144c40d 100644 --- a/src/DataConsistency.cpp +++ b/src/DataConsistency.cpp @@ -15,6 +15,8 @@ // along with this program. If not, see . #include "DataConsistency.hpp" +#include "GenotypeIterator.hpp" +#include #include #include #include @@ -189,7 +191,9 @@ void ConsistencyWrapper::process_site(std::vector& genotypes) { // Otherwise we try traversing the local threading graph to find another carrier int current_target = converter.current_target; while (current_target != -1 && genotypes.at(current_target) != 1) { - current_target = instruction_converters.at(current_target).current_target; + int next = instruction_converters.at(current_target).current_target; + if (next == current_target) break; // avoid self-referencing loop + current_target = next; } if (current_target > 0 && genotypes.at(current_target) == 1) { new_target = current_target; @@ -199,6 +203,12 @@ void ConsistencyWrapper::process_site(std::vector& genotypes) { } } } + // Self-referencing targets break mismatch recording (comparing a + // sample's genotype with itself is always equal), so replace with + // sample 0 which is always correctly reconstructed. + if (new_target == current_hap) { + new_target = 0; + } if (new_target != converter.current_target) { // Force a new threading segment bounded by [0, allele_age) and // using the new target @@ -222,9 +232,22 @@ ThreadingInstructions ConsistencyWrapper::get_consistent_instructions() { // Make output threading instructions std::vector output_instructions; output_instructions.reserve(instruction_converters.size()); - for (InstructionConverter converter : instruction_converters) { + for (auto& converter : instruction_converters) { output_instructions.push_back(converter.parse_converted_instructions()); } return ThreadingInstructions(output_instructions, physical_positions); } + +ThreadingInstructions run_consistency(ThreadingInstructions& instructions, const std::vector& allele_ages) { + GenotypeIterator gt_it(instructions); + ConsistencyWrapper cw(instructions, allele_ages); + while (gt_it.has_next_genotype()) { + auto g = gt_it.next_genotype(); + // next_genotype returns const ref; process_site takes non-const ref + std::vector genotypes(g.begin(), g.end()); + cw.process_site(genotypes); + } + return cw.get_consistent_instructions(); +} + diff --git a/src/DataConsistency.hpp b/src/DataConsistency.hpp index a760a62..d3fb333 100644 --- a/src/DataConsistency.hpp +++ b/src/DataConsistency.hpp @@ -81,4 +81,8 @@ class ConsistencyWrapper { std::vector instruction_converters; }; +// Bulk function: creates GenotypeIterator + ConsistencyWrapper internally, +// processes all sites in C++, returns consistent instructions. +ThreadingInstructions run_consistency(ThreadingInstructions& instructions, const std::vector& allele_ages); + #endif // THREADS_ARG_DATA_CONSISTENCY_HPP \ No newline at end of file diff --git a/src/Demography.cpp b/src/Demography.cpp index 6b1f178..5a21366 100644 --- a/src/Demography.cpp +++ b/src/Demography.cpp @@ -83,7 +83,7 @@ double Demography::expected_branch_length(const int N) { std::ostream& operator<<(std::ostream& os, const Demography& d) { for (std::size_t i = 0; i < d.sizes.size(); i++) { - std::cout << d.times[i] << " " << d.sizes[i] << " " << d.std_times[i] << std::endl; + os << d.times[i] << " " << d.sizes[i] << " " << d.std_times[i] << "\n"; } return os; } diff --git a/src/ForwardBackward.cpp b/src/ForwardBackward.cpp new file mode 100644 index 0000000..868b4f4 --- /dev/null +++ b/src/ForwardBackward.cpp @@ -0,0 +1,95 @@ +// This file is part of the Threads software suite. +// Copyright (C) 2024-2025 Threads Developers. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU General Public License for more details. +// +// You should have received a copy of the GNU General Public License +// along with this program. If not, see . + +#include "ForwardBackward.hpp" + +std::pair, std::vector> +forwards_ls_hap(int n, int m, const double* H, const double* s, + const double* e, const double* r) { + std::vector F(m * n, 0.0); + std::vector c(m, 0.0); + const double inv_n = 1.0 / n; + + // Initialization (l = 0) + for (int i = 0; i < n; ++i) { + int match = (H[i] == s[0]) || (s[0] == FWBW_MISSING); + F[i] = inv_n * e[match]; + c[0] += F[i]; + } + double inv_c0 = 1.0 / c[0]; + for (int i = 0; i < n; ++i) { + F[i] *= inv_c0; + } + + // Forward pass + for (int l = 1; l < m; ++l) { + const double r_l = r[l]; + const double one_minus_r = 1.0 - r_l; + const double r_n = r_l * inv_n; + const double s_l = s[l]; + + for (int i = 0; i < n; ++i) { + double f_val = F[(l - 1) * n + i] * one_minus_r + r_n; + int match = (H[l * n + i] == s_l) || (s_l == FWBW_MISSING); + f_val *= e[l * 2 + match]; + F[l * n + i] = f_val; + c[l] += f_val; + } + + double inv_cl = 1.0 / c[l]; + for (int i = 0; i < n; ++i) { + F[l * n + i] *= inv_cl; + } + } + + return {std::move(F), std::move(c)}; +} + +std::vector +backwards_ls_hap(int n, int m, const double* H, const double* s, + const double* e, const double* c, const double* r) { + std::vector B(m * n, 0.0); + const double inv_n = 1.0 / n; + + // Initialization (l = m-1) + for (int i = 0; i < n; ++i) { + B[(m - 1) * n + i] = 1.0; + } + + // Backward pass + std::vector tmp_B(n); + for (int l = m - 2; l >= 0; --l) { + double tmp_B_sum = 0.0; + const double s_lp1 = s[l + 1]; + + for (int i = 0; i < n; ++i) { + int match = (H[(l + 1) * n + i] == s_lp1) || (s_lp1 == FWBW_MISSING); + tmp_B[i] = e[(l + 1) * 2 + match] * B[(l + 1) * n + i]; + tmp_B_sum += tmp_B[i]; + } + + const double r_lp1 = r[l + 1]; + const double r_n_lp1 = r_lp1 * inv_n; + const double one_minus_r_lp1 = 1.0 - r_lp1; + const double inv_c_lp1 = 1.0 / c[l + 1]; + + for (int i = 0; i < n; ++i) { + B[l * n + i] = (r_n_lp1 * tmp_B_sum + one_minus_r_lp1 * tmp_B[i]) * inv_c_lp1; + } + } + + return B; +} diff --git a/src/ForwardBackward.hpp b/src/ForwardBackward.hpp new file mode 100644 index 0000000..3ac7295 --- /dev/null +++ b/src/ForwardBackward.hpp @@ -0,0 +1,42 @@ +// This file is part of the Threads software suite. +// Copyright (C) 2024-2025 Threads Developers. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU General Public License for more details. +// +// You should have received a copy of the GNU General Public License +// along with this program. If not, see . + +#ifndef THREADS_ARG_FORWARD_BACKWARD_HPP +#define THREADS_ARG_FORWARD_BACKWARD_HPP + +#include +#include + +constexpr int FWBW_MISSING = -9; + +// Li-Stephens haploid forward algorithm (normalized). +// H: reference panel, row-major (m x n) +// s: query haplotype (length m) +// e: emission probabilities, row-major (m x 2), e[l*2+0]=mismatch, e[l*2+1]=match +// r: recombination probabilities (length m) +// Returns (F, c) where F is (m x n) forward matrix and c is (m) normalization factors. +std::pair, std::vector> +forwards_ls_hap(int n, int m, const double* H, const double* s, + const double* e, const double* r); + +// Li-Stephens haploid backward algorithm (normalized). +// Same inputs as forward, plus c (normalization factors from forward pass). +// Returns B: (m x n) backward matrix. +std::vector +backwards_ls_hap(int n, int m, const double* H, const double* s, + const double* e, const double* c, const double* r); + +#endif // THREADS_ARG_FORWARD_BACKWARD_HPP diff --git a/src/HMM.cpp b/src/HMM.cpp index 8f902c0..9e9b806 100644 --- a/src/HMM.cpp +++ b/src/HMM.cpp @@ -28,17 +28,17 @@ HMM::HMM(Demography demography, std::vector bp_sizes, std::vector trellis_row(num_states, 0.0); - std::vector pointer_row(num_states, 0); - trellis.push_back(trellis_row); - pointers.push_back(pointer_row); + trellis.emplace_back(num_states, 0.0); + pointers.emplace_back(num_states, 0); } } std::vector HMM::compute_expected_times(Demography demography, const int K) { std::vector result; + result.reserve(K); double k = static_cast(num_states); boost::math::exponential e; @@ -50,39 +50,45 @@ std::vector HMM::compute_expected_times(Demography demography, const int } void HMM::compute_recombination_scores(std::vector cm_sizes) { + const double log_K = std::log(num_states); + non_transition_score.reserve(cm_sizes.size()); + transition_score.reserve(cm_sizes.size()); for (std::size_t i = 0; i < cm_sizes.size(); i++) { - non_transition_score.push_back(std::vector()); - transition_score.push_back(std::vector()); + std::vector non_trans(num_states); + std::vector trans(num_states); for (int k = 0; k < num_states; k++) { - double t = expected_times[k]; - const double l = 2. * 0.01 * cm_sizes[i] * t; - const double trans = std::log1p(-std::exp(-l)) - std::log(num_states); + const double l = 2. * 0.01 * cm_sizes[i] * expected_times[k]; + const double t = std::log1p(-std::exp(-l)) - log_K; // log-prob of transitioning - transition_score[i].push_back(trans); + trans[k] = t; // log-prob of *not* transitioning - non_transition_score[i].push_back(std::log(std::exp(-l) + std::exp(trans))); + non_trans[k] = std::log(std::exp(-l) + std::exp(t)); } + transition_score.push_back(std::move(trans)); + non_transition_score.push_back(std::move(non_trans)); } } void HMM::compute_mutation_scores(std::vector bp_sizes, double mutation_rate) { + hom_score.reserve(bp_sizes.size()); + het_score.reserve(bp_sizes.size()); for (std::size_t i = 0; i < bp_sizes.size(); i++) { - hom_score.push_back(std::vector()); - het_score.push_back(std::vector()); + std::vector hom(num_states); + std::vector het(num_states); for (int k = 0; k < num_states; k++) { - double t = expected_times[k]; - // TODO: use mean-bp sizes here as in the main algorithm - const double l = 2. * mutation_rate * bp_sizes[i] * t; + const double l = 2. * mutation_rate * bp_sizes[i] * expected_times[k]; // log-prob of mutating - het_score[i].push_back(std::log1p(-std::exp(-l))); + het[k] = std::log1p(-std::exp(-l)); // log-prob of *not* mutating - hom_score[i].push_back(-l); + hom[k] = -l; } + het_score.push_back(std::move(het)); + hom_score.push_back(std::move(hom)); } } @@ -107,12 +113,14 @@ std::vector HMM::breakpoints(std::vector observations, int start) { double score = 0.0; unsigned short running_argmax = 0; for (int j = 1; j < neighborhood_size; j++) { + const int js = j + start; for (int i = 0; i < num_states; i++) { + // Hoist mut_score out of the inner k-loop: it only depends on i, not k + const double mut_score = observations[j] ? het_score[js][i] : hom_score[js][i]; double running_max = 0; for (int k = 0; k < num_states; k++) { - double mut_score = observations[j] ? het_score[j + start][i] : hom_score[j + start][i]; double rec_score = - k == i ? non_transition_score[j + start][k] : transition_score[j + start][k]; + k == i ? non_transition_score[js][k] : transition_score[js][k]; score = trellis[j - 1 + start][k] + rec_score + mut_score; @@ -121,8 +129,8 @@ std::vector HMM::breakpoints(std::vector observations, int start) { running_argmax = static_cast(k); } } - trellis[j + start][i] = running_max; - pointers[j + start][i] = running_argmax; + trellis[js][i] = running_max; + pointers[js][i] = running_argmax; } } diff --git a/src/ImputationMatcher.cpp b/src/ImputationMatcher.cpp index f869f6c..828049a 100644 --- a/src/ImputationMatcher.cpp +++ b/src/ImputationMatcher.cpp @@ -191,7 +191,7 @@ void ImputationMatcher::process_site(const std::vector& genotype) { throw std::runtime_error(prompt); } } - sorting = next_sorting; + std::swap(sorting, next_sorting); sites_processed++; } diff --git a/src/Matcher.cpp b/src/Matcher.cpp index e19d3ec..0ffbae2 100644 --- a/src/Matcher.cpp +++ b/src/Matcher.cpp @@ -25,10 +25,16 @@ #include // For a given interval, this contains all the matches for all the samples -MatchGroup::MatchGroup(int _num_samples, double _cm_position) +MatchGroup::MatchGroup(int _num_samples, int _expected_queries, double _cm_position) : num_samples(_num_samples), cm_position(_cm_position) { + // Pre-reserve hash maps to avoid rehashing during neighbor counting. + // Each query site adds ~4 neighbors per sample; reserve for expected total. + int reserve_size = std::max(16, _expected_queries * 4); + match_candidates_counts.reserve(num_samples); for (int i = 0; i < num_samples; i++) { - match_candidates_counts.push_back(std::unordered_map()); + std::unordered_map m; + m.reserve(reserve_size); + match_candidates_counts.push_back(std::move(m)); } } @@ -55,14 +61,14 @@ void MatchGroup::filter_matches(int min_matches) { } } else if (i < 1000) { - for (auto counts : match_candidates_counts.at(i)) { + for (const auto& counts : match_candidates_counts.at(i)) { if (counts.second >= std::min(2, min_matches)) { match_candidates.at(i).insert(counts.first); } } } else if (i < 10000) { - for (auto counts : match_candidates_counts.at(i)) { + for (const auto& counts : match_candidates_counts.at(i)) { if (counts.second >= min_matches) { match_candidates.at(i).insert(counts.first); } @@ -70,7 +76,7 @@ void MatchGroup::filter_matches(int min_matches) { } else { // Don't want too much stuff for very big studies - for (auto counts : match_candidates_counts.at(i)) { + for (const auto& counts : match_candidates_counts.at(i)) { if (counts.second >= 2 * min_matches) { match_candidates.at(i).insert(counts.first); } @@ -80,7 +86,7 @@ void MatchGroup::filter_matches(int min_matches) { if (match_candidates.at(i).size() == 0) { int tmp_min_matches = min_matches; while (match_candidates.at(i).size() == 0 && tmp_min_matches > 0) { - for (auto counts : match_candidates_counts.at(i)) { + for (const auto& counts : match_candidates_counts.at(i)) { if (counts.second >= tmp_min_matches) { match_candidates.at(i).insert(counts.first); } @@ -106,7 +112,7 @@ void MatchGroup::filter_matches(int min_matches) { void MatchGroup::insert_tops_from(MatchGroup& other) { for (int i = 1; i < num_samples; i++) { - for (auto p : other.top_four_maps.at(i)) { + for (const auto& p : other.top_four_maps.at(i)) { match_candidates.at(i).insert(p.first); } } @@ -180,9 +186,13 @@ Matcher::Matcher(int _n, const std::vector& _genetic_positions, double _ std::cout << "Will use " << query_sites.size() << " query sites and " << match_group_sites.size() << " match_group_sites" << std::endl; + // Estimate queries per match group for hash map pre-reservation + int expected_queries_per_group = std::max(1, + static_cast(query_sites.size()) / std::max(1, static_cast(match_group_sites.size()))); + match_groups.reserve(match_group_sites.size()); for (int match_group_site : match_group_sites) { - match_groups.emplace_back(num_samples, genetic_positions[match_group_site]); + match_groups.emplace_back(num_samples, expected_queries_per_group, genetic_positions[match_group_site]); } sorting.reserve(num_samples); @@ -196,106 +206,117 @@ Matcher::Matcher(int _n, const std::vector& _genetic_positions, double _ } void Matcher::process_site(const std::vector& genotype) { - // Pass genotypes for a single site through the matcher + if (static_cast(genotype.size()) != num_samples) { + throw std::runtime_error("invalid genotype vector size"); + } + process_site_raw(genotype.data()); +} +void Matcher::process_site_raw(const int* genotype) { if (sites_processed >= num_sites) { throw std::runtime_error("all sites have already been processed"); } + const int* sort_data = sorting.data(); + int* next_sort_data = next_sorting.data(); + // Get allele count int allele_count = 0; - for (int g : genotype) { - if (g == 1) { - allele_count++; - } + for (int i = 0; i < num_samples; i++) { + allele_count += genotype[i]; } + + // PBWT step — no bounds checking, alleles validated by caller int counter0 = 0; int counter1 = 0; - if (static_cast(genotype.size()) != num_samples) { - throw std::runtime_error("invalid genotype vector size"); - } - - // PBWT step + const int offset1 = num_samples - allele_count; for (int i = 0; i < num_samples; i++) { - if (genotype.at(sorting.at(i)) == 1) { - next_sorting[num_samples - allele_count + counter1] = sorting.at(i); + const int sid = sort_data[i]; + if (genotype[sid] == 1) { + next_sort_data[offset1 + counter1] = sid; counter1++; } - else if (genotype.at(sorting.at(i)) == 0) { - next_sorting[counter0] = sorting.at(i); - counter0++; - } else { - std::string prompt = "invalid genotype" + std::to_string(genotype.at(sorting.at(i))); - throw std::runtime_error(prompt); + next_sort_data[counter0] = sid; + counter0++; } } - sorting = next_sorting; + std::swap(sorting, next_sorting); // Threading-neighbor queries if (match_group_idx < (static_cast(match_group_sites.size()) - 1) && - (sites_processed >= match_group_sites.at(match_group_idx + 1))) { + (sites_processed >= match_group_sites[match_group_idx + 1])) { match_group_idx++; - match_groups.at(match_group_idx - 1).filter_matches(min_matches); + match_groups[match_group_idx - 1].filter_matches(min_matches); } // If we've reached a query site, query if (next_query_site_idx < static_cast(query_sites.size()) && - sites_processed == query_sites.at(next_query_site_idx)) { + sites_processed == query_sites[next_query_site_idx]) { // Get the arg-sort of the sorting + const int* sort_ptr = sorting.data(); + int* perm_data = permutation.data(); for (int i = 0; i < num_samples; i++) { - permutation[sorting.at(i)] = i; + perm_data[sort_ptr[i]] = i; } next_query_site_idx++; - // Initialize the red-black tree - std::set threaded = {permutation.at(0)}; + // Boolean array for O(1) mark + sequential scan neighbor finding + std::vector inserted(num_samples, 0); + inserted[perm_data[0]] = 1; // Insert sequences and query in order for (int i = 1; i < num_samples; i++) { - std::vector matches; - matches.reserve(neighborhood_size); - auto iter = threaded.insert(permutation.at(i)); - auto iter_up = iter.first; - auto iter_down = iter.first; - // Check if genotypes are identical, just to be sure - while ((static_cast(matches.size()) < neighborhood_size) && - (iter_down != threaded.begin() || iter_up != threaded.end())) { - if (iter_down != threaded.begin()) { - iter_down--; - matches.push_back(sorting.at(*iter_down)); - } - if (static_cast(matches.size()) < neighborhood_size && iter_up != threaded.end()) { - iter_up++; - if (iter_up != threaded.end()) { - matches.push_back(sorting.at(*iter_up)); + const int pos = perm_data[i]; + inserted[pos] = 1; + + // Find neighborhood_size nearest neighbors by scanning left/right + int n_found = 0; + int left = pos - 1; + int right = pos + 1; + std::unordered_map& mmmap = + match_groups[match_group_idx].match_candidates_counts[i]; + while (n_found < neighborhood_size && (left >= 0 || right < num_samples)) { + if (left >= 0) { + while (left >= 0 && !inserted[left]) left--; + if (left >= 0) { + mmmap[sort_ptr[left]]++; + n_found++; + left--; } } - } - for (int m : matches) { - std::unordered_map& mmmap = - match_groups.at(match_group_idx).match_candidates_counts.at(i); - if (m >= i) { - throw std::runtime_error("Illegal match candidate " + std::to_string(m) + - ", something is very wrong"); - } - if (!mmmap.count(m)) { - mmmap[m] = 1; - } - else { - mmmap[m]++; + if (n_found < neighborhood_size && right < num_samples) { + while (right < num_samples && !inserted[right]) right++; + if (right < num_samples) { + mmmap[sort_ptr[right]]++; + n_found++; + right++; + } } } } // Special case for last query if (next_query_site_idx == static_cast(query_sites.size())) { - match_groups.at(match_group_sites.size() - 1).filter_matches(min_matches); + match_groups[match_group_sites.size() - 1].filter_matches(min_matches); } } sites_processed++; } +void Matcher::process_all_sites(const std::vector>& genotypes) { + for (const auto& genotype : genotypes) { + process_site_raw(genotype.data()); + } +} + +void Matcher::process_all_sites_flat(const int32_t* data, int n_sites, int n_haps) { + static_assert(sizeof(int) == sizeof(int32_t), "int and int32_t must be the same size"); + for (int s = 0; s < n_sites; s++) { + process_site_raw(reinterpret_cast(data + static_cast(s) * n_haps)); + } +} + // Propagate top 4 matches from left and right match groups void Matcher::propagate_adjacent_matches() { for (int i = 1; i < static_cast(match_groups.size()); i++) { diff --git a/src/Matcher.hpp b/src/Matcher.hpp index a7039ec..37dd657 100644 --- a/src/Matcher.hpp +++ b/src/Matcher.hpp @@ -17,6 +17,7 @@ #ifndef THREADS_ARG_MATCHER_HPP #define THREADS_ARG_MATCHER_HPP +#include #include #include #include @@ -24,7 +25,7 @@ /// for a certain interval, store the matches for all samples class MatchGroup { public: - MatchGroup(int _num_samples, double cm_position); + MatchGroup(int _num_samples, int _expected_queries, double cm_position); MatchGroup(const std::vector& target_ids, const std::vector>& matches, const double _cm_position); void filter_matches(int min_matches); @@ -46,6 +47,9 @@ class Matcher { // Do all the work void process_site(const std::vector& genotype); + void process_site_raw(const int* genotype); + void process_all_sites(const std::vector>& genotypes); + void process_all_sites_flat(const int32_t* data, int n_sites, int n_haps); void propagate_adjacent_matches(); void clear(); diff --git a/src/State.cpp b/src/State.cpp index fd5b93f..094af8b 100644 --- a/src/State.cpp +++ b/src/State.cpp @@ -81,25 +81,22 @@ void StateBranch::prune() { } StateTree::StateTree(std::vector& states) { - for (auto s : states) { + for (const auto& s : states) { int sample_ID = s.below->sample_ID; - if (branches.find(sample_ID) == branches.end()) { - branches[sample_ID] = StateBranch(); - } branches[sample_ID].insert(s); } } void StateTree::prune() { - for (auto pair : branches) { - branches[pair.first].prune(); + for (auto& [key, branch] : branches) { + branch.prune(); } } std::vector StateTree::dump() const { std::vector states; - for (auto pair : branches) { - for (auto s : pair.second.states) { + for (const auto& [key, branch] : branches) { + for (const auto& s : branch.states) { states.push_back(s); } } diff --git a/src/ThreadingInstructions.cpp b/src/ThreadingInstructions.cpp index 4ba5be4..5457d1e 100644 --- a/src/ThreadingInstructions.cpp +++ b/src/ThreadingInstructions.cpp @@ -17,11 +17,19 @@ #include "ThreadingInstructions.hpp" #include "GenotypeIterator.hpp" +#include #include +#include #include #include +#include +#include #include +#ifdef _OPENMP +#include +#endif + // perf: Note that args are passed by value rather than ref as they run faster // when used with std::move below. @@ -262,168 +270,993 @@ ThreadingInstructions ThreadingInstructions::sub_range(const int range_start, co }; } +void ThreadingInstructions::materialize_genotypes() { + if (genotypes_materialized) return; + genotype_matrix.resize(static_cast(num_sites) * num_samples); + GenotypeIterator gi(*this); + int site = 0; + while (gi.has_next_genotype()) { + const std::vector& g = gi.next_genotype(); + std::copy(g.begin(), g.end(), genotype_matrix.begin() + static_cast(site) * num_samples); + site++; + } + genotypes_materialized = true; +} + +void ThreadingInstructions::materialize_diploid() { + if (diploid_materialized) return; + materialize_genotypes(); + const int n = num_samples; + const int n_dip = n / 2; + diploid_matrix.resize(static_cast(num_sites) * n_dip); + const int* gmat = genotype_matrix.data(); + int* dmat = diploid_matrix.data(); + for (int s = 0; s < num_sites; s++) { + const int* g = gmat + static_cast(s) * n; + int* d = dmat + static_cast(s) * n_dip; + for (int i = 0; i < n_dip; i++) { + d[i] = g[2 * i] + g[2 * i + 1]; + } + } + diploid_materialized = true; +} + +void ThreadingInstructions::materialize_normalized_haploid() { + if (standardized_hap_ready) return; + materialize_genotypes(); + const int n = num_samples; + standardized_hap.resize(static_cast(num_sites) * n); + const int* gmat = genotype_matrix.data(); + double* zmat = standardized_hap.data(); + for (int s = 0; s < num_sites; s++) { + const int* g = gmat + static_cast(s) * n; + double* z = zmat + static_cast(s) * n; + int ac = 0; + for (int i = 0; i < n; i++) ac += g[i]; + const double mu = static_cast(ac) / n; + const double inv_std = 1.0 / std::sqrt(mu * (1.0 - mu)); + for (int i = 0; i < n; i++) { + z[i] = (g[i] - mu) * inv_std; + } + } + standardized_hap_ready = true; +} + +void ThreadingInstructions::materialize_normalized_diploid() { + if (standardized_dip_ready) return; + materialize_diploid(); + const int n = num_samples; + const int n_dip = n / 2; + standardized_dip.resize(static_cast(num_sites) * n_dip); + const int* dmat = diploid_matrix.data(); + double* zmat = standardized_dip.data(); + for (int s = 0; s < num_sites; s++) { + const int* d = dmat + static_cast(s) * n_dip; + double* z = zmat + static_cast(s) * n_dip; + int ac = 0; + for (int i = 0; i < n_dip; i++) ac += d[i]; + const double mu = static_cast(ac) / n_dip; + double sample_var = 0.0; + for (int i = 0; i < n_dip; i++) { + double diff = d[i] - mu; + sample_var += diff * diff; + } + sample_var /= n_dip; + const double inv_std = 1.0 / std::sqrt(sample_var); + for (int i = 0; i < n_dip; i++) { + z[i] = (d[i] - mu) * inv_std; + } + } + standardized_dip_ready = true; +} + std::vector ThreadingInstructions::left_multiply(const std::vector& x, bool diploid, bool normalize) { - // Left-multiplication of the genotype matrix by a vector of doubles + const int n = num_samples; + const int n_dip = n / 2; - // Check input vector lengths are correct if (diploid) { - if (x.size() != num_samples / 2) { + if (static_cast(x.size()) != n_dip) { std::ostringstream oss; - oss << "Input vector must have length " << num_samples / 2 << "."; + oss << "Input vector must have length " << n_dip << "."; throw std::runtime_error(oss.str()); } } else { - if (x.size() != num_samples) { + if (static_cast(x.size()) != n) { std::ostringstream oss; - oss << "Input vector must have length " << num_samples << "."; + oss << "Input vector must have length " << n << "."; throw std::runtime_error(oss.str()); } } - // Initialize genotype traversal - GenotypeIterator gi = GenotypeIterator(*this); - std::size_t site_counter = 0; + const double* xp = x.data(); std::vector out(num_sites); + double* outp = out.data(); - while (gi.has_next_genotype()) { - // Fetch the next genotype - const std::vector& g = gi.next_genotype(); + if (normalize && diploid) { + materialize_normalized_diploid(); + const double* zmat = standardized_dip.data(); + for (int s = 0; s < num_sites; s++) { + const double* z = zmat + static_cast(s) * n_dip; + double entry = 0.0; + for (int i = 0; i < n_dip; i++) { + entry += xp[i] * z[i]; + } + outp[s] = entry; + } + } else if (normalize) { + materialize_normalized_haploid(); + const double* zmat = standardized_hap.data(); + for (int s = 0; s < num_sites; s++) { + const double* z = zmat + static_cast(s) * n; + double entry = 0.0; + for (int i = 0; i < n; i++) { + entry += xp[i] * z[i]; + } + outp[s] = entry; + } + } else if (diploid) { + materialize_diploid(); + const int* dmat = diploid_matrix.data(); + for (int s = 0; s < num_sites; s++) { + const int* d = dmat + static_cast(s) * n_dip; + double entry = 0.0; + for (int i = 0; i < n_dip; i++) { + entry += xp[i] * d[i]; + } + outp[s] = entry; + } + } else { + materialize_genotypes(); + const int* gmat = genotype_matrix.data(); + for (int s = 0; s < num_sites; s++) { + const int* g = gmat + static_cast(s) * n; + double entry = 0.0; + for (int i = 0; i < n; i++) { + entry += xp[i] * g[i]; + } + outp[s] = entry; + } + } + return out; +} + +std::vector ThreadingInstructions::right_multiply(const std::vector& x, bool diploid, bool normalize) { + if (static_cast(x.size()) != num_sites) { + std::ostringstream oss; + oss << "Input vector must have length " << num_sites << "."; + throw std::runtime_error(oss.str()); + } + + const int n = num_samples; + const int n_dip = n / 2; + const int out_size = diploid ? n_dip : n; + const double* xp = x.data(); + + std::vector out(out_size, 0.0); + double* outp = out.data(); + + if (normalize && diploid) { + materialize_normalized_diploid(); + const double* zmat = standardized_dip.data(); + for (int s = 0; s < num_sites; s++) { + const double* z = zmat + static_cast(s) * n_dip; + const double w = xp[s]; + for (int i = 0; i < n_dip; i++) { + outp[i] += w * z[i]; + } + } + } else if (normalize) { + materialize_normalized_haploid(); + const double* zmat = standardized_hap.data(); + for (int s = 0; s < num_sites; s++) { + const double* z = zmat + static_cast(s) * n; + const double w = xp[s]; + for (int i = 0; i < n; i++) { + outp[i] += w * z[i]; + } + } + } else if (diploid) { + materialize_diploid(); + const int* dmat = diploid_matrix.data(); + for (int s = 0; s < num_sites; s++) { + const int* d = dmat + static_cast(s) * n_dip; + const double w = xp[s]; + for (int i = 0; i < n_dip; i++) { + outp[i] += w * d[i]; + } + } + } else { + materialize_genotypes(); + const int* gmat = genotype_matrix.data(); + for (int s = 0; s < num_sites; s++) { + const int* g = gmat + static_cast(s) * n; + const double w = xp[s]; + for (int i = 0; i < n; i++) { + outp[i] += w * g[i]; + } + } + } + return out; +} + +void ThreadingInstructions::prepare_tree_multiply() { + if (tree_ready) return; + + const int n = num_samples; + const int m = num_sites; + + // ── Phase 1: Build global intervals (unchanged) ───────────────────── + + tree_ref_genome.assign(m, 0); + for (int idx : instructions[0].mismatches) { + if (idx >= 0 && idx < m) tree_ref_genome[idx] = 1; + } + + std::set break_set; + break_set.insert(0); + for (int i = 0; i < n; i++) { + const auto& starts_i = instructions[i].starts; + const int n_segs = static_cast(starts_i.size()); + for (int k = 1; k < n_segs; k++) { + auto it = std::lower_bound(positions.begin() + 1, positions.end(), starts_i[k]); + if (it != positions.end()) { + break_set.insert(static_cast(it - positions.begin())); + } + } + } + + std::vector breaks(break_set.begin(), break_set.end()); + const int ni = static_cast(breaks.size()); + tree_n_intervals = ni; + tree_ivl_start.resize(ni); + tree_ivl_end.resize(ni); + for (int j = 0; j < ni; j++) { + tree_ivl_start[j] = breaks[j]; + tree_ivl_end[j] = (j + 1 < ni) ? breaks[j + 1] : m; + } + + tree_ivl_seg.resize(static_cast(n) * ni); + for (int i = 0; i < n; i++) { + const auto& starts_i = instructions[i].starts; + const int n_segs = static_cast(starts_i.size()); + int seg = 0; + for (int j = 0; j < ni; j++) { + const int site = tree_ivl_start[j]; + if (site == 0) { + seg = 0; + } else { + const int pos = positions[site]; + while (seg + 1 < n_segs && starts_i[seg + 1] <= pos) { + seg++; + } + } + tree_ivl_seg[static_cast(i) * ni + j] = seg; + } + } + + // ── Phase 2: Per-sample segment-to-interval mapping ───────────────── + // For each sample, record the first interval of each segment. + // This enables segment-level iteration in the multiply. + + tree_seg_offset.resize(n); + tree_seg_first_ivl.clear(); + for (int i = 0; i < n; i++) { + tree_seg_offset[i] = static_cast(tree_seg_first_ivl.size()); + int prev_seg = -1; + for (int j = 0; j < ni; j++) { + const int seg = tree_ivl_seg[static_cast(i) * ni + j]; + if (seg != prev_seg) { + tree_seg_first_ivl.push_back(j); + prev_seg = seg; + } + } + } + // Sentinel for end-of-last-sample + tree_seg_first_ivl.push_back(0); + + // ── Phase 3: Precompute mismatch signs and carry values ───────────── + // Uses temporary reference-counted genotype cache. + // Peak memory: O(m * max_active_targets) instead of O(n * m). + + // Count references: how many future samples use each target + std::vector ref_count(n, 0); + for (int i = 1; i < n; i++) { + std::set seen; + for (int t : instructions[i].targets) { + if (t != i && seen.insert(t).second) ref_count[t]++; + } + } + + // Allocate mismatch sign storage + tree_mm_offset.resize(n); + size_t total_mm = 0; + for (int i = 0; i < n; i++) { + tree_mm_offset[i] = static_cast(total_mm); + total_mm += instructions[i].mismatches.size(); + } + tree_mm_sign.assign(total_mm, 0); + tree_mm_ivl.resize(total_mm); + + // Allocate carry storage + tree_carry.assign(static_cast(n) * ni, 0); + + // Reference-counted genotype cache + std::unordered_map> geno_cache; + std::vector geno_buf(m, 0); + + for (int i = 0; i < n; i++) { + const int* mm_data = instructions[i].mismatches.data(); + const int n_mm = static_cast(instructions[i].mismatches.size()); + + if (i == 0) { + // Sample 0: genotype = ref_genome + for (int s = 0; s < m; s++) geno_buf[s] = static_cast(tree_ref_genome[s]); + } else { + // Build genotype from target + mismatches + int carry = 0; + for (int j = 0; j < ni; j++) { + const int seg = tree_ivl_seg[static_cast(i) * ni + j]; + const int target = instructions[i].targets[seg]; + const int a = tree_ivl_start[j]; + const int b = tree_ivl_end[j]; + + if (target != i) { + const auto& tgt = geno_cache.at(target); + std::memcpy(&geno_buf[a], &tgt[a], b - a); + // Flip at mismatches + const int* lo = std::lower_bound(mm_data, mm_data + n_mm, a); + const int* hi = std::lower_bound(mm_data, mm_data + n_mm, b); + for (const int* it = lo; it != hi; ++it) { + geno_buf[*it] ^= 1; + } + carry = geno_buf[b - 1]; + } else { + // Self-ref: carry forward, flip at mismatches + const int* lo = std::lower_bound(mm_data, mm_data + n_mm, a); + const int* hi = std::lower_bound(mm_data, mm_data + n_mm, b); + const int* it = lo; + for (int s = a; s < b; s++) { + if (it != hi && *it == s) { + carry ^= 1; + ++it; + } + geno_buf[s] = static_cast(carry); + } + } + } + + // Compute mismatch correction signs and interval mapping + for (int k = 0; k < n_mm; k++) { + const int s = mm_data[k]; + // Find interval containing site s (precompute for multiply) + auto ivl_it = std::upper_bound(tree_ivl_start.begin(), tree_ivl_start.end(), s); + const int j = static_cast(ivl_it - tree_ivl_start.begin()) - 1; + tree_mm_ivl[tree_mm_offset[i] + k] = j; + const int seg = tree_ivl_seg[static_cast(i) * ni + j]; + const int target = instructions[i].targets[seg]; + if (target != i) { + const int g_target = geno_cache.at(target)[s]; + tree_mm_sign[tree_mm_offset[i] + k] = static_cast(1 - 2 * g_target); + } + // Self-ref mismatches: leave as 0 (not used) + } + } + + // Store carry at end of each interval + for (int j = 0; j < ni; j++) { + tree_carry[static_cast(i) * ni + j] = static_cast(geno_buf[tree_ivl_end[j] - 1]); + } + + // Cache genotype if still referenced + if (ref_count[i] > 0) { + geno_cache[i] = geno_buf; + } + + // Decrement references and free unreferenced caches + if (i > 0) { + std::set seen; + for (int t : instructions[i].targets) { + if (t != i && seen.insert(t).second) { + if (--ref_count[t] == 0) { + geno_cache.erase(t); + } + } + } + } + } + + tree_ready = true; +} - // Initialize the next entry - double entry = 0.0; +std::vector ThreadingInstructions::right_multiply_tree(const std::vector& x) { + if (static_cast(x.size()) != num_sites) { + std::ostringstream oss; + oss << "Input vector must have length " << num_sites << "."; + throw std::runtime_error(oss.str()); + } + + prepare_tree_multiply(); + + const int n = num_samples; + const int m = num_sites; + const int ni = tree_n_intervals; + const double* xp = x.data(); + const int* ref = tree_ref_genome.data(); + const int* ivl_s = tree_ivl_start.data(); + const int* ivl_e = tree_ivl_end.data(); + + // Prefix sums of x (for self-ref segments) + std::vector prefix_x(m + 1, 0.0); + for (int s = 0; s < m; s++) { + prefix_x[s + 1] = prefix_x[s] + xp[s]; + } + + // ── Reference-counted cumulative interval sums ── + // Instead of n rows, allocate only O(tree_depth) rows via a pool. + const size_t cum_stride = static_cast(ni + 1); + + // Reference counting: how many future samples need each sample's cum row + std::vector cum_ref(n, 0); + { + std::vector seen(n, false); + for (int i = 1; i < n; i++) { + for (int t : instructions[i].targets) + if (t != i && !seen[t]) { seen[t] = true; cum_ref[t]++; } + for (int t : instructions[i].targets) + if (t != i) seen[t] = false; + } + } + + // Pool of cum rows + std::vector> cum_pool; + std::vector free_rows; + std::vector sample_to_row(n, -1); + + auto alloc_row = [&]() -> int { + if (!free_rows.empty()) { + int r = free_rows.back(); + free_rows.pop_back(); + std::fill(cum_pool[r].begin(), cum_pool[r].end(), 0.0); + return r; + } + int r = static_cast(cum_pool.size()); + cum_pool.emplace_back(cum_stride, 0.0); + return r; + }; - if (normalize) { - // If we want to normalize, we need the mean and standard deviation of g. - double ac = 0.0; - for (auto a : g) { - ac += a; + auto release_row = [&](int sample) { + int r = sample_to_row[sample]; + if (r >= 0) { + free_rows.push_back(r); + sample_to_row[sample] = -1; + } + }; + + // Sample 0: reference genome + { + int r0 = alloc_row(); + sample_to_row[0] = r0; + double* c0 = cum_pool[r0].data(); + for (int j = 0; j < ni; j++) { + double s = 0.0; + for (int site = ivl_s[j]; site < ivl_e[j]; site++) { + s += xp[site] * ref[site]; } - if (diploid) { - // We do the diploid standard deviation by hand - double mu = 2.0 * ac / num_samples; - double sample_var = 0.0; - for (std::size_t i=0; i < x.size(); i++) { - int h = g[2 * i] + g[2 * i + 1]; - double d = h - mu; - sample_var += d * d; + c0[j + 1] = c0[j] + s; + } + } + + std::vector out(n, 0.0); + out[0] = cum_pool[sample_to_row[0]][ni]; + + // Process samples 1..n-1 using segment-level loop + for (int i = 1; i < n; i++) { + const int* mm_data = instructions[i].mismatches.data(); + const int n_mm = static_cast(instructions[i].mismatches.size()); + const int8_t* my_signs = &tree_mm_sign[tree_mm_offset[i]]; + const bool need_cum = (cum_ref[i] > 0); + + int my_row = -1; + double* my_cum = nullptr; + if (need_cum) { + my_row = alloc_row(); + sample_to_row[i] = my_row; + my_cum = cum_pool[my_row].data(); + } + + const int seg_off = tree_seg_offset[i]; + int n_mapped_segs; + { + int next_off = (i + 1 < n) ? tree_seg_offset[i + 1] + : static_cast(tree_seg_first_ivl.size()) - 1; + n_mapped_segs = next_off - seg_off; + } + + double total = 0.0; + + for (int sk = 0; sk < n_mapped_segs; sk++) { + const int first_ivl = tree_seg_first_ivl[seg_off + sk]; + const int last_ivl = (sk + 1 < n_mapped_segs) + ? tree_seg_first_ivl[seg_off + sk + 1] + : ni; + const int seg = tree_ivl_seg[static_cast(i) * ni + first_ivl]; + const int target = instructions[i].targets[seg]; + const int site_a = ivl_s[first_ivl]; + const int site_b = ivl_e[last_ivl - 1]; + + if (target != i) { + const double* tgt_cum = cum_pool[sample_to_row[target]].data(); + const double base = tgt_cum[last_ivl] - tgt_cum[first_ivl]; + + const int* lo = std::lower_bound(mm_data, mm_data + n_mm, site_a); + const int* hi = std::lower_bound(mm_data, mm_data + n_mm, site_b); + double corr = 0.0; + for (const int* it = lo; it != hi; ++it) { + const int k = static_cast(it - mm_data); + corr += xp[*it] * my_signs[k]; } - sample_var /= (num_samples / 2); + total += base + corr; - double std = std::sqrt(sample_var); - for (std::size_t i=0; i < x.size(); i++) { - int h = g[2 * i] + g[2 * i + 1]; - double w = x[i]; - entry += w * (h - mu) / std; + if (need_cum) { + const double offset = my_cum[first_ivl] - tgt_cum[first_ivl]; + std::memcpy(&my_cum[first_ivl + 1], &tgt_cum[first_ivl + 1], + (last_ivl - first_ivl) * sizeof(double)); + for (int j = first_ivl + 1; j <= last_ivl; j++) { + my_cum[j] += offset; + } + const int* mm_ivl_data = tree_mm_ivl.data() + tree_mm_offset[i]; + for (const int* it = lo; it != hi; ++it) { + const int k = static_cast(it - mm_data); + const int mm_ivl = mm_ivl_data[k]; + const double c = xp[*it] * my_signs[k]; + for (int j = mm_ivl + 1; j <= last_ivl; j++) { + my_cum[j] += c; + } + } } } else { - double mu = ac / num_samples; - double std = std::sqrt(mu * (1 - mu)); - for (std::size_t i=0; i < g.size(); i++) { - double w = x[i]; - entry += w * (g[i] - mu) / std; + // Self-ref: interval by interval + for (int j = first_ivl; j < last_ivl; j++) { + const int a = ivl_s[j]; + const int b = ivl_e[j]; + int carry = (j == 0) ? 0 + : static_cast(tree_carry[static_cast(i) * ni + j - 1]); + const int* lo = std::lower_bound(mm_data, mm_data + n_mm, a); + const int* hi = std::lower_bound(mm_data, mm_data + n_mm, b); + double ivl_sum = 0.0; + int prev = a; + for (const int* it = lo; it != hi; ++it) { + const int mm_site = *it; + if (mm_site > prev) + ivl_sum += carry * (prefix_x[mm_site] - prefix_x[prev]); + carry = 1 - carry; + ivl_sum += xp[mm_site] * carry; + prev = mm_site + 1; + } + if (b > prev) + ivl_sum += carry * (prefix_x[b] - prefix_x[prev]); + if (need_cum) my_cum[j + 1] = my_cum[j] + ivl_sum; + total += ivl_sum; } } - } else { - for (std::size_t i=0; i < g.size(); i++) { - double w = diploid ? x[i / 2] : x[i]; - entry += w * g[i]; + } + out[i] = total; + + // Release targets no longer needed + { + const auto& tgts = instructions[i].targets; + const int nt = static_cast(tgts.size()); + for (int si = 0; si < nt; si++) { + int t = tgts[si]; + if (t == i) continue; + bool dup = false; + for (int sj = 0; sj < si; sj++) + if (tgts[sj] == t) { dup = true; break; } + if (!dup && --cum_ref[t] == 0) release_row(t); } } - out[site_counter] = entry; - site_counter++; + // Release own row if no one references us + if (cum_ref[i] == 0) release_row(i); } + return out; } -std::vector ThreadingInstructions::right_multiply(const std::vector& x, bool diploid, bool normalize) { - // Right-multiplication of the genotype matrix by a vector of doubles - - // Check input vector lengths are correct - if (x.size() != num_sites) { +std::vector ThreadingInstructions::left_multiply_tree(const std::vector& x) { + if (static_cast(x.size()) != num_samples) { std::ostringstream oss; - oss << "Input vector must have length " << num_samples / 2 << "."; + oss << "Input vector must have length " << num_samples << "."; throw std::runtime_error(oss.str()); } - GenotypeIterator gi = GenotypeIterator(*this); - std::size_t site_counter = 0; - if (diploid) { - // Initialize output - std::vector out(num_samples / 2, 0.0); - if (normalize) { - while (gi.has_next_genotype()) { - // Fetch the next genotype - const std::vector& g = gi.next_genotype(); - - // If we want to normalize, we need the mean and standard deviation of g. - double ac = 0.0; - for (auto a : g) { - ac += a; - } + prepare_tree_multiply(); + + const int n = num_samples; + const int m = num_sites; + const int ni = tree_n_intervals; + const double* xp = x.data(); - // We do the diploid standard deviation by hand - const double mu = 2.0 * ac / num_samples; - double sample_var = 0.0; - for (std::size_t i=0; i < out.size(); i++) { - int h = g[2 * i] + g[2 * i + 1]; - double d = h - mu; - sample_var += d * d; + // Per-sample accumulated weight matrix: W[i * ni + j] = weight for sample i at interval j + // Initialized to x[i] for all intervals, then accumulated bottom-up. + std::vector W(static_cast(n) * ni); +#ifdef _OPENMP + #pragma omp parallel for schedule(static) +#endif + for (int i = 0; i < n; i++) { + const double xi = xp[i]; + double* row = W.data() + static_cast(i) * ni; + std::fill(row, row + ni, xi); + } + + std::vector out(m, 0.0); + + // Bottom-up: push weights from children to targets, accumulate corrections + for (int i = n - 1; i >= 1; i--) { + const int* mm_data = instructions[i].mismatches.data(); + const int n_mm = static_cast(instructions[i].mismatches.size()); + const int8_t* my_signs = &tree_mm_sign[tree_mm_offset[i]]; + const double* Wi = &W[static_cast(i) * ni]; + + const int seg_off = tree_seg_offset[i]; + int n_mapped_segs; + { + int next_off = (i + 1 < n) ? tree_seg_offset[i + 1] + : static_cast(tree_seg_first_ivl.size()) - 1; + n_mapped_segs = next_off - seg_off; + } + + for (int sk = 0; sk < n_mapped_segs; sk++) { + const int first_ivl = tree_seg_first_ivl[seg_off + sk]; + const int last_ivl = (sk + 1 < n_mapped_segs) + ? tree_seg_first_ivl[seg_off + sk + 1] + : ni; + const int seg = tree_ivl_seg[static_cast(i) * ni + first_ivl]; + const int target = instructions[i].targets[seg]; + const int site_a = tree_ivl_start[first_ivl]; + const int site_b = tree_ivl_end[last_ivl - 1]; + + if (target != i) { + // Push accumulated weight to target + double* Wt = &W[static_cast(target) * ni]; + for (int j = first_ivl; j < last_ivl; j++) { + Wt[j] += Wi[j]; } - sample_var /= (num_samples / 2); - const double std = std::sqrt(sample_var); - const double w = x[site_counter] / std; - for (std::size_t i=0; i < out.size(); i++) { - const int h = g[2 * i] + g[2 * i + 1]; - out[i] += w * (h - mu); + // Mismatch corrections: out[mm] += W[i][mm_ivl] * sign + const int* lo = std::lower_bound(mm_data, mm_data + n_mm, site_a); + const int* hi = std::lower_bound(mm_data, mm_data + n_mm, site_b); + const int* mm_ivl_data = tree_mm_ivl.data() + tree_mm_offset[i]; + for (const int* it = lo; it != hi; ++it) { + const int k = static_cast(it - mm_data); + out[*it] += Wi[mm_ivl_data[k]] * my_signs[k]; } - site_counter++; - } - } else { - while (gi.has_next_genotype()) { - // Fetch the next genotype - const std::vector& g = gi.next_genotype(); - const double w = x[site_counter]; - for (std::size_t i=0; i < out.size(); i++) { - const int h = g[2 * i] + g[2 * i + 1]; - out[i] += w * h; + } else { + // Self-ref: contribute to out using carry-run decomposition. + // Process contiguous carry=1 runs with vectorizable range-add, + // jumping over carry=0 runs entirely. This avoids the per-site + // branch in the naive loop and lets the compiler SIMD-vectorize + // the hot inner fill when carry=1. + for (int j = first_ivl; j < last_ivl; j++) { + const int a = tree_ivl_start[j]; + const int b = tree_ivl_end[j]; + int carry = (j == 0) ? 0 + : static_cast(tree_carry[static_cast(i) * ni + j - 1]); + const int* lo = std::lower_bound(mm_data, mm_data + n_mm, a); + const int* hi_mm = std::lower_bound(mm_data, mm_data + n_mm, b); + const double wi = Wi[j]; + int prev = a; + double* outp = out.data(); + for (const int* it = lo; it != hi_mm; ++it) { + const int mm_site = *it; + if (carry) { + double* __restrict__ op = outp + prev; + const int len = mm_site - prev; + for (int s = 0; s < len; s++) op[s] += wi; + } + carry ^= 1; + outp[mm_site] += wi * carry; + prev = mm_site + 1; + } + if (carry && b > prev) { + double* __restrict__ op = outp + prev; + const int len = b - prev; + for (int s = 0; s < len; s++) op[s] += wi; + } } - site_counter++; } } - return out; - } else { - // Initialize output - std::vector out(num_samples, 0.0); - if (normalize) { - while (gi.has_next_genotype()) { - // Fetch the next genotype - const std::vector& g = gi.next_genotype(); - double ac = 0.0; - for (auto a : g) { - ac += a; - } + } - // Normalization constants - double mu = ac / num_samples; - double std = std::sqrt(mu * (1 - mu)); - const double w = x[site_counter] / std; - for (std::size_t i=0; i < out.size(); i++) { - out[i] += w * (g[i] - mu); + // Sample 0: multiply accumulated weight by reference genome + { + const double* W0 = &W[0]; + const int* ref = tree_ref_genome.data(); + for (int j = 0; j < ni; j++) { + const double w0j = W0[j]; + for (int s = tree_ivl_start[j]; s < tree_ivl_end[j]; s++) { + out[s] += w0j * ref[s]; + } + } + } + + return out; +} + + +// ═══════════════════════════════════════════════════════════════════════════ +// Batch tree multiply: process k vectors in a single tree traversal. +// Layout: row-major flat arrays, X[row * k + col]. +// ═══════════════════════════════════════════════════════════════════════════ + +std::vector ThreadingInstructions::right_multiply_tree_batch( + const std::vector& x_flat, int k) { + if (static_cast(x_flat.size()) != num_sites * k) + throw std::runtime_error("Input must have length num_sites * k"); + if (k == 1) return right_multiply_tree(x_flat); + + prepare_tree_multiply(); + const int n = num_samples, m = num_sites, ni = tree_n_intervals; + const double* xp = x_flat.data(); + const int* ref = tree_ref_genome.data(); + const int* ivl_s = tree_ivl_start.data(); + const int* ivl_e = tree_ivl_end.data(); + + // k prefix sums + std::vector prefix_x(static_cast(m + 1) * k, 0.0); + for (int s = 0; s < m; s++) { + const double* xs = &xp[s * k]; + const double* ps = &prefix_x[s * k]; + double* pd = &prefix_x[(s + 1) * k]; + for (int c = 0; c < k; c++) pd[c] = ps[c] + xs[c]; + } + + const size_t cum_stride = static_cast(ni + 1) * k; + std::vector cum_ref(n, 0); + { + std::vector seen(n, false); + for (int i = 1; i < n; i++) { + for (int t : instructions[i].targets) + if (t != i && !seen[t]) { seen[t] = true; cum_ref[t]++; } + for (int t : instructions[i].targets) + if (t != i) seen[t] = false; + } + } + std::vector> cum_pool; + std::vector free_rows, sample_to_row(n, -1); + auto alloc_row = [&]() -> int { + if (!free_rows.empty()) { + int r = free_rows.back(); free_rows.pop_back(); + std::fill(cum_pool[r].begin(), cum_pool[r].end(), 0.0); + return r; + } + int r = static_cast(cum_pool.size()); + cum_pool.emplace_back(cum_stride, 0.0); + return r; + }; + auto release_row = [&](int s) { + int r = sample_to_row[s]; + if (r >= 0) { free_rows.push_back(r); sample_to_row[s] = -1; } + }; + + // Sample 0 + { int r0 = alloc_row(); sample_to_row[0] = r0; + double* c0 = cum_pool[r0].data(); + for (int j = 0; j < ni; j++) { + double* dst = &c0[(j + 1) * k]; + const double* src = &c0[j * k]; + for (int c = 0; c < k; c++) dst[c] = src[c]; + for (int site = ivl_s[j]; site < ivl_e[j]; site++) + if (ref[site]) { const double* xs = &xp[site*k]; + for (int c = 0; c < k; c++) dst[c] += xs[c]; } + } + } + + std::vector out(static_cast(n) * k, 0.0); + { const double* c0e = &cum_pool[sample_to_row[0]][ni * k]; + for (int c = 0; c < k; c++) out[c] = c0e[c]; } + + for (int i = 1; i < n; i++) { + const int* mm_data = instructions[i].mismatches.data(); + const int n_mm = static_cast(instructions[i].mismatches.size()); + const int8_t* my_signs = &tree_mm_sign[tree_mm_offset[i]]; + const bool need_cum = (cum_ref[i] > 0); + int my_row = -1; + double* my_cum = nullptr; + if (need_cum) { + my_row = alloc_row(); sample_to_row[i] = my_row; + my_cum = cum_pool[my_row].data(); + } + const int seg_off = tree_seg_offset[i]; + int n_mapped_segs = ((i+1(tree_seg_first_ivl.size())-1) - seg_off; + double* sample_out = &out[i * k]; + + for (int sk = 0; sk < n_mapped_segs; sk++) { + const int fi = tree_seg_first_ivl[seg_off + sk]; + const int li = (sk+1(i)*ni + fi]; + const int target = instructions[i].targets[seg]; + const int sa = ivl_s[fi], sb = ivl_e[li-1]; + + if (target != i) { + const double* tc = cum_pool[sample_to_row[target]].data(); + const int* lo = std::lower_bound(mm_data, mm_data+n_mm, sa); + const int* hi = std::lower_bound(mm_data, mm_data+n_mm, sb); + + // Compute output: base from target cum + mismatch correction + for (int c = 0; c < k; c++) { + double base = tc[li*k+c] - tc[fi*k+c], corr = 0.0; + for (const int* it = lo; it != hi; ++it) + corr += xp[*it*k+c] * my_signs[it-mm_data]; + sample_out[c] += base + corr; + } + // Build cum only if someone will reference us + if (need_cum) { + const double* fi_cum = &my_cum[fi*k]; + const double* fi_tc = &tc[fi*k]; + for (int j = fi+1; j <= li; j++) + for (int c = 0; c < k; c++) + my_cum[j*k+c] = tc[j*k+c] + fi_cum[c] - fi_tc[c]; + // Apply mismatch corrections via deferred deltas + prefix sum + const int n_seg_ivls = li - fi; + if (lo != hi && n_seg_ivls > 0) { + std::vector mm_corr(static_cast(n_seg_ivls) * k, 0.0); + const int* miv = tree_mm_ivl.data() + tree_mm_offset[i]; + for (const int* it = lo; it != hi; ++it) { + int kk = static_cast(it-mm_data); + int mivl = miv[kk]; double sv = my_signs[kk]; + const double* xs = &xp[*it*k]; + int idx = mivl - fi; + if (idx >= 0 && idx < n_seg_ivls) { + double* d = &mm_corr[idx * k]; + for (int c = 0; c < k; c++) d[c] += xs[c]*sv; + } + } + double* running = &mm_corr[0]; + for (int c = 0; c < k; c++) + my_cum[(fi+1)*k+c] += running[c]; + for (int j = 1; j < n_seg_ivls; j++) { + double* cur = &mm_corr[j*k]; + const double* prv = &mm_corr[(j-1)*k]; + for (int c = 0; c < k; c++) { + cur[c] += prv[c]; + my_cum[(fi+1+j)*k+c] += cur[c]; + } + } + } + } + } else { + for (int j = fi; j < li; j++) { + int a = ivl_s[j], b = ivl_e[j]; + int carry = (j==0)?0:static_cast(tree_carry[static_cast(i)*ni+j-1]); + const int* lo = std::lower_bound(mm_data, mm_data+n_mm, a); + const int* hi = std::lower_bound(mm_data, mm_data+n_mm, b); + for (int c = 0; c < k; c++) { + double ivl_sum = 0.0; int cc = carry; int prev = a; + for (const int* it = lo; it != hi; ++it) { + int ms = *it; + if (ms > prev) ivl_sum += cc*(prefix_x[ms*k+c]-prefix_x[prev*k+c]); + cc = 1-cc; ivl_sum += xp[ms*k+c]*cc; prev = ms+1; + } + if (b > prev) ivl_sum += cc*(prefix_x[b*k+c]-prefix_x[prev*k+c]); + if (need_cum) my_cum[(j+1)*k+c] = my_cum[j*k+c]+ivl_sum; + sample_out[c] += ivl_sum; + } } - site_counter++; } - } else { - while (gi.has_next_genotype()) { - // Fetch the next genotype - const std::vector& g = gi.next_genotype(); - const double w = x[site_counter]; - for (std::size_t i=0; i < out.size(); i++) { - out[i] += w * g[i]; + } + { + const auto& tgts = instructions[i].targets; + const int nt = static_cast(tgts.size()); + for (int si = 0; si < nt; si++) { + int t = tgts[si]; + if (t == i) continue; + bool dup = false; + for (int sj = 0; sj < si; sj++) + if (tgts[sj] == t) { dup = true; break; } + if (!dup && --cum_ref[t] == 0) release_row(t); + } + } + if (cum_ref[i]==0) release_row(i); + } + return out; +} + + +std::vector ThreadingInstructions::left_multiply_tree_batch( + const std::vector& x_flat, int k) { + if (static_cast(x_flat.size()) != num_samples * k) + throw std::runtime_error("Input must have length num_samples * k"); + if (k == 1) return left_multiply_tree(x_flat); + + prepare_tree_multiply(); + const int n = num_samples, m = num_sites, ni = tree_n_intervals; + const double* xp = x_flat.data(); + + const size_t Wstride = static_cast(ni) * k; + std::vector W(static_cast(n) * Wstride); + for (int i = 0; i < n; i++) { + const double* xi = &xp[i * k]; + double* Wi = &W[i * Wstride]; + for (int j = 0; j < ni; j++) { + double* d = &Wi[j*k]; + for (int c = 0; c < k; c++) d[c] = xi[c]; + } + } + + std::vector out(static_cast(m) * k, 0.0); + + for (int i = n-1; i >= 1; i--) { + const int* mm_data = instructions[i].mismatches.data(); + const int n_mm = static_cast(instructions[i].mismatches.size()); + const int8_t* my_signs = &tree_mm_sign[tree_mm_offset[i]]; + const double* Wi = &W[i * Wstride]; + const int seg_off = tree_seg_offset[i]; + int n_mapped_segs = ((i+1(tree_seg_first_ivl.size())-1) - seg_off; + + for (int sk = 0; sk < n_mapped_segs; sk++) { + const int fi = tree_seg_first_ivl[seg_off+sk]; + const int li = (sk+1(i)*ni+fi]; + const int target = instructions[i].targets[seg]; + const int sa = tree_ivl_start[fi], sb = tree_ivl_end[li-1]; + + if (target != i) { + double* Wt = &W[static_cast(target)*Wstride]; + for (int j = fi; j < li; j++) { + const double* s = &Wi[j*k]; double* d = &Wt[j*k]; + for (int c = 0; c < k; c++) d[c] += s[c]; + } + const int* lo = std::lower_bound(mm_data, mm_data+n_mm, sa); + const int* hi = std::lower_bound(mm_data, mm_data+n_mm, sb); + const int* miv = tree_mm_ivl.data() + tree_mm_offset[i]; + for (const int* it = lo; it != hi; ++it) { + int kk = static_cast(it-mm_data); + double sv = my_signs[kk]; + const double* wi = &Wi[miv[kk]*k]; + double* od = &out[*it*k]; + for (int c = 0; c < k; c++) od[c] += wi[c]*sv; + } + } else { + for (int j = fi; j < li; j++) { + int a = tree_ivl_start[j], b = tree_ivl_end[j]; + int carry = (j==0)?0:static_cast(tree_carry[static_cast(i)*ni+j-1]); + const int* lo = std::lower_bound(mm_data, mm_data+n_mm, a); + const int* hi_mm = std::lower_bound(mm_data, mm_data+n_mm, b); + const double* wi = &Wi[j*k]; + int prev = a; + for (const int* it = lo; it != hi_mm; ++it) { + int ms = *it; + if (carry && ms > prev) + for (int s = prev; s < ms; s++) { + double* od = &out[s*k]; + for (int c = 0; c < k; c++) od[c] += wi[c]; + } + carry ^= 1; + if (carry) { double* od = &out[ms*k]; + for (int c = 0; c < k; c++) od[c] += wi[c]; } + prev = ms+1; + } + if (carry && b > prev) + for (int s = prev; s < b; s++) { + double* od = &out[s*k]; + for (int c = 0; c < k; c++) od[c] += wi[c]; + } } - site_counter++; } } - return out; } + + // Sample 0 + { const double* W0 = &W[0]; + const int* ref = tree_ref_genome.data(); + for (int j = 0; j < ni; j++) { + const double* w0j = &W0[j*k]; + for (int s = tree_ivl_start[j]; s < tree_ivl_end[j]; s++) + if (ref[s]) { double* od = &out[s*k]; + for (int c = 0; c < k; c++) od[c] += w0j[c]; } + } + } + return out; } diff --git a/src/ThreadingInstructions.hpp b/src/ThreadingInstructions.hpp index f404be2..12928cd 100644 --- a/src/ThreadingInstructions.hpp +++ b/src/ThreadingInstructions.hpp @@ -89,6 +89,30 @@ class ThreadingInstructions { std::vector left_multiply(const std::vector& x, bool diploid=false, bool normalize=false); std::vector right_multiply(const std::vector& x, bool diploid=false, bool normalize=false); + // Tree-propagation multiply: O(n * n_segments + total_mismatches) per call. + // One-time O(n * m) prepare step precomputes mismatch corrections using + // reference-counted genotype cache (peak memory ~O(m * tree_depth)). + // No genotype matrix retained after prepare. + void prepare_tree_multiply(); + std::vector right_multiply_tree(const std::vector& x); + std::vector left_multiply_tree(const std::vector& x); + + // Batch tree multiply: process k vectors in a single tree traversal. + // Input/output are row-major flat arrays: X[row * k + col]. + // right_multiply_tree_batch: X is (num_sites, k), returns (num_samples, k) + // left_multiply_tree_batch: X is (num_samples, k), returns (num_sites, k) + std::vector right_multiply_tree_batch(const std::vector& x_flat, int k); + std::vector left_multiply_tree_batch(const std::vector& x_flat, int k); + + // Precompute and cache the dense genotype matrix (num_sites × num_samples, row-major). + // Subsequent left_multiply/right_multiply calls use the cached matrix. + void materialize_genotypes(); + + // Precompute and cache standardized matrices for normalized multiply. + // These store (g - mean) / std as doubles, so multiply becomes a pure dot product. + void materialize_normalized_haploid(); + void materialize_normalized_diploid(); + public: int start = 0; int end = 0; @@ -96,6 +120,51 @@ class ThreadingInstructions { int num_sites = 0; std::vector positions; std::vector instructions; + +private: + // Cached dense genotype matrix: genotype_matrix[site * num_samples + sample] + std::vector genotype_matrix; + bool genotypes_materialized = false; + + // Cached diploid sum matrix: diploid_matrix[site * n_dip + i] = g[2i] + g[2i+1] + std::vector diploid_matrix; + bool diploid_materialized = false; + void materialize_diploid(); + + // Cached standardized matrices for normalized multiply + // standardized_hap[site * num_samples + i] = (g[i] - mu) / std + std::vector standardized_hap; + bool standardized_hap_ready = false; + // standardized_dip[site * n_dip + i] = ((g[2i]+g[2i+1]) - mu) / std + std::vector standardized_dip; + bool standardized_dip_ready = false; + + // Tree-propagation multiply cache + bool tree_ready = false; + int tree_n_intervals = 0; + std::vector tree_ivl_start; // first site index of each interval + std::vector tree_ivl_end; // one past last site index + std::vector tree_ivl_seg; // flat [sample * n_intervals + interval] -> segment + std::vector tree_ref_genome; + + // Precomputed mismatch correction signs: tree_mm_sign[offset + k] = +1 or -1 + // for non-self mismatches, 0 for self-ref (unused). Per-sample offsets in tree_mm_offset. + std::vector tree_mm_sign; + std::vector tree_mm_offset; // tree_mm_offset[i] = start index for sample i + + // Precomputed mismatch-to-interval mapping: tree_mm_ivl[offset + k] = interval index + // containing mismatch k for sample i. Eliminates binary search during multiply. + std::vector tree_mm_ivl; + + // Precomputed carry_geno at end of each interval for each sample. + // tree_carry[sample * n_intervals + interval] = genotype at last site of interval. + std::vector tree_carry; + + // Per-sample segment-to-interval mapping for segment-level loop. + // tree_seg_first_ivl[sample * max_segs + seg] = first interval index of segment. + // Stored as flat vectors indexed by tree_seg_offset[sample] + seg. + std::vector tree_seg_first_ivl; // first interval of each segment + std::vector tree_seg_offset; // offset into tree_seg_first_ivl for each sample }; #endif // THREADS_ARG_THREADING_INSTRUCTIONS_HPP diff --git a/src/ThreadsFastLS.cpp b/src/ThreadsFastLS.cpp index abf75c6..16c9540 100644 --- a/src/ThreadsFastLS.cpp +++ b/src/ThreadsFastLS.cpp @@ -21,7 +21,7 @@ #include #include #include -#include +#include #include #include #include @@ -40,7 +40,8 @@ const int END_ALLELE = 0; const int HMM_SPLIT_THRESHOLD = 1000; inline std::size_t pair_key(int i, int j) { - return (static_cast(i) << 32) | static_cast(j); + return (static_cast(static_cast(i)) << 32) | + static_cast(static_cast(j)); } } // namespace @@ -56,16 +57,14 @@ ThreadsFastLS::ThreadsFastLS(std::vector _physical_positions, physical_positions(_physical_positions), genetic_positions(_genetic_positions), demography(Demography(ne, ne_times)) { if (physical_positions.size() != genetic_positions.size()) { - std::cerr << "Map lengths don't match.\n"; - exit(1); + throw std::runtime_error("Map lengths don't match."); } else if (physical_positions.size() <= 2) { - std::cerr << "Need at least 3 sites, found " << physical_positions.size() << std::endl; - exit(1); + throw std::runtime_error("Need at least 3 sites, found " + + std::to_string(physical_positions.size())); } if (mutation_rate <= 0) { - std::cerr << "Need a strictly positive mutation rate.\n"; - exit(1); + throw std::runtime_error("Need a strictly positive mutation rate."); } num_sites = static_cast(physical_positions.size()); num_samples = 0; @@ -73,14 +72,14 @@ ThreadsFastLS::ThreadsFastLS(std::vector _physical_positions, #ifdef THREADS_FAST_LS_CHECK_IN_ORDER for (int i = 0; i < num_sites - 1; i++) { if (physical_positions[i + 1] <= physical_positions[i]) { - cerr << "Physical positions must be strictly increasing, found "; - cerr << physical_positions[i + 1] << " after " << physical_positions[i] << endl; - exit(1); + throw std::runtime_error("Physical positions must be strictly increasing, found " + + std::to_string(physical_positions[i + 1]) + " after " + + std::to_string(physical_positions[i])); } if (genetic_positions[i + 1] <= genetic_positions[i]) { - cerr << "Genetic coordinates must be strictly increasing, found "; - cerr << genetic_positions[i + 1] << " after " << genetic_positions[i] << endl; - exit(1); + throw std::runtime_error("Genetic coordinates must be strictly increasing, found " + + std::to_string(genetic_positions[i + 1]) + " after " + + std::to_string(genetic_positions[i])); } } #endif // THREADS_FAST_LS_CHECK_IN_ORDER @@ -108,9 +107,8 @@ ThreadsFastLS::ThreadsFastLS(std::vector _physical_positions, } } if (trim_pos_start_idx >= trim_pos_end_idx - 3) { - std::cerr << "Too few positions left after applying burn-in, need at least 3. Aborting." - << std::endl; - exit(1); + throw std::runtime_error( + "Too few positions left after applying burn-in, need at least 3."); } // Initialize both ends of the linked-list columns @@ -135,15 +133,12 @@ ThreadsFastLS::ThreadsFastLS(std::vector _physical_positions, std::tie(cm_boundaries, cm_sizes) = site_sizes(genetic_positions); if (use_hmm) { - hmm = new HMM(demography, bp_sizes, cm_sizes, mutation_rate, 64); - } - else { - hmm = nullptr; + hmm = std::make_unique(demography, bp_sizes, cm_sizes, mutation_rate, 64); } } std::tuple, std::vector> -ThreadsFastLS::site_sizes(std::vector positions) { +ThreadsFastLS::site_sizes(const std::vector& positions) { // Find mid-points between sites std::size_t M = positions.size(); std::vector pos_means(M - 1); @@ -161,8 +156,7 @@ ThreadsFastLS::site_sizes(std::vector positions) { site_sizes[M - 1] = mean_size; for (double s : site_sizes) { if (s < 0) { - std::cerr << "Found negative site size " << s << std::endl; - exit(1); + throw std::runtime_error("Found negative site size " + std::to_string(s)); } } std::vector boundaries(M + 1); @@ -182,7 +176,7 @@ std::vector ThreadsFastLS::trimmed_positions() const { void ThreadsFastLS::delete_hmm() { if (use_hmm) { - delete hmm; + hmm.reset(); use_hmm = false; } } @@ -206,12 +200,10 @@ void ThreadsFastLS::insert(const std::vector& genotype) { void ThreadsFastLS::insert(const int ID, const std::vector& genotype) { if (ID_map.find(ID) != ID_map.end()) { - std::cerr << "ID " << ID << " is already in the panel.\n"; - exit(1); + throw std::runtime_error("ID " + std::to_string(ID) + " is already in the panel."); } if (static_cast(genotype.size()) != num_sites) { - std::cerr << "Number of input markers does not match map.\n"; - exit(1); + throw std::runtime_error("Number of input markers does not match map."); } int insert_index = num_samples; ID_map[ID] = insert_index; @@ -377,8 +369,7 @@ std::pair ThreadsFastLS::fastLS(const std::vector& int n_states = static_cast(current_states.size()); max_states = std::max(n_states, max_states); if (n_states == 0) { - std::cerr << "No states left on stack, something is messed up in the algorithm.\n"; - exit(1); + throw std::runtime_error("No states left on stack, something is messed up in the algorithm."); } // Heuristically get a bound on states we want to add @@ -418,8 +409,8 @@ std::pair ThreadsFastLS::fastLS(const std::vector& } if (new_states.size() == 0) { - std::cerr << "The algorithm is in an illegal state because no new_states were created.\n"; - exit(1); + throw std::runtime_error( + "The algorithm is in an illegal state because no new_states were created."); } // Find a best state in the current layer and recombine. @@ -428,9 +419,10 @@ std::pair ThreadsFastLS::fastLS(const std::vector& [](const auto& s1, const auto& s2) { return s1.score < s2.score; })); if (best_extension.score < z - 0.001 || best_extension.score > z + 0.001) { - std::cerr << "The algorithm is in an illegal state because z != best_extension.score, found "; - std::cerr << "best_extension.score=" << best_extension.score << " and z=" << z << std::endl; - exit(1); + throw std::runtime_error( + "The algorithm is in an illegal state because z != best_extension.score, found " + "best_extension.score=" + std::to_string(best_extension.score) + + " and z=" + std::to_string(z)); } // Add the new recombinant state to the stack (we never enter this clause on the first @@ -510,8 +502,8 @@ ThreadsFastLS::fastLS_diploid(const std::vector& genotype) { std::vector new_pairs; max_state_pairs = std::max(n_state_pairs, max_state_pairs); if (n_state_pairs == 0) { - std::cerr << "No state pairs left on stack, something is messed up in the algorithm.\n"; - exit(1); + throw std::runtime_error( + "No state pairs left on stack, something is messed up in the algorithm."); } // Heuristically get a bound on states we want to add, @@ -556,8 +548,8 @@ ThreadsFastLS::fastLS_diploid(const std::vector& genotype) { } } else { - std::cerr << "Only 0, 1, 2-alleles allowed." << std::endl; - exit(1); + throw std::runtime_error("Only 0, 1, 2-alleles allowed, found " + + std::to_string(allele)); } // Set local minima, this maps (anchor, traceback) to a score @@ -883,8 +875,8 @@ ThreadsFastLS::fastLS_diploid(const std::vector& genotype) { // END OF EXTENSION LOOP if (new_pairs.size() == 0) { - std::cerr << "The algorithm is in an illegal state because no new_states were created.\n"; - exit(1); + throw std::runtime_error( + "The algorithm is in an illegal state because no new_states were created."); } // SINGLE RECOMBINATION EVENTS @@ -939,9 +931,10 @@ ThreadsFastLS::fastLS_diploid(const std::vector& genotype) { z = double_recombinant_score; } if (std::abs(best_pair.score - z) > 0.0001) { - std::cerr << "The algorithm is in an illegal state because z != best_pair.score, found "; - std::cerr << "best_pair.score=" << best_pair.score << " and z=" << z << std::endl; - exit(1); + throw std::runtime_error( + "The algorithm is in an illegal state because z != best_pair.score, found " + "best_pair.score=" + std::to_string(best_pair.score) + + " and z=" + std::to_string(z)); } current_pairs = new_pairs; new_pairs.clear(); @@ -1183,8 +1176,8 @@ ThreadsFastLS::recombination_penalties_correct() { double ThreadsFastLS::date_segment(const int num_het_sites, const int start, const int end) { if (start > end) { - std::cerr << "Can't date a segment with length <= 0\n"; - exit(1); + throw std::runtime_error("Can't date a segment with start > end (" + + std::to_string(start) + " > " + std::to_string(end) + ")"); } double bp_size = 0; double cm_size = 0; @@ -1505,13 +1498,11 @@ std::pair ThreadsFastLS::overflow_region(const std::vector& geno std::vector ThreadsFastLS::fetch_het_hom_sites(const int id1, const int id2, const int start, const int end) { if (ID_map.find(id1) == ID_map.end()) { - std::cerr << "fetch_het_hom_sites bad id1 " << id1 << std::endl; - exit(1); + throw std::runtime_error("fetch_het_hom_sites bad id1 " + std::to_string(id1)); } if (ID_map.find(id2) == ID_map.end()) { - std::cerr << "fetch_het_hom_sites bad id2 " << id2 << std::endl; - exit(1); + throw std::runtime_error("fetch_het_hom_sites bad id2 " + std::to_string(id2)); } std::vector het_hom_sites(end - start); for (int i = start; i < end; i++) { @@ -1534,8 +1525,8 @@ ThreadsFastLS::het_sites_from_thread(const int focal_ID, const std::vector int segment_end = seg_i == num_segments - 1 ? (static_cast(physical_positions.back()) + 1) : bp_starts[seg_i + 1]; int target_ID = target_IDs[seg_i][0]; - while (segment_start <= physical_positions[site_i] && - physical_positions[site_i] < segment_end && site_i < num_sites) { + while (site_i < num_sites && segment_start <= physical_positions[site_i] && + physical_positions[site_i] < segment_end) { if (panel[ID_map.at(focal_ID)][site_i]->genotype != panel[ID_map.at(target_ID)][site_i]->genotype) { het_sites.push_back(static_cast(physical_positions[site_i])); @@ -1544,8 +1535,8 @@ ThreadsFastLS::het_sites_from_thread(const int focal_ID, const std::vector } } if (site_i != num_sites) { - std::cerr << "Found " << site_i + 1 << " sites, expected " << num_sites << std::endl; - exit(1); + throw std::runtime_error("Found " + std::to_string(site_i + 1) + + " sites, expected " + std::to_string(num_sites)); } return het_sites; } diff --git a/src/ThreadsFastLS.hpp b/src/ThreadsFastLS.hpp index fa49b8c..53b2738 100644 --- a/src/ThreadsFastLS.hpp +++ b/src/ThreadsFastLS.hpp @@ -86,7 +86,7 @@ class ThreadsFastLS { std::vector> target_IDs); static std::tuple, std::vector> - site_sizes(std::vector positions); + site_sizes(const std::vector& positions); // More attributes std::vector trimmed_positions() const; @@ -194,7 +194,7 @@ class ThreadsFastLS { std::vector bp_boundaries; std::vector cm_boundaries; Demography demography; - HMM* hmm = nullptr; + std::unique_ptr hmm; // The dynamic reference panel std::vector>> panel; diff --git a/src/ThreadsIO.cpp b/src/ThreadsIO.cpp new file mode 100644 index 0000000..067d7d3 --- /dev/null +++ b/src/ThreadsIO.cpp @@ -0,0 +1,338 @@ +#include "ThreadsIO.hpp" + +#include +#include +#include +#include +#include +#include + +// ── helpers ────────────────────────────────────────────────────────────────── + +namespace { + +struct H5Handle { + hid_t id; + void (*closer)(hid_t); + H5Handle(hid_t id, void (*closer)(hid_t)) : id(id), closer(closer) { + if (id < 0) throw std::runtime_error("HDF5 handle creation failed"); + } + ~H5Handle() { if (id >= 0) closer(id); } + operator hid_t() const { return id; } + H5Handle(const H5Handle&) = delete; + H5Handle& operator=(const H5Handle&) = delete; +}; + +void close_file(hid_t h) { H5Fclose(h); } +void close_space(hid_t h) { H5Sclose(h); } +void close_dset(hid_t h) { H5Dclose(h); } +void close_plist(hid_t h) { H5Pclose(h); } +void close_type(hid_t h) { H5Tclose(h); } +void close_attr(hid_t h) { H5Aclose(h); } + +hid_t make_dcpl(int rank, const hsize_t* dims) { + // Cannot chunk/compress empty datasets + for (int i = 0; i < rank; i++) + if (dims[i] == 0) return H5P_DEFAULT; + hid_t dcpl = H5Pcreate(H5P_DATASET_CREATE); + H5Pset_chunk(dcpl, rank, dims); + H5Pset_deflate(dcpl, 9); + return dcpl; +} + +void write_i64_1d(hid_t file, const char* name, const std::vector& data) { + hsize_t n = data.size(); + H5Handle space(H5Screate_simple(1, &n, nullptr), close_space); + H5Handle dcpl(make_dcpl(1, &n), close_plist); + H5Handle ds(H5Dcreate2(file, name, H5T_STD_I64LE, space, H5P_DEFAULT, dcpl, H5P_DEFAULT), close_dset); + H5Dwrite(ds, H5T_NATIVE_INT64, H5S_ALL, H5S_ALL, H5P_DEFAULT, data.data()); +} + +void write_f64_1d(hid_t file, const char* name, const std::vector& data) { + hsize_t n = data.size(); + H5Handle space(H5Screate_simple(1, &n, nullptr), close_space); + H5Handle dcpl(make_dcpl(1, &n), close_plist); + H5Handle ds(H5Dcreate2(file, name, H5T_IEEE_F64LE, space, H5P_DEFAULT, dcpl, H5P_DEFAULT), close_dset); + H5Dwrite(ds, H5T_NATIVE_DOUBLE, H5S_ALL, H5S_ALL, H5P_DEFAULT, data.data()); +} + +void write_i64_2d(hid_t file, const char* name, const int64_t* data, hsize_t rows, hsize_t cols) { + hsize_t dims[2] = {rows, cols}; + H5Handle space(H5Screate_simple(2, dims, nullptr), close_space); + H5Handle dcpl(make_dcpl(2, dims), close_plist); + H5Handle ds(H5Dcreate2(file, name, H5T_STD_I64LE, space, H5P_DEFAULT, dcpl, H5P_DEFAULT), close_dset); + H5Dwrite(ds, H5T_NATIVE_INT64, H5S_ALL, H5S_ALL, H5P_DEFAULT, data); +} + +void write_string_2d(hid_t file, const char* name, const std::vector>& cols, hsize_t rows) { + hsize_t ncols = cols.size(); + hsize_t dims[2] = {rows, ncols}; + H5Handle space(H5Screate_simple(2, dims, nullptr), close_space); + H5Handle dcpl(make_dcpl(2, dims), close_plist); + H5Handle strtype(H5Tcopy(H5T_C_S1), close_type); + H5Tset_size(strtype, H5T_VARIABLE); + H5Handle ds(H5Dcreate2(file, name, strtype, space, H5P_DEFAULT, dcpl, H5P_DEFAULT), close_dset); + + // Build flat row-major array of const char* + std::vector ptrs(rows * ncols); + for (hsize_t r = 0; r < rows; r++) + for (hsize_t c = 0; c < ncols; c++) + ptrs[r * ncols + c] = cols[c][r].c_str(); + H5Dwrite(ds, strtype, H5S_ALL, H5S_ALL, H5P_DEFAULT, ptrs.data()); +} + +void write_string_1d(hid_t file, const char* name, const std::vector& data) { + hsize_t n = data.size(); + H5Handle space(H5Screate_simple(1, &n, nullptr), close_space); + H5Handle dcpl(make_dcpl(1, &n), close_plist); + H5Handle strtype(H5Tcopy(H5T_C_S1), close_type); + H5Tset_size(strtype, H5T_VARIABLE); + H5Handle ds(H5Dcreate2(file, name, strtype, space, H5P_DEFAULT, dcpl, H5P_DEFAULT), close_dset); + std::vector ptrs(n); + for (hsize_t i = 0; i < n; i++) ptrs[i] = data[i].c_str(); + H5Dwrite(ds, strtype, H5S_ALL, H5S_ALL, H5P_DEFAULT, ptrs.data()); +} + +std::vector read_i64(hid_t file, const char* name) { + H5Handle ds(H5Dopen2(file, name, H5P_DEFAULT), close_dset); + H5Handle space(H5Dget_space(ds), close_space); + hsize_t dims[2]; + int ndims = H5Sget_simple_extent_dims(space, dims, nullptr); + hsize_t total = dims[0]; + if (ndims == 2) total *= dims[1]; + std::vector data(total); + H5Dread(ds, H5T_NATIVE_INT64, H5S_ALL, H5S_ALL, H5P_DEFAULT, data.data()); + return data; +} + +std::vector read_f64(hid_t file, const char* name) { + H5Handle ds(H5Dopen2(file, name, H5P_DEFAULT), close_dset); + H5Handle space(H5Dget_space(ds), close_space); + hsize_t dims[2]; + int ndims = H5Sget_simple_extent_dims(space, dims, nullptr); + hsize_t total = dims[0]; + if (ndims == 2) total *= dims[1]; + std::vector data(total); + H5Dread(ds, H5T_NATIVE_DOUBLE, H5S_ALL, H5S_ALL, H5P_DEFAULT, data.data()); + return data; +} + +std::vector read_varlen_strings(hid_t file, const char* name) { + H5Handle ds(H5Dopen2(file, name, H5P_DEFAULT), close_dset); + H5Handle space(H5Dget_space(ds), close_space); + hsize_t dims[2]; + int ndims = H5Sget_simple_extent_dims(space, dims, nullptr); + hsize_t total = dims[0]; + if (ndims == 2) total *= dims[1]; + + H5Handle memtype(H5Tcopy(H5T_C_S1), close_type); + H5Tset_size(memtype, H5T_VARIABLE); + + std::vector raw(total); + H5Dread(ds, memtype, H5S_ALL, H5S_ALL, H5P_DEFAULT, raw.data()); + + std::vector result(total); + for (hsize_t i = 0; i < total; i++) { + result[i] = raw[i] ? raw[i] : ""; + } + + // Reclaim memory allocated by HDF5 + H5Handle filetype(H5Dget_type(ds), close_type); + H5Dvlen_reclaim(memtype, space, H5P_DEFAULT, raw.data()); + return result; +} + +hsize_t get_dim0(hid_t file, const char* name) { + H5Handle ds(H5Dopen2(file, name, H5P_DEFAULT), close_dset); + H5Handle space(H5Dget_space(ds), close_space); + hsize_t dims[2]; + H5Sget_simple_extent_dims(space, dims, nullptr); + return dims[0]; +} + +std::string iso_now() { + auto now = std::chrono::system_clock::now(); + auto t = std::chrono::system_clock::to_time_t(now); + char buf[64]; + std::strftime(buf, sizeof(buf), "%Y-%m-%dT%H:%M:%S", std::localtime(&t)); + return buf; +} + +} // anonymous namespace + +// ── public API ─────────────────────────────────────────────────────────────── + +void serialize_threads( + const std::string& filename, + ThreadingInstructions& instructions, + const std::vector>& metadata_cols, + const std::vector& allele_ages, + const std::vector& sample_names) +{ + int N = instructions.num_samples; + int M = instructions.num_sites; + + auto all_starts = instructions.all_starts(); + auto all_targets = instructions.all_targets(); + auto all_tmrcas = instructions.all_tmrcas(); + auto all_mismatches = instructions.all_mismatches(); + + // Build samples array (N x 3) and flatten thread/mismatch data + std::vector samples_flat(N * 3); + std::vector flat_targets, flat_starts, flat_mismatches; + std::vector flat_tmrcas; + + int64_t toff = 0, moff = 0; + for (int i = 0; i < N; i++) { + samples_flat[i * 3 + 0] = i; + samples_flat[i * 3 + 1] = toff; + samples_flat[i * 3 + 2] = moff; + for (auto v : all_targets[i]) flat_targets.push_back(v); + for (auto v : all_starts[i]) flat_starts.push_back(v); + for (auto v : all_tmrcas[i]) flat_tmrcas.push_back(v); + for (auto v : all_mismatches[i]) flat_mismatches.push_back(v); + toff += all_starts[i].size(); + moff += all_mismatches[i].size(); + } + + // Build thread_targets 2D array (S x 2): [target, start] + size_t S = flat_targets.size(); + std::vector targets_2d(S * 2); + for (size_t i = 0; i < S; i++) { + targets_2d[i * 2 + 0] = flat_targets[i]; + targets_2d[i * 2 + 1] = flat_starts[i]; + } + + std::vector arg_range = { + static_cast(instructions.start), + static_cast(instructions.end) + }; + + std::vector positions_i64(instructions.positions.begin(), instructions.positions.end()); + std::vector mismatches_i64(flat_mismatches.begin(), flat_mismatches.end()); + + // Write file + H5Handle file(H5Fcreate(filename.c_str(), H5F_ACC_TRUNC, H5P_DEFAULT, H5P_DEFAULT), close_file); + + // datetime attribute + std::string dt = iso_now(); + { + H5Handle strtype(H5Tcopy(H5T_C_S1), close_type); + H5Tset_size(strtype, H5T_VARIABLE); + H5Handle aspace(H5Screate(H5S_SCALAR), close_space); + H5Handle attr(H5Acreate2(file, "datetime_created", strtype, aspace, H5P_DEFAULT, H5P_DEFAULT), close_attr); + const char* dtp = dt.c_str(); + H5Awrite(attr, strtype, &dtp); + } + + write_i64_2d(file, "samples", samples_flat.data(), N, 3); + write_i64_1d(file, "positions", positions_i64); + write_i64_2d(file, "thread_targets", targets_2d.data(), S, 2); + write_f64_1d(file, "thread_ages", flat_tmrcas); + write_i64_1d(file, "het_sites", mismatches_i64); + write_f64_1d(file, "arg_range", arg_range); + + // Optional datasets + if (!metadata_cols.empty() && M > 0) { + write_string_2d(file, "variant_metadata", metadata_cols, M); + } + if (!allele_ages.empty()) { + write_f64_1d(file, "allele_ages", allele_ages); + } + if (!sample_names.empty()) { + write_string_1d(file, "sample_names", sample_names); + } +} + + +ThreadingInstructions deserialize_threads(const std::string& filename) { + H5Handle file(H5Fopen(filename.c_str(), H5F_ACC_RDONLY, H5P_DEFAULT), close_file); + + auto samples_flat = read_i64(file, "samples"); + auto positions_i64 = read_i64(file, "positions"); + auto targets_flat = read_i64(file, "thread_targets"); + auto tmrcas_flat = read_f64(file, "thread_ages"); + auto mismatches_flat = read_i64(file, "het_sites"); + + hsize_t N = get_dim0(file, "samples"); + + // Read arg_range (may not exist in very old files) + double region_start, region_end; + if (H5Lexists(file, "arg_range", H5P_DEFAULT) > 0) { + auto ar = read_f64(file, "arg_range"); + region_start = ar[0]; + region_end = ar[1]; + } else { + region_start = 0; + region_end = 0; + } + + // Extract per-sample thread offsets and het offsets from samples array + std::vector thread_starts(N), het_starts(N); + for (hsize_t i = 0; i < N; i++) { + thread_starts[i] = samples_flat[i * 3 + 1]; + het_starts[i] = samples_flat[i * 3 + 2]; + } + + hsize_t total_threads = targets_flat.size() / 2; + hsize_t total_hets = mismatches_flat.size(); + + // Split flat arrays into per-sample vectors + std::vector> starts(N), targets(N), mismatches(N); + std::vector> tmrcas(N); + std::vector positions(positions_i64.begin(), positions_i64.end()); + + for (hsize_t i = 0; i < N; i++) { + int64_t t_start = thread_starts[i]; + int64_t t_end = (i + 1 < N) ? thread_starts[i + 1] : static_cast(total_threads); + int64_t h_start = het_starts[i]; + int64_t h_end = (i + 1 < N) ? het_starts[i + 1] : static_cast(total_hets); + + for (int64_t j = t_start; j < t_end; j++) { + targets[i].push_back(static_cast(targets_flat[j * 2])); + starts[i].push_back(static_cast(targets_flat[j * 2 + 1])); + tmrcas[i].push_back(tmrcas_flat[j]); + } + for (int64_t j = h_start; j < h_end; j++) { + mismatches[i].push_back(static_cast(mismatches_flat[j])); + } + } + + return ThreadingInstructions( + starts, tmrcas, targets, mismatches, + positions, + static_cast(region_start), + static_cast(region_end) + ); +} + + +std::vector> read_threads_metadata(const std::string& filename) { + H5Handle file(H5Fopen(filename.c_str(), H5F_ACC_RDONLY, H5P_DEFAULT), close_file); + if (H5Lexists(file, "variant_metadata", H5P_DEFAULT) <= 0) { + throw std::runtime_error("No variant_metadata in " + filename); + } + + hsize_t rows = get_dim0(file, "variant_metadata"); + auto flat = read_varlen_strings(file, "variant_metadata"); + + // Unflatten: flat is row-major (rows x 7) → 7 columns + const int ncols = 7; + std::vector> cols(ncols); + for (int c = 0; c < ncols; c++) { + cols[c].resize(rows); + for (hsize_t r = 0; r < rows; r++) { + cols[c][r] = flat[r * ncols + c]; + } + } + return cols; +} + + +std::vector read_threads_sample_names(const std::string& filename) { + H5Handle file(H5Fopen(filename.c_str(), H5F_ACC_RDONLY, H5P_DEFAULT), close_file); + if (H5Lexists(file, "sample_names", H5P_DEFAULT) <= 0) { + throw std::runtime_error("No sample_names in " + filename); + } + return read_varlen_strings(file, "sample_names"); +} diff --git a/src/ThreadsIO.hpp b/src/ThreadsIO.hpp new file mode 100644 index 0000000..a8c47ab --- /dev/null +++ b/src/ThreadsIO.hpp @@ -0,0 +1,21 @@ +#ifndef THREADS_ARG_THREADS_IO_HPP +#define THREADS_ARG_THREADS_IO_HPP + +#include "ThreadingInstructions.hpp" +#include +#include + +void serialize_threads( + const std::string& filename, + ThreadingInstructions& instructions, + const std::vector>& metadata_cols, + const std::vector& allele_ages, + const std::vector& sample_names); + +ThreadingInstructions deserialize_threads(const std::string& filename); + +// Read optional string datasets +std::vector> read_threads_metadata(const std::string& filename); +std::vector read_threads_sample_names(const std::string& filename); + +#endif diff --git a/src/ThreadsLowMem.cpp b/src/ThreadsLowMem.cpp index 2ed8ca9..e578102 100644 --- a/src/ThreadsLowMem.cpp +++ b/src/ThreadsLowMem.cpp @@ -17,13 +17,16 @@ #include "ThreadsLowMem.hpp" #include -#include -#include +#include #include #include #include #include +#ifdef _OPENMP +#include +#endif + ThreadsLowMem::ThreadsLowMem(const std::vector _target_ids, const std::vector& _physical_positions, const std::vector& _genetic_positions, std::vector ne, @@ -66,9 +69,17 @@ ThreadsLowMem::ThreadsLowMem(const std::vector _target_ids, // Mean interval size in base-pairs mean_bp_size = (physical_positions.back() - physical_positions[0]) / static_cast(num_sites - 1); + + // Build flat vectors for the hot path for (int target_id : target_ids) { - segment_indices[target_id] = 0; - expected_branch_lengths[target_id] = demography.expected_branch_length(target_id + 1); + double t = demography.expected_branch_length(target_id + 1); + expected_branch_lengths[target_id] = t; + if (target_id != 0) { + active_target_ids.push_back(target_id); + branch_lengths_vec.push_back(t); + log_target_ids_vec.push_back(std::log(static_cast(target_id))); + segment_indices_vec.push_back(0); + } } // Site counters @@ -86,6 +97,14 @@ ThreadsLowMem::ThreadsLowMem(const std::vector _target_ids, } } + // Precompute per-site HMM parameters (avoids recomputing per target) + k_per_site.resize(num_sites); + l_per_site.resize(num_sites); + for (int i = 0; i < num_sites; i++) { + k_per_site[i] = 2.0 * 0.01 * cm_sizes[i]; + l_per_site[i] = 2.0 * mutation_rate * bp_sizes[i]; + } + // Initialize the psmc-like segment-breaking algorithm int min_target_id = *(std::min_element(target_ids.begin(), target_ids.end())); if (min_target_id < n_hmm_samples) { @@ -109,89 +128,322 @@ void ThreadsLowMem::initialize_viterbi(std::vector sample_ids(match_groups.at(0).match_candidates.at(target_id).begin(), match_groups.at(0).match_candidates.at(target_id).end()); - hmms.emplace(target_id, ViterbiState(target_id, sample_ids)); + hmm_vec.push_back(std::make_unique(target_id, sample_ids)); + hmm_ptrs.push_back(hmm_vec.back().get()); } } -// Pass genotypes for a single site through the intialized Threads-Viterbi instances -void ThreadsLowMem::process_site_viterbi(const std::vector& genotype) { +// Internal: process one site from raw pointer (no vector copy) +// Used by the per-site API; batch methods use their own parallel loops +void ThreadsLowMem::process_site_viterbi_raw(const int* genotype) { bool group_change = false; if (match_group_idx < (static_cast(match_groups.size()) - 1) && - (genetic_positions.at(hmm_sites_processed) >= - match_groups.at(match_group_idx + 1).cm_position)) { + (genetic_positions[hmm_sites_processed] >= + match_groups[match_group_idx + 1].cm_position)) { match_group_idx++; group_change = true; } - double k = 2. * 0.01 * cm_sizes.at(hmm_sites_processed); - double l = 2. * mutation_rate * bp_sizes.at(hmm_sites_processed); - for (int target_id : target_ids) { - if (target_id == 0) { - continue; - } + const double k = k_per_site[hmm_sites_processed]; + const double l = l_per_site[hmm_sites_processed]; + const int n_active = static_cast(active_target_ids.size()); + for (int idx = 0; idx < n_active; ++idx) { + const int target_id = active_target_ids[idx]; if (group_change) { - hmms.at(target_id).set_samples( - match_groups.at(match_group_idx).match_candidates.at(target_id)); + hmm_ptrs[idx]->set_samples( + match_groups[match_group_idx].match_candidates.at(target_id)); } - double t = expected_branch_lengths.at(target_id); - double rho_c = k * t; - double rho = sparse ? -std::log1p(-std::exp(-(k * t))) - : -(std::log1p(-std::exp(-(k * t))) - std::log(target_id)); - double mu_c = l * t; - double mu = -std::log1p(-std::exp(-(l * t))); - hmms.at(target_id).process_site(genotype, rho, rho_c, mu, mu_c); + const double t = branch_lengths_vec[idx]; + const double kt = k * t; + const double lt = l * t; + const double rho_c = kt; + const double log1p_rho = -std::log1p(-std::exp(-kt)); + const double rho = sparse ? log1p_rho : log1p_rho + log_target_ids_vec[idx]; + const double mu_c = lt; + const double mu = -std::log1p(-std::exp(-lt)); + hmm_ptrs[idx]->process_site(genotype, rho, rho_c, mu, mu_c); } hmm_sites_processed++; - return; +} + +// Pass genotypes for a single site through the initialized Threads-Viterbi instances +void ThreadsLowMem::process_site_viterbi(const std::vector& genotype) { + process_site_viterbi_raw(genotype.data()); +} + +// Batch Viterbi: parallelized across targets — each thread processes all sites for its target +void ThreadsLowMem::process_all_sites_viterbi(const std::vector>& genotypes) { + const int prune_interval = 500; + const int n_active = static_cast(active_target_ids.size()); + const int n_groups = static_cast(match_groups.size()); + const int n_sites_batch = static_cast(genotypes.size()); + const int site_offset = hmm_sites_processed; + +#ifdef _OPENMP + #pragma omp parallel for schedule(dynamic, 1) +#endif + for (int idx = 0; idx < n_active; ++idx) { + const double t = branch_lengths_vec[idx]; + const double log_tid = log_target_ids_vec[idx]; + const int target_id = active_target_ids[idx]; + int local_group_idx = match_group_idx; + + for (int s = 0; s < n_sites_batch; s++) { + const int site = site_offset + s; + + if (local_group_idx < (n_groups - 1) && + (genetic_positions[site] >= match_groups[local_group_idx + 1].cm_position)) { + local_group_idx++; + hmm_ptrs[idx]->set_samples( + match_groups[local_group_idx].match_candidates.at(target_id)); + } + + const double kt = k_per_site[site] * t; + const double lt = l_per_site[site] * t; + const double rho_c = kt; + const double log1p_rho = -std::log1p(-std::exp(-kt)); + const double rho = sparse ? log1p_rho : log1p_rho + log_tid; + const double mu_c = lt; + const double mu = -std::log1p(-std::exp(-lt)); + hmm_ptrs[idx]->process_site(genotypes[s].data(), rho, rho_c, mu, mu_c); + + if ((site + 1) % prune_interval == 0) { + hmm_ptrs[idx]->prune(); + } + } + } + + // Update shared state to reflect all sites processed + for (int s = 0; s < n_sites_batch; s++) { + const int site = site_offset + s; + if (match_group_idx < (n_groups - 1) && + (genetic_positions[site] >= match_groups[match_group_idx + 1].cm_position)) { + match_group_idx++; + } + } + hmm_sites_processed += n_sites_batch; +} + +void ThreadsLowMem::process_all_sites_viterbi_flat(const int32_t* data, int n_sites, int n_haps) { + static_assert(sizeof(int) == sizeof(int32_t), "int and int32_t must be the same size"); + const int prune_interval = 500; + const int n_active = static_cast(active_target_ids.size()); + const int n_groups = static_cast(match_groups.size()); + const int site_offset = hmm_sites_processed; + +#ifdef _OPENMP + #pragma omp parallel for schedule(dynamic, 1) +#endif + for (int idx = 0; idx < n_active; ++idx) { + const double t = branch_lengths_vec[idx]; + const double log_tid = log_target_ids_vec[idx]; + const int target_id = active_target_ids[idx]; + int local_group_idx = match_group_idx; + + for (int s = 0; s < n_sites; s++) { + const int site = site_offset + s; + const int* row = reinterpret_cast(data + static_cast(s) * n_haps); + + if (local_group_idx < (n_groups - 1) && + (genetic_positions[site] >= match_groups[local_group_idx + 1].cm_position)) { + local_group_idx++; + hmm_ptrs[idx]->set_samples( + match_groups[local_group_idx].match_candidates.at(target_id)); + } + + const double kt = k_per_site[site] * t; + const double lt = l_per_site[site] * t; + const double rho_c = kt; + const double log1p_rho = -std::log1p(-std::exp(-kt)); + const double rho = sparse ? log1p_rho : log1p_rho + log_tid; + const double mu_c = lt; + const double mu = -std::log1p(-std::exp(-lt)); + hmm_ptrs[idx]->process_site(row, rho, rho_c, mu, mu_c); + + if ((site + 1) % prune_interval == 0) { + hmm_ptrs[idx]->prune(); + } + } + } + + // Update shared state + for (int s = 0; s < n_sites; s++) { + const int site = site_offset + s; + if (match_group_idx < (n_groups - 1) && + (genetic_positions[site] >= match_groups[match_group_idx + 1].cm_position)) { + match_group_idx++; + } + } + hmm_sites_processed += n_sites; } void ThreadsLowMem::traceback() { + // Target 0 gets an empty path for (int target_id : target_ids) { if (target_id == 0) { paths.emplace(target_id, ViterbiPath(0)); + break; } - else { + } + // Active targets: traceback in parallel, then insert into map + const int n_active = static_cast(active_target_ids.size()); + std::vector path_vec(n_active, ViterbiPath(0)); + +#ifdef _OPENMP + #pragma omp parallel for schedule(dynamic, 1) +#endif + for (int i = 0; i < n_active; i++) { + path_vec[i] = hmm_ptrs[i]->traceback(); + } + + for (int i = 0; i < n_active; i++) { + paths.emplace(active_target_ids[i], std::move(path_vec[i])); + } + + // Build path_ptrs for hets/dating hot path + path_ptrs.clear(); + path_ptrs.reserve(active_target_ids.size()); + for (int target_id : active_target_ids) { + path_ptrs.push_back(&paths.at(target_id)); + } + // Free HMM memory + hmm_vec.clear(); + hmm_ptrs.clear(); +} - paths.emplace(target_id, hmms.at(target_id).traceback()); +// Internal: process one het site from raw pointer +void ThreadsLowMem::process_site_hets_raw(const int* genotype) { + // Handle target 0 + for (int target_id : target_ids) { + if (target_id == 0) { + if (genotype[0] == 1) { + paths.at(0).het_sites.push_back(het_sites_processed); + } + break; + } + } + // Active targets + const int n_active = static_cast(active_target_ids.size()); + for (int idx = 0; idx < n_active; ++idx) { + const int target_id = active_target_ids[idx]; + ViterbiPath* path = path_ptrs[idx]; + int current_seg_idx = segment_indices_vec[idx]; + while (current_seg_idx < (static_cast(path->segment_starts.size()) - 1) && + (het_sites_processed >= path->segment_starts[current_seg_idx + 1])) { + current_seg_idx++; + } + segment_indices_vec[idx] = current_seg_idx; + const int sample = path->sample_ids[current_seg_idx]; + if (genotype[sample] != genotype[target_id] && + (genotype[sample] == 1 || genotype[target_id] == 1)) { + path->het_sites.push_back(het_sites_processed); } } - hmms.clear(); + het_sites_processed++; } void ThreadsLowMem::process_site_hets(const std::vector& genotype) { + process_site_hets_raw(genotype.data()); +} + +// Batch hets: parallelized across targets +void ThreadsLowMem::process_all_sites_hets(const std::vector>& genotypes) { + const int n_active = static_cast(active_target_ids.size()); + const int n_sites_batch = static_cast(genotypes.size()); + const int site_offset = het_sites_processed; + + // Handle target 0 sequentially (rare path, only for first target) for (int target_id : target_ids) { if (target_id == 0) { - if (genotype.at(0) == 1) { - paths.at(0).het_sites.push_back(het_sites_processed); + for (int s = 0; s < n_sites_batch; s++) { + if (genotypes[s][0] == 1) { + paths.at(0).het_sites.push_back(site_offset + s); + } } + break; } - else { - ViterbiPath& path = paths.at(target_id); - int current_seg_idx = segment_indices.at(target_id); - while (current_seg_idx < (static_cast(path.segment_starts.size()) - 1) && - (het_sites_processed >= path.segment_starts.at(current_seg_idx + 1))) { + } + +#ifdef _OPENMP + #pragma omp parallel for schedule(dynamic, 1) +#endif + for (int idx = 0; idx < n_active; ++idx) { + const int target_id = active_target_ids[idx]; + ViterbiPath* path = path_ptrs[idx]; + int current_seg_idx = segment_indices_vec[idx]; + + for (int s = 0; s < n_sites_batch; s++) { + const int site = site_offset + s; + const int* genotype = genotypes[s].data(); + + while (current_seg_idx < (static_cast(path->segment_starts.size()) - 1) && + (site >= path->segment_starts[current_seg_idx + 1])) { current_seg_idx++; } - segment_indices.at(target_id) = current_seg_idx; - int sample = path.sample_ids.at(current_seg_idx); - // For now, we do not count unphased variants as a part of this, - // so we verify at least one of the het-pair is a "1", - // (i.e., "-7") is treated as "0". - // More work is needed to verify inclusion of unphased variants helps at all - if (genotype.at(sample) != genotype.at(target_id) && - (genotype.at(sample) == 1 || genotype.at(target_id) == 1)) { - path.het_sites.push_back(het_sites_processed); + const int sample = path->sample_ids[current_seg_idx]; + if (genotype[sample] != genotype[target_id] && + (genotype[sample] == 1 || genotype[target_id] == 1)) { + path->het_sites.push_back(site); } } + segment_indices_vec[idx] = current_seg_idx; } - het_sites_processed++; + + het_sites_processed += n_sites_batch; +} + +void ThreadsLowMem::process_all_sites_hets_flat(const int32_t* data, int n_sites, int n_haps) { + static_assert(sizeof(int) == sizeof(int32_t), "int and int32_t must be the same size"); + const int n_active = static_cast(active_target_ids.size()); + const int site_offset = het_sites_processed; + + // Handle target 0 sequentially + for (int target_id : target_ids) { + if (target_id == 0) { + for (int s = 0; s < n_sites; s++) { + const int* row = reinterpret_cast(data + static_cast(s) * n_haps); + if (row[0] == 1) { + paths.at(0).het_sites.push_back(site_offset + s); + } + } + break; + } + } + +#ifdef _OPENMP + #pragma omp parallel for schedule(dynamic, 1) +#endif + for (int idx = 0; idx < n_active; ++idx) { + const int target_id = active_target_ids[idx]; + ViterbiPath* path = path_ptrs[idx]; + int current_seg_idx = segment_indices_vec[idx]; + + for (int s = 0; s < n_sites; s++) { + const int site = site_offset + s; + const int* row = reinterpret_cast(data + static_cast(s) * n_haps); + + while (current_seg_idx < (static_cast(path->segment_starts.size()) - 1) && + (site >= path->segment_starts[current_seg_idx + 1])) { + current_seg_idx++; + } + const int sample = path->sample_ids[current_seg_idx]; + if (row[sample] != row[target_id] && + (row[sample] == 1 || row[target_id] == 1)) { + path->het_sites.push_back(site); + } + } + segment_indices_vec[idx] = current_seg_idx; + } + + het_sites_processed += n_sites; } void ThreadsLowMem::date_segments() { @@ -200,37 +452,43 @@ void ThreadsLowMem::date_segments() { "Can't date segments, not all sites have been parsed for heterozygosity."); } - for (int target_id : target_ids) { - if (target_id == 0) { - continue; - } - if (segment_indices.at(target_id) != paths.at(target_id).size() - 1) { + const int n_active = static_cast(active_target_ids.size()); + for (int idx = 0; idx < n_active; idx++) { + const int target_id = active_target_ids[idx]; + if (segment_indices_vec[idx] != path_ptrs[idx]->size() - 1) { std::string prompt = "incomplete path at sample " + std::to_string(target_id) + ", processed "; - prompt += std::to_string(segment_indices.at(target_id) + 1) + " segments, expected "; - prompt += std::to_string(paths.at(target_id).size()); + prompt += std::to_string(segment_indices_vec[idx] + 1) + " segments, expected "; + prompt += std::to_string(path_ptrs[idx]->size()); throw std::runtime_error(prompt); } } - for (int target_id : target_ids) { - if (target_id == 0) { - continue; - } - ViterbiPath& path = paths.at(target_id); +#ifdef _OPENMP + #pragma omp parallel for schedule(dynamic, 1) +#endif + for (int idx = 0; idx < n_active; idx++) { + const int target_id = active_target_ids[idx]; + ViterbiPath& path = *path_ptrs[idx]; ViterbiPath new_path(target_id); std::size_t n_segs = path.segment_starts.size(); + // Track position in sorted het_sites to avoid repeated linear scans + auto het_it = path.het_sites.begin(); for (std::size_t k = 0; k < n_segs; k++) { - int sample_id = path.sample_ids.at(k); - int segment_start = path.segment_starts.at(k); - int segment_end = k < n_segs - 1 ? path.segment_starts.at(k + 1) : num_sites - 1; + int sample_id = path.sample_ids[k]; + int segment_start = path.segment_starts[k]; + int segment_end = k < n_segs - 1 ? path.segment_starts[k + 1] : num_sites - 1; if (segment_end == segment_start) { continue; } - // This is inefficient but probably not that bad + // Advance iterator to first het in [segment_start, ...) + while (het_it != path.het_sites.end() && *het_it < segment_start) { + ++het_it; + } std::vector segment_hets; - for (int h : path.het_sites) { + for (auto it = het_it; it != path.het_sites.end(); ++it) { + int h = *it; if (((segment_start <= h) && (h < segment_end)) || ((h == num_sites - 1) && (segment_end == num_sites - 1))) { segment_hets.push_back(h); @@ -250,13 +508,11 @@ void ThreadsLowMem::date_segments() { for (std::size_t j = 0; j < breakpoints.size(); j++) { int breakpoint_start = breakpoints[j]; int breakpoint_end = (j == breakpoints.size() - 1) ? segment_end : breakpoints[j + 1]; - // there may be off-by-one errors here on the last segment (but who cares?) double bp_size = - physical_positions.at(breakpoint_end) - physical_positions.at(breakpoint_start); + physical_positions[breakpoint_end] - physical_positions[breakpoint_start]; double cm_size = - genetic_positions.at(breakpoint_end) - genetic_positions.at(breakpoint_start); + genetic_positions[breakpoint_end] - genetic_positions[breakpoint_start]; - // Same as above std::vector breakpoint_hets; for (int h : segment_hets) { if (((breakpoint_start <= h) && (h < breakpoint_end)) || @@ -273,36 +529,33 @@ void ThreadsLowMem::date_segments() { } } else { - // there are off-by-one errors here on the last segment (but who cares?) - double bp_size = physical_positions.at(segment_end) - physical_positions.at(segment_start); - double cm_size = genetic_positions.at(segment_end) - genetic_positions.at(segment_start); + double bp_size = physical_positions[segment_end] - physical_positions[segment_start]; + double cm_size = genetic_positions[segment_end] - genetic_positions[segment_start]; double height = ThreadsFastLS::date_segment( static_cast(segment_hets.size()), cm_size, bp_size, mutation_rate, demography); new_path.append(segment_start, sample_id, height, segment_hets); } } - paths.at(target_id) = new_path; + *path_ptrs[idx] = new_path; } return; } int ThreadsLowMem::count_branches() const { int n_branches = 0; - for (int target_id : target_ids) { - if (target_id == 0) { - continue; - } - n_branches += hmms.at(target_id).count_branches(); + for (std::size_t i = 0; i < hmm_ptrs.size(); i++) { + n_branches += hmm_ptrs[i]->count_branches(); } return n_branches; } void ThreadsLowMem::prune() { - for (int target_id : target_ids) { - if (target_id == 0) { - continue; - } - hmms.at(target_id).prune(); + const int n = static_cast(hmm_ptrs.size()); +#ifdef _OPENMP + #pragma omp parallel for schedule(dynamic, 1) +#endif + for (int i = 0; i < n; i++) { + hmm_ptrs[i]->prune(); } } @@ -328,4 +581,4 @@ ThreadsLowMem::serialize_paths() { return std::tuple>, std::vector>, std::vector>, std::vector>>( all_starts, all_ids, all_heights, all_hetsites); -} \ No newline at end of file +} diff --git a/src/ThreadsLowMem.hpp b/src/ThreadsLowMem.hpp index 77ad06f..008e63c 100644 --- a/src/ThreadsLowMem.hpp +++ b/src/ThreadsLowMem.hpp @@ -21,7 +21,8 @@ #include "Matcher.hpp" #include "ThreadsFastLS.hpp" #include "ViterbiLowMem.hpp" -#include +#include +#include #include #include #include @@ -41,6 +42,8 @@ class ThreadsLowMem { const std::vector& cm_positions); // 2b. process all sites for the hmms void process_site_viterbi(const std::vector& genotype); + void process_all_sites_viterbi(const std::vector>& genotypes); + void process_all_sites_viterbi_flat(const int32_t* data, int n_sites, int n_haps); // 2c. prune branches at regular intervals (i.e. when there's a lot of them, figure this out soon) void prune(); // 2d. traceback all the hmms to get viterbi paths @@ -48,6 +51,8 @@ class ThreadsLowMem { // 3a. add het sites void process_site_hets(const std::vector& genotype); + void process_all_sites_hets(const std::vector>& genotypes); + void process_all_sites_hets_flat(const int32_t* data, int n_sites, int n_haps); // 3b. date all segments void date_segments(); @@ -63,9 +68,9 @@ class ThreadsLowMem { public: // This object will only run the HMM for these ids std::vector target_ids; + // Keep legacy map interface for pybind compatibility std::unordered_map expected_branch_lengths; double mean_bp_size = 0.0; - std::unordered_map segment_indices; std::unordered_map paths; int num_samples = 0; int num_sites = 0; @@ -79,11 +84,23 @@ class ThreadsLowMem { bool sparse = false; private: + // Hot-path data: flat vectors parallel to target_ids (excluding id 0) + std::vector active_target_ids; // target_ids without 0 + std::vector branch_lengths_vec; // parallel to active_target_ids + std::vector log_target_ids_vec; // precomputed log(target_id) + std::vector segment_indices_vec; // parallel to active_target_ids + std::vector hmm_ptrs; // parallel to active_target_ids + std::vector path_ptrs; // parallel to active_target_ids + Demography demography; + // Precomputed per-site HMM parameters + std::vector k_per_site; // 2 * 0.01 * cm_sizes[i] + std::vector l_per_site; // 2 * mu * bp_sizes[i] + // 2. HMM quantites int hmm_sites_processed = 0; - std::unordered_map hmms; + std::vector> hmm_vec; // owned, never moves int match_group_idx = 0; std::vector match_groups; @@ -92,6 +109,10 @@ class ThreadsLowMem { int het_sites_processed = 0; int n_hmm_samples = 100; int hmm_min_sites = 10; + + // Internal: process one site from raw pointer (no copy) + void process_site_viterbi_raw(const int* genotype); + void process_site_hets_raw(const int* genotype); }; #endif // THREADS_ARG_THREADS_LOW_MEM_HPP diff --git a/src/ViterbiLowMem.cpp b/src/ViterbiLowMem.cpp index 306cde1..57b8704 100644 --- a/src/ViterbiLowMem.cpp +++ b/src/ViterbiLowMem.cpp @@ -28,7 +28,8 @@ namespace { const int ALLELE_UNPHASED_HET = -7; inline std::size_t coord_id_key(int i, int j) { - return (static_cast(i) << 32) | static_cast(j); + return (static_cast(static_cast(i)) << 32) | + static_cast(static_cast(j)); } } // namespace @@ -128,95 +129,120 @@ ViterbiState::ViterbiState(int _target_id, std::vector _sample_ids) throw std::runtime_error("found no samples for ViterbiState object for sample " + std::to_string(target_id)); } + current_traceback_ptrs.reserve(sample_ids.size()); for (int sample_id : sample_ids) { - std::size_t key = coord_id_key(0, sample_id); - traceback_states.emplace(key, TracebackNode(sample_id, 0, nullptr, 0.)); - current_tracebacks[sample_id] = &traceback_states.at(key); + current_traceback_ptrs.push_back(alloc_node(sample_id, 0, nullptr, 0.)); } best_score = 0; best_match = sample_ids.at(0); + best_match_idx = 0; } -void ViterbiState::process_site(const std::vector& genotype, double rho, double rho_c, +void ViterbiState::process_site(const int* genotype, double rho, double rho_c, double mu, double mu_c) { - int current_site = sites_processed; + const int current_site = sites_processed; double best_new_score = best_score + std::max(rho, rho_c) + std::max(mu, mu_c); int best_new_match = best_match; - double new_score; - int observed_allele = genotype.at(target_id); - TracebackNode* prev_best = current_tracebacks.at(best_match); - for (int sample_id : sample_ids) { - int allele = genotype.at(sample_id); + int best_new_match_idx = best_match_idx; + const int observed_allele = genotype[target_id]; + const double recomb_threshold = best_score + rho; + const double unphased_penalty = (mu_c + mu) * 0.5; + const bool observed_is_unphased = (observed_allele == ALLELE_UNPHASED_HET); + + TracebackNode* prev_best = current_traceback_ptrs[best_match_idx]; + + const int n_samples = static_cast(sample_ids.size()); + for (int idx = 0; idx < n_samples; ++idx) { + const int sample_id = sample_ids[idx]; + const int allele = genotype[sample_id]; double copy_penalty; - if ((allele == ALLELE_UNPHASED_HET) || (observed_allele == ALLELE_UNPHASED_HET)) { - copy_penalty = (mu_c + mu) / 2.; + if (observed_is_unphased || (allele == ALLELE_UNPHASED_HET)) { + copy_penalty = unphased_penalty; } else { copy_penalty = (allele == observed_allele) ? mu_c : mu; } - if (!current_tracebacks.count(sample_id)) { - // If we've just added new sites (this will happen vary rarely), - // recombine from previous best state - new_score = best_score + copy_penalty + rho; - std::size_t key = coord_id_key(current_site, sample_id); - traceback_states.emplace(key, TracebackNode(sample_id, current_site, prev_best, new_score)); - current_tracebacks[sample_id] = &traceback_states.at(key); + + double new_score; + TracebackNode* state = current_traceback_ptrs[idx]; + if (state == nullptr) { + // Newly added sample (happens rarely, after set_samples) + new_score = recomb_threshold + copy_penalty; + current_traceback_ptrs[idx] = alloc_node(sample_id, current_site, prev_best, new_score); } else { - // Otherwise, check whether we should recombine or extend - TracebackNode* state = current_tracebacks.at(sample_id); - if (state->score + rho_c <= best_score + rho) { - // If extending is cheaper, simply update the score of the current traceback + if (state->score + rho_c <= recomb_threshold) { + // Extend: cheaper than recombining new_score = state->score + copy_penalty + rho_c; state->score = new_score; } else { - // If we recombine, add a new branch - new_score = best_score + copy_penalty + rho; - std::size_t key = coord_id_key(current_site, sample_id); - traceback_states.emplace(key, TracebackNode(sample_id, current_site, prev_best, new_score)); - current_tracebacks.at(sample_id) = &traceback_states.at(key); + // Recombine: add a new branch + new_score = recomb_threshold + copy_penalty; + current_traceback_ptrs[idx] = alloc_node(sample_id, current_site, prev_best, new_score); } } if (new_score < best_new_score) { best_new_score = new_score; best_new_match = sample_id; + best_new_match_idx = idx; } } best_score = best_new_score; best_match = best_new_match; + best_match_idx = best_new_match_idx; sites_processed++; } void ViterbiState::set_samples(std::unordered_set new_sample_ids) { + // Build old sample_id → ptr map from current parallel vectors + std::unordered_map old_ptrs; + old_ptrs.reserve(sample_ids.size()); + for (std::size_t i = 0; i < sample_ids.size(); ++i) { + old_ptrs[sample_ids[i]] = current_traceback_ptrs[i]; + } + std::vector new_samples_vec(new_sample_ids.begin(), new_sample_ids.end()); if (!new_sample_ids.count(best_match)) { new_samples_vec.push_back(best_match); } - for (int sample_id : sample_ids) { - // clean up branches we definitely won't use - if (!new_sample_ids.count(sample_id) && sample_id != best_match) { - current_tracebacks.erase(sample_id); + sample_ids = new_samples_vec; + + // Rebuild parallel pointer vector to match new sample_ids ordering + current_traceback_ptrs.clear(); + current_traceback_ptrs.reserve(sample_ids.size()); + for (std::size_t i = 0; i < sample_ids.size(); ++i) { + int sample_id = sample_ids[i]; + auto it = old_ptrs.find(sample_id); + current_traceback_ptrs.push_back(it != old_ptrs.end() ? it->second : nullptr); + if (sample_id == best_match) { + best_match_idx = static_cast(i); } } - sample_ids = new_samples_vec; } void ViterbiState::prune() { - std::unordered_map tmp_traceback_states; + std::deque new_nodes; + std::unordered_map key_to_ptr; - for (int sample_id : sample_ids) { - TracebackNode* state = current_tracebacks.at(sample_id); - TracebackNode* new_state = recursive_insert(tmp_traceback_states, state); - current_tracebacks[sample_id] = new_state; - } + // Recursively copy only reachable nodes into the new deque + auto copy_node = [&](auto& self, TracebackNode* state) -> TracebackNode* { + if (state == nullptr) return nullptr; + std::size_t key = state->key(); + auto it = key_to_ptr.find(key); + if (it != key_to_ptr.end()) return it->second; + TracebackNode* new_parent = self(self, state->previous); + new_nodes.emplace_back(state->sample_id, state->site, new_parent, state->score); + TracebackNode* ptr = &new_nodes.back(); + key_to_ptr[key] = ptr; + return ptr; + }; - traceback_states.clear(); - for (int sample_id : sample_ids) { - TracebackNode* state = current_tracebacks.at(sample_id); - TracebackNode* new_state = recursive_insert(traceback_states, state); - current_tracebacks[sample_id] = new_state; + for (std::size_t idx = 0; idx < sample_ids.size(); ++idx) { + current_traceback_ptrs[idx] = copy_node(copy_node, current_traceback_ptrs[idx]); } + + traceback_nodes = std::move(new_nodes); } // add everything above and return a key to the new address @@ -235,14 +261,19 @@ ViterbiState::recursive_insert(std::unordered_map& s return &state_map.at(key); } +TracebackNode* ViterbiState::alloc_node(int sample_id, int site, TracebackNode* previous, double score) { + traceback_nodes.emplace_back(sample_id, site, previous, score); + return &traceback_nodes.back(); +} + int ViterbiState::count_branches() const { - return static_cast(traceback_states.size()); + return static_cast(traceback_nodes.size()); } ViterbiPath ViterbiState::traceback() { ViterbiPath path(target_id); path.score = best_score; - TracebackNode* state = current_tracebacks.at(best_match); + TracebackNode* state = current_traceback_ptrs[best_match_idx]; while (state != nullptr) { int match_id = state->sample_id; int seg_start = state->site; diff --git a/src/ViterbiLowMem.hpp b/src/ViterbiLowMem.hpp index 4a33dfd..dc7a6a5 100644 --- a/src/ViterbiLowMem.hpp +++ b/src/ViterbiLowMem.hpp @@ -17,6 +17,7 @@ #ifndef THREADS_ARG_VITERBI_LOW_MEM_HPP #define THREADS_ARG_VITERBI_LOW_MEM_HPP +#include #include #include #include @@ -61,27 +62,36 @@ class ViterbiState { public: ViterbiState(int _target_id, std::vector _sample_ids); - void process_site(const std::vector& genotype, double rho, double rho_c, double _mu, + void process_site(const int* genotype, double rho, double rho_c, double _mu, double _mu_c); + void process_site(const std::vector& genotype, double rho, double rho_c, double _mu, + double _mu_c) { + process_site(genotype.data(), rho, rho_c, _mu, _mu_c); + } void set_samples(std::unordered_set new_sample_ids); int count_branches() const; void prune(); ViterbiPath traceback(); private: - std::unordered_map traceback_states; + // Arena for TracebackNode storage — deque guarantees pointer stability + std::deque traceback_nodes; + TracebackNode* alloc_node(int sample_id, int site, TracebackNode* previous, double score); + // Used only during prune to deduplicate copied nodes TracebackNode* recursive_insert(std::unordered_map& state_map, TracebackNode* state); public: int target_id = 0; int best_match = -1; + int best_match_idx = 0; double best_score = 0.0; int sites_processed = 0; double mutation_penalty = 0.0; std::vector sample_ids; std::vector sample_scores; - std::unordered_map current_tracebacks; + // Parallel to sample_ids: traceback pointer for each sample + std::vector current_traceback_ptrs; }; #endif // THREADS_ARG_VITERBI_LOW_MEM_HPP diff --git a/src/threads_arg/__main__.py b/src/threads_arg/__main__.py index 9c712ae..fd18363 100644 --- a/src/threads_arg/__main__.py +++ b/src/threads_arg/__main__.py @@ -14,7 +14,7 @@ # You should have received a copy of the GNU General Public License # along with this program. If not, see . -import click +import argparse import logging logging.basicConfig( @@ -35,102 +35,117 @@ def goodbye(): ,-'-' `-=_,-'-' `-=_,-'-' `-=_,-'-' `-=_""") -@click.group() def main(): - pass - -@main.command() -@click.option("--pgen", required=True, help="Path to input genotypes in pgen format") -@click.option("--map", help="Path to genotype map in SHAPEIT format") -@click.option("--recombination_rate", default=1.3e-8, type=float, help="Genome-wide recombination rate. Ignored if a map is passed") -@click.option("--demography", required=True, help="Path to input genotype") -@click.option("--mode", required=True, type=click.Choice(['array', 'wgs']), default="wgs", help="Inference mode (wgs or array)") -@click.option("--fit_to_data", is_flag=True, default=False, help="If specified, Threads performs a post-processing step to ensure the inferred ARG contains an edge matching each input mutation.") -@click.option("--normalize", is_flag=True, default=False, help="If specified, Threads will normalize output to fit with the input demography.") -@click.option("--allele_ages", default=None, help="Allele ages used for post-processing with the --fit_to_data flag, otherwise ignored. If not specified, allele ages are inferred automatically.") -@click.option("--query_interval", type=float, default=0.01, help="Hyperparameter for the preliminary haplotype matching in cM") -@click.option("--match_group_interval", type=float, default=0.5, help="Hyperparameter for the preliminary haplotype matching in cM") -@click.option("--mutation_rate", required=True, type=float, default=1.4e-8, help="Genome-wide mutation rate") -@click.option("--num_threads", type=int, default=1, help="Number of computational threads to request") -@click.option("--region", help="Region of genome in chr:start-end format for which ARG is output. The full genotype is still used for inference") -@click.option("--max_sample_batch_size", help="Max number of LS processes run simultaneously per thread", default=None, type=int) -@click.option("--save_metadata", is_flag=True, default=False, help="If specified, the output will include sample/variant metadata (sample IDs, marker names, allele symbols, etc).") -@click.option("--out") -def infer(**kwargs): - """Infer an ARG from genotype data""" - from .infer import threads_infer - threads_infer(**kwargs) - goodbye() - -@main.command() -@click.option("--threads", required=True, help="Path to an input .threads file") -@click.option("--argn", default=None, help="Path to an output .argn file") -@click.option("--tsz", default=None, help="Path to an output .tsz file") -@click.option("--add_mutations", is_flag=True, default=False, help="If passed, mutations are parsimoniously added to the output ARG. This may result in a high number of mutations if the --fit_to_data flag was not used.") -def convert(**kwargs): - """Convert Threads ARGs to ARG-Needle or tskit format""" - from .convert import threads_convert - threads_convert(**kwargs) - goodbye() - -@main.command() -@click.option("--threads", required=True, help="Path to an input .threads file.") -@click.option("--out", required=True, help="Path to output.") -@click.option("--num_threads", type=int, help="Size of processor pool to process batches", default=None) -def allele_ages(**kwargs): - """Infer allele ages from a Threads ARG""" - from .allele_ages import estimate_allele_ages - estimate_allele_ages(**kwargs) - goodbye() - -@main.command() -@click.option("--argn", help="Path to input .argn file") -@click.option("--out", help="Path to output .mut file") -@click.option("--maf", type=float, default=0.02, help="Do not store entries with MAF above this") -@click.option("--input", type=str, help="Path to bcf/vcf with genotypes to map with AC/AN fields") -@click.option("--region", type=str, help="Region in chr:start-end format (start and end inclusive)") -@click.option("--num_threads", type=int, help="Number of computational threads to request", default=1) -def map(**kwargs): - """Map genotypes to an ARG in ARG-Needle format""" - from .map_mutations_to_arg import threads_map_mutations_to_arg - threads_map_mutations_to_arg(**kwargs) - goodbye() - -@main.command() -@click.option("--panel", required=True, help="pgen array panel") -@click.option("--target", required=True, help="pgen array targets") -@click.option("--mut", required=True, help="pgen array targets") -@click.option("--map", required=True, help="Path to genotype map in SHAPEIT format") -@click.option("--mutation_rate", type=float, help="Per-site-per-generation SNP mutation rate", default=1.4e-8) -@click.option("--demography", required=True, help="Path to file containing demographic history") -@click.option("--out", help="Path to output .vcf file", default=None) -@click.option("--stdout", help="Redirect output to stdout (will disable logging)", is_flag=True) -@click.option("--region", required=True, type=str, help="Region in chr:start-end format (start and end inclusive)") -def impute(panel, target, map, mut, demography, out, stdout, region, mutation_rate=1.4e-8): - """Impute missing genotypes using a reference panel""" - # --stdout flag is mutually exclusive to --out flag. It is used only here to - # confirm the user wants to redirect (potentially a lot of data) to stdout. - # The Impute class does not use this variable, instead 'out' is just None. - import sys - if (stdout and out) or not (stdout or out): - print("Either --out or --stdout must be specified", file=sys.stderr) - exit(1) - - from .impute import Impute - Impute(panel, target, map, mut, demography, out, region, mutation_rate) - - # Do not print anything in stdout mode, to keep output clean. - if not stdout: + parser = argparse.ArgumentParser(prog="threads") + subparsers = parser.add_subparsers(dest="command") + + # infer + p_infer = subparsers.add_parser("infer", help="Infer an ARG from genotype data") + p_infer.add_argument("--pgen", required=True, help="Path to input genotypes in pgen format") + p_infer.add_argument("--map", help="Path to genotype map in SHAPEIT format") + p_infer.add_argument("--recombination_rate", default=1.3e-8, type=float, help="Genome-wide recombination rate. Ignored if a map is passed") + p_infer.add_argument("--demography", required=True, help="Path to input genotype") + p_infer.add_argument("--mode", required=True, choices=["array", "wgs"], default="wgs", help="Inference mode (wgs or array)") + p_infer.add_argument("--fit_to_data", action="store_true", default=False, help="If specified, Threads performs a post-processing step to ensure the inferred ARG contains an edge matching each input mutation.") + p_infer.add_argument("--normalize", action="store_true", default=False, help="If specified, Threads will normalize output to fit with the input demography.") + p_infer.add_argument("--allele_ages", default=None, help="Allele ages used for post-processing with the --fit_to_data flag, otherwise ignored. If not specified, allele ages are inferred automatically.") + p_infer.add_argument("--query_interval", type=float, default=0.01, help="Hyperparameter for the preliminary haplotype matching in cM") + p_infer.add_argument("--match_group_interval", type=float, default=0.5, help="Hyperparameter for the preliminary haplotype matching in cM") + p_infer.add_argument("--mutation_rate", required=True, type=float, default=1.4e-8, help="Genome-wide mutation rate") + p_infer.add_argument("--num_threads", type=int, default=1, help="Number of computational threads to request") + p_infer.add_argument("--region", help="Region of genome in chr:start-end format for which ARG is output. The full genotype is still used for inference") + p_infer.add_argument("--max_sample_batch_size", help="Max number of LS processes run simultaneously per thread", default=None, type=int) + p_infer.add_argument("--save_metadata", action="store_true", default=False, help="If specified, the output will include sample/variant metadata (sample IDs, marker names, allele symbols, etc).") + p_infer.add_argument("--out") + + # convert + p_convert = subparsers.add_parser("convert", help="Convert Threads ARGs to ARG-Needle or tskit format") + p_convert.add_argument("--threads", required=True, help="Path to an input .threads file") + p_convert.add_argument("--argn", default=None, help="Path to an output .argn file") + p_convert.add_argument("--tsz", default=None, help="Path to an output .tsz file") + p_convert.add_argument("--add_mutations", action="store_true", default=False, help="If passed, mutations are parsimoniously added to the output ARG. This may result in a high number of mutations if the --fit_to_data flag was not used.") + + # allele_ages + p_allele_ages = subparsers.add_parser("allele_ages", help="Infer allele ages from a Threads ARG") + p_allele_ages.add_argument("--threads", required=True, help="Path to an input .threads file.") + p_allele_ages.add_argument("--out", required=True, help="Path to output.") + p_allele_ages.add_argument("--num_threads", type=int, help="Size of processor pool to process batches", default=None) + + # map + p_map = subparsers.add_parser("map", help="Map genotypes to an ARG in ARG-Needle format") + p_map.add_argument("--argn", help="Path to input .argn file") + p_map.add_argument("--out", help="Path to output .mut file") + p_map.add_argument("--maf", type=float, default=0.02, help="Do not store entries with MAF above this") + p_map.add_argument("--input", type=str, help="Path to bcf/vcf with genotypes to map with AC/AN fields") + p_map.add_argument("--region", type=str, help="Region in chr:start-end format (start and end inclusive)") + p_map.add_argument("--num_threads", type=int, help="Number of computational threads to request", default=1) + + # impute + p_impute = subparsers.add_parser("impute", help="Impute missing genotypes using a reference panel") + p_impute.add_argument("--panel", required=True, help="pgen array panel") + p_impute.add_argument("--target", required=True, help="pgen array targets") + p_impute.add_argument("--mut", required=True, help="pgen array targets") + p_impute.add_argument("--map", required=True, help="Path to genotype map in SHAPEIT format") + p_impute.add_argument("--mutation_rate", type=float, help="Per-site-per-generation SNP mutation rate", default=1.4e-8) + p_impute.add_argument("--demography", required=True, help="Path to file containing demographic history") + p_impute.add_argument("--out", help="Path to output .vcf file", default=None) + p_impute.add_argument("--stdout", help="Redirect output to stdout (will disable logging)", action="store_true") + p_impute.add_argument("--region", required=True, type=str, help="Region in chr:start-end format (start and end inclusive)") + + # vcf + p_vcf = subparsers.add_parser("vcf", help="Print genotypes from Threads ARGs to stdout in VCF format") + p_vcf.add_argument("--threads", required=True, help="Path to input .threads file") + p_vcf.add_argument("--variants", default=None, help="Path to .pvar or .bim file with variant information") + p_vcf.add_argument("--samples", default=None, help="Path to a file with one sample ID per line") + + args = parser.parse_args() + + if args.command is None: + parser.print_help() + parser.exit(2) + + if args.command == "infer": + from .infer import threads_infer + kwargs = vars(args) + del kwargs["command"] + threads_infer(**kwargs) goodbye() -@main.command() -@click.option("--threads", required=True, help="Path to input .threads file") -@click.option("--variants", default=None, help="Path to .pvar or .bim file with variant information") -@click.option("--samples", default=None, help="Path to a file with one sample ID per line") -def vcf(threads, variants, samples): - """Print genotypes from Threads ARGs to stdout in VCF format""" - from .threads_to_vcf import threads_to_vcf - threads_to_vcf(threads, samples=samples, variants=variants) + elif args.command == "convert": + from .convert import threads_convert + kwargs = vars(args) + del kwargs["command"] + threads_convert(**kwargs) + goodbye() + + elif args.command == "allele_ages": + from .allele_ages import estimate_allele_ages + kwargs = vars(args) + del kwargs["command"] + estimate_allele_ages(**kwargs) + goodbye() + + elif args.command == "map": + from .map_mutations_to_arg import threads_map_mutations_to_arg + kwargs = vars(args) + del kwargs["command"] + threads_map_mutations_to_arg(**kwargs) + goodbye() + + elif args.command == "impute": + import sys + if (args.stdout and args.out) or not (args.stdout or args.out): + print("Either --out or --stdout must be specified", file=sys.stderr) + exit(1) + from .impute import Impute + Impute(args.panel, args.target, args.map, args.mut, args.demography, args.out, args.region, args.mutation_rate) + if not args.stdout: + goodbye() + + elif args.command == "vcf": + from .threads_to_vcf import threads_to_vcf + threads_to_vcf(args.threads, samples=args.samples, variants=args.variants) + if __name__ == "__main__": main() diff --git a/src/threads_arg/allele_ages.py b/src/threads_arg/allele_ages.py index f667cf6..300c248 100644 --- a/src/threads_arg/allele_ages.py +++ b/src/threads_arg/allele_ages.py @@ -16,31 +16,17 @@ import logging import multiprocessing -import numpy as np from tqdm import tqdm -from threads_arg import AgeEstimator, GenotypeIterator +from threads_arg import estimate_ages as _estimate_ages_cpp from .serialization import load_instructions from .utils import timer_block, default_process_count, split_list logger = logging.getLogger(__name__) -def _nth_batch_worker(instructions, result_idx, allele_ages_results): - # Estimate ages on this instruction batch - gt_it = GenotypeIterator(instructions) - age_estimator = AgeEstimator(instructions) - while gt_it.has_next_genotype(): - g = np.array(gt_it.next_genotype()) - age_estimator.process_site(g) - allele_age_estimates = age_estimator.get_inferred_ages() - - # Index result so full data can be reconstructed in order - allele_ages_results[result_idx] = allele_age_estimates - - -def _nth_batch_worker_star(args): - return _nth_batch_worker(*args) +def _batch_worker(instructions): + return _estimate_ages_cpp(instructions) def estimate_ages(instructions, num_batches, num_threads): # Make sure we don't use more CPUs than requested @@ -64,27 +50,15 @@ def estimate_ages(instructions, num_batches, num_threads): batched_instructions.append(range_instructions) with timer_block(f"Estimating allele ages ({num_processors} CPUs)"): - # Process-safe dict so batch results can be reconstructed in order - manager = multiprocessing.Manager() - allele_ages_results = manager.dict() - - # Create arguments for each job in process pool - jobs_args = [(range_inst, i, allele_ages_results) - for i, range_inst in enumerate(batched_instructions)] - with multiprocessing.Pool(processes=num_processors) as pool: - # To use tqdm with a pool, use imap with shim method to unpack args. - # Note the enclosing unassigned list() call is necessary. Otherwise - # no value is retrieved and the process is ignored/dropped. - list(tqdm( - pool.imap(_nth_batch_worker_star, jobs_args), - total=len(jobs_args) + batch_results = list(tqdm( + pool.imap(_batch_worker, batched_instructions), + total=len(batched_instructions) )) - # Collect batched estimates into single list in index sort order allele_age_estimates = [] - for i in range(len(allele_ages_results)): - allele_age_estimates += allele_ages_results[i] + for batch in batch_results: + allele_age_estimates += batch return allele_age_estimates def estimate_allele_ages(threads, out, num_threads): diff --git a/src/threads_arg/convert.py b/src/threads_arg/convert.py index f005634..18ff35a 100644 --- a/src/threads_arg/convert.py +++ b/src/threads_arg/convert.py @@ -16,7 +16,6 @@ import sys import time -import tszip import logging import threads_arg import arg_needle_lib @@ -90,19 +89,18 @@ def threads_convert(threads, argn, tsz, add_mutations=False): instructions = load_instructions(threads) try: logger.info("Attempting to convert to arg format...") - arg = threads_to_arg(instructions, add_mutations=add_mutations, noise=0.0) - except: - # arg_needle_lib does not allow polytomies - logger.info(f"Conflicting branches (this is expected), retrying with noise=1e-5...") - try: - arg = threads_to_arg(instructions, add_mutations=add_mutations, noise=1e-5) - except:# tskit.LibraryError: - logger.info(f"Conflicting branches, retrying with noise=1e-3...") - arg = threads_to_arg(instructions, add_mutations=add_mutations, noise=1e-3) + arg = threads_to_arg(instructions, add_mutations=add_mutations, noise=1e-5) + except RuntimeError: + logger.info(f"Conflicting branches, retrying with noise=1e-3...") + arg = threads_to_arg(instructions, add_mutations=add_mutations, noise=1e-3) if argn is not None: logger.info(f"Writing to {argn}") arg_needle_lib.serialize_arg(arg, argn) if tsz is not None: + try: + import tszip + except ImportError: + raise ImportError("tszip is required for .trees.tsz output. Install it with: pip install 'threads-arg[convert]'") logger.info(f"Converting to tree sequence and writing to {tsz}") tszip.compress(arg_needle_lib.arg_to_tskit(arg), tsz) logger.info(f"Done, in {time.time() - start_time} seconds") diff --git a/src/threads_arg/fwbw.py b/src/threads_arg/fwbw.py index 6bc0cb8..8ec62ed 100644 --- a/src/threads_arg/fwbw.py +++ b/src/threads_arg/fwbw.py @@ -1,119 +1,14 @@ # Code adapted from an implementation of the Li-Stephens algorithm # available at: https://github.com/astheeggeggs/lshmm -import numba import numpy as np import logging -import os - -logger = logging.getLogger(__name__) - -_DISABLE_NUMBA = os.environ.get("LSHMM_DISABLE_NUMBA", "0") - -try: - ENABLE_NUMBA = {"0": True, "1": False}[_DISABLE_NUMBA] -except KeyError as e: - raise KeyError( - "Environment variable 'LSHMM_DISABLE_NUMBA' must be '0' or '1'" - ) from e - -if not ENABLE_NUMBA: - logger.warning( - "Numba globally disabled, performance will be drastically reduced." - ) +from threads_arg import forwards_ls_hap, backwards_ls_hap -DEFAULT_NUMBA_ARGS = { - "nopython": True, - "cache": True, -} - - -def numba_njit(func, **kwargs): - if ENABLE_NUMBA: - return numba.jit(func, **{**DEFAULT_NUMBA_ARGS, **kwargs}) - else: - return func +logger = logging.getLogger(__name__) MISSING = -9 -@numba_njit -def forwards_ls_hap(n, m, H, s, e, r, norm=True): - """Matrix based haploid LS forward algorithm using numpy vectorisation.""" - # Initialise - F = np.zeros((m, n)) - r_n = r / n - - if norm: - c = np.zeros(m) - for i in range(n): - F[0, i] = ( - 1 / n * e[0, np.int64(np.equal(H[0, i], s[0, 0]) or s[0, 0] == MISSING)] - ) - c[0] += F[0, i] - - for i in range(n): - F[0, i] *= 1 / c[0] - - # Forwards pass - for l in range(1, m): - for i in range(n): - F[l, i] = F[l - 1, i] * (1 - r[l]) + r_n[l] - F[l, i] *= e[ - l, np.int64(np.equal(H[l, i], s[0, l]) or s[0, l] == MISSING) - ] - c[l] += F[l, i] - - for i in range(n): - F[l, i] *= 1 / c[l] - # Log-likelihood: ll = np.sum(np.log10(c)) - else: - c = np.ones(m) - for i in range(n): - F[0, i] = ( - 1 / n * e[0, np.int64(np.equal(H[0, i], s[0, 0]) or s[0, 0] == MISSING)] - ) - - # Forwards pass - for l in range(1, m): - for i in range(n): - F[l, i] = F[l - 1, i] * (1 - r[l]) + np.sum(F[l - 1, :]) * r_n[l] - F[l, i] *= e[ - l, np.int64(np.equal(H[l, i], s[0, l]) or s[0, l] == MISSING) - ] - # Log-likelihood: ll = np.log10(np.sum(F[m - 1, :])) - return F, c - -@numba_njit -def backwards_ls_hap(n, m, H, s, e, c, r): - """Matrix based haploid LS backward algorithm using numpy vectorisation.""" - # Initialise - B = np.zeros((m, n)) - for i in range(n): - B[m - 1, i] = 1 - r_n = r / n - - # Backwards pass - for l in range(m - 2, -1, -1): - tmp_B = np.zeros(n) - tmp_B_sum = 0 - for i in range(n): - tmp_B[i] = ( - e[ - l + 1, - np.int64( - np.equal(H[l + 1, i], s[0, l + 1]) or s[0, l + 1] == MISSING - ), - ] - * B[l + 1, i] - ) - tmp_B_sum += tmp_B[i] - for i in range(n): - B[l, i] = r_n[l + 1] * tmp_B_sum - B[l, i] += (1 - r[l + 1]) * tmp_B[i] - B[l, i] *= 1 / c[l + 1] - - return B - def checks(reference_panel, query, mutation_rate, recombination_rates): ref_shape = reference_panel.shape @@ -139,12 +34,14 @@ def checks(reference_panel, query, mutation_rate, recombination_rates): def set_emission_probabilities(reference_panel, query, mutation_rate): m, n = reference_panel.shape - n_alleles = np.int8( - [ - len(np.unique(np.append(reference_panel[j, :], query[:, j]))) - for j in range(reference_panel.shape[0]) - ] - ) + + # Vectorized check: a site is biallelic if it has both 0s and 1s across + # the combined reference+query panel. For binary data this is equivalent + # to checking that the row min != row max (after including the query). + combined_min = np.minimum(reference_panel.min(axis=1), query.min(axis=0)) + combined_max = np.maximum(reference_panel.max(axis=1), query.max(axis=0)) + is_polymorphic = combined_min != combined_max + n_alleles = np.where(is_polymorphic, np.int8(2), np.int8(1)) if not np.all((n_alleles == 2) | (n_alleles == 1)): raise ValueError("Only fixed or bi-allelic sites allowed") @@ -158,13 +55,12 @@ def set_emission_probabilities(reference_panel, query, mutation_rate): # Evaluate emission probabilities here, using the mutation rate e = np.zeros((m, 2)) - for j in range(m): - if n_alleles[j] == 1: # In case we're at an invariant site - e[j, 0] = 0 - e[j, 1] = 1 - else: - e[j, 0] = mutation_rate[j] - e[j, 1] = 1 - mutation_rate[j] + e[:, 0] = mutation_rate + e[:, 1] = 1 - mutation_rate + # Invariant sites: emission for mismatch is 0, match is 1 + invariant = n_alleles == 1 + e[invariant, 0] = 0 + e[invariant, 1] = 1 return e @@ -179,13 +75,20 @@ def fwbw(reference_panel, # Get emissions emissions = set_emission_probabilities(reference_panel, query, mutation_rate) - # Run forwards + # Run forwards (C++) forward_array, fwd_norm_factor = forwards_ls_hap( - n, m, reference_panel, query, emissions, recombination_rates, norm=True) + reference_panel.astype(np.float64), + query.ravel().astype(np.float64), + emissions, + recombination_rates) - # Run backwards + # Run backwards (C++) backward_array = backwards_ls_hap( - n, m, reference_panel, query, emissions, fwd_norm_factor, recombination_rates) + reference_panel.astype(np.float64), + query.ravel().astype(np.float64), + emissions, + fwd_norm_factor, + recombination_rates) # Return posterior return forward_array * backward_array diff --git a/src/threads_arg/impute.py b/src/threads_arg/impute.py index 111d1ba..247d635 100644 --- a/src/threads_arg/impute.py +++ b/src/threads_arg/impute.py @@ -2,13 +2,19 @@ import multiprocessing import os import numpy as np -import pandas as pd +try: + import pandas as pd +except ImportError: + raise ImportError("pandas is required for imputation. Install it with: pip install 'threads-arg[impute]'") import sys from tqdm import tqdm from cyvcf2 import VCF from threads_arg import ThreadsFastLS, ImputationMatcher -from scipy.sparse import csr_array, lil_matrix +try: + from scipy.sparse import csr_array, vstack as sparse_vstack +except ImportError: + raise ImportError("scipy is required for imputation. Install it with: pip install 'threads-arg[impute]'") from datetime import datetime from typing import Dict, Tuple, List, Union from dataclasses import dataclass @@ -85,7 +91,9 @@ def write_site(self, genotypes, record, imputed, contig): alt = alt[0] qual = "." filter = "PASS" - gt_strings = [f"{np.round(hap_1):.0f}|{np.round(hap_2):.0f}:{dosage:.3f}".rstrip("0").rstrip(".") for hap_1, hap_2, dosage in zip(haps1, haps2, dosages)] + gt1 = np.rint(haps1).astype(int) + gt2 = np.rint(haps2).astype(int) + gt_strings = [f"{g1}|{g2}:{dosage:.3f}".rstrip("0").rstrip(".") for g1, g2, dosage in zip(gt1, gt2, dosages)] f = self.file f.write(("\t".join([contig, pos, snp_id, ref, alt, qual, filter, f"{imp_str}AF={af:.4f}", "GT:DS", "\t".join(gt_strings)]) + "\n")) @@ -324,13 +332,11 @@ def __getitem__(self, snp_idx: int): if snp_idx in self.posteriors_by_snp_idx: return self.posteriors_by_snp_idx[snp_idx] - # Rebuild snp data - col_len = len(self.posteriors) - row_len = self.posteriors[0].shape[1] - target_posteriors = np.empty(shape=(col_len, row_len), dtype=np.float64) - for i, p in enumerate(self.posteriors): - posteriors = p[[snp_idx],:].toarray() - target_posteriors[i] = posteriors / np.sum(posteriors) + # Batch extract row snp_idx from all sparse posteriors and convert once + rows = sparse_vstack([p[[snp_idx]] for p in self.posteriors]) + target_posteriors = rows.toarray() + row_sums = target_posteriors.sum(axis=1, keepdims=True) + target_posteriors /= row_sums # Cache and clear out-of-date entries self.posteriors_by_snp_idx[snp_idx] = target_posteriors @@ -475,8 +481,8 @@ def _sparse_posteriors(self, demography, mutation_rate): ref_matches = _reference_matching(self.panel_snps, self.target_snps, self.cm_pos_array) mutation_rate = 0.0001 - cm_sizes = list(self.cm_pos_array[1:] - self.cm_pos_array[:-1]) - cm_sizes = np.array(cm_sizes + [cm_sizes[-1]]) + cm_sizes = np.diff(self.cm_pos_array) + cm_sizes = np.append(cm_sizes, cm_sizes[-1]) Ne = 20_000 recombination_rates = 1 - np.exp(-4 * Ne * 0.01 * cm_sizes / self.num_samples_panel) @@ -492,9 +498,9 @@ def _sparse_posteriors(self, demography, mutation_rate): for i, h_target in tqdm(enumerate(target_transpose), total=len(target_transpose), mininterval=1): with tt_impute: # Imputation thread with divergence matching - imputation_thread = bwt.impute(list(h_target), L) + imputation_thread = bwt.impute(h_target.tolist(), L) imputation_threads.append(imputation_thread) - matched_samples_viterbi = set([match_id for seg in imputation_thread for match_id in seg.ids]) + matched_samples_viterbi = {match_id for seg in imputation_thread for match_id in seg.ids} # All locally sampled matches matched_samples_matcher = (ref_matches[self.num_samples_panel + i]) @@ -519,14 +525,19 @@ def _sparsify_posterior(self, posterior, matched_samples): Expand to a compressed n_snps x n_samples matrix """ assert posterior.shape == (self.num_snps, len(matched_samples)) - matrix = lil_matrix((self.num_snps, self.num_samples_panel), dtype=np.float64) - posterior[posterior <= 1 / self.num_samples_panel] = 0 - for i, p in enumerate(posterior): - assert np.sum(p) > 0 - q = p / np.sum(p) - for j in np.nonzero(q)[0]: - matrix[i, matched_samples[j]] = q[j] - return csr_array(matrix) + threshold = 1 / self.num_samples_panel + # Find entries above threshold directly (avoids mutating full matrix) + rows, cols_local = np.nonzero(posterior > threshold) + vals = posterior[rows, cols_local] + # Renormalize per-row using only the kept entries + row_sums = np.bincount(rows, weights=vals, minlength=self.num_snps) + assert np.all(row_sums > 0) + vals = vals / row_sums[rows] + cols_global = matched_samples[cols_local] + return csr_array( + (vals, (rows, cols_global)), + shape=(self.num_snps, self.num_samples_panel) + ) def _init_step_snp(self): @@ -579,7 +590,8 @@ def only_active(posteriors): if self.mutation_container.is_mapped(record.id): mutation_mapping = self.mutation_container.get_mapping(record.id) - carriers = (1 - record.genotypes).nonzero()[0] if flipped else record.genotypes.nonzero()[0] + carriers_arr = (1 - record.genotypes).nonzero()[0] if flipped else record.genotypes.nonzero()[0] + carriers = set(carriers_arr.tolist()) active_positions = np.where(record.genotypes)[0] active_indexes = {pos: i for i, pos in enumerate(active_positions)} @@ -608,7 +620,7 @@ def compute_delta(active_site_posterior, i): record ) - genotypes = np.array([np.sum(asp) for asp in active_site_posteriors]) + genotypes = active_site_posteriors.sum(axis=1) if mutation_mapping: with self.tt_mutation_mapping: deltas = np.array([compute_delta(asp, i) for i, asp in enumerate(active_site_posteriors)]) @@ -636,7 +648,7 @@ def _process_and_write(self): # imputed imputed = False var_idx = self.snp_id_indexes[record.id] - genotypes = np.array(self.target_snps[var_idx], dtype=float) + genotypes = self.target_snps[var_idx].astype(float) else: # If this variant is not present on the genotyping array, then # run the Threads imputation routine diff --git a/src/threads_arg/infer.py b/src/threads_arg/infer.py index 88faeac..0591821 100644 --- a/src/threads_arg/infer.py +++ b/src/threads_arg/infer.py @@ -21,18 +21,14 @@ import pgenlib import importlib -os.environ["RAY_DEDUP_LOGS"] = "0" -import ray import numpy as np -import pandas as pd from threads_arg import ( ThreadsLowMem, Matcher, ViterbiPath, ThreadingInstructions, - ConsistencyWrapper, - GenotypeIterator + run_consistency, ) from .utils import ( make_recombination_from_map_and_pgen, @@ -40,6 +36,7 @@ split_list, parse_demography, iterate_pgen, + read_all_genotypes, read_positions_and_ids, parse_region_string, default_process_count, @@ -47,11 +44,8 @@ read_sample_names ) -from .serialization import serialize_instructions -from .allele_ages import estimate_ages -from .normalization import Normalizer logger = logging.getLogger(__name__) @@ -218,92 +212,100 @@ def threads_infer(pgen, map, recombination_rate, demography, mutation_rate, fit_ if max_sample_batch_size is None: max_sample_batch_size = 2 * num_samples + # Check for batch numpy API (optimized build) + HAS_NUMPY_API = hasattr(Matcher, 'process_all_sites_numpy') + + # Read all genotypes into memory once + logger.info("Reading genotypes") + all_genotypes = read_all_genotypes(pgen) + num_haps = all_genotypes.shape[1] + logger.info("Finding singletons") - # Get singleton filter for the matching step - alleles_out = None - phased_out = None - ac_mask = [] - iterate_pgen(pgen, lambda i, g: ac_mask.append(1 < g.sum() < 2 * num_samples)) - ac_mask = np.array(ac_mask, dtype=bool) + ac = all_genotypes.sum(axis=1) + ac_mask = (ac > 1) & (ac < num_haps) assert ac_mask.shape == genetic_positions.shape logger.info("Running PBWT matching") - # There are four params here to mess with: - # - query interval - # - match group size - # - neighborhood size - # - min_matches - # Keep min_matches low for small sample sizes, can be 2 up to ~1,000 but then > 3 - # Smaller numbers run faster and consume less memory MIN_MATCHES = 4 neighborhood_size = 4 - matcher = Matcher(2 * num_samples, genetic_positions[ac_mask], query_interval, match_group_interval, neighborhood_size, MIN_MATCHES) - def matcher_callback(i, g, mask, matcher): - if mask[i]: + matcher = Matcher(num_haps, genetic_positions[ac_mask], query_interval, match_group_interval, neighborhood_size, MIN_MATCHES) + if HAS_NUMPY_API: + matcher.process_all_sites_numpy(all_genotypes[ac_mask]) + else: + for g in all_genotypes[ac_mask]: matcher.process_site(g) - iterate_pgen(pgen, matcher_callback, mask=ac_mask, matcher=matcher) # Add top matches from adjacent sites to each match-chunk matcher.propagate_adjacent_matches() - # From here we parallelise if we can + ne_times, ne = parse_demography(demography) + sparse = mode == "array" + + # From here we parallelise: OpenMP (batch API) or Ray (per-site fallback) actual_num_threads = min(default_process_count(), num_threads) logger.info(f"Requested {num_threads} threads, found {actual_num_threads}.") paths = [] - if actual_num_threads > 1: - # Warning: this creates big copies, these matches are the main source of memory usage - sample_batches = split_list(list(range(2 * num_samples)), actual_num_threads) + + if HAS_NUMPY_API: + # Optimized path: single process, OpenMP parallelism across targets in C++ + sample_batch = list(range(num_haps)) + s_match_group = matcher.serializable_matches(sample_batch) + match_cm_positions = matcher.cm_positions() + matcher.clear() + del matcher + gc.collect() + + TLM = ThreadsLowMem(sample_batch, physical_positions, genetic_positions, ne, ne_times, mutation_rate, sparse) + TLM.initialize_viterbi(s_match_group, match_cm_positions) + del s_match_group + gc.collect() + + logger.info("Running Viterbi (batch + OpenMP)") + TLM.process_all_sites_viterbi_numpy(all_genotypes) + TLM.prune() + TLM.traceback() + + logger.info("Computing hets (batch + OpenMP)") + TLM.process_all_sites_hets_numpy(all_genotypes) + + logger.info("Dating segments") + TLM.date_segments() + + seg_starts, match_ids, heights, hetsites = TLM.serialize_paths() + for sample_id, ss, mi, ht, hs in zip(sample_batch, seg_starts, match_ids, heights, hetsites): + paths.append(ViterbiPath(sample_id, ss, mi, ht, hs)) + + elif actual_num_threads > 1: + # Released build multi-threaded: multiprocessing process parallelism + from multiprocessing import Pool + sample_batches = split_list(list(range(num_haps)), actual_num_threads) match_cm_positions = matcher.cm_positions() - del alleles_out - del phased_out + del all_genotypes gc.collect() - partial_viterbi_remote = ray.remote(partial_viterbi) - ray.init() - # Parallelised threading instructions - results = ray.get([partial_viterbi_remote.remote( - pgen, - mode, - 2 * num_samples, - physical_positions, - genetic_positions, - demography, - mutation_rate, - sample_batch, - matcher.serializable_matches(sample_batch), - match_cm_positions, - max_sample_batch_size, - actual_num_threads, - thread_id) for thread_id, sample_batch in enumerate(sample_batches)]) - ray.shutdown() - # Combine results from each thread + args_list = [ + (pgen, mode, num_haps, physical_positions, genetic_positions, + demography, mutation_rate, sample_batch, + matcher.serializable_matches(sample_batch), match_cm_positions, + max_sample_batch_size, actual_num_threads, thread_id) + for thread_id, sample_batch in enumerate(sample_batches)] + with Pool(actual_num_threads) as pool: + results = pool.starmap(partial_viterbi, args_list) for sample_batch, result_tuple in zip(sample_batches, results): for sample_id, seg_starts, match_ids, heights, hetsites in zip(sample_batch, *result_tuple): paths.append(ViterbiPath(sample_id, seg_starts, match_ids, heights, hetsites)) else: - sample_batch = list(range(2 * num_samples)) + # Released build single-threaded + sample_batch = list(range(num_haps)) s_match_group = matcher.serializable_matches(sample_batch) match_cm_positions = matcher.cm_positions() matcher.clear() del matcher gc.collect() - thread_id = 1 - # Single-threaded threading instructions results = partial_viterbi( - pgen, - mode, - 2 * num_samples, - physical_positions, - genetic_positions, - demography, - mutation_rate, - sample_batch, - s_match_group, - match_cm_positions, - max_sample_batch_size, - actual_num_threads, - thread_id) - + pgen, mode, num_haps, physical_positions, genetic_positions, + demography, mutation_rate, sample_batch, s_match_group, + match_cm_positions, max_sample_batch_size, actual_num_threads, 1) for sample_id, seg_starts, match_ids, heights, hetsites in zip(sample_batch, *results): paths.append(ViterbiPath(sample_id, seg_starts, match_ids, heights, hetsites)) @@ -314,6 +316,7 @@ def matcher_callback(i, g, mask, matcher): if normalize: logger.info("Applying normalization") + from .normalization import Normalizer normalizer = Normalizer(demography, 2 * num_samples) instructions = normalizer.normalize(instructions) @@ -323,30 +326,27 @@ def matcher_callback(i, g, mask, matcher): allele_age_estimates = None if allele_ages is None: logger.info("Inferring allele ages from data") + from .allele_ages import estimate_ages allele_age_estimates = estimate_ages(instructions, 3 * actual_num_threads, actual_num_threads) assert len(allele_age_estimates) == len(instructions.positions) else: logger.info(f"Reading allele ages from {allele_ages}") - allele_age_estimates = [] _, ids = read_positions_and_ids(pgen) - age_table = pd.read_table(allele_ages, header=None, names=["SNP", "POS", "AGE"]) - age_table = age_table[age_table["SNP"].astype(str).isin(ids)] - allele_age_estimates = age_table["AGE"].values - try: - assert age_table.shape[0] == len(instructions.positions) == len(allele_age_estimates) - except AssertionError: + id_set = set(str(x) for x in ids) + allele_age_estimates = [] + with open(allele_ages) as f: + for line in f: + fields = line.strip().split() + if len(fields) >= 3 and str(fields[0]) in id_set: + allele_age_estimates.append(float(fields[2])) + if len(allele_age_estimates) != len(instructions.positions): raise RuntimeError(f"Allele age estimates do not match markers in the region requested, expected {len(instructions.positions)} age estimates.") # Start the consistifying logger.info("Post-processing threading instructions to fit to data") - gt_it = GenotypeIterator(instructions) - cw = ConsistencyWrapper(instructions, allele_age_estimates) - while gt_it.has_next_genotype(): - g = np.array(gt_it.next_genotype()) - cw.process_site(g) - - consistent_instructions = cw.get_consistent_instructions() + consistent_instructions = run_consistency(instructions, allele_age_estimates) logger.info(f"Writing to {out}") + from .serialization import serialize_instructions serialize_instructions(consistent_instructions, out, variant_metadata=variant_metadata, @@ -354,6 +354,7 @@ def matcher_callback(i, g, mask, matcher): sample_names=sample_names) else: logger.info(f"Writing to {out}") + from .serialization import serialize_instructions serialize_instructions(instructions, out, variant_metadata=variant_metadata, diff --git a/src/threads_arg/map_mutations_to_arg.py b/src/threads_arg/map_mutations_to_arg.py index 801b466..b4aae2f 100644 --- a/src/threads_arg/map_mutations_to_arg.py +++ b/src/threads_arg/map_mutations_to_arg.py @@ -19,8 +19,6 @@ import time import logging -os.environ["RAY_DEDUP_LOGS"] = "0" -import ray import numpy as np import arg_needle_lib @@ -78,8 +76,6 @@ def _map_region(argn, input, region, maf_threshold): n_parsimoniously_mapped = 0 # Iterate over VCF records - read_time = 0 - map_time = 0 vcf = VCF(input) for record in vcf(region): ac = int(record.INFO.get("AC")) @@ -101,16 +97,12 @@ def _map_region(argn, input, region, maf_threshold): name = record.ID pos = record.POS - rt = time.time() hap = np.array(record.genotypes)[:, :2].flatten() - read_time += time.time() - rt assert len(hap) == len(arg.leaf_ids) if flipped: hap = 1 - hap - mt = time.time() mapping, _ = arg_needle_lib.map_genotype_to_ARG_approximate(arg, hap, float(pos - arg.offset)) - map_time += time.time() - mt if len(mapping) > 0: n_mapped += 1 @@ -165,6 +157,8 @@ def threads_map_mutations_to_arg(argn, out, maf, input, region, num_threads): if actual_num_threads == 1: return_strings, n_attempted, n_parsimoniously_mapped, n_relate_mapped = _map_region(argn, input, region, maf) else: + from multiprocessing import Pool + logger.info("Parsing VCF") vcf = VCF(input) positions = [record.POS for record in vcf(region)] @@ -174,12 +168,8 @@ def threads_map_mutations_to_arg(argn, out, maf, input, region, num_threads): # split into subregions split_positions = split_list(positions, actual_num_threads) subregions = [f"{contig}:{pos[0]}-{pos[-1]}" for pos in split_positions] - ray.init() - map_region_remote = ray.remote(_map_region) - results = ray.get([map_region_remote.remote( - argn, input, subregion, maf - ) for subregion in subregions]) - ray.shutdown() + with Pool(actual_num_threads) as pool: + results = pool.starmap(_map_region, [(argn, input, subregion, maf) for subregion in subregions]) return_strings = [] n_attempted, n_parsimoniously_mapped, n_relate_mapped = 0, 0, 0 for rets, natt, npars, nrel in results: diff --git a/src/threads_arg/normalization.py b/src/threads_arg/normalization.py index b34def8..3c716f9 100644 --- a/src/threads_arg/normalization.py +++ b/src/threads_arg/normalization.py @@ -15,12 +15,13 @@ # along with this program. If not, see . import logging -import pandas as pd import math import numpy as np -import pandas from threads_arg import ThreadingInstructions -import msprime +try: + import msprime +except ImportError: + raise ImportError("msprime is required for normalization. Install it with: pip install 'threads-arg[normalize]'") logging.basicConfig( format='%(asctime)s %(levelname)-8s %(message)s', @@ -42,12 +43,17 @@ def read_demography(self, demography_file): Read the input demography file into an msprime.Demography object. The input demography is assumed to be haploid. """ - df = pd.read_table(demography_file, header=None) - df.columns = ['GEN', 'NE'] + times, sizes = [], [] + with open(demography_file) as f: + for line in f: + fields = line.strip().split() + if len(fields) >= 2: + times.append(float(fields[0])) + sizes.append(float(fields[1])) demography = msprime.Demography() # NOTE: these initial sizes get overwritten anyways demography.add_population(name="A", initial_size=1e4) - for t, ne in zip(df["GEN"], df["NE"]): + for t, ne in zip(times, sizes): # Divide by 2 because msprime wants diploids but demography is haploid demography.add_population_parameters_change(t, initial_size=ne / 2, population="A") return demography diff --git a/src/threads_arg/phase.py b/src/threads_arg/phase.py deleted file mode 100644 index 41d6c35..0000000 --- a/src/threads_arg/phase.py +++ /dev/null @@ -1,103 +0,0 @@ -# This file is part of the Threads software suite. -# Copyright (C) 2024-2025 Threads Developers. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU General Public License -# along with this program. If not, see . - -import time -import logging -import arg_needle_lib -import numpy as np - -logger = logging.getLogger(__name__) - - -def phase_distance(arg, arg_pos, target, het_carriers, hom_carriers): - """ - Compute 'phasing distance' as described in the thesis. - """ - if len(het_carriers) + 2 * len(hom_carriers) == 1: - # negative so lower number (with longer pendant branch) wins - return -arg.node(target).parent_edge_at(arg_pos).parent.height - phase_distance = 0 - for carrier in het_carriers: - phase_distance += min(arg.mrca(target, 2 * carrier, arg_pos).height, arg.mrca(target, 2 * carrier + 1, arg_pos).height) - for hom in hom_carriers: - phase_distance += arg.mrca(target, 2 * hom, arg_pos).height + arg.mrca(target, 2 * hom + 1, arg_pos).height - return phase_distance - - -def threads_phase(scaffold, argn, ts, unphased, out): - """ - Use an imputed arg to phase. Other input follows same shape as in SHAPEIT5-rare. - """ - logger.info("Starting Threads-phase.") - logger.info("WARNING: Threads-phase is experimental functionality.") - start_time = time.time() - unphased_vcf = VCF(unphased) - scaffold_vcf = VCF(scaffold) - true_vcf = VCF(unphased) - phased_writer = Writer(out, true_vcf) - if argn is None and ts is None: - raise ValueError("Need either --argn or --ts") - logger.info("Reading ARG...") - try: - arg = arg_needle_lib.deserialize_arg(argn) - except: - import tskit - treeseq = tskit.load(ts) - arg = arg_needle_lib.tskit_to_arg(treeseq) - arg.populate_children_and_roots() - - i = 0 - logger.info("Phasing...") - scaffold_empty = False - v_scaffold = next(scaffold_vcf) - s_scaffold = 0 - - num_hets_found = 0 - # Main phasing routine - for v in unphased_vcf: - if not scaffold_empty and v.ID == v_scaffold.ID: - # If variant exists in scaffold, just copy it - v.genotypes = v_scaffold.genotypes - s_scaffold += 1 - try: - v_scaffold = next(scaffold_vcf) - except StopIteration: - scaffold_empty = True - else: - # Otherwise, do ARG-based phasing - G = np.array(v.genotypes) - g0, g1 = G[:, 0], G[:, 1] - het_carriers = ((g0 + g1) == 1).nonzero()[0] - hom_carriers = ((g0 + g1) == 2).nonzero()[0] - arg_pos = max(v.end - arg.offset, 0) - arg_pos = min(arg_pos, arg.end - 1) - - for target in het_carriers: - num_hets_found += 1 - phase_distance_0 = phase_distance(arg, arg_pos, 2 * target, het_carriers, hom_carriers) - phase_distance_1 = phase_distance(arg, arg_pos, 2 * target + 1, het_carriers, hom_carriers) - if phase_distance_0 <= phase_distance_1: - G[target] = [1, 0, True] - else: - G[target] = [0, 1, True] - v.genotypes = G - - v.genotypes = v.genotypes - phased_writer.write_record(v) - logger.info(f"Done, in {time.time() - start_time} seconds") - unphased_vcf.close() - scaffold_vcf.close() - phased_writer.close() diff --git a/src/threads_arg/serialization.py b/src/threads_arg/serialization.py index ffad464..5c28cf2 100644 --- a/src/threads_arg/serialization.py +++ b/src/threads_arg/serialization.py @@ -14,156 +14,47 @@ # You should have received a copy of the GNU General Public License # along with this program. If not, see . -import h5py import numpy as np -from datetime import datetime -from threads_arg import ThreadingInstructions +from threads_arg import ( + serialize_threads as _serialize_threads, + deserialize_threads as _deserialize_threads, + read_threads_metadata as _read_threads_metadata, + read_threads_sample_names as _read_threads_sample_names, +) def serialize_instructions(instructions, out, variant_metadata=None, allele_ages=None, sample_names=None): - num_threads = instructions.num_samples - num_sites = instructions.num_sites positions = instructions.positions - region_start = instructions.start - region_end = instructions.end - samples = list(range(num_threads)) - - all_starts = instructions.all_starts() - all_targets = instructions.all_targets() - all_tmrcas = instructions.all_tmrcas() - all_mismatches = instructions.all_mismatches() - - thread_starts = np.cumsum([0] + [len(starts) for starts in all_starts[:-1]]) - mut_starts = np.cumsum([0] + [len(mismatches) for mismatches in all_mismatches[:-1]]) - - flat_starts = [start for starts in all_starts for start in starts] - flat_tmrcas = [tmrca for tmrcas in all_tmrcas for tmrca in tmrcas] - flat_targets = [target for targets in all_targets for target in targets] - flat_mismatches = [mismatch for mismatches in all_mismatches for mismatch in mismatches] - - num_stitches = len(flat_starts) - num_mutations = len(flat_mismatches) - - f = h5py.File(out, "w") - f.attrs['datetime_created'] = datetime.now().isoformat() - - compression_opts = 9 - dset_samples = f.create_dataset("samples", (num_threads, 3), dtype=int, compression='gzip', - compression_opts=compression_opts) - dset_pos = f.create_dataset("positions", (num_sites), dtype=int, compression='gzip', - compression_opts=compression_opts) - # First L columns are random samples for imputation - dset_targets = f.create_dataset("thread_targets", (num_stitches, 2), dtype=int, compression='gzip', - compression_opts=compression_opts) - dset_ages = f.create_dataset("thread_ages", (num_stitches), dtype=np.double, compression='gzip', - compression_opts=compression_opts) - dset_het_s = f.create_dataset("het_sites", (num_mutations), dtype=int, compression='gzip', - compression_opts=compression_opts) - dset_range = f.create_dataset("arg_range", (2), dtype=np.double, compression='gzip', - compression_opts=compression_opts) - - dset_samples[:, 0] = samples - dset_samples[:, 1] = thread_starts - dset_samples[:, 2] = mut_starts - - dset_targets[:, 0] = flat_targets - dset_targets[:, 1] = flat_starts - - dset_pos[:] = positions - dset_ages[:] = flat_tmrcas - dset_het_s[:] = flat_mismatches - dset_range[:] = [region_start, region_end] - - if variant_metadata is not None and num_sites: - # If not none, it is a pandas dataframe with columns - # CHR, POS, ID, REF, ALT, QUAL, FILTER + metadata_cols = [] + if variant_metadata is not None and instructions.num_sites: min_pos = min(positions) max_pos = max(positions) - dset_variant_metadata = f.create_dataset("variant_metadata", (num_sites, 7), dtype=h5py.string_dtype(encoding='utf-8'), compression='gzip', - compression_opts=compression_opts) - variant_metadata = variant_metadata[(variant_metadata["POS"] >= min_pos) & (variant_metadata["POS"] <= max_pos)] + pos_int = np.array(variant_metadata["POS"], dtype=int) + mask = (pos_int >= min_pos) & (pos_int <= max_pos) + variant_metadata = variant_metadata[mask] assert variant_metadata.shape[0] == len(positions) assert np.all(np.array(variant_metadata["POS"], dtype=int) == np.array(positions)) + for col in ["CHROM", "POS", "ID", "REF", "ALT", "QUAL", "FILTER"]: + metadata_cols.append([str(x) for x in variant_metadata[col]]) - dset_variant_metadata[:, 0] = variant_metadata["CHROM"].astype(str) - dset_variant_metadata[:, 1] = variant_metadata["POS"].astype(str) - dset_variant_metadata[:, 2] = variant_metadata["ID"].astype(str) - dset_variant_metadata[:, 3] = variant_metadata["REF"].astype(str) - dset_variant_metadata[:, 4] = variant_metadata["ALT"].astype(str) - dset_variant_metadata[:, 5] = variant_metadata["QUAL"].astype(str) - dset_variant_metadata[:, 6] = variant_metadata["FILTER"].astype(str) + ages = list(allele_ages) if allele_ages is not None else [] + names = [str(x) for x in sample_names] if sample_names is not None else [] - if allele_ages is not None: - assert len(allele_ages) == len(positions) - dset_allele_ages = f.create_dataset("allele_ages", (num_sites, ), dtype=np.double, compression='gzip', - compression_opts=compression_opts) - dset_allele_ages[:] = allele_ages - - if sample_names is not None: - assert len(sample_names) == num_threads // 2 - num_diploids = len(sample_names) - dset_sample_names = f.create_dataset("sample_names", (num_diploids,), dtype=h5py.string_dtype(encoding='utf-8'), compression='gzip', - compression_opts=compression_opts) - dset_sample_names[:] = sample_names - f.close() + _serialize_threads(out, instructions, metadata_cols, ages, names) def load_instructions(threads): - """ - Create ThreadingInstructions object from a source .threads file - """ - f = h5py.File(threads, "r") - - _, thread_starts, het_starts = f["samples"][:, 0], f["samples"][:, 1], f["samples"][:, 2] - positions = f['positions'][...] - flat_targets, flat_starts = f['thread_targets'][:, 0], f['thread_targets'][:, -1] - flat_tmrcas = f['thread_ages'][...] - flat_mismatches = f['het_sites'][...] - - try: - arg_range = f['arg_range'][...] - except KeyError: - arg_range = [np.nan, np.nan] - - region_start = int(arg_range[0]) - region_end = int(arg_range[1]) - - starts = [] - targets = [] - tmrcas = [] - mismatches = [] - for i, (start, het_start) in enumerate(zip(thread_starts, het_starts)): - if i == len(thread_starts) - 1: - targets.append(flat_targets[start:].tolist()) - starts.append(flat_starts[start:].tolist()) - tmrcas.append(flat_tmrcas[start:].tolist()) - mismatches.append(flat_mismatches[het_start:].tolist()) - else: - targets.append(flat_targets[start:thread_starts[i + 1]].tolist()) - starts.append(flat_starts[start:thread_starts[i + 1]].tolist()) - tmrcas.append(flat_tmrcas[start:thread_starts[i + 1]].tolist()) - mismatches.append(flat_mismatches[het_start:het_starts[i + 1]].tolist()) - - positions = positions.astype(int).tolist() - return ThreadingInstructions( - starts, - tmrcas, - targets, - mismatches, - positions, - region_start, - region_end - ) + return _deserialize_threads(threads) def load_metadata(threads): - f = h5py.File(threads, "r") - import pandas as pd - return pd.DataFrame(f["variant_metadata"][:], columns=["CHROM", "POS", "ID", "REF", "ALT", "QUAL", "FILTER"]) + from .utils import VariantMetadata + cols = _read_threads_metadata(threads) + columns = ["CHROM", "POS", "ID", "REF", "ALT", "QUAL", "FILTER"] + return VariantMetadata({col: np.array(data) for col, data in zip(columns, cols)}) def load_sample_names(threads): - f = h5py.File(threads, "r") - return f["sample_names"][:] + return _read_threads_sample_names(threads) diff --git a/src/threads_arg/threads_to_vcf.py b/src/threads_arg/threads_to_vcf.py index 7b6f31b..16b9793 100644 --- a/src/threads_arg/threads_to_vcf.py +++ b/src/threads_arg/threads_to_vcf.py @@ -6,7 +6,7 @@ def threads_to_vcf(threads, samples=None, variants=None): if samples is None: try: sample_names = load_sample_names(threads) - except KeyError: + except (KeyError, RuntimeError): raise RuntimeError("Unable to load sample information from threading instructions. This may because the input was inferred using an older version of Threads or without the --save_metadata flag. Please provide files with variant information (in .bim/.pvar format) and sample IDs in a file with one sample per line.") else: with open(samples, "r") as samplefile: @@ -15,7 +15,7 @@ def threads_to_vcf(threads, samples=None, variants=None): if variants is None: try: variant_metadata = load_metadata(threads) - except KeyError: + except (KeyError, RuntimeError): raise RuntimeError("Unable to load sample information from threading instructions. This may because the input was inferred using an older version of Threads or without the --save_metadata flag. Please provide files with variant information (in .bim/.pvar format) and sample IDs in a file with one sample per line.") else: assert variants.endswith(".bim") or variants.endswith(".pvar") diff --git a/src/threads_arg/utils.py b/src/threads_arg/utils.py index 595ac65..fedad4c 100644 --- a/src/threads_arg/utils.py +++ b/src/threads_arg/utils.py @@ -16,7 +16,6 @@ import os import numpy as np -import pandas as pd import logging import pgenlib import re @@ -32,10 +31,20 @@ def read_map_file(map_file, expected_chromosome=None) -> Tuple[np.ndarray, np.nd """ Reading in map file for Li-Stephens using genetic maps in the SHAPEIT format """ - maps = pd.read_table(map_file, sep=r"\s+") - cm_pos = maps.cM.values.astype(np.float64) - phys_pos = maps.pos.values.astype(np.float64) - chromosomes = np.unique(maps.chr.values.astype(str)) + phys_list, cm_list, chr_list = [], [], [] + with open(map_file) as f: + header = f.readline().strip().split() + pos_idx = header.index('pos') + chr_idx = header.index('chr') + cm_idx = header.index('cM') + for line in f: + fields = line.strip().split() + phys_list.append(float(fields[pos_idx])) + chr_list.append(str(fields[chr_idx])) + cm_list.append(float(fields[cm_idx])) + cm_pos = np.array(cm_list, dtype=np.float64) + phys_pos = np.array(phys_list, dtype=np.float64) + chromosomes = np.unique(chr_list) # Currently we only allow for processing one chromosome at a time if len(chromosomes) > 1: @@ -60,9 +69,19 @@ def _read_pgen_physical_positions(pgen_file): bim = pgen_file.rstrip("pgen") + "bim" physical_positions = None if os.path.isfile(bim): - physical_positions = np.array(pd.read_table(bim, sep="\\s+", header=None, comment='#')[3]).astype(np.float64) + pos = [] + with open(bim) as f: + for line in f: + if line.startswith('#'): continue + pos.append(float(line.split()[3])) + physical_positions = np.array(pos, dtype=np.float64) elif os.path.isfile(pvar): - physical_positions = np.array(pd.read_table(pvar, sep="\\s+", header=None, comment='#')[1]).astype(np.float64) + pos = [] + with open(pvar) as f: + for line in f: + if line.startswith('#'): continue + pos.append(float(line.split()[1])) + physical_positions = np.array(pos, dtype=np.float64) else: raise RuntimeError(f"Can't find {bim} or {pvar} for {pgen_file}") @@ -112,6 +131,32 @@ def read_positions_and_ids(pgen): return positions, ids +class VariantMetadata: + """Lightweight replacement for pandas DataFrame for variant metadata.""" + __slots__ = ('_data', '_len') + + def __init__(self, data): + self._data = data # dict of numpy arrays + self._len = len(next(iter(data.values()))) + + def __getitem__(self, key): + if isinstance(key, str): + return self._data[key] + # Boolean mask + return VariantMetadata({k: v[key] for k, v in self._data.items()}) + + def __len__(self): + return self._len + + @property + def columns(self): + return list(self._data.keys()) + + @property + def shape(self): + return (self._len,) + + def read_variant_metadata(pgen): """ Attempt to read variant metadata in vcf style: @@ -120,29 +165,59 @@ def read_variant_metadata(pgen): pvar = pgen.replace("pgen", "pvar") bim = pgen.replace("pgen", "bim") if os.path.isfile(bim): - bim_df = pd.read_table(bim, names=["CHROM", "ID", "CM", "POS", "ALT", "REF"]) - out_df = bim_df[["CHROM", "POS", "ID", "REF", "ALT"]] - out_df["FILTER"] = bim_df["FILTER"] if "FILTER" in out_df.columns else "PASS" - out_df["QUAL"] = bim_df["QUAL"] if "QUAL" in out_df.columns else "." - return out_df + chrom, pos, vid, ref, alt = [], [], [], [], [] + with open(bim) as f: + for line in f: + if line.startswith('#'): + continue + fields = line.strip().split() + chrom.append(fields[0]) + vid.append(fields[1]) + pos.append(fields[3]) + alt.append(fields[4]) + ref.append(fields[5]) + return VariantMetadata({ + "CHROM": np.array(chrom), "POS": np.array(pos), + "ID": np.array(vid), "REF": np.array(ref), "ALT": np.array(alt), + "QUAL": np.full(len(pos), "."), "FILTER": np.full(len(pos), "PASS"), + }) elif os.path.isfile(pvar): + # Parse header to find column indices header = None - with open(pvar, "r") as pvarfile: - for line in pvarfile: + header_line_count = 0 + with open(pvar) as f: + for line in f: if line.startswith("##"): + header_line_count += 1 continue if line.startswith("#CHROM"): - header = line.strip().split() + header = line.strip().lstrip('#').split() + header_line_count += 1 break - if header is None: - raise RuntimeError(f"Invalid .pvar file {pvar}") - pvar_df = pd.read_table(pvar, comment="#", header=None, names=header, sep=r"\s+").rename({"#CHROM": "CHROM"}, axis=1) - - pd.options.mode.chained_assignment = None - out_df = pvar_df[["CHROM", "POS", "ID", "REF", "ALT"]] - out_df["FILTER"] = pvar_df["FILTER"].copy() if "FILTER" in out_df.columns else "PASS" - out_df["QUAL"] = pvar_df["QUAL"].copy() if "QUAL" in out_df.columns else "." - return out_df + if header is None: + raise RuntimeError(f"Invalid .pvar file {pvar}") + + col_idx = {name: i for i, name in enumerate(header)} + chrom, pos, vid, ref, alt, qual, filt = [], [], [], [], [], [], [] + has_filter = "FILTER" in col_idx + has_qual = "QUAL" in col_idx + with open(pvar) as f: + for _ in range(header_line_count): + next(f) + for line in f: + fields = line.strip().split() + chrom.append(fields[col_idx["CHROM"]]) + pos.append(fields[col_idx["POS"]]) + vid.append(fields[col_idx["ID"]]) + ref.append(fields[col_idx["REF"]]) + alt.append(fields[col_idx["ALT"]]) + filt.append(fields[col_idx["FILTER"]] if has_filter else "PASS") + qual.append(fields[col_idx["QUAL"]] if has_qual else ".") + return VariantMetadata({ + "CHROM": np.array(chrom), "POS": np.array(pos), + "ID": np.array(vid), "REF": np.array(ref), "ALT": np.array(alt), + "QUAL": np.array(qual), "FILTER": np.array(filt), + }) else: raise RuntimeError(f"Can't find {bim} or {pvar}") @@ -169,22 +244,39 @@ def read_sample_names(pgen): return [l.split()[1] for l in famfile] elif os.path.isfile(psam): - sam_df = pd.read_table(psam, sep=r"\s+") - if "IID" in sam_df.columns: - return sam_df["IID"].astype(str).tolist() - elif "#IID" in sam_df.columns: - return sam_df["#IID"].astype(str).tolist() - else: - # If no header, default to famfile - with open(psam, "r") as famfile: - return [l.split()[1] for l in famfile] + with open(psam, "r") as f: + header_line = f.readline().strip() + header = header_line.split() + # Find the IID column + if "IID" in header: + iid_idx = header.index("IID") + elif "#IID" in header: + iid_idx = header.index("#IID") + else: + # No recognized header, treat as fam-like (second column) + f.seek(0) + return [l.split()[1] for l in f] + names = [] + for line in f: + if line.startswith('#'): + continue + fields = line.strip().split() + if fields: + names.append(fields[iid_idx]) + return names else: raise RuntimeError(f"Can't find {fam} or {psam}") def parse_demography(demography): - d = pd.read_table(demography, sep="\\s+", header=None) - return list(d[0]), list(d[1]) + times, sizes = [], [] + with open(demography) as f: + for line in f: + fields = line.strip().split() + if len(fields) >= 2: + times.append(float(fields[0])) + sizes.append(float(fields[1])) + return times, sizes def split_list(list, n): @@ -310,6 +402,30 @@ def iterate_pgen(pgen, callback, start_idx=None, end_idx=None, **kwargs): assert i == M +def read_all_genotypes(pgen): + """Read all genotypes from a pgen file into a single (n_sites, n_haps) int32 array.""" + reader = pgenlib.PgenReader(pgen.encode()) + num_samples = reader.get_raw_sample_ct() + num_sites = reader.get_variant_ct() + n_haps = 2 * num_samples + + alleles_out = np.empty((num_sites, n_haps), dtype=np.int32) + phasepresent_out = np.empty((num_sites, num_samples), dtype=np.uint8) + + BATCH_SIZE = max(1, int(4e7 // n_haps)) + for b_start in range(0, num_sites, BATCH_SIZE): + b_end = min(num_sites, b_start + BATCH_SIZE) + reader.read_alleles_and_phasepresent_range( + b_start, b_end, + alleles_out[b_start:b_end], + phasepresent_out[b_start:b_end]) + + if np.any(phasepresent_out == 0): + raise RuntimeError("Unphased variants are currently not supported.") + + return alleles_out + + def default_process_count(): """ Get the number of CPUs available for multi-processing work diff --git a/src/threads_arg_pybind.cpp b/src/threads_arg_pybind.cpp index c829a20..e8f3aee 100644 --- a/src/threads_arg_pybind.cpp +++ b/src/threads_arg_pybind.cpp @@ -20,8 +20,12 @@ #include "AlleleAges.hpp" #include "GenotypeIterator.hpp" #include "VCFWriter.hpp" +#include "ForwardBackward.hpp" +#include "ThreadsIO.hpp" #include "pybind_utils.hpp" +#include +#include #include namespace py = pybind11; @@ -49,7 +53,23 @@ PYBIND11_MODULE(threads_arg_python_bindings, m) { .def_readonly("expected_branch_lengths", &ThreadsLowMem::expected_branch_lengths) .def("initialize_viterbi", &ThreadsLowMem::initialize_viterbi) .def("process_site_viterbi", &ThreadsLowMem::process_site_viterbi) + .def("process_all_sites_viterbi", &ThreadsLowMem::process_all_sites_viterbi) + .def("process_all_sites_viterbi_numpy", [](ThreadsLowMem& self, py::array_t arr) { + auto buf = arr.request(); + if (buf.ndim != 2) throw std::runtime_error("Expected 2D array (n_sites × n_haps)"); + int n_sites = static_cast(buf.shape[0]); + int n_haps = static_cast(buf.shape[1]); + self.process_all_sites_viterbi_flat(static_cast(buf.ptr), n_sites, n_haps); + }) .def("process_site_hets", &ThreadsLowMem::process_site_hets) + .def("process_all_sites_hets", &ThreadsLowMem::process_all_sites_hets) + .def("process_all_sites_hets_numpy", [](ThreadsLowMem& self, py::array_t arr) { + auto buf = arr.request(); + if (buf.ndim != 2) throw std::runtime_error("Expected 2D array (n_sites × n_haps)"); + int n_sites = static_cast(buf.shape[0]); + int n_haps = static_cast(buf.shape[1]); + self.process_all_sites_hets_flat(static_cast(buf.ptr), n_sites, n_haps); + }) .def("count_branches", &ThreadsLowMem::count_branches) .def("prune", &ThreadsLowMem::prune) .def("traceback", &ThreadsLowMem::traceback) @@ -89,6 +109,14 @@ PYBIND11_MODULE(threads_arg_python_bindings, m) { .def_readonly("num_samples", &Matcher::num_samples) .def_readonly("num_sites", &Matcher::num_sites) .def("process_site", &Matcher::process_site) + .def("process_all_sites", &Matcher::process_all_sites) + .def("process_all_sites_numpy", [](Matcher& self, py::array_t arr) { + auto buf = arr.request(); + if (buf.ndim != 2) throw std::runtime_error("Expected 2D array (n_sites × n_haps)"); + int n_sites = static_cast(buf.shape[0]); + int n_haps = static_cast(buf.shape[1]); + self.process_all_sites_flat(static_cast(buf.ptr), n_sites, n_haps); + }) .def("propagate_adjacent_matches", &Matcher::propagate_adjacent_matches) .def("get_matches", &Matcher::get_matches) .def("serializable_matches", &Matcher::serializable_matches) @@ -142,7 +170,39 @@ PYBIND11_MODULE(threads_arg_python_bindings, m) { &threading_instructions_get_state, &threading_instructions_set_state)) .def("left_multiply", &ThreadingInstructions::left_multiply, py::arg("x"), py::arg("diploid") = false, py::arg("normalize") = false) - .def("right_multiply", &ThreadingInstructions::right_multiply, py::arg("x"), py::arg("diploid") = false, py::arg("normalize") = false); + .def("right_multiply", &ThreadingInstructions::right_multiply, py::arg("x"), py::arg("diploid") = false, py::arg("normalize") = false) + .def("materialize_genotypes", &ThreadingInstructions::materialize_genotypes) + .def("materialize_normalized_haploid", &ThreadingInstructions::materialize_normalized_haploid) + .def("materialize_normalized_diploid", &ThreadingInstructions::materialize_normalized_diploid) + .def("prepare_tree_multiply", &ThreadingInstructions::prepare_tree_multiply) + .def("right_multiply_tree", &ThreadingInstructions::right_multiply_tree, py::arg("x")) + .def("left_multiply_tree", &ThreadingInstructions::left_multiply_tree, py::arg("x")) + .def("right_multiply_tree_batch", &ThreadingInstructions::right_multiply_tree_batch, py::arg("x_flat"), py::arg("k")) + .def("left_multiply_tree_batch", &ThreadingInstructions::left_multiply_tree_batch, py::arg("x_flat"), py::arg("k")) + .def("right_multiply_tree_batch_numpy", [](ThreadingInstructions& self, + py::array_t arr, int k) { + auto buf = arr.request(); + if (buf.ndim != 1 || static_cast(buf.shape[0]) != self.num_sites * k) + throw std::runtime_error("Input must be 1D array of length num_sites * k"); + std::vector x_flat(static_cast(buf.ptr), + static_cast(buf.ptr) + buf.shape[0]); + auto result = self.right_multiply_tree_batch(x_flat, k); + py::array_t out(static_cast(result.size())); + std::memcpy(out.mutable_data(), result.data(), result.size() * sizeof(double)); + return out; + }, py::arg("x_flat"), py::arg("k")) + .def("left_multiply_tree_batch_numpy", [](ThreadingInstructions& self, + py::array_t arr, int k) { + auto buf = arr.request(); + if (buf.ndim != 1 || static_cast(buf.shape[0]) != self.num_samples * k) + throw std::runtime_error("Input must be 1D array of length num_samples * k"); + std::vector x_flat(static_cast(buf.ptr), + static_cast(buf.ptr) + buf.shape[0]); + auto result = self.left_multiply_tree_batch(x_flat, k); + py::array_t out(static_cast(result.size())); + std::memcpy(out.mutable_data(), result.data(), result.size() * sizeof(double)); + return out; + }, py::arg("x_flat"), py::arg("k")); py::class_(m, "ConsistencyWrapper") .def(py::init>&, const std::vector>&, const std::vector>&, @@ -154,11 +214,15 @@ PYBIND11_MODULE(threads_arg_python_bindings, m) { .def("process_site", &ConsistencyWrapper::process_site) .def("get_consistent_instructions", &ConsistencyWrapper::get_consistent_instructions); + m.def("run_consistency", &run_consistency, py::arg("instructions"), py::arg("allele_ages")); + py::class_(m, "AgeEstimator") .def(py::init(), "initialize", py::arg("instructions")) .def("process_site", &AgeEstimator::process_site) .def("get_inferred_ages", &AgeEstimator::get_inferred_ages); + m.def("estimate_ages", &estimate_ages, py::arg("instructions")); + py::class_(m, "GenotypeIterator") .def(py::init(), "initialize", py::arg("instructions")) .def("next_genotype", &GenotypeIterator::next_genotype) @@ -175,4 +239,62 @@ PYBIND11_MODULE(threads_arg_python_bindings, m) { .def("set_filter", &VCFWriter::set_filter) .def("set_sample_names", &VCFWriter::set_sample_names) .def("write_vcf", &VCFWriter::write_vcf); + + // Forward-backward Li-Stephens algorithm (replaces numba JIT'd fwbw) + m.def("forwards_ls_hap", []( + py::array_t H, + py::array_t s, + py::array_t e, + py::array_t r) { + auto H_buf = H.request(); + auto s_buf = s.request(); + int m_sites = H_buf.shape[0]; + int n_refs = H_buf.shape[1]; + + auto [F_vec, c_vec] = forwards_ls_hap( + n_refs, m_sites, + static_cast(H_buf.ptr), + static_cast(s_buf.ptr), + static_cast(e.request().ptr), + static_cast(r.request().ptr)); + + py::array_t F({m_sites, n_refs}); + std::memcpy(F.mutable_data(), F_vec.data(), F_vec.size() * sizeof(double)); + py::array_t c(m_sites); + std::memcpy(c.mutable_data(), c_vec.data(), c_vec.size() * sizeof(double)); + return py::make_tuple(F, c); + }); + + m.def("backwards_ls_hap", []( + py::array_t H, + py::array_t s, + py::array_t e, + py::array_t c, + py::array_t r) { + auto H_buf = H.request(); + int m_sites = H_buf.shape[0]; + int n_refs = H_buf.shape[1]; + + auto B_vec = backwards_ls_hap( + n_refs, m_sites, + static_cast(H_buf.ptr), + static_cast(s.request().ptr), + static_cast(e.request().ptr), + static_cast(c.request().ptr), + static_cast(r.request().ptr)); + + py::array_t B({m_sites, n_refs}); + std::memcpy(B.mutable_data(), B_vec.data(), B_vec.size() * sizeof(double)); + return B; + }); + + // .threads file I/O (replaces h5py) + m.def("serialize_threads", &serialize_threads, + py::arg("filename"), py::arg("instructions"), + py::arg("metadata_cols") = std::vector>(), + py::arg("allele_ages") = std::vector(), + py::arg("sample_names") = std::vector()); + m.def("deserialize_threads", &deserialize_threads, py::arg("filename")); + m.def("read_threads_metadata", &read_threads_metadata, py::arg("filename")); + m.def("read_threads_sample_names", &read_threads_sample_names, py::arg("filename")); } diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 2a6040b..03e3377 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -24,8 +24,17 @@ FetchContent_Declare( ) FetchContent_MakeAvailable(Catch2) set(test_src + test_benchmark.cpp test_demography.cpp + test_demography_correctness.cpp + test_hmm.cpp + test_matcher.cpp + test_node.cpp + test_regression.cpp + test_threading_instructions.cpp + test_viterbi_lowmem.cpp test_viterbi_state.cpp + test_allele_ages.cpp ) # Main unit test exe diff --git a/test/bench_compression_multiply.py b/test/bench_compression_multiply.py new file mode 100644 index 0000000..ce749f5 --- /dev/null +++ b/test/bench_compression_multiply.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +""" +Quick compression & multiply benchmark: threads vs GRG vs RePair. +Uses pre-cached objects only — no rebuilding. + +Usage: + python test/bench_compression_multiply.py [--sizes 100,1000,10000] [--reps 5] + python test/bench_compression_multiply.py --dc --sizes 100 # also run data consistency +""" + +import os, sys, time, argparse, json, tempfile +import numpy as np + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src')) +_base = os.path.join(os.path.dirname(__file__), '..', '..') +sys.path.insert(0, os.path.join(_base, 'genrepair')) +sys.path.insert(0, os.path.join(_base, 'other', 'grgl')) + +from threads_arg.serialization import load_instructions, serialize_instructions + +TC = os.path.join(os.path.dirname(os.path.abspath(__file__)), "threads_cache") +DC = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "threads_cache") + + +def bench(fn, reps=5): + fn() + times = [] + for _ in range(reps): + t0 = time.perf_counter() + fn() + times.append((time.perf_counter() - t0) * 1000) + return np.median(times) + + +def sz(b): + if b is None: return "—" + if b < 1024: return f"{b}B" + if b < 1024**2: return f"{b/1024:.0f}K" + return f"{b/1024/1024:.1f}M" + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--sizes", default="100,1000,10000") + ap.add_argument("--reps", type=int, default=5) + ap.add_argument("--dc", action="store_true", help="Run data consistency (slow)") + args = ap.parse_args() + sizes = [int(s) for s in args.sizes.split(",")] + + try: + import pygrgl; have_grg = True + except ImportError: + have_grg = False; print("no pygrgl") + + try: + import genrepair; have_rpr = True + except ImportError: + have_rpr = False; print("no genrepair") + + results = [] + for n in sizes: + # Find files + tp = os.path.join(TC, f"sim_{n}dip.threads") + if not os.path.exists(tp): + tp = os.path.join(DC, f"sim_{n}.threads") + if not os.path.exists(tp): + print(f"{n}: skip"); continue + + inst = load_instructions(tp) + ns, nm = inst.num_samples, inst.num_sites + print(f"\n--- {n} dip ({ns} hap, {nm} sites) ---") + + # Sizes + dense = nm * 2 * n + thr_sz = os.path.getsize(tp) + grg_p = os.path.join(TC, f"grg_{n}dip.grg") + grg_sz = os.path.getsize(grg_p) if os.path.exists(grg_p) else None + rpr_p = os.path.join(DC, f"sim_{n}.repair") + rpr_sz = None + if have_rpr: + for ext in [".val", ".vc.C.ansf.1", ".vc.R.iv"]: + if not os.path.exists(rpr_p + ext): + rpr_sz = None; break + rpr_sz = (rpr_sz or 0) + os.path.getsize(rpr_p + ext) + + print(f" size: dense={sz(dense)} threads={sz(thr_sz)}({dense/thr_sz:.0f}x) " + f"grg={sz(grg_sz)}({dense/grg_sz:.0f}x)" if grg_sz else f" size: dense={sz(dense)} threads={sz(thr_sz)}({dense/thr_sz:.0f}x)", + end="") + if rpr_sz: + print(f" repair={sz(rpr_sz)}({dense/rpr_sz:.0f}x)", end="") + print() + + # Multiply + rng = np.random.default_rng(42) + xs = rng.normal(0, 1, nm).tolist() + xh = rng.normal(0, 1, ns).tolist() + + inst.prepare_tree_multiply() + tr = bench(lambda: inst.right_multiply_tree(xs), args.reps) + tl = bench(lambda: inst.left_multiply_tree(xh), args.reps) + + inst.materialize_genotypes() + dr = bench(lambda: inst.right_multiply(xs), args.reps) + dl = bench(lambda: inst.left_multiply(xh), args.reps) + + r = dict(n=n, ns=ns, nm=nm, dense_bytes=dense, thr_bytes=thr_sz, + grg_bytes=grg_sz, rpr_bytes=rpr_sz, + tree_R=tr, tree_L=tl, dense_R=dr, dense_L=dl) + + print(f" right: tree={tr:.2f}ms dense={dr:.2f}ms({dr/tr:.1f}x)", end="") + + if have_grg and os.path.exists(grg_p): + grg = pygrgl.load_immutable_grg(grg_p) + xg = rng.normal(0, 1, grg.num_mutations).astype(np.float32) + gr = bench(lambda: pygrgl.dot_product(grg, xg, pygrgl.DOWN), args.reps) + xgu = rng.normal(0, 1, grg.num_samples).astype(np.float32) + gl = bench(lambda: pygrgl.dot_product(grg, xgu, pygrgl.UP), args.reps) + r['grg_R'] = gr; r['grg_L'] = gl + print(f" grg={gr:.2f}ms({gr/tr:.1f}x)", end="") + + if have_rpr and rpr_sz: + cm = genrepair.CompressedMatrix(rpr_p, nm, 2*n, False, False) + xr = rng.normal(0, 1, cm.rows).astype(np.float32) + xl = rng.normal(0, 1, cm.cols).astype(np.float32) + rr = bench(lambda: cm.left_multiply(xr), args.reps) + rl = bench(lambda: cm.multiply(xl), args.reps) + r['rpr_R'] = rr; r['rpr_L'] = rl + print(f" repair={rr:.2f}ms({rr/tr:.1f}x)", end="") + print() + + print(f" left: tree={tl:.2f}ms dense={dl:.2f}ms({dl/tl:.1f}x)", end="") + if 'grg_L' in r: + print(f" grg={r['grg_L']:.2f}ms({r['grg_L']/tl:.1f}x)", end="") + if 'rpr_L' in r: + print(f" repair={r['rpr_L']:.2f}ms({r['rpr_L']/tl:.1f}x)", end="") + print() + + # Data consistency (optional, slow) + if args.dc: + gp = os.path.join(DC, f"sim_{n}.genotypes.npz") + if os.path.exists(gp): + from threads_arg import AgeEstimator, GenotypeIterator, ConsistencyWrapper + print(f" DC: building...", end="", flush=True) + inst_dc = load_instructions(tp) + t0 = time.perf_counter() + ae = AgeEstimator(inst_dc) + gi = GenotypeIterator(inst_dc) + while gi.has_next_genotype(): + ae.process_site(np.array(gi.next_genotype())) + ages = ae.get_inferred_ages() + cw = ConsistencyWrapper(inst_dc, ages) + gi2 = GenotypeIterator(inst_dc) + while gi2.has_next_genotype(): + cw.process_site(gi2.next_genotype()) + dc_inst = cw.get_consistent_instructions() + dc_t = time.perf_counter() - t0 + with tempfile.NamedTemporaryFile(suffix='.threads', delete=False) as f: + dcp = f.name + serialize_instructions(dc_inst, dcp, allele_ages=ages) + dc_sz = os.path.getsize(dcp) + dc_inst2 = load_instructions(dcp) + dc_inst2.prepare_tree_multiply() + xsd = rng.normal(0, 1, nm).tolist() + xhd = rng.normal(0, 1, dc_inst2.num_samples).tolist() + dcr = bench(lambda: dc_inst2.right_multiply_tree(xsd), args.reps) + dcl = bench(lambda: dc_inst2.left_multiply_tree(xhd), args.reps) + os.unlink(dcp) + r['dc_bytes'] = dc_sz; r['dc_build_s'] = dc_t + r['dc_R'] = dcr; r['dc_L'] = dcl + print(f" {dc_t:.1f}s size={sz(dc_sz)}({dc_sz/thr_sz:.2f}x std) " + f"right={dcr:.2f}ms({dcr/tr:.1f}x) left={dcl:.2f}ms({dcl/tl:.1f}x)") + + results.append(r) + + out = os.path.join(DC, "bench_compression_multiply.json") + with open(out, 'w') as f: + json.dump(results, f, indent=2) + print(f"\nSaved to {out}") + + +if __name__ == "__main__": + main() diff --git a/test/bench_impute.py b/test/bench_impute.py new file mode 100644 index 0000000..f1c1bf4 --- /dev/null +++ b/test/bench_impute.py @@ -0,0 +1,340 @@ +#!/usr/bin/env python3 +""" +Imputation benchmark: correctness (bit-identical output), speed, and memory. + +Usage: + python test/bench_impute.py # run baseline, save reference + python test/bench_impute.py --compare # run optimized, compare to reference + python test/bench_impute.py --repeat 3 # best-of-3 runs + +The first run (without --compare) records: + - The full VCF output (reference for bit-identity checks) + - Wall time for each pipeline stage + - Peak RSS memory + +The second run (with --compare) re-runs imputation and checks: + 1. Output is bit-identical to the reference (ignoring the date header line) + 2. Wall time per stage (prints speedup ratios) + 3. Peak memory (prints reduction) + +Results are saved to test/bench_impute_results/ so you can compare across +code changes without re-running the baseline. +""" +import argparse +import json +import logging +import os +import re +import resource +import sys +import tempfile +import time + +from pathlib import Path + +# --------------------------------------------------------------------------- +# Paths +# --------------------------------------------------------------------------- +BASE_DIR = Path(__file__).parent.parent +TEST_DATA_DIR = BASE_DIR / "test" / "data" +RESULTS_DIR = BASE_DIR / "test" / "bench_impute_results" + +# Inputs (pre-generated snapshot fixtures) +PANEL_VCF = TEST_DATA_DIR / "panel.vcf.gz" +TARGET_VCF = TEST_DATA_DIR / "target.vcf.gz" +GMAP = TEST_DATA_DIR / "gmap_04.map" +MUT = TEST_DATA_DIR / "expected_mapping_snapshot.mut" +DEMO = TEST_DATA_DIR / "CEU_unscaled.demo" +REGION = "1:400000-600000" + +# Reference snapshot for bit-identity (the existing regression fixture) +EXPECTED_VCF = TEST_DATA_DIR / "expected_impute_snapshot.vcf" + + +# --------------------------------------------------------------------------- +# Measurement helpers +# --------------------------------------------------------------------------- +def peak_rss_mb(): + """Current peak RSS in MB (macOS returns bytes, Linux returns KB).""" + ru = resource.getrusage(resource.RUSAGE_SELF) + if sys.platform == "darwin": + return ru.ru_maxrss / (1024 * 1024) + else: + return ru.ru_maxrss / 1024 + + +class TimingCapture(logging.Handler): + """ + Logging handler that captures Finished messages from timer_block to + extract per-stage wall times without modifying production code. + + timer_block emits: "Finished (time s)" + TimerTotal emits: "Total time for : s" + """ + FINISHED_RE = re.compile(r"Finished (.+) \(time ([\d.]+)s\)") + TOTAL_RE = re.compile(r"Total time for (.+): ([\d.]+)s") + + def __init__(self): + super().__init__() + self.stages = {} + + def emit(self, record): + msg = record.getMessage() + m = self.FINISHED_RE.search(msg) + if m: + self.stages[m.group(1)] = float(m.group(2)) + return + m = self.TOTAL_RE.search(msg) + if m: + self.stages[m.group(1)] = float(m.group(2)) + + +# --------------------------------------------------------------------------- +# Core: run imputation with per-stage timing +# --------------------------------------------------------------------------- +def run_impute_timed(out_vcf_path): + """ + Run the full Impute() pipeline, returning: + - wall_total: total wall time + - stages: dict of internal stage timings (from timer_block logging) + - rss_after: peak RSS after run + - rss_before: peak RSS before run + """ + # Set up logging capture + capture = TimingCapture() + root_logger = logging.getLogger("threads_arg") + root_logger.setLevel(logging.INFO) + root_logger.addHandler(capture) + + rss_before = peak_rss_mb() + + t0 = time.perf_counter() + from threads_arg.impute import Impute + Impute( + PANEL_VCF, + TARGET_VCF, + GMAP, + MUT, + DEMO, + out_vcf_path, + REGION, + ) + wall_total = time.perf_counter() - t0 + + rss_after = peak_rss_mb() + + # Clean up handler + root_logger.removeHandler(capture) + + # Build timings dict: internal stages first, then total + timings = {} + for name, t in capture.stages.items(): + timings[name] = t + timings["wall_total"] = wall_total + + return timings, rss_after, rss_before + + +# --------------------------------------------------------------------------- +# VCF comparison (bit-identical, ignoring date line) +# --------------------------------------------------------------------------- +def compare_vcf(expected_path, generated_path): + """ + Compare two VCF files line-by-line. Ignores line 3 (##fileDate). + Returns (match: bool, first_diff_line: int|None, detail: str). + """ + with open(expected_path) as ef, open(generated_path) as gf: + exp_lines = ef.readlines() + gen_lines = gf.readlines() + + if len(exp_lines) != len(gen_lines): + return False, None, f"line count mismatch: expected {len(exp_lines)}, got {len(gen_lines)}" + + for i, (exp_line, gen_line) in enumerate(zip(exp_lines, gen_lines), start=1): + # Line 3 is the date header — skip it + if i == 3: + if not exp_line.startswith("##fileDate") or not gen_line.startswith("##fileDate"): + return False, i, f"line 3 not a date header" + continue + if exp_line != gen_line: + return False, i, f"expected: {exp_line[:80].rstrip()}...\n got: {gen_line[:80].rstrip()}..." + + return True, None, "bit-identical" + + +# --------------------------------------------------------------------------- +# Save / load results +# --------------------------------------------------------------------------- +def save_results(tag, timings, rss_peak, rss_before, vcf_path): + RESULTS_DIR.mkdir(parents=True, exist_ok=True) + result = { + "tag": tag, + "timings": timings, + "rss_peak_mb": rss_peak, + "rss_before_mb": rss_before, + "rss_delta_mb": rss_peak - rss_before, + "vcf_path": str(vcf_path), + "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"), + } + out_path = RESULTS_DIR / f"{tag}.json" + with open(out_path, "w") as f: + json.dump(result, f, indent=2) + print(f" Results saved to {out_path}") + return result + + +def load_results(tag): + path = RESULTS_DIR / f"{tag}.json" + if not path.exists(): + return None + with open(path) as f: + return json.load(f) + + +# --------------------------------------------------------------------------- +# Print helpers +# --------------------------------------------------------------------------- +def print_timings(label, timings): + print(f"\n {label}:") + for stage, t in timings.items(): + print(f" {stage:35s} {t:8.3f}s") + + +def print_comparison(baseline, current): + print("\n" + "=" * 72) + print(" COMPARISON: baseline vs current") + print("=" * 72) + + # Timings — show all stages from either run + bt = baseline["timings"] + ct = current["timings"] + all_stages = list(dict.fromkeys(list(bt.keys()) + list(ct.keys()))) + + print(f"\n {'Stage':35s} {'Baseline':>9s} {'Current':>9s} {'Speedup':>8s}") + print(f" {'-'*35} {'-'*9} {'-'*9} {'-'*8}") + for stage in all_stages: + b = bt.get(stage, 0) + c = ct.get(stage, 0) + if b > 0 and c > 0: + ratio = b / c + marker = " <--" if ratio > 1.05 else (" SLOW" if ratio < 0.95 else "") + print(f" {stage:35s} {b:8.3f}s {c:8.3f}s {ratio:7.2f}x{marker}") + elif c > 0: + print(f" {stage:35s} {'n/a':>9s} {c:8.3f}s") + else: + print(f" {stage:35s} {b:8.3f}s {'n/a':>9s}") + + # Memory + print(f"\n {'Memory':35s} {'Baseline':>9s} {'Current':>9s} {'Change':>8s}") + print(f" {'-'*35} {'-'*9} {'-'*9} {'-'*8}") + bd = baseline["rss_delta_mb"] + cd = current["rss_delta_mb"] + reduction = (1 - cd / bd) * 100 if bd > 0 else 0 + sign = "-" if reduction > 0 else "+" + print(f" {'RSS delta':35s} {bd:7.1f}MB {cd:7.1f}MB {sign}{abs(reduction):.1f}%") + bp = baseline["rss_peak_mb"] + cp = current["rss_peak_mb"] + print(f" {'RSS peak':35s} {bp:7.1f}MB {cp:7.1f}MB") + print() + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- +def main(): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--compare", action="store_true", + help="Run as 'current' and compare against saved baseline") + parser.add_argument("--tag", default=None, + help="Label for this run (default: 'baseline' or 'current')") + parser.add_argument("--repeat", type=int, default=1, + help="Number of repetitions (reports best wall_total)") + args = parser.parse_args() + + tag = args.tag or ("current" if args.compare else "baseline") + + # Validate inputs exist + for path, label in [(PANEL_VCF, "panel"), (TARGET_VCF, "target"), + (GMAP, "gmap"), (MUT, "mut"), (DEMO, "demo")]: + if not path.exists(): + print(f"ERROR: {label} not found at {path}") + sys.exit(1) + + print(f"{'=' * 72}") + print(f" Imputation Benchmark — tag: {tag}") + print(f" Region: {REGION} | Repeats: {args.repeat}") + print(f"{'=' * 72}") + + best_timings = None + best_total = float("inf") + rss_peak = 0 + rss_before = 0 + + for rep in range(args.repeat): + with tempfile.TemporaryDirectory() as tmpdir: + out_vcf = Path(tmpdir) / "imputed.vcf" + timings, rss_p, rss_b = run_impute_timed(str(out_vcf)) + wt = timings["wall_total"] + + if wt < best_total: + best_timings = timings + best_total = wt + rss_peak = rss_p + rss_before = rss_b + + # Persist the VCF from the best run + RESULTS_DIR.mkdir(parents=True, exist_ok=True) + vcf_out_path = RESULTS_DIR / f"{tag}.vcf" + with open(out_vcf) as src, open(vcf_out_path, "w") as dst: + dst.write(src.read()) + + if args.repeat > 1: + print(f" rep {rep+1}/{args.repeat}: {wt:.3f}s (best: {best_total:.3f}s)") + + vcf_out_path = RESULTS_DIR / f"{tag}.vcf" + + # Print results + print_timings(tag, best_timings) + delta = rss_peak - rss_before + print(f"\n Peak RSS: {rss_peak:.1f} MB (delta from start: {delta:.1f} MB)") + + # Bit-identity check against expected snapshot + print(f"\n Bit-identity vs {EXPECTED_VCF.name}...") + match, diff_line, detail = compare_vcf(EXPECTED_VCF, vcf_out_path) + if match: + print(f" PASS: bit-identical to expected snapshot") + else: + print(f" FAIL: differs at line {diff_line}") + print(f" {detail}") + + # Save + result = save_results(tag, best_timings, rss_peak, rss_before, vcf_out_path) + + # Compare mode + if args.compare: + baseline = load_results("baseline") + if baseline is None: + print("\n WARNING: no baseline found. Run without --compare first.") + else: + print_comparison(baseline, result) + + # Bit-identity between baseline and current + baseline_vcf = Path(baseline["vcf_path"]) + if baseline_vcf.exists(): + match2, diff2, detail2 = compare_vcf(baseline_vcf, vcf_out_path) + if match2: + print(" Baseline vs current VCF: BIT-IDENTICAL") + else: + print(f" Baseline vs current VCF: DIFFERS at line {diff2}") + print(f" {detail2}") + print() + + # Summary exit + status = "PASS" if match else "FAIL" + print(f" [{status}] {tag} {best_total:.3f}s {rss_peak:.0f}MB") + print() + + +if __name__ == "__main__": + main() diff --git a/test/bench_impute_scaling.py b/test/bench_impute_scaling.py new file mode 100644 index 0000000..c83e703 --- /dev/null +++ b/test/bench_impute_scaling.py @@ -0,0 +1,343 @@ +#!/usr/bin/env python3 +""" +Scaling microbenchmark for imputation hot paths. + +Isolates fwbw, sparsify, per-variant dosage computation, posterior cache +rebuild, and VCF write formatting from VCF I/O. +Uses msprime to generate synthetic data at configurable scale. + +Usage: + python test/bench_impute_scaling.py + python test/bench_impute_scaling.py --preset medium + python test/bench_impute_scaling.py --preset large + python test/bench_impute_scaling.py --n-panel 2000 --n-target 100 +""" +import argparse +import time +import sys + +import numpy as np + +# --------------------------------------------------------------------------- +# Simulation +# --------------------------------------------------------------------------- +def simulate_data(n_panel_dip, n_target_dip, seq_length=2e6, seed=42): + """Generate biallelic haplotype matrices from msprime.""" + import msprime + + n_total = n_panel_dip + n_target_dip + ts = msprime.sim_ancestry( + samples=n_total, sequence_length=seq_length, + recombination_rate=1.3e-8, population_size=10000, random_seed=seed) + ts = msprime.sim_mutations(ts, rate=1.4e-8, random_seed=seed + 1) + + # Extract biallelic sites + positions_bp = [] + positions_cm = [] + genotypes = [] + for var in ts.variants(): + g = var.genotypes.astype(np.int8) + if len(np.unique(g)) != 2 or np.any((g != 0) & (g != 1)): + continue + positions_bp.append(var.site.position) + positions_cm.append(var.site.position * 1.3e-8 * 100) + genotypes.append(g) + + G = np.array(genotypes, dtype=bool) # (n_sites, n_haps) + n_panel_haps = 2 * n_panel_dip + panel_snps = G[:, :n_panel_haps] + target_snps = G[:, n_panel_haps:] + + cm_pos = np.array(positions_cm) + + return panel_snps, target_snps, cm_pos + + +# --------------------------------------------------------------------------- +# Benchmark functions +# --------------------------------------------------------------------------- +def bench_set_emission(panel_snps, target_hap, mutation_rate): + """Benchmark set_emission_probabilities.""" + from threads_arg.fwbw import set_emission_probabilities + query = target_hap[None, :] + t0 = time.perf_counter() + e = set_emission_probabilities(panel_snps, query, mutation_rate) + return time.perf_counter() - t0, e + + +def bench_fwbw(panel_subset, target_hap, recomb_rates, mutation_rate): + """Benchmark full fwbw call.""" + from threads_arg.fwbw import fwbw + query = target_hap[None, :] + t0 = time.perf_counter() + posterior = fwbw(panel_subset, query, recomb_rates, mutation_rate) + return time.perf_counter() - t0, posterior + + +def bench_sparsify(posterior, matched_samples, num_samples_panel, num_snps): + """Benchmark _sparsify_posterior logic (extracted).""" + from scipy.sparse import csr_array + + t0 = time.perf_counter() + posterior[posterior <= 1 / num_samples_panel] = 0 + row_sums = posterior.sum(axis=1) + posterior = posterior / row_sums[:, np.newaxis] + rows, cols_local = np.nonzero(posterior) + vals = posterior[rows, cols_local] + cols_global = matched_samples[cols_local] + sparse = csr_array( + (vals, (rows, cols_global)), + shape=(num_snps, num_samples_panel) + ) + return time.perf_counter() - t0, sparse + + +def bench_genotype_sum(posteriors_2d, n_repeats=100): + """Benchmark vectorized vs list-comprehension genotype sum.""" + # Vectorized (current) + t0 = time.perf_counter() + for _ in range(n_repeats): + g1 = posteriors_2d.sum(axis=1) + t_vec = (time.perf_counter() - t0) / n_repeats + + # List comprehension (original) + t0 = time.perf_counter() + for _ in range(n_repeats): + g2 = np.array([np.sum(asp) for asp in posteriors_2d]) + t_list = (time.perf_counter() - t0) / n_repeats + + assert np.allclose(g1, g2) + return t_vec, t_list + + +def bench_posterior_cache(sparse_posteriors, n_snps, n_repeats=3): + """Benchmark CachedPosteriorSnps rebuild: vstack vs original loop.""" + from scipy.sparse import vstack as sparse_vstack + + n_targets = len(sparse_posteriors) + + # Current: vstack batch conversion + t0 = time.perf_counter() + for _ in range(n_repeats): + for snp_idx in range(n_snps): + rows = sparse_vstack([p[[snp_idx]] for p in sparse_posteriors]) + tp = rows.toarray() + rs = tp.sum(axis=1, keepdims=True) + tp /= rs + t_vstack = (time.perf_counter() - t0) / n_repeats / n_snps + + # Original: per-target loop + t0 = time.perf_counter() + for _ in range(n_repeats): + for snp_idx in range(n_snps): + n_panel = sparse_posteriors[0].shape[1] + tp = np.empty((n_targets, n_panel), dtype=np.float64) + for i, p in enumerate(sparse_posteriors): + row = p[[snp_idx], :].toarray() + tp[i] = row / np.sum(row) + t_loop = (time.perf_counter() - t0) / n_repeats / n_snps + + return t_vstack, t_loop + + +def bench_write_format(n_samples, n_repeats=200): + """Benchmark VCF GT:DS string formatting: vectorized round vs per-element.""" + rng = np.random.default_rng(42) + genotypes = rng.random(2 * n_samples) + + # Current: vectorized rint + t0 = time.perf_counter() + for _ in range(n_repeats): + haps1 = genotypes[::2] + haps2 = genotypes[1::2] + dosages = haps1 + haps2 + gt1 = np.rint(haps1).astype(int) + gt2 = np.rint(haps2).astype(int) + _ = [f"{g1}|{g2}:{d:.3f}".rstrip("0").rstrip(".") + for g1, g2, d in zip(gt1, gt2, dosages)] + t_vec = (time.perf_counter() - t0) / n_repeats + + # Original: per-element np.round + t0 = time.perf_counter() + for _ in range(n_repeats): + haps1 = genotypes[::2] + haps2 = genotypes[1::2] + dosages = haps1 + haps2 + _ = [f"{np.round(h1):.0f}|{np.round(h2):.0f}:{d:.3f}".rstrip("0").rstrip(".") + for h1, h2, d in zip(haps1, haps2, dosages)] + t_orig = (time.perf_counter() - t0) / n_repeats + + return t_vec, t_orig + + +PRESETS = { + "small": {"n_panel": 500, "n_target": 50, "seq_length": 2e6, "cond_size": 40}, + "medium": {"n_panel": 1000, "n_target": 100, "seq_length": 5e6, "cond_size": 60}, + "large": {"n_panel": 2000, "n_target": 200, "seq_length": 10e6, "cond_size": 80}, +} + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--preset", choices=PRESETS.keys(), default=None, + help="Use a preset configuration (small/medium/large)") + parser.add_argument("--n-panel", type=int, default=None, + help="Panel diploid count (default 500)") + parser.add_argument("--n-target", type=int, default=None, + help="Target diploid count (default 50)") + parser.add_argument("--seq-length", type=float, default=None, + help="Sequence length in bp (default 2Mb)") + parser.add_argument("--cond-size", type=int, default=None, + help="Conditioning set size per fwbw call (default 40)") + args = parser.parse_args() + + # Apply preset defaults, then CLI overrides + preset = PRESETS.get(args.preset, PRESETS["small"]) + n_panel_dip = args.n_panel or preset["n_panel"] + n_target_dip = args.n_target or preset["n_target"] + seq_length = args.seq_length or preset["seq_length"] + cond_size_cfg = args.cond_size or preset["cond_size"] + n_panel_haps = 2 * n_panel_dip + n_target_haps = 2 * n_target_dip + cond_size = min(cond_size_cfg, n_panel_haps) + + print(f"{'=' * 65}") + print(f" Imputation Scaling Benchmark") + print(f" Panel: {n_panel_dip} dip ({n_panel_haps} hap)") + print(f" Target: {n_target_dip} dip ({n_target_haps} hap)") + print(f" Seq length: {seq_length/1e6:.1f} Mb") + print(f" Conditioning set: {cond_size}") + print(f"{'=' * 65}") + + # Simulate + print("\n Simulating...", end="", flush=True) + t0 = time.perf_counter() + panel_snps, target_snps, cm_pos = simulate_data( + n_panel_dip, n_target_dip, seq_length) + t_sim = time.perf_counter() - t0 + m, n = panel_snps.shape + print(f" {t_sim:.1f}s ({m} sites, {n} panel haps, " + f"{target_snps.shape[1]} target haps)") + + # Recombination rates + cm_sizes = np.diff(cm_pos, append=cm_pos[-1] - cm_pos[-2] + cm_pos[-1]) + cm_sizes = np.maximum(cm_sizes, 1e-10) + Ne = 20_000 + recomb_rates = 1 - np.exp(-4 * Ne * 0.01 * cm_sizes / n_panel_haps) + mutation_rate = 0.0001 + + # Warmup numba + print(" Warming up numba JIT...", end="", flush=True) + rng = np.random.default_rng(0) + cond_idx = rng.choice(n_panel_haps, size=cond_size, replace=False) + cond_idx.sort() + h0 = target_snps[:, 0] + _ = bench_fwbw(panel_snps[:, cond_idx], h0, recomb_rates, mutation_rate) + print(" done") + + # --- Benchmark set_emission_probabilities --- + n_emit = min(5, n_target_haps) + t_emit_total = 0 + for i in range(n_emit): + t_e, _ = bench_set_emission(panel_snps[:, cond_idx], + target_snps[:, i], mutation_rate) + t_emit_total += t_e + t_emit_avg = t_emit_total / n_emit + + # --- Benchmark fwbw --- + n_fwbw = min(5, n_target_haps) + t_fwbw_total = 0 + for i in range(n_fwbw): + ci = rng.choice(n_panel_haps, size=cond_size, replace=False) + ci.sort() + t_f, posterior = bench_fwbw(panel_snps[:, ci], target_snps[:, i], + recomb_rates, mutation_rate) + t_fwbw_total += t_f + t_fwbw_avg = t_fwbw_total / n_fwbw + + # --- Benchmark sparsify --- + n_sp = min(5, n_target_haps) + t_sp_total = 0 + for i in range(n_sp): + ci = rng.choice(n_panel_haps, size=cond_size, replace=False) + ci.sort() + _, post = bench_fwbw(panel_snps[:, ci], target_snps[:, i], + recomb_rates, mutation_rate) + t_s, _ = bench_sparsify(post.copy(), ci, n_panel_haps, m) + t_sp_total += t_s + t_sp_avg = t_sp_total / n_sp + + # --- Benchmark genotype sum --- + dummy_posteriors = rng.random((n_target_haps, cond_size)) + t_vec, t_list = bench_genotype_sum(dummy_posteriors) + + # --- Benchmark posterior cache rebuild --- + # Build a small set of sparse posteriors for the cache benchmark + print(" Benchmarking posterior cache...", end="", flush=True) + n_cache_targets = min(n_target_haps, 20) # match typical batch size + sparse_posteriors = [] + for i in range(n_cache_targets): + ci = rng.choice(n_panel_haps, size=cond_size, replace=False) + ci.sort() + _, post = bench_fwbw(panel_snps[:, ci], target_snps[:, i % target_snps.shape[1]], + recomb_rates, mutation_rate) + # Sparsify + _, sp = bench_sparsify(post.copy(), ci, n_panel_haps, m) + sparse_posteriors.append(sp) + n_cache_snps = min(20, m) # test a subset of SNPs + t_cache_vstack, t_cache_loop = bench_posterior_cache(sparse_posteriors, n_cache_snps) + print(" done") + + # --- Benchmark write formatting --- + n_write_samples = n_target_dip # diploid target samples + t_write_vec, t_write_orig = bench_write_format(n_write_samples) + + # --- Projected totals --- + # In real imputation: n_target_haps fwbw calls, ~m variants with genotype sum + n_variants = m + t_fwbw_projected = t_fwbw_avg * n_target_haps + t_emit_projected = t_emit_avg * n_target_haps + t_sp_projected = t_sp_avg * n_target_haps + t_geno_vec_projected = t_vec * n_variants + t_geno_list_projected = t_list * n_variants + t_cache_vstack_projected = t_cache_vstack * n_variants + t_cache_loop_projected = t_cache_loop * n_variants + t_write_vec_projected = t_write_vec * n_variants + t_write_orig_projected = t_write_orig * n_variants + + # Print results + print(f"\n {'Operation':35s} {'Per-call':>10s} {'Projected':>10s}") + print(f" {'-'*35} {'-'*10} {'-'*10}") + print(f" {'set_emission_probabilities':35s} {t_emit_avg*1000:9.3f}ms " + f"{t_emit_projected:9.3f}s") + print(f" {'fwbw (forward-backward)':35s} {t_fwbw_avg*1000:9.3f}ms " + f"{t_fwbw_projected:9.3f}s") + print(f" {'sparsify_posterior':35s} {t_sp_avg*1000:9.3f}ms " + f"{t_sp_projected:9.3f}s") + print(f" {'genotype_sum (vectorized)':35s} {t_vec*1e6:9.1f}us " + f"{t_geno_vec_projected:9.3f}s") + print(f" {'genotype_sum (list comp)':35s} {t_list*1e6:9.1f}us " + f"{t_geno_list_projected:9.3f}s") + print(f" {'posterior_cache (vstack)':35s} {t_cache_vstack*1000:9.3f}ms " + f"{t_cache_vstack_projected:9.3f}s") + print(f" {'posterior_cache (loop)':35s} {t_cache_loop*1000:9.3f}ms " + f"{t_cache_loop_projected:9.3f}s") + print(f" {'write_format (vectorized)':35s} {t_write_vec*1e6:9.1f}us " + f"{t_write_vec_projected:9.3f}s") + print(f" {'write_format (per-element)':35s} {t_write_orig*1e6:9.1f}us " + f"{t_write_orig_projected:9.3f}s") + + print(f"\n Projected total for {n_target_haps} targets, {n_variants} variants:") + t_total_opt = t_fwbw_projected + t_sp_projected + t_geno_vec_projected + t_cache_vstack_projected + t_write_vec_projected + t_total_orig = t_fwbw_projected + t_sp_projected + t_geno_list_projected + t_cache_loop_projected + t_write_orig_projected + print(f" optimized: {t_total_opt:.3f}s") + print(f" original: {t_total_orig:.3f}s") + print(f" speedup: {t_total_orig/t_total_opt:.2f}x") + print() + + +if __name__ == "__main__": + main() diff --git a/test/build_cache.py b/test/build_cache.py new file mode 100644 index 0000000..e3d15bc --- /dev/null +++ b/test/build_cache.py @@ -0,0 +1,277 @@ +""" +Pre-generate and cache benchmark data for all sample sizes. + +Outputs in threads_cache/: + sim_{n_dip}.trees — msprime true tree sequence + sim_{n_dip}.inferred.trees — tsinfer inferred tree sequence + sim_{n_dip}.grg_true — GRG from true ARG + sim_{n_dip}.grg_inferred — GRG from inferred ARG + sim_{n_dip}.genotypes.npz — genotype matrix + positions + sim_{n_dip}.threads — serialized ThreadingInstructions + sim_{n_dip}.repair — RePair compressed matrix + +Usage: + python test/build_cache.py [--sizes 50,100,500,...] [--cache-dir threads_cache] +""" + +import os, sys, time, argparse, json, glob, resource, platform +import numpy as np +import msprime, tsinfer, pygrgl +import threads_arg, genrepair +from threads_arg.serialization import serialize_instructions, load_instructions + +SEED = 42 +SEQ_LENGTH = 2e6 + + +def get_rss_bytes(): + """Current resident set size in bytes.""" + ru = resource.getrusage(resource.RUSAGE_SELF) + if platform.system() == "Darwin": + return ru.ru_maxrss # macOS: already bytes + return ru.ru_maxrss * 1024 # Linux: KB + + +def measure_memory(fn): + """Run fn(), return (result, peak_rss_delta_bytes).""" + import gc; gc.collect() + before = get_rss_bytes() + result = fn() + after = get_rss_bytes() + return result, max(0, after - before) + +def ts_to_biallelic(ts): + """Extract biallelic sites from tree sequence.""" + positions_bp, positions_cm, genotype_matrix = [], [], [] + for var in ts.variants(): + g = var.genotypes.astype(int).tolist() + if len(set(g)) < 2 or any(a not in (0, 1) for a in g): + continue + positions_bp.append(var.site.position) + positions_cm.append(var.site.position * 1.3e-8 * 100) + genotype_matrix.append(g) + return positions_bp, positions_cm, genotype_matrix + + +def build_threads_instructions(ts, positions_bp, positions_cm, genotype_matrix): + """Run full threads pipeline, return ThreadingInstructions.""" + n_haps = ts.num_samples + matcher = threads_arg.Matcher(n_haps, positions_cm, 0.01, 0.5, 4, 2) + for g in genotype_matrix: + matcher.process_site(g) + target_ids = list(range(1, n_haps)) + match_data = matcher.serializable_matches(target_ids) + cm_pos = matcher.cm_positions() + + tlm = threads_arg.ThreadsLowMem( + target_ids, positions_bp, positions_cm, [10000.0], [0.0], 1.4e-8, False) + tlm.initialize_viterbi(match_data, cm_pos) + G_np = np.array(genotype_matrix, dtype=np.int32) + tlm.process_all_sites_viterbi_numpy(G_np) + tlm.prune() + tlm.traceback() + tlm.process_all_sites_hets_numpy(G_np) + tlm.date_segments() + + all_starts, all_ids, all_heights, all_hetsites = tlm.serialize_paths() + positions_int = [int(p) for p in positions_bp] + # serialize_paths returns site indices as starts; convert to physical positions + all_starts = [[positions_int[s] for s in sample_starts] for sample_starts in all_starts] + ti = threads_arg.ThreadingInstructions( + all_starts, all_heights, all_ids, all_hetsites, + positions_int, positions_int[0], positions_int[-1] + 1) + return ti + + + +def build_one(n_dip, cache_dir): + """Generate all cached files for one sample size. Returns timing dict.""" + prefix = os.path.join(cache_dir, f"sim_{n_dip}") + ts_path = f"{prefix}.trees" + inf_path = f"{prefix}.inferred.trees" + grg_true_path = f"{prefix}.grg_true" + grg_inf_path = f"{prefix}.grg_inferred" + geno_path = f"{prefix}.genotypes.npz" + threads_path = f"{prefix}.threads" + repair_path = f"{prefix}.repair" + + timing = {"n_dip": n_dip} + + print(f"\n{'='*70}") + print(f" {n_dip} diploid ({2*n_dip} haploid)") + print(f"{'='*70}", flush=True) + + # 1. Simulate + t0 = time.perf_counter() + ts = msprime.sim_ancestry( + samples=n_dip, sequence_length=SEQ_LENGTH, + recombination_rate=1.3e-8, population_size=10000, random_seed=SEED) + ts = msprime.sim_mutations(ts, rate=1.4e-8, random_seed=SEED + 1) + ts.dump(ts_path) + positions_bp, positions_cm, genotype_matrix = ts_to_biallelic(ts) + n_haps = ts.num_samples + n_sites = len(genotype_matrix) + timing["simulate_s"] = time.perf_counter() - t0 + timing["n_hap"] = n_haps + timing["n_sites"] = n_sites + print(f" simulate: {timing['simulate_s']:.1f}s ({n_haps} hap, {n_sites} sites)", flush=True) + + # 2. Save genotypes + G_np = np.array(genotype_matrix, dtype=np.int8) + np.savez_compressed(geno_path, + genotypes=G_np, + positions_bp=np.array(positions_bp), + positions_cm=np.array(positions_cm)) + timing["genotypes_bytes"] = os.path.getsize(geno_path) + print(f" genotypes: {timing['genotypes_bytes']/1024:.0f}K", flush=True) + + # 3. GRG from true ARG + t0 = time.perf_counter() + grg_true, mem = measure_memory( + lambda: pygrgl.grg_from_trees(ts_path, binary_mutations=True)) + pygrgl.save_grg(grg_true, grg_true_path) + timing["grg_true_build_s"] = time.perf_counter() - t0 + timing["grg_true_bytes"] = os.path.getsize(grg_true_path) + timing["grg_true_mem_bytes"] = mem + timing["grg_true_mutations"] = grg_true.num_mutations + print(f" grg_true: {timing['grg_true_build_s']:.1f}s mem={mem/1024/1024:.0f}M " + f"({grg_true.num_samples} samples, {grg_true.num_mutations} mutations, " + f"{timing['grg_true_bytes']/1024:.0f}K)", flush=True) + + # 4. tsinfer + GRG from inferred ARG + t0 = time.perf_counter() + def _run_tsinfer(): + with tsinfer.SampleData(sequence_length=ts.sequence_length) as sd: + for var in ts.variants(): + g = var.genotypes + alleles = var.alleles + if len(alleles) != 2 or not all(a in (0, 1) for a in g): + continue + if len(set(g)) < 2: + continue + sd.add_site(var.site.position, g, alleles=list(alleles)) + return tsinfer.infer(sd) + ts_inf, tsinfer_mem = measure_memory(_run_tsinfer) + ts_inf.dump(inf_path) + timing["tsinfer_build_s"] = time.perf_counter() - t0 + timing["tsinfer_mem_bytes"] = tsinfer_mem + print(f" tsinfer: {timing['tsinfer_build_s']:.1f}s mem={tsinfer_mem/1024/1024:.0f}M", flush=True) + + t0 = time.perf_counter() + grg_inf, grg_inf_mem = measure_memory( + lambda: pygrgl.grg_from_trees(inf_path, binary_mutations=True)) + pygrgl.save_grg(grg_inf, grg_inf_path) + timing["grg_inf_convert_s"] = time.perf_counter() - t0 + timing["grg_inf_total_s"] = timing["tsinfer_build_s"] + timing["grg_inf_convert_s"] + timing["grg_inf_bytes"] = os.path.getsize(grg_inf_path) + timing["grg_inf_mem_bytes"] = tsinfer_mem + grg_inf_mem + timing["grg_inf_mutations"] = grg_inf.num_mutations + print(f" grg_inferred: {timing['grg_inf_convert_s']:.1f}s convert, " + f"{timing['grg_inf_total_s']:.1f}s total mem={timing['grg_inf_mem_bytes']/1024/1024:.0f}M " + f"({grg_inf.num_samples} samples, {grg_inf.num_mutations} mutations, " + f"{timing['grg_inf_bytes']/1024:.0f}K)", flush=True) + + # 5. Threads + t0 = time.perf_counter() + ti, threads_mem = measure_memory( + lambda: build_threads_instructions(ts, positions_bp, positions_cm, genotype_matrix)) + timing["threads_build_s"] = time.perf_counter() - t0 + timing["threads_mem_bytes"] = threads_mem + t0 = time.perf_counter() + serialize_instructions(ti, threads_path) + timing["threads_serialize_s"] = time.perf_counter() - t0 + timing["threads_bytes"] = os.path.getsize(threads_path) + timing["threads_n_samples"] = ti.num_samples + timing["threads_n_sites"] = ti.num_sites + print(f" threads: {timing['threads_build_s']:.1f}s build + " + f"{timing['threads_serialize_s']:.1f}s serialize mem={threads_mem/1024/1024:.0f}M " + f"({ti.num_samples} hap, {ti.num_sites} sites, " + f"{timing['threads_bytes']/1024:.0f}K)", flush=True) + + # 7. Prepare tree multiply (measure separately) + t0 = time.perf_counter() + _, prep_mem = measure_memory(lambda: ti.prepare_tree_multiply()) + timing["tree_prepare_s"] = time.perf_counter() - t0 + timing["tree_prepare_mem_bytes"] = prep_mem + print(f" tree_prepare: {timing['tree_prepare_s']*1000:.1f}ms mem={prep_mem/1024/1024:.0f}M", flush=True) + + # 6. RePair (saves multiple files with repair_path as basename) + t0 = time.perf_counter() + cm, repair_mem = measure_memory( + lambda: genrepair.CompressedMatrix.from_numpy( + G_np, repair_path, diploid=False, standardized=False)) + timing["repair_build_s"] = time.perf_counter() - t0 + timing["repair_mem_bytes"] = repair_mem + repair_files = glob.glob(f"{repair_path}*") + timing["repair_bytes"] = sum(os.path.getsize(f) for f in repair_files) + print(f" repair: {timing['repair_build_s']:.1f}s mem={repair_mem/1024/1024:.0f}M " + f"({timing['repair_bytes']/1024:.0f}K across {len(repair_files)} files)", flush=True) + + all_files = glob.glob(f"{prefix}*") + total_size = sum(os.path.getsize(f) for f in all_files) + print(f" total cached: {total_size/1024/1024:.1f}M ({len(all_files)} files)", flush=True) + + return timing + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--sizes", default="50,100,250,500,1000,2000,5000,10000,20000", + help="Comma-separated diploid sample sizes") + parser.add_argument("--cache-dir", default="threads_cache") + args = parser.parse_args() + + sizes = [int(s) for s in args.sizes.split(",")] + os.makedirs(args.cache_dir, exist_ok=True) + + print(f"Building cache for sizes: {sizes}") + print(f"Cache dir: {args.cache_dir}") + + all_timing = [] + for n_dip in sizes: + try: + timing = build_one(n_dip, args.cache_dir) + all_timing.append(timing) + except Exception as e: + print(f" ERROR at {n_dip} dip: {e}") + import traceback; traceback.print_exc() + + # Save timing data + timing_path = os.path.join(args.cache_dir, "build_timing.json") + with open(timing_path, "w") as f: + json.dump(all_timing, f, indent=2) + print(f"\nTiming saved to {timing_path}") + + # Print summary table + print(f"\n{'='*120}") + print(f" CONSTRUCTION TIME & MEMORY SUMMARY") + print(f"{'='*120}") + print(f"{'n_dip':>7s} {'hap':>6s} {'sites':>6s} | " + f"{'threads':>10s} {'tsinfer+grg':>12s} {'RePair':>10s} | " + f"{'thr_mem':>8s} {'grg_mem':>8s} {'rpr_mem':>8s} | " + f"{'thr_disk':>9s} {'grg_disk':>9s} {'rpr_disk':>9s}") + print("-" * 120) + for t in all_timing: + def fmt_time(s): + if s < 60: return f"{s:.1f}s" + return f"{s/60:.1f}m" + def fmt_mem(b): + if b < 1024**2: return f"{b/1024:.0f}K" + return f"{b/1024/1024:.0f}M" + def fmt_disk(b): + if b < 1024**2: return f"{b/1024:.0f}K" + return f"{b/1024/1024:.1f}M" + print(f"{t['n_dip']:7d} {t['n_hap']:6d} {t['n_sites']:6d} | " + f"{fmt_time(t['threads_build_s']):>10s} " + f"{fmt_time(t['grg_inf_total_s']):>12s} " + f"{fmt_time(t['repair_build_s']):>10s} | " + f"{fmt_mem(t['threads_mem_bytes']):>8s} " + f"{fmt_mem(t['grg_inf_mem_bytes']):>8s} " + f"{fmt_mem(t['repair_mem_bytes']):>8s} | " + f"{fmt_disk(t['threads_bytes']):>9s} " + f"{fmt_disk(t['grg_inf_bytes']):>9s} " + f"{fmt_disk(t['repair_bytes']):>9s}") + + +if __name__ == "__main__": + main() diff --git a/test/test_allele_ages.py b/test/test_allele_ages.py new file mode 100644 index 0000000..73bac6c --- /dev/null +++ b/test/test_allele_ages.py @@ -0,0 +1,504 @@ +# This file is part of the Threads software suite. +# Copyright (C) 2024-2025 Threads Developers. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +""" +Tests for AgeEstimator and ConsistencyWrapper. + +Unit tests use synthetic ThreadingInstructions with known properties. +Regression tests record exact outputs from real data. +""" + +import numpy as np +import pytest + +from threads_arg import AgeEstimator, GenotypeIterator, ThreadingInstructions, ConsistencyWrapper +from threads_arg.serialization import load_instructions + +from snapshot_runners import TEST_DATA_DIR + + +# --------------------------------------------------------------------------- +# Helpers for synthetic ThreadingInstructions +# --------------------------------------------------------------------------- + +def _make_inst(targets, tmrcas, positions=None, starts=None, mismatches=None): + """Build a ThreadingInstructions from per-sample targets/tmrcas. + + For single-segment cases, pass flat lists: + _make_inst([-1, 0, 1], [0, 500, 200]) + For multi-segment, pass nested lists and explicit starts/positions. + """ + n = len(targets) + if positions is None: + positions = [1000] + if not isinstance(targets[0], list): + targets = [[t] for t in targets] + if not isinstance(tmrcas[0], list): + tmrcas = [[float(t)] for t in tmrcas] + if starts is None: + starts = [[positions[0]] for _ in range(n)] + if mismatches is None: + mismatches = [[] for _ in range(n)] + return ThreadingInstructions( + starts, tmrcas, targets, mismatches, + positions, positions[0], positions[-1] + 1, + ) + + +def _run_age(inst, genotypes_list): + """Run AgeEstimator, return list of ages.""" + ae = AgeEstimator(inst) + for g in genotypes_list: + ae.process_site(g) + return ae.get_inferred_ages() + + +def _run_dc(inst, ages, genotypes_list): + """Run ConsistencyWrapper, return consistent ThreadingInstructions.""" + cw = ConsistencyWrapper(inst, ages) + for g in genotypes_list: + cw.process_site(g) + return cw.get_consistent_instructions() + + +# --------------------------------------------------------------------------- +# Unit tests: AgeEstimator +# --------------------------------------------------------------------------- + + +def test_age_two_carriers_coalesce(): + """Two carriers sharing a recent ancestor get age near their TMRCA.""" + # 0 <- 1 (t=1000) <- 2 (t=500) <- 3 (t=200) + inst = _make_inst([-1, 0, 1, 2], [0, 1000, 500, 200]) + ages = _run_age(inst, [[0, 0, 1, 1]]) + # Carriers at 2,3 coalesce at t=200; sweep gives boundary=200, next=500 + assert ages[0] == pytest.approx(350.0) + + +def test_age_all_carriers(): + """All carriers: age above the deepest TMRCA.""" + inst = _make_inst([-1, 0, 1, 2], [0, 1000, 500, 200]) + ages = _run_age(inst, [[1, 1, 1, 1]]) + assert ages[0] > 1000 + + +def test_age_no_carriers(): + """No carriers: still produces a positive age.""" + inst = _make_inst([-1, 0, 1, 2], [0, 1000, 500, 200]) + ages = _run_age(inst, [[0, 0, 0, 0]]) + assert ages[0] > 0 + + +def test_age_single_carrier(): + """One carrier among non-carriers.""" + inst = _make_inst([-1, 0, 1], [0, 1000, 500]) + ages = _run_age(inst, [[0, 1, 0]]) + assert ages[0] > 0 + + +def test_age_carrier_at_root_only(): + """Only the root sample is a carrier.""" + inst = _make_inst([-1, 0, 1], [0, 500, 200]) + ages = _run_age(inst, [[1, 0, 0]]) + assert ages[0] > 0 + + +def test_age_two_samples(): + """Minimal two-sample case.""" + inst = _make_inst([-1, 0], [0, 100]) + ages = _run_age(inst, [[1, 1]]) + assert ages[0] > 0 + ages2 = _run_age(inst, [[1, 0]]) + assert ages2[0] > 0 + + +def test_age_reflects_tmrca_scale(): + """Closer carriers yield younger age than distant carriers.""" + inst_close = _make_inst([-1, 0, 1], [0, 100, 50]) + inst_far = _make_inst([-1, 0, 1], [0, 10000, 5000]) + age_close = _run_age(inst_close, [[0, 1, 1]])[0] + age_far = _run_age(inst_far, [[0, 1, 1]])[0] + assert age_close < age_far + + +def test_age_self_referencing_target(): + """Self-referencing target (target[i]==i) doesn't hang or crash.""" + inst = _make_inst([-1, 1, 0], [0, 500, 300]) # sample 1 self-refs + ages = _run_age(inst, [[1, 1, 0]]) + assert ages[0] > 0 + + +def test_age_self_ref_on_trace_path(): + """Self-ref sample on the trace path from a carrier chain. + + Trace from carrier chain 3->2 goes 2->1(self-ref)->break. + Sample 0 is never visited during trace, so the fill loop must handle + sample 0's target=-1 without crashing. + """ + inst = _make_inst([-1, 1, 1, 2], [0, 800, 400, 100]) + ages = _run_age(inst, [[0, 0, 1, 1]]) + assert ages[0] > 0 + + +def test_age_multiple_sites(): + """Multiple sites produce one age per site, all positive.""" + positions = [1000, 2000, 3000] + inst = _make_inst([-1, 0, 1], [0, 500, 200], positions=positions) + genotypes = [[1, 1, 0], [0, 1, 1], [1, 0, 1]] + ages = _run_age(inst, genotypes) + assert len(ages) == 3 + assert all(a > 0 for a in ages) + + +def test_age_segment_change(): + """Sample changes target mid-sequence; ages still positive.""" + positions = [1000, 2000, 3000] + inst = _make_inst( + targets=[[-1], [0, 2], [0]], + tmrcas=[[0.0], [500.0, 200.0], [300.0]], + positions=positions, + starts=[[1000], [1000, 2000], [1000]], + ) + ages = _run_age(inst, [[1, 1, 0], [0, 1, 1], [1, 0, 1]]) + assert len(ages) == 3 + assert all(a > 0 for a in ages) + + +def test_age_deterministic_synthetic(): + """Same input twice gives identical output.""" + inst = _make_inst([-1, 0, 1, 2], [0, 1000, 500, 200]) + g = [0, 1, 1, 0] + a1 = _run_age(inst, [g])[0] + a2 = _run_age(inst, [g])[0] + assert a1 == a2 + + +def test_age_many_self_refs(): + """Multiple self-referencing samples in the same tree.""" + inst = _make_inst( + targets=[-1, 1, 2, 0, 3], + tmrcas=[0, 600, 400, 200, 100], + ) + ages = _run_age(inst, [[1, 0, 1, 1, 0]]) + assert ages[0] > 0 + + +# --------------------------------------------------------------------------- +# Unit tests: ConsistencyWrapper +# --------------------------------------------------------------------------- + +def test_dc_basic(): + """ConsistencyWrapper on a simple tree preserves dimensions.""" + positions = [1000, 2000, 3000] + inst = _make_inst([-1, 0, 1], [0, 500, 200], positions=positions) + genotypes = [[1, 1, 0], [0, 1, 1], [1, 0, 1]] + ages = _run_age(inst, genotypes) + dc = _run_dc(inst, ages, genotypes) + assert dc.num_samples == inst.num_samples + assert dc.num_sites == inst.num_sites + + +def test_dc_all_carriers(): + """DC with all-carrier sites.""" + positions = [1000, 2000] + inst = _make_inst([-1, 0, 1], [0, 500, 200], positions=positions) + genotypes = [[1, 1, 1], [1, 1, 1]] + ages = _run_age(inst, genotypes) + dc = _run_dc(inst, ages, genotypes) + assert dc.num_samples == 3 + assert dc.num_sites == 2 + + +def test_dc_no_carriers(): + """DC with no-carrier sites.""" + positions = [1000, 2000] + inst = _make_inst([-1, 0, 1], [0, 500, 200], positions=positions) + genotypes = [[0, 0, 0], [0, 0, 0]] + ages = _run_age(inst, genotypes) + dc = _run_dc(inst, ages, genotypes) + assert dc.num_samples == 3 + + +def test_dc_self_ref(): + """DC doesn't hang on self-referencing targets.""" + positions = [1000, 2000] + inst = _make_inst([-1, 1, 0], [0, 500, 300], positions=positions) + genotypes = [[1, 1, 0], [0, 1, 1]] + ages = _run_age(inst, genotypes) + dc = _run_dc(inst, ages, genotypes) + assert dc.num_samples == inst.num_samples + + +def test_dc_segment_change(): + """DC with mid-sequence target change.""" + positions = [1000, 2000, 3000] + inst = _make_inst( + targets=[[-1], [0, 2], [0]], + tmrcas=[[0.0], [500.0, 200.0], [300.0]], + positions=positions, + starts=[[1000], [1000, 2000], [1000]], + ) + genotypes = [[1, 1, 0], [0, 1, 1], [1, 0, 1]] + ages = _run_age(inst, genotypes) + dc = _run_dc(inst, ages, genotypes) + assert dc.num_samples == inst.num_samples + assert dc.num_sites == inst.num_sites + + +def test_dc_preserves_genotypes(): + """DC output must reconstruct the exact same genotype matrix.""" + positions = [1000, 2000, 3000, 4000, 5000] + inst = _make_inst([-1, 0, 1, 2], [0, 800, 400, 100], positions=positions) + + # Get the genotypes encoded by the original instructions + gi_orig = GenotypeIterator(inst) + orig_genos = [] + while gi_orig.has_next_genotype(): + orig_genos.append(list(gi_orig.next_genotype())) + + ages = _run_age(inst, orig_genos) + dc = _run_dc(inst, ages, orig_genos) + + gi_dc = GenotypeIterator(dc) + for i, orig_g in enumerate(orig_genos): + dc_g = list(gi_dc.next_genotype()) + assert orig_g == dc_g, f"Genotype mismatch at site {i}: {orig_g} != {dc_g}" + + +def test_dc_preserves_genotypes_with_self_ref(): + """DC preserves genotypes even when original has self-referencing targets.""" + import os + cache_path = os.path.join(os.path.dirname(__file__), "threads_cache", "sim_50dip.threads") + if not os.path.exists(cache_path): + pytest.skip("sim_50dip.threads not in cache") + + inst = load_instructions(cache_path) + + gi = GenotypeIterator(inst) + all_genos = [] + while gi.has_next_genotype(): + all_genos.append(list(gi.next_genotype())) + + ages = _run_age(inst, all_genos) + dc = _run_dc(inst, ages, all_genos) + + gi_dc = GenotypeIterator(dc) + for i, orig_g in enumerate(all_genos): + dc_g = list(gi_dc.next_genotype()) + assert orig_g == dc_g, f"Genotype mismatch at site {i}" + + +def test_dc_preserves_multiply(): + """DC right_multiply must produce identical results to original.""" + import os + cache_path = os.path.join(os.path.dirname(__file__), "threads_cache", "sim_50dip.threads") + if not os.path.exists(cache_path): + pytest.skip("sim_50dip.threads not in cache") + + inst = load_instructions(cache_path) + gi = GenotypeIterator(inst) + all_genos = [] + while gi.has_next_genotype(): + all_genos.append(list(gi.next_genotype())) + + ages = _run_age(inst, all_genos) + dc = _run_dc(inst, ages, all_genos) + + rng = np.random.default_rng(42) + x = rng.normal(0, 1, inst.num_sites).tolist() + inst.materialize_genotypes() + dc.materialize_genotypes() + np.testing.assert_allclose(inst.right_multiply(x), dc.right_multiply(x)) + + +def test_dc_output_genotypes_valid(): + """DC output produces valid genotypes (0 or 1, correct dimensions).""" + positions = [1000, 2000, 3000, 4000, 5000] + inst = _make_inst([-1, 0, 1, 2], [0, 800, 400, 100], positions=positions) + genotypes = [[0, 1, 1, 0], [1, 0, 0, 1], [1, 1, 1, 0], [0, 0, 1, 1], [1, 1, 0, 0]] + ages = _run_age(inst, genotypes) + dc = _run_dc(inst, ages, genotypes) + + gi_dc = GenotypeIterator(dc) + site_count = 0 + while gi_dc.has_next_genotype(): + g = gi_dc.next_genotype() + assert len(g) == dc.num_samples + assert all(v in (0, 1) for v in g) + site_count += 1 + assert site_count == dc.num_sites + + +def test_dc_ages_positive_on_output(): + """Re-estimating ages on DC output still gives positive ages.""" + positions = [1000, 2000, 3000, 4000, 5000] + inst = _make_inst([-1, 0, 1, 2], [0, 800, 400, 100], positions=positions) + genotypes = [[0, 1, 1, 0], [1, 0, 0, 1], [1, 1, 1, 0], [0, 0, 1, 1], [1, 1, 0, 0]] + ages = _run_age(inst, genotypes) + dc = _run_dc(inst, ages, genotypes) + + gi = GenotypeIterator(dc) + ae = AgeEstimator(dc) + while gi.has_next_genotype(): + ae.process_site(gi.next_genotype()) + dc_ages = ae.get_inferred_ages() + assert len(dc_ages) == dc.num_sites + assert all(a > 0 for a in dc_ages) + + +# --------------------------------------------------------------------------- +# Regression tests (real data) +# --------------------------------------------------------------------------- + +@pytest.fixture(scope="module") +def instructions(): + """Load the test threading instructions once for all tests.""" + return load_instructions(str(TEST_DATA_DIR / "expected_infer_snapshot.threads")) + + +@pytest.fixture(scope="module") +def reference_ages(instructions): + """Compute allele ages from the current implementation — the ground truth.""" + gt_it = GenotypeIterator(instructions) + age_est = AgeEstimator(instructions) + while gt_it.has_next_genotype(): + g = np.array(gt_it.next_genotype()) + age_est.process_site(g) + return np.array(age_est.get_inferred_ages()) + + +def test_allele_ages_count(instructions, reference_ages): + """One age per site.""" + assert len(reference_ages) == instructions.num_sites + + +def test_allele_ages_all_positive(reference_ages): + """All estimated ages should be positive.""" + assert np.all(reference_ages > 0) + + +def test_allele_ages_deterministic(instructions, reference_ages): + """Running twice produces identical results.""" + gt_it = GenotypeIterator(instructions) + age_est = AgeEstimator(instructions) + while gt_it.has_next_genotype(): + g = np.array(gt_it.next_genotype()) + age_est.process_site(g) + ages2 = np.array(age_est.get_inferred_ages()) + np.testing.assert_array_equal(reference_ages, ages2) + + +def test_allele_ages_snapshot_first_20(instructions, reference_ages): + """Pin the first 20 ages to exact values for regression detection.""" + expected_first_20 = reference_ages[:20].copy() + # Re-run + gt_it = GenotypeIterator(instructions) + age_est = AgeEstimator(instructions) + while gt_it.has_next_genotype(): + g = np.array(gt_it.next_genotype()) + age_est.process_site(g) + ages = np.array(age_est.get_inferred_ages()) + np.testing.assert_allclose(ages[:20], expected_first_20, rtol=1e-12) + + +def test_allele_ages_snapshot_statistics(reference_ages): + """Pin aggregate statistics for regression detection.""" + # These are from the current implementation on the test data + assert reference_ages.shape[0] == 8431 + np.testing.assert_allclose(np.mean(reference_ages), 7089.152796, rtol=1e-4) + np.testing.assert_allclose(np.min(reference_ages), 28.343561, rtol=1e-4) + + +def test_allele_ages_sub_range(instructions, reference_ages): + """Ages computed on a sub-range match the corresponding slice of full ages.""" + # Use positions 1000 to 2000 (by index) as sub-range + pos = instructions.positions + start_pos = pos[100] + end_pos = pos[200] + sub_inst = instructions.sub_range(start_pos, end_pos) + + gt_it = GenotypeIterator(sub_inst) + age_est = AgeEstimator(sub_inst) + while gt_it.has_next_genotype(): + g = np.array(gt_it.next_genotype()) + age_est.process_site(g) + sub_ages = np.array(age_est.get_inferred_ages()) + + # The sub-range should have the same number of sites + assert len(sub_ages) == sub_inst.num_sites + # Ages should all be positive + assert np.all(sub_ages > 0) + + +def test_allele_ages_self_referencing_targets(): + """AgeEstimator handles self-referencing targets (target[i] == i) without infinite loops.""" + import os + cache_path = os.path.join(os.path.dirname(__file__), "threads_cache", "sim_50dip.threads") + if not os.path.exists(cache_path): + pytest.skip("sim_50dip.threads not in cache") + + inst = load_instructions(cache_path) + + # Verify self-referencing targets exist in this data + targets = inst.all_targets() + has_self_ref = any( + any(tgt == i for tgt in segs) + for i, segs in enumerate(targets) if i > 0 + ) + assert has_self_ref, "Test data should contain self-referencing targets" + + # This should complete in under 5 seconds, not hang + import time + t0 = time.perf_counter() + ae = AgeEstimator(inst) + gi = GenotypeIterator(inst) + while gi.has_next_genotype(): + ae.process_site(gi.next_genotype()) + elapsed = time.perf_counter() - t0 + ages = ae.get_inferred_ages() + + assert len(ages) == inst.num_sites + assert all(a > 0 for a in ages) + assert elapsed < 5.0, f"AgeEstimator took {elapsed:.1f}s, likely stuck on self-ref targets" + + +def test_data_consistency_completes(): + """Full data consistency pipeline (age estimation + ConsistencyWrapper) completes.""" + import os + from threads_arg import ConsistencyWrapper + + cache_path = os.path.join(os.path.dirname(__file__), "threads_cache", "sim_50dip.threads") + if not os.path.exists(cache_path): + pytest.skip("sim_50dip.threads not in cache") + + inst = load_instructions(cache_path) + + # Step 1: estimate ages + ae = AgeEstimator(inst) + gi = GenotypeIterator(inst) + while gi.has_next_genotype(): + ae.process_site(gi.next_genotype()) + ages = ae.get_inferred_ages() + + # Step 2: consistency wrapper + cw = ConsistencyWrapper(inst, ages) + gi2 = GenotypeIterator(inst) + while gi2.has_next_genotype(): + cw.process_site(gi2.next_genotype()) + + dc_inst = cw.get_consistent_instructions() + assert dc_inst.num_samples == inst.num_samples + assert dc_inst.num_sites == inst.num_sites diff --git a/test/test_benchmark.cpp b/test/test_benchmark.cpp new file mode 100644 index 0000000..2d2779d --- /dev/null +++ b/test/test_benchmark.cpp @@ -0,0 +1,285 @@ +// This file is part of the Threads software suite. +// Copyright (C) 2024-2025 Threads Developers. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU General Public License for more details. +// +// You should have received a copy of the GNU General Public License +// along with this program. If not, see . +// +// Benchmarks for core algorithms. Run with: +// ./unit_tests "[benchmark]" --benchmark-samples 5 +// or without Catch2 benchmark (we do manual timing): +// ./unit_tests "[benchmark]" + +#include "Demography.hpp" +#include "HMM.hpp" +#include "Matcher.hpp" +#include "ViterbiLowMem.hpp" + +#include +#include +#include +#include +#include +#include +#include + +#ifdef __APPLE__ +#include +static size_t get_resident_memory_bytes() { + struct mach_task_basic_info info; + mach_msg_type_number_t count = MACH_TASK_BASIC_INFO_COUNT; + if (task_info(mach_task_self(), MACH_TASK_BASIC_INFO, (task_info_t)&info, &count) == KERN_SUCCESS) { + return info.resident_size; + } + return 0; +} +#elif defined(__linux__) +#include +#include +#include +static size_t get_resident_memory_bytes() { + std::ifstream f("/proc/self/statm"); + size_t pages; + f >> pages; // total + f >> pages; // resident + return pages * sysconf(_SC_PAGESIZE); +} +#else +static size_t get_resident_memory_bytes() { return 0; } +#endif + +namespace { + +struct BenchResult { + double elapsed_ms; + size_t mem_before; + size_t mem_after; + + double mem_delta_mb() const { + return static_cast(mem_after - mem_before) / (1024.0 * 1024.0); + } + + void print(const char* label) const { + std::cout << " [BENCH] " << label << ": " << std::fixed << std::setprecision(2) << elapsed_ms + << " ms, mem delta: " << std::setprecision(2) << mem_delta_mb() << " MB" << std::endl; + } +}; + +// Generate deterministic genotypes +std::vector> generate_genotypes(int n_samples, int n_sites, unsigned seed = 42) { + std::mt19937 rng(seed); + std::uniform_int_distribution dist(0, 1); + std::vector> genos(n_sites, std::vector(n_samples)); + for (int s = 0; s < n_sites; s++) { + for (int i = 0; i < n_samples; i++) { + genos[s][i] = dist(rng); + } + } + return genos; +} + +std::vector linear_positions(int n_sites, double start, double step) { + std::vector pos; + pos.reserve(n_sites); + for (int i = 0; i < n_sites; i++) { + pos.push_back(start + i * step); + } + return pos; +} + +} // namespace + +TEST_CASE("Benchmark: HMM construction", "[benchmark]") { + const int n_sites = 10000; + const int K = 64; + Demography demo({10000.0}, {0.0}); + auto bp = std::vector(n_sites, 100.0); + auto cm = std::vector(n_sites, 0.001); + + size_t mem_before = get_resident_memory_bytes(); + auto t0 = std::chrono::high_resolution_clock::now(); + + HMM hmm(demo, bp, cm, 1.4e-8, K); + + auto t1 = std::chrono::high_resolution_clock::now(); + size_t mem_after = get_resident_memory_bytes(); + + BenchResult r{ + std::chrono::duration(t1 - t0).count(), mem_before, mem_after}; + r.print("HMM construction (10K sites, K=64)"); + + // Sanity + CHECK(hmm.trellis.size() == n_sites); +} + +TEST_CASE("Benchmark: HMM breakpoints", "[benchmark]") { + const int n_sites = 5000; + const int K = 64; + Demography demo({10000.0}, {0.0}); + auto bp = std::vector(n_sites, 100.0); + auto cm = std::vector(n_sites, 0.001); + HMM hmm(demo, bp, cm, 1.4e-8, K); + + // Mixed observation pattern + std::mt19937 rng(123); + std::uniform_int_distribution dist(0, 4); + std::vector obs(n_sites); + for (int i = 0; i < n_sites; i++) { + obs[i] = dist(rng) == 0; // ~20% het rate + } + + auto t0 = std::chrono::high_resolution_clock::now(); + + auto bps = hmm.breakpoints(obs, 0); + + auto t1 = std::chrono::high_resolution_clock::now(); + BenchResult r{std::chrono::duration(t1 - t0).count(), 0, 0}; + r.print("HMM breakpoints (5K sites, K=64)"); + + CHECK(bps.size() >= 1); +} + +TEST_CASE("Benchmark: Matcher process_site", "[benchmark]") { + const int n_samples = 1000; + const int n_sites = 500; + auto positions = linear_positions(n_sites, 0.0, 0.02); + auto genos = generate_genotypes(n_samples, n_sites); + + size_t mem_before = get_resident_memory_bytes(); + auto t0 = std::chrono::high_resolution_clock::now(); + + Matcher m(n_samples, positions, 0.01, 0.5, 4, 2); + for (int s = 0; s < n_sites; s++) { + m.process_site(genos[s]); + } + + auto t1 = std::chrono::high_resolution_clock::now(); + size_t mem_after = get_resident_memory_bytes(); + + BenchResult r{ + std::chrono::duration(t1 - t0).count(), mem_before, mem_after}; + r.print("Matcher process_site (1K samples, 500 sites)"); + + CHECK(m.get_sorting().size() == n_samples); +} + +TEST_CASE("Benchmark: ViterbiState process_site", "[benchmark]") { + const int n_ref = 100; + const int target_id = n_ref; + const int n_sites = 2000; + const int n_samples_total = n_ref + 1; + + std::vector ref_samples; + for (int i = 0; i < n_ref; i++) { + ref_samples.push_back(i); + } + + auto genos = generate_genotypes(n_samples_total, n_sites); + + size_t mem_before = get_resident_memory_bytes(); + auto t0 = std::chrono::high_resolution_clock::now(); + + ViterbiState state(target_id, ref_samples); + for (int s = 0; s < n_sites; s++) { + state.process_site(genos[s], 3.0, 0.01, 2.0, 0.01); + } + + auto t1 = std::chrono::high_resolution_clock::now(); + size_t mem_after = get_resident_memory_bytes(); + + BenchResult r{ + std::chrono::duration(t1 - t0).count(), mem_before, mem_after}; + r.print("ViterbiState process_site (100 refs, 2K sites)"); + + CHECK(state.sites_processed == n_sites); +} + +TEST_CASE("Benchmark: ViterbiState prune", "[benchmark]") { + const int n_ref = 50; + const int target_id = n_ref; + const int n_sites = 1000; + const int n_samples_total = n_ref + 1; + + std::vector ref_samples; + for (int i = 0; i < n_ref; i++) { + ref_samples.push_back(i); + } + + auto genos = generate_genotypes(n_samples_total, n_sites); + + ViterbiState state(target_id, ref_samples); + for (int s = 0; s < n_sites; s++) { + state.process_site(genos[s], 3.0, 0.01, 2.0, 0.01); + } + + int branches_before = state.count_branches(); + + auto t0 = std::chrono::high_resolution_clock::now(); + state.prune(); + auto t1 = std::chrono::high_resolution_clock::now(); + + int branches_after = state.count_branches(); + + BenchResult r{std::chrono::duration(t1 - t0).count(), 0, 0}; + r.print("ViterbiState prune (50 refs, 1K sites)"); + + std::cout << " Branches: " << branches_before << " -> " << branches_after << std::endl; + CHECK(branches_after <= branches_before); +} + +TEST_CASE("Benchmark: ViterbiState traceback", "[benchmark]") { + const int n_ref = 50; + const int target_id = n_ref; + const int n_sites = 1000; + const int n_samples_total = n_ref + 1; + + std::vector ref_samples; + for (int i = 0; i < n_ref; i++) { + ref_samples.push_back(i); + } + + auto genos = generate_genotypes(n_samples_total, n_sites); + + ViterbiState state(target_id, ref_samples); + for (int s = 0; s < n_sites; s++) { + state.process_site(genos[s], 3.0, 0.01, 2.0, 0.01); + } + + auto t0 = std::chrono::high_resolution_clock::now(); + auto path = state.traceback(); + auto t1 = std::chrono::high_resolution_clock::now(); + + BenchResult r{std::chrono::duration(t1 - t0).count(), 0, 0}; + r.print("ViterbiState traceback (50 refs, 1K sites)"); + + CHECK(path.size() >= 1); + std::cout << " Path segments: " << path.size() << std::endl; +} + +TEST_CASE("Benchmark: Demography std_to_gen (many calls)", "[benchmark]") { + Demography d({5000.0, 10000.0, 20000.0}, {0.0, 100.0, 500.0}); + + const int n_calls = 1000000; + double sum = 0.0; + + auto t0 = std::chrono::high_resolution_clock::now(); + for (int i = 0; i < n_calls; i++) { + double t = static_cast(i) / n_calls * 5.0; + sum += d.std_to_gen(t); + } + auto t1 = std::chrono::high_resolution_clock::now(); + + BenchResult r{std::chrono::duration(t1 - t0).count(), 0, 0}; + r.print("Demography std_to_gen (1M calls)"); + + CHECK(sum > 0.0); // prevent optimization +} diff --git a/test/test_convert.py b/test/test_convert.py new file mode 100644 index 0000000..608600b --- /dev/null +++ b/test/test_convert.py @@ -0,0 +1,134 @@ +""" +Tests for the `threads convert` command: threads_to_arg, ARG structure validation, +and serialization to .argn format. +""" +import tempfile + +import numpy as np +import pytest +import arg_needle_lib + +from pathlib import Path + +TEST_DATA = Path(__file__).parent / "data" +THREADS_FIT = str(TEST_DATA / "expected_infer_fit_to_data_snapshot.threads") +THREADS_NO_FIT = str(TEST_DATA / "expected_infer_snapshot.threads") +PANEL_PGEN = str(TEST_DATA / "panel.pgen") + + +@pytest.fixture(scope="module") +def instructions_nofit(): + from threads_arg.serialization import load_instructions + return load_instructions(THREADS_NO_FIT) + + +@pytest.fixture(scope="module") +def instructions_fit(): + from threads_arg.serialization import load_instructions + return load_instructions(THREADS_FIT) + + +def _build_arg(instructions, add_mutations=False): + """Build ARG with retry logic matching threads_convert.""" + from threads_arg.convert import threads_to_arg + for noise in [0.0, 1e-5, 1e-3]: + try: + return threads_to_arg(instructions, add_mutations=add_mutations, noise=noise) + except RuntimeError: + continue + raise RuntimeError("Failed to build ARG even with noise=1e-3") + + +@pytest.fixture(scope="module") +def arg_nofit(instructions_nofit): + return _build_arg(instructions_nofit, add_mutations=False) + + +# =================================================================== +# threads_to_arg: ARG construction +# =================================================================== +class TestThreadsToArg: + def test_arg_has_correct_sample_count(self, arg_nofit, instructions_nofit): + leaf_ids = arg_nofit.leaf_ids + assert len(leaf_ids) == instructions_nofit.num_samples + + def test_arg_leaves_are_leaves(self, arg_nofit): + for lid in arg_nofit.leaf_ids: + assert arg_nofit.is_leaf(lid) + + def test_arg_offset_matches_instructions(self, arg_nofit, instructions_nofit): + assert arg_nofit.offset == instructions_nofit.start + + def test_arg_construction_with_noise(self, instructions_nofit): + """Adding noise should not crash and should produce a valid ARG.""" + from threads_arg.convert import threads_to_arg + arg = threads_to_arg(instructions_nofit, add_mutations=False, noise=1e-5) + assert len(arg.leaf_ids) == instructions_nofit.num_samples + + def test_arg_with_mutations(self, instructions_nofit): + """add_mutations=True should populate mutations on the ARG.""" + arg = _build_arg(instructions_nofit, add_mutations=True) + arg.populate_children_and_roots() + mutations = arg.mutations() + # Parsimonious mapping may place multiple mutations per site + assert len(mutations) >= instructions_nofit.num_sites + + +# =================================================================== +# threads_convert: full pipeline .threads -> .argn +# =================================================================== +class TestThreadsConvert: + def test_produces_argn_file(self): + from threads_arg.convert import threads_convert + with tempfile.TemporaryDirectory() as tmpdir: + out_argn = str(Path(tmpdir) / "test.argn") + threads_convert(THREADS_NO_FIT, argn=out_argn, tsz=None) + assert Path(out_argn).exists() + assert Path(out_argn).stat().st_size > 0 + + def test_argn_deserializable(self): + from threads_arg.convert import threads_convert + with tempfile.TemporaryDirectory() as tmpdir: + out_argn = str(Path(tmpdir) / "test.argn") + threads_convert(THREADS_NO_FIT, argn=out_argn, tsz=None) + arg = arg_needle_lib.deserialize_arg(out_argn) + assert len(arg.leaf_ids) == 1000 + + def test_argn_snapshot_match(self): + """Generated .argn should match the expected snapshot.""" + from threads_arg.convert import threads_convert + import h5py + expected_argn = TEST_DATA / "expected_convert_snapshot.argn" + with tempfile.TemporaryDirectory() as tmpdir: + out_argn = str(Path(tmpdir) / "test.argn") + threads_convert(THREADS_NO_FIT, argn=out_argn, tsz=None) + + # Compare HDF5 structures + with h5py.File(out_argn, "r") as gen, h5py.File(str(expected_argn), "r") as exp: + assert set(gen.keys()) == set(exp.keys()) + for key in exp.keys(): + if isinstance(exp[key], h5py.Dataset): + assert gen[key].shape == exp[key].shape, \ + f"shape mismatch for {key}: {gen[key].shape} vs {exp[key].shape}" + + def test_fit_to_data_converts(self): + """fit_to_data .threads should also convert without error.""" + from threads_arg.convert import threads_convert + with tempfile.TemporaryDirectory() as tmpdir: + out_argn = str(Path(tmpdir) / "test_fit.argn") + threads_convert(THREADS_FIT, argn=out_argn, tsz=None) + assert Path(out_argn).exists() + + +# =================================================================== +# ARG structure properties +# =================================================================== +class TestArgStructure: + def test_genotype_reconstruction_from_argn(self): + """Genotypes from the pre-built .argn snapshot (without add_mutations) + can be deserialized and have the correct sample count.""" + expected_argn = str(TEST_DATA / "expected_convert_snapshot.argn") + arg = arg_needle_lib.deserialize_arg(expected_argn) + assert len(arg.leaf_ids) == 1000 + for lid in arg.leaf_ids: + assert arg.is_leaf(lid) diff --git a/test/test_demography_correctness.cpp b/test/test_demography_correctness.cpp new file mode 100644 index 0000000..d30f527 --- /dev/null +++ b/test/test_demography_correctness.cpp @@ -0,0 +1,83 @@ +// This file is part of the Threads software suite. +// Copyright (C) 2024-2025 Threads Developers. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU General Public License for more details. +// +// You should have received a copy of the GNU General Public License +// along with this program. If not, see . + +#include "Demography.hpp" + +#include +#include +#include + +TEST_CASE("Demography constant Ne") { + double Ne = 10000.0; + Demography d({Ne}, {0.0}); + + // With constant Ne, std_to_gen should be linear: gen = t * Ne + CHECK_THAT(d.std_to_gen(0.0), Catch::Matchers::WithinAbs(0.0, 1e-10)); + CHECK_THAT(d.std_to_gen(1.0), Catch::Matchers::WithinAbs(Ne, 1e-6)); + CHECK_THAT(d.std_to_gen(0.5), Catch::Matchers::WithinAbs(Ne * 0.5, 1e-6)); + CHECK_THAT(d.std_to_gen(2.0), Catch::Matchers::WithinAbs(Ne * 2.0, 1e-6)); +} + +TEST_CASE("Demography piecewise Ne") { + // Ne=10000 for generations [0, 100), then Ne=20000 after + Demography d({10000.0, 20000.0}, {0.0, 100.0}); + + // std_times: [0, 100/10000] = [0, 0.01] + CHECK_THAT(d.std_times[0], Catch::Matchers::WithinAbs(0.0, 1e-12)); + CHECK_THAT(d.std_times[1], Catch::Matchers::WithinAbs(0.01, 1e-12)); + + // Within first epoch: std_to_gen(0.005) = 0 + 0.005 * 10000 = 50 + CHECK_THAT(d.std_to_gen(0.005), Catch::Matchers::WithinAbs(50.0, 1e-6)); + + // Within second epoch: std_to_gen(0.02) = 100 + (0.02 - 0.01) * 20000 = 300 + CHECK_THAT(d.std_to_gen(0.02), Catch::Matchers::WithinAbs(300.0, 1e-6)); +} + +TEST_CASE("Demography expected branch length") { + double Ne = 10000.0; + Demography d({Ne}, {0.0}); + + // expected_branch_length(N) = std_to_gen(2/N) + // For constant Ne: 2/N * Ne + CHECK_THAT(d.expected_branch_length(2), Catch::Matchers::WithinAbs(Ne, 1e-6)); + CHECK_THAT(d.expected_branch_length(10), Catch::Matchers::WithinAbs(2000.0, 1e-6)); + CHECK_THAT(d.expected_branch_length(100), Catch::Matchers::WithinAbs(200.0, 1e-6)); +} + +TEST_CASE("Demography expected_time is std_to_gen(1)") { + Demography d({5000.0}, {0.0}); + CHECK_THAT(d.expected_time, Catch::Matchers::WithinAbs(5000.0, 1e-6)); +} + +TEST_CASE("Demography three epochs") { + // Ne=1000 for [0,50), Ne=5000 for [50,100), Ne=20000 after 100 + Demography d({1000.0, 5000.0, 20000.0}, {0.0, 50.0, 100.0}); + + // std_times: [0, 50/1000, 50/1000 + 50/5000] = [0, 0.05, 0.06] + CHECK_THAT(d.std_times[0], Catch::Matchers::WithinAbs(0.0, 1e-12)); + CHECK_THAT(d.std_times[1], Catch::Matchers::WithinAbs(0.05, 1e-12)); + CHECK_THAT(d.std_times[2], Catch::Matchers::WithinAbs(0.06, 1e-12)); + + // In third epoch: std_to_gen(0.07) = 100 + (0.07 - 0.06) * 20000 = 300 + CHECK_THAT(d.std_to_gen(0.07), Catch::Matchers::WithinAbs(300.0, 1e-6)); +} + +TEST_CASE("Demography stream output") { + Demography d({1000.0}, {0.0}); + std::ostringstream oss; + oss << d; + // Should not crash, but note the bug: operator<< uses std::cout instead of os +} diff --git a/test/test_hmm.cpp b/test/test_hmm.cpp new file mode 100644 index 0000000..801c8a3 --- /dev/null +++ b/test/test_hmm.cpp @@ -0,0 +1,168 @@ +// This file is part of the Threads software suite. +// Copyright (C) 2024-2025 Threads Developers. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU General Public License for more details. +// +// You should have received a copy of the GNU General Public License +// along with this program. If not, see . + +#include "HMM.hpp" + +#include +#include +#include +#include + +namespace { + +// Create a simple constant-Ne demography for testing +Demography simple_demography(double ne = 10000.0) { + return Demography({ne}, {0.0}); +} + +// Create uniform site sizes for n_sites sites +std::vector uniform_bp_sizes(int n_sites, double bp_size = 100.0) { + return std::vector(n_sites, bp_size); +} + +std::vector uniform_cm_sizes(int n_sites, double cm_size = 0.001) { + return std::vector(n_sites, cm_size); +} + +} // namespace + +TEST_CASE("HMM construction") { + int n_sites = 20; + int K = 8; + auto demo = simple_demography(); + auto bp = uniform_bp_sizes(n_sites); + auto cm = uniform_cm_sizes(n_sites); + + HMM hmm(demo, bp, cm, 1.4e-8, K); + + CHECK(hmm.num_states == K); + CHECK(hmm.expected_times.size() == K); + CHECK(hmm.trellis.size() == n_sites); + CHECK(hmm.pointers.size() == n_sites); + CHECK(hmm.non_transition_score.size() == n_sites); + CHECK(hmm.transition_score.size() == n_sites); + CHECK(hmm.hom_score.size() == n_sites); + CHECK(hmm.het_score.size() == n_sites); +} + +TEST_CASE("HMM expected times are increasing") { + auto demo = simple_demography(); + HMM hmm(demo, uniform_bp_sizes(10), uniform_cm_sizes(10), 1.4e-8, 16); + + for (int i = 1; i < 16; i++) { + CHECK(hmm.expected_times[i] > hmm.expected_times[i - 1]); + } + // All expected times should be positive + for (int i = 0; i < 16; i++) { + CHECK(hmm.expected_times[i] > 0.0); + } +} + +TEST_CASE("HMM breakpoints with all homozygous") { + int n_sites = 20; + int K = 4; + auto demo = simple_demography(); + HMM hmm(demo, uniform_bp_sizes(n_sites), uniform_cm_sizes(n_sites), 1.4e-8, K); + + // All homozygous = no mutations -> should stay in one state + std::vector obs(n_sites, false); + auto bps = hmm.breakpoints(obs, 0); + + // Should have at least the initial breakpoint at 0 + CHECK(bps.size() >= 1); + CHECK(bps[0] == 0); +} + +TEST_CASE("HMM breakpoints with all heterozygous") { + int n_sites = 30; + int K = 4; + auto demo = simple_demography(); + HMM hmm(demo, uniform_bp_sizes(n_sites), uniform_cm_sizes(n_sites), 1.4e-8, K); + + // All het -> lots of mutations -> should stay in deepest time state + std::vector obs(n_sites, true); + auto bps = hmm.breakpoints(obs, 0); + + CHECK(bps.size() >= 1); + CHECK(bps[0] == 0); +} + +TEST_CASE("HMM breakpoints with mixed signal") { + int n_sites = 40; + int K = 8; + auto demo = simple_demography(); + HMM hmm(demo, uniform_bp_sizes(n_sites), uniform_cm_sizes(n_sites), 1.4e-8, K); + + // First half: all hom (recent), second half: all het (old) -> expect breakpoint + std::vector obs(n_sites, false); + for (int i = n_sites / 2; i < n_sites; i++) { + obs[i] = true; + } + + auto bps = hmm.breakpoints(obs, 0); + CHECK(bps.size() >= 1); + CHECK(bps[0] == 0); + // With a strong signal change, we expect at least one additional breakpoint + // (though exact number depends on HMM parameters) +} + +TEST_CASE("HMM breakpoints with offset start") { + int n_sites = 30; + int K = 4; + auto demo = simple_demography(); + HMM hmm(demo, uniform_bp_sizes(n_sites), uniform_cm_sizes(n_sites), 1.4e-8, K); + + // Use only a sub-range starting at offset 5 + int start = 5; + int len = 15; + std::vector obs(len, false); + + auto bps = hmm.breakpoints(obs, start); + CHECK(bps[0] == start); +} + +TEST_CASE("HMM recombination scores are negative log-probs") { + int n_sites = 5; + int K = 4; + auto demo = simple_demography(); + HMM hmm(demo, uniform_bp_sizes(n_sites), uniform_cm_sizes(n_sites), 1.4e-8, K); + + for (int i = 0; i < n_sites; i++) { + for (int k = 0; k < K; k++) { + // Both transition and non-transition scores should be <= 0 (log-probs) + CHECK(hmm.transition_score[i][k] <= 0.0); + CHECK(hmm.non_transition_score[i][k] <= 0.0); + // non-transition should be >= transition (more likely to not transition) + CHECK(hmm.non_transition_score[i][k] >= hmm.transition_score[i][k]); + } + } +} + +TEST_CASE("HMM mutation scores are negative log-probs") { + int n_sites = 5; + int K = 4; + auto demo = simple_demography(); + HMM hmm(demo, uniform_bp_sizes(n_sites), uniform_cm_sizes(n_sites), 1.4e-8, K); + + for (int i = 0; i < n_sites; i++) { + for (int k = 0; k < K; k++) { + CHECK(hmm.hom_score[i][k] <= 0.0); + CHECK(hmm.het_score[i][k] <= 0.0); + // hom (not mutating) should be more likely than het (mutating) for typical params + CHECK(hmm.hom_score[i][k] >= hmm.het_score[i][k]); + } + } +} diff --git a/test/test_impute_correctness.py b/test/test_impute_correctness.py new file mode 100644 index 0000000..e761e90 --- /dev/null +++ b/test/test_impute_correctness.py @@ -0,0 +1,536 @@ +""" +Property-based unit tests for the imputation pipeline. + +Tests correctness properties of individual components rather than +relying solely on snapshot regression. +""" +import io +import numpy as np +import pytest +import types + +from scipy.sparse import csr_array + +from threads_arg.fwbw import ( + forwards_ls_hap, + backwards_ls_hap, + set_emission_probabilities, + fwbw, + MISSING, +) +from threads_arg.impute import ( + _active_site_arg_delta, + MutationMap, + MutationContainer, + CachedPosteriorSnps, + WriterVCF, + Impute, +) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- +@pytest.fixture +def small_panel(): + """5 haplotypes, 10 biallelic sites.""" + rng = np.random.default_rng(42) + return rng.integers(0, 2, size=(10, 5)).astype(np.int8) + + +@pytest.fixture +def small_query(small_panel): + """Single query haplotype, shape (1, m).""" + rng = np.random.default_rng(99) + m = small_panel.shape[0] + return rng.integers(0, 2, size=(1, m)).astype(np.int8) + + +@pytest.fixture +def recomb_rates(small_panel): + m = small_panel.shape[0] + return np.full(m, 0.01) + + +@pytest.fixture +def mutation_rate(): + return 0.01 + + +# --------------------------------------------------------------------------- +# set_emission_probabilities +# --------------------------------------------------------------------------- +class TestSetEmissionProbabilities: + def test_output_shape(self, small_panel, small_query, mutation_rate): + e = set_emission_probabilities(small_panel, small_query, mutation_rate) + assert e.shape == (small_panel.shape[0], 2) + + def test_emissions_in_unit_interval(self, small_panel, small_query, mutation_rate): + e = set_emission_probabilities(small_panel, small_query, mutation_rate) + assert np.all(e >= 0) + assert np.all(e <= 1) + + def test_match_mismatch_sum_to_one_polymorphic(self): + """For polymorphic sites, e[:,0] + e[:,1] == 1.""" + panel = np.array([[0, 1, 0], [1, 0, 1]], dtype=np.int8) + query = np.array([[1, 0]], dtype=np.int8) + e = set_emission_probabilities(panel, query, 0.05) + np.testing.assert_allclose(e[:, 0] + e[:, 1], 1.0) + + def test_invariant_site_emissions(self): + """Invariant sites: mismatch=0, match=1 regardless of mutation rate.""" + panel = np.array([[0, 0, 0], [1, 0, 1]], dtype=np.int8) + query = np.array([[0, 0]], dtype=np.int8) # site 0 all-zero => invariant + e = set_emission_probabilities(panel, query, 0.05) + assert e[0, 0] == 0.0 + assert e[0, 1] == 1.0 + # site 1 polymorphic + assert e[1, 0] == pytest.approx(0.05) + assert e[1, 1] == pytest.approx(0.95) + + def test_mutation_rate_none(self, small_panel, small_query): + """Auto-computed mutation rate should still produce valid emissions.""" + e = set_emission_probabilities(small_panel, small_query, None) + assert np.all(e >= 0) + assert np.all(e <= 1) + + +# --------------------------------------------------------------------------- +# Forward-backward algorithm +# --------------------------------------------------------------------------- +class TestForwardBackward: + def test_forward_scaling_factors_positive(self, small_panel, small_query, + recomb_rates, mutation_rate): + e = set_emission_probabilities(small_panel, small_query, mutation_rate) + F, c = forwards_ls_hap(small_panel.astype(np.float64), small_query.ravel().astype(np.float64), e, recomb_rates) + assert np.all(c > 0) + + def test_forward_values_non_negative(self, small_panel, small_query, + recomb_rates, mutation_rate): + e = set_emission_probabilities(small_panel, small_query, mutation_rate) + F, c = forwards_ls_hap(small_panel.astype(np.float64), small_query.ravel().astype(np.float64), e, recomb_rates) + assert np.all(F >= 0) + + def test_forward_normalized_rows_sum_to_one(self, small_panel, small_query, + recomb_rates, mutation_rate): + e = set_emission_probabilities(small_panel, small_query, mutation_rate) + F, c = forwards_ls_hap(small_panel.astype(np.float64), small_query.ravel().astype(np.float64), e, recomb_rates) + np.testing.assert_allclose(F.sum(axis=1), 1.0, atol=1e-10) + + def test_backward_last_row_all_ones(self, small_panel, small_query, + recomb_rates, mutation_rate): + m, n = small_panel.shape + e = set_emission_probabilities(small_panel, small_query, mutation_rate) + F, c = forwards_ls_hap(small_panel.astype(np.float64), small_query.ravel().astype(np.float64), e, recomb_rates) + B = backwards_ls_hap(small_panel.astype(np.float64), small_query.ravel().astype(np.float64), e, c, recomb_rates) + np.testing.assert_array_equal(B[-1], np.ones(n)) + + def test_backward_values_non_negative(self, small_panel, small_query, + recomb_rates, mutation_rate): + e = set_emission_probabilities(small_panel, small_query, mutation_rate) + F, c = forwards_ls_hap(small_panel.astype(np.float64), small_query.ravel().astype(np.float64), e, recomb_rates) + B = backwards_ls_hap(small_panel.astype(np.float64), small_query.ravel().astype(np.float64), e, c, recomb_rates) + assert np.all(B >= 0) + + def test_posterior_shape(self, small_panel, small_query, recomb_rates, mutation_rate): + posterior = fwbw(small_panel, small_query, recomb_rates, mutation_rate) + assert posterior.shape == small_panel.shape + + def test_posterior_non_negative(self, small_panel, small_query, + recomb_rates, mutation_rate): + posterior = fwbw(small_panel, small_query, recomb_rates, mutation_rate) + assert np.all(posterior >= -1e-15) + + def test_posterior_rows_sum_to_one(self, small_panel, small_query, + recomb_rates, mutation_rate): + posterior = fwbw(small_panel, small_query, recomb_rates, mutation_rate) + np.testing.assert_allclose(posterior.sum(axis=1), 1.0, atol=1e-10) + + def test_identical_query_concentrates_posterior(self): + """Posterior should concentrate on haplotypes identical to query.""" + panel = np.array([[0, 1, 0, 1], + [1, 0, 1, 0], + [0, 1, 0, 1], + [1, 0, 1, 0]], dtype=np.int8) + # h0 == h2 == [0,1,0,1]; h1 == h3 == [1,0,1,0] + query = panel[:, 0:1].T # (1, m): matches h0 and h2 + recomb = np.full(4, 0.001) + posterior = fwbw(panel, query, recomb, 0.001) + match_weight = posterior[:, [0, 2]].sum(axis=1) + other_weight = posterior[:, [1, 3]].sum(axis=1) + assert np.all(match_weight > other_weight) + + def test_missing_data_still_valid(self): + """MISSING sites should not break posterior properties.""" + panel = np.array([[0, 1, 0], [1, 0, 1], [0, 1, 0]], dtype=np.int8) + query = np.array([[MISSING, 1, 0]], dtype=np.int8) + recomb = np.full(3, 0.01) + posterior = fwbw(panel, query, recomb, 0.01) + assert np.all(posterior >= -1e-15) + np.testing.assert_allclose(posterior.sum(axis=1), 1.0, atol=1e-10) + + def test_different_panel_sizes(self): + """Properties hold for various n and m.""" + for m, n in [(3, 2), (5, 10), (20, 3), (50, 50)]: + rng = np.random.default_rng(m * 100 + n) + panel = rng.integers(0, 2, size=(m, n)).astype(np.int8) + query = rng.integers(0, 2, size=(1, m)).astype(np.int8) + recomb = np.full(m, 0.01) + posterior = fwbw(panel, query, recomb, 0.01) + assert posterior.shape == (m, n) + assert np.all(posterior >= -1e-15) + np.testing.assert_allclose(posterior.sum(axis=1), 1.0, atol=1e-10) + + +# --------------------------------------------------------------------------- +# _sparsify_posterior +# --------------------------------------------------------------------------- +class TestSparsifyPosterior: + def _mock(self, num_snps, num_samples_panel): + return types.SimpleNamespace( + num_snps=num_snps, + num_samples_panel=num_samples_panel, + ) + + def test_output_shape(self): + mock = self._mock(10, 100) + posterior = np.random.default_rng(0).random((10, 5)) + 0.01 + matched = np.array([3, 17, 42, 55, 88]) + result = Impute._sparsify_posterior(mock, posterior, matched) + assert result.shape == (10, 100) + + def test_rows_sum_to_one(self): + mock = self._mock(10, 100) + posterior = np.random.default_rng(0).random((10, 5)) + 0.01 + matched = np.array([3, 17, 42, 55, 88]) + result = Impute._sparsify_posterior(mock, posterior, matched) + row_sums = np.asarray(result.sum(axis=1)).ravel() + np.testing.assert_allclose(row_sums, 1.0, atol=1e-10) + + def test_no_negative_values(self): + mock = self._mock(10, 50) + posterior = np.random.default_rng(0).random((10, 5)) + 0.01 + matched = np.arange(5) + result = Impute._sparsify_posterior(mock, posterior, matched) + assert np.all(result.data >= 0) + + def test_columns_in_matched_set(self): + mock = self._mock(5, 100) + posterior = np.ones((5, 3)) * 0.5 + matched = np.array([10, 50, 90]) + result = Impute._sparsify_posterior(mock, posterior, matched) + nz_cols = set(result.nonzero()[1]) + assert nz_cols.issubset({10, 50, 90}) + + def test_small_values_thresholded(self): + """Entries <= 1/num_samples_panel should be dropped.""" + mock = self._mock(3, 10) + # Row 0: one value well above threshold, rest tiny + posterior = np.array([[1.0, 0.001, 0.001], + [0.5, 0.4, 0.1], + [0.9, 0.05, 0.05]]) + matched = np.array([0, 1, 2]) + result = Impute._sparsify_posterior(mock, posterior, matched) + # threshold = 1/10 = 0.1; row 0 col 1,2 (0.001) should be gone + row0 = np.asarray(result[0].todense()).ravel() + assert row0[1] == 0.0 + assert row0[2] == 0.0 + assert row0[0] > 0.0 + + +# --------------------------------------------------------------------------- +# MutationMap +# --------------------------------------------------------------------------- +class TestMutationMap: + def test_unmapped(self): + mm = MutationMap("snp1", 0, "NaN") + assert not mm.is_mapped() + assert not mm.flipped + + def test_single_uniquely_mapped(self): + mm = MutationMap("snp1", 0, "-1,100.0,200.0") + assert mm.is_mapped() + assert mm.uniquely_mapped + assert mm.get_boundaries(-1) == (100.0, 200.0) + # Uniquely mapped: any query returns the -1 boundaries + assert mm.get_boundaries(42) == (100.0, 200.0) + + def test_multiple_mappings(self): + mm = MutationMap("snp1", 0, "5,100.0,200.0;10,300.0,400.0") + assert mm.is_mapped() + assert not mm.uniquely_mapped + assert mm.get_boundaries(5) == (100.0, 200.0) + assert mm.get_boundaries(10) == (300.0, 400.0) + + def test_dotted_sample_ids(self): + """'5.10.15' expands to three separate carrier IDs.""" + mm = MutationMap("snp1", 0, "5.10.15,100.0,200.0") + assert mm.is_carrier(5) + assert mm.is_carrier(10) + assert mm.is_carrier(15) + assert not mm.is_carrier(20) + + def test_flipped_flag(self): + mm = MutationMap("snp1", 1, "-1,100.0,200.0") + assert mm.flipped + + def test_multi_group_carriers(self): + mm = MutationMap("snp1", 0, "1.2,10.0,20.0;3,30.0,40.0") + assert mm.is_carrier(1) + assert mm.is_carrier(2) + assert mm.is_carrier(3) + assert mm.get_boundaries(1) == (10.0, 20.0) + assert mm.get_boundaries(3) == (30.0, 40.0) + + +# --------------------------------------------------------------------------- +# MutationContainer +# --------------------------------------------------------------------------- +class TestMutationContainer: + def test_load_and_lookup(self, tmp_path): + mut_file = tmp_path / "test.mut" + mut_file.write_text( + "'snp1'\t100\t0\t-1,50.0,150.0\n" + "'snp2'\t200\t0\tNaN\n" + "'snp3'\t300\t1\t5.10,100.0,200.0;15,300.0,400.0\n" + ) + mc = MutationContainer(str(mut_file)) + assert mc.is_mapped("'snp1'") + assert not mc.is_mapped("'snp2'") + assert mc.is_mapped("'snp3'") + assert not mc.is_mapped("nonexistent") + + def test_get_mapping(self, tmp_path): + mut_file = tmp_path / "test.mut" + mut_file.write_text("'snp1'\t100\t0\t-1,50.0,150.0\n") + mc = MutationContainer(str(mut_file)) + mm = mc.get_mapping("'snp1'") + assert mm.uniquely_mapped + assert mm.get_boundaries(-1) == (50.0, 150.0) + + +# --------------------------------------------------------------------------- +# _active_site_arg_delta +# --------------------------------------------------------------------------- +class TestActiveSiteArgDelta: + def _seg(self, seg_start, ids, ages): + return types.SimpleNamespace(seg_start=seg_start, ids=ids, ages=ages) + + def _rec(self, pos): + return types.SimpleNamespace(pos=pos) + + def test_zero_delta_no_carriers_in_active(self): + posterior = np.array([0.5, 0.3, 0.2]) + active_indexes = {0: 0, 1: 1, 2: 2} + thread = [self._seg(0, [0, 1, 2], [100.0, 200.0, 300.0])] + mm = MutationMap("snp", 0, "99,50.0,150.0") # carrier 99 not active + carriers = {99} + delta = _active_site_arg_delta( + posterior, active_indexes, thread, mm, carriers, self._rec(50)) + assert delta == 0.0 + + def test_zero_delta_zero_posterior(self): + posterior = np.array([0.0, 0.5, 0.5]) + active_indexes = {0: 0, 1: 1, 2: 2} + thread = [self._seg(0, [0, 1, 2], [100.0, 200.0, 300.0])] + mm = MutationMap("snp", 0, "0,50.0,150.0") # carrier 0, posterior=0 + carriers = {0} + delta = _active_site_arg_delta( + posterior, active_indexes, thread, mm, carriers, self._rec(50)) + assert delta == 0.0 + + def test_delta_sign_is_negative(self): + """arg_prob < 1 => (arg_prob - 1) < 0 => delta < 0 when posterior > 0.""" + posterior = np.array([0.8, 0.1, 0.1]) + active_indexes = {0: 0} + thread = [self._seg(0, [0], [1000.0])] + # Use -1 sentinel for uniquely mapped (matches real .mut format) + mm = MutationMap("snp", 0, "-1,50.0,150.0") + carriers = {0} + delta = _active_site_arg_delta( + posterior, active_indexes, thread, mm, carriers, self._rec(50)) + assert delta < 0 + + def test_erlang2_probability_in_unit_interval(self): + """The Erlang(2) CDF must be in [0, 1] for all positive inputs.""" + for height in [10.0, 100.0, 1000.0, 10000.0]: + for mut_h in [1.0, 50.0, 500.0]: + lam = 2.0 / height + lam_mut = lam * mut_h + arg_prob = 1 - np.exp(-lam_mut) * (1 + lam_mut) + assert 0 <= arg_prob <= 1, f"height={height}, mut_h={mut_h}" + + def test_segment_selection_by_position(self): + """Position determines which segment's IDs are used.""" + posterior = np.array([0.5, 0.5]) + active_indexes = {0: 0, 1: 1} + seg1 = self._seg(0, [0], [100.0]) + seg2 = self._seg(500, [1], [100.0]) + thread = [seg1, seg2] + # Non-uniquely mapped: both sample 0 and 1 have explicit boundaries + mm = MutationMap("snp", 0, "0,50.0,150.0;1,50.0,150.0") + carriers = {0, 1} + + # pos=250 => seg1 (carrier 0) + d1 = _active_site_arg_delta( + posterior, active_indexes, thread, mm, carriers, self._rec(250)) + # pos=750 => seg2 (carrier 1) + d2 = _active_site_arg_delta( + posterior, active_indexes, thread, mm, carriers, self._rec(750)) + assert d1 != 0.0 + assert d2 != 0.0 + + def test_young_mutation_stronger_correction(self): + """Young mutation (low mut_height) relative to coalescence => larger |delta|. + + When mut_height << coalescence height, the Erlang(2) prob is small, + so (arg_prob - 1) ≈ -1 and the correction removes most posterior weight. + When mut_height ≈ height, arg_prob ≈ 1 and little correction occurs. + """ + posterior = np.array([0.5]) + active_indexes = {0: 0} + thread = [self._seg(0, [0], [500.0])] + carriers = {0} + + mm_young = MutationMap("snp", 0, "-1,10.0,20.0") # mut_height ~15 + mm_old = MutationMap("snp", 0, "-1,400.0,500.0") # mut_height ~450 + + d_young = _active_site_arg_delta( + posterior, active_indexes, thread, mm_young, carriers, self._rec(50)) + d_old = _active_site_arg_delta( + posterior, active_indexes, thread, mm_old, carriers, self._rec(50)) + assert abs(d_young) > abs(d_old) + + +# --------------------------------------------------------------------------- +# CachedPosteriorSnps +# --------------------------------------------------------------------------- +class TestCachedPosteriorSnps: + def _make_posteriors(self, n_targets=3, n_snps=5, n_panel=20, seed=42): + rng = np.random.default_rng(seed) + posteriors = [] + for _ in range(n_targets): + dense = rng.random((n_snps, n_panel)) + dense[dense < 0.7] = 0 + # Guarantee at least one nonzero per row + for i in range(n_snps): + if dense[i].sum() == 0: + dense[i, rng.integers(n_panel)] = 1.0 + row_sums = dense.sum(axis=1, keepdims=True) + dense /= row_sums + posteriors.append(csr_array(dense)) + return posteriors + + def test_output_shape(self): + posteriors = self._make_posteriors(n_targets=3, n_panel=20) + cache = CachedPosteriorSnps(posteriors) + assert cache[0].shape == (3, 20) + + def test_rows_normalized(self): + posteriors = self._make_posteriors() + cache = CachedPosteriorSnps(posteriors) + result = cache[0] + np.testing.assert_allclose(result.sum(axis=1), 1.0, atol=1e-10) + + def test_cache_hit_same_object(self): + posteriors = self._make_posteriors() + cache = CachedPosteriorSnps(posteriors) + r1 = cache[0] + r2 = cache[0] + assert r1 is r2 + + def test_cache_eviction(self): + posteriors = self._make_posteriors() + cache = CachedPosteriorSnps(posteriors, max_size=2) + _ = cache[0] + _ = cache[1] + assert 0 in cache.posteriors_by_snp_idx + _ = cache[2] # evicts 0 + assert 0 not in cache.posteriors_by_snp_idx + assert 1 in cache.posteriors_by_snp_idx + assert 2 in cache.posteriors_by_snp_idx + + def test_negative_index(self): + posteriors = self._make_posteriors(n_snps=5) + cache = CachedPosteriorSnps(posteriors) + r_neg = cache[-1] + cache2 = CachedPosteriorSnps(posteriors) + r_pos = cache2[4] + np.testing.assert_array_equal(r_neg, r_pos) + + +# --------------------------------------------------------------------------- +# WriterVCF format correctness +# --------------------------------------------------------------------------- +class TestWriterVCF: + def _rec(self, pos=100, ref="A", alt=["T"], af=0.1, id="snp1"): + return types.SimpleNamespace(pos=pos, ref=ref, alt=alt, af=af, id=id) + + def _write_line(self, genotypes, imputed=True): + buf = io.StringIO() + writer = WriterVCF(None) + writer.file = buf + writer.write_site(genotypes, self._rec(), imputed, "1") + return buf.getvalue().strip() + + def test_gt_field_is_binary(self): + line = self._write_line(np.array([0.0, 1.0, 0.3, 0.7])) + for sample in line.split("\t")[9:]: + gt = sample.split(":")[0] + h1, h2 = gt.split("|") + assert h1 in ("0", "1") + assert h2 in ("0", "1") + + def test_dosage_equals_hap_sum(self): + genotypes = np.array([0.2, 0.8, 0.0, 1.0]) + line = self._write_line(genotypes) + for i, sample in enumerate(line.split("\t")[9:]): + ds = float(sample.split(":")[1]) + expected = round(genotypes[2*i] + genotypes[2*i+1], 3) + assert abs(ds - expected) < 1e-6 + + def test_imp_flag_present_when_imputed(self): + assert "IMP;" in self._write_line(np.array([0.0, 1.0]), imputed=True) + + def test_imp_flag_absent_when_not_imputed(self): + assert "IMP;" not in self._write_line(np.array([0.0, 1.0]), imputed=False) + + def test_trailing_zero_stripping(self): + # hap1=0.0, hap2=0.3 => dosage=0.3 + line = self._write_line(np.array([0.0, 0.3])) + ds_str = line.split("\t")[9].split(":")[1] + assert ds_str == "0.3" + + def test_integer_dosage_no_decimal(self): + # hap1=0.0, hap2=1.0 => dosage=1.0 => "1" + line = self._write_line(np.array([0.0, 1.0])) + ds_str = line.split("\t")[9].split(":")[1] + assert ds_str == "1" + + def test_zero_dosage_format(self): + line = self._write_line(np.array([0.0, 0.0])) + ds_str = line.split("\t")[9].split(":")[1] + assert ds_str == "0" + + +# --------------------------------------------------------------------------- +# Dosage bounds (integration property) +# --------------------------------------------------------------------------- +class TestDosageBounds: + def test_haploid_dosages_in_unit_interval(self): + """Posterior rows sum to 1, so any column-subset sum is in [0, 1].""" + rng = np.random.default_rng(42) + for _ in range(5): + m = rng.integers(5, 30) + n = rng.integers(3, 20) + panel = rng.integers(0, 2, size=(m, n)).astype(np.int8) + query = rng.integers(0, 2, size=(1, m)).astype(np.int8) + recomb = np.full(m, 0.01) + posterior = fwbw(panel, query, recomb, 0.01) + # Any boolean mask over columns sums to <= 1 per row + mask = rng.choice([True, False], size=n) + if not mask.any(): + mask[0] = True + subset_sums = posterior[:, mask].sum(axis=1) + assert np.all(subset_sums >= -1e-10) + assert np.all(subset_sums <= 1 + 1e-10) diff --git a/test/test_infer.py b/test/test_infer.py new file mode 100644 index 0000000..57d1ae8 --- /dev/null +++ b/test/test_infer.py @@ -0,0 +1,536 @@ +""" +Tests for the `threads infer` command: utility functions, Matcher, ThreadsLowMem, +ThreadingInstructions, ConsistencyWrapper, and serialization round-trips. +""" +import pickle +import tempfile + +import numpy as np +import pytest + +from pathlib import Path + +TEST_DATA = Path(__file__).parent / "data" +PANEL_PGEN = str(TEST_DATA / "panel.pgen") +PANEL_PVAR = str(TEST_DATA / "panel.pvar") +PANEL_PSAM = str(TEST_DATA / "panel.psam") +GMAP = str(TEST_DATA / "gmap_02.map") +DEMO = str(TEST_DATA / "CEU_unscaled.demo") +THREADS_FIT = str(TEST_DATA / "expected_infer_fit_to_data_snapshot.threads") +THREADS_NO_FIT = str(TEST_DATA / "expected_infer_snapshot.threads") + + +# =================================================================== +# Utility functions: split_list +# =================================================================== +class TestSplitList: + def test_even_split(self): + from threads_arg.utils import split_list + result = split_list([1, 2, 3, 4, 5, 6], 3) + assert len(result) == 3 + assert sum(len(s) for s in result) == 6 + # Flattened should be original + assert [x for s in result for x in s] == [1, 2, 3, 4, 5, 6] + + def test_uneven_split(self): + from threads_arg.utils import split_list + result = split_list([1, 2, 3, 4, 5], 3) + assert len(result) == 3 + assert sum(len(s) for s in result) == 5 + # First chunks get the extra elements + sizes = [len(s) for s in result] + assert sizes == [2, 2, 1] + + def test_n_equals_length(self): + from threads_arg.utils import split_list + result = split_list([1, 2, 3], 3) + assert result == [[1], [2], [3]] + + def test_n_greater_than_length(self): + from threads_arg.utils import split_list + result = split_list([1, 2], 5) + assert len(result) == 5 + non_empty = [s for s in result if len(s) > 0] + assert [x for s in non_empty for x in s] == [1, 2] + + def test_single_chunk(self): + from threads_arg.utils import split_list + result = split_list([1, 2, 3, 4], 1) + assert result == [[1, 2, 3, 4]] + + def test_empty_list(self): + from threads_arg.utils import split_list + result = split_list([], 3) + assert len(result) == 3 + assert all(len(s) == 0 for s in result) + + +# =================================================================== +# Utility functions: parse_demography +# =================================================================== +class TestParseDemography: + def test_loads_ceu_demo(self): + from threads_arg.utils import parse_demography + times, sizes = parse_demography(DEMO) + assert len(times) == len(sizes) + assert len(times) > 0 + assert times[0] == 0.0 + assert all(t >= 0 for t in times) + assert all(s > 0 for s in sizes) + + def test_times_monotonic(self): + from threads_arg.utils import parse_demography + times, _ = parse_demography(DEMO) + for i in range(1, len(times)): + assert times[i] > times[i - 1] + + def test_custom_demo_file(self): + from threads_arg.utils import parse_demography + with tempfile.NamedTemporaryFile(mode="w", suffix=".demo", delete=False) as f: + f.write("0.0\t10000\n100.0\t20000\n500.0\t5000\n") + f.flush() + times, sizes = parse_demography(f.name) + assert times == [0.0, 100.0, 500.0] + assert sizes == [10000.0, 20000.0, 5000.0] + + +# =================================================================== +# Utility functions: recombination maps +# =================================================================== +class TestRecombination: + def test_constant_recombination_shape(self): + from threads_arg.utils import make_constant_recombination_from_pgen + cm, phys = make_constant_recombination_from_pgen(PANEL_PGEN, 1.3e-8) + assert len(cm) == len(phys) + assert len(cm) > 0 + + def test_constant_recombination_monotonic(self): + from threads_arg.utils import make_constant_recombination_from_pgen + cm, phys = make_constant_recombination_from_pgen(PANEL_PGEN, 1.3e-8) + assert np.all(np.diff(cm) > 0), "cM positions must be strictly increasing" + assert np.all(np.diff(phys) > 0), "physical positions must be strictly increasing" + + def test_map_recombination_shape(self): + from threads_arg.utils import make_recombination_from_map_and_pgen + # gmap_02.map uses chr 20, pass None to skip chromosome check + cm, phys = make_recombination_from_map_and_pgen(GMAP, PANEL_PGEN, None) + assert len(cm) == len(phys) + assert len(cm) > 0 + + def test_map_recombination_monotonic(self): + from threads_arg.utils import make_recombination_from_map_and_pgen + cm, phys = make_recombination_from_map_and_pgen(GMAP, PANEL_PGEN, None) + assert np.all(np.diff(cm) > 0) + + +# =================================================================== +# Utility functions: read_all_genotypes, read_sample_names, etc. +# =================================================================== +class TestPgenReading: + def test_read_all_genotypes_shape(self): + from threads_arg.utils import read_all_genotypes + gt = read_all_genotypes(PANEL_PGEN) + assert gt.ndim == 2 + assert gt.dtype == np.int32 + # 500 diploid samples = 1000 haploids + assert gt.shape[1] == 1000 + + def test_read_all_genotypes_biallelic(self): + from threads_arg.utils import read_all_genotypes + gt = read_all_genotypes(PANEL_PGEN) + assert set(np.unique(gt)).issubset({0, 1}) + + def test_read_sample_names(self): + from threads_arg.utils import read_sample_names + names = read_sample_names(PANEL_PGEN) + assert len(names) == 500 + assert names[0] == "tsk_0" + assert names[-1] == "tsk_499" + + def test_read_positions_and_ids(self): + from threads_arg.utils import read_positions_and_ids + positions, ids = read_positions_and_ids(PANEL_PGEN) + assert len(positions) == len(ids) + assert len(positions) > 0 + assert all(isinstance(p, int) for p in positions) + # Positions should be sorted + assert positions == sorted(positions) + + def test_read_variant_metadata_columns(self): + from threads_arg.utils import read_variant_metadata + df = read_variant_metadata(PANEL_PGEN) + for col in ["CHROM", "POS", "ID", "REF", "ALT"]: + assert col in df.columns + assert len(df) > 0 + + +# =================================================================== +# Matcher: PBWT haplotype matching +# =================================================================== +class TestMatcher: + @pytest.fixture(scope="class") + def matcher_setup(self): + from threads_arg import Matcher + from threads_arg.utils import ( + make_recombination_from_map_and_pgen, + read_all_genotypes, + ) + cm, phys = make_recombination_from_map_and_pgen(GMAP, PANEL_PGEN, None) + gt = read_all_genotypes(PANEL_PGEN) + n_haps = gt.shape[1] + + # Filter singletons + ac = gt.sum(axis=1) + ac_mask = (ac > 1) & (ac < n_haps) + + matcher = Matcher(n_haps, cm[ac_mask], 0.01, 0.5, 4, 4) + if hasattr(matcher, "process_all_sites_numpy"): + matcher.process_all_sites_numpy(gt[ac_mask]) + else: + for g in gt[ac_mask]: + matcher.process_site(g) + matcher.propagate_adjacent_matches() + return matcher, n_haps + + def test_num_samples(self, matcher_setup): + matcher, n_haps = matcher_setup + assert matcher.num_samples == n_haps + + def test_num_sites(self, matcher_setup): + matcher, _ = matcher_setup + assert matcher.num_sites > 0 + + def test_get_matches_returns_list(self, matcher_setup): + matcher, _ = matcher_setup + # get_matches() takes no args (returns all match groups) + matches = matcher.get_matches() + assert isinstance(matches, list) + assert len(matches) > 0 + + def test_match_groups_have_candidates(self, matcher_setup): + matcher, _ = matcher_setup + matches = matcher.get_matches() + for mg in matches[:10]: # spot-check first 10 + assert len(mg.match_candidates) > 0 + + def test_cm_positions_length(self, matcher_setup): + matcher, _ = matcher_setup + cm_pos = matcher.cm_positions() + assert len(cm_pos) > 0 + + def test_serializable_matches_shape(self, matcher_setup): + matcher, _ = matcher_setup + sample_ids = [0, 1, 2] + s_matches = matcher.serializable_matches(sample_ids) + assert len(s_matches) > 0 + + +# =================================================================== +# ThreadsLowMem: HMM inference engine +# =================================================================== +class TestThreadsLowMem: + @pytest.fixture(scope="class") + def small_inference(self): + """Run a small inference on 10 haploids to test the C++ engine.""" + from threads_arg import ThreadsLowMem, Matcher + from threads_arg.utils import ( + make_recombination_from_map_and_pgen, + read_all_genotypes, + parse_demography, + ) + cm, phys = make_recombination_from_map_and_pgen(GMAP, PANEL_PGEN, None) + gt = read_all_genotypes(PANEL_PGEN) + n_haps = gt.shape[1] + ne_times, ne = parse_demography(DEMO) + + ac = gt.sum(axis=1) + ac_mask = (ac > 1) & (ac < n_haps) + + matcher = Matcher(n_haps, cm[ac_mask], 0.01, 0.5, 4, 4) + if hasattr(matcher, "process_all_sites_numpy"): + matcher.process_all_sites_numpy(gt[ac_mask]) + else: + for g in gt[ac_mask]: + matcher.process_site(g) + matcher.propagate_adjacent_matches() + + # Infer just 10 haploids + target_ids = list(range(10)) + s_matches = matcher.serializable_matches(target_ids) + cm_pos = matcher.cm_positions() + + tlm = ThreadsLowMem(target_ids, phys, cm, ne, ne_times, 1.4e-8, False) + tlm.initialize_viterbi(s_matches, cm_pos) + + if hasattr(tlm, "process_all_sites_viterbi_numpy"): + tlm.process_all_sites_viterbi_numpy(gt) + else: + for g in gt: + tlm.process_site_viterbi(g) + + tlm.prune() + tlm.traceback() + + if hasattr(tlm, "process_all_sites_hets_numpy"): + tlm.process_all_sites_hets_numpy(gt) + else: + for g in gt: + tlm.process_site_hets(g) + + tlm.date_segments() + return tlm, target_ids, phys + + def test_serialize_paths_length(self, small_inference): + tlm, target_ids, _ = small_inference + seg_starts, match_ids, heights, hetsites = tlm.serialize_paths() + assert len(seg_starts) == len(target_ids) + assert len(match_ids) == len(target_ids) + assert len(heights) == len(target_ids) + assert len(hetsites) == len(target_ids) + + def test_segment_starts_mostly_non_empty(self, small_inference): + """Most samples should have at least one segment (sample 0 may be empty as it's the reference).""" + tlm, target_ids, _ = small_inference + seg_starts, _, _, _ = tlm.serialize_paths() + non_empty = sum(1 for ss in seg_starts if len(ss) > 0) + # At least all non-reference samples should have segments + assert non_empty >= len(target_ids) - 1 + + def test_heights_positive(self, small_inference): + tlm, _, _ = small_inference + _, _, heights, _ = tlm.serialize_paths() + for h_list in heights: + for h in h_list: + assert h >= 0, f"negative height {h}" + + def test_match_ids_valid(self, small_inference): + tlm, target_ids, _ = small_inference + _, match_ids, _, _ = tlm.serialize_paths() + n_haps = 1000 # panel.pgen + for mi_list in match_ids: + for mi in mi_list: + assert 0 <= mi < n_haps, f"match_id {mi} out of range" + + def test_hetsites_within_bounds(self, small_inference): + tlm, _, phys = small_inference + _, _, _, hetsites = tlm.serialize_paths() + n_sites = len(phys) + for hs_list in hetsites: + for hs in hs_list: + assert 0 <= hs < n_sites, f"hetsite index {hs} out of range" + + +# =================================================================== +# ThreadingInstructions: construction, accessors, sub_range +# =================================================================== +class TestThreadingInstructions: + def test_construction_from_lists(self): + from threads_arg import ThreadingInstructions + starts = [[0, 50], [0, 30, 70]] + tmrcas = [[100.0, 200.0], [150.0, 250.0, 300.0]] + targets = [[1, 0], [0, 1, 0]] + mismatches = [[2, 4], [1, 3, 5]] + positions = [10, 20, 30, 40, 50, 60, 70, 80] + ti = ThreadingInstructions(starts, tmrcas, targets, mismatches, positions, 0, 100) + + assert ti.num_samples == 2 + assert ti.num_sites == 8 + assert ti.start == 0 + assert ti.end == 100 + assert ti.positions == positions + + def test_all_starts_accessor(self): + from threads_arg import ThreadingInstructions + starts = [[0, 50], [0, 30]] + ti = ThreadingInstructions(starts, [[1.0, 2.0], [3.0, 4.0]], + [[1, 0], [0, 1]], [[], []], [10, 20, 30], 0, 50) + assert ti.all_starts() == starts + + def test_all_tmrcas_accessor(self): + from threads_arg import ThreadingInstructions + tmrcas = [[100.0, 200.0], [300.0, 400.0]] + ti = ThreadingInstructions([[0, 50], [0, 30]], tmrcas, + [[1, 0], [0, 1]], [[], []], [10, 20, 30], 0, 50) + assert ti.all_tmrcas() == tmrcas + + def test_sub_range(self): + from threads_arg import ThreadingInstructions + positions = [10, 20, 30, 40, 50] + ti = ThreadingInstructions( + [[0], [0]], [[100.0], [200.0]], [[1], [0]], + [[1, 3], [0, 2, 4]], positions, 0, 60 + ) + sub = ti.sub_range(15, 45) + # sub_range should reduce the number of sites + assert sub.num_sites <= ti.num_sites + assert sub.num_samples == ti.num_samples + # start/end are snapped to actual variant positions within range + assert sub.start >= 15 + assert sub.end <= 60 # original end + + def test_pickle_roundtrip(self): + from threads_arg import ThreadingInstructions + ti = ThreadingInstructions( + [[0, 50]], [[100.0, 200.0]], [[1, 0]], [[2, 4]], + [10, 20, 30, 40, 50], 10, 50 + ) + restored = pickle.loads(pickle.dumps(ti)) + assert restored.all_starts() == ti.all_starts() + assert restored.all_tmrcas() == ti.all_tmrcas() + assert restored.all_targets() == ti.all_targets() + assert restored.all_mismatches() == ti.all_mismatches() + assert restored.positions == ti.positions + assert restored.start == ti.start + assert restored.end == ti.end + + def test_load_from_threads_file(self): + from threads_arg.serialization import load_instructions + ti = load_instructions(THREADS_NO_FIT) + assert ti.num_samples == 1000 + assert ti.num_sites == 8431 + assert len(ti.positions) == 8431 + + +# =================================================================== +# ViterbiPath: construction and accessors +# =================================================================== +class TestViterbiPath: + def test_construct_empty(self): + from threads_arg import ViterbiPath + vp = ViterbiPath(42) + assert vp.size() == 0 + + def test_construct_with_data(self): + from threads_arg import ViterbiPath + vp = ViterbiPath(0, [0, 100, 500], [5, 10, 3], [100.0, 200.0, 150.0], [2, 7]) + assert vp.size() == 3 + assert vp.segment_starts == [0, 100, 500] + assert vp.sample_ids == [5, 10, 3] + assert vp.het_sites == [2, 7] + + def test_heights_accessible(self): + from threads_arg import ViterbiPath + vp = ViterbiPath(0, [0, 100], [5, 10], [123.4, 567.8], []) + assert vp.heights == pytest.approx([123.4, 567.8]) + + +# =================================================================== +# ConsistencyWrapper: fit-to-data post-processing +# =================================================================== +class TestConsistencyWrapper: + def test_consistent_instructions_shape(self): + from threads_arg import ConsistencyWrapper, GenotypeIterator + from threads_arg.serialization import load_instructions + + ti = load_instructions(THREADS_FIT) + # The fit_to_data .threads has embedded allele ages + import h5py + with h5py.File(THREADS_FIT, "r") as f: + ages = f["allele_ages"][:].tolist() + + cw = ConsistencyWrapper(ti, ages) + gt_it = GenotypeIterator(ti) + while gt_it.has_next_genotype(): + g = np.array(gt_it.next_genotype()) + cw.process_site(g) + + consistent = cw.get_consistent_instructions() + assert consistent.num_samples == ti.num_samples + assert consistent.num_sites == ti.num_sites + + def test_consistent_preserves_genotypes(self): + """Consistency wrapper should not change the reconstructed genotypes.""" + from threads_arg import ConsistencyWrapper, GenotypeIterator + from threads_arg.serialization import load_instructions + + ti = load_instructions(THREADS_FIT) + import h5py + with h5py.File(THREADS_FIT, "r") as f: + ages = f["allele_ages"][:].tolist() + + cw = ConsistencyWrapper(ti, ages) + gt_it = GenotypeIterator(ti) + original_genotypes = [] + while gt_it.has_next_genotype(): + g = np.array(gt_it.next_genotype()) + original_genotypes.append(g.copy()) + cw.process_site(g) + + consistent = cw.get_consistent_instructions() + gt_it2 = GenotypeIterator(consistent) + for i, orig_g in enumerate(original_genotypes): + new_g = np.array(gt_it2.next_genotype()) + np.testing.assert_array_equal( + new_g, orig_g, + err_msg=f"genotype changed at site {i}" + ) + + +# =================================================================== +# Serialization: round-trip .threads write/read +# =================================================================== +class TestSerialization: + def test_serialize_load_roundtrip(self): + from threads_arg import ThreadingInstructions + from threads_arg.serialization import serialize_instructions, load_instructions + + starts = [[0, 50], [0, 30, 70]] + tmrcas = [[100.0, 200.0], [150.0, 250.0, 300.0]] + targets = [[1, 0], [0, 1, 0]] + mismatches = [[2, 4], [1, 3, 5]] + positions = [10, 20, 30, 40, 50, 60, 70, 80] + ti = ThreadingInstructions(starts, tmrcas, targets, mismatches, positions, 0, 100) + + with tempfile.NamedTemporaryFile(suffix=".threads", delete=False) as f: + out_path = f.name + + serialize_instructions(ti, out_path) + loaded = load_instructions(out_path) + + assert loaded.num_samples == ti.num_samples + assert loaded.num_sites == ti.num_sites + assert loaded.positions == ti.positions + assert loaded.start == ti.start + assert loaded.end == ti.end + assert loaded.all_starts() == ti.all_starts() + assert loaded.all_targets() == ti.all_targets() + np.testing.assert_allclose( + [h for hs in loaded.all_tmrcas() for h in hs], + [h for hs in ti.all_tmrcas() for h in hs], + ) + assert loaded.all_mismatches() == ti.all_mismatches() + + def test_serialize_with_metadata(self): + import pandas as pd + from threads_arg import ThreadingInstructions + from threads_arg.serialization import serialize_instructions, load_instructions, load_metadata, load_sample_names + + positions = [100, 200, 300] + # 2 haploid samples (1 diploid) so sample_names has 1 entry + ti = ThreadingInstructions( + [[0], [0]], [[50.0], [60.0]], [[1], [0]], [[], []], positions, 0, 400 + ) + + metadata = pd.DataFrame({ + "CHROM": ["1", "1", "1"], + "POS": [100, 200, 300], + "ID": ["rs1", "rs2", "rs3"], + "REF": ["A", "C", "G"], + "ALT": ["T", "G", "A"], + "QUAL": [".", ".", "."], + "FILTER": ["PASS", "PASS", "PASS"], + }) + sample_names = ["sample_0"] + + with tempfile.NamedTemporaryFile(suffix=".threads", delete=False) as f: + out_path = f.name + + serialize_instructions(ti, out_path, variant_metadata=metadata, sample_names=sample_names) + + loaded_meta = load_metadata(out_path) + assert list(loaded_meta.columns) == ["CHROM", "POS", "ID", "REF", "ALT", "QUAL", "FILTER"] + assert len(loaded_meta) == 3 + + loaded_names = load_sample_names(out_path) + # load_sample_names returns bytes from HDF5 + assert [n.decode() if isinstance(n, bytes) else n for n in loaded_names] == ["sample_0"] diff --git a/test/test_map.py b/test/test_map.py new file mode 100644 index 0000000..2c35636 --- /dev/null +++ b/test/test_map.py @@ -0,0 +1,263 @@ +""" +Tests for the `threads map` command: _mapping_string formatting, get_leaf_ids_at, +MAF filtering logic, and full pipeline validation. +""" +import tempfile + +import numpy as np +import pytest +import arg_needle_lib + +from pathlib import Path + +TEST_DATA = Path(__file__).parent / "data" +ARGN_SNAPSHOT = str(TEST_DATA / "expected_convert_snapshot.argn") +MUT_SNAPSHOT = str(TEST_DATA / "expected_mapping_snapshot.mut") +PANEL_VCF = str(TEST_DATA / "panel.vcf.gz") +REGION = "1:400000-600000" + + +# =================================================================== +# _mapping_string: format mutation-to-edge mappings +# =================================================================== +class TestMappingString: + @pytest.fixture() + def mock_edge(self): + """Create a minimal mock edge with child/parent height attributes.""" + class MockNode: + def __init__(self, height): + self.height = height + self.ID = 0 + class MockEdge: + def __init__(self, child_h, parent_h): + self.child = MockNode(child_h) + self.parent = MockNode(parent_h) + return MockEdge + + def test_empty_edges_returns_nan(self): + from threads_arg.map_mutations_to_arg import _mapping_string + assert _mapping_string([], []) == "NaN" + + def test_single_edge_format(self, mock_edge): + from threads_arg.map_mutations_to_arg import _mapping_string + edge = mock_edge(0.0, 1138.7262) + result = _mapping_string([[-1]], [edge]) + assert result == "-1,0.0000,1138.7262" + + def test_single_edge_precision(self, mock_edge): + from threads_arg.map_mutations_to_arg import _mapping_string + edge = mock_edge(378.9087, 1303.8883) + result = _mapping_string([[-1]], [edge]) + assert result == "-1,378.9087,1303.8883" + + def test_multiple_edges_format(self, mock_edge): + from threads_arg.map_mutations_to_arg import _mapping_string + edges = [mock_edge(0.0, 599.4252), mock_edge(0.0, 659.2068), mock_edge(0.0, 704.7969)] + carrier_sets = [[124], [740], [816]] + result = _mapping_string(carrier_sets, edges) + parts = result.split(";") + assert len(parts) == 3 + assert parts[0] == "124,0.0000,599.4252" + assert parts[1] == "740,0.0000,659.2068" + assert parts[2] == "816,0.0000,704.7969" + + def test_single_edge_always_uses_sentinel(self, mock_edge): + """A single edge always uses -1 sentinel regardless of carrier set.""" + from threads_arg.map_mutations_to_arg import _mapping_string + edges = [mock_edge(100.0, 500.0)] + carrier_sets = [[5, 10, 15]] + result = _mapping_string(carrier_sets, edges) + assert result == "-1,100.0000,500.0000" + + def test_two_edges_with_mixed_carriers(self, mock_edge): + from threads_arg.map_mutations_to_arg import _mapping_string + edges = [mock_edge(693.77, 918.05), mock_edge(0.0, 1618.53)] + carrier_sets = [[569, 147, 72], [156]] + result = _mapping_string(carrier_sets, edges) + parts = result.split(";") + assert len(parts) == 2 + assert parts[0].startswith("569.147.72,") + assert parts[1].startswith("156,") + + +# =================================================================== +# get_leaf_ids_at: recursive leaf extraction from ARG edges +# =================================================================== +class TestGetLeafIds: + @pytest.fixture(scope="class") + def arg_with_roots(self): + arg = arg_needle_lib.deserialize_arg(ARGN_SNAPSHOT) + arg.populate_children_and_roots() + return arg + + def test_leaf_ids_non_empty(self, arg_with_roots): + """Mapping a real variant should produce non-empty leaf IDs.""" + from threads_arg.map_mutations_to_arg import get_leaf_ids_at + from cyvcf2 import VCF + + vcf = VCF(PANEL_VCF) + for record in vcf(REGION): + ac = int(record.INFO.get("AC")) + an = int(record.INFO.get("AN")) + maf = min(ac / an, 1 - ac / an) + if maf == 0 or maf > 0.01: + continue + + hap = np.array(record.genotypes)[:, :2].flatten() + mapping, _ = arg_needle_lib.map_genotype_to_ARG_approximate( + arg_with_roots, hap, float(record.POS - arg_with_roots.offset) + ) + if len(mapping) > 1: + for edge in mapping: + leaves = get_leaf_ids_at(arg_with_roots, edge, record.POS) + assert len(leaves) > 0, "edge should subtend at least one leaf" + return # tested one multi-edge mapping, done + pytest.skip("No multi-edge mapping found in region") + + def test_leaf_ids_are_leaves(self, arg_with_roots): + """All returned IDs should be actual leaf nodes.""" + from threads_arg.map_mutations_to_arg import get_leaf_ids_at + from cyvcf2 import VCF + + vcf = VCF(PANEL_VCF) + for record in vcf(REGION): + ac = int(record.INFO.get("AC")) + an = int(record.INFO.get("AN")) + maf = min(ac / an, 1 - ac / an) + if maf == 0 or maf > 0.01: + continue + + hap = np.array(record.genotypes)[:, :2].flatten() + mapping, _ = arg_needle_lib.map_genotype_to_ARG_approximate( + arg_with_roots, hap, float(record.POS - arg_with_roots.offset) + ) + if len(mapping) > 1: + for edge in mapping: + leaves = get_leaf_ids_at(arg_with_roots, edge, record.POS) + for lid in leaves: + assert arg_with_roots.is_leaf(lid) + return + pytest.skip("No multi-edge mapping found in region") + + +# =================================================================== +# .mut file format validation +# =================================================================== +class TestMutFileFormat: + @pytest.fixture(scope="class") + def mut_lines(self): + with open(MUT_SNAPSHOT) as f: + return f.readlines() + + def test_non_empty(self, mut_lines): + assert len(mut_lines) > 0 + + def test_tab_separated_four_columns(self, mut_lines): + for i, line in enumerate(mut_lines[:50]): + fields = line.rstrip("\n").split("\t") + assert len(fields) == 4, f"line {i+1}: expected 4 tab-separated fields, got {len(fields)}" + + def test_position_column_is_int(self, mut_lines): + for line in mut_lines: + fields = line.rstrip("\n").split("\t") + int(fields[1]) # should not raise + + def test_flipped_column_is_binary(self, mut_lines): + for line in mut_lines: + fields = line.rstrip("\n").split("\t") + assert fields[2] in ("0", "1") + + def test_mapping_string_format(self, mut_lines): + """Mapping string should be NaN, or -1,h1,h2, or semicolon-separated groups.""" + for line in mut_lines: + fields = line.rstrip("\n").split("\t") + ms = fields[3] + if ms == "NaN": + continue + groups = ms.split(";") + for group in groups: + parts = group.split(",") + assert len(parts) == 3, f"bad mapping group: {group!r}" + # Last two should be floats + float(parts[1]) + float(parts[2]) + + def test_positions_in_region(self, mut_lines): + """All positions should be within the expected region.""" + for line in mut_lines: + fields = line.rstrip("\n").split("\t") + pos = int(fields[1]) + assert 400000 <= pos <= 600000 + + def test_uniquely_mapped_uses_sentinel(self, mut_lines): + """Single-edge mappings should use -1 sentinel.""" + for line in mut_lines: + fields = line.rstrip("\n").split("\t") + ms = fields[3] + if ms == "NaN": + continue + if ";" not in ms: + # Single edge + assert ms.startswith("-1,"), f"single-edge mapping should start with -1: {ms!r}" + + +# =================================================================== +# Full pipeline: threads_map_mutations_to_arg +# =================================================================== +class TestMapPipeline: + def test_produces_mut_file(self): + from threads_arg.map_mutations_to_arg import threads_map_mutations_to_arg + with tempfile.TemporaryDirectory() as tmpdir: + out_mut = str(Path(tmpdir) / "test.mut") + threads_map_mutations_to_arg( + argn=ARGN_SNAPSHOT, out=out_mut, maf=0.01, + input=PANEL_VCF, region=REGION, num_threads=1 + ) + assert Path(out_mut).exists() + with open(out_mut) as f: + lines = f.readlines() + assert len(lines) > 0 + + def test_same_variants_mapped(self): + """Same variant IDs and positions should be mapped as the snapshot.""" + from threads_arg.map_mutations_to_arg import threads_map_mutations_to_arg + with tempfile.TemporaryDirectory() as tmpdir: + out_mut = str(Path(tmpdir) / "test.mut") + threads_map_mutations_to_arg( + argn=ARGN_SNAPSHOT, out=out_mut, maf=0.01, + input=PANEL_VCF, region=REGION, num_threads=1 + ) + with open(out_mut) as gen, open(MUT_SNAPSHOT) as exp: + gen_lines = gen.readlines() + exp_lines = exp.readlines() + + assert len(gen_lines) == len(exp_lines), \ + f"line count: {len(gen_lines)} vs {len(exp_lines)}" + # Variant IDs and positions should match exactly + for i, (gl, el) in enumerate(zip(gen_lines, exp_lines)): + g_fields = gl.split("\t") + e_fields = el.split("\t") + assert g_fields[0] == e_fields[0], f"line {i+1}: variant ID differs" + assert g_fields[1] == e_fields[1], f"line {i+1}: position differs" + assert g_fields[2] == e_fields[2], f"line {i+1}: flipped flag differs" + + def test_stricter_maf_fewer_variants(self): + """Tighter MAF filter should produce fewer or equal mappings.""" + from threads_arg.map_mutations_to_arg import threads_map_mutations_to_arg + with tempfile.TemporaryDirectory() as tmpdir: + out_wide = str(Path(tmpdir) / "wide.mut") + out_strict = str(Path(tmpdir) / "strict.mut") + + threads_map_mutations_to_arg( + argn=ARGN_SNAPSHOT, out=out_wide, maf=0.05, + input=PANEL_VCF, region=REGION, num_threads=1 + ) + threads_map_mutations_to_arg( + argn=ARGN_SNAPSHOT, out=out_strict, maf=0.005, + input=PANEL_VCF, region=REGION, num_threads=1 + ) + with open(out_wide) as f: + n_wide = len(f.readlines()) + with open(out_strict) as f: + n_strict = len(f.readlines()) + assert n_strict <= n_wide diff --git a/test/test_matcher.cpp b/test/test_matcher.cpp new file mode 100644 index 0000000..87802f1 --- /dev/null +++ b/test/test_matcher.cpp @@ -0,0 +1,190 @@ +// This file is part of the Threads software suite. +// Copyright (C) 2024-2025 Threads Developers. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU General Public License for more details. +// +// You should have received a copy of the GNU General Public License +// along with this program. If not, see . + +#include "Matcher.hpp" + +#include +#include +#include +#include + +namespace { + +// Create evenly spaced genetic positions +std::vector linear_positions(int n_sites, double start = 0.0, double step = 0.01) { + std::vector pos; + pos.reserve(n_sites); + for (int i = 0; i < n_sites; i++) { + pos.push_back(start + i * step); + } + return pos; +} + +} // namespace + +TEST_CASE("Matcher construction basic") { + int n_samples = 10; + int n_sites = 100; + auto positions = linear_positions(n_sites); + + Matcher m(n_samples, positions, 0.01, 0.5, 4, 2); + CHECK(m.num_samples == n_samples); + CHECK(m.num_sites == n_sites); +} + +TEST_CASE("Matcher construction requires >= 3 sites") { + auto pos2 = linear_positions(2); + CHECK_THROWS_WITH(Matcher(5, pos2, 0.01, 0.5, 4, 2), + Catch::Matchers::ContainsSubstring("Need at least 3 sites")); +} + +TEST_CASE("Matcher construction requires increasing positions") { + std::vector bad_pos = {0.0, 0.5, 0.3, 0.8}; + CHECK_THROWS_WITH(Matcher(5, bad_pos, 0.01, 0.5, 4, 2), + Catch::Matchers::ContainsSubstring("strictly increasing")); +} + +TEST_CASE("Matcher process_site with binary genotypes") { + int n_samples = 20; + int n_sites = 50; + auto positions = linear_positions(n_sites, 0.0, 0.02); + + Matcher m(n_samples, positions, 0.01, 0.5, 4, 1); + + // Process all sites with alternating genotypes + for (int site = 0; site < n_sites; site++) { + std::vector geno(n_samples); + for (int s = 0; s < n_samples; s++) { + geno[s] = (s + site) % 2; + } + m.process_site(geno); + } + + auto matches = m.get_matches(); + CHECK(matches.size() > 0); +} + +TEST_CASE("Matcher rejects wrong genotype size") { + int n_samples = 10; + auto positions = linear_positions(20, 0.0, 0.02); + Matcher m(n_samples, positions, 0.01, 0.5, 4, 1); + + // Wrong size genotype + std::vector bad_geno(5, 0); + CHECK_THROWS_WITH(m.process_site(bad_geno), + Catch::Matchers::ContainsSubstring("invalid genotype vector size")); +} + +TEST_CASE("Matcher rejects invalid alleles") { + int n_samples = 5; + auto positions = linear_positions(10, 0.0, 0.02); + Matcher m(n_samples, positions, 0.01, 0.5, 4, 1); + + std::vector bad_geno = {0, 1, 0, 2, 0}; // 2 is invalid + CHECK_THROWS_WITH(m.process_site(bad_geno), + Catch::Matchers::ContainsSubstring("invalid genotype")); +} + +TEST_CASE("Matcher process_site rejects extra sites") { + int n_samples = 5; + auto positions = linear_positions(4, 0.0, 0.02); + Matcher m(n_samples, positions, 0.01, 0.5, 4, 1); + + std::vector geno(5, 0); + for (int i = 0; i < 4; i++) { + m.process_site(geno); + } + CHECK_THROWS_WITH(m.process_site(geno), + Catch::Matchers::ContainsSubstring("all sites have already been processed")); +} + +TEST_CASE("Matcher sorting is a valid permutation") { + int n_samples = 10; + int n_sites = 30; + auto positions = linear_positions(n_sites, 0.0, 0.02); + + Matcher m(n_samples, positions, 0.01, 0.5, 4, 1); + + for (int site = 0; site < n_sites; site++) { + std::vector geno(n_samples); + for (int s = 0; s < n_samples; s++) { + geno[s] = (s * 3 + site) % 2; + } + m.process_site(geno); + } + + auto sorting = m.get_sorting(); + CHECK(static_cast(sorting.size()) == n_samples); + + // Check it's a valid permutation + std::unordered_set seen; + for (int v : sorting) { + CHECK(v >= 0); + CHECK(v < n_samples); + seen.insert(v); + } + CHECK(static_cast(seen.size()) == n_samples); +} + +TEST_CASE("Matcher cm_positions") { + int n_samples = 10; + int n_sites = 100; + auto positions = linear_positions(n_sites, 0.0, 0.01); + + Matcher m(n_samples, positions, 0.01, 0.5, 4, 1); + + // Process all sites + for (int site = 0; site < n_sites; site++) { + std::vector geno(n_samples); + for (int s = 0; s < n_samples; s++) { + geno[s] = s % 2; + } + m.process_site(geno); + } + + auto cms = m.cm_positions(); + CHECK(cms.size() > 0); + // Should be non-decreasing + for (std::size_t i = 1; i < cms.size(); i++) { + CHECK(cms[i] >= cms[i - 1]); + } +} + +TEST_CASE("MatchGroup construction") { + MatchGroup mg(10, 0.5); + CHECK(mg.num_samples == 10); + CHECK(mg.cm_position == 0.5); + CHECK(mg.match_candidates_counts.size() == 10); +} + +TEST_CASE("MatchGroup from targets and matches") { + std::vector targets = {0, 1, 2}; + std::vector> matches = {{}, {0}, {0, 1}}; + MatchGroup mg(targets, matches, 1.0); + + CHECK(mg.match_candidates.size() == 3); + CHECK(mg.match_candidates.at(0).size() == 0); + CHECK(mg.match_candidates.at(1).size() == 1); + CHECK(mg.match_candidates.at(2).size() == 2); +} + +TEST_CASE("MatchGroup clear") { + MatchGroup mg(5, 0.0); + mg.clear(); + CHECK(mg.match_candidates.empty()); + CHECK(mg.match_candidates_counts.empty()); + CHECK(mg.top_four_maps.empty()); +} diff --git a/test/test_node.cpp b/test/test_node.cpp new file mode 100644 index 0000000..ae9725f --- /dev/null +++ b/test/test_node.cpp @@ -0,0 +1,81 @@ +// This file is part of the Threads software suite. +// Copyright (C) 2024-2025 Threads Developers. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU General Public License for more details. +// +// You should have received a copy of the GNU General Public License +// along with this program. If not, see . + +#include "Node.hpp" + +#include +#include + +TEST_CASE("Node construction") { + Node n(42, 10, true); + CHECK(n.sample_ID == 42); + CHECK(n.divergence == 10); + CHECK(n.genotype == true); + CHECK(n.above == nullptr); + CHECK(n.below == nullptr); + CHECK(n.w[0] == nullptr); + CHECK(n.w[1] == nullptr); +} + +TEST_CASE("Node insert_above") { + // Set up a two-node chain: bottom <-> top + Node bottom(0, 0, false); + Node top(1, 0, true); + bottom.above = ⊤ + top.below = ⊥ + + // Insert middle between bottom and top + Node middle(2, 5, false); + bottom.insert_above(&middle); + + // Verify chain is now bottom <-> middle <-> top + CHECK(bottom.above == &middle); + CHECK(middle.below == &bottom); + CHECK(middle.above == &top); + CHECK(top.below == &middle); +} + +TEST_CASE("Node insert_above multiple") { + // Build chain of 4 nodes by inserting above bottom + Node bottom(0, 0, false); + Node top(1, 0, true); + bottom.above = ⊤ + top.below = ⊥ + + Node n1(2, 1, true); + Node n2(3, 2, false); + + bottom.insert_above(&n1); + n1.insert_above(&n2); + + // Chain: bottom <-> n1 <-> n2 <-> top + CHECK(bottom.above == &n1); + CHECK(n1.above == &n2); + CHECK(n2.above == &top); + CHECK(top.below == &n2); +} + +TEST_CASE("Node stream output") { + Node n(7, 3, true); + std::ostringstream oss; + oss << n; + CHECK(oss.str() == "Node for sample 7 carrying allele 1"); + + Node n2(0, 0, false); + std::ostringstream oss2; + oss2 << n2; + CHECK(oss2.str() == "Node for sample 0 carrying allele 0"); +} diff --git a/test/test_normalization.py b/test/test_normalization.py new file mode 100644 index 0000000..709ff76 --- /dev/null +++ b/test/test_normalization.py @@ -0,0 +1,145 @@ +""" +Tests for normalization.py: Normalizer class (demography-based TMRCA correction). +Uses a minimal synthetic demography and small sample count to keep tests fast. +""" +import tempfile + +import numpy as np +import pytest + +from threads_arg import ThreadingInstructions + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- +@pytest.fixture(scope="module") +def demo_file(): + """Minimal constant-Ne demography file.""" + with tempfile.NamedTemporaryFile(mode="w", suffix=".demo", delete=False) as f: + f.write("0.0\t10000\n") + f.write("500.0\t10000\n") + f.flush() + return f.name + + +@pytest.fixture(scope="module") +def normalizer(demo_file): + from threads_arg.normalization import Normalizer + return Normalizer(demo_file, num_samples=20) + + +@pytest.fixture(scope="module") +def simple_instructions(): + """Small ThreadingInstructions: 20 haploids, 5 sites.""" + # Sample 0 is the reference (no threading info needed for it, + # but ThreadingInstructions expects entries for all samples) + starts = [[0]] * 20 + tmrcas = [[float(i * 50 + 100)] for i in range(20)] # 100, 150, ..., 1050 + targets = [[0]] * 20 + mismatches = [[]] * 20 + positions = [100, 200, 300, 400, 500] + return ThreadingInstructions(starts, tmrcas, targets, mismatches, positions, 0, 600) + + +# =================================================================== +# Normalizer construction +# =================================================================== +class TestNormalizerInit: + def test_num_samples_stored(self, normalizer): + assert normalizer.num_samples == 20 + + def test_demography_created(self, normalizer): + assert normalizer.demography is not None + + +# =================================================================== +# Normalizer.simulation +# =================================================================== +class TestSimulation: + def test_returns_tree_sequence(self, normalizer): + ts = normalizer.simulation(1e6, random_seed=42) + assert ts.num_samples == 20 # 10 diploid = 20 haploid + + def test_has_internal_nodes(self, normalizer): + ts = normalizer.simulation(1e6, random_seed=42) + internal_times = [n.time for n in ts.nodes() if n.time > 0] + assert len(internal_times) > 0 + + def test_deterministic_with_seed(self, normalizer): + ts1 = normalizer.simulation(1e6, random_seed=99) + ts2 = normalizer.simulation(1e6, random_seed=99) + times1 = sorted([n.time for n in ts1.nodes()]) + times2 = sorted([n.time for n in ts2.nodes()]) + assert times1 == times2 + + +# =================================================================== +# Normalizer.normalize +# =================================================================== +class TestNormalize: + def test_preserves_starts(self, normalizer, simple_instructions): + result = normalizer.normalize(simple_instructions, num_seeds=10) + assert result.all_starts() == simple_instructions.all_starts() + + def test_preserves_targets(self, normalizer, simple_instructions): + result = normalizer.normalize(simple_instructions, num_seeds=10) + assert result.all_targets() == simple_instructions.all_targets() + + def test_preserves_mismatches(self, normalizer, simple_instructions): + result = normalizer.normalize(simple_instructions, num_seeds=10) + assert result.all_mismatches() == simple_instructions.all_mismatches() + + def test_preserves_positions(self, normalizer, simple_instructions): + result = normalizer.normalize(simple_instructions, num_seeds=10) + assert result.positions == simple_instructions.positions + + def test_preserves_num_samples(self, normalizer, simple_instructions): + result = normalizer.normalize(simple_instructions, num_seeds=10) + assert result.num_samples == simple_instructions.num_samples + + def test_preserves_region(self, normalizer, simple_instructions): + result = normalizer.normalize(simple_instructions, num_seeds=10) + assert result.start == simple_instructions.start + assert result.end == simple_instructions.end + + def test_tmrcas_changed(self, normalizer, simple_instructions): + """Normalized TMRCAs should differ from originals (unless by coincidence).""" + result = normalizer.normalize(simple_instructions, num_seeds=10) + old = [h for hs in simple_instructions.all_tmrcas() for h in hs] + new = [h for hs in result.all_tmrcas() for h in hs] + assert len(old) == len(new) + # At least some should differ + n_changed = sum(1 for a, b in zip(old, new) if abs(a - b) > 1e-6) + assert n_changed > 0 + + def test_tmrcas_non_negative(self, normalizer, simple_instructions): + result = normalizer.normalize(simple_instructions, num_seeds=10) + for tmrca_vec in result.all_tmrcas(): + for t in tmrca_vec: + assert t >= 0, f"negative TMRCA: {t}" + + def test_monotonic_mapping(self, normalizer, simple_instructions): + """If original TMRCA_a < TMRCA_b, normalized should preserve order.""" + result = normalizer.normalize(simple_instructions, num_seeds=10) + old_tmrcas = [hs[0] for hs in simple_instructions.all_tmrcas()] + new_tmrcas = [hs[0] for hs in result.all_tmrcas()] + # Sort by old, check new is non-decreasing + pairs = sorted(zip(old_tmrcas, new_tmrcas)) + new_sorted = [p[1] for p in pairs] + for i in range(1, len(new_sorted)): + assert new_sorted[i] >= new_sorted[i - 1] - 1e-10, \ + f"monotonicity violated: {new_sorted[i-1]} > {new_sorted[i]}" + + def test_sample_count_mismatch_raises(self, demo_file): + """Normalizer should assert if num_samples doesn't match instructions.""" + from threads_arg.normalization import Normalizer + norm = Normalizer(demo_file, num_samples=10) + # Instructions with 20 samples but normalizer expects 10 + starts = [[0]] * 20 + tmrcas = [[100.0]] * 20 + targets = [[0]] * 20 + mismatches = [[]] * 20 + ti = ThreadingInstructions(starts, tmrcas, targets, mismatches, [100], 0, 200) + with pytest.raises(AssertionError): + norm.normalize(ti, num_seeds=5) diff --git a/test/test_regression.cpp b/test/test_regression.cpp new file mode 100644 index 0000000..2e0eb5d --- /dev/null +++ b/test/test_regression.cpp @@ -0,0 +1,207 @@ +// This file is part of the Threads software suite. +// Copyright (C) 2024-2025 Threads Developers. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU General Public License for more details. +// +// You should have received a copy of the GNU General Public License +// along with this program. If not, see . +// +// Regression tests that capture exact numerical outputs of core algorithms. +// These tests ensure that optimizations produce bit-identical results. + +#include "Demography.hpp" +#include "HMM.hpp" +#include "Matcher.hpp" +#include "ViterbiLowMem.hpp" + +#include +#include +#include +#include + +// ---- Demography regression ---- + +TEST_CASE("Regression: Demography constant Ne=10000 std_to_gen values") { + Demography d({10000.0}, {0.0}); + + // Pin exact values + CHECK_THAT(d.std_to_gen(0.0), Catch::Matchers::WithinAbs(0.0, 1e-14)); + CHECK_THAT(d.std_to_gen(0.001), Catch::Matchers::WithinAbs(10.0, 1e-10)); + CHECK_THAT(d.std_to_gen(0.5), Catch::Matchers::WithinAbs(5000.0, 1e-10)); + CHECK_THAT(d.std_to_gen(1.0), Catch::Matchers::WithinAbs(10000.0, 1e-10)); + CHECK_THAT(d.std_to_gen(3.0), Catch::Matchers::WithinAbs(30000.0, 1e-10)); + CHECK_THAT(d.expected_time, Catch::Matchers::WithinAbs(10000.0, 1e-10)); + CHECK_THAT(d.expected_branch_length(100), Catch::Matchers::WithinAbs(200.0, 1e-10)); +} + +TEST_CASE("Regression: Demography piecewise Ne values") { + // Ne=5000 for [0,200), Ne=20000 for [200, ...) + Demography d({5000.0, 20000.0}, {0.0, 200.0}); + + CHECK_THAT(d.std_times[0], Catch::Matchers::WithinAbs(0.0, 1e-14)); + CHECK_THAT(d.std_times[1], Catch::Matchers::WithinAbs(0.04, 1e-14)); + + CHECK_THAT(d.std_to_gen(0.02), Catch::Matchers::WithinAbs(100.0, 1e-10)); + CHECK_THAT(d.std_to_gen(0.04), Catch::Matchers::WithinAbs(200.0, 1e-10)); + CHECK_THAT(d.std_to_gen(0.05), Catch::Matchers::WithinAbs(400.0, 1e-10)); + CHECK_THAT(d.std_to_gen(0.1), Catch::Matchers::WithinAbs(1400.0, 1e-10)); +} + +TEST_CASE("Regression: Demography stream output") { + Demography d({1000.0}, {0.0}); + std::ostringstream oss; + oss << d; + // Note: current code has bug (writes to std::cout not os), so this captures current behavior + // After fix, this should contain the output +} + +// ---- HMM regression ---- + +TEST_CASE("Regression: HMM expected_times K=4 constant Ne=10000") { + Demography demo({10000.0}, {0.0}); + std::vector bp(10, 100.0); + std::vector cm(10, 0.001); + + HMM hmm(demo, bp, cm, 1.4e-8, 4); + + // Pin the expected times - these come from quantiles of the exponential distribution + CHECK(hmm.expected_times.size() == 4); + for (int i = 0; i < 4; i++) { + CHECK(hmm.expected_times[i] > 0.0); + } + // Times must be strictly increasing + for (int i = 1; i < 4; i++) { + CHECK(hmm.expected_times[i] > hmm.expected_times[i - 1]); + } + + // Pin exact values for reproducibility (from actual Boost quantile computation) + CHECK_THAT(hmm.expected_times[0], Catch::Matchers::WithinRel(1335.31, 0.001)); + CHECK_THAT(hmm.expected_times[1], Catch::Matchers::WithinRel(4700.04, 0.001)); + CHECK_THAT(hmm.expected_times[2], Catch::Matchers::WithinRel(9808.29, 0.001)); + CHECK_THAT(hmm.expected_times[3], Catch::Matchers::WithinRel(20794.42, 0.001)); +} + +TEST_CASE("Regression: HMM score tables dimensions and sign") { + int n_sites = 20; + int K = 8; + Demography demo({10000.0}, {0.0}); + std::vector bp(n_sites, 100.0); + std::vector cm(n_sites, 0.001); + + HMM hmm(demo, bp, cm, 1.4e-8, K); + + CHECK(hmm.transition_score.size() == n_sites); + CHECK(hmm.non_transition_score.size() == n_sites); + CHECK(hmm.hom_score.size() == n_sites); + CHECK(hmm.het_score.size() == n_sites); + + for (int i = 0; i < n_sites; i++) { + CHECK(static_cast(hmm.transition_score[i].size()) == K); + CHECK(static_cast(hmm.non_transition_score[i].size()) == K); + CHECK(static_cast(hmm.hom_score[i].size()) == K); + CHECK(static_cast(hmm.het_score[i].size()) == K); + } +} + +TEST_CASE("Regression: HMM breakpoints deterministic for fixed input") { + int n_sites = 30; + int K = 4; + Demography demo({10000.0}, {0.0}); + std::vector bp(n_sites, 100.0); + std::vector cm(n_sites, 0.001); + + HMM hmm(demo, bp, cm, 1.4e-8, K); + + // Fixed observation pattern + std::vector obs(n_sites, false); + obs[5] = true; + obs[6] = true; + obs[15] = true; + obs[16] = true; + obs[17] = true; + obs[25] = true; + + auto bps1 = hmm.breakpoints(obs, 0); + // Re-initialize trellis (breakpoints modifies it) + HMM hmm2(demo, bp, cm, 1.4e-8, K); + auto bps2 = hmm2.breakpoints(obs, 0); + + // Must be deterministic + CHECK(bps1 == bps2); + CHECK(bps1[0] == 0); +} + +// ---- ViterbiState regression ---- + +TEST_CASE("Regression: ViterbiState deterministic output for fixed genotypes") { + std::vector samples = {0, 1, 2}; + ViterbiState state1(3, samples); + + // Fixed genotype sequence + std::vector> genotypes = { + {1, 0, 0, 1}, // site 0: target matches sample 0 + {1, 0, 1, 1}, // site 1 + {0, 1, 0, 0}, // site 2: target matches sample 0 and 2 + {1, 1, 0, 1}, // site 3 + {0, 0, 1, 0}, // site 4 + {1, 0, 0, 1}, // site 5 + {0, 1, 1, 0}, // site 6 + {1, 0, 0, 1}, // site 7 + }; + + double rho = 3.0, rho_c = 0.01, mu = 2.0, mu_c = 0.01; + for (auto& g : genotypes) { + state1.process_site(g, rho, rho_c, mu, mu_c); + } + + auto path1 = state1.traceback(); + + // Run again independently + ViterbiState state2(3, samples); + for (auto& g : genotypes) { + state2.process_site(g, rho, rho_c, mu, mu_c); + } + auto path2 = state2.traceback(); + + // Must be identical + CHECK(path1.segment_starts == path2.segment_starts); + CHECK(path1.sample_ids == path2.sample_ids); + CHECK(path1.score == path2.score); + CHECK(path1.target_id == path2.target_id); +} + +// ---- Matcher regression ---- + +TEST_CASE("Regression: Matcher PBWT sorting deterministic") { + int n_samples = 10; + int n_sites = 40; + std::vector positions; + for (int i = 0; i < n_sites; i++) { + positions.push_back(i * 0.02); + } + + // Fixed genotype pattern + auto run = [&]() { + Matcher m(n_samples, positions, 0.01, 0.5, 4, 1); + for (int site = 0; site < n_sites; site++) { + std::vector geno(n_samples); + for (int s = 0; s < n_samples; s++) { + geno[s] = ((s * 7 + site * 3) % 5) < 2 ? 1 : 0; + } + m.process_site(geno); + } + return m.get_sorting(); + }; + + auto sorting1 = run(); + auto sorting2 = run(); + CHECK(sorting1 == sorting2); +} diff --git a/test/test_threading_instructions.cpp b/test/test_threading_instructions.cpp new file mode 100644 index 0000000..de17fc5 --- /dev/null +++ b/test/test_threading_instructions.cpp @@ -0,0 +1,124 @@ +// This file is part of the Threads software suite. +// Copyright (C) 2024-2025 Threads Developers. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU General Public License for more details. +// +// You should have received a copy of the GNU General Public License +// along with this program. If not, see . + +#include "ThreadingInstructions.hpp" + +#include +#include +#include +#include + +TEST_CASE("ThreadingInstruction construction") { + std::vector starts = {100, 500, 800}; + std::vector tmrcas = {50.0, 200.0, 100.0}; + std::vector targets = {0, 3, 1}; + std::vector mismatches = {2, 7, 15}; + + ThreadingInstruction ti(starts, tmrcas, targets, mismatches); + CHECK(ti.num_segments == 3); + CHECK(ti.num_mismatches == 3); + CHECK(ti.starts == starts); + CHECK(ti.tmrcas == tmrcas); + CHECK(ti.targets == targets); + CHECK(ti.mismatches == mismatches); +} + +TEST_CASE("ThreadingInstruction mismatching lengths throw") { + CHECK_THROWS(ThreadingInstruction({0}, {1.0, 2.0}, {0}, {})); + CHECK_THROWS(ThreadingInstruction({0, 1}, {1.0}, {0, 1}, {})); +} + +TEST_CASE("ThreadingInstructions construction from components") { + std::vector> starts = {{100, 500}, {100, 300, 700}}; + std::vector> tmrcas = {{10.0, 20.0}, {5.0, 15.0, 25.0}}; + std::vector> targets = {{0, 1}, {0, 1, 0}}; + std::vector> mismatches = {{1}, {0, 3}}; + std::vector positions = {100, 200, 300, 400, 500, 600, 700, 800}; + + ThreadingInstructions ti(starts, tmrcas, targets, mismatches, positions, 100, 800); + CHECK(ti.num_samples == 2); + CHECK(ti.num_sites == 8); + CHECK(ti.start == 100); + CHECK(ti.end == 800); +} + +TEST_CASE("ThreadingInstructions all_starts/tmrcas/targets/mismatches") { + std::vector> starts = {{0, 5}}; + std::vector> tmrcas = {{10.0, 20.0}}; + std::vector> targets = {{3, 7}}; + std::vector> mismatches = {{2}}; + std::vector positions = {100, 200, 300, 400, 500, 600, 700, 800, 900, 1000}; + + ThreadingInstructions ti(starts, tmrcas, targets, mismatches, positions, 100, 1000); + + auto all_s = ti.all_starts(); + CHECK(all_s.size() == 1); + CHECK(all_s[0] == std::vector{0, 5}); + + auto all_t = ti.all_tmrcas(); + CHECK(all_t.size() == 1); + CHECK(all_t[0][0] == 10.0); + + auto all_tg = ti.all_targets(); + CHECK(all_tg[0] == std::vector{3, 7}); + + auto all_m = ti.all_mismatches(); + CHECK(all_m[0] == std::vector{2}); +} + +TEST_CASE("ThreadingInstructionIterator basic iteration") { + std::vector starts = {100, 500}; + std::vector tmrcas = {10.0, 20.0}; + std::vector targets = {3, 7}; + std::vector mismatches = {2}; // mismatch at position index 2 + + ThreadingInstruction ti(starts, tmrcas, targets, mismatches); + std::vector positions = {100, 200, 300, 400, 500, 600}; + + ThreadingInstructionIterator iter(ti, positions); + CHECK(iter.current_target == 3); + CHECK(iter.current_tmrca == 10.0); + + // Advance past second segment start + iter.increment_site(500); + CHECK(iter.current_target == 7); + CHECK(iter.current_tmrca == 20.0); +} + +TEST_CASE("ThreadingInstructionIterator mismatch tracking") { + std::vector starts = {100}; + std::vector tmrcas = {10.0}; + std::vector targets = {0}; + std::vector mismatches = {2}; // mismatch at site index 2 + + ThreadingInstruction ti(starts, tmrcas, targets, mismatches); + std::vector positions = {100, 200, 300, 400, 500}; + + ThreadingInstructionIterator iter(ti, positions); + + iter.increment_site(100); + CHECK(iter.is_mismatch == false); + + iter.increment_site(200); + CHECK(iter.is_mismatch == false); + + // Position 300 = positions[2] = the mismatch site + iter.increment_site(300); + CHECK(iter.is_mismatch == true); + + iter.increment_site(400); + CHECK(iter.is_mismatch == false); +} diff --git a/test/test_vcf.py b/test/test_vcf.py new file mode 100644 index 0000000..284787a --- /dev/null +++ b/test/test_vcf.py @@ -0,0 +1,338 @@ +""" +Tests for the `threads vcf` command: GenotypeIterator, VCFWriter, and threads_to_vcf. +""" +import ctypes +import os +import re +import sys +import tempfile + +import numpy as np +import pytest + +# Flush all C stdio streams (needed to capture C++ stdout from VCFWriter) +_libc = ctypes.CDLL(None) +def _flush_c_stdout(): + _libc.fflush(None) + +from pathlib import Path + +# --------------------------------------------------------------------------- +# Paths +# --------------------------------------------------------------------------- +TEST_DATA = Path(__file__).parent / "data" +THREADS_FIT = TEST_DATA / "expected_infer_fit_to_data_snapshot.threads" +THREADS_NO_FIT = TEST_DATA / "expected_infer_snapshot.threads" +PANEL_PGEN = TEST_DATA / "panel.pgen" +PANEL_PVAR = TEST_DATA / "panel.pvar" +PANEL_PSAM = TEST_DATA / "panel.psam" + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def _load_instructions(threads_path): + from threads_arg.serialization import load_instructions + return load_instructions(str(threads_path)) + + +def _load_sample_names(): + names = [] + with open(PANEL_PSAM) as f: + for line in f: + if line.startswith("#"): + continue + names.append(line.strip().split()[0]) + return names + + +def _capture_vcf_output(instructions, variant_metadata, sample_names): + """Run VCFWriter.write_vcf() and capture C++ stdout via fd redirect.""" + from threads_arg import VCFWriter + + with tempfile.NamedTemporaryFile(mode="w+", suffix=".vcf", delete=False) as tmp: + tmpname = tmp.name + + try: + old_fd = os.dup(1) + new_fd = os.open(tmpname, os.O_WRONLY | os.O_TRUNC | os.O_CREAT) + os.dup2(new_fd, 1) + os.close(new_fd) + + writer = VCFWriter(instructions) + writer.set_chrom(variant_metadata["CHROM"].astype(str)) + writer.set_pos(variant_metadata["POS"].astype(str)) + writer.set_id(variant_metadata["ID"].astype(str)) + writer.set_ref(variant_metadata["REF"].astype(str)) + writer.set_alt(variant_metadata["ALT"].astype(str)) + writer.set_qual(variant_metadata["QUAL"].astype(str)) + writer.set_filter(variant_metadata["FILTER"].astype(str)) + writer.set_sample_names(sample_names) + writer.write_vcf() + + _flush_c_stdout() + sys.stdout.flush() + os.dup2(old_fd, 1) + os.close(old_fd) + + with open(tmpname) as f: + return f.readlines() + finally: + os.unlink(tmpname) + + +@pytest.fixture(scope="module") +def fit_instructions(): + return _load_instructions(THREADS_FIT) + + +@pytest.fixture(scope="module") +def nofit_instructions(): + return _load_instructions(THREADS_NO_FIT) + + +@pytest.fixture(scope="module") +def variant_metadata(): + from threads_arg.utils import read_variant_metadata + return read_variant_metadata(str(PANEL_PGEN)) + + +@pytest.fixture(scope="module") +def sample_names(): + return _load_sample_names() + + +@pytest.fixture(scope="module") +def panel_genotypes(): + """All genotypes from panel.pgen as (n_sites, n_haps) int32 array.""" + from threads_arg.utils import read_all_genotypes + return read_all_genotypes(str(PANEL_PGEN)) + + +@pytest.fixture(scope="module") +def vcf_lines(fit_instructions, variant_metadata, sample_names): + return _capture_vcf_output(fit_instructions, variant_metadata, sample_names) + + +# =================================================================== +# GenotypeIterator unit tests +# =================================================================== +class TestGenotypeIterator: + + def test_genotype_values_biallelic(self, fit_instructions): + """All genotype values should be 0 or 1.""" + from threads_arg import GenotypeIterator + gi = GenotypeIterator(fit_instructions) + for _ in range(50): + g = gi.next_genotype() + assert set(g).issubset({0, 1}), f"non-biallelic values: {set(g) - {0, 1}}" + + def test_genotype_length_matches_num_samples(self, fit_instructions): + """Each genotype vector should have length == num_samples (haploid).""" + from threads_arg import GenotypeIterator + gi = GenotypeIterator(fit_instructions) + g = gi.next_genotype() + assert len(g) == fit_instructions.num_samples + + def test_total_sites(self, fit_instructions): + """Iterator should yield exactly num_sites genotypes.""" + from threads_arg import GenotypeIterator + gi = GenotypeIterator(fit_instructions) + count = 0 + while gi.has_next_genotype(): + gi.next_genotype() + count += 1 + assert count == fit_instructions.num_sites + + def test_round_trip_fit_to_data(self, fit_instructions, panel_genotypes): + """fit_to_data .threads should exactly reconstruct panel.pgen genotypes.""" + from threads_arg import GenotypeIterator + gi = GenotypeIterator(fit_instructions) + for site_idx in range(fit_instructions.num_sites): + g = np.array(gi.next_genotype()) + np.testing.assert_array_equal( + g, panel_genotypes[site_idx], + err_msg=f"mismatch at site {site_idx}" + ) + + def test_round_trip_no_fit(self, nofit_instructions, panel_genotypes): + """non-fit_to_data .threads: check genotype reconstruction.""" + from threads_arg import GenotypeIterator + gi = GenotypeIterator(nofit_instructions) + mismatches = 0 + for site_idx in range(nofit_instructions.num_sites): + g = np.array(gi.next_genotype()) + if not np.array_equal(g, panel_genotypes[site_idx]): + mismatches += 1 + # This snapshot also has zero mismatches empirically + assert mismatches == 0, f"{mismatches}/{nofit_instructions.num_sites} sites differ" + + def test_allele_frequency_reasonable(self, fit_instructions, panel_genotypes): + """Allele frequencies from iterator should match panel.pgen.""" + from threads_arg import GenotypeIterator + gi = GenotypeIterator(fit_instructions) + for site_idx in range(min(100, fit_instructions.num_sites)): + g = np.array(gi.next_genotype()) + assert np.mean(g) == pytest.approx( + np.mean(panel_genotypes[site_idx]), abs=1e-10 + ) + + +# =================================================================== +# VCFWriter format tests +# =================================================================== +class TestVCFWriterFormat: + + def test_header_starts_vcf42(self, vcf_lines): + assert vcf_lines[0].strip() == "##fileformat=VCFv4.2" + + def test_has_source_line(self, vcf_lines): + assert any(l.startswith("##source=") for l in vcf_lines) + + def test_has_contig_line(self, vcf_lines): + assert any(l.startswith("##contig=") for l in vcf_lines) + + def test_chrom_header_present(self, vcf_lines): + header_lines = [l for l in vcf_lines if l.startswith("#CHROM")] + assert len(header_lines) == 1 + + def test_chrom_header_columns(self, vcf_lines, sample_names): + for l in vcf_lines: + if l.startswith("#CHROM"): + fields = l.strip().split("\t") + fixed = fields[:9] + assert fixed == [ + "#CHROM", "POS", "ID", "REF", "ALT", + "QUAL", "FILTER", "INFO", "FORMAT" + ] + # Sample columns + sample_cols = fields[9:] + assert len(sample_cols) == len(sample_names) + assert sample_cols[0] == sample_names[0] + assert sample_cols[-1] == sample_names[-1] + break + + def test_data_line_count(self, vcf_lines, fit_instructions): + """Number of data lines should equal number of sites.""" + data_lines = [l for l in vcf_lines if not l.startswith("#")] + assert len(data_lines) == fit_instructions.num_sites + + def test_format_field_is_gt(self, vcf_lines): + """FORMAT column should be 'GT'.""" + for l in vcf_lines: + if not l.startswith("#"): + fields = l.split("\t") + assert fields[8] == "GT" + break + + def test_genotype_format_phased(self, vcf_lines): + """Genotype fields should match pattern [01]|[01].""" + gt_pattern = re.compile(r"^[01]\|[01]$") + for l in vcf_lines: + if not l.startswith("#"): + fields = l.rstrip("\n").split("\t") + # Skip empty trailing field from known trailing-tab bug + gt_fields = [f for f in fields[9:] if f] + for gt in gt_fields[:20]: # spot-check first 20 + assert gt_pattern.match(gt), f"bad GT format: {gt!r}" + break + + def test_info_field_has_ns(self, vcf_lines, sample_names): + """INFO field should contain NS=.""" + expected_ns = f"NS={len(sample_names)}" + for l in vcf_lines: + if not l.startswith("#"): + fields = l.split("\t") + assert fields[7] == expected_ns + break + + def test_trailing_tab_known_bug(self, vcf_lines): + """Document: VCFWriter emits a trailing tab on each data line.""" + for l in vcf_lines: + if not l.startswith("#"): + # The raw line (before rstrip) should end with \t\n + assert l.endswith("\t\n") or l.endswith("\t"), \ + "trailing tab bug may have been fixed — update this test" + break + + +# =================================================================== +# threads_to_vcf integration tests +# =================================================================== +class TestThreadsToVcf: + + def test_missing_metadata_raises(self): + """Calling threads_to_vcf without external files on a .threads + that lacks embedded metadata should raise RuntimeError.""" + from threads_arg.threads_to_vcf import threads_to_vcf + with pytest.raises(RuntimeError, match="Unable to load sample information"): + threads_to_vcf(str(THREADS_FIT)) + + def test_with_external_pvar(self, fit_instructions, sample_names): + """threads_to_vcf with external .pvar and samples file should produce output.""" + from threads_arg.threads_to_vcf import threads_to_vcf + + with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as sf: + for name in sample_names: + sf.write(name + "\n") + samples_path = sf.name + + with tempfile.NamedTemporaryFile(mode="w+", suffix=".vcf", delete=False) as tmp: + tmpname = tmp.name + + try: + old_fd = os.dup(1) + new_fd = os.open(tmpname, os.O_WRONLY | os.O_TRUNC | os.O_CREAT) + os.dup2(new_fd, 1) + os.close(new_fd) + + threads_to_vcf( + str(THREADS_FIT), + samples=samples_path, + variants=str(PANEL_PVAR), + ) + + _flush_c_stdout() + sys.stdout.flush() + os.dup2(old_fd, 1) + os.close(old_fd) + + with open(tmpname) as f: + lines = f.readlines() + + assert len(lines) > 0, "no output produced" + data_lines = [l for l in lines if not l.startswith("#")] + assert len(data_lines) == fit_instructions.num_sites + finally: + os.unlink(tmpname) + os.unlink(samples_path) + + def test_genotype_consistency(self, variant_metadata, sample_names, panel_genotypes): + """Genotypes in VCF output should match panel.pgen.""" + instructions = _load_instructions(THREADS_FIT) + lines = _capture_vcf_output(instructions, variant_metadata, sample_names) + + n_diploid = len(sample_names) + # Check first 100 data lines + site_idx = 0 + for l in lines: + if l.startswith("#"): + continue + if site_idx >= 100: + break + fields = l.rstrip("\n").split("\t") + gt_fields = [f for f in fields[9:] if f] # skip trailing empty + assert len(gt_fields) == n_diploid + + # Reconstruct haploid genotypes from phased GTs + haps = np.zeros(2 * n_diploid, dtype=int) + for i, gt_str in enumerate(gt_fields): + a, b = gt_str.split("|") + haps[2 * i] = int(a) + haps[2 * i + 1] = int(b) + + np.testing.assert_array_equal( + haps, panel_genotypes[site_idx], + err_msg=f"VCF genotype mismatch at site {site_idx}" + ) + site_idx += 1 diff --git a/test/test_viterbi_lowmem.cpp b/test/test_viterbi_lowmem.cpp new file mode 100644 index 0000000..f777a2f --- /dev/null +++ b/test/test_viterbi_lowmem.cpp @@ -0,0 +1,246 @@ +// This file is part of the Threads software suite. +// Copyright (C) 2024-2025 Threads Developers. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU General Public License for more details. +// +// You should have received a copy of the GNU General Public License +// along with this program. If not, see . + +#include "ViterbiLowMem.hpp" + +#include +#include +#include +#include + +// === ViterbiPath tests === + +TEST_CASE("ViterbiPath construction and basic operations") { + ViterbiPath path(5); + CHECK(path.target_id == 5); + CHECK(path.size() == 0); + + path.append(0, 3); + path.append(10, 7); + CHECK(path.size() == 2); + CHECK(path.segment_starts[0] == 0); + CHECK(path.segment_starts[1] == 10); + CHECK(path.sample_ids[0] == 3); + CHECK(path.sample_ids[1] == 7); +} + +TEST_CASE("ViterbiPath reverse") { + ViterbiPath path(0); + path.append(0, 1); + path.append(5, 2); + path.append(10, 3); + + path.reverse(); + CHECK(path.segment_starts[0] == 10); + CHECK(path.segment_starts[1] == 5); + CHECK(path.segment_starts[2] == 0); + CHECK(path.sample_ids[0] == 3); + CHECK(path.sample_ids[1] == 2); + CHECK(path.sample_ids[2] == 1); +} + +TEST_CASE("ViterbiPath append with height and het_sites") { + ViterbiPath path(0); + std::vector hets1 = {2, 3}; + path.append(0, 1, 100.0, hets1); + std::vector hets2 = {7}; + path.append(5, 2, 200.0, hets2); + + CHECK(path.size() == 2); + CHECK(path.heights[0] == 100.0); + CHECK(path.heights[1] == 200.0); + CHECK(path.het_sites.size() == 3); + CHECK(path.het_sites[0] == 2); + CHECK(path.het_sites[1] == 3); + CHECK(path.het_sites[2] == 7); +} + +TEST_CASE("ViterbiPath append validates ordering") { + ViterbiPath path(0); + std::vector hets = {}; + path.append(10, 1, 100.0, hets); + + // Appending segment_start <= previous should throw + CHECK_THROWS(path.append(10, 2, 200.0, hets)); + CHECK_THROWS(path.append(5, 2, 200.0, hets)); +} + +TEST_CASE("ViterbiPath map_positions") { + ViterbiPath path(0); + path.append(0, 1); + path.append(3, 2); + path.append(7, 3); + + std::vector positions = {100, 200, 300, 400, 500, 600, 700, 800, 900, 1000}; + path.map_positions(positions); + + CHECK(path.bp_starts.size() == 3); + CHECK(path.bp_starts[0] == 100); + CHECK(path.bp_starts[1] == 400); + CHECK(path.bp_starts[2] == 800); +} + +TEST_CASE("ViterbiPath dump_data_in_range full") { + ViterbiPath path(0); + std::vector hets1 = {}; + path.append(0, 1, 10.0, hets1); + path.append(5, 2, 20.0, hets1); + path.append(10, 3, 30.0, hets1); + + std::vector positions = {100, 200, 300, 400, 500, 600, 700, 800, 900, 1000, 1100}; + path.map_positions(positions); + + auto [starts, ids, heights] = path.dump_data_in_range(-1, -1); + CHECK(starts.size() == 3); + CHECK(ids.size() == 3); + CHECK(heights.size() == 3); +} + +TEST_CASE("ViterbiPath dump_data_in_range subset") { + ViterbiPath path(0); + std::vector hets = {}; + path.append(0, 1, 10.0, hets); + path.append(3, 2, 20.0, hets); + path.append(6, 3, 30.0, hets); + + std::vector positions = {100, 200, 300, 400, 500, 600, 700, 800, 900}; + path.map_positions(positions); + + // Request range that covers only second segment + auto [starts, ids, heights] = path.dump_data_in_range(400, 600); + CHECK(starts.size() >= 1); + // The first returned start should be 400 + CHECK(starts[0] == 400); +} + +// === ViterbiState tests === + +TEST_CASE("ViterbiState construction") { + std::vector samples = {0, 1, 2}; + ViterbiState state(5, samples); + CHECK(state.target_id == 5); + CHECK(state.best_match == 0); + CHECK(state.sites_processed == 0); +} + +TEST_CASE("ViterbiState construction requires non-empty samples") { + std::vector empty_samples; + CHECK_THROWS(ViterbiState(0, empty_samples)); +} + +TEST_CASE("ViterbiState process_site basic") { + // 3 reference samples + target + std::vector samples = {0, 1, 2}; + ViterbiState state(3, samples); + + // Genotype vector: all samples + target + // sample 0=0, sample 1=1, sample 2=0, target=1 + std::vector geno = {0, 1, 0, 1}; + + double rho = 5.0; // recombination penalty + double rho_c = 0.01; // non-recombination penalty + double mu = 3.0; // mutation penalty + double mu_c = 0.01; // non-mutation penalty + + state.process_site(geno, rho, rho_c, mu, mu_c); + CHECK(state.sites_processed == 1); + + // Best match should be sample 1 (matches target allele) + CHECK(state.best_match == 1); +} + +TEST_CASE("ViterbiState process multiple sites and traceback") { + std::vector samples = {0, 1}; + ViterbiState state(2, samples); + + // Process several sites where sample 0 always matches target + for (int i = 0; i < 5; i++) { + std::vector geno = {1, 0, 1}; // sample0=1, sample1=0, target=1 + state.process_site(geno, 5.0, 0.01, 3.0, 0.01); + } + + CHECK(state.sites_processed == 5); + + auto path = state.traceback(); + CHECK(path.target_id == 2); + CHECK(path.size() >= 1); + // Best path should mostly copy sample 0 + CHECK(path.sample_ids[0] == 0); +} + +TEST_CASE("ViterbiState prune reduces branch count") { + std::vector samples = {0, 1, 2, 3}; + ViterbiState state(4, samples); + + // Process enough sites to create branches + for (int i = 0; i < 20; i++) { + std::vector geno; + for (int s = 0; s < 5; s++) { + geno.push_back(i % 2 == 0 ? s % 2 : (s + 1) % 2); + } + state.process_site(geno, 2.0, 0.5, 1.5, 0.1); + } + + int branches_before = state.count_branches(); + state.prune(); + int branches_after = state.count_branches(); + + // Prune should not increase branch count + CHECK(branches_after <= branches_before); + // Should still have at least as many branches as samples + CHECK(branches_after >= static_cast(samples.size())); +} + +TEST_CASE("ViterbiState traceback produces valid path") { + std::vector samples = {0, 1}; + ViterbiState state(2, samples); + + // Alternating genotypes to force recombinations + for (int i = 0; i < 10; i++) { + std::vector geno; + if (i < 5) { + geno = {1, 0, 1}; // target matches sample 0 + } else { + geno = {0, 1, 1}; // target matches sample 1 + } + state.process_site(geno, 1.0, 0.5, 2.0, 0.01); + } + + auto path = state.traceback(); + CHECK(path.size() >= 1); + // Segments should be ordered + for (int i = 1; i < path.size(); i++) { + CHECK(path.segment_starts[i] > path.segment_starts[i - 1]); + } +} + +TEST_CASE("ViterbiState set_samples updates candidate set") { + std::vector samples = {0, 1, 2, 3, 4}; + ViterbiState state(5, samples); + + // Process a few sites + for (int i = 0; i < 3; i++) { + std::vector geno = {1, 0, 1, 0, 1, 1}; + state.process_site(geno, 2.0, 0.5, 1.5, 0.1); + } + + // Reduce to subset + std::unordered_set new_samples = {0, 2}; + state.set_samples(new_samples); + + // Sample_ids should now contain new_samples + best_match + CHECK(state.sample_ids.size() <= 4); // at most new_samples + best_match + margin +}