Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

time windows in statistics #2948

Draft
wants to merge 19 commits into
base: main
Choose a base branch
from
+960 −184
Draft

time windows in statistics #2948

Show file tree
Hide file tree
Changes from 16 commits
Commits
Show all changes
19 commits
Select commit Hold shift + click to select a range
aeda4dc
example test for time windows
petrelharp May 9, 2024
0a30696
a little more flailing at beginning tests
petrelharp May 17, 2024
1a988c0
first simple test for AFS in branch mode with time windows
tforest Jun 18, 2024
18ffda0
AFS non-naive version with time windows + Some tests
tforest Jul 15, 2024
c6f9562
Improve AFS branch mode
tforest Jul 18, 2024
7a3149f
fix AFS tests
tforest Jul 18, 2024
71da7ad
Fix naive branch general stat
tforest Jul 22, 2024
2a44909
intermediate changes to C AFS implementation with time windows
tforest Aug 23, 2024
4460db9
tw afs C implementation
tforest Sep 23, 2024
bbec6a9
Better dimension drop with time windows
tforest Oct 11, 2024
da5f205
AFS branch mode with time windows to review
tforest Oct 21, 2024
6b3ab4f
Adapting some tests with new time_windows parameter
tforest Oct 30, 2024
f2d857b
time_windows addition after code rebase
tforest Oct 30, 2024
b8f4ba5
Fix some tests failing because of time_windows
tforest Oct 31, 2024
c9f6c06
Remove unused parameters in afs function
tforest Oct 31, 2024
59ea266
Fixing a memory issue in tsk_treeseq_update_branch_afs
tforest Oct 31, 2024
96ac0ce
Fixing some time_windows issues with AFS
tforest Nov 19, 2024
26a7f09
Adjusting some stats tests using time_windows
tforest Nov 19, 2024
8a8c05b
Adjusting some tests trying to pass codecov
tforest Dec 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
75 changes: 65 additions & 10 deletions c/tests/test_stats.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -865,6 +865,61 @@ verify_one_way_stat_func_errors(tsk_treeseq_t *ts, one_way_sample_stat_method *m
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_BAD_WINDOWS);
}

// Temporary definition for time_windows in tsk_treeseq_allele_frequency_spectrum
typedef int one_way_sample_stat_method_tw(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 num_windows, const double *windows,
tsk_size_t num_time_windows, const double *time_windows,
tsk_flags_t options, double *result);

static void
verify_one_way_stat_func_errors_tw(tsk_treeseq_t *ts, one_way_sample_stat_method_tw *method)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
verify_one_way_stat_func_errors_tw(tsk_treeseq_t *ts, one_way_sample_stat_method_tw *method)
// Temporary duplicate for time-windows-having methods
verify_one_way_stat_func_errors_tw(tsk_treeseq_t *ts, one_way_sample_stat_method_tw *method)

{
int ret;
tsk_id_t num_nodes = (tsk_id_t) tsk_treeseq_get_num_nodes(ts);
tsk_id_t samples[] = { 0, 1, 2, 3 };
tsk_size_t sample_set_sizes = 4;
double windows[] = { 0, 0, 0 };
double result;

ret = method(ts, 0, &sample_set_sizes, samples, 0, NULL, 0, NULL, 0, &result);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_INSUFFICIENT_SAMPLE_SETS);

samples[0] = TSK_NULL;
ret = method(ts, 1, &sample_set_sizes, samples, 0, NULL, 0, NULL, 0, &result);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_NODE_OUT_OF_BOUNDS);
samples[0] = -10;
ret = method(ts, 1, &sample_set_sizes, samples, 0, NULL, 0, NULL, 0, &result);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_NODE_OUT_OF_BOUNDS);
samples[0] = num_nodes;
ret = method(ts, 1, &sample_set_sizes, samples, 0, NULL, 0, NULL, 0, &result);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_NODE_OUT_OF_BOUNDS);
samples[0] = num_nodes + 1;
ret = method(ts, 1, &sample_set_sizes, samples, 0, NULL, 0, NULL, 0, &result);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_NODE_OUT_OF_BOUNDS);

