From c491d8b1f976cf67dce5086c4192dbefabbeb82f Mon Sep 17 00:00:00 2001 From: lkirk Date: Wed, 18 Dec 2024 14:34:46 -0600 Subject: [PATCH 01/70] C and Python API for two-way two-locus stats This PR implements the C and Python API for computing two-way two-locus statistics. The algorithm is identical to the python version, except during testing I uncovered a small issue with normalisation. We need to handle the case where sample sets are of different sizes. The fix for this was to average the normalisation factor for each sample set. Test coverage has been added to cover C, low-level python and some high-level tests. --- python/tests/test_ld_matrix.py | 53 ++++++++++++++++++++++++++++++++-- 1 file changed, 51 insertions(+), 2 deletions(-) diff --git a/python/tests/test_ld_matrix.py b/python/tests/test_ld_matrix.py index a3eece1e2d..083b467b48 100644 --- a/python/tests/test_ld_matrix.py +++ b/python/tests/test_ld_matrix.py @@ -277,7 +277,6 @@ def norm_hap_weighted_ij( wAB_i = hap_weights[0, i] wAB_j = hap_weights[0, j] result[k] = (wAB_i + wAB_j) / (ni + nj) - # result[k] = (wAB_i / ni / 2) + (wAB_j / nj / 2) def norm_total_weighted( @@ -1053,8 +1052,10 @@ def r2_ij_summary_func( D_j = pAB - pA * pB denom_j = np.sqrt(pA * (1 - pA) * pB * (1 - pB)) + p_A = (w_A_i + w_A_j) / (ni + nj) + p_B = (w_B_i + w_B_j) / (ni + nj) with suppress_overflow_div0_warning(): - result[k] = (D_i * D_j) / (denom_i * denom_j) + result[k] = (D_i * D_j) / (p_A * (1 - p_A) * p_B * (1 - p_B)) def D_summary_func( @@ -2318,12 +2319,20 @@ def test_two_way_site_ld_matrix(ts, stat): ( correlated, (np.array([0, 1, 2, 3, 4, 5]), np.array([6, 7, 8, 9, 10])), +<<<<<<< HEAD np.float64(1.0), +======= + np.float64(0.9708352229780801), +>>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ( correlated, (np.array([0, 1, 2, 3, 4, 5]), np.array([6, 7, 8, 9])), +<<<<<<< HEAD np.float64(1.0), +======= + np.float64(0.9526958931720837), +>>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ( correlated, @@ -2333,12 +2342,20 @@ def test_two_way_site_ld_matrix(ts, stat): ( correlated, (np.array([0, 1, 2, 3, 4, 5]), np.array([6, 7])), +<<<<<<< HEAD np.float64(np.nan), +======= + np.float64(0.7585185185185186), +>>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ( correlated, (np.array([0, 1, 2, 3, 4, 5]), np.array([6])), +<<<<<<< HEAD np.float64(np.nan), +======= + np.float64(0.0), +>>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ( anticorrelated := np.array( @@ -2356,37 +2373,65 @@ def test_two_way_site_ld_matrix(ts, stat): ( anticorrelated, (np.array([0, 2, 4, 6, 8, 10, 12, 14]), np.array([1, 3, 5, 7, 9, 11, 13])), +<<<<<<< HEAD np.float64(1.0), +======= + np.float64(0.9798566895766568), +>>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ( anticorrelated, (np.array([0, 2, 4, 6, 8, 10, 12, 14]), np.array([1, 3, 5, 7, 9, 11])), +<<<<<<< HEAD np.float64(np.nan), +======= + np.float64(0.8574999999999999), +>>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ( anticorrelated, (np.array([0, 2, 4, 6, 8, 10, 12, 14]), np.array([1, 3, 5, 7, 9])), +<<<<<<< HEAD np.float64(np.nan), +======= + np.float64(0.8299777777777777), +>>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ( anticorrelated, (np.array([0, 2, 4, 6, 8, 10, 12, 14]), np.array([1, 3, 5, 7])), +<<<<<<< HEAD np.float64(np.nan), +======= + np.float64(0.6328124999999999), +>>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ( anticorrelated, (np.array([0, 2, 4, 6, 8, 10, 12, 14]), np.array([1, 3, 5])), +<<<<<<< HEAD np.float64(np.nan), +======= + np.float64(0.57179616638322), +>>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ( anticorrelated, (np.array([0, 2, 4, 6, 8, 10, 12, 14]), np.array([1, 3])), +<<<<<<< HEAD np.float64(np.nan), +======= + np.float64(0.0), +>>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ( anticorrelated, (np.array([0, 2, 4, 6, 8, 10, 12, 14]), np.array([1])), +<<<<<<< HEAD np.float64(np.nan), +======= + np.float64(0.0), +>>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ], ) @@ -2405,4 +2450,8 @@ def test_multipopulation_r2_varying_unequal_set_sizes(genotypes, sample_sets, ex r2_ij_summary_func(state_dim, state, 1, result[i, j], params) norm_hap_weighted_ij(1, state, max(a) + 1, max(b) + 1, norm[i, j], params) +<<<<<<< HEAD np.testing.assert_allclose((result * norm).sum(), expected) +======= + np.testing.assert_allclose(expected, (result * norm).sum()) +>>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) From d9991c9f97a88cc6f4a7c86d0c9f7e27d2cfc3a7 Mon Sep 17 00:00:00 2001 From: lkirk Date: Wed, 18 Dec 2024 14:34:46 -0600 Subject: [PATCH 02/70] Two-way two-locus statistics (python proto) This PR implements two-way LD statistics, specified between sample sets. During the development of this functionality, a number of issues with the designation of state_dims/result_dims were discovered. These have been resolved, testing clean for existing code and providing the proper behavior for this new code. The mechanism by which users will specify a multi-population (or two-way) statistic is by providing the `index` argument. This helps us avoid creating another `ld_matrix` method for the TreeSequence object. In other words, for a one-way statistic, a user would specify: ```python ts.ld_matrix(stat="D2", sample_sets=[[ss1, ss2]]) ``` Which would output a 3D ndarray containing one LD matrix per sample set. ```python ts.ld_matrix(stat="D2", sample_sets=[[ss1, ss2]], indexes=[(0, 1)]) ``` Which would output a 2D ndarray containing one LD matrix for the index pair. This would use our `D2_ij_summary_func`, instead of the `D2_summary_func`. Finally, if a user provided ```python ts.ld_matrix(stat="D2", sample_sets=[[ss1, ss2]], indexes=[(0, 1), (1, 1)]) ``` We would output a 3D ndarray containing one LD matrix _per_ index pair provided. Since these are two-way statistics, the indexes must be length 2. We plan on enabling users to implement k-way via a "general_stat" api. We did not implement anything more than two-way statistics here because of the combinatoric explosion of logic required for indexes > 2. I added some basic tests to demonstrate that things were working properly. If we compute two-way statistics on identical sample sets, they should be equal to the one-way statistics. Unfortunately, this does not apply to unbiased statistics, which I've validated manually. I've also cleaned up the docstrings a bit and fixed a bug with the D_prime statistic, which should not be weighted by haplotype frequency. --- python/tests/test_ld_matrix.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/python/tests/test_ld_matrix.py b/python/tests/test_ld_matrix.py index 083b467b48..5e96f0a718 100644 --- a/python/tests/test_ld_matrix.py +++ b/python/tests/test_ld_matrix.py @@ -24,6 +24,7 @@ """ import contextlib import io +import itertools from dataclasses import dataclass from itertools import combinations_with_replacement from itertools import permutations @@ -842,6 +843,7 @@ def two_locus_count_stat( norm_func, polarised, mode, + result_dim, sites=None, positions=None, sample_sets=None, @@ -857,6 +859,9 @@ def two_locus_count_stat( :param polarised: If true, skip the computation of the statistic for the ancestral state. :param mode: Whether or not to compute "site" or "branch" statistics. + :param result_dim: The dimensions of the output array. For one-way stats, + this will be the number of sample sets. For two-way stats, + the number of index tuples. :param sites: List of two lists containing [row_sites, column_sites]. :param positions: List of two lists containing [row_positions, col_positions], which are genomic positions to compute LD on. @@ -1455,6 +1460,7 @@ def ld_matrix( NORM_METHOD[summary_func], POLARIZATION[summary_func], mode, + result_dim, sites=sites, positions=positions, indexes=indexes, From 96525ee24d7b9696d762c0632f22a35d064a439e Mon Sep 17 00:00:00 2001 From: lkirk Date: Thu, 29 May 2025 10:22:53 -0500 Subject: [PATCH 03/70] initial implementation; need to work on output dims/scalar sample sets and indexes --- c/tskit/trees.c | 3 +++ python/tskit/trees.py | 18 +++++++++--------- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index adc8034d4e..dea05a1d97 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -2294,6 +2294,9 @@ get_all_samples_bits(tsk_bit_array_t *all_samples, tsk_size_t n) } } +typedef int norm_func_t(tsk_size_t result_dim, const double *hap_weights, tsk_size_t n_a, + tsk_size_t n_b, double *result, void *params); + static int compute_general_two_site_stat_result(const tsk_bit_array_t *site_a_state, const tsk_bit_array_t *site_b_state, tsk_size_t num_a_alleles, diff --git a/python/tskit/trees.py b/python/tskit/trees.py index 77ccbf4838..14a50d2fde 100644 --- a/python/tskit/trees.py +++ b/python/tskit/trees.py @@ -23,6 +23,7 @@ """ Module responsible for managing trees and tree sequences. """ + from __future__ import annotations import base64 @@ -686,8 +687,7 @@ def __init__( options = 0 if sample_counts is not None: warnings.warn( - "The sample_counts option is not supported since 0.2.4 " - "and is ignored", + "The sample_counts option is not supported since 0.2.4 and is ignored", RuntimeWarning, stacklevel=4, ) @@ -6840,7 +6840,7 @@ def to_macs(self): bytes_genotypes[:] = lookup[variant.genotypes] genotypes = bytes_genotypes.tobytes().decode() output.append( - f"SITE:\t{variant.index}\t{variant.position / m}\t0.0\t" f"{genotypes}" + f"SITE:\t{variant.index}\t{variant.position / m}\t0.0\t{genotypes}" ) return "\n".join(output) + "\n" @@ -9143,9 +9143,9 @@ def pca( if time_windows is None: tree_sequence_low, tree_sequence_high = None, self else: - assert ( - time_windows[0] < time_windows[1] - ), "The second argument should be larger." + assert time_windows[0] < time_windows[1], ( + "The second argument should be larger." + ) tree_sequence_low, tree_sequence_high = ( self.decapitate(time_windows[0]), self.decapitate(time_windows[1]), @@ -9213,9 +9213,9 @@ def _rand_pow_range_finder( """ Algorithm 9 in https://arxiv.org/pdf/2002.01387 """ - assert ( - num_vectors >= rank > 0 - ), "num_vectors should not be smaller than rank" + assert num_vectors >= rank > 0, ( + "num_vectors should not be smaller than rank" + ) for _ in range(depth): Q = np.linalg.qr(Q)[0] Q = operator(Q) From bef9b8c71a09d4e1066a3945cffbfcf765f001b8 Mon Sep 17 00:00:00 2001 From: lkirk Date: Mon, 2 Jun 2025 02:45:34 -0500 Subject: [PATCH 04/70] get rid of disjoint check --- python/tests/test_ld_matrix.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/tests/test_ld_matrix.py b/python/tests/test_ld_matrix.py index 5e96f0a718..95cb48273d 100644 --- a/python/tests/test_ld_matrix.py +++ b/python/tests/test_ld_matrix.py @@ -22,6 +22,7 @@ """ Test cases for two-locus statistics """ + import contextlib import io import itertools From b3f61b99e2586a6653c72da57257a7a62dbc5dcb Mon Sep 17 00:00:00 2001 From: lkirk Date: Thu, 17 Jul 2025 02:34:24 -0500 Subject: [PATCH 05/70] fix normalisation and a small indexing bug. dimension tests --- c/tskit/trees.c | 9 +++------ python/tests/test_ld_matrix.py | 11 ++++------- 2 files changed, 7 insertions(+), 13 deletions(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index dea05a1d97..6e692919a3 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -3244,17 +3244,12 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl tsk_memset(&sample_sets_bits, 0, sizeof(sample_sets_bits)); - // We do not support two-locus node stats - if (!!(options & TSK_STAT_NODE)) { - ret = tsk_trace_error(TSK_ERR_UNSUPPORTED_STAT_MODE); - goto out; - } // If no mode is specified, we default to site mode if (!(stat_site || stat_branch)) { stat_site = true; } // It's an error to specify more than one mode - if (stat_site + stat_branch > 1) { + if (stat_site + stat_branch + stat_node > 1) { ret = tsk_trace_error(TSK_ERR_MULTIPLE_STAT_MODES); goto out; } @@ -3299,6 +3294,8 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl ret = tsk_treeseq_two_branch_count_stat(self, state_dim, &sample_sets_bits, result_dim, f, &f_params, norm_f, out_rows, row_positions, out_cols, col_positions, options, result); + } else { + ret = tsk_trace_error(TSK_ERR_UNSUPPORTED_STAT_MODE); } out: diff --git a/python/tests/test_ld_matrix.py b/python/tests/test_ld_matrix.py index 95cb48273d..3142c349f1 100644 --- a/python/tests/test_ld_matrix.py +++ b/python/tests/test_ld_matrix.py @@ -22,10 +22,8 @@ """ Test cases for two-locus statistics """ - import contextlib import io -import itertools from dataclasses import dataclass from itertools import combinations_with_replacement from itertools import permutations @@ -592,6 +590,9 @@ def compute_general_two_site_stat_result( for k in range(result_dim): result[k] += result_tmp[k] * norm[k] + # for k in range(result_dim): + # print(mut_a, mut_b, k, weights[0, k], weights[1, k], weights[2, k], sep="\t") + def two_site_count_stat( ts: tskit.TreeSequence, @@ -844,7 +845,6 @@ def two_locus_count_stat( norm_func, polarised, mode, - result_dim, sites=None, positions=None, sample_sets=None, @@ -860,9 +860,6 @@ def two_locus_count_stat( :param polarised: If true, skip the computation of the statistic for the ancestral state. :param mode: Whether or not to compute "site" or "branch" statistics. - :param result_dim: The dimensions of the output array. For one-way stats, - this will be the number of sample sets. For two-way stats, - the number of index tuples. :param sites: List of two lists containing [row_sites, column_sites]. :param positions: List of two lists containing [row_positions, col_positions], which are genomic positions to compute LD on. @@ -1461,7 +1458,6 @@ def ld_matrix( NORM_METHOD[summary_func], POLARIZATION[summary_func], mode, - result_dim, sites=sites, positions=positions, indexes=indexes, @@ -1784,6 +1780,7 @@ def test_ld_empty_examples(ts): def test_input_validation(): + # TODO ts = get_paper_ex_ts() with pytest.raises(ValueError, match="Unknown two-locus statistic"): ts.ld_matrix(stat="bad_stat") From 0ad3e67611473613b4cde049b693b18aee4a056f Mon Sep 17 00:00:00 2001 From: lkirk Date: Thu, 17 Jul 2025 02:47:06 -0500 Subject: [PATCH 06/70] lints --- python/tests/test_ld_matrix.py | 3 --- python/tskit/trees.py | 1 - 2 files changed, 4 deletions(-) diff --git a/python/tests/test_ld_matrix.py b/python/tests/test_ld_matrix.py index 3142c349f1..90492123c2 100644 --- a/python/tests/test_ld_matrix.py +++ b/python/tests/test_ld_matrix.py @@ -590,9 +590,6 @@ def compute_general_two_site_stat_result( for k in range(result_dim): result[k] += result_tmp[k] * norm[k] - # for k in range(result_dim): - # print(mut_a, mut_b, k, weights[0, k], weights[1, k], weights[2, k], sep="\t") - def two_site_count_stat( ts: tskit.TreeSequence, diff --git a/python/tskit/trees.py b/python/tskit/trees.py index 14a50d2fde..f6c7d51007 100644 --- a/python/tskit/trees.py +++ b/python/tskit/trees.py @@ -23,7 +23,6 @@ """ Module responsible for managing trees and tree sequences. """ - from __future__ import annotations import base64 From 698a1f7aea976976a5fa5fd627da18a25928eca9 Mon Sep 17 00:00:00 2001 From: lkirk Date: Thu, 17 Jul 2025 14:41:41 -0500 Subject: [PATCH 07/70] fix logic error for mode testing --- c/tskit/trees.c | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index 6e692919a3..5d55689d94 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -3245,7 +3245,7 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl tsk_memset(&sample_sets_bits, 0, sizeof(sample_sets_bits)); // If no mode is specified, we default to site mode - if (!(stat_site || stat_branch)) { + if (!(stat_site || stat_branch || stat_node)) { stat_site = true; } // It's an error to specify more than one mode From 13c93989b084d26c1f97986a40395b1bd2ddf73a Mon Sep 17 00:00:00 2001 From: lkirk Date: Mon, 21 Jul 2025 20:26:21 -0500 Subject: [PATCH 08/70] fix norm_hap_weighted_ij. needs to get wAB instead of using indices From 5788f70bdc675e77f45c0c90ac58b59574e1fb8e Mon Sep 17 00:00:00 2001 From: lkirk Date: Tue, 22 Jul 2025 22:13:52 -0500 Subject: [PATCH 09/70] Create tests with pragmatic test coverage I had to make tsk_treeseq_two_locus_count_stat public to test it. This is a precursor to creating the general two locus stat api. --- c/tskit/trees.c | 3 --- 1 file changed, 3 deletions(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index 5d55689d94..cd76ebd2d6 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -2294,9 +2294,6 @@ get_all_samples_bits(tsk_bit_array_t *all_samples, tsk_size_t n) } } -typedef int norm_func_t(tsk_size_t result_dim, const double *hap_weights, tsk_size_t n_a, - tsk_size_t n_b, double *result, void *params); - static int compute_general_two_site_stat_result(const tsk_bit_array_t *site_a_state, const tsk_bit_array_t *site_b_state, tsk_size_t num_a_alleles, From 0e14329281799f0a35bf4d18358abd8f3b065264 Mon Sep 17 00:00:00 2001 From: lkirk Date: Tue, 22 Jul 2025 22:16:01 -0500 Subject: [PATCH 10/70] trivial cleanup --- python/tests/test_ld_matrix.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/tests/test_ld_matrix.py b/python/tests/test_ld_matrix.py index 90492123c2..10cb6a38d3 100644 --- a/python/tests/test_ld_matrix.py +++ b/python/tests/test_ld_matrix.py @@ -22,6 +22,7 @@ """ Test cases for two-locus statistics """ + import contextlib import io from dataclasses import dataclass From 48b9b36715195cdd8fa1988bbcd6f3a95fbae313 Mon Sep 17 00:00:00 2001 From: lkirk Date: Tue, 22 Jul 2025 22:18:06 -0500 Subject: [PATCH 11/70] small lint grr --- python/tests/test_ld_matrix.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/tests/test_ld_matrix.py b/python/tests/test_ld_matrix.py index 10cb6a38d3..90492123c2 100644 --- a/python/tests/test_ld_matrix.py +++ b/python/tests/test_ld_matrix.py @@ -22,7 +22,6 @@ """ Test cases for two-locus statistics """ - import contextlib import io from dataclasses import dataclass From 8a475219fef097e1aa833e732e04d01f9dc979d7 Mon Sep 17 00:00:00 2001 From: lkirk Date: Wed, 23 Jul 2025 05:46:15 -0500 Subject: [PATCH 12/70] lowlevel test coverage --- python/tests/test_lowlevel.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/python/tests/test_lowlevel.py b/python/tests/test_lowlevel.py index d6416c22b8..8f9cbb83e3 100644 --- a/python/tests/test_lowlevel.py +++ b/python/tests/test_lowlevel.py @@ -23,6 +23,7 @@ """ Test cases for the low level C interface to tskit. """ + import collections import gc import inspect @@ -3242,7 +3243,10 @@ def f(x): for bad_array in [[1, 1], range(10)]: with pytest.raises(ValueError): ts.general_stat( - W, lambda x: bad_array, 1, ts.get_breakpoints() # noqa:B023 + W, + lambda x: bad_array, + 1, + ts.get_breakpoints(), # noqa:B023 ) with pytest.raises(ValueError): ts.general_stat(W, lambda x: [1], 2, ts.get_breakpoints()) @@ -3251,7 +3255,10 @@ def f(x): for bad_array in [["sdf"], 0, "w4", None]: with pytest.raises(ValueError): ts.general_stat( - W, lambda x: bad_array, 1, ts.get_breakpoints() # noqa:B023 + W, + lambda x: bad_array, + 1, + ts.get_breakpoints(), # noqa:B023 ) From 5502e9d8404c4bd3e8c507b2210bfbcff7460e8f Mon Sep 17 00:00:00 2001 From: lkirk Date: Wed, 23 Jul 2025 05:57:29 -0500 Subject: [PATCH 13/70] concat is not compatible with numpy<2 From e3bb8f50a89bfbf4c33a982505358fe7783fd219 Mon Sep 17 00:00:00 2001 From: lkirk Date: Wed, 23 Jul 2025 11:10:57 -0500 Subject: [PATCH 14/70] remove unreachable code path (we switch on the presence of indexes) From bdcdeafa721b918af4ad62f5173792059f5c3c12 Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 25 Jul 2025 17:42:13 -0500 Subject: [PATCH 15/70] just error early if node stat requested --- c/tskit/trees.c | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index cd76ebd2d6..adc8034d4e 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -3241,12 +3241,17 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl tsk_memset(&sample_sets_bits, 0, sizeof(sample_sets_bits)); + // We do not support two-locus node stats + if (!!(options & TSK_STAT_NODE)) { + ret = tsk_trace_error(TSK_ERR_UNSUPPORTED_STAT_MODE); + goto out; + } // If no mode is specified, we default to site mode - if (!(stat_site || stat_branch || stat_node)) { + if (!(stat_site || stat_branch)) { stat_site = true; } // It's an error to specify more than one mode - if (stat_site + stat_branch + stat_node > 1) { + if (stat_site + stat_branch > 1) { ret = tsk_trace_error(TSK_ERR_MULTIPLE_STAT_MODES); goto out; } @@ -3291,8 +3296,6 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl ret = tsk_treeseq_two_branch_count_stat(self, state_dim, &sample_sets_bits, result_dim, f, &f_params, norm_f, out_rows, row_positions, out_cols, col_positions, options, result); - } else { - ret = tsk_trace_error(TSK_ERR_UNSUPPORTED_STAT_MODE); } out: From 5328faaff97b420de367931d25d98b70a72fdb28 Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 25 Jul 2025 17:49:04 -0500 Subject: [PATCH 16/70] better docstring for norm func From a39dc70dc88d9aa80ae2b5421be1e5e8e883364c Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 25 Jul 2025 17:55:22 -0500 Subject: [PATCH 17/70] remove todo --- python/tests/test_ld_matrix.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/tests/test_ld_matrix.py b/python/tests/test_ld_matrix.py index 90492123c2..083b467b48 100644 --- a/python/tests/test_ld_matrix.py +++ b/python/tests/test_ld_matrix.py @@ -1777,7 +1777,6 @@ def test_ld_empty_examples(ts): def test_input_validation(): - # TODO ts = get_paper_ex_ts() with pytest.raises(ValueError, match="Unknown two-locus statistic"): ts.ld_matrix(stat="bad_stat") From 8f6f89b2544591b9717585993104d775e0e856ba Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 25 Jul 2025 17:57:17 -0500 Subject: [PATCH 18/70] remove todo From d31ca45013944bf22298525fc1b00041e8895214 Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 25 Jul 2025 17:59:00 -0500 Subject: [PATCH 19/70] remove todo From 87d92aa41f006f6fab8e6c27327bb4890a49ded9 Mon Sep 17 00:00:00 2001 From: lkirk Date: Thu, 24 Jul 2025 01:11:59 -0500 Subject: [PATCH 20/70] move allocations out of the hot path; create a structure to hold the intermediate data --- c/tskit/trees.c | 134 +++++++++++++++++++++++++++++++----------------- 1 file changed, 87 insertions(+), 47 deletions(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index adc8034d4e..b832ae7515 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -2294,61 +2294,98 @@ get_all_samples_bits(tsk_bit_array_t *all_samples, tsk_size_t n) } } +typedef struct { + double *weights; + double *norm; + double *result_tmp; + tsk_bit_array_t AB_samples; + tsk_bit_array_t ss_A_samples; + tsk_bit_array_t ss_B_samples; + tsk_bit_array_t ss_AB_samples; +} two_locus_work_t; + static int -compute_general_two_site_stat_result(const tsk_bit_array_t *site_a_state, - const tsk_bit_array_t *site_b_state, tsk_size_t num_a_alleles, - tsk_size_t num_b_alleles, tsk_size_t num_samples, tsk_size_t state_dim, - const tsk_bit_array_t *sample_sets, tsk_size_t result_dim, general_stat_func_t *f, - sample_count_stat_params_t *f_params, norm_func_t *norm_f, bool polarised, - double *result) +two_locus_work_init(tsk_size_t max_alleles, tsk_size_t result_dim, tsk_size_t state_dim, + tsk_size_t num_samples, two_locus_work_t *out) { int ret = 0; - tsk_bit_array_t A_samples, B_samples; - // ss_ prefix refers to a sample set - tsk_bit_array_t ss_row; - tsk_bit_array_t ss_A_samples, ss_B_samples, ss_AB_samples, AB_samples; - // Sample sets and b sites are rows, a sites are columns - // b1 b2 b3 - // a1 [s1, s2, s3] [s1, s2, s3] [s1, s2, s3] - // a2 [s1, s2, s3] [s1, s2, s3] [s1, s2, s3] - // a3 [s1, s2, s3] [s1, s2, s3] [s1, s2, s3] - tsk_size_t k, mut_a, mut_b; - tsk_size_t result_row_len = num_b_alleles * result_dim; - tsk_size_t w_A = 0, w_B = 0, w_AB = 0; - uint8_t polarised_val = polarised ? 1 : 0; - double *hap_weight_row; - double *result_tmp_row; - double *weights = tsk_malloc(3 * state_dim * sizeof(*weights)); - double *norm = tsk_malloc(result_dim * sizeof(*norm)); - double *result_tmp - = tsk_malloc(result_row_len * num_a_alleles * sizeof(*result_tmp)); - - tsk_memset(&ss_A_samples, 0, sizeof(ss_A_samples)); - tsk_memset(&ss_B_samples, 0, sizeof(ss_B_samples)); - tsk_memset(&ss_AB_samples, 0, sizeof(ss_AB_samples)); - tsk_memset(&AB_samples, 0, sizeof(AB_samples)); + out->weights = tsk_malloc(3 * state_dim * sizeof(*out->weights)); + out->norm = tsk_malloc(result_dim * sizeof(*out->norm)); + out->result_tmp + = tsk_malloc(result_dim * max_alleles * max_alleles * sizeof(*out->result_tmp)); - if (weights == NULL || norm == NULL || result_tmp == NULL) { + tsk_memset(&out->ss_A_samples, 0, sizeof(out->ss_A_samples)); + tsk_memset(&out->ss_B_samples, 0, sizeof(out->ss_B_samples)); + tsk_memset(&out->ss_AB_samples, 0, sizeof(out->ss_AB_samples)); + tsk_memset(&out->AB_samples, 0, sizeof(out->AB_samples)); + + if (out->weights == NULL || out->norm == NULL || out->result_tmp == NULL) { ret = tsk_trace_error(TSK_ERR_NO_MEMORY); goto out; } - ret = tsk_bit_array_init(&ss_A_samples, num_samples, 1); + ret = tsk_bit_array_init(&out->AB_samples, num_samples, 1); if (ret != 0) { goto out; } - ret = tsk_bit_array_init(&ss_B_samples, num_samples, 1); + ret = tsk_bit_array_init(&out->ss_A_samples, num_samples, 1); if (ret != 0) { goto out; } - ret = tsk_bit_array_init(&ss_AB_samples, num_samples, 1); + ret = tsk_bit_array_init(&out->ss_B_samples, num_samples, 1); if (ret != 0) { goto out; } - ret = tsk_bit_array_init(&AB_samples, num_samples, 1); + ret = tsk_bit_array_init(&out->ss_AB_samples, num_samples, 1); if (ret != 0) { goto out; } +out: + return ret; +} + +static void +two_locus_work_free(two_locus_work_t *work) +{ + tsk_safe_free(work->weights); + tsk_safe_free(work->norm); + tsk_safe_free(work->result_tmp); + tsk_bit_array_free(&work->AB_samples); + tsk_bit_array_free(&work->ss_A_samples); + tsk_bit_array_free(&work->ss_B_samples); + tsk_bit_array_free(&work->ss_AB_samples); +} + +static int +compute_general_two_site_stat_result(const tsk_bit_array_t *site_a_state, + const tsk_bit_array_t *site_b_state, tsk_size_t num_a_alleles, + tsk_size_t num_b_alleles, tsk_size_t state_dim, const tsk_bit_array_t *sample_sets, + tsk_size_t result_dim, general_stat_func_t *f, sample_count_stat_params_t *f_params, + norm_func_t *norm_f, bool polarised, two_locus_work_t *restrict work, double *result) +{ + int ret = 0; + tsk_bit_array_t A_samples, B_samples; + // ss_ prefix refers to a sample set + tsk_bit_array_t ss_row; + // Sample sets and b sites are rows, a sites are columns + // b1 b2 b3 + // a1 [s1, s2, s3] [s1, s2, s3] [s1, s2, s3] + // a2 [s1, s2, s3] [s1, s2, s3] [s1, s2, s3] + // a3 [s1, s2, s3] [s1, s2, s3] [s1, s2, s3] + tsk_size_t k, mut_a, mut_b; + tsk_size_t result_row_len = num_b_alleles * result_dim; + tsk_size_t w_A = 0, w_B = 0, w_AB = 0; + uint8_t polarised_val = polarised ? 1 : 0; + double *hap_weight_row; + double *result_tmp_row; + + double *norm = work->norm; + double *weights = work->weights; + double *result_tmp = work->result_tmp; + tsk_bit_array_t AB_samples = work->AB_samples; + tsk_bit_array_t ss_A_samples = work->ss_A_samples; + tsk_bit_array_t ss_B_samples = work->ss_B_samples; + tsk_bit_array_t ss_AB_samples = work->ss_AB_samples; for (mut_a = polarised_val; mut_a < num_a_alleles; mut_a++) { result_tmp_row = GET_2D_ROW(result_tmp, result_row_len, mut_a); @@ -2389,13 +2426,6 @@ compute_general_two_site_stat_result(const tsk_bit_array_t *site_a_state, } out: - tsk_safe_free(weights); - tsk_safe_free(norm); - tsk_safe_free(result_tmp); - tsk_bit_array_free(&ss_A_samples); - tsk_bit_array_free(&ss_B_samples); - tsk_bit_array_free(&ss_AB_samples); - tsk_bit_array_free(&AB_samples); return ret; } @@ -2544,12 +2574,12 @@ tsk_treeseq_two_site_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim, tsk_bit_array_t allele_samples, c_state, r_state; bool polarised = false; tsk_id_t *sites; - tsk_size_t r, c, s, n_alleles, n_sites, *row_idx, *col_idx; + tsk_size_t r, c, s, max_alleles, n_alleles, n_sites, *row_idx, *col_idx; double *result_row; const tsk_size_t num_samples = self->num_samples; tsk_size_t *num_alleles = NULL, *site_offsets = NULL; tsk_size_t result_row_len = n_cols * result_dim; - + two_locus_work_t work = { 0 }; tsk_memset(&allele_samples, 0, sizeof(allele_samples)); sites = tsk_malloc(self->tables->sites.num_rows * sizeof(*sites)); @@ -2570,11 +2600,20 @@ tsk_treeseq_two_site_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim, ret = tsk_trace_error(TSK_ERR_NO_MEMORY); goto out; } - n_alleles = 0; + max_alleles = 0; for (s = 0; s < n_sites; s++) { site_offsets[s] = n_alleles; n_alleles += self->site_mutations_length[sites[s]] + 1; + if (self->site_mutations_length[sites[s]] > max_alleles) { + max_alleles = self->site_mutations_length[sites[s]]; + } + } + max_alleles++; + + ret = two_locus_work_init(max_alleles, result_dim, state_dim, num_samples, &work); + if (ret != 0) { + goto out; } ret = tsk_bit_array_init(&allele_samples, num_samples, n_alleles); if (ret != 0) { @@ -2596,8 +2635,8 @@ tsk_treeseq_two_site_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim, tsk_bit_array_get_row(&allele_samples, site_offsets[row_idx[r]], &r_state); tsk_bit_array_get_row(&allele_samples, site_offsets[col_idx[c]], &c_state); ret = compute_general_two_site_stat_result(&r_state, &c_state, - num_alleles[row_idx[r]], num_alleles[col_idx[c]], num_samples, state_dim, - sample_sets, result_dim, f, f_params, norm_f, polarised, + num_alleles[row_idx[r]], num_alleles[col_idx[c]], state_dim, sample_sets, + result_dim, f, f_params, norm_f, polarised, &work, &(result_row[c * result_dim])); if (ret != 0) { goto out; @@ -2611,6 +2650,7 @@ tsk_treeseq_two_site_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim, tsk_safe_free(col_idx); tsk_safe_free(num_alleles); tsk_safe_free(site_offsets); + two_locus_work_free(&work); tsk_bit_array_free(&allele_samples); return ret; } From 01498e2813ff0938fc92f359b4350c4214c216f0 Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 8 Aug 2025 14:28:21 -0500 Subject: [PATCH 21/70] fix ruff formatting --- python/tskit/trees.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/python/tskit/trees.py b/python/tskit/trees.py index f6c7d51007..77ccbf4838 100644 --- a/python/tskit/trees.py +++ b/python/tskit/trees.py @@ -686,7 +686,8 @@ def __init__( options = 0 if sample_counts is not None: warnings.warn( - "The sample_counts option is not supported since 0.2.4 and is ignored", + "The sample_counts option is not supported since 0.2.4 " + "and is ignored", RuntimeWarning, stacklevel=4, ) @@ -6839,7 +6840,7 @@ def to_macs(self): bytes_genotypes[:] = lookup[variant.genotypes] genotypes = bytes_genotypes.tobytes().decode() output.append( - f"SITE:\t{variant.index}\t{variant.position / m}\t0.0\t{genotypes}" + f"SITE:\t{variant.index}\t{variant.position / m}\t0.0\t" f"{genotypes}" ) return "\n".join(output) + "\n" @@ -9142,9 +9143,9 @@ def pca( if time_windows is None: tree_sequence_low, tree_sequence_high = None, self else: - assert time_windows[0] < time_windows[1], ( - "The second argument should be larger." - ) + assert ( + time_windows[0] < time_windows[1] + ), "The second argument should be larger." tree_sequence_low, tree_sequence_high = ( self.decapitate(time_windows[0]), self.decapitate(time_windows[1]), @@ -9212,9 +9213,9 @@ def _rand_pow_range_finder( """ Algorithm 9 in https://arxiv.org/pdf/2002.01387 """ - assert num_vectors >= rank > 0, ( - "num_vectors should not be smaller than rank" - ) + assert ( + num_vectors >= rank > 0 + ), "num_vectors should not be smaller than rank" for _ in range(depth): Q = np.linalg.qr(Q)[0] Q = operator(Q) From 4da94925518f54749ecb1814ddd9b9b0d84a1a15 Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 8 Aug 2025 14:29:40 -0500 Subject: [PATCH 22/70] oops forgot to resolve all conflicts --- python/tests/test_ld_matrix.py | 49 +--------------------------------- 1 file changed, 1 insertion(+), 48 deletions(-) diff --git a/python/tests/test_ld_matrix.py b/python/tests/test_ld_matrix.py index 083b467b48..324838d70c 100644 --- a/python/tests/test_ld_matrix.py +++ b/python/tests/test_ld_matrix.py @@ -22,6 +22,7 @@ """ Test cases for two-locus statistics """ + import contextlib import io from dataclasses import dataclass @@ -2319,20 +2320,12 @@ def test_two_way_site_ld_matrix(ts, stat): ( correlated, (np.array([0, 1, 2, 3, 4, 5]), np.array([6, 7, 8, 9, 10])), -<<<<<<< HEAD np.float64(1.0), -======= - np.float64(0.9708352229780801), ->>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ( correlated, (np.array([0, 1, 2, 3, 4, 5]), np.array([6, 7, 8, 9])), -<<<<<<< HEAD np.float64(1.0), -======= - np.float64(0.9526958931720837), ->>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ( correlated, @@ -2342,20 +2335,12 @@ def test_two_way_site_ld_matrix(ts, stat): ( correlated, (np.array([0, 1, 2, 3, 4, 5]), np.array([6, 7])), -<<<<<<< HEAD np.float64(np.nan), -======= - np.float64(0.7585185185185186), ->>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ( correlated, (np.array([0, 1, 2, 3, 4, 5]), np.array([6])), -<<<<<<< HEAD np.float64(np.nan), -======= - np.float64(0.0), ->>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ( anticorrelated := np.array( @@ -2373,65 +2358,37 @@ def test_two_way_site_ld_matrix(ts, stat): ( anticorrelated, (np.array([0, 2, 4, 6, 8, 10, 12, 14]), np.array([1, 3, 5, 7, 9, 11, 13])), -<<<<<<< HEAD np.float64(1.0), -======= - np.float64(0.9798566895766568), ->>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ( anticorrelated, (np.array([0, 2, 4, 6, 8, 10, 12, 14]), np.array([1, 3, 5, 7, 9, 11])), -<<<<<<< HEAD np.float64(np.nan), -======= - np.float64(0.8574999999999999), ->>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ( anticorrelated, (np.array([0, 2, 4, 6, 8, 10, 12, 14]), np.array([1, 3, 5, 7, 9])), -<<<<<<< HEAD np.float64(np.nan), -======= - np.float64(0.8299777777777777), ->>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ( anticorrelated, (np.array([0, 2, 4, 6, 8, 10, 12, 14]), np.array([1, 3, 5, 7])), -<<<<<<< HEAD np.float64(np.nan), -======= - np.float64(0.6328124999999999), ->>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ( anticorrelated, (np.array([0, 2, 4, 6, 8, 10, 12, 14]), np.array([1, 3, 5])), -<<<<<<< HEAD np.float64(np.nan), -======= - np.float64(0.57179616638322), ->>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ( anticorrelated, (np.array([0, 2, 4, 6, 8, 10, 12, 14]), np.array([1, 3])), -<<<<<<< HEAD np.float64(np.nan), -======= - np.float64(0.0), ->>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ( anticorrelated, (np.array([0, 2, 4, 6, 8, 10, 12, 14]), np.array([1])), -<<<<<<< HEAD np.float64(np.nan), -======= - np.float64(0.0), ->>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) ), ], ) @@ -2450,8 +2407,4 @@ def test_multipopulation_r2_varying_unequal_set_sizes(genotypes, sample_sets, ex r2_ij_summary_func(state_dim, state, 1, result[i, j], params) norm_hap_weighted_ij(1, state, max(a) + 1, max(b) + 1, norm[i, j], params) -<<<<<<< HEAD np.testing.assert_allclose((result * norm).sum(), expected) -======= - np.testing.assert_allclose(expected, (result * norm).sum()) ->>>>>>> c5bdbc7e (C and Python API for two-way two-locus stats) From 7d084be03dd512cac312ada4f8d547d9c6e86fbf Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 8 Aug 2025 14:30:32 -0500 Subject: [PATCH 23/70] Fix another ruff formatting change that snuck by --- python/tests/test_lowlevel.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/python/tests/test_lowlevel.py b/python/tests/test_lowlevel.py index 8f9cbb83e3..d6416c22b8 100644 --- a/python/tests/test_lowlevel.py +++ b/python/tests/test_lowlevel.py @@ -23,7 +23,6 @@ """ Test cases for the low level C interface to tskit. """ - import collections import gc import inspect @@ -3243,10 +3242,7 @@ def f(x): for bad_array in [[1, 1], range(10)]: with pytest.raises(ValueError): ts.general_stat( - W, - lambda x: bad_array, - 1, - ts.get_breakpoints(), # noqa:B023 + W, lambda x: bad_array, 1, ts.get_breakpoints() # noqa:B023 ) with pytest.raises(ValueError): ts.general_stat(W, lambda x: [1], 2, ts.get_breakpoints()) @@ -3255,10 +3251,7 @@ def f(x): for bad_array in [["sdf"], 0, "w4", None]: with pytest.raises(ValueError): ts.general_stat( - W, - lambda x: bad_array, - 1, - ts.get_breakpoints(), # noqa:B023 + W, lambda x: bad_array, 1, ts.get_breakpoints() # noqa:B023 ) From 944334f0f9f48026776ecb462be0f4b7c8cd72b8 Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 8 Aug 2025 14:32:03 -0500 Subject: [PATCH 24/70] a few more things that snuck by merge conflict resolution --- python/tests/test_ld_matrix.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/python/tests/test_ld_matrix.py b/python/tests/test_ld_matrix.py index 324838d70c..5131b02822 100644 --- a/python/tests/test_ld_matrix.py +++ b/python/tests/test_ld_matrix.py @@ -22,7 +22,6 @@ """ Test cases for two-locus statistics """ - import contextlib import io from dataclasses import dataclass @@ -1053,10 +1052,8 @@ def r2_ij_summary_func( D_j = pAB - pA * pB denom_j = np.sqrt(pA * (1 - pA) * pB * (1 - pB)) - p_A = (w_A_i + w_A_j) / (ni + nj) - p_B = (w_B_i + w_B_j) / (ni + nj) with suppress_overflow_div0_warning(): - result[k] = (D_i * D_j) / (p_A * (1 - p_A) * p_B * (1 - p_B)) + result[k] = (D_i * D_j) / (denom_i * denom_j) def D_summary_func( From c60110726b8733923fcde4db55707914c65895b4 Mon Sep 17 00:00:00 2001 From: lkirk Date: Wed, 18 Dec 2024 14:34:46 -0600 Subject: [PATCH 25/70] C and Python API for two-way two-locus stats This PR implements the C and Python API for computing two-way two-locus statistics. The algorithm is identical to the python version, except during testing I uncovered a small issue with normalisation. We need to handle the case where sample sets are of different sizes. The fix for this was to average the normalisation factor for each sample set. Test coverage has been added to cover C, low-level python and some high-level tests. --- python/tests/test_ld_matrix.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/tests/test_ld_matrix.py b/python/tests/test_ld_matrix.py index 5131b02822..6f3074c049 100644 --- a/python/tests/test_ld_matrix.py +++ b/python/tests/test_ld_matrix.py @@ -1052,8 +1052,10 @@ def r2_ij_summary_func( D_j = pAB - pA * pB denom_j = np.sqrt(pA * (1 - pA) * pB * (1 - pB)) + p_A = (w_A_i + w_A_j) / (ni + nj) + p_B = (w_B_i + w_B_j) / (ni + nj) with suppress_overflow_div0_warning(): - result[k] = (D_i * D_j) / (denom_i * denom_j) + result[k] = (D_i * D_j) / (p_A * (1 - p_A) * p_B * (1 - p_B)) def D_summary_func( From ab60e502935dbad0a1f93f8c0e920b4b620266bb Mon Sep 17 00:00:00 2001 From: lkirk Date: Mon, 21 Jul 2025 20:26:21 -0500 Subject: [PATCH 26/70] fix norm_hap_weighted_ij. needs to get wAB instead of using indices From 81aef93a258405dae96e92d0dd2938b5d60d74e9 Mon Sep 17 00:00:00 2001 From: lkirk Date: Wed, 23 Jul 2025 05:57:29 -0500 Subject: [PATCH 27/70] concat is not compatible with numpy<2 From 0dac45d6ad6762249cf0941927d9d84f9fed8c92 Mon Sep 17 00:00:00 2001 From: lkirk Date: Wed, 23 Jul 2025 11:10:57 -0500 Subject: [PATCH 28/70] remove unreachable code path (we switch on the presence of indexes) From f918f80b56e2aee74bd495dcb41a8326c6e3b390 Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 25 Jul 2025 17:49:04 -0500 Subject: [PATCH 29/70] better docstring for norm func From 0c169923e793216461968d837bfca71dd6ce2410 Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 25 Jul 2025 17:57:17 -0500 Subject: [PATCH 30/70] remove todo From 312fbef1740533f62ac14b9ee8d11dd8f95d64a2 Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 25 Jul 2025 17:59:00 -0500 Subject: [PATCH 31/70] remove todo From d13d6964bbd24085e088b1fa516fb54102b8acd5 Mon Sep 17 00:00:00 2001 From: lkirk Date: Thu, 29 May 2025 10:22:53 -0500 Subject: [PATCH 32/70] initial implementation; need to work on output dims/scalar sample sets and indexes --- c/tskit/trees.c | 63 ++----------------------------------------- python/tskit/trees.py | 1 + 2 files changed, 3 insertions(+), 61 deletions(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index b832ae7515..253188808e 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -2294,67 +2294,8 @@ get_all_samples_bits(tsk_bit_array_t *all_samples, tsk_size_t n) } } -typedef struct { - double *weights; - double *norm; - double *result_tmp; - tsk_bit_array_t AB_samples; - tsk_bit_array_t ss_A_samples; - tsk_bit_array_t ss_B_samples; - tsk_bit_array_t ss_AB_samples; -} two_locus_work_t; - -static int -two_locus_work_init(tsk_size_t max_alleles, tsk_size_t result_dim, tsk_size_t state_dim, - tsk_size_t num_samples, two_locus_work_t *out) -{ - int ret = 0; - out->weights = tsk_malloc(3 * state_dim * sizeof(*out->weights)); - out->norm = tsk_malloc(result_dim * sizeof(*out->norm)); - out->result_tmp - = tsk_malloc(result_dim * max_alleles * max_alleles * sizeof(*out->result_tmp)); - - tsk_memset(&out->ss_A_samples, 0, sizeof(out->ss_A_samples)); - tsk_memset(&out->ss_B_samples, 0, sizeof(out->ss_B_samples)); - tsk_memset(&out->ss_AB_samples, 0, sizeof(out->ss_AB_samples)); - tsk_memset(&out->AB_samples, 0, sizeof(out->AB_samples)); - - if (out->weights == NULL || out->norm == NULL || out->result_tmp == NULL) { - ret = tsk_trace_error(TSK_ERR_NO_MEMORY); - goto out; - } - - ret = tsk_bit_array_init(&out->AB_samples, num_samples, 1); - if (ret != 0) { - goto out; - } - ret = tsk_bit_array_init(&out->ss_A_samples, num_samples, 1); - if (ret != 0) { - goto out; - } - ret = tsk_bit_array_init(&out->ss_B_samples, num_samples, 1); - if (ret != 0) { - goto out; - } - ret = tsk_bit_array_init(&out->ss_AB_samples, num_samples, 1); - if (ret != 0) { - goto out; - } -out: - return ret; -} - -static void -two_locus_work_free(two_locus_work_t *work) -{ - tsk_safe_free(work->weights); - tsk_safe_free(work->norm); - tsk_safe_free(work->result_tmp); - tsk_bit_array_free(&work->AB_samples); - tsk_bit_array_free(&work->ss_A_samples); - tsk_bit_array_free(&work->ss_B_samples); - tsk_bit_array_free(&work->ss_AB_samples); -} +typedef int norm_func_t(tsk_size_t result_dim, const double *hap_weights, tsk_size_t n_a, + tsk_size_t n_b, double *result, void *params); static int compute_general_two_site_stat_result(const tsk_bit_array_t *site_a_state, diff --git a/python/tskit/trees.py b/python/tskit/trees.py index 77ccbf4838..699debc0c3 100644 --- a/python/tskit/trees.py +++ b/python/tskit/trees.py @@ -23,6 +23,7 @@ """ Module responsible for managing trees and tree sequences. """ + from __future__ import annotations import base64 From ec39a322ed3b9e0374279f397ecb986677018fdf Mon Sep 17 00:00:00 2001 From: lkirk Date: Thu, 17 Jul 2025 02:34:24 -0500 Subject: [PATCH 33/70] fix normalisation and a small indexing bug. dimension tests --- c/tskit/trees.c | 9 +++------ python/tests/test_ld_matrix.py | 4 ++++ 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index 253188808e..6769e82dba 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -3222,17 +3222,12 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl tsk_memset(&sample_sets_bits, 0, sizeof(sample_sets_bits)); - // We do not support two-locus node stats - if (!!(options & TSK_STAT_NODE)) { - ret = tsk_trace_error(TSK_ERR_UNSUPPORTED_STAT_MODE); - goto out; - } // If no mode is specified, we default to site mode if (!(stat_site || stat_branch)) { stat_site = true; } // It's an error to specify more than one mode - if (stat_site + stat_branch > 1) { + if (stat_site + stat_branch + stat_node > 1) { ret = tsk_trace_error(TSK_ERR_MULTIPLE_STAT_MODES); goto out; } @@ -3277,6 +3272,8 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl ret = tsk_treeseq_two_branch_count_stat(self, state_dim, &sample_sets_bits, result_dim, f, &f_params, norm_f, out_rows, row_positions, out_cols, col_positions, options, result); + } else { + ret = tsk_trace_error(TSK_ERR_UNSUPPORTED_STAT_MODE); } out: diff --git a/python/tests/test_ld_matrix.py b/python/tests/test_ld_matrix.py index 6f3074c049..c69ecd4dad 100644 --- a/python/tests/test_ld_matrix.py +++ b/python/tests/test_ld_matrix.py @@ -590,6 +590,9 @@ def compute_general_two_site_stat_result( for k in range(result_dim): result[k] += result_tmp[k] * norm[k] + # for k in range(result_dim): + # print(mut_a, mut_b, k, weights[0, k], weights[1, k], weights[2, k], sep="\t") + def two_site_count_stat( ts: tskit.TreeSequence, @@ -1777,6 +1780,7 @@ def test_ld_empty_examples(ts): def test_input_validation(): + # TODO ts = get_paper_ex_ts() with pytest.raises(ValueError, match="Unknown two-locus statistic"): ts.ld_matrix(stat="bad_stat") From c09e5d1edc4af37be74d62e79277080923ebb4ae Mon Sep 17 00:00:00 2001 From: lkirk Date: Thu, 17 Jul 2025 23:57:00 -0500 Subject: [PATCH 34/70] k should be unsigned From 0ef7cd6860885e5747ec4fe065a9486b622e6b98 Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 25 Jul 2025 17:42:13 -0500 Subject: [PATCH 35/70] just error early if node stat requested --- c/tskit/trees.c | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index 6769e82dba..253188808e 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -3222,12 +3222,17 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl tsk_memset(&sample_sets_bits, 0, sizeof(sample_sets_bits)); + // We do not support two-locus node stats + if (!!(options & TSK_STAT_NODE)) { + ret = tsk_trace_error(TSK_ERR_UNSUPPORTED_STAT_MODE); + goto out; + } // If no mode is specified, we default to site mode if (!(stat_site || stat_branch)) { stat_site = true; } // It's an error to specify more than one mode - if (stat_site + stat_branch + stat_node > 1) { + if (stat_site + stat_branch > 1) { ret = tsk_trace_error(TSK_ERR_MULTIPLE_STAT_MODES); goto out; } @@ -3272,8 +3277,6 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl ret = tsk_treeseq_two_branch_count_stat(self, state_dim, &sample_sets_bits, result_dim, f, &f_params, norm_f, out_rows, row_positions, out_cols, col_positions, options, result); - } else { - ret = tsk_trace_error(TSK_ERR_UNSUPPORTED_STAT_MODE); } out: From 751c948021c1d644dd207c1d1aaf35f357f18b0b Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 25 Jul 2025 17:59:00 -0500 Subject: [PATCH 36/70] remove todo From a73f736091b078d29e0fac3bf8f555ae4ad9ccd2 Mon Sep 17 00:00:00 2001 From: lkirk Date: Thu, 24 Jul 2025 18:11:56 -0500 Subject: [PATCH 37/70] first pass at bit array refactor --- c/tests/test_core.c | 56 ++++----- c/tskit/core.c | 93 +++++++++----- c/tskit/core.h | 43 ++++--- c/tskit/trees.c | 301 +++++++++++++++++++++++++------------------- 4 files changed, 283 insertions(+), 210 deletions(-) diff --git a/c/tests/test_core.c b/c/tests/test_core.c index 6a14ecc655..9fa1b576b3 100644 --- a/c/tests/test_core.c +++ b/c/tests/test_core.c @@ -531,22 +531,22 @@ test_bit_arrays(void) { // NB: This test is only valid for the 32 bit implementation of bit arrays. If we // were to change the chunk size of a bit array, we'd need to update these tests - tsk_bit_array_t arr; + tsk_bitset_t arr; tsk_id_t items_truth[64] = { 0 }, items[64] = { 0 }; tsk_size_t n_items = 0, n_items_truth = 0; // test item retrieval - tsk_bit_array_init(&arr, 90, 1); - tsk_bit_array_get_items(&arr, items, &n_items); + tsk_bitset_init(&arr, 90, 1); + tsk_bitset_get_items(&arr, items, &n_items); assert_arrays_equal(n_items_truth, items, items_truth); - for (tsk_bit_array_value_t i = 0; i < 20; i++) { - tsk_bit_array_add_bit(&arr, i); + for (tsk_bitset_val_t i = 0; i < 20; i++) { + tsk_bitset_set_bit(&arr, 0, i); items_truth[n_items_truth] = (tsk_id_t) i; n_items_truth++; } - tsk_bit_array_add_bit(&arr, 63); - tsk_bit_array_add_bit(&arr, 65); + tsk_bitset_set_bit(&arr, 0, 63); + tsk_bitset_set_bit(&arr, 0, 65); // these assertions are only valid for 32-bit values CU_ASSERT_EQUAL_FATAL(arr.data[0], 1048575); @@ -554,32 +554,32 @@ test_bit_arrays(void) CU_ASSERT_EQUAL_FATAL(arr.data[2], 2); // verify our assumptions about bit array counting - CU_ASSERT_EQUAL_FATAL(tsk_bit_array_count(&arr), 22); + CU_ASSERT_EQUAL_FATAL(tsk_bitset_count(&arr, 0), 22); - tsk_bit_array_get_items(&arr, items, &n_items); + tsk_bitset_get_items(&arr, items, &n_items); assert_arrays_equal(n_items_truth, items, items_truth); tsk_memset(items, 0, 64); tsk_memset(items_truth, 0, 64); n_items = n_items_truth = 0; - tsk_bit_array_free(&arr); + tsk_bitset_free(&arr); // create a length-2 array with 64 bit capacity - tsk_bit_array_init(&arr, 64, 2); - tsk_bit_array_t arr_row1, arr_row2; + tsk_bitset_init(&arr, 64, 2); + tsk_bitset_t arr_row1, arr_row2; // select the first and second row - tsk_bit_array_get_row(&arr, 0, &arr_row1); - tsk_bit_array_get_row(&arr, 1, &arr_row2); + tsk_bitset_get_row(&arr, 0, &arr_row1); + tsk_bitset_get_row(&arr, 1, &arr_row2); // fill the first 50 bits of the first row - for (tsk_bit_array_value_t i = 0; i < 50; i++) { - tsk_bit_array_add_bit(&arr_row1, i); + for (tsk_bitset_val_t i = 0; i < 50; i++) { + tsk_bitset_set_bit(&arr_row1, 0, i); items_truth[n_items_truth] = (tsk_id_t) i; n_items_truth++; } - tsk_bit_array_get_items(&arr_row1, items, &n_items); + tsk_bitset_get_items(&arr_row1, items, &n_items); assert_arrays_equal(n_items_truth, items, items_truth); tsk_memset(items, 0, 64); @@ -587,13 +587,13 @@ test_bit_arrays(void) n_items = n_items_truth = 0; // fill bits 20-40 of the second row - for (tsk_bit_array_value_t i = 20; i < 40; i++) { - tsk_bit_array_add_bit(&arr_row2, i); + for (tsk_bitset_val_t i = 20; i < 40; i++) { + tsk_bitset_set_bit(&arr_row2, 0, i); items_truth[n_items_truth] = (tsk_id_t) i; n_items_truth++; } - tsk_bit_array_get_items(&arr_row2, items, &n_items); + tsk_bitset_get_items(&arr_row2, items, &n_items); assert_arrays_equal(n_items_truth, items, items_truth); tsk_memset(items, 0, 64); @@ -612,30 +612,30 @@ test_bit_arrays(void) CU_ASSERT_EQUAL_FATAL(arr_row2.data[1], 255); // subtract the second from the first row, store in first - tsk_bit_array_subtract(&arr_row1, &arr_row2); + tsk_bitset_subtract(&arr_row1, 0, &arr_row2, 0); // verify our assumptions about subtraction CU_ASSERT_EQUAL_FATAL(arr_row1.data[0], 1048575); CU_ASSERT_EQUAL_FATAL(arr_row1.data[1], 261888); - tsk_bit_array_t int_result; - tsk_bit_array_init(&int_result, 64, 1); + tsk_bitset_t int_result; + tsk_bitset_init(&int_result, 64, 1); // their intersection should be zero - tsk_bit_array_intersect(&arr_row1, &arr_row2, &int_result); + tsk_bitset_intersect(&arr_row1, 0, &arr_row2, 0, &int_result); CU_ASSERT_EQUAL_FATAL(int_result.data[0], 0); CU_ASSERT_EQUAL_FATAL(int_result.data[1], 0); // now, add them back together, storing back in a - tsk_bit_array_add(&arr_row1, &arr_row2); + tsk_bitset_union(&arr_row1, 0, &arr_row2, 0); // now, their intersection should be the subtracted chunk (20-40) - tsk_bit_array_intersect(&arr_row1, &arr_row2, &int_result); + tsk_bitset_intersect(&arr_row1, 0, &arr_row2, 0, &int_result); CU_ASSERT_EQUAL_FATAL(int_result.data[0], 4293918720); CU_ASSERT_EQUAL_FATAL(int_result.data[1], 255); - tsk_bit_array_free(&int_result); - tsk_bit_array_free(&arr); + tsk_bitset_free(&int_result); + tsk_bitset_free(&arr); } static void diff --git a/c/tskit/core.c b/c/tskit/core.c index 175505439e..785f590133 100644 --- a/c/tskit/core.c +++ b/c/tskit/core.c @@ -1253,13 +1253,14 @@ tsk_avl_tree_int_ordered_nodes(const tsk_avl_tree_int_t *self, tsk_avl_node_int_ // Currently implemented as an array of 32-bit unsigned integers for ease of counting. int -tsk_bit_array_init(tsk_bit_array_t *self, tsk_size_t num_bits, tsk_size_t length) +tsk_bitset_init(tsk_bitset_t *self, tsk_size_t num_bits, tsk_size_t length) { int ret = 0; - self->size = (num_bits >> TSK_BIT_ARRAY_CHUNK) - + (num_bits % TSK_BIT_ARRAY_NUM_BITS ? 1 : 0); - self->data = tsk_calloc(self->size * length, sizeof(*self->data)); + self->row_len = (num_bits >> TSK_BIT_ARRAY_CHUNK) + + (num_bits % TSK_BIT_ARRAY_NUM_BITS ? 1 : 0); + self->len = length; + self->data = tsk_calloc(self->row_len * length, sizeof(*self->data)); if (self->data == NULL) { ret = tsk_trace_error(TSK_ERR_NO_MEMORY); goto out; @@ -1269,54 +1270,77 @@ tsk_bit_array_init(tsk_bit_array_t *self, tsk_size_t num_bits, tsk_size_t length } void -tsk_bit_array_get_row(const tsk_bit_array_t *self, tsk_size_t row, tsk_bit_array_t *out) +tsk_bitset_get_row(const tsk_bitset_t *self, tsk_size_t row, tsk_bitset_t *out) { - out->size = self->size; - out->data = self->data + (row * self->size); + out->row_len = self->row_len; + out->data = self->data + (row * self->row_len); + out->len = (tsk_size_t)(out->data - self->data) / self->row_len; + out->len++; } void -tsk_bit_array_intersect( - const tsk_bit_array_t *self, const tsk_bit_array_t *other, tsk_bit_array_t *out) +tsk_bitset_intersect(const tsk_bitset_t *self, tsk_size_t self_row, + const tsk_bitset_t *other, tsk_size_t other_row, tsk_bitset_t *out) { - for (tsk_size_t i = 0; i < self->size; i++) { - out->data[i] = self->data[i] & other->data[i]; + tsk_bitset_val_t *self_d = self->data + (self_row * self->row_len); + tsk_bitset_val_t *other_d = other->data + (other_row * self->row_len); + if (self_row >= self->len || other_row >= other->len) { + printf("isect: OUTSIDE_OF_BOUNDS"); + } + for (tsk_size_t i = 0; i < self->row_len; i++) { + out->data[i] = self_d[i] & other_d[i]; } } void -tsk_bit_array_subtract(tsk_bit_array_t *self, const tsk_bit_array_t *other) +tsk_bitset_subtract(tsk_bitset_t *self, tsk_size_t self_row, const tsk_bitset_t *other, + tsk_size_t other_row) { - for (tsk_size_t i = 0; i < self->size; i++) { - self->data[i] &= ~(other->data[i]); + tsk_bitset_val_t *self_d = self->data + (self_row * self->row_len); + tsk_bitset_val_t *other_d = other->data + (other_row * self->row_len); + if (self_row >= self->len || other_row >= other->len) { + printf("subtr: OUTSIDE_OF_BOUNDS"); + } + for (tsk_size_t i = 0; i < self->row_len; i++) { + self_d[i] &= ~(other_d[i]); } } void -tsk_bit_array_add(tsk_bit_array_t *self, const tsk_bit_array_t *other) +tsk_bitset_union(tsk_bitset_t *self, tsk_size_t self_row, const tsk_bitset_t *other, + tsk_size_t other_row) { - for (tsk_size_t i = 0; i < self->size; i++) { - self->data[i] |= other->data[i]; + tsk_bitset_val_t *self_d = self->data + (self_row * self->row_len); + tsk_bitset_val_t *other_d = other->data + (other_row * self->row_len); + if (self_row >= self->len || other_row >= other->len) { + printf("union: OUTSIDE_OF_BOUNDS"); + } + for (tsk_size_t i = 0; i < self->row_len; i++) { + self_d[i] |= other_d[i]; } } void -tsk_bit_array_add_bit(tsk_bit_array_t *self, const tsk_bit_array_value_t bit) +tsk_bitset_set_bit(tsk_bitset_t *self, tsk_size_t row, const tsk_bitset_val_t bit) { - tsk_bit_array_value_t i = bit >> TSK_BIT_ARRAY_CHUNK; - self->data[i] |= (tsk_bit_array_value_t) 1 << (bit - (TSK_BIT_ARRAY_NUM_BITS * i)); + tsk_bitset_val_t i = (bit >> TSK_BIT_ARRAY_CHUNK); + self->data[i + row * self->row_len] |= (tsk_bitset_val_t) 1 + << (bit - (TSK_BIT_ARRAY_NUM_BITS * i)); } bool -tsk_bit_array_contains(const tsk_bit_array_t *self, const tsk_bit_array_value_t bit) +tsk_bitset_contains(const tsk_bitset_t *self, tsk_size_t row, const tsk_bitset_val_t bit) { - tsk_bit_array_value_t i = bit >> TSK_BIT_ARRAY_CHUNK; - return self->data[i] - & ((tsk_bit_array_value_t) 1 << (bit - (TSK_BIT_ARRAY_NUM_BITS * i))); + tsk_bitset_val_t i = (bit >> TSK_BIT_ARRAY_CHUNK); + if (row >= self->len) { + printf("contains: OUTSIDE_OF_BOUNDS"); + } + return self->data[i + row * self->row_len] + & ((tsk_bitset_val_t) 1 << (bit - (TSK_BIT_ARRAY_NUM_BITS * i))); } tsk_size_t -tsk_bit_array_count(const tsk_bit_array_t *self) +tsk_bitset_count(const tsk_bitset_t *self, tsk_size_t row) { // Utilizes 12 operations per bit array. NB this only works on 32 bit integers. // Taken from: @@ -1330,11 +1354,15 @@ tsk_bit_array_count(const tsk_bit_array_t *self) // This option is relatively simple, but requires a 64 bit bit array and also // involves some compiler flag plumbing (-mpopcnt) - tsk_bit_array_value_t tmp; + tsk_bitset_val_t tmp; tsk_size_t i, count = 0; + tsk_bitset_val_t *self_d = self->data + (row * self->row_len); - for (i = 0; i < self->size; i++) { - tmp = self->data[i] - ((self->data[i] >> 1) & 0x55555555); + if (row >= self->len) { + printf("count: OUTSIDE_OF_BOUNDS"); + } + for (i = 0; i < self->row_len; i++) { + tmp = self_d[i] - ((self_d[i] >> 1) & 0x55555555); tmp = (tmp & 0x33333333) + ((tmp >> 2) & 0x33333333); count += (((tmp + (tmp >> 4)) & 0xF0F0F0F) * 0x1010101) >> 24; } @@ -1342,20 +1370,19 @@ tsk_bit_array_count(const tsk_bit_array_t *self) } void -tsk_bit_array_get_items( - const tsk_bit_array_t *self, tsk_id_t *items, tsk_size_t *n_items) +tsk_bitset_get_items(const tsk_bitset_t *self, tsk_id_t *items, tsk_size_t *n_items) { // Get the items stored in the row of a bitset. // Uses a de Bruijn sequence lookup table to determine the lowest bit set. See the // wikipedia article for more info: https://w.wiki/BYiF tsk_size_t i, n, off; - tsk_bit_array_value_t v, lsb; // least significant bit + tsk_bitset_val_t v, lsb; // least significant bit static const tsk_id_t lookup[32] = { 0, 1, 28, 2, 29, 14, 24, 3, 30, 22, 20, 15, 25, 17, 4, 8, 31, 27, 13, 23, 21, 19, 16, 7, 26, 12, 18, 6, 11, 5, 10, 9 }; n = 0; - for (i = 0; i < self->size; i++) { + for (i = 0; i < self->row_len; i++) { v = self->data[i]; off = i * ((tsk_size_t) TSK_BIT_ARRAY_NUM_BITS); if (v == 0) { @@ -1371,7 +1398,7 @@ tsk_bit_array_get_items( } void -tsk_bit_array_free(tsk_bit_array_t *self) +tsk_bitset_free(tsk_bitset_t *self) { tsk_safe_free(self->data); } diff --git a/c/tskit/core.h b/c/tskit/core.h index a4d022c2dd..c7f1c6ea89 100644 --- a/c/tskit/core.h +++ b/c/tskit/core.h @@ -1094,29 +1094,32 @@ FILE *tsk_get_debug_stream(void); /* Bit Array functionality */ -typedef uint32_t tsk_bit_array_value_t; +typedef uint32_t tsk_bitset_val_t; typedef struct { - tsk_size_t size; // Number of chunks per row - tsk_bit_array_value_t *data; // Array data -} tsk_bit_array_t; + tsk_size_t row_len; // Number of chunks per row + tsk_size_t len; // Number of rows + tsk_bitset_val_t *data; +} tsk_bitset_t; -#define TSK_BIT_ARRAY_CHUNK 5U +#define TSK_BIT_ARRAY_CHUNK 5 #define TSK_BIT_ARRAY_NUM_BITS (1U << TSK_BIT_ARRAY_CHUNK) - -int tsk_bit_array_init(tsk_bit_array_t *self, tsk_size_t num_bits, tsk_size_t length); -void tsk_bit_array_free(tsk_bit_array_t *self); -void tsk_bit_array_get_row( - const tsk_bit_array_t *self, tsk_size_t row, tsk_bit_array_t *out); -void tsk_bit_array_intersect( - const tsk_bit_array_t *self, const tsk_bit_array_t *other, tsk_bit_array_t *out); -void tsk_bit_array_subtract(tsk_bit_array_t *self, const tsk_bit_array_t *other); -void tsk_bit_array_add(tsk_bit_array_t *self, const tsk_bit_array_t *other); -void tsk_bit_array_add_bit(tsk_bit_array_t *self, const tsk_bit_array_value_t bit); -bool tsk_bit_array_contains( - const tsk_bit_array_t *self, const tsk_bit_array_value_t bit); -tsk_size_t tsk_bit_array_count(const tsk_bit_array_t *self); -void tsk_bit_array_get_items( - const tsk_bit_array_t *self, tsk_id_t *items, tsk_size_t *n_items); +#define TSK_BITSET_BITS 32U + +int tsk_bitset_init(tsk_bitset_t *self, tsk_size_t num_bits, tsk_size_t length); +void tsk_bitset_free(tsk_bitset_t *self); +void tsk_bitset_get_row(const tsk_bitset_t *self, tsk_size_t row, tsk_bitset_t *out); +void tsk_bitset_intersect(const tsk_bitset_t *self, tsk_size_t self_row, + const tsk_bitset_t *other, tsk_size_t other_row, tsk_bitset_t *out); +void tsk_bitset_subtract(tsk_bitset_t *self, tsk_size_t self_row, + const tsk_bitset_t *other, tsk_size_t other_row); +void tsk_bitset_union(tsk_bitset_t *self, tsk_size_t self_row, const tsk_bitset_t *other, + tsk_size_t other_row); +void tsk_bitset_set_bit(tsk_bitset_t *self, tsk_size_t row, const tsk_bitset_val_t bit); +bool tsk_bitset_contains( + const tsk_bitset_t *self, tsk_size_t row, const tsk_bitset_val_t bit); +tsk_size_t tsk_bitset_count(const tsk_bitset_t *self, tsk_size_t row); +void tsk_bitset_get_items( + const tsk_bitset_t *self, tsk_id_t *items, tsk_size_t *n_items); #ifdef __cplusplus } diff --git a/c/tskit/trees.c b/c/tskit/trees.c index 253188808e..fd759b05c5 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -2149,20 +2149,18 @@ tsk_treeseq_sample_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sample_s ***********************************/ static int -get_allele_samples(const tsk_site_t *site, const tsk_bit_array_t *state, - tsk_bit_array_t *out_allele_samples, tsk_size_t *out_num_alleles) +get_allele_samples(const tsk_site_t *site, tsk_size_t site_offset, + const tsk_bitset_t *state, tsk_bitset_t *out_allele_samples, + tsk_size_t *out_num_alleles) { int ret = 0; tsk_mutation_t mutation, parent_mut; - tsk_size_t mutation_index, allele, alt_allele_length; + tsk_size_t mutation_index, allele, alt_allele, alt_allele_length; /* The allele table */ tsk_size_t max_alleles = site->mutations_length + 1; const char **alleles = tsk_malloc(max_alleles * sizeof(*alleles)); tsk_size_t *allele_lengths = tsk_calloc(max_alleles, sizeof(*allele_lengths)); - const char *alt_allele; - tsk_bit_array_t state_row; - tsk_bit_array_t allele_samples_row; - tsk_bit_array_t alt_allele_samples_row; + const char *alt_allele_state; tsk_size_t num_alleles = 1; if (alleles == NULL || allele_lengths == NULL) { @@ -2193,29 +2191,29 @@ get_allele_samples(const tsk_site_t *site, const tsk_bit_array_t *state, } /* Add the mutation's samples to this allele */ - tsk_bit_array_get_row(out_allele_samples, allele, &allele_samples_row); - tsk_bit_array_get_row(state, mutation_index, &state_row); - tsk_bit_array_add(&allele_samples_row, &state_row); + tsk_bitset_union( + out_allele_samples, allele + site_offset, state, mutation_index); /* Get the index for the alternate allele that we must subtract from */ - alt_allele = site->ancestral_state; + alt_allele_state = site->ancestral_state; alt_allele_length = site->ancestral_state_length; if (mutation.parent != TSK_NULL) { parent_mut = site->mutations[mutation.parent - site->mutations[0].id]; - alt_allele = parent_mut.derived_state; + alt_allele_state = parent_mut.derived_state; alt_allele_length = parent_mut.derived_state_length; } - for (allele = 0; allele < num_alleles; allele++) { - if (alt_allele_length == allele_lengths[allele] - && tsk_memcmp(alt_allele, alleles[allele], allele_lengths[allele]) + for (alt_allele = 0; alt_allele < num_alleles; alt_allele++) { + if (alt_allele_length == allele_lengths[alt_allele] + && tsk_memcmp( + alt_allele_state, alleles[alt_allele], allele_lengths[alt_allele]) == 0) { break; } } tsk_bug_assert(allele < num_alleles); - tsk_bit_array_get_row(out_allele_samples, allele, &alt_allele_samples_row); - tsk_bit_array_subtract(&alt_allele_samples_row, &allele_samples_row); + tsk_bitset_subtract(out_allele_samples, alt_allele + site_offset, + out_allele_samples, allele + site_offset); } *out_num_alleles = num_alleles; out: @@ -2281,33 +2279,89 @@ norm_total_weighted(tsk_size_t result_dim, const double *TSK_UNUSED(hap_weights) } static void -get_all_samples_bits(tsk_bit_array_t *all_samples, tsk_size_t n) +get_all_samples_bits(tsk_bitset_t *all_samples, tsk_size_t n) { tsk_size_t i; - const tsk_bit_array_value_t all = ~((tsk_bit_array_value_t) 0); - const tsk_bit_array_value_t remainder_samples = n % TSK_BIT_ARRAY_NUM_BITS; + const tsk_bitset_val_t all = ~((tsk_bitset_val_t) 0); + const tsk_bitset_val_t remainder_samples = n % TSK_BIT_ARRAY_NUM_BITS; - all_samples->data[all_samples->size - 1] + all_samples->data[all_samples->row_len - 1] = remainder_samples ? ~(all << remainder_samples) : all; - for (i = 0; i < all_samples->size - 1; i++) { + for (i = 0; i < all_samples->row_len - 1; i++) { all_samples->data[i] = all; } } -typedef int norm_func_t(tsk_size_t result_dim, const double *hap_weights, tsk_size_t n_a, - tsk_size_t n_b, double *result, void *params); +typedef struct { + double *weights; + double *norm; + double *result_tmp; + tsk_bitset_t AB_samples; + tsk_bitset_t ss_A_samples; + tsk_bitset_t ss_B_samples; + tsk_bitset_t ss_AB_samples; +} two_locus_work_t; static int -compute_general_two_site_stat_result(const tsk_bit_array_t *site_a_state, - const tsk_bit_array_t *site_b_state, tsk_size_t num_a_alleles, - tsk_size_t num_b_alleles, tsk_size_t state_dim, const tsk_bit_array_t *sample_sets, - tsk_size_t result_dim, general_stat_func_t *f, sample_count_stat_params_t *f_params, - norm_func_t *norm_f, bool polarised, two_locus_work_t *restrict work, double *result) +two_locus_work_init(tsk_size_t max_alleles, tsk_size_t result_dim, tsk_size_t state_dim, + tsk_size_t num_samples, two_locus_work_t *out) +{ + int ret = 0; + out->weights = tsk_malloc(3 * state_dim * sizeof(*out->weights)); + out->norm = tsk_malloc(result_dim * sizeof(*out->norm)); + out->result_tmp + = tsk_malloc(result_dim * max_alleles * max_alleles * sizeof(*out->result_tmp)); + + tsk_memset(&out->ss_A_samples, 0, sizeof(out->ss_A_samples)); + tsk_memset(&out->ss_B_samples, 0, sizeof(out->ss_B_samples)); + tsk_memset(&out->ss_AB_samples, 0, sizeof(out->ss_AB_samples)); + tsk_memset(&out->AB_samples, 0, sizeof(out->AB_samples)); + + if (out->weights == NULL || out->norm == NULL || out->result_tmp == NULL) { + ret = tsk_trace_error(TSK_ERR_NO_MEMORY); + goto out; + } + + ret = tsk_bitset_init(&out->AB_samples, num_samples, 1); + if (ret != 0) { + goto out; + } + ret = tsk_bitset_init(&out->ss_A_samples, num_samples, 1); + if (ret != 0) { + goto out; + } + ret = tsk_bitset_init(&out->ss_B_samples, num_samples, 1); + if (ret != 0) { + goto out; + } + ret = tsk_bitset_init(&out->ss_AB_samples, num_samples, 1); + if (ret != 0) { + goto out; + } +out: + return ret; +} + +static void +two_locus_work_free(two_locus_work_t *work) +{ + tsk_safe_free(work->weights); + tsk_safe_free(work->norm); + tsk_safe_free(work->result_tmp); + tsk_bitset_free(&work->AB_samples); + tsk_bitset_free(&work->ss_A_samples); + tsk_bitset_free(&work->ss_B_samples); + tsk_bitset_free(&work->ss_AB_samples); +} + +static int +compute_general_two_site_stat_result(const tsk_bitset_t *state, tsk_size_t a_off, + tsk_size_t b_off, tsk_size_t num_a_alleles, tsk_size_t num_b_alleles, + tsk_size_t state_dim, const tsk_bitset_t *sample_sets, tsk_size_t result_dim, + general_stat_func_t *f, sample_count_stat_params_t *f_params, norm_func_t *norm_f, + bool polarised, two_locus_work_t *restrict work, double *result) { int ret = 0; - tsk_bit_array_t A_samples, B_samples; - // ss_ prefix refers to a sample set - tsk_bit_array_t ss_row; // Sample sets and b sites are rows, a sites are columns // b1 b2 b3 // a1 [s1, s2, s3] [s1, s2, s3] [s1, s2, s3] @@ -2315,7 +2369,7 @@ compute_general_two_site_stat_result(const tsk_bit_array_t *site_a_state, // a3 [s1, s2, s3] [s1, s2, s3] [s1, s2, s3] tsk_size_t k, mut_a, mut_b; tsk_size_t result_row_len = num_b_alleles * result_dim; - tsk_size_t w_A = 0, w_B = 0, w_AB = 0; + double w_AB = 0; uint8_t polarised_val = polarised ? 1 : 0; double *hap_weight_row; double *result_tmp_row; @@ -2323,32 +2377,26 @@ compute_general_two_site_stat_result(const tsk_bit_array_t *site_a_state, double *norm = work->norm; double *weights = work->weights; double *result_tmp = work->result_tmp; - tsk_bit_array_t AB_samples = work->AB_samples; - tsk_bit_array_t ss_A_samples = work->ss_A_samples; - tsk_bit_array_t ss_B_samples = work->ss_B_samples; - tsk_bit_array_t ss_AB_samples = work->ss_AB_samples; + tsk_bitset_t AB_samples = work->AB_samples; + tsk_bitset_t ss_A_samples = work->ss_A_samples; + tsk_bitset_t ss_B_samples = work->ss_B_samples; + tsk_bitset_t ss_AB_samples = work->ss_AB_samples; for (mut_a = polarised_val; mut_a < num_a_alleles; mut_a++) { result_tmp_row = GET_2D_ROW(result_tmp, result_row_len, mut_a); for (mut_b = polarised_val; mut_b < num_b_alleles; mut_b++) { - tsk_bit_array_get_row(site_a_state, mut_a, &A_samples); - tsk_bit_array_get_row(site_b_state, mut_b, &B_samples); - tsk_bit_array_intersect(&A_samples, &B_samples, &AB_samples); + tsk_bitset_intersect( + state, mut_a + a_off, state, mut_b + b_off, &AB_samples); for (k = 0; k < state_dim; k++) { - tsk_bit_array_get_row(sample_sets, k, &ss_row); hap_weight_row = GET_2D_ROW(weights, 3, k); - - tsk_bit_array_intersect(&A_samples, &ss_row, &ss_A_samples); - tsk_bit_array_intersect(&B_samples, &ss_row, &ss_B_samples); - tsk_bit_array_intersect(&AB_samples, &ss_row, &ss_AB_samples); - - w_AB = tsk_bit_array_count(&ss_AB_samples); - w_A = tsk_bit_array_count(&ss_A_samples); - w_B = tsk_bit_array_count(&ss_B_samples); - - hap_weight_row[0] = (double) w_AB; - hap_weight_row[1] = (double) (w_A - w_AB); // w_Ab - hap_weight_row[2] = (double) (w_B - w_AB); // w_aB + tsk_bitset_intersect( + state, mut_a + a_off, sample_sets, k, &ss_A_samples); + tsk_bitset_intersect( + state, mut_b + b_off, sample_sets, k, &ss_B_samples); + tsk_bitset_intersect(&AB_samples, 0, sample_sets, k, &ss_AB_samples); + w_AB = hap_weight_row[0] = (double) tsk_bitset_count(&ss_AB_samples, 0); + hap_weight_row[1] = ((double) tsk_bitset_count(&ss_A_samples, 0) - w_AB); + hap_weight_row[2] = ((double) tsk_bitset_count(&ss_B_samples, 0) - w_AB); } ret = f(state_dim, weights, result_dim, result_tmp_row, f_params); if (ret != 0) { @@ -2417,7 +2465,7 @@ get_site_row_col_indices(tsk_size_t n_rows, const tsk_id_t *row_sites, tsk_size_ static int get_mutation_samples(const tsk_treeseq_t *ts, const tsk_id_t *sites, tsk_size_t n_sites, - tsk_size_t *num_alleles, tsk_bit_array_t *allele_samples) + tsk_size_t *num_alleles, tsk_bitset_t *allele_samples) { int ret = 0; const tsk_flags_t *restrict flags = ts->tables->nodes.flags; @@ -2425,7 +2473,7 @@ get_mutation_samples(const tsk_treeseq_t *ts, const tsk_id_t *sites, tsk_size_t const tsk_size_t *restrict site_muts_len = ts->site_mutations_length; tsk_site_t site; tsk_tree_t tree; - tsk_bit_array_t all_samples_bits, mut_samples, mut_samples_row, out_row; + tsk_bitset_t all_samples_bits, mut_samples; tsk_size_t max_muts_len, site_offset, num_nodes, site_idx, s, m, n; tsk_id_t node, *nodes = NULL; void *tmp_nodes; @@ -2440,11 +2488,11 @@ get_mutation_samples(const tsk_treeseq_t *ts, const tsk_id_t *sites, tsk_size_t } } // Allocate a bit array of size max alleles for all sites - ret = tsk_bit_array_init(&mut_samples, num_samples, max_muts_len); + ret = tsk_bitset_init(&mut_samples, num_samples, max_muts_len); if (ret != 0) { goto out; } - ret = tsk_bit_array_init(&all_samples_bits, num_samples, 1); + ret = tsk_bitset_init(&all_samples_bits, num_samples, 1); if (ret != 0) { goto out; } @@ -2469,15 +2517,11 @@ get_mutation_samples(const tsk_treeseq_t *ts, const tsk_id_t *sites, tsk_size_t goto out; } nodes = tmp_nodes; - - tsk_bit_array_get_row(allele_samples, site_offset, &out_row); - tsk_bit_array_add(&out_row, &all_samples_bits); - + tsk_bitset_union(allele_samples, site_offset, &all_samples_bits, 0); // Zero out results before the start of each iteration tsk_memset(mut_samples.data, 0, - mut_samples.size * max_muts_len * sizeof(tsk_bit_array_value_t)); + mut_samples.row_len * max_muts_len * sizeof(tsk_bitset_val_t)); for (m = 0; m < site.mutations_length; m++) { - tsk_bit_array_get_row(&mut_samples, m, &mut_samples_row); node = site.mutations[m].node; ret = tsk_tree_preorder_from(&tree, node, nodes, &num_nodes); if (ret != 0) { @@ -2486,33 +2530,34 @@ get_mutation_samples(const tsk_treeseq_t *ts, const tsk_id_t *sites, tsk_size_t for (n = 0; n < num_nodes; n++) { node = nodes[n]; if (flags[node] & TSK_NODE_IS_SAMPLE) { - tsk_bit_array_add_bit(&mut_samples_row, - (tsk_bit_array_value_t) ts->sample_index_map[node]); + tsk_bitset_set_bit( + &mut_samples, m, (tsk_bitset_val_t) ts->sample_index_map[node]); } } } + get_allele_samples( + &site, site_offset, &mut_samples, allele_samples, &(num_alleles[site_idx])); site_offset += site.mutations_length + 1; - get_allele_samples(&site, &mut_samples, &out_row, &(num_alleles[site_idx])); } // if adding code below, check ret before continuing out: tsk_safe_free(nodes); tsk_tree_free(&tree); - tsk_bit_array_free(&mut_samples); - tsk_bit_array_free(&all_samples_bits); + tsk_bitset_free(&mut_samples); + tsk_bitset_free(&all_samples_bits); return ret == TSK_TREE_OK ? 0 : ret; } static int tsk_treeseq_two_site_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim, - const tsk_bit_array_t *sample_sets, tsk_size_t result_dim, general_stat_func_t *f, + const tsk_bitset_t *sample_sets, tsk_size_t result_dim, general_stat_func_t *f, sample_count_stat_params_t *f_params, norm_func_t *norm_f, tsk_size_t n_rows, const tsk_id_t *row_sites, tsk_size_t n_cols, const tsk_id_t *col_sites, tsk_flags_t options, double *result) { int ret = 0; - tsk_bit_array_t allele_samples, c_state, r_state; + tsk_bitset_t allele_samples; bool polarised = false; tsk_id_t *sites; tsk_size_t r, c, s, max_alleles, n_alleles, n_sites, *row_idx, *col_idx; @@ -2556,7 +2601,7 @@ tsk_treeseq_two_site_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim, if (ret != 0) { goto out; } - ret = tsk_bit_array_init(&allele_samples, num_samples, n_alleles); + ret = tsk_bitset_init(&allele_samples, num_samples, n_alleles); if (ret != 0) { goto out; } @@ -2573,9 +2618,8 @@ tsk_treeseq_two_site_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim, for (r = 0; r < n_rows; r++) { result_row = GET_2D_ROW(result, result_row_len, r); for (c = 0; c < n_cols; c++) { - tsk_bit_array_get_row(&allele_samples, site_offsets[row_idx[r]], &r_state); - tsk_bit_array_get_row(&allele_samples, site_offsets[col_idx[c]], &c_state); - ret = compute_general_two_site_stat_result(&r_state, &c_state, + ret = compute_general_two_site_stat_result(&allele_samples, + site_offsets[row_idx[r]], site_offsets[col_idx[c]], num_alleles[row_idx[r]], num_alleles[col_idx[c]], state_dim, sample_sets, result_dim, f, f_params, norm_f, polarised, &work, &(result_row[c * result_dim])); @@ -2592,37 +2636,37 @@ tsk_treeseq_two_site_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim, tsk_safe_free(num_alleles); tsk_safe_free(site_offsets); two_locus_work_free(&work); - tsk_bit_array_free(&allele_samples); + tsk_bitset_free(&allele_samples); return ret; } static int -sample_sets_to_bit_array(const tsk_treeseq_t *self, const tsk_size_t *sample_set_sizes, +sample_sets_to_bitset(const tsk_treeseq_t *self, const tsk_size_t *sample_set_sizes, const tsk_id_t *sample_sets, tsk_size_t num_sample_sets, - tsk_bit_array_t *sample_sets_bits) + tsk_bitset_t *sample_sets_bits) { int ret; - tsk_bit_array_t bits_row; + tsk_bitset_t bits_row; tsk_size_t j, k, l; tsk_id_t u, sample_index; - ret = tsk_bit_array_init(sample_sets_bits, self->num_samples, num_sample_sets); + ret = tsk_bitset_init(sample_sets_bits, self->num_samples, num_sample_sets); if (ret != 0) { return ret; } j = 0; for (k = 0; k < num_sample_sets; k++) { - tsk_bit_array_get_row(sample_sets_bits, k, &bits_row); + tsk_bitset_get_row(sample_sets_bits, k, &bits_row); for (l = 0; l < sample_set_sizes[k]; l++) { u = sample_sets[j]; sample_index = self->sample_index_map[u]; - if (tsk_bit_array_contains( - &bits_row, (tsk_bit_array_value_t) sample_index)) { + if (tsk_bitset_contains( + sample_sets_bits, k, (tsk_bitset_val_t) sample_index)) { ret = tsk_trace_error(TSK_ERR_DUPLICATE_SAMPLE); goto out; } - tsk_bit_array_add_bit(&bits_row, (tsk_bit_array_value_t) sample_index); + tsk_bitset_set_bit(sample_sets_bits, k, (tsk_bitset_val_t) sample_index); j++; } } @@ -2759,7 +2803,7 @@ get_index_counts( typedef struct { tsk_tree_t tree; - tsk_bit_array_t *node_samples; + tsk_bitset_t *node_samples; tsk_id_t *parent; tsk_id_t *edges_out; tsk_id_t *edges_in; @@ -2783,7 +2827,7 @@ iter_state_init(iter_state *self, const tsk_treeseq_t *ts, tsk_size_t state_dim) ret = tsk_trace_error(TSK_ERR_NO_MEMORY); goto out; } - ret = tsk_bit_array_init(self->node_samples, ts->num_samples, state_dim * num_nodes); + ret = tsk_bitset_init(self->node_samples, ts->num_samples, state_dim * num_nodes); if (ret != 0) { goto out; } @@ -2802,29 +2846,29 @@ iter_state_init(iter_state *self, const tsk_treeseq_t *ts, tsk_size_t state_dim) static int get_node_samples(const tsk_treeseq_t *ts, tsk_size_t state_dim, - const tsk_bit_array_t *sample_sets, tsk_bit_array_t *node_samples) + const tsk_bitset_t *sample_sets, tsk_bitset_t *node_samples) { int ret = 0; tsk_size_t n, k; - tsk_bit_array_t sample_set_row, node_samples_row; + tsk_bitset_t sample_set_row, node_samples_row; tsk_size_t num_nodes = ts->tables->nodes.num_rows; - tsk_bit_array_value_t sample; + tsk_bitset_val_t sample; const tsk_id_t *restrict sample_index_map = ts->sample_index_map; const tsk_flags_t *restrict flags = ts->tables->nodes.flags; - ret = tsk_bit_array_init(node_samples, ts->num_samples, num_nodes * state_dim); + ret = tsk_bitset_init(node_samples, ts->num_samples, num_nodes * state_dim); if (ret != 0) { goto out; } for (k = 0; k < state_dim; k++) { - tsk_bit_array_get_row(sample_sets, k, &sample_set_row); + tsk_bitset_get_row(sample_sets, k, &sample_set_row); for (n = 0; n < num_nodes; n++) { if (flags[n] & TSK_NODE_IS_SAMPLE) { - sample = (tsk_bit_array_value_t) sample_index_map[n]; - if (tsk_bit_array_contains(&sample_set_row, sample)) { - tsk_bit_array_get_row( + sample = (tsk_bitset_val_t) sample_index_map[n]; + if (tsk_bitset_contains(&sample_set_row, 0, sample)) { + tsk_bitset_get_row( node_samples, (state_dim * n) + k, &node_samples_row); - tsk_bit_array_add_bit(&node_samples_row, sample); + tsk_bitset_set_bit(&node_samples_row, 0, sample); } } } @@ -2835,7 +2879,7 @@ get_node_samples(const tsk_treeseq_t *ts, tsk_size_t state_dim, static void iter_state_clear(iter_state *self, tsk_size_t state_dim, tsk_size_t num_nodes, - const tsk_bit_array_t *node_samples) + const tsk_bitset_t *node_samples) { self->n_edges_out = 0; self->n_edges_in = 0; @@ -2845,14 +2889,14 @@ iter_state_clear(iter_state *self, tsk_size_t state_dim, tsk_size_t num_nodes, tsk_memset(self->edges_in, TSK_NULL, num_nodes * sizeof(*self->edges_in)); tsk_memset(self->branch_len, 0, num_nodes * sizeof(*self->branch_len)); tsk_memcpy(self->node_samples->data, node_samples->data, - node_samples->size * state_dim * num_nodes * sizeof(*node_samples->data)); + node_samples->row_len * state_dim * num_nodes * sizeof(*node_samples->data)); } static void iter_state_free(iter_state *self) { tsk_tree_free(&self->tree); - tsk_bit_array_free(self->node_samples); + tsk_bitset_free(self->node_samples); tsk_safe_free(self->node_samples); tsk_safe_free(self->parent); tsk_safe_free(self->edges_out); @@ -2939,10 +2983,10 @@ compute_two_tree_branch_state_update(const tsk_treeseq_t *ts, tsk_id_t c, double *restrict B_branch_len = B_state->branch_len; double *weights = NULL, *weights_row, *result_tmp = NULL; tsk_size_t n, k, a_row, b_row; - tsk_bit_array_t A_samples, B_samples, AB_samples, B_samples_tmp; + tsk_bitset_t AB_samples, B_samples_tmp; const double *restrict A_branch_len = A_state->branch_len; - const tsk_bit_array_t *restrict A_state_samples = A_state->node_samples; - const tsk_bit_array_t *restrict B_state_samples = B_state->node_samples; + const tsk_bitset_t *restrict A_state_samples = A_state->node_samples; + const tsk_bitset_t *restrict B_state_samples = B_state->node_samples; tsk_size_t num_samples = ts->num_samples; tsk_size_t num_nodes = ts->tables->nodes.num_rows; b_len = B_branch_len[c] * sign; @@ -2959,11 +3003,11 @@ compute_two_tree_branch_state_update(const tsk_treeseq_t *ts, tsk_id_t c, ret = tsk_trace_error(TSK_ERR_NO_MEMORY); goto out; } - ret = tsk_bit_array_init(&AB_samples, num_samples, 1); + ret = tsk_bitset_init(&AB_samples, num_samples, 1); if (ret != 0) { goto out; } - ret = tsk_bit_array_init(&B_samples_tmp, num_samples, 1); + ret = tsk_bitset_init(&B_samples_tmp, num_samples, 1); if (ret != 0) { goto out; } @@ -2975,15 +3019,14 @@ compute_two_tree_branch_state_update(const tsk_treeseq_t *ts, tsk_id_t c, for (k = 0; k < state_dim; k++) { a_row = (state_dim * n) + k; b_row = (state_dim * (tsk_size_t) c) + k; - tsk_bit_array_get_row(A_state_samples, a_row, &A_samples); - tsk_bit_array_get_row(B_state_samples, b_row, &B_samples); - tsk_bit_array_intersect(&A_samples, &B_samples, &AB_samples); + tsk_bitset_intersect( + A_state_samples, a_row, B_state_samples, b_row, &AB_samples); weights_row = GET_2D_ROW(weights, 3, k); - weights_row[0] = (double) tsk_bit_array_count(&AB_samples); // w_AB + weights_row[0] = (double) tsk_bitset_count(&AB_samples, 0); // w_AB weights_row[1] - = (double) tsk_bit_array_count(&A_samples) - weights_row[0]; // w_Ab + = (double) tsk_bitset_count(A_state_samples, a_row) - weights_row[0]; weights_row[2] - = (double) tsk_bit_array_count(&B_samples) - weights_row[0]; // w_aB + = (double) tsk_bitset_count(B_state_samples, b_row) - weights_row[0]; } ret = f(state_dim, weights, result_dim, result_tmp, f_params); if (ret != 0) { @@ -2996,8 +3039,8 @@ compute_two_tree_branch_state_update(const tsk_treeseq_t *ts, tsk_id_t c, out: tsk_safe_free(weights); tsk_safe_free(result_tmp); - tsk_bit_array_free(&AB_samples); - tsk_bit_array_free(&B_samples_tmp); + tsk_bitset_free(&AB_samples); + tsk_bitset_free(&B_samples_tmp); return ret; } @@ -3013,10 +3056,10 @@ compute_two_tree_branch_stat(const tsk_treeseq_t *ts, const iter_state *l_state, const tsk_id_t *restrict edges_child = ts->tables->edges.child; const tsk_id_t *restrict edges_parent = ts->tables->edges.parent; const tsk_size_t num_nodes = ts->tables->nodes.num_rows; - tsk_bit_array_t updates, row, ec_row, *r_samples = r_state->node_samples; + tsk_bitset_t updates, row, ec_row, *r_samples = r_state->node_samples; tsk_memset(&updates, 0, sizeof(updates)); - ret = tsk_bit_array_init(&updates, num_nodes, 1); + ret = tsk_bitset_init(&updates, num_nodes, 1); if (ret != 0) { goto out; } @@ -3033,13 +3076,13 @@ compute_two_tree_branch_stat(const tsk_treeseq_t *ts, const iter_state *l_state, c = edges_child[e]; // Identify affected nodes above child while (p != TSK_NULL) { - tsk_bit_array_add_bit(&updates, (tsk_bit_array_value_t) c); + tsk_bitset_set_bit(&updates, 0, (tsk_bitset_val_t) c); c = p; p = r_state->parent[p]; } } // Subtract the whole contribution from the child node - tsk_bit_array_get_items(&updates, updated_nodes, &n_updates); + tsk_bitset_get_items(&updates, updated_nodes, &n_updates); while (n_updates != 0) { n_updates--; c = updated_nodes[n_updates]; @@ -3053,10 +3096,10 @@ compute_two_tree_branch_stat(const tsk_treeseq_t *ts, const iter_state *l_state, ec = edges_child[e]; // edge child while (p != TSK_NULL) { for (k = 0; k < state_dim; k++) { - tsk_bit_array_get_row( + tsk_bitset_get_row( r_samples, (state_dim * (tsk_size_t) ec) + k, &ec_row); - tsk_bit_array_get_row(r_samples, (state_dim * (tsk_size_t) p) + k, &row); - tsk_bit_array_subtract(&row, &ec_row); + tsk_bitset_get_row(r_samples, (state_dim * (tsk_size_t) p) + k, &row); + tsk_bitset_subtract(&row, 0, &ec_row, 0); } p = r_state->parent[p]; } @@ -3071,12 +3114,12 @@ compute_two_tree_branch_stat(const tsk_treeseq_t *ts, const iter_state *l_state, r_state->branch_len[c] = time[p] - time[c]; r_state->parent[c] = p; while (p != TSK_NULL) { - tsk_bit_array_add_bit(&updates, (tsk_bit_array_value_t) c); + tsk_bitset_set_bit(&updates, 0, (tsk_bitset_val_t) c); for (k = 0; k < state_dim; k++) { - tsk_bit_array_get_row( + tsk_bitset_get_row( r_samples, (state_dim * (tsk_size_t) ec) + k, &ec_row); - tsk_bit_array_get_row(r_samples, (state_dim * (tsk_size_t) p) + k, &row); - tsk_bit_array_add(&row, &ec_row); + tsk_bitset_get_row(r_samples, (state_dim * (tsk_size_t) p) + k, &row); + tsk_bitset_union(&row, 0, &ec_row, 0); } c = p; p = r_state->parent[p]; @@ -3084,7 +3127,7 @@ compute_two_tree_branch_stat(const tsk_treeseq_t *ts, const iter_state *l_state, } // Update all affected child nodes (fully subtracted, deferred from addition) n_updates = 0; - tsk_bit_array_get_items(&updates, updated_nodes, &n_updates); + tsk_bitset_get_items(&updates, updated_nodes, &n_updates); while (n_updates != 0) { n_updates--; c = updated_nodes[n_updates]; @@ -3093,13 +3136,13 @@ compute_two_tree_branch_stat(const tsk_treeseq_t *ts, const iter_state *l_state, } out: tsk_safe_free(updated_nodes); - tsk_bit_array_free(&updates); + tsk_bitset_free(&updates); return ret; } static int tsk_treeseq_two_branch_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim, - const tsk_bit_array_t *sample_sets, tsk_size_t result_dim, general_stat_func_t *f, + const tsk_bitset_t *sample_sets, tsk_size_t result_dim, general_stat_func_t *f, sample_count_stat_params_t *f_params, norm_func_t *TSK_UNUSED(norm_f), tsk_size_t n_rows, const double *row_positions, tsk_size_t n_cols, const double *col_positions, tsk_flags_t TSK_UNUSED(options), double *result) @@ -3108,7 +3151,7 @@ tsk_treeseq_two_branch_count_stat(const tsk_treeseq_t *self, tsk_size_t state_di int r, c; tsk_id_t *row_indexes = NULL, *col_indexes = NULL; tsk_size_t i, j, k, row, col, *row_repeats = NULL, *col_repeats = NULL; - tsk_bit_array_t node_samples; + tsk_bitset_t node_samples; iter_state l_state, r_state; double *result_tmp = NULL, *result_row; const tsk_size_t num_nodes = self->tables->nodes.num_rows; @@ -3196,7 +3239,7 @@ tsk_treeseq_two_branch_count_stat(const tsk_treeseq_t *self, tsk_size_t state_di tsk_safe_free(col_repeats); iter_state_free(&l_state); iter_state_free(&r_state); - tsk_bit_array_free(&node_samples); + tsk_bitset_free(&node_samples); return ret; } @@ -3211,7 +3254,7 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl // TODO: generalize this function if we ever decide to do weighted two_locus stats. // We only implement count stats and therefore we don't handle weights. int ret = 0; - tsk_bit_array_t sample_sets_bits; + tsk_bitset_t sample_sets_bits; bool stat_site = !!(options & TSK_STAT_SITE); bool stat_branch = !!(options & TSK_STAT_BRANCH); tsk_size_t state_dim = num_sample_sets; @@ -3245,7 +3288,7 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl ret = tsk_trace_error(TSK_ERR_BAD_RESULT_DIMS); goto out; } - ret = sample_sets_to_bit_array( + ret = sample_sets_to_bitset( self, sample_set_sizes, sample_sets, num_sample_sets, &sample_sets_bits); if (ret != 0) { goto out; @@ -3280,7 +3323,7 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl } out: - tsk_bit_array_free(&sample_sets_bits); + tsk_bitset_free(&sample_sets_bits); return ret; } From 6f059668ce90156c4f4b13a48a332b9d24f338a0 Mon Sep 17 00:00:00 2001 From: lkirk Date: Thu, 24 Jul 2025 18:58:06 -0500 Subject: [PATCH 38/70] remove temporary references to rows --- c/tskit/core.c | 33 ++++++++++-------- c/tskit/core.h | 2 ++ c/tskit/trees.c | 91 ++++++++++--------------------------------------- 3 files changed, 38 insertions(+), 88 deletions(-) diff --git a/c/tskit/core.c b/c/tskit/core.c index 785f590133..d6b226415b 100644 --- a/c/tskit/core.c +++ b/c/tskit/core.c @@ -1284,9 +1284,6 @@ tsk_bitset_intersect(const tsk_bitset_t *self, tsk_size_t self_row, { tsk_bitset_val_t *self_d = self->data + (self_row * self->row_len); tsk_bitset_val_t *other_d = other->data + (other_row * self->row_len); - if (self_row >= self->len || other_row >= other->len) { - printf("isect: OUTSIDE_OF_BOUNDS"); - } for (tsk_size_t i = 0; i < self->row_len; i++) { out->data[i] = self_d[i] & other_d[i]; } @@ -1298,9 +1295,6 @@ tsk_bitset_subtract(tsk_bitset_t *self, tsk_size_t self_row, const tsk_bitset_t { tsk_bitset_val_t *self_d = self->data + (self_row * self->row_len); tsk_bitset_val_t *other_d = other->data + (other_row * self->row_len); - if (self_row >= self->len || other_row >= other->len) { - printf("subtr: OUTSIDE_OF_BOUNDS"); - } for (tsk_size_t i = 0; i < self->row_len; i++) { self_d[i] &= ~(other_d[i]); } @@ -1312,9 +1306,6 @@ tsk_bitset_union(tsk_bitset_t *self, tsk_size_t self_row, const tsk_bitset_t *ot { tsk_bitset_val_t *self_d = self->data + (self_row * self->row_len); tsk_bitset_val_t *other_d = other->data + (other_row * self->row_len); - if (self_row >= self->len || other_row >= other->len) { - printf("union: OUTSIDE_OF_BOUNDS"); - } for (tsk_size_t i = 0; i < self->row_len; i++) { self_d[i] |= other_d[i]; } @@ -1332,13 +1323,28 @@ bool tsk_bitset_contains(const tsk_bitset_t *self, tsk_size_t row, const tsk_bitset_val_t bit) { tsk_bitset_val_t i = (bit >> TSK_BIT_ARRAY_CHUNK); - if (row >= self->len) { - printf("contains: OUTSIDE_OF_BOUNDS"); - } return self->data[i + row * self->row_len] & ((tsk_bitset_val_t) 1 << (bit - (TSK_BIT_ARRAY_NUM_BITS * i))); } +tsk_size_t +tsk_bitset_isect_and_count(const tsk_bitset_t *self, tsk_size_t self_row, + const tsk_bitset_t *other, tsk_size_t other_row) +{ + tsk_bitset_val_t tmp; + tsk_size_t i, count = 0; + tsk_bitset_val_t *self_d = self->data + (self_row * self->row_len); + tsk_bitset_val_t *other_d = other->data + (other_row * self->row_len); + + for (i = 0; i < self->row_len; i++) { + tmp = self_d[i] & other_d[i]; + tmp = tmp - ((tmp >> 1) & 0x55555555); + tmp = (tmp & 0x33333333) + ((tmp >> 2) & 0x33333333); + count += (((tmp + (tmp >> 4)) & 0xF0F0F0F) * 0x1010101) >> 24; + } + return count; +} + tsk_size_t tsk_bitset_count(const tsk_bitset_t *self, tsk_size_t row) { @@ -1358,9 +1364,6 @@ tsk_bitset_count(const tsk_bitset_t *self, tsk_size_t row) tsk_size_t i, count = 0; tsk_bitset_val_t *self_d = self->data + (row * self->row_len); - if (row >= self->len) { - printf("count: OUTSIDE_OF_BOUNDS"); - } for (i = 0; i < self->row_len; i++) { tmp = self_d[i] - ((self_d[i] >> 1) & 0x55555555); tmp = (tmp & 0x33333333) + ((tmp >> 2) & 0x33333333); diff --git a/c/tskit/core.h b/c/tskit/core.h index c7f1c6ea89..272d514749 100644 --- a/c/tskit/core.h +++ b/c/tskit/core.h @@ -1120,6 +1120,8 @@ bool tsk_bitset_contains( tsk_size_t tsk_bitset_count(const tsk_bitset_t *self, tsk_size_t row); void tsk_bitset_get_items( const tsk_bitset_t *self, tsk_id_t *items, tsk_size_t *n_items); +tsk_size_t tsk_bitset_isect_and_count(const tsk_bitset_t *self, tsk_size_t self_row, + const tsk_bitset_t *other, tsk_size_t other_row); #ifdef __cplusplus } diff --git a/c/tskit/trees.c b/c/tskit/trees.c index fd759b05c5..8ac19ddd89 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -2297,9 +2297,6 @@ typedef struct { double *norm; double *result_tmp; tsk_bitset_t AB_samples; - tsk_bitset_t ss_A_samples; - tsk_bitset_t ss_B_samples; - tsk_bitset_t ss_AB_samples; } two_locus_work_t; static int @@ -2311,33 +2308,15 @@ two_locus_work_init(tsk_size_t max_alleles, tsk_size_t result_dim, tsk_size_t st out->norm = tsk_malloc(result_dim * sizeof(*out->norm)); out->result_tmp = tsk_malloc(result_dim * max_alleles * max_alleles * sizeof(*out->result_tmp)); - - tsk_memset(&out->ss_A_samples, 0, sizeof(out->ss_A_samples)); - tsk_memset(&out->ss_B_samples, 0, sizeof(out->ss_B_samples)); - tsk_memset(&out->ss_AB_samples, 0, sizeof(out->ss_AB_samples)); tsk_memset(&out->AB_samples, 0, sizeof(out->AB_samples)); - if (out->weights == NULL || out->norm == NULL || out->result_tmp == NULL) { ret = tsk_trace_error(TSK_ERR_NO_MEMORY); goto out; } - ret = tsk_bitset_init(&out->AB_samples, num_samples, 1); if (ret != 0) { goto out; } - ret = tsk_bitset_init(&out->ss_A_samples, num_samples, 1); - if (ret != 0) { - goto out; - } - ret = tsk_bitset_init(&out->ss_B_samples, num_samples, 1); - if (ret != 0) { - goto out; - } - ret = tsk_bitset_init(&out->ss_AB_samples, num_samples, 1); - if (ret != 0) { - goto out; - } out: return ret; } @@ -2349,9 +2328,6 @@ two_locus_work_free(two_locus_work_t *work) tsk_safe_free(work->norm); tsk_safe_free(work->result_tmp); tsk_bitset_free(&work->AB_samples); - tsk_bitset_free(&work->ss_A_samples); - tsk_bitset_free(&work->ss_B_samples); - tsk_bitset_free(&work->ss_AB_samples); } static int @@ -2369,7 +2345,6 @@ compute_general_two_site_stat_result(const tsk_bitset_t *state, tsk_size_t a_off // a3 [s1, s2, s3] [s1, s2, s3] [s1, s2, s3] tsk_size_t k, mut_a, mut_b; tsk_size_t result_row_len = num_b_alleles * result_dim; - double w_AB = 0; uint8_t polarised_val = polarised ? 1 : 0; double *hap_weight_row; double *result_tmp_row; @@ -2378,9 +2353,6 @@ compute_general_two_site_stat_result(const tsk_bitset_t *state, tsk_size_t a_off double *weights = work->weights; double *result_tmp = work->result_tmp; tsk_bitset_t AB_samples = work->AB_samples; - tsk_bitset_t ss_A_samples = work->ss_A_samples; - tsk_bitset_t ss_B_samples = work->ss_B_samples; - tsk_bitset_t ss_AB_samples = work->ss_AB_samples; for (mut_a = polarised_val; mut_a < num_a_alleles; mut_a++) { result_tmp_row = GET_2D_ROW(result_tmp, result_row_len, mut_a); @@ -2389,14 +2361,14 @@ compute_general_two_site_stat_result(const tsk_bitset_t *state, tsk_size_t a_off state, mut_a + a_off, state, mut_b + b_off, &AB_samples); for (k = 0; k < state_dim; k++) { hap_weight_row = GET_2D_ROW(weights, 3, k); - tsk_bitset_intersect( - state, mut_a + a_off, sample_sets, k, &ss_A_samples); - tsk_bitset_intersect( - state, mut_b + b_off, sample_sets, k, &ss_B_samples); - tsk_bitset_intersect(&AB_samples, 0, sample_sets, k, &ss_AB_samples); - w_AB = hap_weight_row[0] = (double) tsk_bitset_count(&ss_AB_samples, 0); - hap_weight_row[1] = ((double) tsk_bitset_count(&ss_A_samples, 0) - w_AB); - hap_weight_row[2] = ((double) tsk_bitset_count(&ss_B_samples, 0) - w_AB); + hap_weight_row[0] = (double) tsk_bitset_isect_and_count( + &AB_samples, 0, sample_sets, k); + hap_weight_row[1] = (double) tsk_bitset_isect_and_count( + state, mut_a + a_off, sample_sets, k) + - hap_weight_row[0]; + hap_weight_row[2] = (double) tsk_bitset_isect_and_count( + state, mut_b + b_off, sample_sets, k) + - hap_weight_row[0]; } ret = f(state_dim, weights, result_dim, result_tmp_row, f_params); if (ret != 0) { @@ -2646,7 +2618,6 @@ sample_sets_to_bitset(const tsk_treeseq_t *self, const tsk_size_t *sample_set_si tsk_bitset_t *sample_sets_bits) { int ret; - tsk_bitset_t bits_row; tsk_size_t j, k, l; tsk_id_t u, sample_index; @@ -2654,10 +2625,8 @@ sample_sets_to_bitset(const tsk_treeseq_t *self, const tsk_size_t *sample_set_si if (ret != 0) { return ret; } - j = 0; for (k = 0; k < num_sample_sets; k++) { - tsk_bitset_get_row(sample_sets_bits, k, &bits_row); for (l = 0; l < sample_set_sizes[k]; l++) { u = sample_sets[j]; sample_index = self->sample_index_map[u]; @@ -2850,7 +2819,6 @@ get_node_samples(const tsk_treeseq_t *ts, tsk_size_t state_dim, { int ret = 0; tsk_size_t n, k; - tsk_bitset_t sample_set_row, node_samples_row; tsk_size_t num_nodes = ts->tables->nodes.num_rows; tsk_bitset_val_t sample; const tsk_id_t *restrict sample_index_map = ts->sample_index_map; @@ -2861,14 +2829,11 @@ get_node_samples(const tsk_treeseq_t *ts, tsk_size_t state_dim, goto out; } for (k = 0; k < state_dim; k++) { - tsk_bitset_get_row(sample_sets, k, &sample_set_row); for (n = 0; n < num_nodes; n++) { if (flags[n] & TSK_NODE_IS_SAMPLE) { sample = (tsk_bitset_val_t) sample_index_map[n]; - if (tsk_bitset_contains(&sample_set_row, 0, sample)) { - tsk_bitset_get_row( - node_samples, (state_dim * n) + k, &node_samples_row); - tsk_bitset_set_bit(&node_samples_row, 0, sample); + if (tsk_bitset_contains(sample_sets, k, sample)) { + tsk_bitset_set_bit(node_samples, (state_dim * n) + k, sample); } } } @@ -2983,34 +2948,21 @@ compute_two_tree_branch_state_update(const tsk_treeseq_t *ts, tsk_id_t c, double *restrict B_branch_len = B_state->branch_len; double *weights = NULL, *weights_row, *result_tmp = NULL; tsk_size_t n, k, a_row, b_row; - tsk_bitset_t AB_samples, B_samples_tmp; const double *restrict A_branch_len = A_state->branch_len; const tsk_bitset_t *restrict A_state_samples = A_state->node_samples; const tsk_bitset_t *restrict B_state_samples = B_state->node_samples; - tsk_size_t num_samples = ts->num_samples; tsk_size_t num_nodes = ts->tables->nodes.num_rows; + b_len = B_branch_len[c] * sign; if (b_len == 0) { return ret; } - - tsk_memset(&AB_samples, 0, sizeof(AB_samples)); - tsk_memset(&B_samples_tmp, 0, sizeof(B_samples_tmp)); - weights = tsk_calloc(3 * state_dim, sizeof(*weights)); result_tmp = tsk_calloc(result_dim, sizeof(*result_tmp)); if (weights == NULL || result_tmp == NULL) { ret = tsk_trace_error(TSK_ERR_NO_MEMORY); goto out; } - ret = tsk_bitset_init(&AB_samples, num_samples, 1); - if (ret != 0) { - goto out; - } - ret = tsk_bitset_init(&B_samples_tmp, num_samples, 1); - if (ret != 0) { - goto out; - } for (n = 0; n < num_nodes; n++) { a_len = A_branch_len[n]; if (a_len == 0) { @@ -3019,10 +2971,9 @@ compute_two_tree_branch_state_update(const tsk_treeseq_t *ts, tsk_id_t c, for (k = 0; k < state_dim; k++) { a_row = (state_dim * n) + k; b_row = (state_dim * (tsk_size_t) c) + k; - tsk_bitset_intersect( - A_state_samples, a_row, B_state_samples, b_row, &AB_samples); weights_row = GET_2D_ROW(weights, 3, k); - weights_row[0] = (double) tsk_bitset_count(&AB_samples, 0); // w_AB + weights_row[0] = (double) tsk_bitset_isect_and_count( + A_state_samples, a_row, B_state_samples, b_row); weights_row[1] = (double) tsk_bitset_count(A_state_samples, a_row) - weights_row[0]; weights_row[2] @@ -3039,8 +2990,6 @@ compute_two_tree_branch_state_update(const tsk_treeseq_t *ts, tsk_id_t c, out: tsk_safe_free(weights); tsk_safe_free(result_tmp); - tsk_bitset_free(&AB_samples); - tsk_bitset_free(&B_samples_tmp); return ret; } @@ -3056,7 +3005,7 @@ compute_two_tree_branch_stat(const tsk_treeseq_t *ts, const iter_state *l_state, const tsk_id_t *restrict edges_child = ts->tables->edges.child; const tsk_id_t *restrict edges_parent = ts->tables->edges.parent; const tsk_size_t num_nodes = ts->tables->nodes.num_rows; - tsk_bitset_t updates, row, ec_row, *r_samples = r_state->node_samples; + tsk_bitset_t updates, *r_samples = r_state->node_samples; tsk_memset(&updates, 0, sizeof(updates)); ret = tsk_bitset_init(&updates, num_nodes, 1); @@ -3096,10 +3045,8 @@ compute_two_tree_branch_stat(const tsk_treeseq_t *ts, const iter_state *l_state, ec = edges_child[e]; // edge child while (p != TSK_NULL) { for (k = 0; k < state_dim; k++) { - tsk_bitset_get_row( - r_samples, (state_dim * (tsk_size_t) ec) + k, &ec_row); - tsk_bitset_get_row(r_samples, (state_dim * (tsk_size_t) p) + k, &row); - tsk_bitset_subtract(&row, 0, &ec_row, 0); + tsk_bitset_subtract(r_samples, (state_dim * (tsk_size_t) p) + k, + r_samples, (state_dim * (tsk_size_t) ec) + k); } p = r_state->parent[p]; } @@ -3116,10 +3063,8 @@ compute_two_tree_branch_stat(const tsk_treeseq_t *ts, const iter_state *l_state, while (p != TSK_NULL) { tsk_bitset_set_bit(&updates, 0, (tsk_bitset_val_t) c); for (k = 0; k < state_dim; k++) { - tsk_bitset_get_row( - r_samples, (state_dim * (tsk_size_t) ec) + k, &ec_row); - tsk_bitset_get_row(r_samples, (state_dim * (tsk_size_t) p) + k, &row); - tsk_bitset_union(&row, 0, &ec_row, 0); + tsk_bitset_union(r_samples, (state_dim * (tsk_size_t) p) + k, r_samples, + (state_dim * (tsk_size_t) ec) + k); } c = p; p = r_state->parent[p]; From 83a1c856b07311a9cb406b131619dccb6b4b53ab Mon Sep 17 00:00:00 2001 From: lkirk Date: Thu, 24 Jul 2025 19:16:22 -0500 Subject: [PATCH 39/70] remove alloc in branch hot path --- c/tskit/trees.c | 38 ++++++++++++++++++++++---------------- 1 file changed, 22 insertions(+), 16 deletions(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index 8ac19ddd89..1187e61503 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -2537,9 +2537,10 @@ tsk_treeseq_two_site_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim, const tsk_size_t num_samples = self->num_samples; tsk_size_t *num_alleles = NULL, *site_offsets = NULL; tsk_size_t result_row_len = n_cols * result_dim; - two_locus_work_t work = { 0 }; - tsk_memset(&allele_samples, 0, sizeof(allele_samples)); + two_locus_work_t work; + tsk_memset(&work, 0, sizeof(work)); + tsk_memset(&allele_samples, 0, sizeof(allele_samples)); sites = tsk_malloc(self->tables->sites.num_rows * sizeof(*sites)); row_idx = tsk_malloc(self->tables->sites.num_rows * sizeof(*row_idx)); col_idx = tsk_malloc(self->tables->sites.num_rows * sizeof(*col_idx)); @@ -2577,6 +2578,7 @@ tsk_treeseq_two_site_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim, if (ret != 0) { goto out; } + // we track the number of alleles to account for backmutations ret = get_mutation_samples(self, sites, n_sites, num_alleles, &allele_samples); if (ret != 0) { goto out; @@ -2941,28 +2943,25 @@ static int compute_two_tree_branch_state_update(const tsk_treeseq_t *ts, tsk_id_t c, const iter_state *A_state, const iter_state *B_state, tsk_size_t state_dim, tsk_size_t result_dim, int sign, general_stat_func_t *f, - sample_count_stat_params_t *f_params, double *result) + sample_count_stat_params_t *f_params, two_locus_work_t *restrict work, + double *result) { int ret = 0; double a_len, b_len; double *restrict B_branch_len = B_state->branch_len; - double *weights = NULL, *weights_row, *result_tmp = NULL; + double *weights_row; tsk_size_t n, k, a_row, b_row; const double *restrict A_branch_len = A_state->branch_len; const tsk_bitset_t *restrict A_state_samples = A_state->node_samples; const tsk_bitset_t *restrict B_state_samples = B_state->node_samples; tsk_size_t num_nodes = ts->tables->nodes.num_rows; + double *weights = work->weights; + double *result_tmp = work->result_tmp; b_len = B_branch_len[c] * sign; if (b_len == 0) { return ret; } - weights = tsk_calloc(3 * state_dim, sizeof(*weights)); - result_tmp = tsk_calloc(result_dim, sizeof(*result_tmp)); - if (weights == NULL || result_tmp == NULL) { - ret = tsk_trace_error(TSK_ERR_NO_MEMORY); - goto out; - } for (n = 0; n < num_nodes; n++) { a_len = A_branch_len[n]; if (a_len == 0) { @@ -2988,8 +2987,6 @@ compute_two_tree_branch_state_update(const tsk_treeseq_t *ts, tsk_id_t c, } } out: - tsk_safe_free(weights); - tsk_safe_free(result_tmp); return ret; } @@ -3005,9 +3002,17 @@ compute_two_tree_branch_stat(const tsk_treeseq_t *ts, const iter_state *l_state, const tsk_id_t *restrict edges_child = ts->tables->edges.child; const tsk_id_t *restrict edges_parent = ts->tables->edges.parent; const tsk_size_t num_nodes = ts->tables->nodes.num_rows; + const tsk_size_t num_samples = ts->num_samples; tsk_bitset_t updates, *r_samples = r_state->node_samples; + two_locus_work_t work; + tsk_memset(&work, 0, sizeof(work)); tsk_memset(&updates, 0, sizeof(updates)); + // only two alleles are possible for branch stats + ret = two_locus_work_init(2, result_dim, state_dim, num_samples, &work); + if (ret != 0) { + goto out; + } ret = tsk_bitset_init(&updates, num_nodes, 1); if (ret != 0) { goto out; @@ -3035,8 +3040,8 @@ compute_two_tree_branch_stat(const tsk_treeseq_t *ts, const iter_state *l_state, while (n_updates != 0) { n_updates--; c = updated_nodes[n_updates]; - compute_two_tree_branch_state_update( - ts, c, l_state, r_state, state_dim, result_dim, -1, f, f_params, result); + compute_two_tree_branch_state_update(ts, c, l_state, r_state, state_dim, + result_dim, -1, f, f_params, &work, result); } // Remove samples under nodes from removed edges to parent nodes for (j = 0; j < r_state->n_edges_out; j++) { @@ -3076,11 +3081,12 @@ compute_two_tree_branch_stat(const tsk_treeseq_t *ts, const iter_state *l_state, while (n_updates != 0) { n_updates--; c = updated_nodes[n_updates]; - compute_two_tree_branch_state_update( - ts, c, l_state, r_state, state_dim, result_dim, +1, f, f_params, result); + compute_two_tree_branch_state_update(ts, c, l_state, r_state, state_dim, + result_dim, +1, f, f_params, &work, result); } out: tsk_safe_free(updated_nodes); + two_locus_work_free(&work); tsk_bitset_free(&updates); return ret; } From 19b5f74725b8971293d03b3d29d587c26339a254 Mon Sep 17 00:00:00 2001 From: lkirk Date: Mon, 28 Jul 2025 00:01:19 -0500 Subject: [PATCH 40/70] remove isect_and_count; it adds data dependencies and we reduce the potential for instruction level parallelism; also add some comments --- c/tskit/core.c | 35 +++++++++------------------ c/tskit/core.h | 2 -- c/tskit/trees.c | 63 ++++++++++++++++++++++++++++++++++--------------- 3 files changed, 55 insertions(+), 45 deletions(-) diff --git a/c/tskit/core.c b/c/tskit/core.c index d6b226415b..e396ab0626 100644 --- a/c/tskit/core.c +++ b/c/tskit/core.c @@ -1327,28 +1327,10 @@ tsk_bitset_contains(const tsk_bitset_t *self, tsk_size_t row, const tsk_bitset_v & ((tsk_bitset_val_t) 1 << (bit - (TSK_BIT_ARRAY_NUM_BITS * i))); } -tsk_size_t -tsk_bitset_isect_and_count(const tsk_bitset_t *self, tsk_size_t self_row, - const tsk_bitset_t *other, tsk_size_t other_row) -{ - tsk_bitset_val_t tmp; - tsk_size_t i, count = 0; - tsk_bitset_val_t *self_d = self->data + (self_row * self->row_len); - tsk_bitset_val_t *other_d = other->data + (other_row * self->row_len); - - for (i = 0; i < self->row_len; i++) { - tmp = self_d[i] & other_d[i]; - tmp = tmp - ((tmp >> 1) & 0x55555555); - tmp = (tmp & 0x33333333) + ((tmp >> 2) & 0x33333333); - count += (((tmp + (tmp >> 4)) & 0xF0F0F0F) * 0x1010101) >> 24; - } - return count; -} - tsk_size_t tsk_bitset_count(const tsk_bitset_t *self, tsk_size_t row) { - // Utilizes 12 operations per bit array. NB this only works on 32 bit integers. + // Utilizes 12 operations per chunk. NB this only works on 32 bit integers. // Taken from: // https://graphics.stanford.edu/~seander/bithacks.html#CountBitsSetParallel // There's a nice breakdown of this algorithm here: @@ -1356,12 +1338,17 @@ tsk_bitset_count(const tsk_bitset_t *self, tsk_size_t row) // Could probably do better with explicit SIMD (instead of SWAR), but not as // portable: https://arxiv.org/pdf/1611.07612.pdf // - // There is one solution to explore further, which uses __builtin_popcountll. - // This option is relatively simple, but requires a 64 bit bit array and also - // involves some compiler flag plumbing (-mpopcnt) + // The gcc/clang compiler flag will -mpopcnt will convert this code to a + // popcnt instruction (most if not all modern CPUs will support this). The + // popcnt instruction will yield some speed improvements, which depend on + // the tree sequence. + // + // NB: 32bit counting is typically faster than 64bit counting for this task. + // (at least on x86-64) tsk_bitset_val_t tmp; - tsk_size_t i, count = 0; + tsk_size_t i = 0; + uint32_t count = 0; tsk_bitset_val_t *self_d = self->data + (row * self->row_len); for (i = 0; i < self->row_len; i++) { @@ -1369,7 +1356,7 @@ tsk_bitset_count(const tsk_bitset_t *self, tsk_size_t row) tmp = (tmp & 0x33333333) + ((tmp >> 2) & 0x33333333); count += (((tmp + (tmp >> 4)) & 0xF0F0F0F) * 0x1010101) >> 24; } - return count; + return (tsk_size_t) count; } void diff --git a/c/tskit/core.h b/c/tskit/core.h index 272d514749..c7f1c6ea89 100644 --- a/c/tskit/core.h +++ b/c/tskit/core.h @@ -1120,8 +1120,6 @@ bool tsk_bitset_contains( tsk_size_t tsk_bitset_count(const tsk_bitset_t *self, tsk_size_t row); void tsk_bitset_get_items( const tsk_bitset_t *self, tsk_id_t *items, tsk_size_t *n_items); -tsk_size_t tsk_bitset_isect_and_count(const tsk_bitset_t *self, tsk_size_t self_row, - const tsk_bitset_t *other, tsk_size_t other_row); #ifdef __cplusplus } diff --git a/c/tskit/trees.c b/c/tskit/trees.c index 1187e61503..e410418a3e 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -2297,6 +2297,9 @@ typedef struct { double *norm; double *result_tmp; tsk_bitset_t AB_samples; + tsk_bitset_t ss_A_samples; + tsk_bitset_t ss_B_samples; + tsk_bitset_t ss_AB_samples; } two_locus_work_t; static int @@ -2308,11 +2311,28 @@ two_locus_work_init(tsk_size_t max_alleles, tsk_size_t result_dim, tsk_size_t st out->norm = tsk_malloc(result_dim * sizeof(*out->norm)); out->result_tmp = tsk_malloc(result_dim * max_alleles * max_alleles * sizeof(*out->result_tmp)); + + tsk_memset(&out->ss_A_samples, 0, sizeof(out->ss_A_samples)); + tsk_memset(&out->ss_B_samples, 0, sizeof(out->ss_B_samples)); + tsk_memset(&out->ss_AB_samples, 0, sizeof(out->ss_AB_samples)); tsk_memset(&out->AB_samples, 0, sizeof(out->AB_samples)); + if (out->weights == NULL || out->norm == NULL || out->result_tmp == NULL) { ret = tsk_trace_error(TSK_ERR_NO_MEMORY); goto out; } + ret = tsk_bitset_init(&out->ss_A_samples, num_samples, 1); + if (ret != 0) { + goto out; + } + ret = tsk_bitset_init(&out->ss_B_samples, num_samples, 1); + if (ret != 0) { + goto out; + } + ret = tsk_bitset_init(&out->ss_AB_samples, num_samples, 1); + if (ret != 0) { + goto out; + } ret = tsk_bitset_init(&out->AB_samples, num_samples, 1); if (ret != 0) { goto out; @@ -2343,39 +2363,42 @@ compute_general_two_site_stat_result(const tsk_bitset_t *state, tsk_size_t a_off // a1 [s1, s2, s3] [s1, s2, s3] [s1, s2, s3] // a2 [s1, s2, s3] [s1, s2, s3] [s1, s2, s3] // a3 [s1, s2, s3] [s1, s2, s3] [s1, s2, s3] - tsk_size_t k, mut_a, mut_b; - tsk_size_t result_row_len = num_b_alleles * result_dim; - uint8_t polarised_val = polarised ? 1 : 0; - double *hap_weight_row; - double *result_tmp_row; + tsk_size_t k, mut_a, mut_b, result_row_len = num_b_alleles * result_dim; + uint8_t is_polarised = polarised ? 1 : 0; + double *hap_weight_row, *result_tmp_row; double *norm = work->norm; double *weights = work->weights; double *result_tmp = work->result_tmp; tsk_bitset_t AB_samples = work->AB_samples; + tsk_bitset_t ss_A_samples = work->ss_A_samples; + tsk_bitset_t ss_B_samples = work->ss_B_samples; + tsk_bitset_t ss_AB_samples = work->ss_AB_samples; - for (mut_a = polarised_val; mut_a < num_a_alleles; mut_a++) { + for (mut_a = is_polarised; mut_a < num_a_alleles; mut_a++) { result_tmp_row = GET_2D_ROW(result_tmp, result_row_len, mut_a); - for (mut_b = polarised_val; mut_b < num_b_alleles; mut_b++) { + for (mut_b = is_polarised; mut_b < num_b_alleles; mut_b++) { tsk_bitset_intersect( state, mut_a + a_off, state, mut_b + b_off, &AB_samples); for (k = 0; k < state_dim; k++) { hap_weight_row = GET_2D_ROW(weights, 3, k); - hap_weight_row[0] = (double) tsk_bitset_isect_and_count( - &AB_samples, 0, sample_sets, k); - hap_weight_row[1] = (double) tsk_bitset_isect_and_count( - state, mut_a + a_off, sample_sets, k) - - hap_weight_row[0]; - hap_weight_row[2] = (double) tsk_bitset_isect_and_count( - state, mut_b + b_off, sample_sets, k) - - hap_weight_row[0]; + tsk_bitset_intersect( + state, mut_a + a_off, sample_sets, k, &ss_A_samples); + tsk_bitset_intersect( + state, mut_b + b_off, sample_sets, k, &ss_B_samples); + tsk_bitset_intersect(&AB_samples, 0, sample_sets, k, &ss_AB_samples); + hap_weight_row[0] = (double) tsk_bitset_count(&ss_AB_samples, 0); + hap_weight_row[1] + = (double) tsk_bitset_count(&ss_A_samples, 0) - hap_weight_row[0]; + hap_weight_row[2] + = (double) tsk_bitset_count(&ss_B_samples, 0) - hap_weight_row[0]; } ret = f(state_dim, weights, result_dim, result_tmp_row, f_params); if (ret != 0) { goto out; } - ret = norm_f(result_dim, weights, num_a_alleles - polarised_val, - num_b_alleles - polarised_val, norm, f_params); + ret = norm_f(result_dim, weights, num_a_alleles - is_polarised, + num_b_alleles - is_polarised, norm, f_params); if (ret != 0) { goto out; } @@ -2957,6 +2980,7 @@ compute_two_tree_branch_state_update(const tsk_treeseq_t *ts, tsk_id_t c, tsk_size_t num_nodes = ts->tables->nodes.num_rows; double *weights = work->weights; double *result_tmp = work->result_tmp; + tsk_bitset_t AB_samples = work->AB_samples; b_len = B_branch_len[c] * sign; if (b_len == 0) { @@ -2971,8 +2995,9 @@ compute_two_tree_branch_state_update(const tsk_treeseq_t *ts, tsk_id_t c, a_row = (state_dim * n) + k; b_row = (state_dim * (tsk_size_t) c) + k; weights_row = GET_2D_ROW(weights, 3, k); - weights_row[0] = (double) tsk_bitset_isect_and_count( - A_state_samples, a_row, B_state_samples, b_row); + tsk_bitset_intersect( + A_state_samples, a_row, B_state_samples, b_row, &AB_samples); + weights_row[0] = (double) tsk_bitset_count(&AB_samples, 0); weights_row[1] = (double) tsk_bitset_count(A_state_samples, a_row) - weights_row[0]; weights_row[2] From e5143b959fb0ad97e88e1cf20ee31e7314686398 Mon Sep 17 00:00:00 2001 From: lkirk Date: Wed, 30 Jul 2025 14:05:59 -0500 Subject: [PATCH 41/70] fix bitset_get_items and bitset_contains --- c/tests/test_core.c | 8 ++++---- c/tskit/core.c | 6 ++++-- c/tskit/core.h | 2 +- c/tskit/trees.c | 4 ++-- 4 files changed, 11 insertions(+), 9 deletions(-) diff --git a/c/tests/test_core.c b/c/tests/test_core.c index 9fa1b576b3..5ac779efa0 100644 --- a/c/tests/test_core.c +++ b/c/tests/test_core.c @@ -537,7 +537,7 @@ test_bit_arrays(void) // test item retrieval tsk_bitset_init(&arr, 90, 1); - tsk_bitset_get_items(&arr, items, &n_items); + tsk_bitset_get_items(&arr, 0, items, &n_items); assert_arrays_equal(n_items_truth, items, items_truth); for (tsk_bitset_val_t i = 0; i < 20; i++) { @@ -556,7 +556,7 @@ test_bit_arrays(void) // verify our assumptions about bit array counting CU_ASSERT_EQUAL_FATAL(tsk_bitset_count(&arr, 0), 22); - tsk_bitset_get_items(&arr, items, &n_items); + tsk_bitset_get_items(&arr, 0, items, &n_items); assert_arrays_equal(n_items_truth, items, items_truth); tsk_memset(items, 0, 64); @@ -579,7 +579,7 @@ test_bit_arrays(void) n_items_truth++; } - tsk_bitset_get_items(&arr_row1, items, &n_items); + tsk_bitset_get_items(&arr_row1, 0, items, &n_items); assert_arrays_equal(n_items_truth, items, items_truth); tsk_memset(items, 0, 64); @@ -593,7 +593,7 @@ test_bit_arrays(void) n_items_truth++; } - tsk_bitset_get_items(&arr_row2, items, &n_items); + tsk_bitset_get_items(&arr_row2, 0, items, &n_items); assert_arrays_equal(n_items_truth, items, items_truth); tsk_memset(items, 0, 64); diff --git a/c/tskit/core.c b/c/tskit/core.c index e396ab0626..2366de38b5 100644 --- a/c/tskit/core.c +++ b/c/tskit/core.c @@ -1360,7 +1360,8 @@ tsk_bitset_count(const tsk_bitset_t *self, tsk_size_t row) } void -tsk_bitset_get_items(const tsk_bitset_t *self, tsk_id_t *items, tsk_size_t *n_items) +tsk_bitset_get_items( + const tsk_bitset_t *self, tsk_size_t row, tsk_id_t *items, tsk_size_t *n_items) { // Get the items stored in the row of a bitset. // Uses a de Bruijn sequence lookup table to determine the lowest bit set. See the @@ -1370,10 +1371,11 @@ tsk_bitset_get_items(const tsk_bitset_t *self, tsk_id_t *items, tsk_size_t *n_it tsk_bitset_val_t v, lsb; // least significant bit static const tsk_id_t lookup[32] = { 0, 1, 28, 2, 29, 14, 24, 3, 30, 22, 20, 15, 25, 17, 4, 8, 31, 27, 13, 23, 21, 19, 16, 7, 26, 12, 18, 6, 11, 5, 10, 9 }; + tsk_bitset_val_t *self_d = self->data + (row * self->row_len); n = 0; for (i = 0; i < self->row_len; i++) { - v = self->data[i]; + v = self_d[i]; off = i * ((tsk_size_t) TSK_BIT_ARRAY_NUM_BITS); if (v == 0) { continue; diff --git a/c/tskit/core.h b/c/tskit/core.h index c7f1c6ea89..cd95ad596b 100644 --- a/c/tskit/core.h +++ b/c/tskit/core.h @@ -1119,7 +1119,7 @@ bool tsk_bitset_contains( const tsk_bitset_t *self, tsk_size_t row, const tsk_bitset_val_t bit); tsk_size_t tsk_bitset_count(const tsk_bitset_t *self, tsk_size_t row); void tsk_bitset_get_items( - const tsk_bitset_t *self, tsk_id_t *items, tsk_size_t *n_items); + const tsk_bitset_t *self, tsk_size_t row, tsk_id_t *items, tsk_size_t *n_items); #ifdef __cplusplus } diff --git a/c/tskit/trees.c b/c/tskit/trees.c index e410418a3e..cf4a132f62 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -3061,7 +3061,7 @@ compute_two_tree_branch_stat(const tsk_treeseq_t *ts, const iter_state *l_state, } } // Subtract the whole contribution from the child node - tsk_bitset_get_items(&updates, updated_nodes, &n_updates); + tsk_bitset_get_items(&updates, 0, updated_nodes, &n_updates); while (n_updates != 0) { n_updates--; c = updated_nodes[n_updates]; @@ -3102,7 +3102,7 @@ compute_two_tree_branch_stat(const tsk_treeseq_t *ts, const iter_state *l_state, } // Update all affected child nodes (fully subtracted, deferred from addition) n_updates = 0; - tsk_bitset_get_items(&updates, updated_nodes, &n_updates); + tsk_bitset_get_items(&updates, 0, updated_nodes, &n_updates); while (n_updates != 0) { n_updates--; c = updated_nodes[n_updates]; From 949c7ba9b23b7e4f11a2414a092f6d6b6700c042 Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 8 Aug 2025 14:38:13 -0500 Subject: [PATCH 42/70] a bit of manual merge conflict resolution --- python/tests/test_ld_matrix.py | 5 +---- python/tskit/trees.py | 18 ++++++++---------- 2 files changed, 9 insertions(+), 14 deletions(-) diff --git a/python/tests/test_ld_matrix.py b/python/tests/test_ld_matrix.py index c69ecd4dad..5e6f572f41 100644 --- a/python/tests/test_ld_matrix.py +++ b/python/tests/test_ld_matrix.py @@ -1055,10 +1055,8 @@ def r2_ij_summary_func( D_j = pAB - pA * pB denom_j = np.sqrt(pA * (1 - pA) * pB * (1 - pB)) - p_A = (w_A_i + w_A_j) / (ni + nj) - p_B = (w_B_i + w_B_j) / (ni + nj) with suppress_overflow_div0_warning(): - result[k] = (D_i * D_j) / (p_A * (1 - p_A) * p_B * (1 - p_B)) + result[k] = result[k] = (D_i * D_j) / (denom_i * denom_j) def D_summary_func( @@ -1780,7 +1778,6 @@ def test_ld_empty_examples(ts): def test_input_validation(): - # TODO ts = get_paper_ex_ts() with pytest.raises(ValueError, match="Unknown two-locus statistic"): ts.ld_matrix(stat="bad_stat") diff --git a/python/tskit/trees.py b/python/tskit/trees.py index 699debc0c3..f6c7d51007 100644 --- a/python/tskit/trees.py +++ b/python/tskit/trees.py @@ -23,7 +23,6 @@ """ Module responsible for managing trees and tree sequences. """ - from __future__ import annotations import base64 @@ -687,8 +686,7 @@ def __init__( options = 0 if sample_counts is not None: warnings.warn( - "The sample_counts option is not supported since 0.2.4 " - "and is ignored", + "The sample_counts option is not supported since 0.2.4 and is ignored", RuntimeWarning, stacklevel=4, ) @@ -6841,7 +6839,7 @@ def to_macs(self): bytes_genotypes[:] = lookup[variant.genotypes] genotypes = bytes_genotypes.tobytes().decode() output.append( - f"SITE:\t{variant.index}\t{variant.position / m}\t0.0\t" f"{genotypes}" + f"SITE:\t{variant.index}\t{variant.position / m}\t0.0\t{genotypes}" ) return "\n".join(output) + "\n" @@ -9144,9 +9142,9 @@ def pca( if time_windows is None: tree_sequence_low, tree_sequence_high = None, self else: - assert ( - time_windows[0] < time_windows[1] - ), "The second argument should be larger." + assert time_windows[0] < time_windows[1], ( + "The second argument should be larger." + ) tree_sequence_low, tree_sequence_high = ( self.decapitate(time_windows[0]), self.decapitate(time_windows[1]), @@ -9214,9 +9212,9 @@ def _rand_pow_range_finder( """ Algorithm 9 in https://arxiv.org/pdf/2002.01387 """ - assert ( - num_vectors >= rank > 0 - ), "num_vectors should not be smaller than rank" + assert num_vectors >= rank > 0, ( + "num_vectors should not be smaller than rank" + ) for _ in range(depth): Q = np.linalg.qr(Q)[0] Q = operator(Q) From f3cc3323c7bc147d834884fce423b3d8006eca82 Mon Sep 17 00:00:00 2001 From: lkirk Date: Wed, 18 Dec 2024 14:34:46 -0600 Subject: [PATCH 43/70] C and Python API for two-way two-locus stats This PR implements the C and Python API for computing two-way two-locus statistics. The algorithm is identical to the python version, except during testing I uncovered a small issue with normalisation. We need to handle the case where sample sets are of different sizes. The fix for this was to average the normalisation factor for each sample set. Test coverage has been added to cover C, low-level python and some high-level tests. --- python/tests/test_ld_matrix.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/tests/test_ld_matrix.py b/python/tests/test_ld_matrix.py index 5e6f572f41..4480f603d3 100644 --- a/python/tests/test_ld_matrix.py +++ b/python/tests/test_ld_matrix.py @@ -1055,6 +1055,8 @@ def r2_ij_summary_func( D_j = pAB - pA * pB denom_j = np.sqrt(pA * (1 - pA) * pB * (1 - pB)) + p_A = (w_A_i + w_A_j) / (ni + nj) + p_B = (w_B_i + w_B_j) / (ni + nj) with suppress_overflow_div0_warning(): result[k] = result[k] = (D_i * D_j) / (denom_i * denom_j) From ca5cca07871beffa7d2696e1431832b0abf316ca Mon Sep 17 00:00:00 2001 From: lkirk Date: Mon, 21 Jul 2025 20:26:21 -0500 Subject: [PATCH 44/70] fix norm_hap_weighted_ij. needs to get wAB instead of using indices From 87ac1f7b5eeecb7def07be1a0300a1fc17d343b8 Mon Sep 17 00:00:00 2001 From: lkirk Date: Wed, 23 Jul 2025 05:57:29 -0500 Subject: [PATCH 45/70] concat is not compatible with numpy<2 From 9291f9dd5b97d90a9c752da30b1404b1575fde05 Mon Sep 17 00:00:00 2001 From: lkirk Date: Wed, 23 Jul 2025 11:10:57 -0500 Subject: [PATCH 46/70] remove unreachable code path (we switch on the presence of indexes) From 7b60d4b62bb974951afa9f3a856dddfb8b563fe4 Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 25 Jul 2025 17:49:04 -0500 Subject: [PATCH 47/70] better docstring for norm func From d57ad9257cfa87618ad3e63fce52dc3ea1c4d148 Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 25 Jul 2025 17:57:17 -0500 Subject: [PATCH 48/70] remove todo From f125781d93f6f887fa880d223eacaa87eab8a534 Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 25 Jul 2025 17:59:00 -0500 Subject: [PATCH 49/70] remove todo From 029e6203be9e8346fc2162a0857c8541c3696da0 Mon Sep 17 00:00:00 2001 From: lkirk Date: Thu, 17 Jul 2025 02:34:24 -0500 Subject: [PATCH 50/70] fix normalisation and a small indexing bug. dimension tests --- c/tskit/trees.c | 9 +++------ python/tests/test_ld_matrix.py | 1 + 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index cf4a132f62..97b510d923 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -3241,17 +3241,12 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl tsk_memset(&sample_sets_bits, 0, sizeof(sample_sets_bits)); - // We do not support two-locus node stats - if (!!(options & TSK_STAT_NODE)) { - ret = tsk_trace_error(TSK_ERR_UNSUPPORTED_STAT_MODE); - goto out; - } // If no mode is specified, we default to site mode if (!(stat_site || stat_branch)) { stat_site = true; } // It's an error to specify more than one mode - if (stat_site + stat_branch > 1) { + if (stat_site + stat_branch + stat_node > 1) { ret = tsk_trace_error(TSK_ERR_MULTIPLE_STAT_MODES); goto out; } @@ -3296,6 +3291,8 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl ret = tsk_treeseq_two_branch_count_stat(self, state_dim, &sample_sets_bits, result_dim, f, &f_params, norm_f, out_rows, row_positions, out_cols, col_positions, options, result); + } else { + ret = tsk_trace_error(TSK_ERR_UNSUPPORTED_STAT_MODE); } out: diff --git a/python/tests/test_ld_matrix.py b/python/tests/test_ld_matrix.py index 4480f603d3..fcec932b05 100644 --- a/python/tests/test_ld_matrix.py +++ b/python/tests/test_ld_matrix.py @@ -1780,6 +1780,7 @@ def test_ld_empty_examples(ts): def test_input_validation(): + # TODO ts = get_paper_ex_ts() with pytest.raises(ValueError, match="Unknown two-locus statistic"): ts.ld_matrix(stat="bad_stat") From ac8646dbe512a1efe68c2a3b99515abfb95edc66 Mon Sep 17 00:00:00 2001 From: lkirk Date: Thu, 17 Jul 2025 23:57:00 -0500 Subject: [PATCH 51/70] k should be unsigned From 0c56d6d4e2ce19714c30fac51b0caba315bde56b Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 25 Jul 2025 17:59:00 -0500 Subject: [PATCH 52/70] remove todo From ec28f2fd1a95ee3098b0116f90f1388bea6ca722 Mon Sep 17 00:00:00 2001 From: lkirk Date: Thu, 24 Jul 2025 18:11:56 -0500 Subject: [PATCH 53/70] first pass at bit array refactor --- c/tskit/trees.c | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index 97b510d923..5ca81d14e1 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -2321,19 +2321,20 @@ two_locus_work_init(tsk_size_t max_alleles, tsk_size_t result_dim, tsk_size_t st ret = tsk_trace_error(TSK_ERR_NO_MEMORY); goto out; } - ret = tsk_bitset_init(&out->ss_A_samples, num_samples, 1); + + ret = tsk_bitset_init(&out->AB_samples, num_samples, 1); if (ret != 0) { goto out; } - ret = tsk_bitset_init(&out->ss_B_samples, num_samples, 1); + ret = tsk_bitset_init(&out->ss_A_samples, num_samples, 1); if (ret != 0) { goto out; } - ret = tsk_bitset_init(&out->ss_AB_samples, num_samples, 1); + ret = tsk_bitset_init(&out->ss_B_samples, num_samples, 1); if (ret != 0) { goto out; } - ret = tsk_bitset_init(&out->AB_samples, num_samples, 1); + ret = tsk_bitset_init(&out->ss_AB_samples, num_samples, 1); if (ret != 0) { goto out; } @@ -2348,6 +2349,9 @@ two_locus_work_free(two_locus_work_t *work) tsk_safe_free(work->norm); tsk_safe_free(work->result_tmp); tsk_bitset_free(&work->AB_samples); + tsk_bitset_free(&work->ss_A_samples); + tsk_bitset_free(&work->ss_B_samples); + tsk_bitset_free(&work->ss_AB_samples); } static int From 88b6a886acab03496767f40057fcc4ddb6a63c6a Mon Sep 17 00:00:00 2001 From: lkirk Date: Thu, 29 May 2025 10:22:53 -0500 Subject: [PATCH 54/70] initial implementation; need to work on output dims/scalar sample sets and indexes --- c/tskit/trees.c | 3 +++ python/tskit/trees.py | 1 + 2 files changed, 4 insertions(+) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index 5ca81d14e1..85fedaea81 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -2302,6 +2302,9 @@ typedef struct { tsk_bitset_t ss_AB_samples; } two_locus_work_t; +typedef int norm_func_t(tsk_size_t result_dim, const double *hap_weights, tsk_size_t n_a, + tsk_size_t n_b, double *result, void *params); + static int two_locus_work_init(tsk_size_t max_alleles, tsk_size_t result_dim, tsk_size_t state_dim, tsk_size_t num_samples, two_locus_work_t *out) diff --git a/python/tskit/trees.py b/python/tskit/trees.py index f6c7d51007..14a50d2fde 100644 --- a/python/tskit/trees.py +++ b/python/tskit/trees.py @@ -23,6 +23,7 @@ """ Module responsible for managing trees and tree sequences. """ + from __future__ import annotations import base64 From eb55592f5bacbf515a45b4b5c6157359c3d22aff Mon Sep 17 00:00:00 2001 From: lkirk Date: Thu, 17 Jul 2025 02:34:24 -0500 Subject: [PATCH 55/70] fix normalisation and a small indexing bug. dimension tests From da3b0abc7b450838a45627fb3648fc337070614f Mon Sep 17 00:00:00 2001 From: lkirk Date: Thu, 17 Jul 2025 23:57:00 -0500 Subject: [PATCH 56/70] k should be unsigned From 41aa9b7ae3f1908a861a9049b8cb9f630de5b2b3 Mon Sep 17 00:00:00 2001 From: lkirk Date: Mon, 21 Jul 2025 20:26:21 -0500 Subject: [PATCH 57/70] fix norm_hap_weighted_ij. needs to get wAB instead of using indices From 6c29c804a06f27a7e538b5f9b3a1c5d09ca63d6b Mon Sep 17 00:00:00 2001 From: lkirk Date: Wed, 23 Jul 2025 05:57:29 -0500 Subject: [PATCH 58/70] concat is not compatible with numpy<2 From 3fb5523b081b486ac7aea478b0cbe412ee15618f Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 25 Jul 2025 17:49:04 -0500 Subject: [PATCH 59/70] better docstring for norm func From 58b866c960bded7d1ca7a0fff7cc7e339baa165d Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 25 Jul 2025 17:57:17 -0500 Subject: [PATCH 60/70] remove todo From 43d38b6fa16d3a691b7943fda2d5c9e912f9cef2 Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 25 Jul 2025 17:59:00 -0500 Subject: [PATCH 61/70] remove todo From 73671bfd26931e0c2eb67c4db6b61a7e6b33f673 Mon Sep 17 00:00:00 2001 From: lkirk Date: Thu, 24 Jul 2025 18:11:56 -0500 Subject: [PATCH 62/70] first pass at bit array refactor From 6bba438d291350ec17baef3addc76abc4ba4e81b Mon Sep 17 00:00:00 2001 From: lkirk Date: Wed, 30 Jul 2025 14:07:17 -0500 Subject: [PATCH 63/70] Further optimizations Compute A/B counts upfront for each sample set. We were computing them redundantly each for each site pair. This is performed in the new function get_mutation_sample_sets, which takes the samples for each site and computes the sample for each site for each sample set. During this operation, we compute the number of samples containing the given allele for each site. Add a non-normalized version of compute_general_two_site_stat_result for situations where we're computing stats from biallelic loci. For site statistics, we do not convert sample sets to bitsets, opting to defer that action to get_mutation_sample_sets. We should consider doing this for branch stats as well, but it will be trickier. Finally, produce an optimized version of r2_summary_function to reduce the number of divisions we're doing. Ultimately, the issue we were running into was memory latency when accessing elements within the state_row. This has been mitigated by providing the entire expression to the compiler, allowing instruction level parallelism to compute values as the memory we depend on becomes available. The expression we've chosen also reduces the amount of memory dependencies, reducing the latency of these computations. This all works in the case of a single sample set, tests are not passing for multiple sample sets. I'll need to streamline the computation of site offsets to handle this, I first wanted to see how much of an improvement we would see from these optimizations. --- c/tskit/trees.c | 332 +++++++++++++++++++++++++++++++++--------------- 1 file changed, 230 insertions(+), 102 deletions(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index 85fedaea81..08fa4717c0 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -2292,34 +2292,72 @@ get_all_samples_bits(tsk_bitset_t *all_samples, tsk_size_t n) } } +// typedef struct { +// double *weights; +// double *norm; +// double *result_tmp; +// tsk_bitset_t AB_samples; +// tsk_bitset_t **ss_mutations; +// tsk_size_t *ss_allele_counts; +// } two_locus_work_t; + +// static int +// two_locus_work_init(tsk_size_t max_alleles, tsk_size_t result_dim, tsk_size_t +// state_dim, +// tsk_size_t num_sample_sets, const tsk_size_t *sample_set_sizes, +// two_locus_work_t *out) +// { +// int ret = 0; +// tsk_size_t i, max_ss_size = 0; +// out->weights = tsk_malloc(3 * state_dim * sizeof(*out->weights)); +// out->norm = tsk_malloc(result_dim * sizeof(*out->norm)); +// out->result_tmp +// = tsk_malloc(result_dim * max_alleles * max_alleles * +// sizeof(*out->result_tmp)); +// out->ss_mutations = num_sample_sets > 0 +// ? tsk_calloc(num_sample_sets, sizeof(out->ss_mutations)) +// : NULL; +// for (i = 0; i < num_sample_sets; i++) { +// if (sample_set_sizes[i] > max_ss_size) { +// max_ss_size = sample_set_sizes[i]; +// } +// } +// tsk_memset(&out->AB_samples, 0, sizeof(out->AB_samples)); +// if (out->weights == NULL || out->norm == NULL || out->result_tmp == NULL) { +// ret = tsk_trace_error(TSK_ERR_NO_MEMORY); +// goto out; +// } +// ret = tsk_bitset_init(&out->AB_samples, max_ss_size, 1); +// if (ret != 0) { +// goto out; +// } +// for (i = 0; i < num_sample_sets; i++) { +// ret = tsk_bitset_init(out->ss_mutations[i], sample_set_sizes[i], 1); +// if (ret != 0) { +// goto out; +// } +// } +// out: +// return ret; +// } + typedef struct { double *weights; double *norm; double *result_tmp; tsk_bitset_t AB_samples; - tsk_bitset_t ss_A_samples; - tsk_bitset_t ss_B_samples; - tsk_bitset_t ss_AB_samples; } two_locus_work_t; -typedef int norm_func_t(tsk_size_t result_dim, const double *hap_weights, tsk_size_t n_a, - tsk_size_t n_b, double *result, void *params); - static int -two_locus_work_init(tsk_size_t max_alleles, tsk_size_t result_dim, tsk_size_t state_dim, - tsk_size_t num_samples, two_locus_work_t *out) +two_locus_work_init(tsk_size_t max_alleles, tsk_size_t num_samples, + tsk_size_t result_dim, tsk_size_t state_dim, two_locus_work_t *out) { int ret = 0; out->weights = tsk_malloc(3 * state_dim * sizeof(*out->weights)); out->norm = tsk_malloc(result_dim * sizeof(*out->norm)); out->result_tmp = tsk_malloc(result_dim * max_alleles * max_alleles * sizeof(*out->result_tmp)); - - tsk_memset(&out->ss_A_samples, 0, sizeof(out->ss_A_samples)); - tsk_memset(&out->ss_B_samples, 0, sizeof(out->ss_B_samples)); - tsk_memset(&out->ss_AB_samples, 0, sizeof(out->ss_AB_samples)); tsk_memset(&out->AB_samples, 0, sizeof(out->AB_samples)); - if (out->weights == NULL || out->norm == NULL || out->result_tmp == NULL) { ret = tsk_trace_error(TSK_ERR_NO_MEMORY); goto out; @@ -2329,18 +2367,6 @@ two_locus_work_init(tsk_size_t max_alleles, tsk_size_t result_dim, tsk_size_t st if (ret != 0) { goto out; } - ret = tsk_bitset_init(&out->ss_A_samples, num_samples, 1); - if (ret != 0) { - goto out; - } - ret = tsk_bitset_init(&out->ss_B_samples, num_samples, 1); - if (ret != 0) { - goto out; - } - ret = tsk_bitset_init(&out->ss_AB_samples, num_samples, 1); - if (ret != 0) { - goto out; - } out: return ret; } @@ -2352,17 +2378,14 @@ two_locus_work_free(two_locus_work_t *work) tsk_safe_free(work->norm); tsk_safe_free(work->result_tmp); tsk_bitset_free(&work->AB_samples); - tsk_bitset_free(&work->ss_A_samples); - tsk_bitset_free(&work->ss_B_samples); - tsk_bitset_free(&work->ss_AB_samples); } static int -compute_general_two_site_stat_result(const tsk_bitset_t *state, tsk_size_t a_off, - tsk_size_t b_off, tsk_size_t num_a_alleles, tsk_size_t num_b_alleles, - tsk_size_t state_dim, const tsk_bitset_t *sample_sets, tsk_size_t result_dim, - general_stat_func_t *f, sample_count_stat_params_t *f_params, norm_func_t *norm_f, - bool polarised, two_locus_work_t *restrict work, double *result) +compute_general_normed_two_site_stat_result(const tsk_bitset_t *state, + const tsk_size_t *allele_counts, tsk_size_t a_off, tsk_size_t b_off, + tsk_size_t num_a_alleles, tsk_size_t num_b_alleles, tsk_size_t state_dim, + tsk_size_t result_dim, general_stat_func_t *f, sample_count_stat_params_t *f_params, + norm_func_t *norm_f, bool polarised, two_locus_work_t *restrict work, double *result) { int ret = 0; // Sample sets and b sites are rows, a sites are columns @@ -2374,31 +2397,23 @@ compute_general_two_site_stat_result(const tsk_bitset_t *state, tsk_size_t a_off uint8_t is_polarised = polarised ? 1 : 0; double *hap_weight_row, *result_tmp_row; - double *norm = work->norm; - double *weights = work->weights; - double *result_tmp = work->result_tmp; + double *restrict norm = work->norm; + double *restrict weights = work->weights; + double *restrict result_tmp = work->result_tmp; tsk_bitset_t AB_samples = work->AB_samples; - tsk_bitset_t ss_A_samples = work->ss_A_samples; - tsk_bitset_t ss_B_samples = work->ss_B_samples; - tsk_bitset_t ss_AB_samples = work->ss_AB_samples; for (mut_a = is_polarised; mut_a < num_a_alleles; mut_a++) { result_tmp_row = GET_2D_ROW(result_tmp, result_row_len, mut_a); for (mut_b = is_polarised; mut_b < num_b_alleles; mut_b++) { - tsk_bitset_intersect( - state, mut_a + a_off, state, mut_b + b_off, &AB_samples); for (k = 0; k < state_dim; k++) { - hap_weight_row = GET_2D_ROW(weights, 3, k); tsk_bitset_intersect( - state, mut_a + a_off, sample_sets, k, &ss_A_samples); - tsk_bitset_intersect( - state, mut_b + b_off, sample_sets, k, &ss_B_samples); - tsk_bitset_intersect(&AB_samples, 0, sample_sets, k, &ss_AB_samples); - hap_weight_row[0] = (double) tsk_bitset_count(&ss_AB_samples, 0); + state, mut_a + a_off, state, mut_b + b_off, &AB_samples); + hap_weight_row = GET_2D_ROW(weights, 3, k); + hap_weight_row[0] = (double) tsk_bitset_count(&AB_samples, 0); hap_weight_row[1] - = (double) tsk_bitset_count(&ss_A_samples, 0) - hap_weight_row[0]; + = (double) allele_counts[mut_a + a_off] - hap_weight_row[0]; hap_weight_row[2] - = (double) tsk_bitset_count(&ss_B_samples, 0) - hap_weight_row[0]; + = (double) allele_counts[mut_b + b_off] - hap_weight_row[0]; } ret = f(state_dim, weights, result_dim, result_tmp_row, f_params); if (ret != 0) { @@ -2420,6 +2435,34 @@ compute_general_two_site_stat_result(const tsk_bitset_t *state, tsk_size_t a_off return ret; } +static int +compute_general_two_site_stat_result(const tsk_bitset_t *state, + const tsk_size_t *allele_counts, tsk_size_t a_off, tsk_size_t b_off, + tsk_size_t state_dim, tsk_size_t result_dim, general_stat_func_t *f, + sample_count_stat_params_t *f_params, bool polarised, + two_locus_work_t *restrict work, double *result) +{ + int ret = 0; + tsk_size_t k; + tsk_bitset_t AB_samples = work->AB_samples; + tsk_size_t mut_a = polarised ? 1 : 0, mut_b = 1; + double *restrict hap_row, *restrict weights = work->weights; + + for (k = 0; k < state_dim; k++) { + tsk_bitset_intersect(state, mut_a + a_off, state, mut_b + b_off, &AB_samples); + hap_row = GET_2D_ROW(weights, 3, k); + hap_row[0] = (double) tsk_bitset_count(&AB_samples, 0); + hap_row[1] = (double) allele_counts[mut_a + a_off] - hap_row[0]; + hap_row[2] = (double) allele_counts[mut_b + b_off] - hap_row[0]; + } + ret = f(state_dim, weights, result_dim, result, f_params); + if (ret != 0) { + goto out; + } +out: + return ret; +} + static void get_site_row_col_indices(tsk_size_t n_rows, const tsk_id_t *row_sites, tsk_size_t n_cols, const tsk_id_t *col_sites, tsk_id_t *sites, tsk_size_t *n_sites, tsk_size_t *row_idx, @@ -2465,6 +2508,34 @@ get_site_row_col_indices(tsk_size_t n_rows, const tsk_id_t *row_sites, tsk_size_ *n_sites = s; } +// static int +// get_mutation_sample_sets(const tsk_bitset_t *allele_samples, +// const tsk_bitset_t *sample_sets, two_locus_work_t *restrict work) +// { +// int ret = 0; +// tsk_size_t i, j, k, n_ss_samples; +// tsk_id_t sample, *ss_samples = NULL; + +// ss_samples = tsk_calloc( +// allele_samples->row_len * TSK_BIT_ARRAY_NUM_BITS, sizeof(*ss_samples)); +// if (ss_samples == NULL) { +// ret = tsk_trace_error(TSK_ERR_NO_MEMORY); +// goto out; +// } + +// for (i = 0; i < sample_sets->len; i++) { +// tsk_bitset_get_items(sample_sets, i, ss_samples, &n_ss_samples); +// if (n_ss_samples == 0) { +// continue; +// } +// for (j = 0; j < n_ss_samples; j++) { +// sample = ss_samples[0]; +// } +// } +// out: +// return ret; +// } + static int get_mutation_samples(const tsk_treeseq_t *ts, const tsk_id_t *sites, tsk_size_t n_sites, tsk_size_t *num_alleles, tsk_bitset_t *allele_samples) @@ -2550,27 +2621,75 @@ get_mutation_samples(const tsk_treeseq_t *ts, const tsk_id_t *sites, tsk_size_t return ret == TSK_TREE_OK ? 0 : ret; } +static int +get_mutation_sample_sets(const tsk_bitset_t *allele_samples, tsk_size_t num_sample_sets, + const tsk_size_t *sample_set_sizes, const tsk_id_t *sample_sets, + tsk_size_t *max_ss_size, tsk_bitset_t *allele_sample_sets, + tsk_size_t **allele_sample_set_counts) +{ + int ret = 0; + tsk_bitset_val_t k; + tsk_size_t i, j, ss_off; + + *max_ss_size = 0; + for (i = 0; i < num_sample_sets; i++) { + if (sample_set_sizes[i] > *max_ss_size) { + *max_ss_size = sample_set_sizes[i]; + } + } + + *allele_sample_set_counts = tsk_calloc( + allele_samples->len * num_sample_sets, sizeof(**allele_sample_set_counts)); + if (*allele_sample_set_counts == NULL) { + ret = tsk_trace_error(TSK_ERR_NO_MEMORY); + goto out; + } + ret = tsk_bitset_init( + allele_sample_sets, *max_ss_size, allele_samples->len * num_sample_sets); + if (ret != 0) { + goto out; + } + + for (i = 0; i < allele_samples->len; i++) { + ss_off = 0; + for (j = 0; j < num_sample_sets; j++) { + for (k = 0; k < sample_set_sizes[j]; k++) { + if (tsk_bitset_contains( + allele_samples, i, (tsk_bitset_val_t) sample_sets[k + ss_off])) { + tsk_bitset_set_bit(allele_sample_sets, j + i * num_sample_sets, k); + (*allele_sample_set_counts)[j + i * num_sample_sets]++; + } + } + ss_off += sample_set_sizes[j]; + } + } +out: + return ret; +} + static int tsk_treeseq_two_site_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim, - const tsk_bitset_t *sample_sets, tsk_size_t result_dim, general_stat_func_t *f, + tsk_size_t num_sample_sets, const tsk_size_t *sample_set_sizes, + const tsk_id_t *sample_sets, tsk_size_t result_dim, general_stat_func_t *f, sample_count_stat_params_t *f_params, norm_func_t *norm_f, tsk_size_t n_rows, const tsk_id_t *row_sites, tsk_size_t n_cols, const tsk_id_t *col_sites, tsk_flags_t options, double *result) { - int ret = 0; - tsk_bitset_t allele_samples; + tsk_bitset_t allele_samples, allele_sample_sets; bool polarised = false; tsk_id_t *sites; - tsk_size_t r, c, s, max_alleles, n_alleles, n_sites, *row_idx, *col_idx; + tsk_size_t r, c, s, max_ss_size, max_alleles, n_alleles, n_sites, *row_idx, *col_idx; double *result_row; const tsk_size_t num_samples = self->num_samples; - tsk_size_t *num_alleles = NULL, *site_offsets = NULL; + tsk_size_t *num_alleles = NULL, *site_offsets = NULL, + *allele_sample_set_counts = NULL; tsk_size_t result_row_len = n_cols * result_dim; two_locus_work_t work; tsk_memset(&work, 0, sizeof(work)); tsk_memset(&allele_samples, 0, sizeof(allele_samples)); + tsk_memset(&allele_sample_sets, 0, sizeof(allele_sample_sets)); sites = tsk_malloc(self->tables->sites.num_rows * sizeof(*sites)); row_idx = tsk_malloc(self->tables->sites.num_rows * sizeof(*row_idx)); col_idx = tsk_malloc(self->tables->sites.num_rows * sizeof(*col_idx)); @@ -2598,18 +2717,22 @@ tsk_treeseq_two_site_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim, max_alleles = self->site_mutations_length[sites[s]]; } } - max_alleles++; + max_alleles++; // add 1 for the ancestral allele - ret = two_locus_work_init(max_alleles, result_dim, state_dim, num_samples, &work); + ret = tsk_bitset_init(&allele_samples, num_samples, n_alleles); if (ret != 0) { goto out; } - ret = tsk_bitset_init(&allele_samples, num_samples, n_alleles); + ret = get_mutation_samples(self, sites, n_sites, num_alleles, &allele_samples); if (ret != 0) { goto out; } - // we track the number of alleles to account for backmutations - ret = get_mutation_samples(self, sites, n_sites, num_alleles, &allele_samples); + ret = get_mutation_sample_sets(&allele_samples, num_sample_sets, sample_set_sizes, + sample_sets, &max_ss_size, &allele_sample_sets, &allele_sample_set_counts); + if (ret != 0) { + goto out; + } + ret = two_locus_work_init(max_alleles, max_ss_size, result_dim, state_dim, &work); if (ret != 0) { goto out; } @@ -2622,11 +2745,18 @@ tsk_treeseq_two_site_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim, for (r = 0; r < n_rows; r++) { result_row = GET_2D_ROW(result, result_row_len, r); for (c = 0; c < n_cols; c++) { - ret = compute_general_two_site_stat_result(&allele_samples, - site_offsets[row_idx[r]], site_offsets[col_idx[c]], - num_alleles[row_idx[r]], num_alleles[col_idx[c]], state_dim, sample_sets, - result_dim, f, f_params, norm_f, polarised, &work, - &(result_row[c * result_dim])); + if (num_alleles[row_idx[r]] == 2 && num_alleles[col_idx[c]] == 2) { + ret = compute_general_two_site_stat_result(&allele_sample_sets, + allele_sample_set_counts, site_offsets[row_idx[r]], + site_offsets[col_idx[c]], state_dim, result_dim, f, f_params, + polarised, &work, &(result_row[c * result_dim])); + } else { + ret = compute_general_normed_two_site_stat_result(&allele_sample_sets, + allele_sample_set_counts, site_offsets[row_idx[r]], + site_offsets[col_idx[c]], num_alleles[row_idx[r]], + num_alleles[col_idx[c]], state_dim, result_dim, f, f_params, norm_f, + polarised, &work, &(result_row[c * result_dim])); + } if (ret != 0) { goto out; } @@ -3034,14 +3164,13 @@ compute_two_tree_branch_stat(const tsk_treeseq_t *ts, const iter_state *l_state, const tsk_id_t *restrict edges_child = ts->tables->edges.child; const tsk_id_t *restrict edges_parent = ts->tables->edges.parent; const tsk_size_t num_nodes = ts->tables->nodes.num_rows; - const tsk_size_t num_samples = ts->num_samples; tsk_bitset_t updates, *r_samples = r_state->node_samples; two_locus_work_t work; tsk_memset(&work, 0, sizeof(work)); tsk_memset(&updates, 0, sizeof(updates)); // only two alleles are possible for branch stats - ret = two_locus_work_init(2, result_dim, state_dim, num_samples, &work); + ret = two_locus_work_init(2, ts->num_samples, result_dim, state_dim, &work); if (ret != 0) { goto out; } @@ -3125,7 +3254,8 @@ compute_two_tree_branch_stat(const tsk_treeseq_t *ts, const iter_state *l_state, static int tsk_treeseq_two_branch_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim, - const tsk_bitset_t *sample_sets, tsk_size_t result_dim, general_stat_func_t *f, + tsk_size_t num_sample_sets, const tsk_size_t *sample_set_sizes, + const tsk_id_t *sample_sets, tsk_size_t result_dim, general_stat_func_t *f, sample_count_stat_params_t *f_params, norm_func_t *TSK_UNUSED(norm_f), tsk_size_t n_rows, const double *row_positions, tsk_size_t n_cols, const double *col_positions, tsk_flags_t TSK_UNUSED(options), double *result) @@ -3134,11 +3264,12 @@ tsk_treeseq_two_branch_count_stat(const tsk_treeseq_t *self, tsk_size_t state_di int r, c; tsk_id_t *row_indexes = NULL, *col_indexes = NULL; tsk_size_t i, j, k, row, col, *row_repeats = NULL, *col_repeats = NULL; - tsk_bitset_t node_samples; + tsk_bitset_t node_samples, sample_sets_bits; iter_state l_state, r_state; double *result_tmp = NULL, *result_row; const tsk_size_t num_nodes = self->tables->nodes.num_rows; + tsk_memset(&sample_sets_bits, 0, sizeof(sample_sets_bits)); tsk_memset(&node_samples, 0, sizeof(node_samples)); tsk_memset(&l_state, 0, sizeof(l_state)); tsk_memset(&r_state, 0, sizeof(r_state)); @@ -3155,6 +3286,11 @@ tsk_treeseq_two_branch_count_stat(const tsk_treeseq_t *self, tsk_size_t state_di if (ret != 0) { goto out; } + ret = sample_sets_to_bitset( + self, sample_set_sizes, sample_sets, num_sample_sets, &sample_sets_bits); + if (ret != 0) { + goto out; + } ret = positions_to_tree_indexes(self, row_positions, n_rows, &row_indexes); if (ret != 0) { goto out; @@ -3171,7 +3307,7 @@ tsk_treeseq_two_branch_count_stat(const tsk_treeseq_t *self, tsk_size_t state_di if (ret != 0) { goto out; } - ret = get_node_samples(self, state_dim, sample_sets, &node_samples); + ret = get_node_samples(self, state_dim, &sample_sets_bits, &node_samples); if (ret != 0) { goto out; } @@ -3237,7 +3373,6 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl // TODO: generalize this function if we ever decide to do weighted two_locus stats. // We only implement count stats and therefore we don't handle weights. int ret = 0; - tsk_bitset_t sample_sets_bits; bool stat_site = !!(options & TSK_STAT_SITE); bool stat_branch = !!(options & TSK_STAT_BRANCH); tsk_size_t state_dim = num_sample_sets; @@ -3246,8 +3381,11 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl .sample_set_sizes = sample_set_sizes, .set_indexes = set_indexes }; - tsk_memset(&sample_sets_bits, 0, sizeof(sample_sets_bits)); - + // We do not support two-locus node stats + if (!!(options & TSK_STAT_NODE)) { + ret = tsk_trace_error(TSK_ERR_UNSUPPORTED_STAT_MODE); + goto out; + } // If no mode is specified, we default to site mode if (!(stat_site || stat_branch)) { stat_site = true; @@ -3266,11 +3404,6 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl ret = tsk_trace_error(TSK_ERR_BAD_RESULT_DIMS); goto out; } - ret = sample_sets_to_bitset( - self, sample_set_sizes, sample_sets, num_sample_sets, &sample_sets_bits); - if (ret != 0) { - goto out; - } if (stat_site) { ret = check_sites(row_sites, out_rows, self->tables->sites.num_rows); @@ -3281,9 +3414,10 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl if (ret != 0) { goto out; } - ret = tsk_treeseq_two_site_count_stat(self, state_dim, &sample_sets_bits, - result_dim, f, &f_params, norm_f, out_rows, row_sites, out_cols, col_sites, - options, result); + // TODO: result dim/state dim can be set internally now. + ret = tsk_treeseq_two_site_count_stat(self, state_dim, num_sample_sets, + sample_set_sizes, sample_sets, result_dim, f, &f_params, norm_f, out_rows, + row_sites, out_cols, col_sites, options, result); } else if (stat_branch) { ret = check_positions( row_positions, out_rows, tsk_treeseq_get_sequence_length(self)); @@ -3295,15 +3429,12 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl if (ret != 0) { goto out; } - ret = tsk_treeseq_two_branch_count_stat(self, state_dim, &sample_sets_bits, - result_dim, f, &f_params, norm_f, out_rows, row_positions, out_cols, - col_positions, options, result); - } else { - ret = tsk_trace_error(TSK_ERR_UNSUPPORTED_STAT_MODE); + ret = tsk_treeseq_two_branch_count_stat(self, state_dim, num_sample_sets, + sample_set_sizes, sample_sets, result_dim, f, &f_params, norm_f, out_rows, + row_positions, out_cols, col_positions, options, result); } out: - tsk_bitset_free(&sample_sets_bits); return ret; } @@ -4179,28 +4310,25 @@ tsk_treeseq_D2(const tsk_treeseq_t *self, tsk_size_t num_sample_sets, } static int -r2_summary_func(tsk_size_t state_dim, const double *state, - tsk_size_t TSK_UNUSED(result_dim), double *result, void *params) +r2_summary_func(tsk_size_t state_dim, const double *restrict state, + tsk_size_t TSK_UNUSED(result_dim), double *restrict result, void *restrict params) { sample_count_stat_params_t args = *(sample_count_stat_params_t *) params; - double n; - const double *state_row; + double n, wA, wB, wAwB, wAB, wAb, waB; tsk_size_t j; + const tsk_size_t *ns = args.sample_set_sizes; + const double *state_row; for (j = 0; j < state_dim; j++) { - n = (double) args.sample_set_sizes[j]; state_row = GET_2D_ROW(state, 3, j); - double p_AB = state_row[0] / n; - double p_Ab = state_row[1] / n; - double p_aB = state_row[2] / n; - - double p_A = p_AB + p_Ab; - double p_B = p_AB + p_aB; - - double D = p_AB - (p_A * p_B); - double denom = p_A * p_B * (1 - p_A) * (1 - p_B); - - result[j] = (D * D) / denom; + wAB = state_row[0]; + wAb = state_row[1]; + waB = state_row[2]; + n = (double) ns[j]; + wA = wAB + wAb; + wB = wAB + waB; + wAwB = wA * wB; + result[j] = pow((n * wAB - wAwB), 2) / ((n - wB) * (n - wA) * wAwB); } return 0; } From 650275870137a7c2c93a61bf6471ae4bdde482a9 Mon Sep 17 00:00:00 2001 From: lkirk Date: Wed, 30 Jul 2025 14:19:08 -0500 Subject: [PATCH 64/70] clean up unused routines --- c/tskit/trees.c | 77 ------------------------------------------------- 1 file changed, 77 deletions(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index 08fa4717c0..aa5ad14c3f 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -2292,55 +2292,6 @@ get_all_samples_bits(tsk_bitset_t *all_samples, tsk_size_t n) } } -// typedef struct { -// double *weights; -// double *norm; -// double *result_tmp; -// tsk_bitset_t AB_samples; -// tsk_bitset_t **ss_mutations; -// tsk_size_t *ss_allele_counts; -// } two_locus_work_t; - -// static int -// two_locus_work_init(tsk_size_t max_alleles, tsk_size_t result_dim, tsk_size_t -// state_dim, -// tsk_size_t num_sample_sets, const tsk_size_t *sample_set_sizes, -// two_locus_work_t *out) -// { -// int ret = 0; -// tsk_size_t i, max_ss_size = 0; -// out->weights = tsk_malloc(3 * state_dim * sizeof(*out->weights)); -// out->norm = tsk_malloc(result_dim * sizeof(*out->norm)); -// out->result_tmp -// = tsk_malloc(result_dim * max_alleles * max_alleles * -// sizeof(*out->result_tmp)); -// out->ss_mutations = num_sample_sets > 0 -// ? tsk_calloc(num_sample_sets, sizeof(out->ss_mutations)) -// : NULL; -// for (i = 0; i < num_sample_sets; i++) { -// if (sample_set_sizes[i] > max_ss_size) { -// max_ss_size = sample_set_sizes[i]; -// } -// } -// tsk_memset(&out->AB_samples, 0, sizeof(out->AB_samples)); -// if (out->weights == NULL || out->norm == NULL || out->result_tmp == NULL) { -// ret = tsk_trace_error(TSK_ERR_NO_MEMORY); -// goto out; -// } -// ret = tsk_bitset_init(&out->AB_samples, max_ss_size, 1); -// if (ret != 0) { -// goto out; -// } -// for (i = 0; i < num_sample_sets; i++) { -// ret = tsk_bitset_init(out->ss_mutations[i], sample_set_sizes[i], 1); -// if (ret != 0) { -// goto out; -// } -// } -// out: -// return ret; -// } - typedef struct { double *weights; double *norm; @@ -2508,34 +2459,6 @@ get_site_row_col_indices(tsk_size_t n_rows, const tsk_id_t *row_sites, tsk_size_ *n_sites = s; } -// static int -// get_mutation_sample_sets(const tsk_bitset_t *allele_samples, -// const tsk_bitset_t *sample_sets, two_locus_work_t *restrict work) -// { -// int ret = 0; -// tsk_size_t i, j, k, n_ss_samples; -// tsk_id_t sample, *ss_samples = NULL; - -// ss_samples = tsk_calloc( -// allele_samples->row_len * TSK_BIT_ARRAY_NUM_BITS, sizeof(*ss_samples)); -// if (ss_samples == NULL) { -// ret = tsk_trace_error(TSK_ERR_NO_MEMORY); -// goto out; -// } - -// for (i = 0; i < sample_sets->len; i++) { -// tsk_bitset_get_items(sample_sets, i, ss_samples, &n_ss_samples); -// if (n_ss_samples == 0) { -// continue; -// } -// for (j = 0; j < n_ss_samples; j++) { -// sample = ss_samples[0]; -// } -// } -// out: -// return ret; -// } - static int get_mutation_samples(const tsk_treeseq_t *ts, const tsk_id_t *sites, tsk_size_t n_sites, tsk_size_t *num_alleles, tsk_bitset_t *allele_samples) From 05c89e55f8a0401180f74c6be1e808e78d3d4b1e Mon Sep 17 00:00:00 2001 From: lkirk Date: Thu, 31 Jul 2025 02:15:45 -0500 Subject: [PATCH 65/70] fix stats for >1 sample set; need a dup sample check --- c/tskit/trees.c | 73 +++++++++++++++++++++++++------------------------ 1 file changed, 37 insertions(+), 36 deletions(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index aa5ad14c3f..3ea7d1e75c 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -2346,7 +2346,7 @@ compute_general_normed_two_site_stat_result(const tsk_bitset_t *state, // a3 [s1, s2, s3] [s1, s2, s3] [s1, s2, s3] tsk_size_t k, mut_a, mut_b, result_row_len = num_b_alleles * result_dim; uint8_t is_polarised = polarised ? 1 : 0; - double *hap_weight_row, *result_tmp_row; + double *restrict hap_row, *restrict result_tmp_row; double *restrict norm = work->norm; double *restrict weights = work->weights; @@ -2357,14 +2357,14 @@ compute_general_normed_two_site_stat_result(const tsk_bitset_t *state, result_tmp_row = GET_2D_ROW(result_tmp, result_row_len, mut_a); for (mut_b = is_polarised; mut_b < num_b_alleles; mut_b++) { for (k = 0; k < state_dim; k++) { - tsk_bitset_intersect( - state, mut_a + a_off, state, mut_b + b_off, &AB_samples); - hap_weight_row = GET_2D_ROW(weights, 3, k); - hap_weight_row[0] = (double) tsk_bitset_count(&AB_samples, 0); - hap_weight_row[1] - = (double) allele_counts[mut_a + a_off] - hap_weight_row[0]; - hap_weight_row[2] - = (double) allele_counts[mut_b + b_off] - hap_weight_row[0]; + tsk_bitset_intersect(state, a_off + (mut_a * state_dim) + k, state, + b_off + (mut_b * state_dim) + k, &AB_samples); + hap_row = GET_2D_ROW(weights, 3, k); + hap_row[0] = (double) tsk_bitset_count(&AB_samples, 0); + hap_row[1] = (double) allele_counts[a_off + (mut_a * state_dim) + k] + - hap_row[0]; + hap_row[2] = (double) allele_counts[b_off + (mut_b * state_dim) + k] + - hap_row[0]; } ret = f(state_dim, weights, result_dim, result_tmp_row, f_params); if (ret != 0) { @@ -2381,7 +2381,6 @@ compute_general_normed_two_site_stat_result(const tsk_bitset_t *state, result_tmp_row += result_dim; // Advance to the next column } } - out: return ret; } @@ -2390,21 +2389,24 @@ static int compute_general_two_site_stat_result(const tsk_bitset_t *state, const tsk_size_t *allele_counts, tsk_size_t a_off, tsk_size_t b_off, tsk_size_t state_dim, tsk_size_t result_dim, general_stat_func_t *f, - sample_count_stat_params_t *f_params, bool polarised, - two_locus_work_t *restrict work, double *result) + sample_count_stat_params_t *f_params, two_locus_work_t *restrict work, + double *result) { int ret = 0; tsk_size_t k; tsk_bitset_t AB_samples = work->AB_samples; - tsk_size_t mut_a = polarised ? 1 : 0, mut_b = 1; + tsk_size_t mut_a = 1, mut_b = 1; double *restrict hap_row, *restrict weights = work->weights; for (k = 0; k < state_dim; k++) { - tsk_bitset_intersect(state, mut_a + a_off, state, mut_b + b_off, &AB_samples); + tsk_bitset_intersect(state, a_off + (mut_a * state_dim) + k, state, + b_off + (mut_b * state_dim) + k, &AB_samples); hap_row = GET_2D_ROW(weights, 3, k); hap_row[0] = (double) tsk_bitset_count(&AB_samples, 0); - hap_row[1] = (double) allele_counts[mut_a + a_off] - hap_row[0]; - hap_row[2] = (double) allele_counts[mut_b + b_off] - hap_row[0]; + hap_row[1] + = (double) allele_counts[a_off + (mut_a * state_dim) + k] - hap_row[0]; + hap_row[2] + = (double) allele_counts[b_off + (mut_b * state_dim) + k] - hap_row[0]; } ret = f(state_dim, weights, result_dim, result, f_params); if (ret != 0) { @@ -2602,11 +2604,10 @@ tsk_treeseq_two_site_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim, tsk_bitset_t allele_samples, allele_sample_sets; bool polarised = false; tsk_id_t *sites; - tsk_size_t r, c, s, max_ss_size, max_alleles, n_alleles, n_sites, *row_idx, *col_idx; + tsk_size_t i, j, max_ss_size, max_alleles, n_alleles, n_sites, *row_idx, *col_idx; double *result_row; const tsk_size_t num_samples = self->num_samples; - tsk_size_t *num_alleles = NULL, *site_offsets = NULL, - *allele_sample_set_counts = NULL; + tsk_size_t *num_alleles = NULL, *site_offsets = NULL, *allele_counts = NULL; tsk_size_t result_row_len = n_cols * result_dim; two_locus_work_t work; @@ -2633,11 +2634,11 @@ tsk_treeseq_two_site_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim, } n_alleles = 0; max_alleles = 0; - for (s = 0; s < n_sites; s++) { - site_offsets[s] = n_alleles; - n_alleles += self->site_mutations_length[sites[s]] + 1; - if (self->site_mutations_length[sites[s]] > max_alleles) { - max_alleles = self->site_mutations_length[sites[s]]; + for (i = 0; i < n_sites; i++) { + site_offsets[i] = n_alleles * num_sample_sets; + n_alleles += self->site_mutations_length[sites[i]] + 1; + if (self->site_mutations_length[sites[i]] > max_alleles) { + max_alleles = self->site_mutations_length[sites[i]]; } } max_alleles++; // add 1 for the ancestral allele @@ -2651,7 +2652,7 @@ tsk_treeseq_two_site_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim, goto out; } ret = get_mutation_sample_sets(&allele_samples, num_sample_sets, sample_set_sizes, - sample_sets, &max_ss_size, &allele_sample_sets, &allele_sample_set_counts); + sample_sets, &max_ss_size, &allele_sample_sets, &allele_counts); if (ret != 0) { goto out; } @@ -2665,20 +2666,20 @@ tsk_treeseq_two_site_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim, } // For each row/column pair, fill in the sample set in the result matrix. - for (r = 0; r < n_rows; r++) { - result_row = GET_2D_ROW(result, result_row_len, r); - for (c = 0; c < n_cols; c++) { - if (num_alleles[row_idx[r]] == 2 && num_alleles[col_idx[c]] == 2) { + for (i = 0; i < n_rows; i++) { + result_row = GET_2D_ROW(result, result_row_len, i); + for (j = 0; j < n_cols; j++) { + if (num_alleles[row_idx[i]] == 2 && num_alleles[col_idx[j]] == 2) { ret = compute_general_two_site_stat_result(&allele_sample_sets, - allele_sample_set_counts, site_offsets[row_idx[r]], - site_offsets[col_idx[c]], state_dim, result_dim, f, f_params, - polarised, &work, &(result_row[c * result_dim])); + allele_counts, site_offsets[row_idx[i]], site_offsets[col_idx[j]], + state_dim, result_dim, f, f_params, &work, + &(result_row[j * result_dim])); } else { ret = compute_general_normed_two_site_stat_result(&allele_sample_sets, - allele_sample_set_counts, site_offsets[row_idx[r]], - site_offsets[col_idx[c]], num_alleles[row_idx[r]], - num_alleles[col_idx[c]], state_dim, result_dim, f, f_params, norm_f, - polarised, &work, &(result_row[c * result_dim])); + allele_counts, site_offsets[row_idx[i]], site_offsets[col_idx[j]], + num_alleles[row_idx[i]], num_alleles[col_idx[j]], state_dim, + result_dim, f, f_params, norm_f, polarised, &work, + &(result_row[j * result_dim])); } if (ret != 0) { goto out; From 89b75fe1dc2f3cf1f16ad65c34389f37585387f1 Mon Sep 17 00:00:00 2001 From: lkirk Date: Thu, 31 Jul 2025 03:05:59 -0500 Subject: [PATCH 66/70] add dup sample check; now all tests pass --- c/tskit/trees.c | 43 +++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 41 insertions(+), 2 deletions(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index 3ea7d1e75c..3081071293 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -2234,7 +2234,6 @@ norm_hap_weighted(tsk_size_t result_dim, const double *hap_weights, for (k = 0; k < result_dim; k++) { weight_row = GET_2D_ROW(hap_weights, 3, k); n = (double) args.sample_set_sizes[k]; - // TODO: what to do when n = 0 result[k] = weight_row[0] / n; } return 0; @@ -2271,9 +2270,10 @@ norm_total_weighted(tsk_size_t result_dim, const double *TSK_UNUSED(hap_weights) tsk_size_t n_a, tsk_size_t n_b, double *result, void *TSK_UNUSED(params)) { tsk_size_t k; + double norm = 1 / (double) (n_a * n_b); for (k = 0; k < result_dim; k++) { - result[k] = 1 / (double) (n_a * n_b); + result[k] = norm; } return 0; } @@ -3286,6 +3286,40 @@ tsk_treeseq_two_branch_count_stat(const tsk_treeseq_t *self, tsk_size_t state_di return ret; } +static int +check_sample_set_dups(tsk_size_t num_sample_sets, const tsk_size_t *sample_set_sizes, + const tsk_id_t *sample_sets, const tsk_id_t *restrict sample_index_map, + tsk_size_t num_samples) +{ + int ret; + tsk_size_t j, k, l; + tsk_id_t u, sample_index; + tsk_bitset_t tmp; + + tsk_memset(&tmp, 0, sizeof(tmp)); + ret = tsk_bitset_init(&tmp, num_samples, 1); + if (ret != 0) { + return ret; + } + j = 0; + for (k = 0; k < num_sample_sets; k++) { + tsk_memset(tmp.data, 0, sizeof(*tmp.data) * tmp.len); + for (l = 0; l < sample_set_sizes[k]; l++) { + u = sample_sets[j]; + sample_index = sample_index_map[u]; + if (tsk_bitset_contains(&tmp, k, (tsk_bitset_val_t) sample_index)) { + ret = tsk_trace_error(TSK_ERR_DUPLICATE_SAMPLE); + goto out; + } + tsk_bitset_set_bit(&tmp, k, (tsk_bitset_val_t) sample_index); + j++; + } + } +out: + tsk_bitset_free(&tmp); + return ret; +} + int tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sample_sets, const tsk_size_t *sample_set_sizes, const tsk_id_t *sample_sets, @@ -3338,6 +3372,11 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl if (ret != 0) { goto out; } + ret = check_sample_set_dups(num_sample_sets, sample_set_sizes, sample_sets, + self->sample_index_map, self->num_samples); + if (ret != 0) { + goto out; + } // TODO: result dim/state dim can be set internally now. ret = tsk_treeseq_two_site_count_stat(self, state_dim, num_sample_sets, sample_set_sizes, sample_sets, result_dim, f, &f_params, norm_f, out_rows, From bd4c9ff1cb01b95049fd719085c89d293a4bd325 Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 8 Aug 2025 14:50:23 -0500 Subject: [PATCH 67/70] merge cleanup --- c/tskit/trees.c | 1 + python/tests/test_ld_matrix.py | 8 +------- python/tskit/trees.py | 18 +++++++++--------- 3 files changed, 11 insertions(+), 16 deletions(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index 3081071293..25e8ceee30 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -3333,6 +3333,7 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl int ret = 0; bool stat_site = !!(options & TSK_STAT_SITE); bool stat_branch = !!(options & TSK_STAT_BRANCH); + bool stat_node = !!(options & TSK_STAT_NODE); tsk_size_t state_dim = num_sample_sets; sample_count_stat_params_t f_params = { .sample_sets = sample_sets, .num_sample_sets = num_sample_sets, diff --git a/python/tests/test_ld_matrix.py b/python/tests/test_ld_matrix.py index fcec932b05..5131b02822 100644 --- a/python/tests/test_ld_matrix.py +++ b/python/tests/test_ld_matrix.py @@ -590,9 +590,6 @@ def compute_general_two_site_stat_result( for k in range(result_dim): result[k] += result_tmp[k] * norm[k] - # for k in range(result_dim): - # print(mut_a, mut_b, k, weights[0, k], weights[1, k], weights[2, k], sep="\t") - def two_site_count_stat( ts: tskit.TreeSequence, @@ -1055,10 +1052,8 @@ def r2_ij_summary_func( D_j = pAB - pA * pB denom_j = np.sqrt(pA * (1 - pA) * pB * (1 - pB)) - p_A = (w_A_i + w_A_j) / (ni + nj) - p_B = (w_B_i + w_B_j) / (ni + nj) with suppress_overflow_div0_warning(): - result[k] = result[k] = (D_i * D_j) / (denom_i * denom_j) + result[k] = (D_i * D_j) / (denom_i * denom_j) def D_summary_func( @@ -1780,7 +1775,6 @@ def test_ld_empty_examples(ts): def test_input_validation(): - # TODO ts = get_paper_ex_ts() with pytest.raises(ValueError, match="Unknown two-locus statistic"): ts.ld_matrix(stat="bad_stat") diff --git a/python/tskit/trees.py b/python/tskit/trees.py index 14a50d2fde..77ccbf4838 100644 --- a/python/tskit/trees.py +++ b/python/tskit/trees.py @@ -23,7 +23,6 @@ """ Module responsible for managing trees and tree sequences. """ - from __future__ import annotations import base64 @@ -687,7 +686,8 @@ def __init__( options = 0 if sample_counts is not None: warnings.warn( - "The sample_counts option is not supported since 0.2.4 and is ignored", + "The sample_counts option is not supported since 0.2.4 " + "and is ignored", RuntimeWarning, stacklevel=4, ) @@ -6840,7 +6840,7 @@ def to_macs(self): bytes_genotypes[:] = lookup[variant.genotypes] genotypes = bytes_genotypes.tobytes().decode() output.append( - f"SITE:\t{variant.index}\t{variant.position / m}\t0.0\t{genotypes}" + f"SITE:\t{variant.index}\t{variant.position / m}\t0.0\t" f"{genotypes}" ) return "\n".join(output) + "\n" @@ -9143,9 +9143,9 @@ def pca( if time_windows is None: tree_sequence_low, tree_sequence_high = None, self else: - assert time_windows[0] < time_windows[1], ( - "The second argument should be larger." - ) + assert ( + time_windows[0] < time_windows[1] + ), "The second argument should be larger." tree_sequence_low, tree_sequence_high = ( self.decapitate(time_windows[0]), self.decapitate(time_windows[1]), @@ -9213,9 +9213,9 @@ def _rand_pow_range_finder( """ Algorithm 9 in https://arxiv.org/pdf/2002.01387 """ - assert num_vectors >= rank > 0, ( - "num_vectors should not be smaller than rank" - ) + assert ( + num_vectors >= rank > 0 + ), "num_vectors should not be smaller than rank" for _ in range(depth): Q = np.linalg.qr(Q)[0] Q = operator(Q) From 20bf634ddc8e5a2d7e42289d45d31d7d767fc94f Mon Sep 17 00:00:00 2001 From: lkirk Date: Fri, 8 Aug 2025 17:09:42 -0500 Subject: [PATCH 68/70] zero out whole row, not just one element --- c/tskit/trees.c | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index 25e8ceee30..48af3843ef 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -3303,7 +3303,7 @@ check_sample_set_dups(tsk_size_t num_sample_sets, const tsk_size_t *sample_set_s } j = 0; for (k = 0; k < num_sample_sets; k++) { - tsk_memset(tmp.data, 0, sizeof(*tmp.data) * tmp.len); + tsk_memset(tmp.data, 0, sizeof(*tmp.data) * tmp.row_len); for (l = 0; l < sample_set_sizes[k]; l++) { u = sample_sets[j]; sample_index = sample_index_map[u]; From a8a6c221427fdeeeeaf3262dbad3bbc4b9b92aaa Mon Sep 17 00:00:00 2001 From: lkirk Date: Sat, 9 Aug 2025 10:49:55 -0500 Subject: [PATCH 69/70] fix sample dup checking; wrong tmp row; not using index map --- c/tskit/trees.c | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index 48af3843ef..f6f01e7fbb 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -2549,11 +2549,11 @@ get_mutation_samples(const tsk_treeseq_t *ts, const tsk_id_t *sites, tsk_size_t static int get_mutation_sample_sets(const tsk_bitset_t *allele_samples, tsk_size_t num_sample_sets, const tsk_size_t *sample_set_sizes, const tsk_id_t *sample_sets, - tsk_size_t *max_ss_size, tsk_bitset_t *allele_sample_sets, - tsk_size_t **allele_sample_set_counts) + const tsk_id_t *sample_index_map, tsk_size_t *max_ss_size, + tsk_bitset_t *allele_sample_sets, tsk_size_t **allele_sample_set_counts) { int ret = 0; - tsk_bitset_val_t k; + tsk_bitset_val_t k, sample; tsk_size_t i, j, ss_off; *max_ss_size = 0; @@ -2579,8 +2579,8 @@ get_mutation_sample_sets(const tsk_bitset_t *allele_samples, tsk_size_t num_samp ss_off = 0; for (j = 0; j < num_sample_sets; j++) { for (k = 0; k < sample_set_sizes[j]; k++) { - if (tsk_bitset_contains( - allele_samples, i, (tsk_bitset_val_t) sample_sets[k + ss_off])) { + sample = (tsk_bitset_val_t) sample_index_map[sample_sets[k + ss_off]]; + if (tsk_bitset_contains(allele_samples, i, sample)) { tsk_bitset_set_bit(allele_sample_sets, j + i * num_sample_sets, k); (*allele_sample_set_counts)[j + i * num_sample_sets]++; } @@ -2652,7 +2652,8 @@ tsk_treeseq_two_site_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim, goto out; } ret = get_mutation_sample_sets(&allele_samples, num_sample_sets, sample_set_sizes, - sample_sets, &max_ss_size, &allele_sample_sets, &allele_counts); + sample_sets, self->sample_index_map, &max_ss_size, &allele_sample_sets, + &allele_counts); if (ret != 0) { goto out; } @@ -3307,11 +3308,11 @@ check_sample_set_dups(tsk_size_t num_sample_sets, const tsk_size_t *sample_set_s for (l = 0; l < sample_set_sizes[k]; l++) { u = sample_sets[j]; sample_index = sample_index_map[u]; - if (tsk_bitset_contains(&tmp, k, (tsk_bitset_val_t) sample_index)) { + if (tsk_bitset_contains(&tmp, 0, (tsk_bitset_val_t) sample_index)) { ret = tsk_trace_error(TSK_ERR_DUPLICATE_SAMPLE); goto out; } - tsk_bitset_set_bit(&tmp, k, (tsk_bitset_val_t) sample_index); + tsk_bitset_set_bit(&tmp, 0, (tsk_bitset_val_t) sample_index); j++; } } From 9bb79fe5f7725529b8f81dd8fc5b2678f8043551 Mon Sep 17 00:00:00 2001 From: lkirk Date: Thu, 14 Aug 2025 23:08:36 -0500 Subject: [PATCH 70/70] clean commit to remove noise; will get on PR --- python/tests/test_ld_matrix.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/tests/test_ld_matrix.py b/python/tests/test_ld_matrix.py index 5131b02822..a3eece1e2d 100644 --- a/python/tests/test_ld_matrix.py +++ b/python/tests/test_ld_matrix.py @@ -277,6 +277,7 @@ def norm_hap_weighted_ij( wAB_i = hap_weights[0, i] wAB_j = hap_weights[0, j] result[k] = (wAB_i + wAB_j) / (ni + nj) + # result[k] = (wAB_i / ni / 2) + (wAB_j / nj / 2) def norm_total_weighted(