Two locus general matrix

.github/workflows/lint.yml

@@ -2,6 +2,7 @@ name: Lint
 
 on:
   pull_request:
+  merge_group:
 
 jobs:
   Lint:

c/tests/test_core.c

@@ -101,19 +101,19 @@ test_json_struct_metadata_get_blob(void)
 {
     int ret;
     char metadata[128];
-    const char *json;
+    char *json;
     tsk_size_t json_buffer_length;
-    const char *blob;
+    char *blob;
     tsk_size_t blob_length;
     uint8_t *bytes;
     tsk_size_t metadata_length;
     size_t header_length;
     size_t json_length;
     size_t payload_length;
     size_t total_length;
-    const char json_payload[] = "{\"a\":1}";
-    const uint8_t binary_payload[] = { 0x01, 0x02, 0x03, 0x04 };
-    const uint8_t empty_payload[] = { 0 };
+    char json_payload[] = "{\"a\":1}";
+    uint8_t binary_payload[] = { 0x01, 0x02, 0x03, 0x04 };
+    uint8_t empty_payload[] = { 0 };
 
     bytes = (uint8_t *) metadata;
     header_length = 4 + 1 + 8 + 8;
@@ -135,7 +135,7 @@ test_json_struct_metadata_get_blob(void)
     ret = tsk_json_struct_metadata_get_blob(
         metadata, metadata_length, &json, &json_buffer_length, &blob, &blob_length);
     CU_ASSERT_EQUAL(ret, 0);
-    CU_ASSERT_PTR_EQUAL(json, (const char *) bytes + header_length);
+    CU_ASSERT_PTR_EQUAL(json, (char *) bytes + header_length);
     CU_ASSERT_EQUAL(json_buffer_length, (tsk_size_t) json_length);
     if (json_length > 0) {
         CU_ASSERT_EQUAL(memcmp(json, json_payload, json_length), 0);
@@ -152,7 +152,7 @@ test_json_struct_metadata_get_blob(void)
     ret = tsk_json_struct_metadata_get_blob(
         metadata, metadata_length, &json, &json_buffer_length, &blob, &blob_length);
     CU_ASSERT_EQUAL(ret, 0);
-    CU_ASSERT_PTR_EQUAL(json, (const char *) bytes + header_length);
+    CU_ASSERT_PTR_EQUAL(json, (char *) bytes + header_length);
     CU_ASSERT_EQUAL(json_buffer_length, (tsk_size_t) json_length);
     CU_ASSERT_EQUAL(blob_length, (tsk_size_t) payload_length);
     CU_ASSERT_PTR_EQUAL(blob, bytes + header_length + json_length);
@@ -168,7 +168,7 @@ test_json_struct_metadata_get_blob(void)
     ret = tsk_json_struct_metadata_get_blob(
         metadata, metadata_length, &json, &json_buffer_length, &blob, &blob_length);
     CU_ASSERT_EQUAL(ret, 0);
-    CU_ASSERT_PTR_EQUAL(json, (const char *) bytes + header_length);
+    CU_ASSERT_PTR_EQUAL(json, (char *) bytes + header_length);
     CU_ASSERT_EQUAL(json_buffer_length, (tsk_size_t) json_length);
     CU_ASSERT_EQUAL(blob_length, (tsk_size_t) payload_length);
     CU_ASSERT_PTR_EQUAL(blob, bytes + header_length + json_length);

c/tests/test_stats.c

@@ -3680,9 +3680,9 @@ test_two_locus_stat_input_errors(void)
         row_sites, NULL, num_sites, col_sites, NULL, 0, result);
     CU_ASSERT_EQUAL_FATAL(ret, 0);
 
-    ret = tsk_treeseq_two_locus_count_stat(&ts, num_sample_sets, sample_set_sizes,
-        sample_sets, 0, NULL, NULL, NULL, num_sites, row_sites, NULL, num_sites,
-        col_sites, NULL, 0, result);
+    ret = tsk_treeseq_two_locus_general_count_stat(&ts, num_sample_sets,
+        sample_set_sizes, sample_sets, 0, NULL, NULL, NULL, num_sites, row_sites, NULL,
+        num_sites, col_sites, NULL, 0, result);
     CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_BAD_RESULT_DIMS);
 
     ret = tsk_treeseq_r2(&ts, 1, sample_set_sizes, sample_sets, num_sites, row_sites,

c/tskit/core.c

@@ -142,26 +142,25 @@ tsk_generate_uuid(char *dest, int TSK_UNUSED(flags))
 }
 
 int
