Skip to content

Commit

Permalink
Update tests, remove unused window code
Browse files Browse the repository at this point in the history
  • Loading branch information
@lkirk
lkirk committed Sep 4, 2023
1 parent 5fdf850 commit 2991d92
Show file tree
Hide file tree
Showing 2 changed files with 133 additions and 44 deletions.
110 changes: 110 additions & 0 deletions c/tests/test_stats.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2305,6 +2305,114 @@ test_two_site_uncorrelated_multiallelic(void)
tsk_treeseq_free(&ts);
}

static void
test_two_site_backmutation(void)
{
const char *nodes
= "1 0 -1\n1 0 -1\n1 0 -1\n1 0 -1\n1 0 -1\n1 0 -1\n1 0 -1\n1 0 -1\n"
"1 0 -1\n1 0 -1\n1 0 -1\n1 0 -1\n1 0 -1\n1 0 -1\n1 0 -1\n1 0 -1\n"
"1 0 -1\n1 0 -1\n1 0 -1\n1 0 -1\n1 0 -1\n1 0 -1\n1 0 -1\n1 0 -1\n"
"1 0 -1\n1 0 -1\n1 0 -1\n1 0 -1\n1 0 -1\n1 0 -1\n1 0 -1\n1 0 -1\n"
"1 0 -1\n1 0 -1\n1 0 -1\n0 2 -1\n0 4 -1\n0 6 -1\n0 8 -1\n0 10 -1\n"
"0 12 -1\n0 14 -1\n0 16 -1\n0 18 -1\n0 20 -1\n0 22 -1\n0 24 -1\n0 26 -1\n"
"0 28 -1\n0 30 -1\n0 32 -1\n0 34 -1\n0 36 -1\n0 38 -1\n0 40 -1\n0 42 -1\n"
"0 44 -1\n0 46 -1\n0 48 -1\n0 50 -1\n0 52 -1\n0 54 -1\n0 56 -1\n0 58 -1\n"
"0 60 -1\n0 62 -1\n0 64 -1\n0 66 -1\n0 68 -1\n";

const char *edges
= "0 10 35 0,1\n0 10 36 2,35\n0 10 37 3,36\n0 10 38 4,37\n0 10 39 5,38\n"
"0 10 40 6,39\n0 10 41 7,40\n0 10 42 8,41\n0 10 43 9,42\n0 10 44 10,43\n"
"0 10 45 11,44\n0 10 46 12,45\n0 10 47 13,46\n0 10 48 14,47\n0 10 49 15,48\n"
"0 10 50 16,49\n0 10 51 17,50\n0 10 52 18,51\n0 10 53 19,52\n0 10 54 20,53\n"
"0 10 55 21,54\n0 10 56 22,55\n0 10 57 23,56\n0 10 58 24,57\n0 10 59 25,58\n"
"0 10 60 26,59\n0 10 61 27,60\n0 10 62 28,61\n0 10 63 29,62\n0 10 64 30,63\n"
"0 10 65 31,64\n0 10 66 32,65\n0 10 67 33,66\n0 10 68 34,67\n";

const char *sites = "1 A\n"
"4.5 T\n";

const char *mutations = "0 50 T -1\n"
"0 48 G 0\n"
"0 46 A 1\n"
"1 62 G -1\n"
"1 60 T 3\n"
"1 58 A 4\n";

int ret;
double *result;
tsk_size_t result_size;

tsk_treeseq_t ts;
tsk_treeseq_from_text(&ts, 10, nodes, edges, NULL, sites, mutations, NULL, NULL, 0);

tsk_size_t sample_set_sizes[1] = { ts.num_samples };
tsk_size_t num_sample_sets = 1;
tsk_id_t sample_sets[ts.num_samples];

for (tsk_size_t s = 0; s < ts.num_samples; s++) {
sample_sets[s] = (tsk_id_t) s;
}

ret = tsk_treeseq_r2(&ts, num_sample_sets, sample_set_sizes, sample_sets, 0, NULL, 0,
NULL, 0, &result_size, &result);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL(result_size, 3);
/* assert_arrays_almost_equal(result_size, result, truth_r2); */
tsk_safe_free(result);
}

static void
test_two_locus_stat_input_errors(void)
{
tsk_treeseq_t ts;
double *result;
tsk_size_t s, result_size;
int ret;

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);

tsk_size_t sample_set_sizes[1];
tsk_size_t num_sample_sets;
tsk_id_t sample_sets[ts.num_samples];

sample_set_sizes[0] = ts.num_samples;
num_sample_sets = 1;
for (s = 0; s < ts.num_samples; s++) {
sample_sets[s] = (tsk_id_t) s;
}

sample_sets[1] = 0;
ret = tsk_treeseq_r2(&ts, num_sample_sets, sample_set_sizes, sample_sets, 0, NULL, 0,
NULL, 0, &result_size, &result);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_DUPLICATE_SAMPLE);
sample_sets[1] = 1;

ret = tsk_treeseq_r2(&ts, num_sample_sets, sample_set_sizes, sample_sets, 0, NULL, 0,
NULL, TSK_STAT_SITE | TSK_STAT_BRANCH, &result_size, &result);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_MULTIPLE_STAT_MODES);

ret = tsk_treeseq_r2(&ts, num_sample_sets, sample_set_sizes, sample_sets, 0, NULL, 0,
NULL, TSK_STAT_BRANCH, &result_size, &result);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_UNSUPPORTED_STAT_MODE);