samples[0] = num_nodes - 1;
ret = method(ts, 1, &sample_set_sizes, samples, 0, NULL, 0, NULL, 0, &result);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_BAD_SAMPLES);

samples[0] = 1;
ret = method(ts, 1, &sample_set_sizes, samples, 0, NULL, 0, NULL, 0, &result);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_DUPLICATE_SAMPLE);

samples[0] = 0;
sample_set_sizes = 0;
ret = method(ts, 1, &sample_set_sizes, samples, 0, NULL, 0, NULL, 0, &result);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_EMPTY_SAMPLE_SET);

sample_set_sizes = 4;
/* Window errors */
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should also have time window errors here?

ret = method(ts, 1, &sample_set_sizes, samples, 0, windows, 0, NULL, 0, &result);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_BAD_NUM_WINDOWS);

ret = method(ts, 1, &sample_set_sizes, samples, 2, windows, 0, NULL, 0, &result);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_BAD_WINDOWS);
}

static void
verify_two_way_stat_func_errors(
tsk_treeseq_t *ts, general_sample_stat_method *method, tsk_flags_t options)
Expand Down Expand Up @@ -1203,23 +1258,23 @@ verify_afs(tsk_treeseq_t *ts)
sample_set_sizes[0] = n - 2;
sample_set_sizes[1] = 2;
ret = tsk_treeseq_allele_frequency_spectrum(
ts, 2, sample_set_sizes, samples, 0, NULL, 0, result);
ts, 2, sample_set_sizes, samples, 0, NULL, 0, NULL, 0, result);
CU_ASSERT_EQUAL_FATAL(ret, 0);

ret = tsk_treeseq_allele_frequency_spectrum(
ts, 2, sample_set_sizes, samples, 0, NULL, TSK_STAT_POLARISED, result);
ts, 2, sample_set_sizes, samples, 0, NULL, 0, NULL, TSK_STAT_POLARISED, result);
CU_ASSERT_EQUAL_FATAL(ret, 0);

ret = tsk_treeseq_allele_frequency_spectrum(ts, 2, sample_set_sizes, samples, 0,
NULL, TSK_STAT_POLARISED | TSK_STAT_SPAN_NORMALISE, result);
NULL, 0, NULL, TSK_STAT_POLARISED | TSK_STAT_SPAN_NORMALISE, result);
CU_ASSERT_EQUAL_FATAL(ret, 0);

ret = tsk_treeseq_allele_frequency_spectrum(ts, 2, sample_set_sizes, samples, 0,
NULL, TSK_STAT_BRANCH | TSK_STAT_POLARISED | TSK_STAT_SPAN_NORMALISE, result);
NULL, 0, NULL, TSK_STAT_BRANCH | TSK_STAT_POLARISED | TSK_STAT_SPAN_NORMALISE, result);
CU_ASSERT_EQUAL_FATAL(ret, 0);

ret = tsk_treeseq_allele_frequency_spectrum(ts, 2, sample_set_sizes, samples, 0,
NULL, TSK_STAT_BRANCH | TSK_STAT_SPAN_NORMALISE, result);
NULL, 0, NULL, TSK_STAT_BRANCH | TSK_STAT_SPAN_NORMALISE, result);
CU_ASSERT_EQUAL_FATAL(ret, 0);

free(result);
Expand Down Expand Up @@ -2418,14 +2473,14 @@ test_paper_ex_afs_errors(void)
tsk_treeseq_from_text(&ts, 10, paper_ex_nodes, paper_ex_edges, NULL, paper_ex_sites,
paper_ex_mutations, paper_ex_individuals, NULL, 0);

verify_one_way_stat_func_errors(&ts, tsk_treeseq_allele_frequency_spectrum);
verify_one_way_stat_func_errors_tw(&ts, tsk_treeseq_allele_frequency_spectrum);

ret = tsk_treeseq_allele_frequency_spectrum(
&ts, 2, sample_set_sizes, samples, 0, NULL, TSK_STAT_NODE, result);
&ts, 2, sample_set_sizes, samples, 0, NULL, 0, NULL, TSK_STAT_NODE, result);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_UNSUPPORTED_STAT_MODE);