-tsk_json_struct_metadata_get_blob(const char *metadata, tsk_size_t metadata_length,
-    const char **json, tsk_size_t *json_length, const char **blob,
-    tsk_size_t *blob_length)
+tsk_json_struct_metadata_get_blob(char *metadata, tsk_size_t metadata_length,
+    char **json, tsk_size_t *json_length, char **blob, tsk_size_t *blob_length)
 {
     int ret;
     uint8_t version;
     uint64_t json_length_u64;
     uint64_t binary_length_u64;
     uint64_t header_and_json_length;
     uint64_t total_length;
-    const uint8_t *bytes;
-    const char *blob_start;
-    const char *json_start;
+    uint8_t *bytes;
+    char *blob_start;
+    char *json_start;
 
     if (metadata == NULL || json == NULL || json_length == NULL || blob == NULL
         || blob_length == NULL) {
         ret = tsk_trace_error(TSK_ERR_BAD_PARAM_VALUE);
         goto out;
     }
-    bytes = (const uint8_t *) metadata;
+    bytes = (uint8_t *) metadata;
     if (metadata_length < TSK_JSON_BINARY_HEADER_SIZE) {
         ret = tsk_trace_error(TSK_ERR_JSON_STRUCT_METADATA_TRUNCATED);
         goto out;
@@ -191,8 +190,8 @@ tsk_json_struct_metadata_get_blob(const char *metadata, tsk_size_t metadata_leng
         ret = tsk_trace_error(TSK_ERR_JSON_STRUCT_METADATA_TRUNCATED);
         goto out;
     }
-    json_start = (const char *) bytes + TSK_JSON_BINARY_HEADER_SIZE;
-    blob_start = (const char *) bytes + TSK_JSON_BINARY_HEADER_SIZE + json_length_u64;
+    json_start = (char *) bytes + TSK_JSON_BINARY_HEADER_SIZE;
+    blob_start = (char *) bytes + TSK_JSON_BINARY_HEADER_SIZE + json_length_u64;
     *json = json_start;
     *json_length = (tsk_size_t) json_length_u64;
     *blob = blob_start;

c/tskit/core.h

@@ -1153,9 +1153,8 @@ the original metadata buffer is alive.
 @param[out] blob_length On success, set to the payload length in bytes.
 @return Return 0 on success or a negative value on failure.
 */
-int tsk_json_struct_metadata_get_blob(const char *metadata, tsk_size_t metadata_length,
-    const char **json, tsk_size_t *json_length, const char **blob,
-    tsk_size_t *blob_length);
+int tsk_json_struct_metadata_get_blob(char *metadata, tsk_size_t metadata_length,
+    char **json, tsk_size_t *json_length, char **blob, tsk_size_t *blob_length);
 
 /* TODO most of these can probably be macros so they compile out as no-ops.
  * Lets do the 64 bit tsk_size_t switch first though. */

c/tskit/trees.c

@@ -2298,8 +2298,9 @@ get_allele_samples(const tsk_site_t *site, tsk_size_t site_offset,
 }
 
 static int
-norm_hap_weighted(tsk_size_t result_dim, const double *hap_weights,
-    tsk_size_t TSK_UNUSED(n_a), tsk_size_t TSK_UNUSED(n_b), double *result, void *params)
+norm_hap_weighted(tsk_size_t TSK_UNUSED(state_dim), const double *hap_weights,
+    tsk_size_t result_dim, tsk_size_t TSK_UNUSED(n_a), tsk_size_t TSK_UNUSED(n_b),
+    double *result, void *params)
 {
     sample_count_stat_params_t args = *(sample_count_stat_params_t *) params;
     const double *weight_row;
@@ -2315,8 +2316,9 @@ norm_hap_weighted(tsk_size_t result_dim, const double *hap_weights,
 }
 
 static int
-norm_hap_weighted_ij(tsk_size_t result_dim, const double *hap_weights,
-    tsk_size_t TSK_UNUSED(n_a), tsk_size_t TSK_UNUSED(n_b), double *result, void *params)
+norm_hap_weighted_ij(tsk_size_t TSK_UNUSED(state_dim), const double *hap_weights,
+    tsk_size_t result_dim, tsk_size_t TSK_UNUSED(n_a), tsk_size_t TSK_UNUSED(n_b),
+    double *result, void *params)
 {
     sample_count_stat_params_t args = *(sample_count_stat_params_t *) params;
     const double *weight_row;
@@ -2341,8 +2343,9 @@ norm_hap_weighted_ij(tsk_size_t result_dim, const double *hap_weights,
 }
 
 static int
-norm_total_weighted(tsk_size_t result_dim, const double *TSK_UNUSED(hap_weights),
-    tsk_size_t n_a, tsk_size_t n_b, double *result, void *TSK_UNUSED(params))
+norm_total_weighted(tsk_size_t TSK_UNUSED(state_dim),
+    const double *TSK_UNUSED(hap_weights), tsk_size_t result_dim, tsk_size_t n_a,
+    tsk_size_t n_b, double *result, void *TSK_UNUSED(params))
 {
     tsk_size_t k;
     double norm = 1 / (double) (n_a * n_b);
@@ -2411,8 +2414,8 @@ static int
 compute_general_normed_two_site_stat_result(const tsk_bitset_t *state,
     const tsk_size_t *allele_counts, tsk_size_t a_off, tsk_size_t b_off,
     tsk_size_t num_a_alleles, tsk_size_t num_b_alleles, tsk_size_t state_dim,
-    tsk_size_t result_dim, general_stat_func_t *f, sample_count_stat_params_t *f_params,
-    norm_func_t *norm_f, bool polarised, two_locus_work_t *restrict work, double *result)
+    tsk_size_t result_dim, general_stat_func_t *f, void *f_params, norm_func_t *norm_f,
+    bool polarised, two_locus_work_t *restrict work, double *result)
 {
     int ret = 0;
     // Sample sets and b sites are rows, a sites are columns
@@ -2445,7 +2448,7 @@ compute_general_normed_two_site_stat_result(const tsk_bitset_t *state,
             if (ret != 0) {
                 goto out;
             }
-            ret = norm_f(result_dim, weights, num_a_alleles - is_polarised,
+            ret = norm_f(state_dim, weights, result_dim, num_a_alleles - is_polarised,
                 num_b_alleles - is_polarised, norm, f_params);
             if (ret != 0) {
                 goto out;
@@ -2463,9 +2466,8 @@ compute_general_normed_two_site_stat_result(const tsk_bitset_t *state,
 static int
 compute_general_two_site_stat_result(const tsk_bitset_t *state,
     const tsk_size_t *allele_counts, tsk_size_t a_off, tsk_size_t b_off,
-    tsk_size_t state_dim, tsk_size_t result_dim, general_stat_func_t *f,
-    sample_count_stat_params_t *f_params, two_locus_work_t *restrict work,
-    double *result)
+    tsk_size_t state_dim, tsk_size_t result_dim, general_stat_func_t *f, void *f_params,
+    two_locus_work_t *restrict work, double *result)
 {
     int ret = 0;
     tsk_size_t k;
@@ -2653,9 +2655,8 @@ static int
 tsk_treeseq_two_site_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim,
     tsk_size_t num_sample_sets, const tsk_size_t *sample_set_sizes,
     const tsk_id_t *sample_sets, tsk_size_t result_dim, general_stat_func_t *f,
-    sample_count_stat_params_t *f_params, norm_func_t *norm_f, tsk_size_t n_rows,
-    const tsk_id_t *row_sites, tsk_size_t n_cols, const tsk_id_t *col_sites,
-    tsk_flags_t options, double *result)
+    void *f_params, norm_func_t *norm_f, tsk_size_t n_rows, const tsk_id_t *row_sites,
+    tsk_size_t n_cols, const tsk_id_t *col_sites, tsk_flags_t options, double *result)
 {
     int ret = 0;
     tsk_bitset_t allele_samples, allele_sample_sets;
@@ -3089,9 +3090,8 @@ advance_collect_edges(iter_state *s, tsk_id_t index)
 static int
 compute_two_tree_branch_state_update(const tsk_treeseq_t *ts, tsk_id_t c,
     const iter_state *A_state, const iter_state *B_state, tsk_size_t state_dim,
-    tsk_size_t result_dim, int sign, general_stat_func_t *f,
-    sample_count_stat_params_t *f_params, two_locus_work_t *restrict work,
-    double *result)
+    tsk_size_t result_dim, int sign, general_stat_func_t *f, void *f_params,
+    two_locus_work_t *restrict work, double *result)
 {
     int ret = 0;
     double a_len, b_len;
@@ -3141,8 +3141,8 @@ compute_two_tree_branch_state_update(const tsk_treeseq_t *ts, tsk_id_t c,
 
 static int
 compute_two_tree_branch_stat(const tsk_treeseq_t *ts, const iter_state *l_state,
-    iter_state *r_state, general_stat_func_t *f, sample_count_stat_params_t *f_params,
-    tsk_size_t result_dim, tsk_size_t state_dim, double *result)
+    iter_state *r_state, general_stat_func_t *f, void *f_params, tsk_size_t result_dim,
+    tsk_size_t state_dim, double *result)
 {
     int ret = 0;
     tsk_id_t e, c, ec, p, *updated_nodes = NULL;
@@ -3188,8 +3188,11 @@ compute_two_tree_branch_stat(const tsk_treeseq_t *ts, const iter_state *l_state,
     while (n_updates != 0) {
         n_updates--;
         c = updated_nodes[n_updates];
-        compute_two_tree_branch_state_update(ts, c, l_state, r_state, state_dim,
+        ret = compute_two_tree_branch_state_update(ts, c, l_state, r_state, state_dim,
             result_dim, -1, f, f_params, &work, result);
+        if (ret != 0) {
+            goto out;
+        }
     }
     // Remove samples under nodes from removed edges to parent nodes
     for (j = 0; j < r_state->n_edges_out; j++) {
@@ -3229,8 +3232,11 @@ compute_two_tree_branch_stat(const tsk_treeseq_t *ts, const iter_state *l_state,
     while (n_updates != 0) {
         n_updates--;
         c = updated_nodes[n_updates];
-        compute_two_tree_branch_state_update(ts, c, l_state, r_state, state_dim,
+        ret = compute_two_tree_branch_state_update(ts, c, l_state, r_state, state_dim,
             result_dim, +1, f, f_params, &work, result);
+        if (ret != 0) {
+            goto out;
+        }
     }
 out:
     tsk_safe_free(updated_nodes);
@@ -3243,9 +3249,9 @@ static int
 tsk_treeseq_two_branch_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim,
     tsk_size_t num_sample_sets, const tsk_size_t *sample_set_sizes,
     const tsk_id_t *sample_sets, tsk_size_t result_dim, general_stat_func_t *f,
-    sample_count_stat_params_t *f_params, norm_func_t *TSK_UNUSED(norm_f),
-    tsk_size_t n_rows, const double *row_positions, tsk_size_t n_cols,
-    const double *col_positions, tsk_flags_t TSK_UNUSED(options), double *result)
+    void *f_params, norm_func_t *TSK_UNUSED(norm_f), tsk_size_t n_rows,
+    const double *row_positions, tsk_size_t n_cols, const double *col_positions,
+    tsk_flags_t TSK_UNUSED(options), double *result)
 {
     int ret = 0;
     int r, c;
@@ -3384,11 +3390,12 @@ check_sample_set_dups(tsk_size_t num_sample_sets, const tsk_size_t *sample_set_s
     return ret;
 }
 
+/* Called directly by C python interface `two_locus_count_stat` */
 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 out_rows, const tsk_id_t *row_sites,
+tsk_treeseq_two_locus_general_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, general_stat_func_t *f,
+    void *f_params, norm_func_t *norm_f, tsk_size_t out_rows, const tsk_id_t *row_sites,
     const double *row_positions, tsk_size_t out_cols, const tsk_id_t *col_sites,
     const double *col_positions, tsk_flags_t options, double *result)
 {
@@ -3398,10 +3405,6 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl
     bool stat_site = !!(options & TSK_STAT_SITE);
     bool stat_branch = !!(options & TSK_STAT_BRANCH);
     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,
-        .sample_set_sizes = sample_set_sizes,
-        .set_indexes = set_indexes };
 
     // We do not support two-locus node stats
     if (!!(options & TSK_STAT_NODE)) {
@@ -3441,9 +3444,10 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl
             goto out;
         }
         ret = tsk_treeseq_two_site_count_stat(self, state_dim, num_sample_sets,
-            sample_set_sizes, sample_sets, result_dim, f, &f_params, norm_f, out_rows,
+            sample_set_sizes, sample_sets, result_dim, f, f_params, norm_f, out_rows,
             row_sites, out_cols, col_sites, options, result);
-    } else if (stat_branch) {
+    } else {
+        tsk_bug_assert(stat_branch);
         ret = check_positions(
             row_positions, out_rows, tsk_treeseq_get_sequence_length(self));
         if (ret != 0) {
@@ -3455,13 +3459,31 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl
             goto out;
         }
         ret = tsk_treeseq_two_branch_count_stat(self, state_dim, num_sample_sets,
-            sample_set_sizes, sample_sets, result_dim, f, &f_params, norm_f, out_rows,
+            sample_set_sizes, sample_sets, result_dim, f, f_params, norm_f, out_rows,
             row_positions, out_cols, col_positions, options, result);
     }
 out:
     return ret;
 }
 
+/* Wrapper of `tsk_treeseq_two_locus_general_count_stat` for C summary Functions */
+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 out_rows, const tsk_id_t *row_sites,
+    const double *row_positions, tsk_size_t out_cols, const tsk_id_t *col_sites,
+    const double *col_positions, tsk_flags_t options, double *result)
+{
+    sample_count_stat_params_t f_params = { .sample_sets = sample_sets,
+        .num_sample_sets = num_sample_sets,
+        .sample_set_sizes = sample_set_sizes,
+        .set_indexes = set_indexes };
+    return tsk_treeseq_two_locus_general_count_stat(self, num_sample_sets,
+        sample_set_sizes, sample_sets, result_dim, f, &f_params, norm_f, out_rows,
+        row_sites, row_positions, out_cols, col_sites, col_positions, options, result);
+}
+
 /***********************************
  * Allele frequency spectrum
  ***********************************/

c/tskit/trees.h

@@ -1036,16 +1036,8 @@ int tsk_treeseq_general_stat(const tsk_treeseq_t *self, tsk_size_t K, const doub
     tsk_size_t M, general_stat_func_t *f, void *f_params, tsk_size_t num_windows,
     const double *windows, tsk_flags_t options, double *result);
 
-typedef int norm_func_t(tsk_size_t result_dim, const double *hap_weights, tsk_size_t n_a,
-    tsk_size_t n_b, double *result, void *params);
-
-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 out_rows,
-    const tsk_id_t *row_sites, const double *row_positions, tsk_size_t out_cols,
-    const tsk_id_t *col_sites, const double *col_positions, tsk_flags_t options,
-    double *result);
+typedef int norm_func_t(tsk_size_t state_dim, const double *hap_weights,
+    tsk_size_t result_dim, tsk_size_t n_a, tsk_size_t n_b, double *result, void *params);
 
 /* One way weighted stats */
 
@@ -1120,6 +1112,13 @@ typedef int general_sample_stat_method(const tsk_treeseq_t *self,
     const tsk_id_t *sample_sets, tsk_size_t num_indexes, const tsk_id_t *indexes,
     tsk_size_t num_windows, const double *windows, tsk_flags_t options, double *result);
 
+int tsk_treeseq_two_locus_general_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, general_stat_func_t *f,
+    void *f_params, norm_func_t *norm_f, tsk_size_t out_rows, const tsk_id_t *row_sites,
+    const double *row_positions, tsk_size_t out_cols, const tsk_id_t *col_sites,
+    const double *col_positions, tsk_flags_t options, double *result);
+
 typedef int two_locus_count_stat_method(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_rows, const tsk_id_t *row_sites,