ret = tsk_treeseq_r2(&ts, 0, sample_set_sizes, sample_sets, 0, NULL, 0, NULL, 0,
&result_size, &result);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_BAD_STATE_DIMS);

sample_set_sizes[0] = 0;
ret = tsk_treeseq_r2(&ts, num_sample_sets, sample_set_sizes, sample_sets, 0, NULL, 0,
NULL, 0, &result_size, &result);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_EMPTY_SAMPLE_SET);
sample_set_sizes[0] = ts.num_samples;

sample_sets[1] = 10;
ret = tsk_treeseq_r2(&ts, num_sample_sets, sample_set_sizes, sample_sets, 0, NULL, 0,
NULL, 0, &result_size, &result);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_NODE_OUT_OF_BOUNDS);
sample_sets[1] = 1;
}

static void
test_simplest_divergence_matrix(void)
{
Expand Down Expand Up @@ -2573,6 +2681,8 @@ main(int argc, char **argv)
test_two_site_correlated_multiallelic },
{ "test_two_site_uncorrelated_multiallelic",
test_two_site_uncorrelated_multiallelic },
{ "test_two_site_backmutation", test_two_site_backmutation },
{ "test_two_locus_stat_input_errors", test_two_locus_stat_input_errors },

{ "test_simplest_divergence_matrix", test_simplest_divergence_matrix },
{ "test_simplest_divergence_matrix_windows",
Expand Down
67 changes: 23 additions & 44 deletions c/tskit/trees.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2562,17 +2562,18 @@ static int
tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sample_sets,
const tsk_size_t *sample_set_sizes, const tsk_id_t *sample_sets,
tsk_size_t result_dim, const tsk_id_t *set_indexes, general_stat_func_t *f,
norm_func_t *norm_f, tsk_size_t num_left_windows, const double *left_windows,
tsk_size_t num_right_windows, const double *right_windows, tsk_flags_t options,
tsk_size_t *result_size, double **result)
norm_func_t *norm_f, tsk_size_t TSK_UNUSED(num_left_windows),
const double *left_windows, tsk_size_t TSK_UNUSED(num_right_windows),
const double *right_windows, tsk_flags_t options, tsk_size_t *result_size,
double **result)
{
// TODO: generalize this function if we ever decide to do weighted two_locus stats.
// We only implement count stats and therefore we don't handle weights.
int ret = 0;
tsk_bit_array_t sample_sets_bits;
bool stat_site = !!(options & TSK_STAT_SITE);
bool stat_branch = !!(options & TSK_STAT_BRANCH);
double default_windows[] = { 0, self->tables->sequence_length };
// double default_windows[] = { 0, self->tables->sequence_length };
tsk_size_t state_dim = num_sample_sets;
sample_count_stat_params_t f_params = { .sample_sets = sample_sets,
.num_sample_sets = num_sample_sets,
Expand All @@ -2581,66 +2582,44 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl

tsk_memset(&sample_sets_bits, 0, sizeof(sample_sets_bits));

ret = tsk_treeseq_check_sample_sets(
self, num_sample_sets, sample_set_sizes, sample_sets);
if (ret != 0) {
goto out;
}
ret = sample_sets_to_bit_array(
self, sample_set_sizes, sample_sets, num_sample_sets, &sample_sets_bits);
if (ret != 0) {
goto out;
}

/* If no mode is specified, we default to site mode */
// If no mode is specified, we default to site mode
if (!(stat_site || stat_branch)) {
stat_site = true;
}
/* It's an error to specify more than one mode */
// It's an error to specify more than one mode
if (stat_site + stat_branch > 1) {
ret = TSK_ERR_MULTIPLE_STAT_MODES;
goto out;
}

if (state_dim < 1) {
ret = TSK_ERR_BAD_STATE_DIMS;
goto out;
}
if (result_dim < 1) {
ret = TSK_ERR_BAD_RESULT_DIMS;
// TODO: impossible until we implement branch/windows
// if (result_dim < 1) {
// ret = TSK_ERR_BAD_RESULT_DIMS;
// goto out;
// }

tsk_bug_assert(left_windows == NULL && right_windows == NULL);

ret = tsk_treeseq_check_sample_sets(
self, num_sample_sets, sample_set_sizes, sample_sets);
if (ret != 0) {
goto out;
}
if (left_windows == NULL) {
num_left_windows = 1;
left_windows = default_windows;
num_right_windows = 1;
right_windows = default_windows;
} else if (right_windows == NULL) {
ret = tsk_treeseq_check_windows(self, num_left_windows, left_windows, 0);
if (ret != 0) {
goto out;
}
num_right_windows = num_left_windows;
right_windows = left_windows;
} else {
ret = tsk_treeseq_check_windows(self, num_left_windows, left_windows, 0);
if (ret != 0) {
goto out;
}
ret = tsk_treeseq_check_windows(self, num_right_windows, right_windows, 0);
if (ret != 0) {
goto out;
}
ret = sample_sets_to_bit_array(
self, sample_set_sizes, sample_sets, num_sample_sets, &sample_sets_bits);
if (ret != 0) {
goto out;
}

if (stat_site) {
// TODO: assert 1 window left/right
ret = tsk_treeseq_two_site_count_stat(self, state_dim, &sample_sets_bits,
result_dim, f, &f_params, norm_f, left_windows, right_windows, options,
result_size, result);
} else {
// TODO: branch statistics
ret = TSK_ERR_GENERIC;
ret = TSK_ERR_UNSUPPORTED_STAT_MODE;
}

out:
Expand Down

0 comments on commit 2991d92

Please sign in to comment.