ret = tsk_treeseq_allele_frequency_spectrum(&ts, 2, sample_set_sizes, samples, 0,
NULL, TSK_STAT_BRANCH | TSK_STAT_SITE, result);
NULL, 0, NULL, TSK_STAT_BRANCH | TSK_STAT_SITE, result);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_MULTIPLE_STAT_MODES);

tsk_treeseq_free(&ts);
Expand All @@ -2445,14 +2500,14 @@ test_paper_ex_afs(void)
/* we have two singletons and one tripleton */

ret = tsk_treeseq_allele_frequency_spectrum(
&ts, 1, sample_set_sizes, samples, 0, NULL, 0, result);
&ts, 1, sample_set_sizes, samples, 0, NULL, 0, NULL, 0, result);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL_FATAL(result[0], 0);
CU_ASSERT_EQUAL_FATAL(result[1], 3.0);
CU_ASSERT_EQUAL_FATAL(result[2], 0);

ret = tsk_treeseq_allele_frequency_spectrum(
&ts, 1, sample_set_sizes, samples, 0, NULL, TSK_STAT_POLARISED, result);
&ts, 1, sample_set_sizes, samples, 0, NULL, 0, NULL, TSK_STAT_POLARISED, result);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL_FATAL(result[0], 0);
CU_ASSERT_EQUAL_FATAL(result[1], 2.0);
Expand Down
6 changes: 3 additions & 3 deletions c/tests/test_trees.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8182,13 +8182,13 @@ test_time_uncalibrated(void)
CU_ASSERT_EQUAL_FATAL(ret, 0);

ret = tsk_treeseq_allele_frequency_spectrum(
&ts2, 2, sample_set_sizes, samples, 0, NULL, TSK_STAT_SITE, result);
&ts2, 2, sample_set_sizes, samples, 0, NULL, 0, NULL, TSK_STAT_SITE, result);
CU_ASSERT_EQUAL_FATAL(ret, 0);
ret = tsk_treeseq_allele_frequency_spectrum(
&ts2, 2, sample_set_sizes, samples, 0, NULL, TSK_STAT_BRANCH, result);
&ts2, 2, sample_set_sizes, samples, 0, NULL, 0, NULL, TSK_STAT_BRANCH, result);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_TIME_UNCALIBRATED);
ret = tsk_treeseq_allele_frequency_spectrum(&ts2, 2, sample_set_sizes, samples, 0,
NULL, TSK_STAT_BRANCH | TSK_STAT_ALLOW_TIME_UNCALIBRATED, result);
NULL, 0, NULL, TSK_STAT_BRANCH | TSK_STAT_ALLOW_TIME_UNCALIBRATED, result);
CU_ASSERT_EQUAL_FATAL(ret, 0);

sigma = tsk_calloc(tsk_treeseq_get_num_nodes(&ts2), sizeof(double));
Expand Down
138 changes: 104 additions & 34 deletions c/tskit/trees.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1232,6 +1232,35 @@
return ret;
}

static int
tsk_treeseq_check_time_windows(tsk_size_t num_windows,

Check warning on line 1236 in c/tskit/trees.c

View check run for this annotation

Codecov / codecov/patch

c/tskit/trees.c#L1236 

Added line #L1236 was not covered by tests
const double *windows)
{
int ret = TSK_ERR_BAD_WINDOWS;

Check warning on line 1239 in c/tskit/trees.c

View check run for this annotation

Codecov / codecov/patch

c/tskit/trees.c#L1239 

Added line #L1239 was not covered by tests
tsk_size_t j;

if (num_windows < 1) {
ret = TSK_ERR_BAD_NUM_WINDOWS;
goto out;

Check warning on line 1244 in c/tskit/trees.c

View check run for this annotation

Codecov / codecov/patch

c/tskit/trees.c#L1243-L1244 

Added lines #L1243 - L1244 were not covered by tests
}

if (windows[0] < 0) {
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if currently the code assumes this is 0, should check for == here

goto out;

Check warning on line 1248 in c/tskit/trees.c

View check run for this annotation

Codecov / codecov/patch

c/tskit/trees.c#L1248 

Added line #L1248 was not covered by tests
}
if (windows[num_windows] > INFINITY) {
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm pretty sure this can't happen; I think we should either check if this is < INFINITY (if it's assume this is == infinity in the code) or skip this check?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correct, you're right, it would be difficult to initialize something larger than INFINITY! I'm removing this.

goto out;

Check warning on line 1251 in c/tskit/trees.c

View check run for this annotation

Codecov / codecov/patch

c/tskit/trees.c#L1251 

Added line #L1251 was not covered by tests
}

