Skip to content

Implement split_edges #2296

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
May 25, 2022
Merged

Implement split_edges #2296

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
203 changes: 203 additions & 0 deletions c/tests/test_trees.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7249,6 +7249,206 @@ test_reference_sequence(void)
tsk_table_collection_free(&tables);
}

static void
test_split_edges_no_populations(void)
{
int ret;
tsk_treeseq_t ts, split_ts;
tsk_table_collection_t tables;
tsk_id_t new_nodes[] = { 9, 10, 11 };
tsk_size_t num_new_nodes = 3;
const char *metadata = "some metadata";
tsk_size_t j;
tsk_node_t node;
double time = 0.09;
tsk_id_t ret_id;

tsk_treeseq_from_text(&ts, 10, paper_ex_nodes, paper_ex_edges, NULL, paper_ex_sites,
paper_ex_mutations, paper_ex_individuals, NULL, 0);

ret_id = tsk_table_collection_copy(ts.tables, &tables, 0);
CU_ASSERT_EQUAL_FATAL(ret_id, 0);
tsk_treeseq_free(&ts);
ret_id = tsk_population_table_add_row(&tables.populations, NULL, 0);
CU_ASSERT_EQUAL_FATAL(ret_id, 0);
ret = tsk_table_collection_compute_mutation_times(&tables, NULL, 0);
CU_ASSERT_EQUAL_FATAL(ret_id, 0);
ret_id = tsk_treeseq_init(&ts, &tables, 0);
CU_ASSERT_EQUAL_FATAL(ret_id, 0);

/* NOTE: haven't worked out the exact IDs on the branches here, just
* for illustration.

0.25┊ 8 ┊ ┊ ┊
┊ ┏━┻━┓ ┊ ┊ ┊
0.20┊ ┃ ┃ ┊ ┊ 7 ┊
┊ ┃ ┃ ┊ ┊ ┏━┻━┓ ┊
0.17┊ 6 ┃ ┊ 6 ┊ ┃ ┃ ┊
┊ ┏━┻┓ ┃ ┊ ┏━┻━┓ ┊ ┃ ┃ ┊
0.09┊ 9 5 10┊ 9 5 ┊ 11 5 ┊
┊ ┃ ┏┻┓ ┃ ┊ ┃ ┏━┻┓ ┊ ┃ ┏━┻┓ ┊
0.07┊ ┃ ┃ ┃ ┃ ┊ ┃ ┃ 4 ┊ ┃ ┃ 4 ┊
┊ ┃ ┃ ┃ ┃ ┊ ┃ ┃ ┏┻┓ ┊ ┃ ┃ ┏┻┓ ┊
0.00┊ 0 1 3 2 ┊ 0 1 2 3 ┊ 0 1 2 3 ┊
0.00 2.00 7.00 10.00
*/
ret = tsk_treeseq_split_edges(
&ts, time, 1234, 0, metadata, strlen(metadata), 0, &split_ts);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL(tsk_treeseq_get_num_trees(&split_ts), 3);
CU_ASSERT_EQUAL(tsk_treeseq_get_num_nodes(&split_ts), 12);

for (j = 0; j < num_new_nodes; j++) {
ret = tsk_treeseq_get_node(&split_ts, new_nodes[j], &node);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL(node.time, time);
CU_ASSERT_EQUAL(node.flags, 1234);
CU_ASSERT_EQUAL(node.individual, TSK_NULL);
CU_ASSERT_EQUAL(node.population, 0);
CU_ASSERT_EQUAL(node.metadata_length, strlen(metadata));
CU_ASSERT_EQUAL(strncmp(node.metadata, metadata, strlen(metadata)), 0);
}
tsk_treeseq_free(&split_ts);

/* And again with imputed population value */
ret = tsk_treeseq_split_edges(&ts, time, 1234, 0, metadata, strlen(metadata),
TSK_SPLIT_EDGES_IMPUTE_POPULATION, &split_ts);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL(tsk_treeseq_get_num_trees(&split_ts), 3);
CU_ASSERT_EQUAL(tsk_treeseq_get_num_nodes(&split_ts), 12);

for (j = 0; j < num_new_nodes; j++) {
ret = tsk_treeseq_get_node(&split_ts, new_nodes[j], &node);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL(node.time, time);
CU_ASSERT_EQUAL(node.flags, 1234);
CU_ASSERT_EQUAL(node.individual, TSK_NULL);
CU_ASSERT_EQUAL(node.population, TSK_NULL);
CU_ASSERT_EQUAL(node.metadata_length, strlen(metadata));
CU_ASSERT_EQUAL(strncmp(node.metadata, metadata, strlen(metadata)), 0);
}
tsk_treeseq_free(&split_ts);

tsk_table_collection_free(&tables);
tsk_treeseq_free(&ts);
}

static void
test_split_edges_populations(void)
{
int ret;
tsk_treeseq_t ts, split_ts;
tsk_table_collection_t tables;
double time = 0.5;
tsk_node_t node;
tsk_id_t valid_pops[] = { -1, 0, 1 };
tsk_id_t num_valid_pops = 3;
tsk_id_t j, population, ret_id;

ret = tsk_table_collection_init(&tables, 0);
CU_ASSERT_EQUAL_FATAL(ret, 0);
tables.sequence_length = 1;

ret_id = tsk_population_table_add_row(&tables.populations, NULL, 0);
CU_ASSERT_EQUAL_FATAL(ret_id, 0);
ret_id = tsk_population_table_add_row(&tables.populations, NULL, 0);
CU_ASSERT_EQUAL_FATAL(ret_id, 1);
ret_id = tsk_node_table_add_row(&tables.nodes, 0, 0, 0, TSK_NULL, NULL, 0);
CU_ASSERT_EQUAL_FATAL(ret_id, 0);
ret_id = tsk_node_table_add_row(&tables.nodes, 0, 1, 1, TSK_NULL, NULL, 0);
CU_ASSERT_EQUAL_FATAL(ret_id, 1);
ret_id = tsk_edge_table_add_row(&tables.edges, 0, 1, 1, 0, NULL, 0);
CU_ASSERT_EQUAL_FATAL(ret_id, 0);

ret = tsk_treeseq_init(&ts, &tables, TSK_TS_INIT_BUILD_INDEXES);
CU_ASSERT_EQUAL_FATAL(ret, 0);

for (j = 0; j < num_valid_pops; j++) {
population = valid_pops[j];
ret = tsk_treeseq_split_edges(&ts, time, 0, population, NULL, 0, 0, &split_ts);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL(tsk_treeseq_get_num_trees(&split_ts), 1);
CU_ASSERT_EQUAL(tsk_treeseq_get_num_nodes(&split_ts), 3);
ret = tsk_treeseq_get_node(&split_ts, 2, &node);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL(node.population, population);
tsk_treeseq_free(&split_ts);

ret = tsk_treeseq_split_edges(&ts, time, 0, population, NULL, 0,
TSK_SPLIT_EDGES_IMPUTE_POPULATION, &split_ts);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL(tsk_treeseq_get_num_trees(&split_ts), 1);
CU_ASSERT_EQUAL(tsk_treeseq_get_num_nodes(&split_ts), 3);
ret = tsk_treeseq_get_node(&split_ts, 2, &node);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL(node.population, 0);
tsk_treeseq_free(&split_ts);
}

tsk_table_collection_free(&tables);
tsk_treeseq_free(&ts);
}

static void
test_split_edges_errors(void)
{
int ret;
tsk_treeseq_t ts, split_ts;
tsk_table_collection_t tables;
double time = 0.5;
tsk_id_t invalid_pops[] = { -2, 2, 3 };
tsk_id_t num_invalid_pops = 3;
tsk_id_t j, population, ret_id;

ret = tsk_table_collection_init(&tables, 0);
CU_ASSERT_EQUAL_FATAL(ret, 0);
tables.sequence_length = 1;

ret_id = tsk_population_table_add_row(&tables.populations, NULL, 0);
CU_ASSERT_EQUAL_FATAL(ret_id, 0);
ret_id = tsk_population_table_add_row(&tables.populations, NULL, 0);
CU_ASSERT_EQUAL_FATAL(ret_id, 1);
ret_id = tsk_node_table_add_row(&tables.nodes, 0, 0, 0, TSK_NULL, NULL, 0);
CU_ASSERT_EQUAL_FATAL(ret_id, 0);
ret_id = tsk_node_table_add_row(&tables.nodes, 0, 1, 1, TSK_NULL, NULL, 0);
CU_ASSERT_EQUAL_FATAL(ret_id, 1);
ret_id = tsk_edge_table_add_row(&tables.edges, 0, 1, 1, 0, NULL, 0);
CU_ASSERT_EQUAL_FATAL(ret_id, 0);

ret = tsk_treeseq_init(&ts, &tables, TSK_TS_INIT_BUILD_INDEXES);
CU_ASSERT_EQUAL_FATAL(ret, 0);

ret = tsk_treeseq_split_edges(
&ts, TSK_UNKNOWN_TIME, 0, TSK_NULL, NULL, 0, 0, &split_ts);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_TIME_NONFINITE);

for (j = 0; j < num_invalid_pops; j++) {
population = invalid_pops[j];
ret = tsk_treeseq_split_edges(&ts, time, 0, population, NULL, 0, 0, &split_ts);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_POPULATION_OUT_OF_BOUNDS);
tsk_treeseq_free(&split_ts);

/* We always check population values, even if they aren't used */
ret = tsk_treeseq_split_edges(&ts, time, 0, population, NULL, 0,
TSK_SPLIT_EDGES_IMPUTE_POPULATION, &split_ts);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_POPULATION_OUT_OF_BOUNDS);
tsk_treeseq_free(&split_ts);
}
tsk_treeseq_free(&ts);

ret_id
= tsk_migration_table_add_row(&tables.migrations, 0, 1, 0, 0, 1, 1.0, NULL, 0);
CU_ASSERT_EQUAL_FATAL(ret_id, 0);
ret = tsk_treeseq_init(&ts, &tables, TSK_TS_INIT_BUILD_INDEXES);
CU_ASSERT_EQUAL_FATAL(ret, 0);
ret = tsk_treeseq_split_edges(&ts, time, 0, population, NULL, 0, 0, &split_ts);
CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_MIGRATIONS_NOT_SUPPORTED);
tsk_treeseq_free(&split_ts);

tsk_table_collection_free(&tables);
tsk_treeseq_free(&ts);
}

static void
test_init_take_ownership_no_edge_metadata(void)
{
Expand Down Expand Up @@ -7449,6 +7649,9 @@ main(int argc, char **argv)
{ "test_tree_sequence_metadata", test_tree_sequence_metadata },
{ "test_time_uncalibrated", test_time_uncalibrated },
{ "test_reference_sequence", test_reference_sequence },
{ "test_split_edges_no_populations", test_split_edges_no_populations },
{ "test_split_edges_populations", test_split_edges_populations },
{ "test_split_edges_errors", test_split_edges_errors },
{ "test_init_take_ownership_no_edge_metadata",
test_init_take_ownership_no_edge_metadata },
{ NULL, NULL },
Expand Down
1 change: 1 addition & 0 deletions c/tskit/tables.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4252,6 +4252,7 @@ int tsk_provenance_table_takeset_columns(tsk_provenance_table_t *self,
tsk_size_t *record_offset);

bool tsk_table_collection_has_reference_sequence(const tsk_table_collection_t *self);

int tsk_reference_sequence_init(tsk_reference_sequence_t *self, tsk_flags_t options);
int tsk_reference_sequence_free(tsk_reference_sequence_t *self);
bool tsk_reference_sequence_is_null(const tsk_reference_sequence_t *self);
Expand Down
120 changes: 120 additions & 0 deletions c/tskit/trees.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3253,6 +3253,126 @@ tsk_treeseq_simplify(const tsk_treeseq_t *self, const tsk_id_t *samples,
return ret;
}

int TSK_WARN_UNUSED
tsk_treeseq_split_edges(const tsk_treeseq_t *self, double time, tsk_flags_t flags,
tsk_id_t population, const char *metadata, tsk_size_t metadata_length,
tsk_flags_t options, tsk_treeseq_t *output)
{
int ret = 0;
tsk_table_collection_t *tables = tsk_malloc(sizeof(*tables));
const double *restrict node_time = self->tables->nodes.time;
const tsk_id_t *restrict node_population = self->tables->nodes.population;
const tsk_size_t num_edges = self->tables->edges.num_rows;
const tsk_size_t num_mutations = self->tables->mutations.num_rows;
tsk_id_t *split_edge = tsk_malloc(num_edges * sizeof(*split_edge));
tsk_id_t j, u, mapped_node, ret_id;
double mutation_time;
tsk_edge_t edge;
tsk_mutation_t mutation;
tsk_bookmark_t sort_start;
bool impute_population = options & TSK_SPLIT_EDGES_IMPUTE_POPULATION;

memset(output, 0, sizeof(*output));
if (split_edge == NULL) {
ret = TSK_ERR_NO_MEMORY;
goto out;
}
ret = tsk_treeseq_copy_tables(self, tables, 0);
if (ret != 0) {
goto out;
}
if (tables->migrations.num_rows > 0) {
ret = TSK_ERR_MIGRATIONS_NOT_SUPPORTED;
goto out;
}
if (population < -1 || population >= (tsk_id_t) self->tables->populations.num_rows) {
ret = TSK_ERR_POPULATION_OUT_OF_BOUNDS;
goto out;
}
if (!tsk_isfinite(time)) {
ret = TSK_ERR_TIME_NONFINITE;
goto out;
}

tsk_edge_table_clear(&tables->edges);
tsk_memset(split_edge, TSK_NULL, num_edges * sizeof(*split_edge));

for (j = 0; j < (tsk_id_t) num_edges; j++) {
/* Would prefer to use tsk_edge_table_get_row_unsafe, but it's
* currently static to tables.c */
ret = tsk_edge_table_get_row(&self->tables->edges, j, &edge);
tsk_bug_assert(ret == 0);
if (node_time[edge.child] < time && time < node_time[edge.parent]) {
if (impute_population) {
population = TSK_NULL;
if (node_population[edge.child] != TSK_NULL) {
population = node_population[edge.child];
}
}
u = tsk_node_table_add_row(&tables->nodes, flags, time, population, TSK_NULL,
metadata, metadata_length);
if (u < 0) {
ret = (int) u;
goto out;
}
ret_id = tsk_edge_table_add_row(&tables->edges, edge.left, edge.right, u,
edge.child, edge.metadata, edge.metadata_length);
if (ret_id < 0) {
ret = (int) ret_id;
goto out;
}
edge.child = u;
split_edge[j] = u;
}
ret_id = tsk_edge_table_add_row(&tables->edges, edge.left, edge.right,
edge.parent, edge.child, edge.metadata, edge.metadata_length);
if (ret_id < 0) {
ret = (int) ret_id;
goto out;
}
}

for (j = 0; j < (tsk_id_t) num_mutations; j++) {
/* Note: we could speed this up a bit by accessing the local
* memory for mutations directly. */
ret = tsk_treeseq_get_mutation(self, j, &mutation);
tsk_bug_assert(ret == 0);
mapped_node = TSK_NULL;
if (mutation.edge != TSK_NULL) {
mapped_node = split_edge[mutation.edge];
}
mutation_time = tsk_is_unknown_time(mutation.time) ? node_time[mutation.node]
: mutation.time;
if (mapped_node != TSK_NULL && mutation_time >= time) {
/* Update the column in-place to save a bit of time. */
tables->mutations.node[j] = mapped_node;
}
}

/* Skip mutations and sites as they haven't been altered */
/* Note we can probably optimise the edge sort a bit here also by
* reasoning about when the first edge gets altered in the table.
*/
memset(&sort_start, 0, sizeof(sort_start));
sort_start.sites = tables->sites.num_rows;
sort_start.mutations = tables->mutations.num_rows;
ret = tsk_table_collection_sort(tables, &sort_start, 0);
if (ret != 0) {
goto out;
}

ret = tsk_treeseq_init(
output, tables, TSK_TS_INIT_BUILD_INDEXES | TSK_TAKE_OWNERSHIP);
tables = NULL;
out:
if (tables != NULL) {
tsk_table_collection_free(tables);
tsk_safe_free(tables);
}
tsk_safe_free(split_edge);
return ret;
}

/* ======================================================== *
* Tree
* ======================================================== */
Expand Down
6 changes: 6 additions & 0 deletions c/tskit/trees.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -882,6 +882,12 @@ int tsk_treeseq_simplify(const tsk_treeseq_t *self, const tsk_id_t *samples,

/** @} */

#define TSK_SPLIT_EDGES_IMPUTE_POPULATION (1 << 1)

int tsk_treeseq_split_edges(const tsk_treeseq_t *self, double time, tsk_flags_t flags,
tsk_id_t population, const char *metadata, tsk_size_t metadata_length,
tsk_flags_t options, tsk_treeseq_t *output);

bool tsk_treeseq_has_reference_sequence(const tsk_treeseq_t *self);

int tsk_treeseq_kc_distance(const tsk_treeseq_t *self, const tsk_treeseq_t *other,
Expand Down
1 change: 1 addition & 0 deletions docs/python-api.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -218,6 +218,7 @@ which perform the same actions but modify the {class}`TableCollection` in place.
TreeSequence.delete_intervals
TreeSequence.delete_sites
TreeSequence.trim
TreeSequence.split_edges
```

(sec_python_api_tree_sequences_ibd)=
Expand Down
7 changes: 6 additions & 1 deletion python/CHANGELOG.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,8 @@

**Changes**

- ``VcfWriter.write`` now prints the site ID of variants in the ID field of the output VCF files.
- ``VcfWriter.write`` now prints the site ID of variants in the ID field of the
output VCF files.
(:user:`roohy`, :issue:`2103`, :pr:`2107`)

- Make dumping of tables and tree sequences to disk a zero-copy operation.
Expand All @@ -21,6 +22,10 @@
edge that the mutation falls on.
(:user:`jeromekelleher`, :issue:`685`, :pr:`2279`).

- Add the ``TreeSequence.split_edges`` operation which inserts nodes into
edges at a specific time.
(:user:`jeromekelleher`, :issue:`2276`, :pr:`2296`).

**Breaking Changes**

- The JSON metadata codec now interprets the empty string as an empty object. This means
Expand Down
Loading