Implementation of time windows for the AFS,

and testing version of time windows for general stat

Author: petrelharp

c/tests/test_stats.c

@@ -23,6 +23,7 @@
  */
 
 #include "testlib.h"
+#include <math.h>
 #include <tskit/stats.h>
 
 #include <unistd.h>
@@ -602,7 +603,7 @@ verify_pair_coalescence_rates(tsk_treeseq_t *ts)
     epochs[B] = DBL_MAX;
     ret = tsk_treeseq_pair_coalescence_rates(ts, P, sample_set_sizes, sample_sets, I,
         index_tuples, T, breakpoints, B, node_bin_map, epochs, 0, C);
-    CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_BAD_TIME_WINDOWS);
+    CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_BAD_TIME_WINDOWS_END);
     epochs[B] = INFINITY;
 
     node_bin_map[0] = (tsk_id_t) B;
@@ -868,6 +869,84 @@ 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);
+
+// Temporary duplicate for time-windows-having methods
+static void
+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 time_windows[] = { -1, 0.5, INFINITY };
+    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 */
+    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);
+
+    /* Time window errors */
+    ret = method(
+        ts, 1, &sample_set_sizes, samples, 0, NULL, 0, time_windows, 0, &result);
+    CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_BAD_TIME_WINDOWS_DIM);
+
+    ret = method(
+        ts, 1, &sample_set_sizes, samples, 0, NULL, 2, time_windows, 0, &result);
+    CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_BAD_TIME_WINDOWS);
+
+    time_windows[0] = 0.1;
+    ret = method(
+        ts, 1, &sample_set_sizes, samples, 0, NULL, 2, time_windows, 0, &result);
+    CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_BAD_TIME_WINDOWS);
+
+    time_windows[0] = 0;
+    time_windows[1] = 0;
+    ret = method(
+        ts, 1, &sample_set_sizes, samples, 0, NULL, 2, time_windows, 0, &result);
+    CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_BAD_TIME_WINDOWS);
+}
+
 static void
 verify_two_way_stat_func_errors(
     tsk_treeseq_t *ts, general_sample_stat_method *method, tsk_flags_t options)
@@ -1206,23 +1285,24 @@ 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);
@@ -2416,21 +2496,26 @@ test_paper_ex_afs_errors(void)
     tsk_size_t sample_set_sizes[] = { 2, 2 };
     tsk_id_t samples[] = { 0, 1, 2, 3 };
     double result[10]; /* not thinking too hard about the actual value needed */
+    double time_windows[] = { 0, 1 };
     int ret;
 
     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);
 
+    ret = tsk_treeseq_allele_frequency_spectrum(&ts, 2, sample_set_sizes, samples, 0,
+        NULL, 1, time_windows, TSK_STAT_SITE, result);
+    CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_UNSUPPORTED_STAT_MODE);
+
     tsk_treeseq_free(&ts);
 }
 
@@ -2448,14 +2533,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);

c/tests/test_trees.c

@@ -8026,13 +8026,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));

c/tskit/core.c

@@ -526,9 +526,13 @@ tsk_strerror_internal(int err)
                   "(TSK_ERR_BAD_SAMPLE_PAIR_TIMES)";
             break;
         case TSK_ERR_BAD_TIME_WINDOWS:
-            ret = "Time windows must be strictly increasing and end at infinity. "
+            ret = "Time windows must be strictly increasing. "
                   "(TSK_ERR_BAD_TIME_WINDOWS)";
             break;
+        case TSK_ERR_BAD_TIME_WINDOWS_END:
+            ret = "Time windows must end at infinity for this method. "
+                  "(TSK_ERR_BAD_TIME_WINDOWS_END)";
+            break;
         case TSK_ERR_BAD_NODE_TIME_WINDOW:
             ret = "Node time does not fall within assigned time window. "
                   "(TSK_ERR_BAD_NODE_TIME_WINDOW)";

c/tskit/core.h

@@ -746,13 +746,17 @@ Sample times do not all equal the start of first time window
 */
 #define TSK_ERR_BAD_SAMPLE_PAIR_TIMES                               -923
 /**
-Time windows are not strictly increasing ending at infinity
+Time windows are not strictly increasing
 */
 #define TSK_ERR_BAD_TIME_WINDOWS                                    -924
 /**
+Time windows do not end at infinity
+*/
+#define TSK_ERR_BAD_TIME_WINDOWS_END                                -925
+/**
 Node time does not fall within assigned time window
 */
-#define TSK_ERR_BAD_NODE_TIME_WINDOW                                -925
+#define TSK_ERR_BAD_NODE_TIME_WINDOW                                -926
 /** @} */
 
 /**