for (j = 0; j < num_windows; j++) {
if (windows[j] >= windows[j + 1]) {
goto out;

Check warning on line 1256 in c/tskit/trees.c

View check run for this annotation

Codecov / codecov/patch

c/tskit/trees.c#L1256 

Added line #L1256 was not covered by tests
}
}
ret = 0;
out:
return ret;

Check warning on line 1261 in c/tskit/trees.c

View check run for this annotation

Codecov / codecov/patch

c/tskit/trees.c#L1259-L1261 

Added lines #L1259 - L1261 were not covered by tests
}

/* TODO make these functions more consistent in how the arguments are ordered */

static inline void
Expand Down Expand Up @@ -3484,10 +3513,25 @@
return ret;
}

#define MAX(a,b) ((a) > (b) ? (a) : (b))
#define MIN(a,b) ((a) < (b) ? (a) : (b))
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

you can use TSK_MIN and TSK_MAX for these? (defined in core.h)

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

YES! Actually. Thanks!


/* int getValue_nDimensions( int * baseAddress, int * indexes, int nDimensions ) { */
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is this needed? remove if not?

/* int i; */
/* int offset = 0; */
/* for( i = 0; i < nDimensions; i++ ) { */
/* offset += pow(LEN,i) * indexes[nDimensions - (i + 1)]; */
/* } */

/* return *(baseAddress + offset); */
/* } */

static int TSK_WARN_UNUSED
tsk_treeseq_update_branch_afs(const tsk_treeseq_t *self, tsk_id_t u, double right,
const double *restrict branch_length, double *restrict last_update,
const double *counts, tsk_size_t num_sample_sets, tsk_size_t window_index,
double *restrict last_update,
const double *restrict time, tsk_id_t *restrict parent, const double *time_windows,
const double *counts, tsk_size_t num_sample_sets,
tsk_size_t num_time_windows, tsk_size_t window_index, tsk_size_t time_window_index,
const tsk_size_t *result_dims, tsk_flags_t options, double *result)
{
int ret = 0;
Expand All @@ -3497,24 +3541,32 @@
tsk_size_t *coordinate = tsk_malloc(num_sample_sets * sizeof(*coordinate));
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Gee, wouldn't it be better to malloc this outside this function, and pass it in? (I honestly don't know...)

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

HA! IDK, I'll let it there for now, but yeah maybe!

bool polarised = !!(options & TSK_STAT_POLARISED);
const double *count_row = GET_2D_ROW(counts, num_sample_sets + 1, u);
double x = (right - last_update[u]) * branch_length[u];
/* double x = (right - last_update[u]) * branch_length[u]; */
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
/* double x = (right - last_update[u]) * branch_length[u]; */

double x = 0;
double t_v = 0;
double tw_branch_length = 0;
const tsk_size_t all_samples = (tsk_size_t) count_row[num_sample_sets];

if (coordinate == NULL) {
ret = TSK_ERR_NO_MEMORY;
goto out;
}

