Skip to content

Commit

Permalink
more restricts
Browse files Browse the repository at this point in the history
  • Loading branch information
petrelharp committed Sep 25, 2024
1 parent 39174d2 commit 93a147a
Showing 1 changed file with 21 additions and 15 deletions.
36 changes: 21 additions & 15 deletions c/tskit/trees.c
Original file line number Diff line number Diff line change
Expand Up @@ -10108,54 +10108,60 @@ tsk_matvec_calculator_write_output(tsk_matvec_calculator_t *self, double *restri
tsk_id_t u;
tsk_size_t j, k;
tsk_size_t n = tsk_treeseq_get_num_samples(self->ts);
const tsk_size_t num_weights = self->num_weights;
const double tree_left = self->tree_left;
double *u_row, *out_row;
double *out_means = tsk_malloc(self->num_weights * sizeof(*out_means));
double *out_means = tsk_malloc(num_weights * sizeof(*out_means));
const tsk_id_t *restrict parent = self->parent;
const double *restrict nodes_time = self->ts->tables->nodes.time;
double *restrict x = self->x;
double *restrict w = self->w;
double *restrict v = self->v;
const tsk_id_t *restrict samples = self->ts->samples;

if (out_means == NULL) {
ret = TSK_ERR_NO_MEMORY;
goto out;

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

View check run for this annotation

Codecov / codecov/patch

c/tskit/trees.c#L10123-L10124

Added lines #L10123 - L10124 were not covered by tests
}

for (j = 0; j < n; j++) {
out_row = GET_2D_ROW(y, self->num_weights, j);
u = self->ts->samples[j];
out_row = GET_2D_ROW(y, num_weights, j);
u = samples[j];
while (u != TSK_NULL) {
if (self->x[u] != self->tree_left) {
tsk_matvec_calculator_add_z(u, parent[u], self->tree_left, self->x,
self->num_weights, self->w, self->v, nodes_time);
if (x[u] != tree_left) {
tsk_matvec_calculator_add_z(
u, parent[u], tree_left, x, num_weights, w, v, nodes_time);
}
u_row = GET_2D_ROW(self->v, self->num_weights, u);
for (k = 0; k < self->num_weights; k++) {
u_row = GET_2D_ROW(v, num_weights, u);
for (k = 0; k < num_weights; k++) {
out_row[k] += u_row[k];
}
u = parent[u];
}
}

if (!(self->options & TSK_STAT_NONCENTRED)) {
for (k = 0; k < self->num_weights; k++) {
for (k = 0; k < num_weights; k++) {
out_means[k] = 0.0;
}
for (j = 0; j < n; j++) {
out_row = GET_2D_ROW(y, self->num_weights, j);
for (k = 0; k < self->num_weights; k++) {
out_row = GET_2D_ROW(y, num_weights, j);
for (k = 0; k < num_weights; k++) {
out_means[k] += out_row[k];
}
}
for (k = 0; k < self->num_weights; k++) {
for (k = 0; k < num_weights; k++) {
out_means[k] /= (double) n;
}
for (j = 0; j < n; j++) {
out_row = GET_2D_ROW(y, self->num_weights, j);
for (k = 0; k < self->num_weights; k++) {
out_row = GET_2D_ROW(y, num_weights, j);
for (k = 0; k < num_weights; k++) {
out_row[k] -= out_means[k];
}
}
}
/* zero out v */
tsk_memset(self->v, 0, self->num_nodes * self->num_weights * sizeof(*self->v));
tsk_memset(self->v, 0, self->num_nodes * num_weights * sizeof(*self->v));
out:
tsk_safe_free(out_means);
return ret;
Expand Down

0 comments on commit 93a147a

Please sign in to comment.