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 = [