if (0 < all_samples && all_samples < self->num_samples) {
afs_size = result_dims[num_sample_sets];
afs = result + afs_size * window_index;
for (k = 0; k < num_sample_sets; k++) {
coordinate[k] = (tsk_size_t) count_row[k];
}
if (!polarised) {
fold(coordinate, result_dims, num_sample_sets);
}
increment_nd_array_value(afs, num_sample_sets, result_dims, coordinate, x);
if (parent[u] != -1){
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
if (parent[u] != -1){
if (parent[u] != TSK_NULL){

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hah - it took me a while to see why we needed this, but now I see it

t_v = time[parent[u]];
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

perhaps also t_u here?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There's no t_u variable, do you think that's necessary? I'm using time[u] here.

if (0 < all_samples && all_samples < self->num_samples) {
for (time_window_index = 0; time_window_index < num_time_windows; time_window_index++){
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A lot of edges are recent, so we might avoid substantial work if we do like

	    time_window_index = 0;
	    while (time_window_index < num_time_windows && time_windows[time_window_index] < t_v){
                   ...
                 time_window_index++;
             }

afs_size = result_dims[num_sample_sets];
afs = result + afs_size * (window_index * num_time_windows + time_window_index);
for (k = 0; k < num_sample_sets; k++) {
coordinate[k] = (tsk_size_t) count_row[k];
}
if (!polarised){
fold(coordinate, result_dims, num_sample_sets);
}
tw_branch_length = MIN(time_windows[time_window_index + 1], t_v) - MAX(time_windows[0], time[u]);
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

shouldn't this be

Suggested change
tw_branch_length = MIN(time_windows[time_window_index + 1], t_v) - MAX(time_windows[0], time[u]);
tw_branch_length = MIN(time_windows[time_window_index + 1], t_v) - MAX(time_windows[time_window_index], time[u]);

?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hm - the tests below should be catching this if it is indeed wrong, but it sure looks wrong to me - I'm not sure what's going on?

x = (right - last_update[u]) * tw_branch_length;
increment_nd_array_value(afs, num_sample_sets, result_dims, coordinate, x);
}
}
}
last_update[u] = right;
out:
Expand All @@ -3525,12 +3577,12 @@
static int
tsk_treeseq_branch_allele_frequency_spectrum(const tsk_treeseq_t *self,
tsk_size_t num_sample_sets, double *counts, tsk_size_t num_windows,
const double *windows, const tsk_size_t *result_dims, tsk_flags_t options,
double *result)
tsk_size_t num_time_windows, const double *windows, const double *time_windows,
const tsk_size_t *result_dims, tsk_flags_t options, double *result)
{
int ret = 0;
tsk_id_t u, v;
tsk_size_t window_index;
tsk_size_t window_index, time_window_index;
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think time_window_index is actually used in this function (since the update function iterates over all time windows).

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correct. So the allocation should be done locally too (in the update function), I presume.

tsk_size_t num_nodes = self->tables->nodes.num_rows;
const tsk_id_t num_edges = (tsk_id_t) self->tables->edges.num_rows;
const tsk_id_t *restrict I = self->tables->indexes.edge_insertion_order;
Expand Down Expand Up @@ -3564,23 +3616,26 @@
tk = 0;
t_left = 0;
window_index = 0;
time_window_index = 0;
while (tj < num_edges || t_left < sequence_length) {
tsk_bug_assert(window_index < num_windows);
while (tk < num_edges && edge_right[O[tk]] == t_left) {
h = O[tk];
tk++;
u = edge_child[h];
v = edge_parent[h];
ret = tsk_treeseq_update_branch_afs(self, u, t_left, branch_length,
last_update, counts, num_sample_sets, window_index, result_dims, options,
result);
ret = tsk_treeseq_update_branch_afs(self, u, t_left,
last_update, node_time, parent, time_windows, counts, num_sample_sets,
num_time_windows, window_index, time_window_index,
result_dims, options, result);
if (ret != 0) {
goto out;
}
while (v != TSK_NULL) {
ret = tsk_treeseq_update_branch_afs(self, v, t_left, branch_length,
last_update, counts, num_sample_sets, window_index, result_dims,
options, result);
ret = tsk_treeseq_update_branch_afs(self, v, t_left,
last_update, node_time, parent, time_windows, counts,
num_sample_sets, num_time_windows, window_index,
time_window_index, result_dims, options, result);
if (ret != 0) {
goto out;
}
Expand All @@ -3599,9 +3654,10 @@
parent[u] = v;
branch_length[u] = node_time[v] - node_time[u];
while (v != TSK_NULL) {
ret = tsk_treeseq_update_branch_afs(self, v, t_left, branch_length,
last_update, counts, num_sample_sets, window_index, result_dims,
options, result);
ret = tsk_treeseq_update_branch_afs(self, v, t_left,
last_update, node_time, parent, time_windows, counts,
num_sample_sets, num_time_windows, window_index,
time_window_index, result_dims, options, result);
if (ret != 0) {
goto out;
}
Expand All @@ -3623,9 +3679,10 @@
/* Flush the contributions of all nodes to the current window */
for (u = 0; u < (tsk_id_t) num_nodes; u++) {
tsk_bug_assert(last_update[u] < w_right);
ret = tsk_treeseq_update_branch_afs(self, u, w_right, branch_length,
last_update, counts, num_sample_sets, window_index, result_dims,
options, result);
ret = tsk_treeseq_update_branch_afs(self, u, w_right,
last_update, node_time, parent, time_windows, counts,
num_sample_sets, num_time_windows, window_index,
time_window_index, result_dims, options, result);
if (ret != 0) {
goto out;
}
Expand Down Expand Up @@ -3653,13 +3710,15 @@
tsk_treeseq_allele_frequency_spectrum(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 num_windows, const double *windows,
tsk_flags_t options, double *result)
tsk_size_t num_time_windows, const double *time_windows, tsk_flags_t options,
double *result)
{
int ret = 0;
bool stat_site = !!(options & TSK_STAT_SITE);
bool stat_branch = !!(options & TSK_STAT_BRANCH);
bool stat_node = !!(options & TSK_STAT_NODE);
const double default_windows[] = { 0, self->tables->sequence_length };
const double default_time_windows[] = { 0, INFINITY };
const tsk_size_t num_nodes = self->tables->nodes.num_rows;
const tsk_size_t K = num_sample_sets + 1;
tsk_size_t j, k, l, afs_size;
Expand All @@ -3669,7 +3728,6 @@
* reuse code from the general_stats code paths. */
double *counts = NULL;
double *count_row;

if (stat_node) {
ret = TSK_ERR_UNSUPPORTED_STAT_MODE;
goto out;
Expand All @@ -3693,6 +3751,16 @@
goto out;
}
}
if (time_windows == NULL) {
num_time_windows = 1;
time_windows = default_time_windows;
} else {
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

After this line is probably the right place to check if it's mode="site" and throw an error?

ret = tsk_treeseq_check_time_windows(

Check warning on line 3758 in c/tskit/trees.c

View check run for this annotation

Codecov / codecov/patch

c/tskit/trees.c#L3758 

Added line #L3758 was not covered by tests
num_time_windows, time_windows);
if (ret != 0) {
goto out;

Check warning on line 3761 in c/tskit/trees.c

View check run for this annotation

Codecov / codecov/patch

c/tskit/trees.c#L3761 

Added line #L3761 was not covered by tests
}
}
ret = tsk_treeseq_check_sample_sets(
self, num_sample_sets, sample_set_sizes, sample_sets);
if (ret != 0) {
Expand Down Expand Up @@ -3728,15 +3796,17 @@
count_row[num_sample_sets] = 1;
}
result_dims[num_sample_sets] = (tsk_size_t) afs_size;
// Initiate memory for result array
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
// Initiate memory for result array

tsk_memset(result, 0, num_windows * num_time_windows * afs_size * sizeof(*result));

tsk_memset(result, 0, num_windows * afs_size * sizeof(*result));
if (stat_site) {
ret = tsk_treeseq_site_allele_frequency_spectrum(self, num_sample_sets,
sample_set_sizes, counts, num_windows, windows, result_dims, options,
result);
} else {
ret = tsk_treeseq_branch_allele_frequency_spectrum(self, num_sample_sets, counts,
num_windows, windows, result_dims, options, result);
num_windows, num_time_windows, windows, time_windows, result_dims, options,
result);
}

if (options & TSK_STAT_SPAN_NORMALISE) {
Expand Down
Loading
Loading