Skip to content

Commit

Permalink
Allocate the result array within the site stat fn
Browse files Browse the repository at this point in the history 
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.
  • Loading branch information
@lkirk
lkirk committed Jul 18, 2023
1 parent 0e380df commit f4ae923
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 28 deletions.
42 changes: 21 additions & 21 deletions c/tskit/trees.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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));

Expand All @@ -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++;
}
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
14 changes: 7 additions & 7 deletions c/tskit/trees.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down

0 comments on commit f4ae923

Please sign in to comment.