From f4ae923f016d4d5ec9ad231bef84e508400c0b72 Mon Sep 17 00:00:00 2001 From: lkirk Date: Tue, 18 Jul 2023 14:12:35 -0500 Subject: [PATCH] Allocate the result array within the site stat fn We'll have to think carefully about the python ownership of this data and how it gets cleaned up. This is the only way it can work, especially once we have windows that restrict the array size even further. --- c/tskit/trees.c | 42 +++++++++++++++++++++--------------------- c/tskit/trees.h | 14 +++++++------- 2 files changed, 28 insertions(+), 28 deletions(-) diff --git a/c/tskit/trees.c b/c/tskit/trees.c index 10cce4ab68..35fbb4f6e7 100644 --- a/c/tskit/trees.c +++ b/c/tskit/trees.c @@ -2564,9 +2564,8 @@ tsk_treeseq_two_site_general_stat(const tsk_treeseq_t *self, tsk_size_t state_di const double *sample_weights, tsk_size_t result_dim, general_stat_func_t *f, sample_count_stat_params_t *f_params, norm_func_t *norm_f, const double *TSK_UNUSED(left_window), const double *TSK_UNUSED(right_window), - tsk_flags_t options, tsk_size_t num_result, double *result) + tsk_flags_t options, tsk_size_t *result_size, double **result) { - // TODO: implement windows! int ret = 0; const tsk_size_t num_sites = self->tables->sites.num_rows; const tsk_size_t num_samples = self->num_samples; @@ -2612,9 +2611,9 @@ tsk_treeseq_two_site_general_stat(const tsk_treeseq_t *self, tsk_size_t state_di } // Number of pairs w/ replacement (sites) - num_result = (num_sites * (1 + num_sites)) >> (tsk_size_t) 1; + *result_size = (num_sites * (1 + num_sites)) >> (tsk_size_t) 1; + *result = tsk_calloc(*result_size * result_dim, sizeof(*result)); - result = tsk_calloc(num_result * result_dim, sizeof(*result)); tsk_bit_array_t *sample_bits = tsk_calloc(num_sample_chunks * state_dim, sizeof(*sample_bits)); @@ -2631,19 +2630,19 @@ tsk_treeseq_two_site_general_stat(const tsk_treeseq_t *self, tsk_size_t state_di polarised = true; } - tsk_size_t inner = 0, site_a_offset = 0, site_b_offset = 0, result_offset = 0; + // TODO: implement windows! + tsk_size_t inner = 0; + tsk_size_t result_offset = 0; for (tsk_size_t site_a = 0; site_a < num_sites; site_a++) { - site_a_offset = site_offsets[site_a]; for (tsk_size_t site_b = inner; site_b < num_sites; site_b++) { - site_b_offset = site_offsets[site_b]; - ret = compute_general_two_site_stat_result(site_a, site_a_offset, site_b, - site_b_offset, num_sample_chunks, num_alleles, allele_samples, state_dim, - sample_bits, result_dim, f, f_params, norm_f, polarised, - &(result[result_offset])); + ret = compute_general_two_site_stat_result(site_a, site_offsets[site_a], + site_b, site_offsets[site_b], num_sample_chunks, num_alleles, + allele_samples, state_dim, sample_bits, result_dim, f, f_params, norm_f, + polarised, &((*result)[result_offset])); if (ret != 0) { goto out; } - result_offset++; + result_offset += state_dim; } inner++; } @@ -2671,7 +2670,7 @@ tsk_treeseq_two_locus_general_stat(const tsk_treeseq_t *self, tsk_size_t num_sam 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 num_result, double *result) + tsk_size_t *result_size, double **result) { int ret = 0; bool stat_site = !!(options & TSK_STAT_SITE); @@ -2749,7 +2748,8 @@ tsk_treeseq_two_locus_general_stat(const tsk_treeseq_t *self, tsk_size_t num_sam if (stat_site) { // TODO: assert 1 window left/right ret = tsk_treeseq_two_site_general_stat(self, state_dim, weights, result_dim, f, - &f_params, norm_f, left_windows, right_windows, options, num_result, result); + &f_params, norm_f, left_windows, right_windows, options, result_size, + result); } else { // TODO: branch statistics ret = TSK_ERR_GENERIC; @@ -3585,7 +3585,7 @@ tsk_treeseq_D(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_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) + tsk_size_t *result_size, double **result) { options |= TSK_STAT_POLARISED; // TODO: allow user to pick? return tsk_treeseq_two_locus_general_stat(self, num_sample_sets, sample_set_sizes, @@ -3622,7 +3622,7 @@ tsk_treeseq_D2(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_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) + tsk_size_t *result_size, double **result) { return tsk_treeseq_two_locus_general_stat(self, num_sample_sets, sample_set_sizes, sample_sets, num_sample_sets, NULL, D2_summary_func, norm_total_weighted, @@ -3664,7 +3664,7 @@ tsk_treeseq_r2(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_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) + tsk_size_t *result_size, double **result) { return tsk_treeseq_two_locus_general_stat(self, num_sample_sets, sample_set_sizes, sample_sets, num_sample_sets, NULL, r2_summary_func, norm_hap_weighted, @@ -3704,7 +3704,7 @@ tsk_treeseq_D_prime(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_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) + tsk_size_t *result_size, double **result) { options |= TSK_STAT_POLARISED; // TODO: allow user to pick? return tsk_treeseq_two_locus_general_stat(self, num_sample_sets, sample_set_sizes, @@ -3747,7 +3747,7 @@ tsk_treeseq_r(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_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) + tsk_size_t *result_size, double **result) { options |= TSK_STAT_POLARISED; // TODO: allow user to pick? return tsk_treeseq_two_locus_general_stat(self, num_sample_sets, sample_set_sizes, @@ -3785,7 +3785,7 @@ tsk_treeseq_Dz(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_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) + tsk_size_t *result_size, double **result) { return tsk_treeseq_two_locus_general_stat(self, num_sample_sets, sample_set_sizes, sample_sets, num_sample_sets, NULL, Dz_summary_func, norm_total_weighted, @@ -3819,7 +3819,7 @@ tsk_treeseq_pi2(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_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) + tsk_size_t *result_size, double **result) { return tsk_treeseq_two_locus_general_stat(self, num_sample_sets, sample_set_sizes, sample_sets, num_sample_sets, NULL, pi2_summary_func, norm_total_weighted, diff --git a/c/tskit/trees.h b/c/tskit/trees.h index a031166cd0..e3129e40d4 100644 --- a/c/tskit/trees.h +++ b/c/tskit/trees.h @@ -1011,37 +1011,37 @@ int tsk_treeseq_D(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_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); + tsk_size_t *result_size, double **result); int tsk_treeseq_D2(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_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); + tsk_size_t *result_size, double **result); int tsk_treeseq_r2(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_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); + tsk_size_t *result_size, double **result); int tsk_treeseq_D_prime(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_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); + tsk_size_t *result_size, double **result); int tsk_treeseq_r(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_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); + tsk_size_t *result_size, double **result); int tsk_treeseq_Dz(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_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); + tsk_size_t *result_size, double **result); int tsk_treeseq_pi2(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_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); + tsk_size_t *result_size, double **result); /* Three way sample set stats */ int tsk_treeseq_Y3(const tsk_treeseq_t *self, tsk_size_t num_sample_sets,