From d93738646eb28cc8b4600af38044cf46c5226a28 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 30 Oct 2026 10:25:52 +0000 Subject: [PATCH] Remove haplotype matching from the C library The Li and Stephens HMM (tsk_ls_hmm_t and the compressed and Viterbi matrices) is replaced by the tsls package, which tsinfer and sc2ts use. Nothing else in the library uses it. Remove the low-level haplotype matching classes _tskit.LsHmm, CompressedMatrix and ViterbiMatrix wrapped the C haplotype matching code removed in the previous commit, and were never public API. Their tests compared against the lshmm package, which is no longer needed. The tests that were worth keeping are ported to tsls. --- c/CHANGELOG.rst | 8 + c/meson.build | 9 +- c/tests/test_haplotype_matching.c | 579 -------- c/tskit.h | 1 - c/tskit/core.c | 16 - c/tskit/core.h | 22 - c/tskit/haplotype_matching.c | 1770 ----------------------- c/tskit/haplotype_matching.h | 169 --- python/CHANGELOG.rst | 2 +- python/_tskitmodule.c | 816 ----------- python/lwt_interface/setup.py | 1 - python/pyproject.toml | 1 - python/setup.py | 1 - python/tests/test_genotype_matching.py | 1504 ------------------- python/tests/test_haplotype_matching.py | 1327 ----------------- python/tests/test_python_c.py | 178 --- python/uv.lock | 36 +- 17 files changed, 21 insertions(+), 6419 deletions(-) delete mode 100644 c/tests/test_haplotype_matching.c delete mode 100644 c/tskit/haplotype_matching.c delete mode 100644 c/tskit/haplotype_matching.h delete mode 100644 python/tests/test_genotype_matching.py delete mode 100644 python/tests/test_haplotype_matching.py diff --git a/c/CHANGELOG.rst b/c/CHANGELOG.rst index 05e12d021b..253d64316e 100644 --- a/c/CHANGELOG.rst +++ b/c/CHANGELOG.rst @@ -10,6 +10,14 @@ In development - Update to kastore 2.1.3 to include file security fix (:user:`jeromekelleher`, :pr:`3478`). +**Breaking changes** + +- Remove the ``haplotype_matching.c`` file which implemented an experimental + Li and Stephens model (replaced by the + `tsls `_ package). This code was never + part of the public API, but users may need to edit build recipes to + remove the missing file. (:pr:`3499`, :user:`jeromekelleher`). + **Bugfixes** - Fix off-by-one node ID bounds check in ``tsk_table_collection_link_ancestors`` diff --git a/c/meson.build b/c/meson.build index fb89c06907..d37efd815a 100644 --- a/c/meson.build +++ b/c/meson.build @@ -25,10 +25,10 @@ extra_c_args = [ lib_sources = [ 'tskit/core.c', 'tskit/tables.c', 'tskit/trees.c', - 'tskit/genotypes.c', 'tskit/stats.c', 'tskit/convert.c', 'tskit/haplotype_matching.c'] + 'tskit/genotypes.c', 'tskit/stats.c', 'tskit/convert.c'] lib_headers = [ 'tskit/core.h', 'tskit/tables.h', 'tskit/trees.h', - 'tskit/genotypes.h', 'tskit/stats.h', 'tskit/convert.h', 'tskit/haplotype_matching.h'] + 'tskit/genotypes.h', 'tskit/stats.h', 'tskit/convert.h'] # Subprojects use the static library for simplicity. tskit_inc = [kastore_inc, include_directories(['.'])] @@ -82,11 +82,6 @@ if not meson.is_subproject() link_with: [tskit_lib, test_lib], c_args: extra_c_args, dependencies: kastore_dep) test('stats', test_stats) - test_haplotype_matching = executable('test_haplotype_matching', - sources: ['tests/test_haplotype_matching.c'], - link_with: [tskit_lib, test_lib], c_args: extra_c_args, dependencies: kastore_dep) - test('haplotype_matching', test_haplotype_matching) - test_file_format = executable('test_file_format', sources: ['tests/test_file_format.c'], link_with: [tskit_lib, test_lib], c_args: extra_c_args, dependencies: kastore_dep) diff --git a/c/tests/test_haplotype_matching.c b/c/tests/test_haplotype_matching.c deleted file mode 100644 index 943654ff91..0000000000 --- a/c/tests/test_haplotype_matching.c +++ /dev/null @@ -1,579 +0,0 @@ -/* - * MIT License - * - * Copyright (c) 2019-2023 Tskit Developers - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in all - * copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE - * SOFTWARE. - */ - -#include "testlib.h" -#include - -#include -#include - -static void -test_single_tree_missing_alleles(void) -{ - int ret = 0; - tsk_treeseq_t ts; - tsk_ls_hmm_t ls_hmm; - tsk_compressed_matrix_t forward; - tsk_viterbi_matrix_t viterbi; - - double rho[] = { 0, 0.25, 0.25 }; - double mu[] = { 0.125, 0.125, 0.125 }; - int32_t h[] = { 0, 0, 0, 0 }; - - tsk_treeseq_from_text(&ts, 1, single_tree_ex_nodes, single_tree_ex_edges, NULL, - single_tree_ex_sites, single_tree_ex_mutations, NULL, NULL, 0); - - ret = tsk_ls_hmm_init(&ls_hmm, &ts, rho, mu, TSK_ALLELES_ACGT); - CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = tsk_ls_hmm_forward(&ls_hmm, h, &forward, 0); - CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_ALLELE_NOT_FOUND); - ret = tsk_ls_hmm_viterbi(&ls_hmm, h, &viterbi, 0); - CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_ALLELE_NOT_FOUND); - - tsk_ls_hmm_free(&ls_hmm); - tsk_compressed_matrix_free(&forward); - tsk_viterbi_matrix_free(&viterbi); - tsk_treeseq_free(&ts); -} - -static void -test_single_tree_exact_match(void) -{ - int ret = 0; - tsk_treeseq_t ts; - tsk_ls_hmm_t ls_hmm; - tsk_compressed_matrix_t forward; - tsk_viterbi_matrix_t viterbi; - - double rho[] = { 0.0, 0.25, 0.25 }; - double mu[] = { 0, 0, 0 }; - int32_t h[] = { 1, 1, 1 }; - tsk_id_t path[3]; - double decoded_compressed_matrix[12]; - unsigned int precision; - - tsk_treeseq_from_text(&ts, 1, single_tree_ex_nodes, single_tree_ex_edges, NULL, - single_tree_ex_sites, single_tree_ex_mutations, NULL, NULL, 0); - - ret = tsk_ls_hmm_init(&ls_hmm, &ts, rho, mu, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = tsk_ls_hmm_forward(&ls_hmm, h, &forward, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - tsk_compressed_matrix_print_state(&forward, _devnull); - tsk_ls_hmm_print_state(&ls_hmm, _devnull); - ret = tsk_compressed_matrix_decode(&forward, decoded_compressed_matrix); - CU_ASSERT_EQUAL_FATAL(ret, 0); - - ret = tsk_ls_hmm_viterbi(&ls_hmm, h, &viterbi, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - tsk_viterbi_matrix_print_state(&viterbi, _devnull); - tsk_ls_hmm_print_state(&ls_hmm, _devnull); - ret = tsk_viterbi_matrix_traceback(&viterbi, path, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - CU_ASSERT_EQUAL(path[0], 2); - CU_ASSERT_EQUAL(path[1], 1); - CU_ASSERT_EQUAL(path[2], 1); - - /* Should get the same answer at lower precision */ - for (precision = 1; precision < 24; precision++) { - ret = tsk_ls_hmm_set_precision(&ls_hmm, precision); - CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = tsk_ls_hmm_viterbi(&ls_hmm, h, &viterbi, TSK_NO_INIT); - CU_ASSERT_EQUAL_FATAL(ret, 0); - tsk_viterbi_matrix_print_state(&viterbi, _devnull); - tsk_ls_hmm_print_state(&ls_hmm, _devnull); - ret = tsk_viterbi_matrix_traceback(&viterbi, path, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - CU_ASSERT_EQUAL(path[0], 2); - CU_ASSERT_EQUAL(path[1], 1); - CU_ASSERT_EQUAL(path[2], 1); - } - - tsk_ls_hmm_free(&ls_hmm); - tsk_compressed_matrix_free(&forward); - tsk_viterbi_matrix_free(&viterbi); - tsk_treeseq_free(&ts); -} - -static void -test_single_tree_missing_haplotype_data(void) -{ - int ret = 0; - tsk_treeseq_t ts; - tsk_ls_hmm_t ls_hmm; - tsk_compressed_matrix_t forward; - tsk_viterbi_matrix_t viterbi; - - double rho[] = { 0.0, 0.25, 0.25 }; - double mu[] = { 0, 0, 0 }; - int32_t h[] = { 1, TSK_MISSING_DATA, 1 }; - tsk_id_t path[3]; - double decoded_compressed_matrix[12]; - - tsk_treeseq_from_text(&ts, 1, single_tree_ex_nodes, single_tree_ex_edges, NULL, - single_tree_ex_sites, single_tree_ex_mutations, NULL, NULL, 0); - - ret = tsk_ls_hmm_init(&ls_hmm, &ts, rho, mu, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = tsk_ls_hmm_forward(&ls_hmm, h, &forward, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - tsk_compressed_matrix_print_state(&forward, _devnull); - tsk_ls_hmm_print_state(&ls_hmm, _devnull); - ret = tsk_compressed_matrix_decode(&forward, decoded_compressed_matrix); - CU_ASSERT_EQUAL_FATAL(ret, 0); - - ret = tsk_ls_hmm_viterbi(&ls_hmm, h, &viterbi, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - tsk_viterbi_matrix_print_state(&viterbi, _devnull); - tsk_ls_hmm_print_state(&ls_hmm, _devnull); - ret = tsk_viterbi_matrix_traceback(&viterbi, path, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - CU_ASSERT_EQUAL(path[0], 2); - CU_ASSERT_EQUAL(path[1], 2); - CU_ASSERT_EQUAL(path[2], 2); - - tsk_ls_hmm_free(&ls_hmm); - tsk_compressed_matrix_free(&forward); - tsk_viterbi_matrix_free(&viterbi); - tsk_treeseq_free(&ts); -} - -static void -test_single_tree_match_impossible(void) -{ - int ret = 0; - tsk_treeseq_t ts; - tsk_ls_hmm_t ls_hmm; - tsk_compressed_matrix_t forward; - tsk_compressed_matrix_t backward; - tsk_viterbi_matrix_t viterbi; - - double rho[] = { 0.0, 0.25, 0.25 }; - double mu[] = { 0, 0, 0 }; - /* This haplotype can't happen with a mutation rate of 0 */ - int32_t h[] = { 0, 0, 0 }; - - tsk_treeseq_from_text(&ts, 1, single_tree_ex_nodes, single_tree_ex_edges, NULL, - single_tree_ex_sites, single_tree_ex_mutations, NULL, NULL, 0); - - ret = tsk_ls_hmm_init(&ls_hmm, &ts, rho, mu, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = tsk_ls_hmm_forward(&ls_hmm, h, &forward, 0); - CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_MATCH_IMPOSSIBLE); - tsk_compressed_matrix_print_state(&forward, _devnull); - tsk_ls_hmm_print_state(&ls_hmm, _devnull); - - ret = tsk_ls_hmm_viterbi(&ls_hmm, h, &viterbi, 0); - CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_MATCH_IMPOSSIBLE); - tsk_viterbi_matrix_print_state(&viterbi, _devnull); - tsk_ls_hmm_print_state(&ls_hmm, _devnull); - - ret = tsk_ls_hmm_backward(&ls_hmm, h, forward.normalisation_factor, &backward, 0); - CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_MATCH_IMPOSSIBLE); - tsk_compressed_matrix_print_state(&backward, _devnull); - /* tsk_compressed_matrix_print_state(&forward, stdout); */ - /* tsk_compressed_matrix_print_state(&backward, stdout); */ - tsk_ls_hmm_print_state(&ls_hmm, _devnull); - - tsk_ls_hmm_free(&ls_hmm); - tsk_compressed_matrix_free(&forward); - tsk_compressed_matrix_free(&backward); - tsk_viterbi_matrix_free(&viterbi); - tsk_treeseq_free(&ts); -} - -static void -test_single_tree_errors(void) -{ - int ret = 0; - tsk_treeseq_t ts; - tsk_ls_hmm_t ls_hmm; - tsk_compressed_matrix_t forward; - tsk_viterbi_matrix_t viterbi; - tsk_value_transition_t T[1]; - double decoded[3][4]; - - double rho[] = { 0.0, 0.25, 0.25 }; - double mu[] = { 0, 0, 0 }; - int32_t h[] = { 0, 0, 0 }; - - tsk_treeseq_from_text(&ts, 1, single_tree_ex_nodes, single_tree_ex_edges, NULL, - single_tree_ex_sites, single_tree_ex_mutations, NULL, NULL, 0); - - ret = tsk_viterbi_matrix_init(&viterbi, &ts, 0, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = tsk_compressed_matrix_init(&forward, &ts, 0, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - - ret = tsk_ls_hmm_init(&ls_hmm, &ts, rho, mu, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - - viterbi.matrix.tree_sequence = NULL; - ret = tsk_ls_hmm_viterbi(&ls_hmm, h, &viterbi, TSK_NO_INIT); - CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_BAD_PARAM_VALUE); - viterbi.matrix.tree_sequence = &ts; - - forward.tree_sequence = NULL; - ret = tsk_ls_hmm_forward(&ls_hmm, h, &forward, TSK_NO_INIT); - CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_BAD_PARAM_VALUE); - forward.tree_sequence = &ts; - - ret = tsk_compressed_matrix_store_site(&forward, 3, 0, 0, NULL); - CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_SITE_OUT_OF_BOUNDS); - ret = tsk_compressed_matrix_store_site(&forward, 4, 0, 0, NULL); - CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_SITE_OUT_OF_BOUNDS); - - /* FIXME disabling this tests for now because we filter out negative - * nodes when storing now, to accomodate some oddness in the initial - * conditions of the backward matrix. */ - /* T[0].tree_node = -1; */ - /* T[0].value = 0; */ - /* ret = tsk_compressed_matrix_store_site(&forward, 0, 1, 1, T); */ - /* CU_ASSERT_EQUAL_FATAL(ret, 0); */ - /* ret = tsk_compressed_matrix_decode(&forward, (double *) decoded); */ - /* CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_NODE_OUT_OF_BOUNDS); */ - - T[0].tree_node = 7; - T[0].value = 0; - ret = tsk_compressed_matrix_store_site(&forward, 0, 1, 1, T); - CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = tsk_compressed_matrix_decode(&forward, (double *) decoded); - CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_NODE_OUT_OF_BOUNDS); - - tsk_ls_hmm_free(&ls_hmm); - tsk_compressed_matrix_free(&forward); - tsk_viterbi_matrix_free(&viterbi); - tsk_treeseq_free(&ts); -} - -static void -test_single_tree_compressed_matrix(void) -{ - int ret = 0; - tsk_treeseq_t ts; - tsk_compressed_matrix_t matrix; - tsk_ls_hmm_t ls_hmm; - tsk_size_t max_transitions = 1024; - tsk_value_transition_t T[max_transitions]; - double decoded[3][4]; - int j; - - double rho[] = { 0.0, 0.25, 0.25 }; - double mu[] = { 0.1, 0.1, 0.1 }; - int32_t h[] = { 0, 0, 0 }; - - tsk_treeseq_from_text(&ts, 1, single_tree_ex_nodes, single_tree_ex_edges, NULL, - single_tree_ex_sites, single_tree_ex_mutations, NULL, NULL, 0); - - ret = tsk_compressed_matrix_init(&matrix, &ts, 0, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - - tsk_compressed_matrix_print_state(&matrix, _devnull); - - T[0].tree_node = 6; - T[0].value = 0; - for (j = 0; j < 3; j++) { - T[1].tree_node = j; - T[1].value = 1; - ret = tsk_compressed_matrix_store_site(&matrix, j, 1.0, 2, T); - CU_ASSERT_EQUAL_FATAL(ret, 0); - } - tsk_compressed_matrix_print_state(&matrix, _devnull); - - ret = tsk_compressed_matrix_decode(&matrix, (double *) decoded); - CU_ASSERT_EQUAL_FATAL(ret, 0); - CU_ASSERT_EQUAL(decoded[0][0], 1.0); - CU_ASSERT_EQUAL(decoded[0][1], 0.0); - CU_ASSERT_EQUAL(decoded[0][2], 0.0); - CU_ASSERT_EQUAL(decoded[1][0], 0.0); - CU_ASSERT_EQUAL(decoded[1][1], 1.0); - CU_ASSERT_EQUAL(decoded[1][2], 0.0); - CU_ASSERT_EQUAL(decoded[2][0], 0.0); - CU_ASSERT_EQUAL(decoded[2][1], 0.0); - CU_ASSERT_EQUAL(decoded[2][2], 1.0); - - /* Cleared matrix should be zero everywhere */ - tsk_compressed_matrix_clear(&matrix); - ret = tsk_compressed_matrix_decode(&matrix, (double *) decoded); - CU_ASSERT_EQUAL_FATAL(ret, 0); - for (j = 0; j < 3; j++) { - CU_ASSERT_EQUAL(decoded[j][0], 0.0); - CU_ASSERT_EQUAL(decoded[j][1], 0.0); - CU_ASSERT_EQUAL(decoded[j][2], 0.0); - } - - ret = tsk_ls_hmm_init(&ls_hmm, &ts, rho, mu, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = tsk_ls_hmm_forward(&ls_hmm, h, &matrix, TSK_NO_INIT); - CU_ASSERT_EQUAL_FATAL(ret, 0); - - tsk_compressed_matrix_print_state(&matrix, _devnull); - ret = tsk_compressed_matrix_decode(&matrix, (double *) decoded); - CU_ASSERT_EQUAL_FATAL(ret, 0); - - tsk_compressed_matrix_free(&matrix); - tsk_ls_hmm_free(&ls_hmm); - tsk_treeseq_free(&ts); -} - -static void -test_single_tree_viterbi_matrix(void) -{ - int ret = 0; - tsk_treeseq_t ts; - tsk_viterbi_matrix_t viterbi; - tsk_ls_hmm_t ls_hmm; - double rho[] = { 0.0, 0.25, 0.25 }; - double mu[] = { 0, 0, 0 }; - int32_t h[] = { 1, 1, 1 }; - tsk_id_t path[3]; - tsk_value_transition_t T[2]; - int j; - - tsk_treeseq_from_text(&ts, 1, single_tree_ex_nodes, single_tree_ex_edges, NULL, - single_tree_ex_sites, single_tree_ex_mutations, NULL, NULL, 0); - - ret = tsk_viterbi_matrix_init(&viterbi, &ts, 0, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - - tsk_viterbi_matrix_print_state(&viterbi, _devnull); - ret = tsk_viterbi_matrix_traceback(&viterbi, path, 0); - CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_NULL_VITERBI_MATRIX); - - T[0].tree_node = 6; - T[0].value = 0; - T[1].tree_node = 1; - T[1].value = 1; - for (j = 0; j < 3; j++) { - ret = tsk_compressed_matrix_store_site(&viterbi.matrix, j, 1.0, 2, T); - CU_ASSERT_EQUAL_FATAL(ret, 0); - /* We need to have one record per site, so we put in a record - * at the root saying we don't need to recombine */ - ret = tsk_viterbi_matrix_add_recombination_required(&viterbi, j, 6, false); - CU_ASSERT_EQUAL_FATAL(ret, 0); - } - ret = tsk_viterbi_matrix_traceback(&viterbi, path, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - CU_ASSERT_EQUAL_FATAL(path[0], 1); - CU_ASSERT_EQUAL_FATAL(path[1], 1); - CU_ASSERT_EQUAL_FATAL(path[2], 1); - - ret = tsk_ls_hmm_init(&ls_hmm, &ts, rho, mu, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = tsk_ls_hmm_viterbi(&ls_hmm, h, &viterbi, TSK_NO_INIT); - CU_ASSERT_EQUAL_FATAL(ret, 0); - - tsk_viterbi_matrix_print_state(&viterbi, _devnull); - ret = tsk_viterbi_matrix_traceback(&viterbi, path, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - - ret = tsk_viterbi_matrix_clear(&viterbi); - CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = tsk_viterbi_matrix_traceback(&viterbi, path, 0); - CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_NULL_VITERBI_MATRIX); - - tsk_viterbi_matrix_free(&viterbi); - - ret = tsk_viterbi_matrix_init(&viterbi, &ts, 1, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - /* Make sure we hit the realloc case for recombination records */ - for (j = 0; j < 100; j++) { - ret = tsk_viterbi_matrix_add_recombination_required(&viterbi, 0, 6, false); - CU_ASSERT_EQUAL_FATAL(ret, 0); - } - tsk_viterbi_matrix_print_state(&viterbi, _devnull); - - tsk_viterbi_matrix_free(&viterbi); - tsk_ls_hmm_free(&ls_hmm); - tsk_treeseq_free(&ts); -} - -static void -test_multi_tree_exact_match(void) -{ - int ret = 0; - tsk_treeseq_t ts; - tsk_ls_hmm_t ls_hmm; - tsk_compressed_matrix_t forward, backward; - tsk_viterbi_matrix_t viterbi; - - double rho[] = { 0.0, 0.25, 0.25 }; - double mu[] = { 0, 0, 0 }; - int32_t h[] = { 1, 1, 1 }; - tsk_id_t path[3]; - double decoded_compressed_matrix[12]; - unsigned int precision; - - tsk_treeseq_from_text(&ts, 10, paper_ex_nodes, paper_ex_edges, NULL, paper_ex_sites, - paper_ex_mutations, paper_ex_individuals, NULL, 0); - - ret = tsk_ls_hmm_init(&ls_hmm, &ts, rho, mu, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = tsk_ls_hmm_forward(&ls_hmm, h, &forward, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - tsk_ls_hmm_print_state(&ls_hmm, _devnull); - tsk_compressed_matrix_print_state(&forward, _devnull); - ret = tsk_compressed_matrix_decode(&forward, decoded_compressed_matrix); - CU_ASSERT_EQUAL_FATAL(ret, 0); - - ret = tsk_ls_hmm_backward(&ls_hmm, h, forward.normalisation_factor, &backward, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - tsk_ls_hmm_print_state(&ls_hmm, _devnull); - tsk_compressed_matrix_print_state(&backward, _devnull); - ret = tsk_compressed_matrix_decode(&backward, decoded_compressed_matrix); - CU_ASSERT_EQUAL_FATAL(ret, 0); - - ret = tsk_ls_hmm_viterbi(&ls_hmm, h, &viterbi, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - tsk_viterbi_matrix_print_state(&viterbi, _devnull); - tsk_ls_hmm_print_state(&ls_hmm, _devnull); - ret = tsk_viterbi_matrix_traceback(&viterbi, path, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - CU_ASSERT_EQUAL(path[0], 2); - CU_ASSERT_EQUAL(path[1], 0); - CU_ASSERT_EQUAL(path[2], 1); - - /* Should get the same answer at lower precision */ - for (precision = 4; precision < 24; precision++) { - ret = tsk_ls_hmm_set_precision(&ls_hmm, precision); - CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = tsk_ls_hmm_viterbi(&ls_hmm, h, &viterbi, TSK_NO_INIT); - CU_ASSERT_EQUAL_FATAL(ret, 0); - tsk_viterbi_matrix_print_state(&viterbi, _devnull); - tsk_ls_hmm_print_state(&ls_hmm, _devnull); - ret = tsk_viterbi_matrix_traceback(&viterbi, path, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - CU_ASSERT_EQUAL(path[0], 2); - CU_ASSERT_EQUAL(path[1], 0); - CU_ASSERT_EQUAL(path[2], 1); - } - - tsk_ls_hmm_free(&ls_hmm); - tsk_compressed_matrix_free(&forward); - tsk_compressed_matrix_free(&backward); - tsk_viterbi_matrix_free(&viterbi); - tsk_treeseq_free(&ts); -} - -static void -test_multi_tree_errors(void) -{ - int ret = 0; - tsk_treeseq_t ts; - tsk_compressed_matrix_t forward; - tsk_value_transition_t T[1]; - double decoded[3][4]; - - tsk_treeseq_from_text(&ts, 10, paper_ex_nodes, paper_ex_edges, NULL, paper_ex_sites, - paper_ex_mutations, paper_ex_individuals, NULL, 0); - - ret = tsk_compressed_matrix_init(&forward, &ts, 0, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - - /* We want a tree node that is not in the first tree */ - T[0].tree_node = 7; - T[0].value = 0; - ret = tsk_compressed_matrix_store_site(&forward, 0, 1, 1, T); - CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = tsk_compressed_matrix_decode(&forward, (double *) decoded); - CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_BAD_COMPRESSED_MATRIX_NODE); - - tsk_compressed_matrix_free(&forward); - tsk_treeseq_free(&ts); -} - -static void -test_caterpillar_tree_many_values(void) -{ - int ret = 0; - tsk_ls_hmm_t ls_hmm; - tsk_compressed_matrix_t matrix; - double rho[] = { 0.1, 0.1, 0.1, 0.1, 0.1 }; - double mu[] = { 0.0, 0.0, 0.0, 0.0, 0.0 }; - int32_t h[] = { 0, 0, 0, 0, 0 }; - tsk_size_t n[] = { - 8, - 16, - 32, - 64, - }; - tsk_treeseq_t *ts; - tsk_size_t j; - - for (j = 0; j < sizeof(n) / sizeof(*n); j++) { - ts = caterpillar_tree(n[j], 5, n[j] - 2); - ret = tsk_ls_hmm_init(&ls_hmm, ts, rho, mu, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = tsk_compressed_matrix_init(&matrix, ts, 1 << 10, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = tsk_ls_hmm_forward(&ls_hmm, h, &matrix, TSK_NO_INIT); - CU_ASSERT_EQUAL_FATAL(ret, 0); - tsk_compressed_matrix_print_state(&matrix, _devnull); - tsk_ls_hmm_print_state(&ls_hmm, _devnull); - - tsk_ls_hmm_free(&ls_hmm); - tsk_compressed_matrix_free(&matrix); - tsk_treeseq_free(ts); - free(ts); - } - - j = 40; - ts = caterpillar_tree(j, 5, j - 2); - ret = tsk_ls_hmm_init(&ls_hmm, ts, rho, mu, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - ret = tsk_compressed_matrix_init(&matrix, ts, 1 << 20, 0); - CU_ASSERT_EQUAL_FATAL(ret, 0); - /* Short circuit this value so we can run the test */ - ls_hmm.max_parsimony_words = 0; - ret = tsk_ls_hmm_forward(&ls_hmm, h, &matrix, TSK_NO_INIT); - CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_TOO_MANY_VALUES); - - tsk_ls_hmm_free(&ls_hmm); - tsk_compressed_matrix_free(&matrix); - tsk_treeseq_free(ts); - free(ts); -} - -int -main(int argc, char **argv) -{ - CU_TestInfo tests[] = { - { "test_single_tree_missing_alleles", test_single_tree_missing_alleles }, - { "test_single_tree_exact_match", test_single_tree_exact_match }, - { "test_single_tree_missing_haplotype_data", - test_single_tree_missing_haplotype_data }, - { "test_single_tree_match_impossible", test_single_tree_match_impossible }, - { "test_single_tree_errors", test_single_tree_errors }, - { "test_single_tree_compressed_matrix", test_single_tree_compressed_matrix }, - { "test_single_tree_viterbi_matrix", test_single_tree_viterbi_matrix }, - - { "test_multi_tree_exact_match", test_multi_tree_exact_match }, - { "test_multi_tree_errors", test_multi_tree_errors }, - - { "test_caterpillar_tree_many_values", test_caterpillar_tree_many_values }, - { NULL, NULL }, - }; - - return test_main(tests, argc, argv); -} diff --git a/c/tskit.h b/c/tskit.h index e8f1b69f60..7f977c9d23 100644 --- a/c/tskit.h +++ b/c/tskit.h @@ -34,6 +34,5 @@ #include #include #include -#include #endif diff --git a/c/tskit/core.c b/c/tskit/core.c index 91e684c52d..2bc575b93e 100644 --- a/c/tskit/core.c +++ b/c/tskit/core.c @@ -620,22 +620,6 @@ tsk_strerror_internal(int err) "the tree. (TSK_ERR_NO_SAMPLE_LISTS)"; break; - /* Haplotype matching errors */ - case TSK_ERR_NULL_VITERBI_MATRIX: - ret = "Viterbi matrix has not filled. (TSK_ERR_NULL_VITERBI_MATRIX)"; - break; - case TSK_ERR_MATCH_IMPOSSIBLE: - ret = "No matching haplotype exists with current parameters. " - "(TSK_ERR_MATCH_IMPOSSIBLE)"; - break; - case TSK_ERR_BAD_COMPRESSED_MATRIX_NODE: - ret = "The compressed matrix contains a node that subtends no samples. " - "(TSK_ERR_BAD_COMPRESSED_MATRIX_NODE)"; - break; - case TSK_ERR_TOO_MANY_VALUES: - ret = "Too many values to compress. (TSK_ERR_TOO_MANY_VALUES)"; - break; - /* Union errors */ case TSK_ERR_UNION_BAD_MAP: ret = "Node map contains an entry of a node not present in this table " diff --git a/c/tskit/core.h b/c/tskit/core.h index 9512a905e7..bd958129c0 100644 --- a/c/tskit/core.h +++ b/c/tskit/core.h @@ -844,28 +844,6 @@ enabled (TSK_SAMPLE_LISTS). #define TSK_ERR_NO_SAMPLE_LISTS -1205 /** @} */ -/** -@defgroup HAPLOTYPE_ERROR_GROUP Haplotype matching errors. -@{ -*/ -/** -The Viterbi matrix has not filled (it has zero transitions). -*/ -#define TSK_ERR_NULL_VITERBI_MATRIX -1300 -/** -There was no matching haplotype. -*/ -#define TSK_ERR_MATCH_IMPOSSIBLE -1301 -/** -The compressed matrix has a node that has no samples in it's descendants. -*/ -#define TSK_ERR_BAD_COMPRESSED_MATRIX_NODE -1302 -/** -There are too many values to compress. -*/ -#define TSK_ERR_TOO_MANY_VALUES -1303 -/** @} */ - /** @defgroup UNION_ERROR_GROUP Union errors. @{ diff --git a/c/tskit/haplotype_matching.c b/c/tskit/haplotype_matching.c deleted file mode 100644 index d68776a2a8..0000000000 --- a/c/tskit/haplotype_matching.c +++ /dev/null @@ -1,1770 +0,0 @@ -/* - * MIT License - * - * Copyright (c) 2019-2025 Tskit Developers - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in all - * copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE - * SOFTWARE. - */ - -#include -#include -#include -#include -#include -#include - -#include - -#define MAX_PARSIMONY_WORDS 256 - -const char *_zero_one_alleles[] = { "0", "1", NULL }; -const char *_acgt_alleles[] = { "A", "C", "G", "T", NULL }; - -static int -cmp_double(const void *a, const void *b) -{ - const double *ia = (const double *) a; - const double *ib = (const double *) b; - return (*ia > *ib) - (*ia < *ib); -} - -static int -cmp_argsort(const void *a, const void *b) -{ - const tsk_argsort_t *ia = (const tsk_argsort_t *) a; - const tsk_argsort_t *ib = (const tsk_argsort_t *) b; - int ret = (ia->value > ib->value) - (ia->value < ib->value); - /* Break any ties using the index to ensure consistency */ - if (ret == 0) { - ret = (ia->index > ib->index) - (ia->index < ib->index); - } - return ret; -} - -static void -tsk_ls_hmm_check_state(tsk_ls_hmm_t *self) -{ - tsk_id_t *T_index = self->transition_index; - tsk_value_transition_t *T = self->transitions; - tsk_id_t j; - - for (j = 0; j < (tsk_id_t) self->num_transitions; j++) { - if (T[j].tree_node != TSK_NULL) { - tsk_bug_assert(T_index[T[j].tree_node] == j); - } - } - /* tsk_bug_assert(self->num_transitions <= self->num_samples); */ - - if (self->num_transitions > 0) { - for (j = 0; j < (tsk_id_t) self->num_nodes; j++) { - if (T_index[j] != TSK_NULL) { - tsk_bug_assert(T[T_index[j]].tree_node == j); - } - tsk_bug_assert(self->tree.parent[j] == self->parent[j]); - } - } -} - -void -tsk_ls_hmm_print_state(tsk_ls_hmm_t *self, FILE *out) -{ - tsk_size_t j, l; - - fprintf(out, "tree_sequence = %p\n", (void *) self->tree_sequence); - fprintf(out, "num_sites = %lld\n", (long long) self->num_sites); - fprintf(out, "num_samples = %lld\n", (long long) self->num_samples); - fprintf(out, "num_values = %lld\n", (long long) self->num_values); - fprintf(out, "max_values = %lld\n", (long long) self->max_values); - fprintf(out, "num_optimal_value_set_words = %lld\n", - (long long) self->num_optimal_value_set_words); - - fprintf(out, "sites::\n"); - for (l = 0; l < self->num_sites; l++) { - fprintf(out, "%lld\t%lld\t[", (long long) l, (long long) self->num_alleles[l]); - for (j = 0; j < self->num_alleles[l]; j++) { - fprintf(out, "%s,", self->alleles[l][j]); - } - fprintf(out, "]\n"); - } - fprintf(out, "transitions::%lld\n", (long long) self->num_transitions); - for (j = 0; j < self->num_transitions; j++) { - fprintf(out, "tree_node=%lld\tvalue=%.14f\tvalue_index=%lld\n", - (long long) self->transitions[j].tree_node, self->transitions[j].value, - (long long) self->transitions[j].value_index); - } - if (self->num_transitions > 0) { - fprintf(out, "tree::%lld\n", (long long) self->num_nodes); - for (j = 0; j < self->num_nodes; j++) { - fprintf(out, "%lld\tparent=%lld\ttransition=%lld\n", (long long) j, - (long long) self->parent[j], (long long) self->transition_index[j]); - } - } - tsk_ls_hmm_check_state(self); -} - -int TSK_WARN_UNUSED -tsk_ls_hmm_init(tsk_ls_hmm_t *self, tsk_treeseq_t *tree_sequence, - double *recombination_rate, double *mutation_rate, tsk_flags_t options) -{ - int ret = TSK_ERR_GENERIC; - tsk_size_t l; - - tsk_memset(self, 0, sizeof(tsk_ls_hmm_t)); - self->tree_sequence = tree_sequence; - self->precision = 6; /* Seems like a safe value, but probably not ideal for perf */ - self->num_sites = tsk_treeseq_get_num_sites(tree_sequence); - self->num_samples = tsk_treeseq_get_num_samples(tree_sequence); - self->num_alleles = tsk_malloc(self->num_sites * sizeof(*self->num_alleles)); - self->num_nodes = tsk_treeseq_get_num_nodes(tree_sequence); - self->parent = tsk_malloc(self->num_nodes * sizeof(*self->parent)); - self->allelic_state = tsk_malloc(self->num_nodes * sizeof(*self->allelic_state)); - self->transition_index - = tsk_malloc(self->num_nodes * sizeof(*self->transition_index)); - self->transition_stack - = tsk_malloc(self->num_nodes * sizeof(*self->transition_stack)); - /* We can't have more than 2 * num_samples transitions, so we use this as the - * upper bound. Because of the implementation, we'll also have to worry about - * the extra mutations at the first site, which in worst case involves all - * mutations. We can definitely save some memory here if we want to.*/ - self->max_transitions - = 2 * self->num_samples + tsk_treeseq_get_num_mutations(tree_sequence); - /* FIXME Arbitrarily doubling this after hitting problems */ - self->max_transitions *= 2; - self->transitions = tsk_malloc(self->max_transitions * sizeof(*self->transitions)); - self->transitions_copy - = tsk_malloc(self->max_transitions * sizeof(*self->transitions)); - self->num_transition_samples - = tsk_malloc(self->max_transitions * sizeof(*self->num_transition_samples)); - self->transition_parent - = tsk_malloc(self->max_transitions * sizeof(*self->transition_parent)); - self->transition_time_order - = tsk_malloc(self->max_transitions * sizeof(*self->transition_time_order)); - self->values = tsk_malloc(self->max_transitions * sizeof(*self->values)); - self->recombination_rate - = tsk_malloc(self->num_sites * sizeof(*self->recombination_rate)); - self->mutation_rate = tsk_malloc(self->num_sites * sizeof(*self->mutation_rate)); - self->alleles = tsk_calloc(self->num_sites, sizeof(*self->alleles)); - if (self->num_alleles == NULL || self->parent == NULL || self->allelic_state == NULL - || self->transition_index == NULL || self->transition_stack == NULL - || self->transitions == NULL || self->transitions_copy == NULL - || self->num_transition_samples == NULL || self->transition_parent == NULL - || self->transition_time_order == NULL || self->values == NULL - || self->recombination_rate == NULL || self->mutation_rate == NULL - || self->alleles == NULL) { - ret = tsk_trace_error(TSK_ERR_NO_MEMORY); - goto out; - } - for (l = 0; l < self->num_sites; l++) { - /* TODO check these inputs */ - self->recombination_rate[l] = recombination_rate[l]; - self->mutation_rate[l] = mutation_rate[l]; - if (options & TSK_ALLELES_ACGT) { - self->num_alleles[l] = 4; - self->alleles[l] = _acgt_alleles; - } else { - /* Default to the 0/1 alleles */ - self->num_alleles[l] = 2; - self->alleles[l] = _zero_one_alleles; - } - } - ret = tsk_tree_init(&self->tree, self->tree_sequence, 0); - if (ret != 0) { - goto out; - } - self->num_values = 0; - self->max_values = 0; - /* Keep this as a struct variable so that we can test overflow, but this - * should never be set to more than MAX_PARSIMONY_WORDS as we're doing - * a bunch of stack allocations based on this. */ - self->max_parsimony_words = MAX_PARSIMONY_WORDS; - ret = 0; -out: - return ret; -} - -int -tsk_ls_hmm_set_precision(tsk_ls_hmm_t *self, unsigned int precision) -{ - self->precision = precision; - return 0; -} - -int -tsk_ls_hmm_free(tsk_ls_hmm_t *self) -{ - tsk_tree_free(&self->tree); - tsk_safe_free(self->recombination_rate); - tsk_safe_free(self->mutation_rate); - tsk_safe_free(self->recombination_rate); - tsk_safe_free(self->alleles); - tsk_safe_free(self->num_alleles); - tsk_safe_free(self->parent); - tsk_safe_free(self->allelic_state); - tsk_safe_free(self->transition_index); - tsk_safe_free(self->transition_stack); - tsk_safe_free(self->transitions); - tsk_safe_free(self->transitions_copy); - tsk_safe_free(self->transition_time_order); - tsk_safe_free(self->values); - tsk_safe_free(self->num_transition_samples); - tsk_safe_free(self->transition_parent); - tsk_safe_free(self->optimal_value_sets); - return 0; -} - -static int -tsk_ls_hmm_reset(tsk_ls_hmm_t *self, double value) -{ - int ret = 0; - tsk_size_t j; - tsk_id_t u; - const tsk_id_t *samples; - tsk_size_t N = self->num_nodes; - - tsk_memset(self->parent, 0xff, N * sizeof(*self->parent)); - tsk_memset(self->transition_index, 0xff, N * sizeof(*self->transition_index)); - tsk_memset(self->allelic_state, 0xff, N * sizeof(*self->allelic_state)); - tsk_memset(self->transitions, 0, self->max_transitions * sizeof(*self->transitions)); - tsk_memset(self->num_transition_samples, 0, - self->max_transitions * sizeof(*self->num_transition_samples)); - tsk_memset(self->transition_parent, 0xff, - self->max_transitions * sizeof(*self->transition_parent)); - - samples = tsk_treeseq_get_samples(self->tree_sequence); - for (j = 0; j < self->num_samples; j++) { - u = samples[j]; - self->transitions[j].tree_node = u; - self->transitions[j].value = value; - self->transition_index[u] = (tsk_id_t) j; - } - self->num_transitions = self->num_samples; - return ret; -} - -/* After we have moved on to a new tree we can have transitions still associated - * with the old roots, which are now disconnected. Remove. */ -static int -tsk_ls_hmm_remove_dead_roots(tsk_ls_hmm_t *self) -{ - tsk_id_t *restrict T_index = self->transition_index; - tsk_value_transition_t *restrict T = self->transitions; - const tsk_id_t *restrict right_sib = self->tree.right_sib; - const tsk_id_t left_root = tsk_tree_get_left_root(&self->tree); - const tsk_id_t *restrict parent = self->parent; - tsk_id_t root, u; - tsk_size_t j; - const tsk_id_t root_marker = -2; - - for (root = left_root; root != TSK_NULL; root = right_sib[root]) { - if (T_index[root] != TSK_NULL) { - /* Use the value_index slot as a marker. We don't use this between - * iterations, so it's safe to appropriate here */ - T[T_index[root]].value_index = root_marker; - } - } - for (j = 0; j < self->num_transitions; j++) { - u = T[j].tree_node; - if (u != TSK_NULL) { - if (parent[u] == TSK_NULL && T[j].value_index != root_marker) { - T_index[u] = TSK_NULL; - T[j].tree_node = TSK_NULL; - } - T[j].value_index = -1; - } - } - return 0; -} - -static int -tsk_ls_hmm_update_tree(tsk_ls_hmm_t *self, int direction) -{ - int ret = 0; - tsk_id_t *restrict parent = self->parent; - tsk_id_t *restrict T_index = self->transition_index; - const tsk_id_t *restrict edges_child = self->tree_sequence->tables->edges.child; - const tsk_id_t *restrict edges_parent = self->tree_sequence->tables->edges.parent; - tsk_value_transition_t *restrict T = self->transitions; - tsk_id_t u, c, p, j, e; - tsk_value_transition_t *vt; - tsk_tree_position_t tree_pos; - - tree_pos = self->tree.tree_pos; - for (j = tree_pos.out.start; j != tree_pos.out.stop; j += direction) { - e = tree_pos.out.order[j]; - c = edges_child[e]; - u = c; - if (T_index[u] == TSK_NULL) { - /* Ensure the subtree we're detaching has a transition at the root */ - while (T_index[u] == TSK_NULL) { - u = parent[u]; - tsk_bug_assert(u != TSK_NULL); - } - tsk_bug_assert(self->num_transitions < self->max_transitions); - T_index[c] = (tsk_id_t) self->num_transitions; - T[self->num_transitions].tree_node = c; - T[self->num_transitions].value = T[T_index[u]].value; - self->num_transitions++; - } - parent[c] = TSK_NULL; - } - - for (j = tree_pos.in.start; j != tree_pos.in.stop; j += direction) { - e = tree_pos.in.order[j]; - c = edges_child[e]; - p = edges_parent[e]; - parent[c] = p; - u = p; - if (parent[p] == TSK_NULL) { - /* Grafting onto a new root. */ - if (T_index[p] == TSK_NULL) { - T_index[p] = (tsk_id_t) self->num_transitions; - tsk_bug_assert(self->num_transitions < self->max_transitions); - T[self->num_transitions].tree_node = p; - T[self->num_transitions].value = T[T_index[c]].value; - self->num_transitions++; - } - } else { - /* Grafting into an existing subtree. */ - while (T_index[u] == TSK_NULL) { - u = parent[u]; - } - tsk_bug_assert(u != TSK_NULL); - } - tsk_bug_assert(T_index[u] != -1 && T_index[c] != -1); - if (T[T_index[u]].value == T[T_index[c]].value) { - vt = &T[T_index[c]]; - /* Mark the value transition as unusued */ - vt->value = -1; - vt->tree_node = TSK_NULL; - T_index[c] = TSK_NULL; - } - } - - ret = tsk_ls_hmm_remove_dead_roots(self); - return ret; -} - -static int -tsk_ls_hmm_get_allele_index(tsk_ls_hmm_t *self, tsk_id_t site, const char *allele_state, - const tsk_size_t allele_length) -{ - /* Note we're not doing tsk_trace_error here because it would require changing - * the logic of the function. Could be done easily enough, though */ - int ret = TSK_ERR_ALLELE_NOT_FOUND; - const char **alleles = self->alleles[site]; - const tsk_id_t num_alleles = (tsk_id_t) self->num_alleles[site]; - - tsk_id_t j; - - for (j = 0; j < num_alleles; j++) { - if (strlen(alleles[j]) != allele_length) { - break; - } - if (strncmp(alleles[j], allele_state, (size_t) allele_length) == 0) { - ret = (int) j; - break; - } - } - return ret; -} - -static int -tsk_ls_hmm_update_probabilities( - tsk_ls_hmm_t *self, const tsk_site_t *site, int32_t haplotype_state) -{ - int ret = 0; - tsk_id_t root; - tsk_tree_t *tree = &self->tree; - tsk_id_t *restrict parent = self->parent; - tsk_id_t *restrict T_index = self->transition_index; - tsk_value_transition_t *restrict T = self->transitions; - int32_t *restrict allelic_state = self->allelic_state; - const tsk_id_t left_root = tsk_tree_get_left_root(tree); - tsk_mutation_t mut; - tsk_id_t j, u, v; - double x; - bool match; - - /* Set the allelic states */ - ret = tsk_ls_hmm_get_allele_index( - self, site->id, site->ancestral_state, site->ancestral_state_length); - if (ret < 0) { - goto out; - } - for (root = left_root; root != TSK_NULL; root = tree->right_sib[root]) { - allelic_state[root] = (int32_t) ret; - } - - for (j = 0; j < (tsk_id_t) site->mutations_length; j++) { - mut = site->mutations[j]; - ret = tsk_ls_hmm_get_allele_index( - self, site->id, mut.derived_state, mut.derived_state_length); - if (ret < 0) { - goto out; - } - u = mut.node; - allelic_state[u] = (int32_t) ret; - if (T_index[u] == TSK_NULL) { - while (T_index[u] == TSK_NULL) { - u = parent[u]; - } - tsk_bug_assert(self->num_transitions < self->max_transitions); - T_index[mut.node] = (tsk_id_t) self->num_transitions; - T[self->num_transitions].tree_node = mut.node; - T[self->num_transitions].value = T[T_index[u]].value; - self->num_transitions++; - } - } - - for (j = 0; j < (tsk_id_t) self->num_transitions; j++) { - u = T[j].tree_node; - if (u != TSK_NULL) { - /* Get the allelic_state at u. */ - v = u; - while (allelic_state[v] == TSK_MISSING_DATA) { - v = parent[v]; - tsk_bug_assert(v != -1); - } - match = haplotype_state == TSK_MISSING_DATA - || haplotype_state == allelic_state[v]; - ret = self->next_probability(self, site->id, T[j].value, match, u, &x); - if (ret != 0) { - goto out; - } - T[j].value = x; - } - } - - /* Unset the allelic states */ - for (root = left_root; root != TSK_NULL; root = tree->right_sib[root]) { - allelic_state[root] = TSK_MISSING_DATA; - } - for (j = 0; j < (tsk_id_t) site->mutations_length; j++) { - mut = site->mutations[j]; - allelic_state[mut.node] = TSK_MISSING_DATA; - } - ret = 0; -out: - return ret; -} - -static int -tsk_ls_hmm_discretise_values(tsk_ls_hmm_t *self) -{ - int ret = 0; - tsk_value_transition_t *T = self->transitions; - double *values = self->values; - tsk_size_t j, k, num_values; - - num_values = 0; - for (j = 0; j < self->num_transitions; j++) { - if (T[j].tree_node != TSK_NULL) { - values[num_values] = T[j].value; - num_values++; - } - } - tsk_bug_assert(num_values > 0); - - qsort(values, (size_t) num_values, sizeof(double), cmp_double); - - k = 0; - for (j = 1; j < num_values; j++) { - if (values[j] != values[k]) { - k++; - values[k] = values[j]; - } - } - num_values = k + 1; - self->num_values = num_values; - for (j = 0; j < self->num_transitions; j++) { - if (T[j].tree_node != TSK_NULL) { - T[j].value_index - = (tsk_id_t) tsk_search_sorted(values, num_values, T[j].value); - tsk_bug_assert(T[j].value == self->values[T[j].value_index]); - } - } - return ret; -} - -/* - * TODO We also have these function in tree.c where they're used in the - * parsimony calculations (which are slightly different). It would be good to bring - * these together, or at least avoid having the same function in two - * files. Keeping it as it is for now so that it can be inlined, since - * it's perf-sensitive. */ - -static inline tsk_id_t -get_smallest_set_bit(uint64_t v) -{ - /* This is an inefficient implementation, there are several better - * approaches. On GCC we can use - * return (uint8_t) (__builtin_ffsll((long long) v) - 1); - */ - uint64_t t = 1; - tsk_id_t r = 0; - assert(v != 0); - - while ((v & t) == 0) { - t <<= 1; - r++; - } - return r; -} - -static inline uint64_t -set_bit(uint64_t value, uint8_t bit) -{ - return value | (1ULL << bit); -} - -static inline bool -bit_is_set(uint64_t value, uint8_t bit) -{ - return (value & (1ULL << bit)) != 0; -} - -static inline tsk_id_t -get_smallest_element(const uint64_t *restrict A, tsk_size_t u, tsk_size_t num_words) -{ - tsk_size_t base = u * num_words; - const uint64_t *restrict a = A + base; - tsk_id_t j = 0; - - while (a[j] == 0) { - j++; - tsk_bug_assert(j < (tsk_id_t) num_words); - } - return j * 64 + get_smallest_set_bit(a[j]); -} - -/* static variables are zero-initialised by default. */ -static const uint64_t zero_block[MAX_PARSIMONY_WORDS]; - -static inline bool -all_zero(const uint64_t *restrict A, tsk_id_t u, tsk_size_t num_words) -{ - if (num_words == 1) { - return A[u] == 0; - } else { - return tsk_memcmp( - zero_block, A + (tsk_size_t) u * num_words, num_words * sizeof(*A)) - == 0; - } -} - -static inline bool -element_in( - const uint64_t *restrict A, tsk_id_t u, const tsk_id_t state, tsk_size_t num_words) -{ - tsk_size_t index = ((tsk_size_t) u) * num_words + (tsk_size_t) (state / 64); - return (A[index] & (1ULL << (state % 64))) != 0; -} - -static inline void -set_optimal_value( - uint64_t *restrict A, tsk_id_t u, const tsk_size_t num_words, tsk_id_t state) -{ - tsk_size_t index = ((tsk_size_t) u) * num_words + (tsk_size_t) (state / 64); - tsk_bug_assert(((tsk_size_t) state) / 64 < num_words); - A[index] |= 1ULL << (state % 64); -} - -/* TODO the implementation here isn't particularly optimal and the way things - * were organised was really driven by the old Fitch parsimony algorithm - * (which only worked on binary trees. In particular, we should be working - * word-by-word where possible rather than iterating by values like we do here. - * Needs to be reworked when we're documenting/writing up this algorithm. - */ - -static void -compute_optimal_value_1(uint64_t *restrict A, const tsk_id_t *restrict left_child, - const tsk_id_t *restrict right_sib, const tsk_id_t u, const tsk_id_t parent_state, - const tsk_size_t num_values) -{ - tsk_id_t v; - uint64_t child; - tsk_size_t value_count[64], max_value_count; - uint8_t j; - - assert(num_values < 64); - - tsk_memset(value_count, 0, num_values * sizeof(*value_count)); - for (v = left_child[u]; v != TSK_NULL; v = right_sib[v]) { - child = A[v]; - /* If the set for a given child is empty, then we know it inherits - * directly from the parent state and must be a singleton set. */ - if (child == 0) { - child = 1ULL << parent_state; - } - for (j = 0; j < num_values; j++) { - value_count[j] += bit_is_set(child, j); - } - } - max_value_count = 0; - for (j = 0; j < num_values; j++) { - max_value_count = TSK_MAX(max_value_count, value_count[j]); - } - A[u] = 0; - for (j = 0; j < num_values; j++) { - if (value_count[j] == max_value_count) { - A[u] = set_bit(A[u], j); - } - } -} - -static void -compute_optimal_value_general(uint64_t *restrict A, const tsk_id_t *restrict left_child, - const tsk_id_t *restrict right_sib, const tsk_id_t u, const tsk_id_t parent_state, - const tsk_size_t num_values, const tsk_size_t num_words) -{ - tsk_id_t v; - uint64_t child[MAX_PARSIMONY_WORDS]; - uint64_t *Au; - tsk_size_t base, word, bit; - bool child_all_zero; - const tsk_id_t state_index = parent_state / 64; - const uint64_t state_word = 1ULL << (parent_state % 64); - tsk_size_t value_count[64 * MAX_PARSIMONY_WORDS], max_value_count; - tsk_size_t j; - - tsk_bug_assert(num_values < 64 * MAX_PARSIMONY_WORDS); - tsk_bug_assert(num_words <= MAX_PARSIMONY_WORDS); - for (j = 0; j < num_values; j++) { - value_count[j] = 0; - } - - for (v = left_child[u]; v != TSK_NULL; v = right_sib[v]) { - child_all_zero = true; - base = ((tsk_size_t) v) * num_words; - for (word = 0; word < num_words; word++) { - child[word] = A[base + word]; - child_all_zero = child_all_zero && (child[word] == 0); - } - /* If the set for a given child is empty, then we know it inherits - * directly from the parent state and must be a singleton set. */ - if (child_all_zero) { - child[state_index] = state_word; - } - for (j = 0; j < num_values; j++) { - word = j / 64; - bit = j % 64; - assert(word < num_words); - value_count[j] += bit_is_set(child[word], (uint8_t) bit); - } - } - max_value_count = 0; - for (j = 0; j < num_values; j++) { - max_value_count = TSK_MAX(max_value_count, value_count[j]); - } - - Au = A + ((size_t) u * num_words); - for (word = 0; word < num_words; word++) { - Au[word] = 0; - } - for (j = 0; j < num_values; j++) { - if (value_count[j] == max_value_count) { - word = j / 64; - bit = j % 64; - Au[word] = set_bit(Au[word], (uint8_t) bit); - } - } -} - -static void -compute_optimal_value(uint64_t *restrict A, const tsk_id_t *restrict left_child, - const tsk_id_t *restrict right_sib, const tsk_id_t u, const tsk_id_t parent_state, - const tsk_size_t num_values, const tsk_size_t num_words) -{ - if (num_words == 1) { - compute_optimal_value_1(A, left_child, right_sib, u, parent_state, num_values); - } else { - compute_optimal_value_general( - A, left_child, right_sib, u, parent_state, num_values, num_words); - } -} - -static int -tsk_ls_hmm_setup_optimal_value_sets(tsk_ls_hmm_t *self) -{ - int ret = 0; - - /* We expect that most of the time there will be one word per optimal_value set, - * but there will be times when we need more than one word. This approach - * lets us expand the memory if we need to, but when the number of - * values goes back below 64 we revert to using one word per set. We - * could in principle release back the memory as well, but it doesn't seem - * worth the bother. */ - self->num_optimal_value_set_words = (self->num_values / 64) + 1; - if (self->num_optimal_value_set_words > self->max_parsimony_words) { - ret = tsk_trace_error(TSK_ERR_TOO_MANY_VALUES); - goto out; - } - if (self->num_values >= self->max_values) { - self->max_values = self->num_optimal_value_set_words * 64; - tsk_safe_free(self->optimal_value_sets); - self->optimal_value_sets - = tsk_calloc(self->num_nodes * self->num_optimal_value_set_words, - sizeof(*self->optimal_value_sets)); - if (self->optimal_value_sets == NULL) { - ret = tsk_trace_error(TSK_ERR_NO_MEMORY); - goto out; - } - } -out: - return ret; -} - -static int -tsk_ls_hmm_build_optimal_value_sets(tsk_ls_hmm_t *self) -{ - int ret = 0; - const double *restrict node_time = self->tree_sequence->tables->nodes.time; - const tsk_id_t *restrict left_child = self->tree.left_child; - const tsk_id_t *restrict right_sib = self->tree.right_sib; - const tsk_id_t *restrict parent = self->parent; - const tsk_value_transition_t *restrict T = self->transitions; - const tsk_id_t *restrict T_index = self->transition_index; - tsk_argsort_t *restrict order = self->transition_time_order; - const tsk_size_t num_optimal_value_set_words = self->num_optimal_value_set_words; - uint64_t *restrict A = self->optimal_value_sets; - tsk_size_t j; - tsk_id_t u, v, state, parent_state; - - /* argsort the transitions by node time so we can visit them in the - * correct order */ - for (j = 0; j < self->num_transitions; j++) { - order[j].index = j; - order[j].value = DBL_MAX; - if (T[j].tree_node != TSK_NULL) { - order[j].value = node_time[T[j].tree_node]; - } - } - qsort(order, (size_t) self->num_transitions, sizeof(*order), cmp_argsort); - - for (j = 0; j < self->num_transitions; j++) { - u = T[order[j].index].tree_node; - if (u != TSK_NULL) { - state = T[order[j].index].value_index; - if (left_child[u] == TSK_NULL) { - /* leaf node */ - set_optimal_value(A, u, num_optimal_value_set_words, state); - } else { - compute_optimal_value(A, left_child, right_sib, u, state, - self->num_values, num_optimal_value_set_words); - } - v = parent[u]; - if (v != TSK_NULL) { - while (T_index[v] == TSK_NULL) { - v = parent[v]; - tsk_bug_assert(v != TSK_NULL); - } - parent_state = T[T_index[v]].value_index; - v = parent[u]; - while (T_index[v] == TSK_NULL) { - compute_optimal_value(A, left_child, right_sib, v, parent_state, - self->num_values, num_optimal_value_set_words); - v = parent[v]; - tsk_bug_assert(v != TSK_NULL); - } - } - } - } - return ret; -} - -static int -tsk_ls_hmm_redistribute_transitions(tsk_ls_hmm_t *self) -{ - int ret = 0; - const tsk_id_t *restrict left_child = self->tree.left_child; - const tsk_id_t *restrict right_sib = self->tree.right_sib; - const tsk_id_t *restrict parent = self->parent; - tsk_id_t *restrict T_index = self->transition_index; - tsk_id_t *restrict T_parent = self->transition_parent; - tsk_value_transition_t *restrict T = self->transitions; - tsk_value_transition_t *restrict T_old = self->transitions_copy; - tsk_transition_stack_t *stack = self->transition_stack; - uint64_t *restrict A = self->optimal_value_sets; - const tsk_size_t num_optimal_value_set_words = self->num_optimal_value_set_words; - tsk_transition_stack_t s, child_s; - tsk_id_t root, u, v; - int stack_top = 0; - tsk_size_t j, old_num_transitions; - - tsk_memcpy(T_old, T, self->num_transitions * sizeof(*T)); - old_num_transitions = self->num_transitions; - self->num_transitions = 0; - - /* TODO refactor this to push the virtual root onto the stack rather then - * iterating over the roots. See the existing parsimony implementations - * for an example. */ - for (root = tsk_tree_get_left_root(&self->tree); root != TSK_NULL; - root = right_sib[root]) { - stack[0].tree_node = root; - stack[0].old_state = T_old[T_index[root]].value_index; - stack[0].new_state - = get_smallest_element(A, (tsk_size_t) root, num_optimal_value_set_words); - stack[0].transition_parent = 0; - stack_top = 0; - - tsk_bug_assert(self->num_transitions < self->max_transitions); - T_parent[self->num_transitions] = TSK_NULL; - T[self->num_transitions].tree_node = stack[0].tree_node; - T[self->num_transitions].value_index = stack[0].new_state; - self->num_transitions++; - - while (stack_top >= 0) { - s = stack[stack_top]; - stack_top--; - for (v = left_child[s.tree_node]; v != TSK_NULL; v = right_sib[v]) { - child_s = s; - child_s.tree_node = v; - if (T_index[v] != TSK_NULL) { - child_s.old_state = T_old[T_index[v]].value_index; - } - if (!all_zero(A, v, num_optimal_value_set_words)) { - if (!element_in(A, v, s.new_state, num_optimal_value_set_words)) { - child_s.new_state = get_smallest_element( - A, (tsk_size_t) v, num_optimal_value_set_words); - child_s.transition_parent = (tsk_id_t) self->num_transitions; - /* Add a new transition */ - tsk_bug_assert(self->num_transitions < self->max_transitions); - T_parent[self->num_transitions] = s.transition_parent; - T[self->num_transitions].tree_node = v; - T[self->num_transitions].value_index = child_s.new_state; - self->num_transitions++; - } - stack_top++; - stack[stack_top] = child_s; - } else { - /* Node that we didn't visit when moving up the tree */ - if (s.old_state != s.new_state) { - tsk_bug_assert(self->num_transitions < self->max_transitions); - T_parent[self->num_transitions] = s.transition_parent; - T[self->num_transitions].tree_node = v; - T[self->num_transitions].value_index = s.old_state; - self->num_transitions++; - } - } - } - } - } - - /* Unset the old T_index pointers and optimal_value sets. */ - for (j = 0; j < old_num_transitions; j++) { - u = T_old[j].tree_node; - if (u != TSK_NULL) { - T_index[u] = TSK_NULL; - while (u != TSK_NULL && !all_zero(A, u, num_optimal_value_set_words)) { - tsk_memset(A + ((tsk_size_t) u) * num_optimal_value_set_words, 0, - num_optimal_value_set_words * sizeof(uint64_t)); - u = parent[u]; - } - } - } - /* Set the new pointers for transition nodes and the values.*/ - for (j = 0; j < self->num_transitions; j++) { - T_index[T[j].tree_node] = (tsk_id_t) j; - T[j].value = self->values[T[j].value_index]; - } - return ret; -} - -static int -tsk_ls_hmm_compress(tsk_ls_hmm_t *self) -{ - int ret = 0; - - ret = tsk_ls_hmm_discretise_values(self); - if (ret != 0) { - goto out; - } - ret = tsk_ls_hmm_setup_optimal_value_sets(self); - if (ret != 0) { - goto out; - } - ret = tsk_ls_hmm_build_optimal_value_sets(self); - if (ret != 0) { - goto out; - } - ret = tsk_ls_hmm_redistribute_transitions(self); - if (ret != 0) { - goto out; - } -out: - return ret; -} - -static int -tsk_ls_hmm_process_site_forward( - tsk_ls_hmm_t *self, const tsk_site_t *site, int32_t haplotype_state) -{ - int ret = 0; - double x, normalisation_factor; - tsk_compressed_matrix_t *output = (tsk_compressed_matrix_t *) self->output; - tsk_value_transition_t *restrict T = self->transitions; - const unsigned int precision = (unsigned int) self->precision; - tsk_size_t j; - - ret = tsk_ls_hmm_update_probabilities(self, site, haplotype_state); - if (ret != 0) { - goto out; - } - /* See notes in the Python implementation on why we don't want to compress - * here, but rather should be doing it after rounding. */ - ret = tsk_ls_hmm_compress(self); - if (ret != 0) { - goto out; - } - tsk_bug_assert(self->num_transitions <= self->num_samples); - normalisation_factor = self->compute_normalisation_factor(self); - - if (normalisation_factor == 0) { - ret = tsk_trace_error(TSK_ERR_MATCH_IMPOSSIBLE); - goto out; - } - for (j = 0; j < self->num_transitions; j++) { - tsk_bug_assert(T[j].tree_node != TSK_NULL); - x = T[j].value / normalisation_factor; - T[j].value = tsk_round(x, precision); - } - - ret = tsk_compressed_matrix_store_site( - output, site->id, normalisation_factor, (tsk_size_t) self->num_transitions, T); -out: - return ret; -} - -static int -tsk_ls_hmm_run_forward(tsk_ls_hmm_t *self, int32_t *haplotype) -{ - int ret = 0; - int t_ret; - const tsk_site_t *sites; - tsk_size_t j, num_sites; - const double n = (double) self->num_samples; - - ret = tsk_ls_hmm_reset(self, 1 / n); - if (ret != 0) { - goto out; - } - - for (t_ret = tsk_tree_first(&self->tree); t_ret == TSK_TREE_OK; - t_ret = tsk_tree_next(&self->tree)) { - ret = tsk_ls_hmm_update_tree(self, TSK_DIR_FORWARD); - if (ret != 0) { - goto out; - } - /* tsk_ls_hmm_check_state(self); */ - ret = tsk_tree_get_sites(&self->tree, &sites, &num_sites); - if (ret != 0) { - goto out; - } - for (j = 0; j < num_sites; j++) { - ret = tsk_ls_hmm_process_site_forward( - self, &sites[j], haplotype[sites[j].id]); - if (ret != 0) { - goto out; - } - } - } - /* Set to zero so we can print and check the state OK. */ - self->num_transitions = 0; - if (t_ret != 0) { - ret = t_ret; - goto out; - } -out: - return ret; -} - -/**************************************************************** - * Forward Algorithm - ****************************************************************/ - -static double -tsk_ls_hmm_compute_normalisation_factor_forward(tsk_ls_hmm_t *self) -{ - tsk_size_t *restrict N = self->num_transition_samples; - tsk_value_transition_t *restrict T = self->transitions; - const tsk_id_t *restrict T_parent = self->transition_parent; - const tsk_size_t *restrict num_samples = self->tree.num_samples; - const tsk_id_t num_transitions = (tsk_id_t) self->num_transitions; - double normalisation_factor; - tsk_id_t j; - - /* Compute the number of samples directly inheriting from each transition */ - for (j = 0; j < num_transitions; j++) { - tsk_bug_assert(T[j].tree_node != TSK_NULL); - N[j] = num_samples[T[j].tree_node]; - } - for (j = 0; j < num_transitions; j++) { - if (T_parent[j] != TSK_NULL) { - N[T_parent[j]] -= N[j]; - } - } - - /* Compute the normalising constant used to avoid underflow */ - normalisation_factor = 0; - for (j = 0; j < num_transitions; j++) { - normalisation_factor += (double) N[j] * T[j].value; - } - return normalisation_factor; -} - -static int -tsk_ls_hmm_next_probability_forward(tsk_ls_hmm_t *self, tsk_id_t site_id, double p_last, - bool is_match, tsk_id_t TSK_UNUSED(node), double *result) -{ - const double rho = self->recombination_rate[site_id]; - const double mu = self->mutation_rate[site_id]; - const double n = (double) self->num_samples; - const double num_alleles = self->num_alleles[site_id]; - double p_t, p_e; - - p_t = p_last * (1 - rho) + rho / n; - p_e = mu; - if (is_match) { - p_e = 1 - (num_alleles - 1) * mu; - } - *result = p_t * p_e; - return 0; -} - -int -tsk_ls_hmm_forward(tsk_ls_hmm_t *self, int32_t *haplotype, - tsk_compressed_matrix_t *output, tsk_flags_t options) -{ - int ret = 0; - - if (!(options & TSK_NO_INIT)) { - ret = tsk_compressed_matrix_init(output, self->tree_sequence, 0, 0); - if (ret != 0) { - goto out; - } - } else { - if (output->tree_sequence != self->tree_sequence) { - ret = tsk_trace_error(TSK_ERR_BAD_PARAM_VALUE); - goto out; - } - ret = tsk_compressed_matrix_clear(output); - if (ret != 0) { - goto out; - } - } - - self->next_probability = tsk_ls_hmm_next_probability_forward; - self->compute_normalisation_factor = tsk_ls_hmm_compute_normalisation_factor_forward; - self->output = output; - - ret = tsk_ls_hmm_run_forward(self, haplotype); -out: - return ret; -} - -/**************************************************************** - * Backward Algorithm - ****************************************************************/ - -static int -tsk_ls_hmm_next_probability_backward(tsk_ls_hmm_t *self, tsk_id_t site_id, double p_last, - bool is_match, tsk_id_t TSK_UNUSED(node), double *result) -{ - const double mu = self->mutation_rate[site_id]; - const double num_alleles = self->num_alleles[site_id]; - double p_e; - - p_e = mu; - if (is_match) { - p_e = 1 - (num_alleles - 1) * mu; - } - *result = p_last * p_e; - return 0; -} - -static int -tsk_ls_hmm_process_site_backward(tsk_ls_hmm_t *self, const tsk_site_t *site, - const int32_t haplotype_state, const double normalisation_factor) -{ - int ret = 0; - double x, b_last_sum; - tsk_compressed_matrix_t *output = (tsk_compressed_matrix_t *) self->output; - tsk_value_transition_t *restrict T = self->transitions; - const unsigned int precision = (unsigned int) self->precision; - const double rho = self->recombination_rate[site->id]; - const double n = (double) self->num_samples; - tsk_size_t j; - - /* FIXME!!! We are calling compress twice here because we need to compress - * immediately before calling store_site in order to filter out -1 nodes, - * and also (crucially) to ensure that the value transitions are listed - * in preorder, which we rely on later for decoding. - * - * https://github.com/tskit-dev/tskit/issues/2803 - */ - ret = tsk_ls_hmm_compress(self); - if (ret != 0) { - goto out; - } - ret = tsk_compressed_matrix_store_site( - output, site->id, normalisation_factor, (tsk_size_t) self->num_transitions, T); - if (ret != 0) { - goto out; - } - - ret = tsk_ls_hmm_update_probabilities(self, site, haplotype_state); - if (ret != 0) { - goto out; - } - /* DO WE NEED THIS compress?? See above */ - ret = tsk_ls_hmm_compress(self); - if (ret != 0) { - goto out; - } - tsk_bug_assert(self->num_transitions <= self->num_samples); - b_last_sum = self->compute_normalisation_factor(self); - for (j = 0; j < self->num_transitions; j++) { - tsk_bug_assert(T[j].tree_node != TSK_NULL); - x = rho * b_last_sum / n + (1 - rho) * T[j].value; - x /= normalisation_factor; - T[j].value = tsk_round(x, precision); - } -out: - return ret; -} - -static int -tsk_ls_hmm_run_backward( - tsk_ls_hmm_t *self, int32_t *haplotype, const double *forward_norm) -{ - int ret = 0; - int t_ret; - const tsk_site_t *sites; - double s; - tsk_size_t num_sites; - tsk_id_t j; - - ret = tsk_ls_hmm_reset(self, 1); - if (ret != 0) { - goto out; - } - - for (t_ret = tsk_tree_last(&self->tree); t_ret == TSK_TREE_OK; - t_ret = tsk_tree_prev(&self->tree)) { - ret = tsk_ls_hmm_update_tree(self, TSK_DIR_REVERSE); - if (ret != 0) { - goto out; - } - /* tsk_ls_hmm_check_state(self); */ - ret = tsk_tree_get_sites(&self->tree, &sites, &num_sites); - if (ret != 0) { - goto out; - } - for (j = (tsk_id_t) num_sites - 1; j >= 0; j--) { - s = forward_norm[sites[j].id]; - if (s <= 0) { - /* NOTE: I'm not sure if this is the correct interpretation, - * but norm values of 0 do lead to problems, and this seems - * like a simple way of guarding against it. We do seem to - * get norm values of 0 with impossible matches from the fwd - * matrix. - */ - ret = tsk_trace_error(TSK_ERR_MATCH_IMPOSSIBLE); - goto out; - } - ret = tsk_ls_hmm_process_site_backward( - self, &sites[j], haplotype[sites[j].id], s); - if (ret != 0) { - goto out; - } - } - } - /* Set to zero so we can print and check the state OK. */ - self->num_transitions = 0; - if (t_ret != 0) { - ret = t_ret; - goto out; - } -out: - return ret; -} - -int -tsk_ls_hmm_backward(tsk_ls_hmm_t *self, int32_t *haplotype, const double *forward_norm, - tsk_compressed_matrix_t *output, tsk_flags_t options) -{ - int ret = 0; - - if (!(options & TSK_NO_INIT)) { - ret = tsk_compressed_matrix_init(output, self->tree_sequence, 0, 0); - if (ret != 0) { - goto out; - } - } else { - if (output->tree_sequence != self->tree_sequence) { - ret = tsk_trace_error(TSK_ERR_BAD_PARAM_VALUE); - goto out; - } - ret = tsk_compressed_matrix_clear(output); - if (ret != 0) { - goto out; - } - } - - self->next_probability = tsk_ls_hmm_next_probability_backward; - self->compute_normalisation_factor = tsk_ls_hmm_compute_normalisation_factor_forward; - self->output = output; - - ret = tsk_ls_hmm_run_backward(self, haplotype, forward_norm); -out: - return ret; -} - -/**************************************************************** - * Viterbi Algorithm - ****************************************************************/ - -static double -tsk_ls_hmm_compute_normalisation_factor_viterbi(tsk_ls_hmm_t *self) -{ - tsk_value_transition_t *restrict T = self->transitions; - const tsk_id_t num_transitions = (tsk_id_t) self->num_transitions; - tsk_value_transition_t max_vt; - tsk_id_t j; - - max_vt.value = -1; - max_vt.tree_node = 0; /* keep compiler happy */ - tsk_bug_assert(num_transitions > 0); - for (j = 0; j < num_transitions; j++) { - tsk_bug_assert(T[j].tree_node != TSK_NULL); - if (T[j].value > max_vt.value) { - max_vt = T[j]; - } - } - return max_vt.value; -} - -static int -tsk_ls_hmm_next_probability_viterbi(tsk_ls_hmm_t *self, tsk_id_t site, double p_last, - bool is_match, tsk_id_t node, double *result) -{ - const double rho = self->recombination_rate[site]; - const double mu = self->mutation_rate[site]; - const double num_alleles = self->num_alleles[site]; - const double n = (double) self->num_samples; - double p_recomb, p_no_recomb, p_t, p_e; - bool recombination_required = false; - - p_no_recomb = p_last * (1 - rho + rho / n); - p_recomb = rho / n; - if (p_no_recomb > p_recomb) { - p_t = p_no_recomb; - } else { - p_t = p_recomb; - recombination_required = true; - } - p_e = mu; - if (is_match) { - p_e = 1 - (num_alleles - 1) * mu; - } - *result = p_t * p_e; - return tsk_viterbi_matrix_add_recombination_required( - self->output, site, node, recombination_required); -} - -int -tsk_ls_hmm_viterbi(tsk_ls_hmm_t *self, int32_t *haplotype, tsk_viterbi_matrix_t *output, - tsk_flags_t options) -{ - int ret = 0; - - if (!(options & TSK_NO_INIT)) { - ret = tsk_viterbi_matrix_init(output, self->tree_sequence, 0, 0); - if (ret != 0) { - goto out; - } - } else { - if (output->matrix.tree_sequence != self->tree_sequence) { - ret = tsk_trace_error(TSK_ERR_BAD_PARAM_VALUE); - goto out; - } - ret = tsk_viterbi_matrix_clear(output); - if (ret != 0) { - goto out; - } - } - - self->next_probability = tsk_ls_hmm_next_probability_viterbi; - self->compute_normalisation_factor = tsk_ls_hmm_compute_normalisation_factor_viterbi; - self->output = output; - - ret = tsk_ls_hmm_run_forward(self, haplotype); -out: - return ret; -} - -/**************************************************************** - * Compressed matrix - ****************************************************************/ - -int -tsk_compressed_matrix_init(tsk_compressed_matrix_t *self, tsk_treeseq_t *tree_sequence, - tsk_size_t block_size, tsk_flags_t options) -{ - int ret = 0; - - tsk_memset(self, 0, sizeof(*self)); - self->tree_sequence = tree_sequence; - self->options = options; - self->num_sites = tsk_treeseq_get_num_sites(tree_sequence); - self->num_samples = tsk_treeseq_get_num_samples(tree_sequence); - self->num_transitions = tsk_malloc(self->num_sites * sizeof(*self->num_transitions)); - self->normalisation_factor - = tsk_malloc(self->num_sites * sizeof(*self->normalisation_factor)); - self->values = tsk_malloc(self->num_sites * sizeof(*self->values)); - self->nodes = tsk_malloc(self->num_sites * sizeof(*self->nodes)); - if (self->num_transitions == NULL || self->values == NULL || self->nodes == NULL) { - ret = tsk_trace_error(TSK_ERR_NO_MEMORY); - goto out; - } - if (block_size == 0) { - block_size = 1 << 20; - } - ret = tsk_blkalloc_init(&self->memory, (size_t) block_size); - if (ret != 0) { - goto out; - } - ret = tsk_compressed_matrix_clear(self); -out: - return ret; -} - -int -tsk_compressed_matrix_free(tsk_compressed_matrix_t *self) -{ - tsk_blkalloc_free(&self->memory); - tsk_safe_free(self->num_transitions); - tsk_safe_free(self->normalisation_factor); - tsk_safe_free(self->values); - tsk_safe_free(self->nodes); - return 0; -} - -int -tsk_compressed_matrix_clear(tsk_compressed_matrix_t *self) -{ - tsk_blkalloc_reset(&self->memory); - tsk_memset( - self->num_transitions, 0, self->num_sites * sizeof(*self->num_transitions)); - tsk_memset(self->normalisation_factor, 0, - self->num_sites * sizeof(*self->normalisation_factor)); - return 0; -} - -void -tsk_compressed_matrix_print_state(tsk_compressed_matrix_t *self, FILE *out) -{ - tsk_size_t l, j; - - fprintf(out, "Compressed matrix for %p\n", (void *) self->tree_sequence); - fprintf(out, "num_sites = %lld\n", (long long) self->num_sites); - fprintf(out, "num_samples = %lld\n", (long long) self->num_samples); - for (l = 0; l < self->num_sites; l++) { - fprintf(out, "%lld\ts=%f\tv=%lld [", (long long) l, - self->normalisation_factor[l], (long long) self->num_transitions[l]); - for (j = 0; j < self->num_transitions[l]; j++) { - fprintf( - out, "(%lld, %f)", (long long) self->nodes[l][j], self->values[l][j]); - if (j < self->num_transitions[l] - 1) { - fprintf(out, ","); - } else { - fprintf(out, "]\n"); - } - } - } - fprintf(out, "Memory:\n"); - tsk_blkalloc_print_state(&self->memory, out); -} - -int -tsk_compressed_matrix_store_site(tsk_compressed_matrix_t *self, tsk_id_t site, - double normalisation_factor, tsk_size_t num_transitions, - const tsk_value_transition_t *transitions) -{ - int ret = 0; - tsk_size_t j; - - if (site < 0 || site >= (tsk_id_t) self->num_sites) { - ret = tsk_trace_error(TSK_ERR_SITE_OUT_OF_BOUNDS); - goto out; - } - - self->num_transitions[site] = num_transitions; - self->normalisation_factor[site] = normalisation_factor; - self->nodes[site] - = tsk_blkalloc_get(&self->memory, (size_t) num_transitions * sizeof(tsk_id_t)); - self->values[site] - = tsk_blkalloc_get(&self->memory, (size_t) num_transitions * sizeof(double)); - if (self->nodes[site] == NULL || self->values[site] == NULL) { - ret = tsk_trace_error(TSK_ERR_NO_MEMORY); - goto out; - } - - for (j = 0; j < num_transitions; j++) { - tsk_bug_assert(transitions[j].tree_node >= 0); - self->values[site][j] = transitions[j].value; - self->nodes[site][j] = transitions[j].tree_node; - } - -out: - return ret; -} - -static int -tsk_compressed_matrix_decode_site(tsk_compressed_matrix_t *self, const tsk_tree_t *tree, - const tsk_id_t site, double *values) -{ - int ret = 0; - const tsk_id_t *restrict list_left = tree->left_sample; - const tsk_id_t *restrict list_right = tree->right_sample; - const tsk_id_t *restrict list_next = tree->next_sample; - const tsk_id_t num_nodes = (tsk_id_t) tsk_treeseq_get_num_nodes(self->tree_sequence); - tsk_size_t j; - tsk_id_t node, index, stop; - double value; - - for (j = 0; j < self->num_transitions[site]; j++) { - node = self->nodes[site][j]; - if (node < 0 || node >= num_nodes) { - ret = tsk_trace_error(TSK_ERR_NODE_OUT_OF_BOUNDS); - goto out; - } - value = self->values[site][j]; - index = list_left[node]; - if (index == TSK_NULL) { - /* It's an error if there are nodes that don't subtend any samples */ - ret = tsk_trace_error(TSK_ERR_BAD_COMPRESSED_MATRIX_NODE); - goto out; - } - stop = list_right[node]; - while (true) { - values[index] = value; - if (index == stop) { - break; - } - index = list_next[index]; - } - } -out: - return ret; -} - -int -tsk_compressed_matrix_decode(tsk_compressed_matrix_t *self, double *values) -{ - int ret = 0; - int t_ret; - tsk_tree_t tree; - tsk_size_t j, num_tree_sites; - const tsk_site_t *sites = NULL; - tsk_id_t site_id; - double *site_array; - - ret = tsk_tree_init(&tree, self->tree_sequence, TSK_SAMPLE_LISTS); - if (ret != 0) { - goto out; - } - - for (t_ret = tsk_tree_first(&tree); t_ret == TSK_TREE_OK; - t_ret = tsk_tree_next(&tree)) { - ret = tsk_tree_get_sites(&tree, &sites, &num_tree_sites); - if (ret != 0) { - goto out; - } - for (j = 0; j < num_tree_sites; j++) { - site_id = sites[j].id; - site_array = values + ((tsk_size_t) site_id) * self->num_samples; - if (self->num_transitions[site_id] == 0) { - tsk_memset(site_array, 0, self->num_samples * sizeof(*site_array)); - } else { - ret = tsk_compressed_matrix_decode_site( - self, &tree, site_id, site_array); - if (ret != 0) { - goto out; - } - } - } - } - if (t_ret < 0) { - ret = t_ret; - goto out; - } -out: - tsk_tree_free(&tree); - return ret; -} - -/**************************************************************** - * Viterbi matrix - ****************************************************************/ - -static int -tsk_viterbi_matrix_expand_recomb_records(tsk_viterbi_matrix_t *self) -{ - int ret = 0; - tsk_recomb_required_record *tmp = tsk_realloc( - self->recombination_required, self->max_recomb_records * sizeof(*tmp)); - - if (tmp == NULL) { - ret = tsk_trace_error(TSK_ERR_NO_MEMORY); - goto out; - } - self->recombination_required = tmp; -out: - return ret; -} - -int -tsk_viterbi_matrix_init(tsk_viterbi_matrix_t *self, tsk_treeseq_t *tree_sequence, - tsk_size_t block_size, tsk_flags_t options) -{ - int ret = 0; - - tsk_memset(self, 0, sizeof(*self)); - if (block_size == 0) { - block_size = 1 << 20; /* 1MiB */ - } - ret = tsk_compressed_matrix_init(&self->matrix, tree_sequence, block_size, options); - if (ret != 0) { - goto out; - } - - self->max_recomb_records - = TSK_MAX(1, block_size / sizeof(tsk_recomb_required_record)); - ret = tsk_viterbi_matrix_expand_recomb_records(self); - if (ret != 0) { - goto out; - } - /* Add the sentinel at the start to simplify traceback */ - self->recombination_required[0].site = -1; - - ret = tsk_viterbi_matrix_clear(self); -out: - return ret; -} - -int -tsk_viterbi_matrix_free(tsk_viterbi_matrix_t *self) -{ - tsk_compressed_matrix_free(&self->matrix); - tsk_safe_free(self->recombination_required); - return 0; -} - -int -tsk_viterbi_matrix_clear(tsk_viterbi_matrix_t *self) -{ - self->num_recomb_records = 1; - tsk_compressed_matrix_clear(&self->matrix); - return 0; -} - -void -tsk_viterbi_matrix_print_state(tsk_viterbi_matrix_t *self, FILE *out) -{ - tsk_id_t l, j; - - fprintf(out, "viterbi_matrix\n"); - fprintf(out, "num_recomb_records = %lld\n", (long long) self->num_recomb_records); - fprintf(out, "max_recomb_records = %lld\n", (long long) self->max_recomb_records); - - j = 1; - for (l = 0; l < (tsk_id_t) self->matrix.num_sites; l++) { - fprintf(out, "%lld\t[", (long long) l); - while (j < (tsk_id_t) self->num_recomb_records - && self->recombination_required[j].site == l) { - fprintf(out, "(%lld, %d) ", (long long) self->recombination_required[j].node, - self->recombination_required[j].required); - j++; - } - fprintf(out, "]\n"); - } - tsk_compressed_matrix_print_state(&self->matrix, out); -} - -TSK_WARN_UNUSED int -tsk_viterbi_matrix_add_recombination_required( - tsk_viterbi_matrix_t *self, tsk_id_t site, tsk_id_t node, bool required) -{ - int ret = 0; - tsk_recomb_required_record *record; - - if (self->num_recomb_records == self->max_recomb_records) { - self->max_recomb_records *= 2; - ret = tsk_viterbi_matrix_expand_recomb_records(self); - if (ret != 0) { - goto out; - } - } - record = self->recombination_required + self->num_recomb_records; - record->site = site; - record->node = node; - record->required = required; - self->num_recomb_records++; -out: - return ret; -} - -static tsk_id_t -tsk_viterbi_matrix_choose_sample( - tsk_viterbi_matrix_t *self, tsk_id_t site, tsk_tree_t *tree) -{ - tsk_id_t ret; - tsk_id_t u = TSK_NULL; - const tsk_flags_t *node_flags = self->matrix.tree_sequence->tables->nodes.flags; - const tsk_size_t num_transitions = self->matrix.num_transitions[site]; - const tsk_id_t *transition_nodes = self->matrix.nodes[site]; - const double *transition_values = self->matrix.values[site]; - double max_value = -1; - tsk_size_t j; - tsk_id_t v; - bool found; - - if (num_transitions == 0) { - ret = tsk_trace_error(TSK_ERR_NULL_VITERBI_MATRIX); - goto out; - } - for (j = 0; j < num_transitions; j++) { - if (max_value < transition_values[j]) { - u = transition_nodes[j]; - max_value = transition_values[j]; - } - } - tsk_bug_assert(u != TSK_NULL); - - while (!(node_flags[u] & TSK_NODE_IS_SAMPLE)) { - found = false; - for (v = tree->left_child[u]; v != TSK_NULL; v = tree->right_sib[v]) { - /* Choose the first child that is not in the list of transition nodes */ - for (j = 0; j < num_transitions; j++) { - if (transition_nodes[j] == v) { - break; - } - } - if (j == num_transitions) { - u = v; - found = true; - break; - } - } - /* TODO: should remove this once we're sure this is robust */ - tsk_bug_assert(found); - } - ret = u; -out: - return ret; -} - -int -tsk_viterbi_matrix_traceback( - tsk_viterbi_matrix_t *self, tsk_id_t *path, tsk_flags_t TSK_UNUSED(options)) -{ - int ret = 0; - tsk_site_t site; - tsk_id_t u, site_id, current_node; - tsk_recomb_required_record *rr_record, *rr_record_tmp; - const tsk_id_t num_sites = (tsk_id_t) self->matrix.num_sites; - const tsk_id_t num_nodes - = (tsk_id_t) tsk_treeseq_get_num_nodes(self->matrix.tree_sequence); - tsk_tree_t tree; - tsk_id_t *recombination_tree - = tsk_malloc((size_t) num_nodes * sizeof(*recombination_tree)); - - ret = tsk_tree_init(&tree, self->matrix.tree_sequence, 0); - if (ret != 0) { - goto out; - } - if (recombination_tree == NULL) { - ret = tsk_trace_error(TSK_ERR_NO_MEMORY); - goto out; - } - /* Initialise the path an recombination_tree to contain TSK_NULL */ - tsk_memset(path, 0xff, ((size_t) num_sites) * sizeof(*path)); - tsk_memset(recombination_tree, 0xff, ((size_t) num_nodes) * sizeof(*path)); - - current_node = TSK_NULL; - rr_record = &self->recombination_required[self->num_recomb_records - 1]; - ret = tsk_tree_last(&tree); - if (ret < 0) { - goto out; - } - - for (site_id = num_sites - 1; site_id >= 0; site_id--) { - ret = tsk_treeseq_get_site(self->matrix.tree_sequence, site_id, &site); - if (ret != 0) { - goto out; - } - while (tree.interval.left > site.position) { - ret = tsk_tree_prev(&tree); - if (ret < 0) { - goto out; - } - } - tsk_bug_assert(tree.interval.left <= site.position); - tsk_bug_assert(site.position < tree.interval.right); - - /* Fill in the recombination tree */ - rr_record_tmp = rr_record; - while (rr_record->site == site.id) { - recombination_tree[rr_record->node] = rr_record->required; - rr_record--; - } - if (current_node == TSK_NULL) { - current_node = tsk_viterbi_matrix_choose_sample(self, site.id, &tree); - if (current_node < 0) { - ret = (int) current_node; - goto out; - } - } - path[site.id] = current_node; - /* Now traverse up the tree from the current node. The - * first marked node tells us whether we need to recombine */ - u = current_node; - while (u != TSK_NULL && recombination_tree[u] == TSK_NULL) { - u = tree.parent[u]; - } - tsk_bug_assert(u != TSK_NULL); - if (recombination_tree[u] == 1) { - /* Switch at the next site */ - current_node = TSK_NULL; - } - - /* Reset in the recombination tree */ - rr_record = rr_record_tmp; - while (rr_record->site == site.id) { - recombination_tree[rr_record->node] = TSK_NULL; - rr_record--; - } - } - ret = 0; -out: - tsk_tree_free(&tree); - tsk_safe_free(recombination_tree); - return ret; -} diff --git a/c/tskit/haplotype_matching.h b/c/tskit/haplotype_matching.h deleted file mode 100644 index 151eb321b8..0000000000 --- a/c/tskit/haplotype_matching.h +++ /dev/null @@ -1,169 +0,0 @@ -/* - * MIT License - * - * Copyright (c) 2019-2024 Tskit Developers - * - * Permission is hereby granted, free of charge, to any person obtaining a copy - * of this software and associated documentation files (the "Software"), to deal - * in the Software without restriction, including without limitation the rights - * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell - * copies of the Software, and to permit persons to whom the Software is - * furnished to do so, subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in all - * copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, - * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE - * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER - * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, - * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE - * SOFTWARE. - */ - -#ifndef TSK_HAPLOTYPE_MATCHING_H -#define TSK_HAPLOTYPE_MATCHING_H - -#ifdef __cplusplus -extern "C" { -#endif - -#include - -/* Seems like we might use this somewhere else as well, so putting it into the middle - * of the flags space */ -#define TSK_ALLELES_ACGT (1 << 16) - -typedef struct { - tsk_id_t tree_node; - tsk_id_t value_index; - double value; -} tsk_value_transition_t; - -typedef struct { - tsk_size_t index; - double value; -} tsk_argsort_t; - -typedef struct { - tsk_id_t tree_node; - tsk_id_t old_state; - tsk_id_t new_state; - tsk_id_t transition_parent; -} tsk_transition_stack_t; - -typedef struct { - double normalisation_factor; - double *value; - tsk_id_t *node; - tsk_size_t num_values; -} tsk_site_probability_t; - -typedef struct { - tsk_treeseq_t *tree_sequence; - tsk_flags_t options; - tsk_size_t num_sites; - tsk_size_t num_samples; - double *normalisation_factor; - tsk_size_t *num_transitions; - double **values; - tsk_id_t **nodes; - tsk_blkalloc_t memory; -} tsk_compressed_matrix_t; - -typedef struct { - tsk_id_t site; - tsk_id_t node; - bool required; -} tsk_recomb_required_record; - -typedef struct { - tsk_compressed_matrix_t matrix; - tsk_recomb_required_record *recombination_required; - tsk_size_t num_recomb_records; - tsk_size_t max_recomb_records; -} tsk_viterbi_matrix_t; - -typedef struct _tsk_ls_hmm_t { - /* input */ - tsk_treeseq_t *tree_sequence; - double *recombination_rate; - double *mutation_rate; - const char ***alleles; - unsigned int precision; - uint32_t *num_alleles; - tsk_size_t num_samples; - tsk_size_t num_sites; - tsk_size_t num_nodes; - /* state */ - tsk_tree_t tree; - tsk_id_t *parent; - /* The probability value transitions on the tree */ - tsk_value_transition_t *transitions; - tsk_value_transition_t *transitions_copy; - /* Stack used when distributing transitions on the tree */ - tsk_transition_stack_t *transition_stack; - /* Map of node_id to index in the transitions list */ - tsk_id_t *transition_index; - /* Buffer used to argsort the transitions by node time */ - tsk_argsort_t *transition_time_order; - tsk_size_t num_transitions; - tsk_size_t max_transitions; - /* The distinct values in the transitions */ - double *values; - tsk_size_t num_values; - tsk_size_t max_values; - tsk_size_t max_parsimony_words; - /* Number of machine words per node optimal value set. */ - tsk_size_t num_optimal_value_set_words; - uint64_t *optimal_value_sets; - /* The parent transition; used during compression */ - tsk_id_t *transition_parent; - /* The number of samples directly subtended by a transition */ - tsk_size_t *num_transition_samples; - int32_t *allelic_state; - /* Algorithms set these values before they are run */ - int (*next_probability)( - struct _tsk_ls_hmm_t *, tsk_id_t, double, bool, tsk_id_t, double *); - double (*compute_normalisation_factor)(struct _tsk_ls_hmm_t *); - void *output; -} tsk_ls_hmm_t; - -/* TODO constify these APIs */ -int tsk_ls_hmm_init(tsk_ls_hmm_t *self, tsk_treeseq_t *tree_sequence, - double *recombination_rate, double *mutation_rate, tsk_flags_t options); -int tsk_ls_hmm_set_precision(tsk_ls_hmm_t *self, unsigned int precision); -int tsk_ls_hmm_free(tsk_ls_hmm_t *self); -void tsk_ls_hmm_print_state(tsk_ls_hmm_t *self, FILE *out); -int tsk_ls_hmm_forward(tsk_ls_hmm_t *self, int32_t *haplotype, - tsk_compressed_matrix_t *output, tsk_flags_t options); -int tsk_ls_hmm_backward(tsk_ls_hmm_t *self, int32_t *haplotype, - const double *forward_norm, tsk_compressed_matrix_t *output, tsk_flags_t options); -int tsk_ls_hmm_viterbi(tsk_ls_hmm_t *self, int32_t *haplotype, - tsk_viterbi_matrix_t *output, tsk_flags_t options); - -int tsk_compressed_matrix_init(tsk_compressed_matrix_t *self, - tsk_treeseq_t *tree_sequence, tsk_size_t block_size, tsk_flags_t options); -int tsk_compressed_matrix_free(tsk_compressed_matrix_t *self); -int tsk_compressed_matrix_clear(tsk_compressed_matrix_t *self); -void tsk_compressed_matrix_print_state(tsk_compressed_matrix_t *self, FILE *out); -int tsk_compressed_matrix_store_site(tsk_compressed_matrix_t *self, tsk_id_t site, - double normalisation_factor, tsk_size_t num_transitions, - const tsk_value_transition_t *transitions); -int tsk_compressed_matrix_decode(tsk_compressed_matrix_t *self, double *values); - -int tsk_viterbi_matrix_init(tsk_viterbi_matrix_t *self, tsk_treeseq_t *tree_sequence, - tsk_size_t block_size, tsk_flags_t options); -int tsk_viterbi_matrix_free(tsk_viterbi_matrix_t *self); -int tsk_viterbi_matrix_clear(tsk_viterbi_matrix_t *self); -void tsk_viterbi_matrix_print_state(tsk_viterbi_matrix_t *self, FILE *out); -int tsk_viterbi_matrix_add_recombination_required( - tsk_viterbi_matrix_t *self, tsk_id_t site, tsk_id_t node, bool required); -int tsk_viterbi_matrix_traceback( - tsk_viterbi_matrix_t *self, tsk_id_t *path, tsk_flags_t options); - -#ifdef __cplusplus -} -#endif -#endif diff --git a/python/CHANGELOG.rst b/python/CHANGELOG.rst index fbce38d205..1be12c4ed3 100644 --- a/python/CHANGELOG.rst +++ b/python/CHANGELOG.rst @@ -20,7 +20,7 @@ **Bugfixes** -- Update to kastore C version 2.1.3 to include file security fix (:user:`jeromekelleher`, +- Update to kastore C version 2.1.3 to include file security fix (:user:`jeromekelleher`, :pr:`3478`). - ``TableCollection.delete_older`` now checks table integrity first, raising an error for invalid node or mutation references instead of reading out of bounds diff --git a/python/_tskitmodule.c b/python/_tskitmodule.c index 01bcf9fb7a..bc4d3ab00a 100644 --- a/python/_tskitmodule.c +++ b/python/_tskitmodule.c @@ -166,24 +166,6 @@ typedef struct { tsk_ld_calc_t *ld_calc; } LdCalculator; -typedef struct { - PyObject_HEAD - TreeSequence *tree_sequence; - tsk_ls_hmm_t *ls_hmm; -} LsHmm; - -typedef struct { - PyObject_HEAD - TreeSequence *tree_sequence; - tsk_compressed_matrix_t *compressed_matrix; -} CompressedMatrix; - -typedef struct { - PyObject_HEAD - TreeSequence *tree_sequence; - tsk_viterbi_matrix_t *viterbi_matrix; -} ViterbiMatrix; - typedef struct { PyObject_HEAD PyObject *owner; @@ -683,67 +665,6 @@ convert_transitions(tsk_state_transition_t *transitions, tsk_size_t num_transiti return ret; } -/* TODO: this should really be a dict we're returning */ -static PyObject * -convert_compressed_matrix_site(tsk_compressed_matrix_t *matrix, unsigned int site) -{ - PyObject *ret = NULL; - PyObject *list = NULL; - PyObject *item = NULL; - tsk_size_t j, num_values; - - if (site >= matrix->num_sites) { - PyErr_SetString(PyExc_ValueError, "Site index out of bounds"); - goto out; - } - - num_values = matrix->num_transitions[site]; - list = PyList_New(num_values); - if (list == NULL) { - goto out; - } - for (j = 0; j < num_values; j++) { - item = Py_BuildValue("id", matrix->nodes[site][j], matrix->values[site][j]); - if (item == NULL) { - goto out; - } - PyList_SET_ITEM(list, j, item); - item = NULL; - } - ret = list; - list = NULL; -out: - Py_XDECREF(item); - Py_XDECREF(list); - return ret; -} - -static PyObject * -decode_compressed_matrix(tsk_compressed_matrix_t *matrix) -{ - int err; - PyObject *ret = NULL; - PyArrayObject *decoded = NULL; - npy_intp dims[2]; - - dims[0] = tsk_treeseq_get_num_sites(matrix->tree_sequence); - dims[1] = tsk_treeseq_get_num_samples(matrix->tree_sequence); - decoded = (PyArrayObject *) PyArray_SimpleNew(2, dims, NPY_FLOAT64); - if (decoded == NULL) { - goto out; - } - err = tsk_compressed_matrix_decode(matrix, PyArray_DATA(decoded)); - if (err != 0) { - handle_library_error(err); - goto out; - } - ret = (PyObject *) decoded; - decoded = NULL; -out: - Py_XDECREF(decoded); - return ret; -} - static const char ** parse_allele_list(PyObject *allele_tuple) { @@ -11261,722 +11182,6 @@ static PyTypeObject LdCalculatorType = { // clang-format on }; -/*=================================================================== - * CompressedMatrix - *=================================================================== - */ - -static int -CompressedMatrix_check_state(CompressedMatrix *self) -{ - int ret = -1; - if (self->compressed_matrix == NULL) { - PyErr_SetString(PyExc_SystemError, "CompressedMatrix not initialised"); - goto out; - } - ret = 0; -out: - return ret; -} - -static void -CompressedMatrix_dealloc(CompressedMatrix *self) -{ - if (self->compressed_matrix != NULL) { - tsk_compressed_matrix_free(self->compressed_matrix); - PyMem_Free(self->compressed_matrix); - self->compressed_matrix = NULL; - } - Py_XDECREF(self->tree_sequence); - Py_TYPE(self)->tp_free((PyObject *) self); -} - -static int -CompressedMatrix_init(CompressedMatrix *self, PyObject *args, PyObject *kwds) -{ - int ret = -1; - int err; - static char *kwlist[] = { "tree_sequence", "block_size", NULL }; - TreeSequence *tree_sequence = NULL; - Py_ssize_t block_size = 0; - - self->compressed_matrix = NULL; - self->tree_sequence = NULL; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "O!|n", kwlist, &TreeSequenceType, - &tree_sequence, &block_size)) { - goto out; - } - self->tree_sequence = tree_sequence; - Py_INCREF(self->tree_sequence); - if (TreeSequence_check_state(self->tree_sequence) != 0) { - goto out; - } - self->compressed_matrix = PyMem_Malloc(sizeof(tsk_compressed_matrix_t)); - if (self->compressed_matrix == NULL) { - PyErr_NoMemory(); - goto out; - } - memset(self->compressed_matrix, 0, sizeof(tsk_compressed_matrix_t)); - - err = tsk_compressed_matrix_init( - self->compressed_matrix, self->tree_sequence->tree_sequence, block_size, 0); - if (err != 0) { - handle_library_error(err); - goto out; - } - ret = 0; -out: - return ret; -} - -static PyObject * -CompressedMatrix_get_num_sites(CompressedMatrix *self, void *closure) -{ - PyObject *ret = NULL; - - if (CompressedMatrix_check_state(self) != 0) { - goto out; - } - ret = Py_BuildValue("n", (Py_ssize_t) self->compressed_matrix->num_sites); -out: - return ret; -} - -static PyObject * -CompressedMatrix_get_normalisation_factor(CompressedMatrix *self, void *closure) -{ - PyObject *ret = NULL; - PyArrayObject *array; - tsk_size_t num_sites; - npy_intp dims; - - if (CompressedMatrix_check_state(self) != 0) { - goto out; - } - num_sites = self->compressed_matrix->num_sites; - dims = (npy_intp) num_sites; - array = (PyArrayObject *) PyArray_EMPTY(1, &dims, NPY_FLOAT64, 0); - if (array == NULL) { - goto out; - } - memcpy(PyArray_DATA(array), self->compressed_matrix->normalisation_factor, - num_sites * sizeof(*self->compressed_matrix->normalisation_factor)); - ret = (PyObject *) array; -out: - return ret; -} - -static PyObject * -CompressedMatrix_get_num_transitions(CompressedMatrix *self, void *closure) -{ - PyObject *ret = NULL; - PyArrayObject *array; - tsk_size_t num_sites; - npy_intp dims; - - if (CompressedMatrix_check_state(self) != 0) { - goto out; - } - num_sites = self->compressed_matrix->num_sites; - dims = (npy_intp) num_sites; - array = (PyArrayObject *) PyArray_EMPTY(1, &dims, NPY_UINT64, 0); - if (array == NULL) { - goto out; - } - memcpy(PyArray_DATA(array), self->compressed_matrix->num_transitions, - num_sites * sizeof(*self->compressed_matrix->num_transitions)); - ret = (PyObject *) array; -out: - return ret; -} - -static PyObject * -CompressedMatrix_get_site(CompressedMatrix *self, PyObject *args) -{ - PyObject *ret = NULL; - unsigned int site; - - if (CompressedMatrix_check_state(self) != 0) { - goto out; - } - if (!PyArg_ParseTuple(args, "I", &site)) { - goto out; - } - ret = convert_compressed_matrix_site(self->compressed_matrix, site); -out: - return ret; -} - -static PyObject * -CompressedMatrix_decode(CompressedMatrix *self) -{ - PyObject *ret = NULL; - if (CompressedMatrix_check_state(self) != 0) { - goto out; - } - ret = decode_compressed_matrix(self->compressed_matrix); -out: - return ret; -} - -static PyGetSetDef CompressedMatrix_getsetters[] = { - { .name = "num_sites", - .get = (getter) CompressedMatrix_get_num_sites, - .doc = "The number of sites." }, - { .name = "normalisation_factor", - .get = (getter) CompressedMatrix_get_normalisation_factor, - .doc = "The per-site normalisation factor." }, - { .name = "num_transitions", - .get = (getter) CompressedMatrix_get_num_transitions, - .doc = "The per-site number of transitions in the compressed matrix." }, - { NULL } /* Sentinel */ -}; - -static PyMethodDef CompressedMatrix_methods[] = { - { .ml_name = "get_site", - .ml_meth = (PyCFunction) CompressedMatrix_get_site, - .ml_flags = METH_VARARGS, - .ml_doc = "Returns the list of (node, value) tuples for the specified site." }, - { .ml_name = "decode", - .ml_meth = (PyCFunction) CompressedMatrix_decode, - .ml_flags = METH_NOARGS, - .ml_doc = "Returns the full decoded forward matrix." }, - { NULL } /* Sentinel */ -}; - -static PyTypeObject CompressedMatrixType = { - // clang-format off - PyVarObject_HEAD_INIT(NULL, 0) - .tp_name = "_tskit.CompressedMatrix", - .tp_basicsize = sizeof(CompressedMatrix), - .tp_dealloc = (destructor) CompressedMatrix_dealloc, - .tp_flags = Py_TPFLAGS_DEFAULT, - .tp_doc = "CompressedMatrix objects", - .tp_methods = CompressedMatrix_methods, - .tp_getset = CompressedMatrix_getsetters, - .tp_init = (initproc) CompressedMatrix_init, - .tp_new = PyType_GenericNew, - // clang-format on -}; - -/*=================================================================== - * ViterbiMatrix - *=================================================================== - */ - -static int -ViterbiMatrix_check_state(ViterbiMatrix *self) -{ - int ret = -1; - if (self->viterbi_matrix == NULL) { - PyErr_SetString(PyExc_SystemError, "ViterbiMatrix not initialised"); - goto out; - } - ret = 0; -out: - return ret; -} - -static void -ViterbiMatrix_dealloc(ViterbiMatrix *self) -{ - if (self->viterbi_matrix != NULL) { - tsk_viterbi_matrix_free(self->viterbi_matrix); - PyMem_Free(self->viterbi_matrix); - self->viterbi_matrix = NULL; - } - Py_XDECREF(self->tree_sequence); - Py_TYPE(self)->tp_free((PyObject *) self); -} - -static int -ViterbiMatrix_init(ViterbiMatrix *self, PyObject *args, PyObject *kwds) -{ - int ret = -1; - int err; - static char *kwlist[] = { "tree_sequence", "num_records", NULL }; - TreeSequence *tree_sequence = NULL; - Py_ssize_t num_records = 0; - - self->viterbi_matrix = NULL; - self->tree_sequence = NULL; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "O!|n", kwlist, &TreeSequenceType, - &tree_sequence, &num_records)) { - goto out; - } - self->tree_sequence = tree_sequence; - Py_INCREF(self->tree_sequence); - if (TreeSequence_check_state(self->tree_sequence) != 0) { - goto out; - } - self->viterbi_matrix = PyMem_Malloc(sizeof(tsk_viterbi_matrix_t)); - if (self->viterbi_matrix == NULL) { - PyErr_NoMemory(); - goto out; - } - memset(self->viterbi_matrix, 0, sizeof(tsk_viterbi_matrix_t)); - - err = tsk_viterbi_matrix_init( - self->viterbi_matrix, self->tree_sequence->tree_sequence, num_records, 0); - if (err != 0) { - handle_library_error(err); - goto out; - } - ret = 0; -out: - return ret; -} - -static PyObject * -ViterbiMatrix_traceback(ViterbiMatrix *self) -{ - PyObject *ret = NULL; - PyArrayObject *path = NULL; - npy_intp dims; - int err; - - if (ViterbiMatrix_check_state(self) != 0) { - goto out; - } - dims = self->viterbi_matrix->matrix.num_sites; - path = (PyArrayObject *) PyArray_SimpleNew(1, &dims, NPY_INT32); - if (path == NULL) { - goto out; - } - - err = tsk_viterbi_matrix_traceback(self->viterbi_matrix, PyArray_DATA(path), 0); - if (err != 0) { - handle_library_error(err); - goto out; - } - ret = (PyObject *) path; - path = NULL; -out: - Py_XDECREF(path); - return ret; -} - -static PyObject * -ViterbiMatrix_get_num_sites(ViterbiMatrix *self, void *closure) -{ - PyObject *ret = NULL; - - if (ViterbiMatrix_check_state(self) != 0) { - goto out; - } - ret = Py_BuildValue("n", (Py_ssize_t) self->viterbi_matrix->matrix.num_sites); -out: - return ret; -} - -/* NOTE: We're doing something pretty ugly here in that we're duplicating the - * methods from the CompressedMatrix class to provide access to the - * viterbi_matrix struct's embedded compressed_matrix. It would be more - * elegant if the ViterbiMatrix class had a CompressedMatrix member, - * but the memory management is tricky, so it doesn't seem worth the - * hassle. - */ - -static PyObject * -ViterbiMatrix_get_normalisation_factor(ViterbiMatrix *self, void *closure) -{ - PyObject *ret = NULL; - PyArrayObject *array; - tsk_size_t num_sites; - npy_intp dims; - - if (ViterbiMatrix_check_state(self) != 0) { - goto out; - } - num_sites = self->viterbi_matrix->matrix.num_sites; - dims = (npy_intp) num_sites; - array = (PyArrayObject *) PyArray_EMPTY(1, &dims, NPY_FLOAT64, 0); - if (array == NULL) { - goto out; - } - memcpy(PyArray_DATA(array), self->viterbi_matrix->matrix.normalisation_factor, - num_sites * sizeof(*self->viterbi_matrix->matrix.normalisation_factor)); - ret = (PyObject *) array; -out: - return ret; -} - -static PyObject * -ViterbiMatrix_get_num_transitions(ViterbiMatrix *self, void *closure) -{ - PyObject *ret = NULL; - PyArrayObject *array; - tsk_size_t num_sites; - npy_intp dims; - - if (ViterbiMatrix_check_state(self) != 0) { - goto out; - } - num_sites = self->viterbi_matrix->matrix.num_sites; - dims = (npy_intp) num_sites; - - array = (PyArrayObject *) PyArray_EMPTY(1, &dims, NPY_UINT64, 0); - if (array == NULL) { - goto out; - } - memcpy(PyArray_DATA(array), self->viterbi_matrix->matrix.num_transitions, - num_sites * sizeof(*self->viterbi_matrix->matrix.num_transitions)); - ret = (PyObject *) array; -out: - return ret; -} - -static PyObject * -ViterbiMatrix_get_site(ViterbiMatrix *self, PyObject *args) -{ - PyObject *ret = NULL; - unsigned int site; - - if (ViterbiMatrix_check_state(self) != 0) { - goto out; - } - if (!PyArg_ParseTuple(args, "I", &site)) { - goto out; - } - ret = convert_compressed_matrix_site(&self->viterbi_matrix->matrix, site); -out: - return ret; -} - -static PyObject * -ViterbiMatrix_decode(ViterbiMatrix *self) -{ - PyObject *ret = NULL; - if (ViterbiMatrix_check_state(self) != 0) { - goto out; - } - ret = decode_compressed_matrix(&self->viterbi_matrix->matrix); -out: - return ret; -} - -static PyGetSetDef ViterbiMatrix_getsetters[] = { - { .name = "num_sites", - .get = (getter) ViterbiMatrix_get_num_sites, - .doc = "The number of sites." }, - { .name = "normalisation_factor", - .get = (getter) ViterbiMatrix_get_normalisation_factor, - .doc = "The per-site normalisation factor." }, - { .name = "num_transitions", - .get = (getter) ViterbiMatrix_get_num_transitions, - .doc = "The per-site number of transitions in the compressed matrix." }, - { NULL } /* Sentinel */ -}; - -static PyMethodDef ViterbiMatrix_methods[] = { - { .ml_name = "traceback", - .ml_meth = (PyCFunction) ViterbiMatrix_traceback, - .ml_flags = METH_NOARGS, - .ml_doc = "Returns a path for a given haplotype." }, - { .ml_name = "get_site", - .ml_meth = (PyCFunction) ViterbiMatrix_get_site, - .ml_flags = METH_VARARGS, - .ml_doc = "Returns the list of (node, value) tuples for the specified site." }, - { .ml_name = "decode", - .ml_meth = (PyCFunction) ViterbiMatrix_decode, - .ml_flags = METH_NOARGS, - .ml_doc = "Returns the full decoded forward matrix." }, - { NULL } /* Sentinel */ -}; - -static PyTypeObject ViterbiMatrixType = { - // clang-format off - PyVarObject_HEAD_INIT(NULL, 0) - .tp_name = "_tskit.ViterbiMatrix", - .tp_basicsize = sizeof(ViterbiMatrix), - .tp_dealloc = (destructor) ViterbiMatrix_dealloc, - .tp_flags = Py_TPFLAGS_DEFAULT, - .tp_doc = "ViterbiMatrix objects", - .tp_methods = ViterbiMatrix_methods, - .tp_getset = ViterbiMatrix_getsetters, - .tp_init = (initproc) ViterbiMatrix_init, - .tp_new = PyType_GenericNew, - // clang-format on -}; - -/*=================================================================== - * LsHmm - *=================================================================== - */ - -static int -LsHmm_check_state(LsHmm *self) -{ - int ret = -1; - if (self->ls_hmm == NULL) { - PyErr_SetString(PyExc_SystemError, "LsHmm not initialised"); - goto out; - } - ret = 0; -out: - return ret; -} - -static void -LsHmm_dealloc(LsHmm *self) -{ - if (self->ls_hmm != NULL) { - tsk_ls_hmm_free(self->ls_hmm); - PyMem_Free(self->ls_hmm); - self->ls_hmm = NULL; - } - Py_XDECREF(self->tree_sequence); - Py_TYPE(self)->tp_free((PyObject *) self); -} - -static int -LsHmm_init(LsHmm *self, PyObject *args, PyObject *kwds) -{ - int ret = -1; - int err; - static char *kwlist[] = { "tree_sequence", "recombination_rate", "mutation_rate", - "precision", "acgt_alleles", NULL }; - PyObject *recombination_rate = NULL; - PyArrayObject *recombination_rate_array = NULL; - PyObject *mutation_rate = NULL; - PyArrayObject *mutation_rate_array = NULL; - TreeSequence *tree_sequence = NULL; - unsigned int precision = 23; - int acgt_alleles = 0; - tsk_flags_t options = 0; - npy_intp *shape, num_sites; - - self->ls_hmm = NULL; - self->tree_sequence = NULL; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "O!OO|Ii", kwlist, &TreeSequenceType, - &tree_sequence, &recombination_rate, &mutation_rate, &precision, - &acgt_alleles)) { - goto out; - } - self->tree_sequence = tree_sequence; - Py_INCREF(self->tree_sequence); - if (TreeSequence_check_state(self->tree_sequence) != 0) { - goto out; - } - self->ls_hmm = PyMem_Malloc(sizeof(tsk_ls_hmm_t)); - if (self->ls_hmm == NULL) { - PyErr_NoMemory(); - goto out; - } - memset(self->ls_hmm, 0, sizeof(tsk_ls_hmm_t)); - - num_sites = (npy_intp) tsk_treeseq_get_num_sites(self->tree_sequence->tree_sequence); - recombination_rate_array = (PyArrayObject *) PyArray_FROMANY( - recombination_rate, NPY_FLOAT64, 1, 1, NPY_ARRAY_IN_ARRAY); - if (recombination_rate_array == NULL) { - goto out; - } - shape = PyArray_DIMS(recombination_rate_array); - if (shape[0] != num_sites) { - PyErr_SetString(PyExc_ValueError, - "recombination_rate array must have dimension (num_sites,)"); - goto out; - } - mutation_rate_array = (PyArrayObject *) PyArray_FROMANY( - mutation_rate, NPY_FLOAT64, 1, 1, NPY_ARRAY_IN_ARRAY); - if (mutation_rate_array == NULL) { - goto out; - } - shape = PyArray_DIMS(mutation_rate_array); - if (shape[0] != num_sites) { - PyErr_SetString( - PyExc_ValueError, "mutation_rate array must have dimension (num_sites,)"); - goto out; - } - if (acgt_alleles) { - options |= TSK_ALLELES_ACGT; - } - - err = tsk_ls_hmm_init(self->ls_hmm, self->tree_sequence->tree_sequence, - PyArray_DATA(recombination_rate_array), PyArray_DATA(mutation_rate_array), - options); - if (err != 0) { - handle_library_error(err); - goto out; - } - tsk_ls_hmm_set_precision(self->ls_hmm, precision); - ret = 0; -out: - Py_XDECREF(recombination_rate_array); - Py_XDECREF(mutation_rate_array); - return ret; -} - -static PyObject * -LsHmm_forward_matrix(LsHmm *self, PyObject *args) -{ - int err; - PyObject *ret = NULL; - PyObject *haplotype = NULL; - CompressedMatrix *compressed_matrix = NULL; - PyArrayObject *haplotype_array = NULL; - npy_intp *shape, num_sites; - - if (LsHmm_check_state(self) != 0) { - goto out; - } - if (!PyArg_ParseTuple( - args, "OO!", &haplotype, &CompressedMatrixType, &compressed_matrix)) { - goto out; - } - num_sites = (npy_intp) tsk_treeseq_get_num_sites(self->tree_sequence->tree_sequence); - haplotype_array = (PyArrayObject *) PyArray_FROMANY( - haplotype, NPY_INT32, 1, 1, NPY_ARRAY_IN_ARRAY); - if (haplotype_array == NULL) { - goto out; - } - shape = PyArray_DIMS(haplotype_array); - if (shape[0] != num_sites) { - PyErr_SetString( - PyExc_ValueError, "haplotype array must have dimension (num_sites,)"); - goto out; - } - err = tsk_ls_hmm_forward(self->ls_hmm, PyArray_DATA(haplotype_array), - compressed_matrix->compressed_matrix, TSK_NO_INIT); - if (err != 0) { - handle_library_error(err); - goto out; - } - ret = Py_BuildValue(""); -out: - Py_XDECREF(haplotype_array); - return ret; -} - -static PyObject * -LsHmm_backward_matrix(LsHmm *self, PyObject *args) -{ - int err; - PyObject *ret = NULL; - PyObject *haplotype = NULL; - PyObject *forward_norm = NULL; - CompressedMatrix *compressed_matrix = NULL; - PyArrayObject *haplotype_array = NULL; - PyArrayObject *forward_norm_array = NULL; - npy_intp *shape, num_sites; - - if (LsHmm_check_state(self) != 0) { - goto out; - } - if (!PyArg_ParseTuple(args, "OOO!", &haplotype, &forward_norm, &CompressedMatrixType, - &compressed_matrix)) { - goto out; - } - num_sites = (npy_intp) tsk_treeseq_get_num_sites(self->tree_sequence->tree_sequence); - - haplotype_array = (PyArrayObject *) PyArray_FROMANY( - haplotype, NPY_INT32, 1, 1, NPY_ARRAY_IN_ARRAY); - if (haplotype_array == NULL) { - goto out; - } - shape = PyArray_DIMS(haplotype_array); - if (shape[0] != num_sites) { - PyErr_SetString( - PyExc_ValueError, "haplotype array must have dimension (num_sites,)"); - goto out; - } - - forward_norm_array = (PyArrayObject *) PyArray_FROMANY( - forward_norm, NPY_FLOAT64, 1, 1, NPY_ARRAY_IN_ARRAY); - if (forward_norm_array == NULL) { - goto out; - } - shape = PyArray_DIMS(forward_norm_array); - if (shape[0] != num_sites) { - PyErr_SetString( - PyExc_ValueError, "forward_norm array must have dimension (num_sites,)"); - goto out; - } - err = tsk_ls_hmm_backward(self->ls_hmm, PyArray_DATA(haplotype_array), - PyArray_DATA(forward_norm_array), compressed_matrix->compressed_matrix, - TSK_NO_INIT); - if (err != 0) { - handle_library_error(err); - goto out; - } - ret = Py_BuildValue(""); -out: - Py_XDECREF(haplotype_array); - Py_XDECREF(forward_norm_array); - return ret; -} - -static PyObject * -LsHmm_viterbi_matrix(LsHmm *self, PyObject *args) -{ - int err; - PyObject *ret = NULL; - PyObject *haplotype = NULL; - ViterbiMatrix *viterbi_matrix = NULL; - PyArrayObject *haplotype_array = NULL; - npy_intp *shape, num_sites; - - if (LsHmm_check_state(self) != 0) { - goto out; - } - if (!PyArg_ParseTuple( - args, "OO!", &haplotype, &ViterbiMatrixType, &viterbi_matrix)) { - goto out; - } - num_sites = (npy_intp) tsk_treeseq_get_num_sites(self->tree_sequence->tree_sequence); - haplotype_array = (PyArrayObject *) PyArray_FROMANY( - haplotype, NPY_INT32, 1, 1, NPY_ARRAY_IN_ARRAY); - if (haplotype_array == NULL) { - goto out; - } - shape = PyArray_DIMS(haplotype_array); - if (shape[0] != num_sites) { - PyErr_SetString( - PyExc_ValueError, "haplotype array must have dimension (num_sites,)"); - goto out; - } - err = tsk_ls_hmm_viterbi(self->ls_hmm, PyArray_DATA(haplotype_array), - viterbi_matrix->viterbi_matrix, TSK_NO_INIT); - if (err != 0) { - handle_library_error(err); - goto out; - } - ret = Py_BuildValue(""); -out: - Py_XDECREF(haplotype_array); - return ret; -} - -static PyMethodDef LsHmm_methods[] = { - { .ml_name = "forward_matrix", - .ml_meth = (PyCFunction) LsHmm_forward_matrix, - .ml_flags = METH_VARARGS, - .ml_doc = "Returns the tree encoded forward matrix for a given haplotype" }, - { .ml_name = "backward_matrix", - .ml_meth = (PyCFunction) LsHmm_backward_matrix, - .ml_flags = METH_VARARGS, - .ml_doc = "Returns the tree encoded backward matrix for a given haplotype" }, - { .ml_name = "viterbi_matrix", - .ml_meth = (PyCFunction) LsHmm_viterbi_matrix, - .ml_flags = METH_VARARGS, - .ml_doc = "Returns the tree encoded Viterbi matrix for a given haplotype" }, - { NULL } /* Sentinel */ -}; - -static PyTypeObject LsHmmType = { - // clang-format off - PyVarObject_HEAD_INIT(NULL, 0) - .tp_name = "_tskit.LsHmm", - .tp_basicsize = sizeof(LsHmm), - .tp_dealloc = (destructor) LsHmm_dealloc, - .tp_flags = Py_TPFLAGS_DEFAULT, - .tp_doc = "LsHmm objects", - .tp_methods = LsHmm_methods, - .tp_init = (initproc) LsHmm_init, - .tp_new = PyType_GenericNew, - // clang-format on -}; - /*=================================================================== * Module level functions *=================================================================== @@ -12131,27 +11336,6 @@ PyInit__tskit(void) Py_INCREF(&LdCalculatorType); PyModule_AddObject(module, "LdCalculator", (PyObject *) &LdCalculatorType); - /* CompressedMatrix type */ - if (PyType_Ready(&CompressedMatrixType) < 0) { - return NULL; - } - Py_INCREF(&CompressedMatrixType); - PyModule_AddObject(module, "CompressedMatrix", (PyObject *) &CompressedMatrixType); - - /* ViterbiMatrix type */ - if (PyType_Ready(&ViterbiMatrixType) < 0) { - return NULL; - } - Py_INCREF(&ViterbiMatrixType); - PyModule_AddObject(module, "ViterbiMatrix", (PyObject *) &ViterbiMatrixType); - - /* LsHmm type */ - if (PyType_Ready(&LsHmmType) < 0) { - return NULL; - } - Py_INCREF(&LsHmmType); - PyModule_AddObject(module, "LsHmm", (PyObject *) &LsHmmType); - /* IdentitySegments type */ if (PyType_Ready(&IdentitySegmentsType) < 0) { return NULL; diff --git a/python/lwt_interface/setup.py b/python/lwt_interface/setup.py index ff2389ddae..a6856e63b1 100644 --- a/python/lwt_interface/setup.py +++ b/python/lwt_interface/setup.py @@ -31,7 +31,6 @@ def finalize_options(self): "genotypes.c", "stats.c", "convert.c", - "haplotype_matching.c", ] sources = ( ["example_c_module.c"] diff --git a/python/pyproject.toml b/python/pyproject.toml index 33cb69f55d..09b9080528 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -66,7 +66,6 @@ test = [ "coverage", "dendropy", "kastore", - "lshmm", "msgpack", "msprime", "networkx", diff --git a/python/setup.py b/python/setup.py index b284e65b6b..8e663ca778 100644 --- a/python/setup.py +++ b/python/setup.py @@ -16,7 +16,6 @@ "genotypes.c", "stats.c", "convert.c", - "haplotype_matching.c", ] sources = ( ["_tskitmodule.c"] diff --git a/python/tests/test_genotype_matching.py b/python/tests/test_genotype_matching.py deleted file mode 100644 index 091b0bde42..0000000000 --- a/python/tests/test_genotype_matching.py +++ /dev/null @@ -1,1504 +0,0 @@ -import copy -import itertools - -import lshmm as ls -import msprime -import numpy as np -import pytest - -import tskit - -EQUAL_BOTH_HOM = 4 -UNEQUAL_BOTH_HOM = 0 -BOTH_HET = 7 -REF_HOM_OBS_HET = 1 -REF_HET_OBS_HOM = 2 - -MISSING = -1 - - -def mirror_coordinates(ts): - """ - Returns a copy of the specified tree sequence in which all - coordinates x are transformed into L - x. - """ - L = ts.sequence_length - tables = ts.dump_tables() - left = tables.edges.left - right = tables.edges.right - tables.edges.left = L - right - tables.edges.right = L - left - tables.sites.position = L - tables.sites.position # + 1 - # TODO migrations. - tables.sort() - return tables.tree_sequence() - - -class ValueTransition: - """Simple struct holding value transition values.""" - - def __init__(self, tree_node=-1, inner_summation=-1, value_list=-1, value_index=-1): - self.tree_node = tree_node - self.value_list = value_list - self.inner_summation = inner_summation - self.value_index = value_index - - def copy(self): - return ValueTransition( - self.tree_node, - self.inner_summation, - self.value_list, - self.value_index, - ) - - def __repr__(self): - return repr(self.__dict__) - - def __str__(self): - return repr(self) - - -class InternalValueTransition: - """Simple struct holding the internal value transition values.""" - - def __init__(self, tree_node=-1, value=-1, inner_summation=-1): - self.tree_node = tree_node - self.value = value - self.inner_summation = inner_summation - - def __repr__(self): - return repr(self.__dict__) - - def __str__(self): - return repr(self) - - -class LsHmmAlgorithm: - """ - Abstract superclass of Li and Stephens HMM algorithm. - """ - - def __init__(self, ts, rho, mu, precision=30): - self.ts = ts - self.mu = mu - self.rho = rho - self.precision = precision - # The array of ValueTransitions. - self.T = [] - # indexes in to the T array for each node. - self.T_index = np.zeros(ts.num_nodes, dtype=int) - 1 - # The number of nodes underneath each element in the T array. - self.N = np.zeros(ts.num_nodes, dtype=int) - # Efficiently compute the allelic state at a site - self.allelic_state = np.zeros(ts.num_nodes, dtype=int) - 1 - # Diffs so we can can update T and T_index between trees. - self.edge_diffs = self.ts.edge_diffs() - self.parent = np.zeros(self.ts.num_nodes, dtype=int) - 1 - self.tree = tskit.Tree(self.ts) - self.output = None - - def decode_site_dict(self): - """ - Decodes the tree encoding of the values into an explicit - matrix. - """ - A = np.zeros((self.ts.num_samples, self.ts.num_samples)) - # To look at the inner summations too. - B = np.zeros((self.ts.num_samples, self.ts.num_samples)) - f = {st.tree_node: st for st in self.T} - - for j1, u1 in enumerate(self.ts.samples()): - while u1 not in f: - u1 = self.tree.parent(u1) - f1 = {st.tree_node: st for st in f[u1].value_list} - for j2, u2 in enumerate(self.ts.samples()): - while u2 not in f1: - u2 = self.tree.parent(u2) - A[j1, j2] = f1[u2].value - B[j1, j2] = f1[u2].inner_summation - return A, B - - def check_integrity(self): - M = [st.tree_node for st in self.T if st.tree_node != -1] - assert np.all(self.T_index[M] >= 0) - index = np.ones_like(self.T_index, dtype=bool) - index[M] = 0 - assert np.all(self.T_index[index] == -1) - for j, st in enumerate(self.T): - if st.tree_node != -1: - assert j == self.T_index[st.tree_node] - - def stupid_compress_dict(self): - """ - Duncan created a compression that just runs parsimony so is - guaranteed to work. - """ - tree = self.tree - T = self.T - alleles_string_vec = np.zeros(tree.num_samples()).astype("object") - genotypes = np.zeros(tree.num_samples(), dtype=int) - genotype_index = 0 - mapping_back = {} - - node_map = {st.tree_node: st for st in self.T} - - for st1 in T: - if st1.tree_node != -1: - alleles_string_tmp = [ - f"{st2.tree_node}:{st2.value:.16f}" for st2 in st1.value_list - ] - alleles_string = ",".join(alleles_string_tmp) - # Add an extra element that tells me the alleles_string there. - st1.alleles_string = alleles_string - st1.genotype_index = genotype_index - # assert alleles_string not in mapping_back - if alleles_string not in mapping_back: - mapping_back[alleles_string] = { - "tree_node": st1.tree_node, - "value_list": st1.value_list, - "inner_summation": st1.inner_summation, - } - genotype_index += 1 - - for leaf in tree.samples(): - u = leaf - while u not in node_map: - u = tree.parent(u) - genotypes[leaf] = node_map[u].genotype_index - - alleles_string_vec = [] - for st in T: - if st.tree_node != -1: - alleles_string_vec.append(st.alleles_string) - - ancestral_allele, mutations = tree.map_mutations(genotypes, alleles_string_vec) - - # Retain the old T_index, because the internal T that's passed up the tree will - # retain this ordering. - old_T_index = copy.deepcopy(self.T_index) - self.T_index = np.zeros(tree.tree_sequence.num_nodes, dtype=int) - 1 - self.N = np.zeros(tree.tree_sequence.num_nodes, dtype=int) - self.T.clear() - - # First, create T root. - self.T_index[tree.root] = 0 - self.T.append( - ValueTransition( - tree_node=tree.root, - value_list=[ - InternalValueTransition( - tree_node=tree.root, - value=mapping_back[ancestral_allele]["value_list"][ - old_T_index[mapping_back[ancestral_allele]["tree_node"]] - ].value, - ) - ], - ) - ) - - # Then create the rest of T, adding the root each time to value_list - for i, mut in enumerate(mutations): - self.T_index[mut.node] = i + 1 - self.T.append( - ValueTransition( - tree_node=mut.node, - value_list=[ - InternalValueTransition( - tree_node=tree.root, - value=mapping_back[mut.derived_state]["value_list"][ - old_T_index[mapping_back[ancestral_allele]["tree_node"]] - ].value, - ) - ], - ) - ) - - # First add to the root - for mut in mutations: - self.T[self.T_index[tree.root]].value_list.append( - InternalValueTransition( - tree_node=mut.node, - value=mapping_back[ancestral_allele]["value_list"][ - old_T_index[mapping_back[mut.derived_state]["tree_node"]] - ].value, - ) - ) - - # Then add the rest of T_internal to each internal T. - for mut1 in mutations: - for mut2 in mutations: - self.T[self.T_index[mut1.node]].value_list.append( - InternalValueTransition( - tree_node=mut2.node, - value=mapping_back[mut1.derived_state]["value_list"][ - old_T_index[mapping_back[mut2.derived_state]["tree_node"]] - ].value, - ) - ) - - # General approach here is to use - # mapping_back[mut.derived_state]['value_list'][ - # old_T_index[mapping_back[mut2.derived_state]["tree_node"] - # ] and append this to the T_inner. - - node_map = {st.tree_node: st for st in self.T} - - for u in tree.samples(): - while u not in node_map: - u = tree.parent(u) - self.N[self.T_index[u]] += 1 - - def update_tree(self): - """ - Update the internal data structures to move on to the next tree. - """ - parent = self.parent - T_index = self.T_index - T = self.T - _, edges_out, edges_in = next(self.edge_diffs) - - for edge in edges_out: - u = edge.child - if T_index[u] == -1: - # Make sure the subtree we're detaching has an T_index-value at the root. - while T_index[u] == -1: - u = parent[u] - assert u != -1 - T_index[edge.child] = len(T) - T.append( - ValueTransition( - tree_node=edge.child, - value_list=copy.deepcopy(T[T_index[u]].value_list), - ) - ) - # Add on this extra node to each of the internal lists - for st in T: - if not (st.value_list == tskit.NULL): - st.value_list.append( - InternalValueTransition( - tree_node=edge.child, - value=st.value_list.copy()[T_index[u]].value, - ) - ) - parent[edge.child] = -1 - - for edge in edges_in: - parent[edge.child] = edge.parent - u = edge.parent - if parent[edge.parent] == -1: - # Grafting onto a new root. - if T_index[edge.parent] == -1: - T_index[edge.parent] = len(T) - T.append( - ValueTransition( - tree_node=edge.parent, - value_list=copy.deepcopy(T[T_index[edge.child]].value_list), - ) - ) - # Add on this extra node to each of the internal lists - for st in T: - if not (st.value_list == tskit.NULL): - st.value_list.append( - InternalValueTransition( - tree_node=edge.parent, - value=st.value_list.copy()[ - T_index[edge.child] - ].value, - ) - ) - else: - # Grafting into an existing subtree. - while T_index[u] == -1: - u = parent[u] - assert u != -1 - assert T_index[u] != -1 and T_index[edge.child] != -1 - if ( - T[T_index[u]].value_list == T[T_index[edge.child]].value_list - ): # DEV: is this fine? - st = T[T_index[edge.child]] - # Mark the lower ValueTransition as unused. - st.value_list = -1 - # Also need to mark the corresponding InternalValueTransition as - # unused for the remaining states - for st2 in T: - if not (st2.value_list == tskit.NULL): - st2.value_list[T_index[edge.child]].value = -1 - st2.value_list[T_index[edge.child]].tree_node = -1 - - st.tree_node = -1 - T_index[edge.child] = -1 - - # We can have values left over still pointing to old roots. Remove - for root in self.tree.roots: - if T_index[root] != -1: - # Use a special marker here to designate the real roots. - T[T_index[root]].value_index = -2 - - for vt in T: - if vt.tree_node != -1: - if parent[vt.tree_node] == -1 and vt.value_index != -2: - # Also need to mark the corresponding InternalValueTransition as - # unused for the remaining states - for st2 in T: - if not (st2.value_list == tskit.NULL): - st2.value_list[T_index[vt.tree_node]].value = -1 - st2.value_list[T_index[vt.tree_node]].tree_node = -1 - T_index[vt.tree_node] = -1 - vt.tree_node = -1 - vt.value_index = -1 - - self.N = np.zeros(self.tree.tree_sequence.num_nodes, dtype=int) - node_map = {st.tree_node: st for st in self.T} - - for u in self.tree.samples(): - while u not in node_map: - u = self.tree.parent(u) - self.N[self.T_index[u]] += 1 - - def update_probabilities(self, site, genotype_state): - tree = self.tree - T_index = self.T_index - T = self.T - alleles = ["0", "1"] - allelic_state = self.allelic_state - # Set the allelic_state for this site. - allelic_state[self.tree.root] = alleles.index(site.ancestral_state) - normalisation_factor_inner = {} - - for st1 in T: - if st1.tree_node != -1: - normalisation_factor_inner[st1.tree_node] = ( - self.compute_normalisation_factor_inner_dict(st1.tree_node) - ) - - for st1 in T: - if st1.tree_node != -1: - for st2 in st1.value_list: - if st2.tree_node != -1: - self.T[self.T_index[st1.tree_node]].value_list[ - self.T_index[st2.tree_node] - ].inner_summation = self.inner_summation_evaluation( - normalisation_factor_inner, - st1.tree_node, - st2.tree_node, - ) - # ( - # normalisation_factor_inner[st1.tree_node] - # + normalisation_factor_inner[st2.tree_node] - # ) - - for mutation in site.mutations: - u = mutation.node - allelic_state[u] = alleles.index(mutation.derived_state) - if T_index[u] == -1: - while T_index[u] == tskit.NULL: - u = tree.parent(u) - T_index[mutation.node] = len(T) - T.append( - ValueTransition( - tree_node=mutation.node, - value_list=copy.deepcopy(T[T_index[u]].value_list), - ) # DEV: is it possible to not use deepcopies? - ) - for st in T: - if not (st.value_list == tskit.NULL): - st.value_list.append( - InternalValueTransition( - tree_node=mutation.node, - value=st.value_list.copy()[T_index[u]].value, - inner_summation=st.value_list.copy()[ - T_index[u] - ].inner_summation, - ) - ) - - # Get the allelic state at the leaves. - allelic_state[: tree.num_samples()] = tree.tree_sequence.genotype_matrix()[ - site.id, : - ] - - query_is_het = genotype_state == 1 - query_is_missing = genotype_state == MISSING - - for st1 in T: - u1 = st1.tree_node - - if u1 != -1: - # Get the allelic_state at u. TODO we can cache these states to - # avoid some upward traversals. - v1 = u1 - while allelic_state[v1] == -1: - v1 = tree.parent(v1) - assert v1 != -1 - - for st2 in st1.value_list: - u2 = st2.tree_node - if u2 != -1: - # Get the allelic_state at u. TODO we can cache these states to - # avoid some upward traversals. - v2 = u2 - while allelic_state[v2] == -1: - v2 = tree.parent(v2) - assert v2 != -1 - - genotype_template_state = allelic_state[v1] + allelic_state[v2] - match = genotype_state == genotype_template_state - template_is_het = genotype_template_state == 1 - # Fill in the value at the combination of states: (s1, s2) - st2.value = self.compute_next_probability_dict( - site.id, - st2.value, - st2.inner_summation, - match, - template_is_het, - query_is_het, - query_is_missing, - u1, - u2, - # Last two are not used by forward-backward - # but required by Viterbi - ) - - # This will ensure that allelic_state[:n] is filled - genotype_template_state = ( - allelic_state[v1] + allelic_state[: tree.num_samples()] - ) - # These are vectors of length n (at internal nodes). - match = genotype_state == genotype_template_state - template_is_het = genotype_template_state == 1 - - # Unset the states - allelic_state[tree.root] = -1 - for mutation in site.mutations: - allelic_state[mutation.node] = -1 - - def process_site( - self, site, genotype_state, forwards=True - ): # Note: forwards turned on for Viterbi - if forwards: - # Forwards algorithm - self.update_probabilities(site, genotype_state) - self.stupid_compress_dict() - s1 = self.compute_normalisation_factor_dict() - T = self.T - - for st in T: - if st.tree_node != tskit.NULL: - # Need to loop through value copy, and normalise - for st2 in st.value_list: - st2.value /= s1 - st2.value = np.round(st2.value, self.precision) - - self.output.store_site( - site.id, s1, [(st.tree_node, st.value_list) for st in self.T] - ) - else: - # Backwards algorithm - self.output.store_site( - site.id, - self.output.normalisation_factor[site.id], - [(st.tree_node, st.value_list) for st in self.T], - ) - self.update_probabilities(site, genotype_state) - self.stupid_compress_dict() - b_last_sum = self.compute_normalisation_factor_dict() # (Double sum) - - normalisation_factor_inner = {} - - for st1 in self.T: - if st1.tree_node != -1: - normalisation_factor_inner[st1.tree_node] = ( - self.compute_normalisation_factor_inner_dict(st1.tree_node) - ) - - for st1 in self.T: - if st1.tree_node != -1: - for st2 in st1.value_list: - if st2.tree_node != -1: - self.T[self.T_index[st1.tree_node]].value_list[ - self.T_index[st2.tree_node] - ].inner_summation = ( - normalisation_factor_inner[st1.tree_node] - + normalisation_factor_inner[st2.tree_node] - ) - - s = self.output.normalisation_factor[site.id] - for st1 in self.T: - if st1.tree_node != tskit.NULL: - for st2 in st1.value_list: - st2.value = ( - ((self.rho[site.id] / self.ts.num_samples) ** 2) * b_last_sum - + (1 - self.rho[site.id]) - * (self.rho[site.id] / self.ts.num_samples) - * st2.inner_summation - + (1 - self.rho[site.id]) ** 2 * st2.value - ) - st2.value /= s - st2.value = np.round(st2.value, self.precision) - - def run_forward(self, g): - n = self.ts.num_samples - self.tree.clear() - for u in self.ts.samples(): - self.T_index[u] = len(self.T) - self.T.append(ValueTransition(tree_node=u, value_list=[])) - for v in self.ts.samples(): - self.T[self.T_index[u]].value_list.append( - InternalValueTransition(tree_node=v, value=(1 / n) ** 2) - ) - - while self.tree.next(): - self.update_tree() - for site in self.tree.sites(): - self.process_site(site, g[site.id]) - - return self.output - - def run_backward(self, g): - self.tree.clear() - for u in self.ts.samples(): - self.T_index[u] = len(self.T) - self.T.append(ValueTransition(tree_node=u, value_list=[])) - for v in self.ts.samples(): - self.T[self.T_index[u]].value_list.append( - InternalValueTransition(tree_node=v, value=1) - ) - - while self.tree.next(): - self.update_tree() - for site in self.tree.sites(): - self.process_site(site, g[site.id], forwards=False) - return self.output - - def compute_normalisation_factor_dict(self): - raise NotImplementedError() - - def compute_next_probability_dict( - self, - site_id, - p_last, - inner_summation, - is_match, - template_is_het, - query_is_het, - query_is_missing, - node_1, # Note: these are only used in Viterbi (node_1 and node_2) - node_2, - ): - raise NotImplementedError() - - -class CompressedMatrix: - """ - Class representing a num_samples x num_sites matrix compressed by a - tree sequence. Each site is represented by a set of (node, value) - pairs, which act as "mutations", i.e., any sample that descends - from a particular node will inherit that value (unless any other - values are on the path). - """ - - def __init__(self, ts, normalisation_factor=None): - self.ts = ts - self.num_sites = ts.num_sites - self.num_samples = ts.num_samples - self.value_transitions = [None for _ in range(self.num_sites)] - if normalisation_factor is None: - self.normalisation_factor = np.zeros(self.num_sites) - else: - self.normalisation_factor = normalisation_factor - assert len(self.normalisation_factor) == self.num_sites - - def store_site(self, site, normalisation_factor, value_transitions): - self.normalisation_factor[site] = normalisation_factor - self.value_transitions[site] = copy.deepcopy(value_transitions) - - # Expose the same API as the low-level classes - - @property - def num_transitions(self): - a = [len(self.value_transitions[j]) for j in range(self.num_sites)] - return np.array(a, dtype=np.int32) - - def get_site(self, site): - return self.value_transitions[site] - - def decode_site_dict(self, tree, site_id): - """ - Decodes the tree encoding of the values into an explicit - matrix. - """ - A = np.zeros((self.num_samples, self.num_samples)) - f = dict(self.value_transitions[site_id]) - - for j1, u1 in enumerate(self.ts.samples()): - while u1 not in f: - u1 = tree.parent(u1) - f1 = {st.tree_node: st for st in f[u1]} - for j2, u2 in enumerate(self.ts.samples()): - while u2 not in f1: - u2 = tree.parent(u2) - A[j1, j2] = f1[u2].value - return A - - def decode(self): - """ - Decodes the tree encoding of the values into an explicit - matrix. - """ - A = np.zeros((self.num_sites, self.num_samples, self.num_samples)) - for tree in self.ts.trees(): - for site in tree.sites(): - A[site.id] = self.decode_site_dict(tree, site.id) - return A - - -class ForwardMatrix(CompressedMatrix): - """Class representing a compressed forward matrix.""" - - -class BackwardMatrix(CompressedMatrix): - """Class representing a compressed forward matrix.""" - - -class ViterbiMatrix(CompressedMatrix): - """ - Class representing the compressed Viterbi matrix. - """ - - def __init__(self, ts): - super().__init__(ts) - # Tuples containing the site, the pair of nodes in the tree, - # and whether recombination is required - self.double_recombination_required = [(-1, 0, 0, False)] - self.single_recombination_required = [(-1, 0, 0, False)] - - def add_single_recombination_required(self, site, node_s1, node_s2, required): - self.single_recombination_required.append((site, node_s1, node_s2, required)) - - def add_double_recombination_required(self, site, node_s1, node_s2, required): - self.double_recombination_required.append((site, node_s1, node_s2, required)) - - def choose_sample_double(self, site_id, tree): - max_value = -1 - u1 = -1 - u2 = -1 - - for node_s1, value_outer in self.value_transitions[site_id]: - for value_list in value_outer: - value_tmp = value_list - if value_tmp.value > max_value: - max_value = value_tmp.value - u1 = node_s1 - u2 = value_tmp.tree_node - - assert u1 != -1 - assert u2 != -1 - - transition_nodes = [u_tmp for (u_tmp, _) in self.value_transitions[site_id]] - - while not tree.is_sample(u1): - for v in tree.children(u1): - if v not in transition_nodes: - u1 = v - break - else: - raise AssertionError("could not find path") - - while not tree.is_sample(u2): - for v in tree.children(u2): - if v not in transition_nodes: - u2 = v - break - else: - raise AssertionError("could not find path") - - return (u1, u2) - - def choose_sample_single(self, site_id, tree, current_nodes): - # I want to find which is the max between any choice if I switch just u1, - # and any choice if I switch just u2. - node_map = {st[0]: st for st in self.value_transitions[site_id]} - to_compute = ( - np.zeros(2, dtype=int) - 1 - ) # We have two to compute - one for each single switch set of possibilities. - - for i, v in enumerate(current_nodes): # (u1, u2) - while v not in node_map: - v = tree.parent(v) - to_compute[i] = v - - # Need to go to the (j1 :)th entries, and the (:,j2)the entries, - # and pick the best. - T_index = np.zeros(self.ts.num_nodes, dtype=int) - 1 - for j, st in enumerate(self.value_transitions[site_id]): - T_index[st[0]] = j - - node_single_switch_maxes = np.zeros(2, dtype=int) - 1 - single_switch = np.zeros(2) - 1 - - for i, node in enumerate(to_compute): - value_list = self.value_transitions[site_id][T_index[node]][1] - s_inner = 0 - for st in value_list: - j = st.tree_node - if j != -1: - max_st = st.value - max_arg = st.tree_node - if max_st > s_inner: - s_inner = max_st - s_arg = max_arg - node_single_switch_maxes[i] = s_arg - single_switch[i] = s_inner - - if np.argmax(single_switch) == 0: - # u1 is fixed, and we switch u2 - u1 = current_nodes[0] - current_nodes = (u1, node_single_switch_maxes[0]) - else: - # u2 is fixed, and we switch u1. - u2 = current_nodes[1] - current_nodes = (node_single_switch_maxes[1], u2) - - u1 = current_nodes[0] - u2 = current_nodes[1] - - # Find the collection of transition nodes to use to descend down the tree - transition_nodes = [u for (u, _) in self.value_transitions[site_id]] - - # Traverse down to find a leaves. - while not tree.is_sample(u1): - for v in tree.children(u1): - if v not in transition_nodes: - u1 = v - break - else: - raise AssertionError("could not find path") - - while not tree.is_sample(u2): - for v in tree.children(u2): - if v not in transition_nodes: - u2 = v - break - else: - raise AssertionError("could not find path") - - current_nodes = (u1, u2) - - return current_nodes - - def traceback(self): - # Run the traceback. - m = self.ts.num_sites - match = np.zeros((m, 2), dtype=int) - - single_recombination_tree = ( - np.zeros((self.ts.num_nodes, self.ts.num_nodes), dtype=int) - 1 - ) - double_recombination_tree = ( - np.zeros((self.ts.num_nodes, self.ts.num_nodes), dtype=int) - 1 - ) - - tree = tskit.Tree(self.ts) - tree.last() - double_switch = True - current_nodes = (-1, -1) - current_node_outer = current_nodes[0] - - rr_single_index = len(self.single_recombination_required) - 1 - rr_double_index = len(self.double_recombination_required) - 1 - - for site in reversed(self.ts.sites()): - while tree.interval.left > site.position: - tree.prev() - assert tree.interval.left <= site.position < tree.interval.right - - # Fill in the recombination single tree - j_single = rr_single_index - # The above starts from the end of all the recombination required - # information, and includes all the information for the current site. - while self.single_recombination_required[j_single][0] == site.id: - u1, u2, required = self.single_recombination_required[j_single][1:] - single_recombination_tree[u1, u2] = required - j_single -= 1 - - # Fill in the recombination double tree - j_double = rr_double_index - # The above starts from the end of all the recombination required - # information, and includes all the information for the current site. - while self.double_recombination_required[j_double][0] == site.id: - u1, u2, required = self.double_recombination_required[j_double][1:] - double_recombination_tree[u1, u2] = required - j_double -= 1 - - # Note - current nodes are the leaf nodes. - if current_node_outer == -1: - if double_switch: - current_nodes = self.choose_sample_double(site.id, tree) - else: - current_nodes = self.choose_sample_single( - site.id, tree, current_nodes - ) - - match[site.id, :] = current_nodes - - # Now traverse up the tree from the current node. The first marked node - # we meet tells us whether we need to recombine. - current_node_outer = current_nodes[0] - u1 = current_node_outer - u2 = current_nodes[1] - - # Just need to move up the tree to evaluate u1 and u2. - if double_switch: - while u1 != -1 and double_recombination_tree[u1, u1] == -1: - u1 = tree.parent(u1) - - while u2 != -1 and double_recombination_tree[u1, u2] == -1: - u2 = tree.parent(u2) - else: - while u1 != -1 and single_recombination_tree[u1, u1] == -1: - u1 = tree.parent(u1) - - while u2 != -1 and single_recombination_tree[u1, u2] == -1: - u2 = tree.parent(u2) - - assert u1 != -1 - assert u2 != -1 - - if double_recombination_tree[u1, u2] == 1: - # Need to double switch at the next site. - current_node_outer = -1 - double_switch = True - elif single_recombination_tree[u1, u2] == 1: - # Need to single switch at the next site - current_node_outer = -1 - double_switch = False - - # Reset the nodes in the double recombination tree. - j = rr_single_index - while self.single_recombination_required[j][0] == site.id: - u1_tmp, u2_tmp, _ = self.single_recombination_required[j][1:] - single_recombination_tree[u1_tmp, u2_tmp] = -1 - j -= 1 - rr_single_index = j - - # Reset the nodes in the single recombination tree. - j = rr_double_index - while self.double_recombination_required[j][0] == site.id: - u1_tmp, u2_tmp, _ = self.double_recombination_required[j][1:] - double_recombination_tree[u1_tmp, u2_tmp] = -1 - j -= 1 - rr_double_index = j - - return match - - -class ForwardAlgorithm(LsHmmAlgorithm): - """Runs the Li and Stephens forward algorithm.""" - - def __init__(self, ts, rho, mu, precision=30): - super().__init__(ts, rho, mu, precision) - self.output = ForwardMatrix(ts) - - def inner_summation_evaluation( - self, normalisation_factor_inner, st1_tree_node, st2_tree_node - ): - return ( - normalisation_factor_inner[st1_tree_node] - + normalisation_factor_inner[st2_tree_node] - ) - - def compute_normalisation_factor_dict(self): - s = 0 - for j, st in enumerate(self.T): - assert st.tree_node != tskit.NULL - assert self.N[j] > 0 - s += self.N[j] * self.compute_normalisation_factor_inner_dict(st.tree_node) - return s - - def compute_normalisation_factor_inner_dict(self, node): - s_inner = 0 - F_previous = self.T[self.T_index[node]].value_list - for st in F_previous: - j = st.tree_node - if j != -1: - s_inner += self.N[self.T_index[j]] * st.value - return s_inner - - def compute_next_probability_dict( - self, - site_id, - p_last, - inner_normalisation_factor, - is_match, - template_is_het, - query_is_het, - query_is_missing, - node_1, - node_2, - ): - rho = self.rho[site_id] - mu = self.mu[site_id] - n = self.ts.num_samples - - p_t = ( - (rho / n) ** 2 - + ((1 - rho) * (rho / n)) * inner_normalisation_factor - + (1 - rho) ** 2 * p_last - ) - - if query_is_missing: - p_e = 1 - else: - query_is_hom = np.logical_not(query_is_het) - template_is_hom = np.logical_not(template_is_het) - - equal_both_hom = np.logical_and( - np.logical_and(is_match, template_is_hom), query_is_hom - ) - unequal_both_hom = np.logical_and( - np.logical_and(np.logical_not(is_match), template_is_hom), query_is_hom - ) - both_het = np.logical_and(template_is_het, query_is_het) - ref_hom_obs_het = np.logical_and(template_is_hom, query_is_het) - ref_het_obs_hom = np.logical_and(template_is_het, query_is_hom) - - p_e = ( - equal_both_hom * (1 - mu) ** 2 - + unequal_both_hom * (mu**2) - + ref_hom_obs_het * (2 * mu * (1 - mu)) - + ref_het_obs_hom * (mu * (1 - mu)) - + both_het * ((1 - mu) ** 2 + mu**2) - ) - - return p_t * p_e - - -class BackwardAlgorithm(LsHmmAlgorithm): - """Runs the Li and Stephens forward algorithm.""" - - def __init__(self, ts, rho, mu, normalisation_factor, precision=10): - super().__init__(ts, rho, mu, precision) - self.output = BackwardMatrix(ts, normalisation_factor) - - def inner_summation_evaluation( - self, normalisation_factor_inner, st1_tree_node, st2_tree_node - ): - return ( - normalisation_factor_inner[st1_tree_node] - + normalisation_factor_inner[st2_tree_node] - ) - - def compute_normalisation_factor_dict(self): - s = 0 - for j, st in enumerate(self.T): - assert st.tree_node != tskit.NULL - assert self.N[j] > 0 - s += self.N[j] * self.compute_normalisation_factor_inner_dict(st.tree_node) - return s - - def compute_normalisation_factor_inner_dict(self, node): - s_inner = 0 - F_previous = self.T[self.T_index[node]].value_list - for st in F_previous: - j = st.tree_node - if j != -1: - s_inner += self.N[self.T_index[j]] * st.value - return s_inner - - def compute_next_probability_dict( - self, - site_id, - p_next, - inner_normalisation_factor, - is_match, - template_is_het, - query_is_het, - query_is_missing, - node_1, - node_2, - ): - mu = self.mu[site_id] - template_is_hom = np.logical_not(template_is_het) - - if query_is_missing: - p_e = 1 - else: - query_is_hom = np.logical_not(query_is_het) - - equal_both_hom = np.logical_and( - np.logical_and(is_match, template_is_hom), query_is_hom - ) - unequal_both_hom = np.logical_and( - np.logical_and(np.logical_not(is_match), template_is_hom), query_is_hom - ) - both_het = np.logical_and(template_is_het, query_is_het) - ref_hom_obs_het = np.logical_and(template_is_hom, query_is_het) - ref_het_obs_hom = np.logical_and(template_is_het, query_is_hom) - - p_e = ( - equal_both_hom * (1 - mu) ** 2 - + unequal_both_hom * (mu**2) - + ref_hom_obs_het * (2 * mu * (1 - mu)) - + ref_het_obs_hom * (mu * (1 - mu)) - + both_het * ((1 - mu) ** 2 + mu**2) - ) - - return p_next * p_e - - -class ViterbiAlgorithm(LsHmmAlgorithm): - """ - Runs the Li and Stephens Viterbi algorithm. - """ - - def __init__(self, ts, rho, mu, precision=10): - super().__init__(ts, rho, mu, precision) - self.output = ViterbiMatrix(ts) - - def inner_summation_evaluation( - self, normalisation_factor_inner, st1_tree_node, st2_tree_node - ): - return max( - normalisation_factor_inner[st1_tree_node], - normalisation_factor_inner[st2_tree_node], - ) - - def compute_normalisation_factor_dict(self): - s = 0 - for st in self.T: - assert st.tree_node != tskit.NULL - max_st = self.compute_normalisation_factor_inner_dict(st.tree_node) - if max_st > s: - s = max_st - if s == 0: - raise ValueError( - "Trying to match non-existent allele with zero mutation rate" - ) - return s - - def compute_normalisation_factor_inner_dict(self, node): - s_inner = 0 - V_previous = self.T[self.T_index[node]].value_list - for st in V_previous: - j = st.tree_node - if j != -1: - max_st = st.value - if max_st > s_inner: - s_inner = max_st - - return s_inner - - def compute_next_probability_dict( - self, - site_id, - p_last, - inner_normalisation_factor, - is_match, - template_is_het, - query_is_het, - query_is_missing, - node_1, - node_2, - ): - r = self.rho[site_id] - mu = self.mu[site_id] - n = self.ts.num_samples - r_n = r / n - - double_recombination_required = False - single_recombination_required = False - - if query_is_missing: - p_e = 1 - else: - template_is_hom = np.logical_not(template_is_het) - query_is_hom = np.logical_not(query_is_het) - equal_both_hom = np.logical_and( - np.logical_and(is_match, template_is_hom), query_is_hom - ) - unequal_both_hom = np.logical_and( - np.logical_and(np.logical_not(is_match), template_is_hom), query_is_hom - ) - both_het = np.logical_and(template_is_het, query_is_het) - ref_hom_obs_het = np.logical_and(template_is_hom, query_is_het) - ref_het_obs_hom = np.logical_and(template_is_het, query_is_hom) - - p_e = ( - equal_both_hom * (1 - mu) ** 2 - + unequal_both_hom * (mu**2) - + ref_hom_obs_het * (2 * mu * (1 - mu)) - + ref_het_obs_hom * (mu * (1 - mu)) - + both_het * ((1 - mu) ** 2 + mu**2) - ) - - no_switch = (1 - r) ** 2 + 2 * (r_n * (1 - r)) + r_n**2 - single_switch = r_n * (1 - r) + r_n**2 - double_switch = r_n**2 - - V_single_switch = inner_normalisation_factor - p_t = p_last * no_switch - single_switch_tmp = single_switch * V_single_switch - - if single_switch_tmp > double_switch: - # Then single switch is the alternative - if p_t < single_switch * V_single_switch: - p_t = single_switch * V_single_switch - single_recombination_required = True - else: - # Double switch is the alternative - if p_t < double_switch: - p_t = double_switch - double_recombination_required = True - - self.output.add_single_recombination_required( - site_id, node_1, node_2, single_recombination_required - ) - self.output.add_double_recombination_required( - site_id, node_1, node_2, double_recombination_required - ) - - return p_t * p_e - - -def ls_forward_tree(g, ts, rho, mu, precision=30): - """Forward matrix computation based on a tree sequence.""" - fa = ForwardAlgorithm(ts, rho, mu, precision=precision) - return fa.run_forward(g) - - -def ls_backward_tree(g, ts_mirror, rho, mu, normalisation_factor, precision=30): - """Backward matrix computation based on a tree sequence.""" - ba = BackwardAlgorithm(ts_mirror, rho, mu, normalisation_factor, precision=precision) - return ba.run_backward(g) - - -def ls_viterbi_tree(g, ts, rho, mu, precision=30): - """ - Viterbi path computation based on a tree sequence. - """ - va = ViterbiAlgorithm(ts, rho, mu, precision=precision) - return va.run_forward(g) - - -class LSBase: - """Superclass of Li and Stephens tests.""" - - def genotype_emission(self, mu, m): - # Define the emission probability matrix - e = np.zeros((m, 8)) - e[:, EQUAL_BOTH_HOM] = (1 - mu) ** 2 - e[:, UNEQUAL_BOTH_HOM] = mu**2 - e[:, BOTH_HET] = (1 - mu) ** 2 + mu**2 - e[:, REF_HOM_OBS_HET] = 2 * mu * (1 - mu) - e[:, REF_HET_OBS_HOM] = mu * (1 - mu) - - return e - - def example_genotypes(self, ts): - H = ts.genotype_matrix() - s = H[:, 0].reshape(1, H.shape[0]) + H[:, 1].reshape(1, H.shape[0]) - H = H[:, 2:] - - genotypes = [ - s, - H[:, -1].reshape(1, H.shape[0]) + H[:, -2].reshape(1, H.shape[0]), - ] - - s_tmp = s.copy() - s_tmp[0, -1] = MISSING - genotypes.append(s_tmp) - s_tmp = s.copy() - s_tmp[0, ts.num_sites // 2] = MISSING - genotypes.append(s_tmp) - s_tmp = s.copy() - s_tmp[0, :] = MISSING - genotypes.append(s_tmp) - - m = ts.get_num_sites() - n = H.shape[1] - - G = np.zeros((m, n, n)) - for i in range(m): - G[i, :, :] = np.add.outer(H[i, :], H[i, :]) - - return H, G, genotypes - - def example_parameters_genotypes(self, ts, seed=42): - np.random.seed(seed) - H, G, genotypes = self.example_genotypes(ts) - n = H.shape[1] - m = ts.get_num_sites() - - # Here we have equal mutation and recombination - r = np.zeros(m) + 0.01 - mu = np.zeros(m) + 0.01 - r[0] = 0 - - e = self.genotype_emission(mu, m) - - for s in genotypes: - yield n, m, G, s, e, r, mu - - # Mixture of random and extremes - rs = [np.zeros(m) + 0.999, np.zeros(m) + 1e-6, np.random.rand(m)] - mus = [np.zeros(m) + 0.33, np.zeros(m) + 1e-6, np.random.rand(m) * 0.33] - - s = genotypes[0] - for r, mu in itertools.product(rs, mus): - r[0] = 0 - e = self.genotype_emission(mu, m) - yield n, m, G, s, e, r, mu - - def assertAllClose(self, A, B): - """Assert that all entries of two matrices are 'close'""" - assert np.allclose(A, B, rtol=1e-5, atol=1e-8) - - # Define a bunch of very small tree-sequences for testing a collection of - # parameters on - @pytest.mark.skip(reason="No plans to implement diploid LS HMM yet.") - def test_simple_n_10_no_recombination(self): - ts = msprime.simulate( - 10, recombination_rate=0, mutation_rate=0.5, random_seed=42 - ) - assert ts.num_sites > 3 - self.verify(ts) - - @pytest.mark.skip(reason="No plans to implement diploid LS HMM yet.") - def test_simple_n_6(self): - ts = msprime.simulate(6, recombination_rate=2, mutation_rate=7, random_seed=42) - assert ts.num_sites > 5 - self.verify(ts) - - @pytest.mark.skip(reason="No plans to implement diploid LS HMM yet.") - def test_simple_n_8_high_recombination(self): - ts = msprime.simulate(8, recombination_rate=20, mutation_rate=5, random_seed=42) - assert ts.num_trees > 15 - assert ts.num_sites > 5 - self.verify(ts) - - # FIXME Reducing the number of test cases here as they take a long time to run, - # and we will want to refactor the test infrastructure when implementing these - # diploid methods in the library. - - # def test_simple_n_10_no_recombination_high_mut(self): - # ts = msprime.simulate( - # 10, recombination_rate=0, mutation_rate=3, random_seed=42) - # assert ts.num_sites > 3 - # self.verify(ts) - - # def test_simple_n_10_no_recombination_higher_mut(self): - # ts = msprime.simulate( - # 20, recombination_rate=0, mutation_rate=3, random_seed=42) - # assert ts.num_sites > 3 - # self.verify(ts) - - # def test_simple_n_8(self): - # ts = msprime.simulate( - # 8, recombination_rate=2, mutation_rate=5, random_seed=42) - # assert ts.num_sites > 5 - # self.verify(ts) - - # def test_simple_n_16(self): - # ts = msprime.simulate( - # 16, recombination_rate=2, mutation_rate=5, random_seed=42) - # assert ts.num_sites > 5 - # self.verify(ts) - - def verify(self, ts): - raise NotImplementedError() - - -class FBAlgorithmBase(LSBase): - """Base for forwards backwards algorithm tests.""" - - -class VitAlgorithmBase(LSBase): - """Base for viterbi algoritm tests.""" - - -class TestMirroringDipdict(FBAlgorithmBase): - """Tests that mirroring the tree sequence and running forwards and backwards - algorithms give the same log-likelihood of observing the data.""" - - def verify(self, ts): - for n, m, _, s, _, r, mu in self.example_parameters_genotypes(ts): - # Note, need to remove the first sample from the ts, and ensure that - # invariant sites aren't removed. - ts_check, mapping = ts.simplify( - range(1, n + 1), filter_sites=False, map_nodes=True - ) - H_check = ts_check.genotype_matrix() - G_check = np.zeros((m, n, n)) - for i in range(m): - G_check[i, :, :] = np.add.outer(H_check[i, :], H_check[i, :]) - - cm_d = ls_forward_tree(s[0, :], ts_check, r, mu) - ll_tree = np.sum(np.log10(cm_d.normalisation_factor)) - - ts_check_mirror = mirror_coordinates(ts_check) - r_flip = np.insert(np.flip(r)[:-1], 0, 0) - cm_mirror = ls_forward_tree( - np.flip(s[0, :]), ts_check_mirror, r_flip, np.flip(mu) - ) - ll_mirror_tree_dict = np.sum(np.log10(cm_mirror.normalisation_factor)) - - self.assertAllClose(ll_tree, ll_mirror_tree_dict) - - # Ensure that the decoded matrices are the same - flipped_H_check = np.flip(H_check, axis=0) - flipped_s = np.flip(s, axis=1) - - F_mirror_matrix, c, ll = ls.forwards( - flipped_H_check, - flipped_s, - ploidy=2, - prob_recombination=r_flip, - prob_mutation=np.flip(mu), - scale_mutation_rate=False, - ) - - self.assertAllClose(F_mirror_matrix, cm_mirror.decode()) - - -class TestForwardDipTree(FBAlgorithmBase): - """Tests that the tree algorithm computes the same forward matrix as the simple - method.""" - - def verify(self, ts): - for n, m, _, s, _, r, mu in self.example_parameters_genotypes(ts): - # Note, need to remove the first sample from the ts, and ensure that - # invariant sites aren't removed. - ts_check, mapping = ts.simplify( - range(1, n + 1), filter_sites=False, map_nodes=True - ) - H_check = ts_check.genotype_matrix() - G_check = np.zeros((m, n, n)) - for i in range(m): - G_check[i, :, :] = np.add.outer(H_check[i, :], H_check[i, :]) - - F, c, ll = ls.forwards( - reference_panel=H_check, - query=s, - ploidy=2, - prob_recombination=r, - prob_mutation=mu, - scale_mutation_rate=False, - ) - cm_d = ls_forward_tree(s[0, :], ts_check, r, mu) - self.assertAllClose(cm_d.decode(), F) - ll_tree = np.sum(np.log10(cm_d.normalisation_factor)) - self.assertAllClose(ll, ll_tree) - - -class TestForwardBackwardTree(FBAlgorithmBase): - """Tests that the tree algorithm computes the same forward matrix as the simple - method.""" - - def verify(self, ts): - for n, m, _, s, _, r, mu in self.example_parameters_genotypes(ts): - # Note, need to remove the first sample from the ts, and ensure that - # invariant sites aren't removed. - ts_check, mapping = ts.simplify( - range(1, n + 1), filter_sites=False, map_nodes=True - ) - H_check = ts_check.genotype_matrix() - G_check = np.zeros((m, n, n)) - for i in range(m): - G_check[i, :, :] = np.add.outer(H_check[i, :], H_check[i, :]) - - F, c, ll = ls.forwards( - reference_panel=H_check, - query=s, - ploidy=2, - prob_recombination=r, - prob_mutation=mu, - scale_mutation_rate=False, - ) - B = ls.backwards( - reference_panel=H_check, - query=s, - ploidy=2, - normalisation_factor_from_forward=c, - prob_recombination=r, - prob_mutation=mu, - scale_mutation_rate=False, - ) - - # Note, need to remove the first sample from the ts, and ensure that - # invariant sites aren't removed. - ts_check = ts.simplify(range(1, n + 1), filter_sites=False) - c_f = ls_forward_tree(s[0, :], ts_check, r, mu) - ll_tree = np.sum(np.log10(c_f.normalisation_factor)) - - ts_check_mirror = mirror_coordinates(ts_check) - r_flip = np.flip(r) - c_b = ls_backward_tree( - np.flip(s[0, :]), - ts_check_mirror, - r_flip, - np.flip(mu), - np.flip(c_f.normalisation_factor), - ) - B_tree = np.flip(c_b.decode(), axis=0) - F_tree = c_f.decode() - - self.assertAllClose(B, B_tree) - self.assertAllClose(F, F_tree) - self.assertAllClose(ll, ll_tree) - - -class TestTreeViterbiDip(VitAlgorithmBase): - """ - Test that we have the same log-likelihood between tree and matrix - implementations - """ - - def verify(self, ts): - for n, m, _, s, _, r, mu in self.example_parameters_genotypes(ts): - # Note, need to remove the first sample from the ts, and ensure that - # invariant sites aren't removed. - ts_check, mapping = ts.simplify( - range(1, n + 1), filter_sites=False, map_nodes=True - ) - H_check = ts_check.genotype_matrix() - G_check = np.zeros((m, n, n)) - for i in range(m): - G_check[i, :, :] = np.add.outer(H_check[i, :], H_check[i, :]) - ts_check = ts.simplify(range(1, n + 1), filter_sites=False) - - phased_path, ll = ls.viterbi( - reference_panel=H_check, - query=s, - ploidy=2, - prob_recombination=r, - prob_mutation=mu, - scale_mutation_rate=False, - ) - path_ll_matrix = ls.path_loglik( - reference_panel=H_check, - query=s, - ploidy=2, - path=phased_path, - prob_recombination=r, - prob_mutation=mu, - scale_mutation_rate=False, - ) - - c_v = ls_viterbi_tree(s[0, :], ts_check, r, mu) - ll_tree = np.sum(np.log10(c_v.normalisation_factor)) - - # Attempt to get the path - path_tree_dict = c_v.traceback() - # Work out the likelihood of the proposed path - path_ll_tree = ls.path_loglik( - reference_panel=H_check, - query=s, - ploidy=2, - path=np.transpose(path_tree_dict), - prob_recombination=r, - prob_mutation=mu, - scale_mutation_rate=False, - ) - - self.assertAllClose(ll, ll_tree) - self.assertAllClose(path_ll_tree, path_ll_matrix) diff --git a/python/tests/test_haplotype_matching.py b/python/tests/test_haplotype_matching.py deleted file mode 100644 index 92bdf63058..0000000000 --- a/python/tests/test_haplotype_matching.py +++ /dev/null @@ -1,1327 +0,0 @@ -# MIT License -# -# Copyright (c) 2019-2024 Tskit Developers -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -""" -Python implementation of the Li and Stephens forwards and backwards algorithms. -""" - -import warnings - -import lshmm as ls -import msprime -import numpy as np -import numpy.testing as nt -import pytest - -import _tskit -import tskit -from tests import tsutil - -MISSING = -1 - - -def check_alleles(alleles, m): - """ - Checks the specified allele list and returns a list of lists - of alleles of length num_sites. - If alleles is a 1D list of strings, assume that this list is used - for each site and return num_sites copies of this list. - Otherwise, raise a ValueError if alleles is not a list of length - num_sites. - """ - if isinstance(alleles[0], str): - return [alleles for _ in range(m)], np.int8([len(alleles) for _ in range(m)]) - if len(alleles) != m: - raise ValueError("Malformed alleles list") - n_alleles = np.int8([(len(alleles_site)) for alleles_site in alleles]) - return alleles, n_alleles - - -def mirror_coordinates(ts): - """ - Returns a copy of the specified tree sequence in which all - coordinates x are transformed into L - x. - """ - L = ts.sequence_length - tables = ts.dump_tables() - left = tables.edges.left - right = tables.edges.right - tables.edges.left = L - right - tables.edges.right = L - left - tables.sites.position = L - tables.sites.position # + 1 - # TODO migrations. - tables.sort() - return tables.tree_sequence() - - -class ValueTransition: - """Simple struct holding value transition values.""" - - def __init__(self, tree_node=-1, value=-1, value_index=-1): - self.tree_node = tree_node - self.value = value - self.value_index = value_index - - def copy(self): - return ValueTransition( - self.tree_node, - self.value, - self.value_index, - ) - - def __repr__(self): - return repr(self.__dict__) - - def __str__(self): - return repr(self) - - -class LsHmmAlgorithm: - """ - Abstract superclass of Li and Stephens HMM algorithm. - """ - - def __init__( - self, ts, rho, mu, alleles, n_alleles, precision=10, scale_mutation=False - ): - self.ts = ts - self.mu = mu - self.rho = rho - self.precision = precision - # The array of ValueTransitions. - self.T = [] - # indexes in to the T array for each node. - self.T_index = np.zeros(ts.num_nodes, dtype=int) - 1 - # The number of nodes underneath each element in the T array. - self.N = np.zeros(ts.num_nodes, dtype=int) - # Efficiently compute the allelic state at a site - self.allelic_state = np.zeros(ts.num_nodes, dtype=int) - 1 - # TreeIndexes so we can can update T and T_index between trees. - self.tree_pos = tsutil.TreeIndexes(ts) - self.parent = np.zeros(self.ts.num_nodes, dtype=int) - 1 - self.tree = tskit.Tree(self.ts) - self.output = None - # Vector of the number of alleles at each site - self.n_alleles = n_alleles - self.alleles = alleles - self.scale_mutation_based_on_n_alleles = scale_mutation - - def check_integrity(self): - M = [st.tree_node for st in self.T if st.tree_node != -1] - assert np.all(self.T_index[M] >= 0) - index = np.ones_like(self.T_index, dtype=bool) - index[M] = 0 - assert np.all(self.T_index[index] == -1) - for j, st in enumerate(self.T): - if st.tree_node != -1: - assert j == self.T_index[st.tree_node] - - def compress(self): - tree = self.tree - T = self.T - T_index = self.T_index - - values = np.unique(list(st.value if st.tree_node != -1 else 1e200 for st in T)) - for st in T: - if st.tree_node != -1: - st.value_index = np.searchsorted(values, st.value) - - child = np.zeros(len(values), dtype=int) - num_values = len(values) - value_count = np.zeros(num_values, dtype=int) - - def compute(u, parent_state): - value_count[:] = 0 - for v in tree.children(u): - child[:] = optimal_set[v] - # If the set for a given child is empty, then we know it inherits - # directly from the parent state and must be a singleton set. - if np.sum(child) == 0: - child[parent_state] = 1 - for j in range(num_values): - value_count[j] += child[j] - max_value_count = np.max(value_count) - optimal_set[u, :] = 0 - optimal_set[u, value_count == max_value_count] = 1 - - optimal_set = np.zeros((tree.tree_sequence.num_nodes, len(values)), dtype=int) - t_node_time = [-1 if st.tree_node == -1 else tree.time(st.tree_node) for st in T] - order = np.argsort(t_node_time) - for j in order: - st = T[j] - u = st.tree_node - if u != -1: - # Compute the value at this node - state = st.value_index - if tree.is_internal(u): - compute(u, state) - else: - # A[u, state] = 1 - optimal_set[u, state] = 1 - # Find parent state - v = tree.parent(u) - if v != -1: - while T_index[v] == -1: - v = tree.parent(v) - parent_state = T[T_index[v]].value_index - v = tree.parent(u) - while T_index[v] == -1: - compute(v, parent_state) - v = tree.parent(v) - - T_old = [st.copy() for st in T] - T.clear() - T_parent = [] - - old_state = T_old[T_index[tree.root]].value_index - new_state = np.argmax(optimal_set[tree.root]) - - T.append(ValueTransition(tree_node=tree.root, value=values[new_state])) - T_parent.append(-1) - stack = [(tree.root, old_state, new_state, 0)] - while len(stack) > 0: - u, old_state, new_state, t_parent = stack.pop() - for v in tree.children(u): - old_child_state = old_state - if T_index[v] != -1: - old_child_state = T_old[T_index[v]].value_index - if np.sum(optimal_set[v]) > 0: - new_child_state = new_state - child_t_parent = t_parent - - if optimal_set[v, new_state] == 0: - new_child_state = np.argmax(optimal_set[v]) - child_t_parent = len(T) - T_parent.append(t_parent) - T.append( - ValueTransition(tree_node=v, value=values[new_child_state]) - ) - stack.append((v, old_child_state, new_child_state, child_t_parent)) - else: - if old_child_state != new_state: - T_parent.append(t_parent) - T.append( - ValueTransition(tree_node=v, value=values[old_child_state]) - ) - - for st in T_old: - if st.tree_node != -1: - T_index[st.tree_node] = -1 - for j, st in enumerate(T): - T_index[st.tree_node] = j - self.N[j] = tree.num_samples(st.tree_node) - for j in range(len(T)): - if T_parent[j] != -1: - self.N[T_parent[j]] -= self.N[j] - - def update_tree(self, direction=tskit.FORWARD): - """ - Update the internal data structures to move on to the next tree. - """ - parent = self.parent - T_index = self.T_index - T = self.T - if direction == tskit.FORWARD: - self.tree_pos.next() - else: - self.tree_pos.prev() - assert self.tree_pos.index == self.tree.index - - for j in range( - self.tree_pos.out_range.start, self.tree_pos.out_range.stop, direction - ): - e = self.tree_pos.out_range.order[j] - edge = self.ts.edge(e) - u = edge.child - if T_index[u] == -1: - # Make sure the subtree we're detaching has an T_index-value at the root. - while T_index[u] == -1: - u = parent[u] - assert u != -1 - T_index[edge.child] = len(T) - T.append( - ValueTransition(tree_node=edge.child, value=T[T_index[u]].value) - ) - parent[edge.child] = -1 - - for j in range( - self.tree_pos.in_range.start, self.tree_pos.in_range.stop, direction - ): - e = self.tree_pos.in_range.order[j] - edge = self.ts.edge(e) - parent[edge.child] = edge.parent - u = edge.parent - if parent[edge.parent] == -1: - # Grafting onto a new root. - if T_index[edge.parent] == -1: - T_index[edge.parent] = len(T) - T.append( - ValueTransition( - tree_node=edge.parent, value=T[T_index[edge.child]].value - ) - ) - else: - # Grafting into an existing subtree. - while T_index[u] == -1: - u = parent[u] - assert u != -1 - assert T_index[u] != -1 and T_index[edge.child] != -1 - if T[T_index[u]].value == T[T_index[edge.child]].value: - st = T[T_index[edge.child]] - # Mark the lower ValueTransition as unused. - st.value = -1 - st.tree_node = -1 - T_index[edge.child] = -1 - - # We can have values left over still pointing to old roots. Remove - for root in self.tree.roots: - if T_index[root] != -1: - # Use a special marker here to designate the real roots. - T[T_index[root]].value_index = -2 - for vt in T: - if vt.tree_node != -1: - if parent[vt.tree_node] == -1 and vt.value_index != -2: - T_index[vt.tree_node] = -1 - vt.tree_node = -1 - vt.value_index = -1 - - def update_probabilities(self, site, haplotype_state): - tree = self.tree - T_index = self.T_index - T = self.T - alleles = self.alleles[site.id] - allelic_state = self.allelic_state - # Set the allelic_state for this site. - allelic_state[tree.root] = alleles.index(site.ancestral_state) - - for mutation in site.mutations: - u = mutation.node - allelic_state[u] = alleles.index(mutation.derived_state) - if T_index[u] == -1: - while T_index[u] == tskit.NULL: - u = tree.parent(u) - T_index[mutation.node] = len(T) - T.append( - ValueTransition(tree_node=mutation.node, value=T[T_index[u]].value) - ) - - for st in T: - u = st.tree_node - if u != -1: - # Get the allelic_state at u. TODO we can cache these states to - # avoid some upward traversals. - v = u - while allelic_state[v] == -1: - v = tree.parent(v) - assert v != -1 - match = haplotype_state == allelic_state[v] - is_query_missing = haplotype_state == MISSING - # Note that the node u is used only by Viterbi - st.value = self.compute_next_probability( - site.id, st.value, match, is_query_missing, u - ) - - # Unset the states - allelic_state[tree.root] = -1 - for mutation in site.mutations: - allelic_state[mutation.node] = -1 - - def process_site(self, site, haplotype_state): - self.update_probabilities(site, haplotype_state) - self.compress() - s = self.compute_normalisation_factor() - for st in self.T: - assert st.tree_node != tskit.NULL - # if st.tree_node != tskit.NULL: - st.value /= s - st.value = round(st.value, self.precision) - self.output.store_site(site.id, s, [(st.tree_node, st.value) for st in self.T]) - - def compute_emission_proba(self, site_id, is_match): - mu = self.mu[site_id] - n_alleles = self.n_alleles[site_id] - if self.scale_mutation_based_on_n_alleles: - if is_match: - # Scale mutation based on the number of alleles - # - so the mutation rate is the mutation rate to one of the - # alleles. The overall mutation rate is then - # (n_alleles - 1) * mutation_rate. - p_e = 1 - (n_alleles - 1) * mu - else: - p_e = mu - mu * (n_alleles == 1) - # Added boolean in case we're at an invariant site - else: - # No scaling based on the number of alleles - # - so the mutation rate is the mutation rate to anything. - # This means that we must rescale the mutation rate to a different - # allele, by the number of alleles. - if n_alleles == 1: # In case we're at an invariant site - if is_match: - p_e = 1 - else: - p_e = 0 - else: - if is_match: - p_e = 1 - mu - else: - p_e = mu / (n_alleles - 1) - return p_e - - def initialise(self, value): - self.tree.clear() - for u in self.ts.samples(): - j = len(self.T) - self.T_index[u] = j - self.T.append(ValueTransition(tree_node=u, value=value)) - - def run(self, h): - n = self.ts.num_samples - self.initialise(1 / n) - while self.tree.next(): - self.update_tree() - for site in self.tree.sites(): - self.process_site(site, h[site.id]) - return self.output - - def compute_normalisation_factor(self): - raise NotImplementedError() - - def compute_next_probability( - self, site_id, p_last, is_match, is_query_missing, node - ): - raise NotImplementedError() - - -class ForwardAlgorithm(LsHmmAlgorithm): - """ - The Li and Stephens forward algorithm. - """ - - def __init__( - self, ts, rho, mu, alleles, n_alleles, scale_mutation=False, precision=10 - ): - super().__init__( - ts, - rho, - mu, - alleles, - n_alleles, - precision=precision, - scale_mutation=scale_mutation, - ) - self.output = CompressedMatrix(ts) - - def compute_normalisation_factor(self): - s = 0 - for j, st in enumerate(self.T): - assert st.tree_node != tskit.NULL - # assert self.N[j] > 0 - s += self.N[j] * st.value - return s - - def compute_next_probability( - self, site_id, p_last, is_match, is_query_missing, node - ): - rho = self.rho[site_id] - n = self.ts.num_samples - if is_query_missing: - p_e = 1.0 - else: - p_e = self.compute_emission_proba(site_id, is_match) - p_t = p_last * (1 - rho) + rho / n - return p_t * p_e - - -class BackwardAlgorithm(ForwardAlgorithm): - """ - The Li and Stephens backward algorithm. - """ - - def compute_next_probability( - self, site_id, p_next, is_match, is_query_missing, node - ): - if is_query_missing: - p_e = 1.0 - else: - p_e = self.compute_emission_proba(site_id, is_match) - return p_next * p_e - - def process_site(self, site, haplotype_state, s): - # FIXME see nodes in the C code for why we have two calls to - # compress - # https://github.com/tskit-dev/tskit/issues/2803 - self.compress() - self.output.store_site( - site.id, - s, - [(st.tree_node, st.value) for st in self.T], - ) - self.update_probabilities(site, haplotype_state) - # FIXME see nodes in the C code for why we have two calls to - # compress - self.compress() - b_last_sum = self.compute_normalisation_factor() - n = self.ts.num_samples - rho = self.rho[site.id] - for st in self.T: - if st.tree_node != tskit.NULL: - st.value = rho * b_last_sum / n + (1 - rho) * st.value - st.value /= s - st.value = round(st.value, self.precision) - - def run(self, h, normalisation_factor): - self.initialise(value=1) - while self.tree.prev(): - self.update_tree(direction=tskit.REVERSE) - for site in reversed(list(self.tree.sites())): - self.process_site(site, h[site.id], normalisation_factor[site.id]) - return self.output - - -class ViterbiAlgorithm(LsHmmAlgorithm): - """ - Runs the Li and Stephens Viterbi algorithm. - """ - - def __init__( - self, ts, rho, mu, alleles, n_alleles, scale_mutation=False, precision=10 - ): - super().__init__( - ts, - rho, - mu, - alleles, - n_alleles, - precision=precision, - scale_mutation=scale_mutation, - ) - self.output = ViterbiMatrix(ts) - - def compute_normalisation_factor(self): - max_st = ValueTransition(value=-1) - for st in self.T: - assert st.tree_node != tskit.NULL - if st.value > max_st.value: - max_st = st - if max_st.value == 0: - raise ValueError( - "Trying to match non-existent allele with zero mutation rate" - ) - return max_st.value - - def compute_next_probability( - self, site_id, p_last, is_match, is_query_missing, node - ): - rho = self.rho[site_id] - n = self.ts.num_samples - - p_no_recomb = p_last * (1 - rho + rho / n) - p_recomb = rho / n - recombination_required = False - if p_no_recomb > p_recomb: - p_t = p_no_recomb - else: - p_t = p_recomb - recombination_required = True - self.output.add_recombination_required(site_id, node, recombination_required) - - if is_query_missing: - p_e = 1.0 - else: - p_e = self.compute_emission_proba(site_id, is_match) - - return p_t * p_e - - -def assert_compressed_matrices_equal(cm1, cm2): - nt.assert_array_almost_equal(cm1.normalisation_factor, cm2.normalisation_factor) - - for j in range(cm1.num_sites): - site1 = cm1.get_site(j) - site2 = cm2.get_site(j) - assert len(site1) == len(site2) - site1 = dict(site1) - site2 = dict(site2) - - assert set(site1.keys()) == set(site2.keys()) - for node in site1.keys(): - # TODO the precision value should be used as a parameter here - nt.assert_allclose(site1[node], site2[node], rtol=1e-5, atol=1e-8) - - -class CompressedMatrix: - """ - Class representing a num_samples x num_sites matrix compressed by a - tree sequence. Each site is represented by a set of (node, value) - pairs, which act as "mutations", i.e., any sample that descends - from a particular node will inherit that value (unless any other - values are on the path). - """ - - def __init__(self, ts): - self.ts = ts - self.num_sites = ts.num_sites - self.num_samples = ts.num_samples - self.value_transitions = [None for _ in range(self.num_sites)] - self.normalisation_factor = np.zeros(self.num_sites) - - def store_site(self, site, normalisation_factor, value_transitions): - assert all(u >= 0 for u, _ in value_transitions) - self.normalisation_factor[site] = normalisation_factor - self.value_transitions[site] = value_transitions - - # Expose the same API as the low-level classes - - @property - def num_transitions(self): - a = [len(self.value_transitions[j]) for j in range(self.num_sites)] - return np.array(a, dtype=np.int32) - - def get_site(self, site): - return self.value_transitions[site] - - def decode(self): - """ - Decodes the tree encoding of the values into an explicit - matrix. - """ - sample_index_map = np.zeros(self.ts.num_nodes, dtype=int) - 1 - sample_index_map[self.ts.samples()] = np.arange(self.ts.num_samples) - A = np.zeros((self.num_sites, self.num_samples)) - for tree in self.ts.trees(): - for site in tree.sites(): - for node, value in self.value_transitions[site.id]: - for u in tree.samples(node): - j = sample_index_map[u] - A[site.id, j] = value - return A - - -class ViterbiMatrix(CompressedMatrix): - """ - Class representing the compressed Viterbi matrix. - """ - - def __init__(self, ts): - super().__init__(ts) - # Tuple containing the site, the node in the tree, and whether - # recombination is required - self.recombination_required = [(-1, 0, False)] - - def add_recombination_required(self, site, node, required): - self.recombination_required.append((site, node, required)) - - def choose_sample(self, site_id, tree): - max_value = -1 - u = -1 - for node, value in self.value_transitions[site_id]: - if value > max_value: - max_value = value - u = node - assert u != -1 - - transition_nodes = [u for (u, _) in self.value_transitions[site_id]] - while not tree.is_sample(u): - for v in tree.children(u): - if v not in transition_nodes: - u = v - break - else: - raise AssertionError("could not find path") - return u - - def traceback(self): - # Run the traceback. - m = self.ts.num_sites - match = np.zeros(m, dtype=int) - recombination_tree = np.zeros(self.ts.num_nodes, dtype=int) - 1 - tree = tskit.Tree(self.ts) - tree.last() - current_node = -1 - - rr_index = len(self.recombination_required) - 1 - for site in reversed(self.ts.sites()): - while tree.interval.left > site.position: - tree.prev() - assert tree.interval.left <= site.position < tree.interval.right - - # Fill in the recombination tree - j = rr_index - while self.recombination_required[j][0] == site.id: - u, required = self.recombination_required[j][1:] - recombination_tree[u] = required - j -= 1 - - if current_node == -1: - current_node = self.choose_sample(site.id, tree) - match[site.id] = current_node - - # Now traverse up the tree from the current node. The first marked node - # we meet tells us whether we need to recombine. - u = current_node - while u != -1 and recombination_tree[u] == -1: - u = tree.parent(u) - - assert u != -1 - if recombination_tree[u] == 1: - # Need to switch at the next site. - current_node = -1 - # Reset the nodes in the recombination tree. - j = rr_index - while self.recombination_required[j][0] == site.id: - u, required = self.recombination_required[j][1:] - recombination_tree[u] = -1 - j -= 1 - rr_index = j - - return match - - -def get_site_alleles(ts, h, alleles): - if alleles is None: - n_alleles = np.zeros(ts.num_sites, dtype=np.int8) - 1 - for j in range(ts.num_sites): - uniq_alleles = np.unique(np.append(ts.genotype_matrix()[j, :], h[j])) - uniq_alleles = uniq_alleles[uniq_alleles != MISSING] - n_alleles[j] = len(uniq_alleles) - assert np.all(n_alleles > 0) - alleles = tskit.ALLELES_ACGT - if len(set(alleles).intersection(next(ts.variants()).alleles)) == 0: - alleles = tskit.ALLELES_01 - if len(set(alleles).intersection(next(ts.variants()).alleles)) == 0: - raise ValueError( - """Alleles list could not be identified. - Please pass a list of lists of alleles of length m, - or a list of alleles (e.g. tskit.ALLELES_ACGT)""" - ) - alleles = [alleles for _ in range(ts.num_sites)] - else: - alleles, n_alleles = check_alleles(alleles, ts.num_sites) - return alleles, n_alleles - - -def ls_forward_tree( - h, ts, rho, mu, precision=30, alleles=None, scale_mutation_based_on_n_alleles=False -): - alleles, n_alleles = get_site_alleles(ts, h, alleles) - fa = ForwardAlgorithm( - ts, - rho, - mu, - alleles, - n_alleles, - precision=precision, - scale_mutation=scale_mutation_based_on_n_alleles, - ) - return fa.run(h) - - -def ls_backward_tree(h, ts, rho, mu, normalisation_factor, precision=30, alleles=None): - alleles, n_alleles = get_site_alleles(ts, h, alleles) - ba = BackwardAlgorithm( - ts, - rho, - mu, - alleles, - n_alleles, - precision=precision, - ) - return ba.run(h, normalisation_factor) - - -def ls_viterbi_tree( - h, ts, rho, mu, precision=30, alleles=None, scale_mutation_based_on_n_alleles=False -): - alleles, n_alleles = get_site_alleles(ts, h, alleles) - va = ViterbiAlgorithm( - ts, - rho, - mu, - alleles, - n_alleles, - precision=precision, - scale_mutation=scale_mutation_based_on_n_alleles, - ) - return va.run(h) - - -class LSBase: - """Superclass of Li and Stephens tests.""" - - def example_haplotypes(self, ts): - H = ts.genotype_matrix() - s = H[:, 0].reshape(1, H.shape[0]) - H = H[:, 1:] - - haplotypes = [ - s, - H[:, -1].reshape(1, H.shape[0]), - ] - s_tmp = s.copy() - s_tmp[0, -1] = MISSING - haplotypes.append(s_tmp) - s_tmp = s.copy() - s_tmp[0, ts.num_sites // 2] = MISSING - haplotypes.append(s_tmp) - s_tmp = s.copy() - s_tmp[0, :] = MISSING - haplotypes.append(s_tmp) - - return H, haplotypes - - def example_parameters_haplotypes(self, ts, seed=42): - """Returns an iterator over combinations of haplotype, - recombination and mutation rates.""" - np.random.seed(seed) - H, haplotypes = self.example_haplotypes(ts) - n = H.shape[1] - m = ts.get_num_sites() - - # Here we have equal mutation and recombination - r = np.zeros(m) + 0.01 - mu = np.zeros(m) + 0.01 - r[0] = 0 - - for s in haplotypes: - yield n, H, s, r, mu - - # FIXME removing these as tests are abominably slow. - # We'll be refactoring all this to use pytest anyway, so let's not - # worry too much about coverage for now. - # # Mixture of random and extremes - # rs = [np.zeros(m) + 0.999, np.zeros(m) + 1e-6, np.random.rand(m)] - # mus = [np.zeros(m) + 0.33, np.zeros(m) + 1e-6, np.random.rand(m) * 0.33] - - # import itertools - # for s, r, mu in itertools.product(haplotypes, rs, mus): - # r[0] = 0 - # yield n, H, s, r, mu - - def assertAllClose(self, A, B): - """Assert that all entries of two matrices are 'close'""" - assert np.allclose(A, B, rtol=1e-5, atol=1e-8) - - # Define a bunch of very small tree-sequences for testing a collection - # of parameters on - def test_simple_n_10_no_recombination(self): - ts = msprime.simulate( - 10, recombination_rate=0, mutation_rate=0.5, random_seed=42 - ) - assert ts.num_sites > 3 - self.verify(ts) - - def test_simple_n_10_no_recombination_high_mut(self): - ts = msprime.simulate(10, recombination_rate=0, mutation_rate=3, random_seed=42) - assert ts.num_sites > 3 - self.verify(ts) - - def test_simple_n_10_no_recombination_higher_mut(self): - ts = msprime.simulate(20, recombination_rate=0, mutation_rate=3, random_seed=42) - assert ts.num_sites > 3 - self.verify(ts) - - def test_simple_n_6(self): - ts = msprime.simulate(6, recombination_rate=2, mutation_rate=7, random_seed=42) - assert ts.num_sites > 5 - self.verify(ts) - - def test_simple_n_8(self): - ts = msprime.simulate(8, recombination_rate=2, mutation_rate=5, random_seed=42) - assert ts.num_sites > 5 - self.verify(ts) - - def test_simple_n_8_high_recombination(self): - ts = msprime.simulate(8, recombination_rate=20, mutation_rate=5, random_seed=42) - assert ts.num_trees > 15 - assert ts.num_sites > 5 - self.verify(ts) - - def test_simple_n_16(self): - ts = msprime.simulate(16, recombination_rate=2, mutation_rate=5, random_seed=42) - assert ts.num_sites > 5 - self.verify(ts) - - # # Define a bunch of very small tree-sequences for testing a collection - # # of parameters on - # def test_simple_n_10_no_recombination_blah(self): - # ts = msprime.sim_ancestry( - # samples=10, - # recombination_rate=0, - # random_seed=42, - # sequence_length=10, - # population_size=10000, - # ) - # ts = msprime.sim_mutations(ts, rate=1e-5, random_seed=42) - # assert ts.num_sites > 3 - # self.verify(ts) - - # def test_simple_n_6_blah(self): - # ts = msprime.sim_ancestry( - # samples=6, - # recombination_rate=1e-4, - # random_seed=42, - # sequence_length=40, - # population_size=10000, - # ) - # ts = msprime.sim_mutations(ts, rate=1e-3, random_seed=42) - # assert ts.num_sites > 5 - # self.verify(ts) - - # def test_simple_n_8_blah(self): - # ts = msprime.sim_ancestry( - # samples=8, - # recombination_rate=1e-4, - # random_seed=42, - # sequence_length=20, - # population_size=10000, - # ) - # ts = msprime.sim_mutations(ts, rate=1e-4, random_seed=42) - # assert ts.num_sites > 5 - # assert ts.num_trees > 15 - # self.verify(ts) - - # def test_simple_n_16_blah(self): - # ts = msprime.sim_ancestry( - # samples=16, - # recombination_rate=1e-2, - # random_seed=42, - # sequence_length=20, - # population_size=10000, - # ) - # ts = msprime.sim_mutations(ts, rate=1e-4, random_seed=42) - # assert ts.num_sites > 5 - # self.verify(ts) - - def verify(self, ts): - raise NotImplementedError() - - -class FBAlgorithmBase(LSBase): - """Base for forwards backwards algorithm tests.""" - - -class VitAlgorithmBase(LSBase): - """Base for viterbi algoritm tests.""" - - -class TestMirroringHap(FBAlgorithmBase): - """Tests that mirroring the tree sequence and running forwards and backwards - algorithms gives the same log-likelihood of observing the data.""" - - def verify(self, ts): - for n, H, s, r, mu in self.example_parameters_haplotypes(ts): - # Note, need to remove the first sample from the ts, and ensure that - # invariant sites aren't removed. - ts_check = ts.simplify(range(1, n + 1), filter_sites=False) - cm = ls_forward_tree(s[0, :], ts_check, r, mu) - ll_tree = np.sum(np.log10(cm.normalisation_factor)) - - ts_check_mirror = mirror_coordinates(ts_check) - r_flip = np.insert(np.flip(r)[:-1], 0, 0) - cm_mirror = ls_forward_tree( - np.flip(s[0, :]), ts_check_mirror, r_flip, np.flip(mu) - ) - ll_mirror_tree = np.sum(np.log10(cm_mirror.normalisation_factor)) - self.assertAllClose(ll_tree, ll_mirror_tree) - - # Ensure that the decoded matrices are the same - flipped_H = np.flip(H, axis=0) - flipped_s = np.flip(s, axis=1) - F_mirror_matrix, c, ll = ls.forwards( - reference_panel=flipped_H, - query=flipped_s, - ploidy=1, - prob_recombination=r_flip, - prob_mutation=np.flip(mu), - scale_mutation_rate=False, - ) - - self.assertAllClose(F_mirror_matrix, cm_mirror.decode()) - self.assertAllClose(ll, ll_tree) - - -class TestForwardHapTree(FBAlgorithmBase): - """Tests that the tree algorithm computes the same forward matrix as the - simple method.""" - - def verify(self, ts): - for n, H, s, r, mu in self.example_parameters_haplotypes(ts): - for scale_mutation in [False, True]: - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - # Warning from lshmm: - # Passed a vector of mutation rates, but rescaling each mutation - # rate conditional on the number of alleles - F, c, ll = ls.forwards( - reference_panel=H, - query=s, - ploidy=1, - prob_recombination=r, - prob_mutation=mu, - scale_mutation_rate=scale_mutation, - ) - # Note, need to remove the first sample from the ts, and ensure - # that invariant sites aren't removed. - ts_check = ts.simplify(range(1, n + 1), filter_sites=False) - cm = ls_forward_tree( - s[0, :], - ts_check, - r, - mu, - scale_mutation_based_on_n_alleles=scale_mutation, - ) - self.assertAllClose(cm.decode(), F) - ll_tree = np.sum(np.log10(cm.normalisation_factor)) - self.assertAllClose(ll, ll_tree) - - -class TestForwardBackwardTree(FBAlgorithmBase): - """Tests that the tree algorithm computes the same forward matrix as the - simple method.""" - - def verify(self, ts): - for n, H, s, r, mu in self.example_parameters_haplotypes(ts): - F, c, ll = ls.forwards( - reference_panel=H, - query=s, - ploidy=1, - prob_recombination=r, - prob_mutation=mu, - scale_mutation_rate=False, - ) - B = ls.backwards( - reference_panel=H, - query=s, - ploidy=1, - normalisation_factor_from_forward=c, - prob_recombination=r, - prob_mutation=mu, - scale_mutation_rate=False, - ) - - # Note, need to remove the first sample from the ts, and ensure that - # invariant sites aren't removed. - ts_check = ts.simplify(range(1, n + 1), filter_sites=False) - c_f = ls_forward_tree(s[0, :], ts_check, r, mu) - ll_tree = np.sum(np.log10(c_f.normalisation_factor)) - - c_b = ls_backward_tree( - s[0, :], - ts_check, - r, - mu, - c_f.normalisation_factor, - ) - B_tree = c_b.decode() - - F_tree = c_f.decode() - - self.assertAllClose(B, B_tree) - self.assertAllClose(F, F_tree) - self.assertAllClose(ll, ll_tree) - - -class TestTreeViterbiHap(VitAlgorithmBase): - """Test that we have the same log-likelihood between tree and matrix - implementations""" - - def verify(self, ts): - for n, H, s, r, mu in self.example_parameters_haplotypes(ts): - path, ll = ls.viterbi( - reference_panel=H, - query=s, - ploidy=1, - prob_recombination=r, - prob_mutation=mu, - scale_mutation_rate=False, - ) - ts_check = ts.simplify(range(1, n + 1), filter_sites=False) - cm = ls_viterbi_tree(s[0, :], ts_check, r, mu) - ll_tree = np.sum(np.log10(cm.normalisation_factor)) - self.assertAllClose(ll, ll_tree) - - # Now, need to ensure that the likelihood of the preferred path is - # the same as ll_tree (and ll). - path_tree = cm.traceback() - ll_check = ls.path_loglik( - reference_panel=H, - query=s, - ploidy=1, - path=path_tree, - prob_recombination=r, - prob_mutation=mu, - scale_mutation_rate=False, - ) - self.assertAllClose(ll, ll_check) - - -# TODO add params to run the various checks -def check_viterbi(ts, h, recombination=None, mutation=None): - h = np.array(h).astype(np.int8) - m = ts.num_sites - assert len(h) == m - if recombination is None: - recombination = np.zeros(ts.num_sites) + 1e-9 - if mutation is None: - mutation = np.zeros(ts.num_sites) - precision = 22 - - G = ts.genotype_matrix() - s = h.reshape(1, m) - - path, ll = ls.viterbi( - reference_panel=G, - query=s, - ploidy=1, - prob_recombination=recombination, - prob_mutation=mutation, - scale_mutation_rate=False, - ) - assert np.isscalar(ll) - - cm = ls_viterbi_tree(h, ts, rho=recombination, mu=mutation) - ll_tree = np.sum(np.log10(cm.normalisation_factor)) - assert np.isscalar(ll_tree) - nt.assert_allclose(ll_tree, ll) - - # Check that the likelihood of the preferred path is - # the same as ll_tree (and ll). - path_tree = cm.traceback() - ll_check = ls.path_loglik( - reference_panel=G, - query=s, - ploidy=1, - path=path_tree, - prob_recombination=recombination, - prob_mutation=mutation, - scale_mutation_rate=False, - ) - nt.assert_allclose(ll_check, ll) - - ll_ts = ts._ll_tree_sequence - ls_hmm = _tskit.LsHmm(ll_ts, recombination, mutation, precision=precision) - cm_lib = _tskit.ViterbiMatrix(ll_ts) - ls_hmm.viterbi_matrix(h, cm_lib) - path_lib = cm_lib.traceback() - - # Not true in general, but let's see how far it goes - nt.assert_array_equal(path_lib, path_tree) - - nt.assert_allclose(cm_lib.normalisation_factor, cm.normalisation_factor) - - return path - - -# TODO add params to run the various checks -def check_forward_matrix(ts, h, recombination=None, mutation=None): - precision = 22 - h = np.array(h).astype(np.int8) - n = ts.num_samples - m = ts.num_sites - assert len(h) == m - if recombination is None: - recombination = np.zeros(ts.num_sites) + 1e-9 - if mutation is None: - mutation = np.zeros(ts.num_sites) - - G = ts.genotype_matrix() - s = h.reshape(1, m) - - F, c, ll = ls.forwards( - reference_panel=G, - query=s, - ploidy=1, - prob_recombination=recombination, - prob_mutation=mutation, - scale_mutation_rate=False, - ) - assert F.shape == (m, n) - assert c.shape == (m,) - assert np.isscalar(ll) - - cm = ls_forward_tree( - h, ts, recombination, mutation, scale_mutation_based_on_n_alleles=False - ) - F2 = cm.decode() - nt.assert_allclose(F, F2) - nt.assert_allclose(c, cm.normalisation_factor) - ll_tree = np.sum(np.log10(cm.normalisation_factor)) - nt.assert_allclose(ll_tree, ll) - - ll_ts = ts._ll_tree_sequence - ls_hmm = _tskit.LsHmm(ll_ts, recombination, mutation, precision=precision) - cm_lib = _tskit.CompressedMatrix(ll_ts) - ls_hmm.forward_matrix(h, cm_lib) - F3 = cm_lib.decode() - - assert_compressed_matrices_equal(cm, cm_lib) - - nt.assert_allclose(F, F3) - nt.assert_allclose(c, cm_lib.normalisation_factor) - return cm_lib - - -def check_backward_matrix(ts, h, forward_cm, recombination=None, mutation=None): - precision = 22 - h = np.array(h).astype(np.int8) - m = ts.num_sites - assert len(h) == m - if recombination is None: - recombination = np.zeros(ts.num_sites) + 1e-9 - if mutation is None: - mutation = np.zeros(ts.num_sites) - - G = ts.genotype_matrix() - s = h.reshape(1, m) - - B = ls.backwards( - reference_panel=G, - query=s, - ploidy=1, - normalisation_factor_from_forward=forward_cm.normalisation_factor, - prob_recombination=recombination, - prob_mutation=mutation, - scale_mutation_rate=False, - ) - - backward_cm = ls_backward_tree( - h, - ts, - recombination, - mutation, - forward_cm.normalisation_factor, - precision=precision, - ) - nt.assert_array_equal( - backward_cm.normalisation_factor, forward_cm.normalisation_factor - ) - - ll_ts = ts._ll_tree_sequence - ls_hmm = _tskit.LsHmm(ll_ts, recombination, mutation, precision=precision) - cm_lib = _tskit.CompressedMatrix(ll_ts) - ls_hmm.backward_matrix(h, forward_cm.normalisation_factor, cm_lib) - - assert_compressed_matrices_equal(backward_cm, cm_lib) - - B_lib = cm_lib.decode() - B_tree = backward_cm.decode() - nt.assert_allclose(B_tree, B_lib) - nt.assert_allclose(B, B_lib) - - -def add_unique_sample_mutations(ts, start=0): - """ - Adds a mutation for each of the samples at equally spaced locations - along the genome. - """ - tables = ts.dump_tables() - L = int(ts.sequence_length) - assert L % ts.num_samples == 0 - gap = L // ts.num_samples - x = start - for u in ts.samples(): - site = tables.sites.add_row(position=x, ancestral_state="0") - tables.mutations.add_row(site=site, derived_state="1", node=u) - x += gap - return tables.tree_sequence() - - -class TestSingleBalancedTreeExample: - # 3.00┊ 6 ┊ - # ┊ ┏━┻━┓ ┊ - # 2.00┊ 4 5 ┊ - # ┊ ┏┻┓ ┏┻┓ ┊ - # 1.00┊ 0 1 2 3 ┊ - # 0 8 - - @staticmethod - def ts(): - return add_unique_sample_mutations( - tskit.Tree.generate_balanced(4, span=8).tree_sequence, - start=1, - ) - - @pytest.mark.parametrize("j", [0, 1, 2, 3]) - def test_match_sample(self, j): - ts = self.ts() - h = np.zeros(4) - h[j] = 1 - path = check_viterbi(ts, h) - nt.assert_array_equal([j, j, j, j], path) - cm = check_forward_matrix(ts, h) - check_backward_matrix(ts, h, cm) - - @pytest.mark.parametrize("j", [1, 2]) - def test_match_sample_missing_flanks(self, j): - ts = self.ts() - h = np.zeros(4) - h[0] = -1 - h[-1] = -1 - h[j] = 1 - path = check_viterbi(ts, h) - nt.assert_array_equal([j, j, j, j], path) - cm = check_forward_matrix(ts, h) - check_backward_matrix(ts, h, cm) - - def test_switch_each_sample(self): - ts = self.ts() - h = np.ones(4) - path = check_viterbi(ts, h) - nt.assert_array_equal([0, 1, 2, 3], path) - cm = check_forward_matrix(ts, h) - check_backward_matrix(ts, h, cm) - - def test_switch_each_sample_missing_flanks(self): - ts = self.ts() - h = np.ones(4) - h[0] = -1 - h[-1] = -1 - path = check_viterbi(ts, h) - nt.assert_array_equal([1, 1, 2, 2], path) - cm = check_forward_matrix(ts, h) - check_backward_matrix(ts, h, cm) - - def test_switch_each_sample_missing_middle(self): - ts = self.ts() - h = np.ones(4) - h[1:3] = -1 - path = check_viterbi(ts, h) - # Implementation of Viterbi switches at right-most position - nt.assert_array_equal([0, 3, 3, 3], path) - cm = check_forward_matrix(ts, h) - check_backward_matrix(ts, h, cm) - - -class TestSimulationExamples: - @pytest.mark.slow - @pytest.mark.parametrize("n", [3, 10, 50]) - @pytest.mark.parametrize("L", [1, 10, 100]) - def test_continuous_genome(self, n, L): - ts = msprime.simulate( - n, length=L, recombination_rate=1, mutation_rate=1, random_seed=42 - ) - h = np.zeros(ts.num_sites, dtype=np.int8) - # NOTE this is a bit slow at the moment but we can disable the Python - # implementation once testing has been improved on smaller examples. - # Add ``compare_py=False``to these calls. - check_viterbi(ts, h) - cm = check_forward_matrix(ts, h) - check_backward_matrix(ts, h, cm) diff --git a/python/tests/test_python_c.py b/python/tests/test_python_c.py index 75d52d571b..5f18699d82 100644 --- a/python/tests/test_python_c.py +++ b/python/tests/test_python_c.py @@ -3846,184 +3846,6 @@ def test_r2_array_lifetime(self): assert a3 is not a2 -class TestLsHmm(LowLevelTestCase): - """ - Tests for the LsHmm class. - """ - - def test_uninitialised_tree_sequence(self): - ts = _tskit.TreeSequence() - with pytest.raises(ValueError): - _tskit.LsHmm(ts, None, None) - - def test_constructor(self): - ts = self.get_example_tree_sequence() - with pytest.raises(TypeError): - _tskit.LsHmm() - with pytest.raises(TypeError): - _tskit.LsHmm(None) - values = np.zeros(ts.get_num_sites()) - for bad_array in ["asdf", [[], []], None]: - with pytest.raises(ValueError): - _tskit.LsHmm(ts, bad_array, values) - with pytest.raises(ValueError): - _tskit.LsHmm(ts, values, bad_array) - - def test_bad_rate_arrays(self): - ts = self.get_example_tree_sequence() - m = ts.get_num_sites() - assert m > 0 - values = np.zeros(m) - for bad_size in [0, m - 1, m + 1, m + 2]: - bad_array = np.zeros(bad_size) - with pytest.raises(ValueError): - _tskit.LsHmm(ts, bad_array, values) - with pytest.raises(ValueError): - _tskit.LsHmm(ts, values, bad_array) - - def test_haplotype_input(self): - ts = self.get_example_tree_sequence() - m = ts.get_num_sites() - fm = _tskit.CompressedMatrix(ts) - vm = _tskit.ViterbiMatrix(ts) - norm = np.ones(m) - ls_hmm = _tskit.LsHmm(ts, np.zeros(m), np.zeros(m)) - for bad_size in [0, m - 1, m + 1, m + 2]: - bad_array = np.zeros(bad_size, dtype=np.int8) - with pytest.raises(ValueError, match="haplotype array"): - ls_hmm.forward_matrix(bad_array, fm) - with pytest.raises(ValueError, match="haplotype array"): - ls_hmm.backward_matrix(bad_array, norm, fm) - with pytest.raises(ValueError, match="haplotype array"): - ls_hmm.viterbi_matrix(bad_array, vm) - for bad_array in [[0.002], [[], []], None]: - with pytest.raises(ValueError): - ls_hmm.forward_matrix(bad_array, fm) - with pytest.raises(ValueError): - ls_hmm.viterbi_matrix(bad_array, vm) - with pytest.raises(ValueError): - ls_hmm.backward_matrix(bad_array, norm, fm) - - def test_norm_input(self): - ts = self.get_example_tree_sequence() - m = ts.get_num_sites() - cm = _tskit.CompressedMatrix(ts) - h = np.zeros(m, dtype=np.int32) - ls_hmm = _tskit.LsHmm(ts, np.zeros(m), np.zeros(m)) - for bad_size in [0, m - 1, m + 1, m + 2]: - bad_array = np.zeros(bad_size) - with pytest.raises(ValueError, match="forward_norm array"): - ls_hmm.backward_matrix(h, bad_array, cm) - - for bad_array in [[0.002], [[], []], None]: - with pytest.raises(ValueError): - ls_hmm.backward_matrix(h, bad_array, cm) - - def test_output_type_errors(self): - ts = self.get_example_tree_sequence() - m = ts.get_num_sites() - h = np.zeros(m, dtype=np.int8) - norm = np.ones(m) - ls_hmm = _tskit.LsHmm(ts, np.zeros(m), np.zeros(m)) - for bad_type in [ls_hmm, None, m, []]: - with pytest.raises(TypeError): - ls_hmm.forward_matrix(h, bad_type) - with pytest.raises(TypeError): - ls_hmm.viterbi_matrix(h, bad_type) - with pytest.raises(TypeError): - ls_hmm.backward_matrix(h, norm, bad_type) - - other_ts = self.get_example_tree_sequence() - output = _tskit.CompressedMatrix(other_ts) - with pytest.raises(_tskit.LibraryError): - ls_hmm.forward_matrix(h, output) - with pytest.raises(_tskit.LibraryError): - ls_hmm.backward_matrix(h, norm, output) - output = _tskit.ViterbiMatrix(other_ts) - with pytest.raises(_tskit.LibraryError): - ls_hmm.viterbi_matrix(h, output) - - def test_empty_forward_matrix(self): - for mu in [0, 1]: - ts = self.get_example_tree_sequence(mutation_rate=mu) - m = ts.get_num_sites() - fm = _tskit.CompressedMatrix(ts) - assert fm.num_sites == m - assert np.array_equal(np.zeros(m), fm.normalisation_factor) - assert np.array_equal(np.zeros(m, dtype=np.uint32), fm.num_transitions) - F = fm.decode() - assert np.all(F >= 0) - for j in range(m): - assert fm.get_site(j) == [] - - def test_empty_viterbi_matrix(self): - for mu in [0, 1]: - ts = self.get_example_tree_sequence(mutation_rate=mu) - m = ts.get_num_sites() - vm = _tskit.ViterbiMatrix(ts) - assert vm.num_sites == m - # TODO we should have the same semantics for 0 sites - if m == 0: - h = vm.traceback() - assert len(h) == 0 - else: - with pytest.raises(_tskit.LibraryError): - vm.traceback() - - def verify_compressed_matrix(self, ts, output): - S = output.normalisation_factor - N = output.num_transitions - assert np.all(0 < S) - assert np.all(S < 1) - assert np.all(N > 0) - F = output.decode() - assert F.shape == (ts.get_num_sites(), ts.get_num_samples()) - assert np.all(F >= 0) - m = ts.get_num_sites() - for j in range(m): - site_list = output.get_site(j) - assert len(site_list) == N[j] - for item in site_list: - assert len(item) == 2 - node, value = item - assert 0 <= node < ts.get_num_nodes() - assert value >= 0 - for site in [m, m + 1, 2 * m]: - with pytest.raises(ValueError): - output.get_site(site) - - def test_forward_matrix(self): - ts = self.get_example_tree_sequence() - m = ts.get_num_sites() - output = _tskit.CompressedMatrix(ts) - ls_hmm = _tskit.LsHmm(ts, np.zeros(m) + 0.1, np.zeros(m) + 0.1) - rv = ls_hmm.forward_matrix([0 for _ in range(m)], output) - assert rv is None - self.verify_compressed_matrix(ts, output) - - def test_backward_matrix(self): - ts = self.get_example_tree_sequence() - m = ts.get_num_sites() - fm = _tskit.CompressedMatrix(ts) - bm = _tskit.CompressedMatrix(ts) - h = np.zeros(m, dtype=np.int32) - ls_hmm = _tskit.LsHmm(ts, np.zeros(m) + 0.1, np.zeros(m) + 0.1) - ls_hmm.forward_matrix(h, fm) - ls_hmm.backward_matrix(h, fm.normalisation_factor, bm) - self.verify_compressed_matrix(ts, bm) - - def test_viterbi_matrix(self): - ts = self.get_example_tree_sequence() - m = ts.get_num_sites() - output = _tskit.ViterbiMatrix(ts) - ls_hmm = _tskit.LsHmm(ts, np.zeros(m) + 0.1, np.zeros(m) + 0.1) - rv = ls_hmm.viterbi_matrix([0 for _ in range(m)], output) - assert rv is None - self.verify_compressed_matrix(ts, output) - h = output.traceback() - assert isinstance(h, np.ndarray) - - class TestTree(LowLevelTestCase): """ Tests on the low-level tree interface. diff --git a/python/uv.lock b/python/uv.lock index f13b16de61..4654830091 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -9,10 +9,10 @@ resolution-markers = [ "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform == 'emscripten'", "python_full_version >= '3.12' and python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'win32'", - "python_full_version < '3.11.5' and sys_platform == 'win32'", "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'emscripten'", - "python_full_version < '3.11.5' and sys_platform == 'emscripten'", "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version < '3.11.5' and sys_platform == 'win32'", + "python_full_version < '3.11.5' and sys_platform == 'emscripten'", "python_full_version < '3.11.5' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] @@ -1287,18 +1287,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/93/73/72553170eada174775d9a738c471c7be4ab3dc2c06368beeee89e002345c/llvmlite-0.50.0-cp315-cp315t-win_amd64.whl", hash = "sha256:4da0e8c6e6f144b433672a632f75d6b4da7bd4fdb5c3e9981d6ea6741319aeae", size = 42986722, upload-time = "2026-09-29T18:44:44.491Z" }, ] -[[package]] -name = "lshmm" -version = "0.0.8" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "numba" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/bc/d1/627b2dd2572e4ccc615b015121b89993fd9b57948ad5c8b185465f59a856/lshmm-0.0.8.tar.gz", hash = "sha256:0d24eb5ad0a79e17532b85e9971d684ec8895828ca26b450cf63f1627ba7645f", size = 26783, upload-time = "2024-06-24T10:46:00.145Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ef/03/6f0eb918cf7f5367b06554847f3174373e3e8d5582b2b2049172c8aacb24/lshmm-0.0.8-py3-none-any.whl", hash = "sha256:74bac9e8abb71dfa7d9b26bc88bca09397d92e3baa045b134ddf3d814c469bc3", size = 17325, upload-time = "2024-06-24T10:45:58.462Z" }, -] - [[package]] name = "lxml" version = "6.1.3" @@ -1894,10 +1882,10 @@ version = "3.6.1" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'win32'", - "python_full_version < '3.11.5' and sys_platform == 'win32'", "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'emscripten'", - "python_full_version < '3.11.5' and sys_platform == 'emscripten'", "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version < '3.11.5' and sys_platform == 'win32'", + "python_full_version < '3.11.5' and sys_platform == 'emscripten'", "python_full_version < '3.11.5' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" } @@ -2011,10 +1999,10 @@ version = "0.16.6" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'win32'", - "python_full_version < '3.11.5' and sys_platform == 'win32'", "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'emscripten'", - "python_full_version < '3.11.5' and sys_platform == 'emscripten'", "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version < '3.11.5' and sys_platform == 'win32'", + "python_full_version < '3.11.5' and sys_platform == 'emscripten'", "python_full_version < '3.11.5' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ @@ -2091,10 +2079,10 @@ version = "2.4.6" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'win32'", - "python_full_version < '3.11.5' and sys_platform == 'win32'", "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'emscripten'", - "python_full_version < '3.11.5' and sys_platform == 'emscripten'", "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version < '3.11.5' and sys_platform == 'win32'", + "python_full_version < '3.11.5' and sys_platform == 'emscripten'", "python_full_version < '3.11.5' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] sdist = { url = "https://files.pythonhosted.org/packages/d0/ad/fed0499ce6a338d2a03ebae59cd15093910c8875328855781952abf6c2fe/numpy-2.4.6.tar.gz", hash = "sha256:f3a3570c4a2a16746ac2c31a7c7c7b0c186b95ce902e33db6f28094ed7387dda", size = 20735807, upload-time = "2026-05-18T23:37:14.07Z" } @@ -3499,7 +3487,6 @@ dev = [ { name = "dendropy" }, { name = "jupyter-book" }, { name = "kastore" }, - { name = "lshmm" }, { name = "msgpack" }, { name = "msprime" }, { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, @@ -3552,7 +3539,6 @@ test = [ { name = "coverage" }, { name = "dendropy" }, { name = "kastore" }, - { name = "lshmm" }, { name = "msgpack" }, { name = "msprime" }, { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, @@ -3593,7 +3579,6 @@ dev = [ { name = "dendropy" }, { name = "jupyter-book", specifier = "<2" }, { name = "kastore" }, - { name = "lshmm" }, { name = "msgpack" }, { name = "msprime" }, { name = "networkx" }, @@ -3644,7 +3629,6 @@ test = [ { name = "coverage" }, { name = "dendropy" }, { name = "kastore" }, - { name = "lshmm" }, { name = "msgpack" }, { name = "msprime" }, { name = "networkx" }, @@ -3814,10 +3798,10 @@ version = "3.1.6" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'win32'", - "python_full_version < '3.11.5' and sys_platform == 'win32'", "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform == 'emscripten'", - "python_full_version < '3.11.5' and sys_platform == 'emscripten'", "python_full_version >= '3.11.5' and python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version < '3.11.5' and sys_platform == 'win32'", + "python_full_version < '3.11.5' and sys_platform == 'emscripten'", "python_full_version < '3.11.5' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [