diff --git a/c/examples/haploid_wright_fisher.c b/c/examples/haploid_wright_fisher.c index 285e939d7f..5689fefc7a 100644 --- a/c/examples/haploid_wright_fisher.c +++ b/c/examples/haploid_wright_fisher.c @@ -92,6 +92,12 @@ main(int argc, char **argv) check_tsk_error(ret); srand((unsigned)atoi(argv[5])); simulate(&tables, atoi(argv[1]), atoi(argv[2]), atoi(argv[3])); + + /* Sort and index so that the result can be opened as a tree sequence */ + ret = tsk_table_collection_sort(&tables, NULL, 0); + check_tsk_error(ret); + ret = tsk_table_collection_build_index(&tables, 0); + check_tsk_error(ret); ret = tsk_table_collection_dump(&tables, argv[4], 0); check_tsk_error(ret); diff --git a/c/examples/multichrom_wright_fisher.c b/c/examples/multichrom_wright_fisher.c new file mode 100644 index 0000000000..9dafe5b524 --- /dev/null +++ b/c/examples/multichrom_wright_fisher.c @@ -0,0 +1,269 @@ +#include +#include +#include +#include +#include + +#include +#include + +#define check_tsk_error(val) \ + if (val < 0) { \ + errx(EXIT_FAILURE, "line %d: %s\n", __LINE__, tsk_strerror(val)); \ + } + +static void +init_tables(tsk_table_collection_t *tcs, int num_chroms) +{ + int j, ret; + + for (j = 0; j < num_chroms; j++) { + ret = tsk_table_collection_init(&tcs[j], 0); + check_tsk_error(ret); + if (j > 0) { + tsk_node_table_free(&tcs[j].nodes); + } + } +} + +static void +free_tables(tsk_table_collection_t *tcs, int num_chroms) +{ + int j; + + for (j = 0; j < num_chroms; j++) { + if (j > 0) { + /* Must not double free node table columns. */ + memset(&tcs[j].nodes, 0, sizeof(tcs[j].nodes)); + } + tsk_table_collection_free(&tcs[j]); + } +} + +static void +join_tables(tsk_table_collection_t *tcs, int num_chroms) +{ + int j, ret; + + for (j = 1; j < num_chroms; j++) { + ret = tsk_edge_table_extend( + &tcs[0].edges, &tcs[j].edges, tcs[j].edges.num_rows, NULL, 0); + check_tsk_error(ret); + } + /* Get all the squashable edges next to each other */ + ret = tsk_table_collection_sort(&tcs[0], NULL, 0); + check_tsk_error(ret); + ret = tsk_edge_table_squash(&tcs[0].edges); + check_tsk_error(ret); + /* We need to sort again after squash */ + ret = tsk_table_collection_sort(&tcs[0], NULL, 0); + check_tsk_error(ret); + ret = tsk_table_collection_build_index(&tcs[0], 0); + check_tsk_error(ret); +} + +struct chunk_work { + int chunk; + tsk_table_collection_t *tc; + int *samples; + int N; +}; + +void * +simplify_chunk(void *arg) +{ + int ret; + struct chunk_work *work = (struct chunk_work *) arg; + tsk_size_t edges_before = work->tc->edges.num_rows; + + ret = tsk_table_collection_sort(work->tc, NULL, 0); + check_tsk_error(ret); + ret = tsk_table_collection_simplify(work->tc, work->samples, work->N, + TSK_SIMPLIFY_NO_FILTER_NODES | TSK_SIMPLIFY_NO_UPDATE_SAMPLE_FLAGS, NULL); + check_tsk_error(ret); + /* NOTE: this printf makes helgrind complain */ + printf("\tchunk %d: %lld -> %lld\n", work->chunk, (long long) edges_before, + (long long) work->tc->edges.num_rows); + + return NULL; +} + +void +sort_and_simplify_all(tsk_table_collection_t *tcs, int num_chroms, int *samples, int N) +{ + int j, ret; + struct chunk_work work[num_chroms]; + pthread_t threads[num_chroms]; + + for (j = 1; j < num_chroms; j++) { + tcs[j].nodes = tcs[0].nodes; + } + + for (j = 0; j < num_chroms; j++) { + work[j].chunk = j; + work[j].tc = &tcs[j]; + work[j].samples = samples; + work[j].N = N; + + ret = pthread_create(&threads[j], NULL, simplify_chunk, (void *) &work[j]); + if (ret != 0) { + errx(EXIT_FAILURE, "Pthread create failed"); + } + /* simplify_chunk((void *) &work[j]); */ + } + for (j = 0; j < num_chroms; j++) { + ret = pthread_join(threads[j], NULL); + if (ret != 0) { + errx(EXIT_FAILURE, "Pthread join failed"); + } + } +} + +void +simplify_tables(tsk_table_collection_t *tcs, int num_chroms, int *samples, int N) +{ + int j, k, num_edges, ret; + const tsk_size_t num_nodes = tcs[0].nodes.num_rows; + tsk_bool_t *keep_nodes = malloc(num_nodes * sizeof(*keep_nodes)); + tsk_id_t *node_id_map = malloc(num_nodes * sizeof(*node_id_map)); + tsk_id_t *edge_child, *edge_parent; + + if (keep_nodes == NULL || node_id_map == NULL) { + errx(EXIT_FAILURE, "Out of memory"); + } + + printf("Simplify %lld nodes\n", (long long) tcs[0].nodes.num_rows); + sort_and_simplify_all(tcs, num_chroms, samples, N); + + for (j = 0; j < num_nodes; j++) { + keep_nodes[j] = false; + tcs[0].nodes.flags[j] &= (~TSK_NODE_IS_SAMPLE); + } + for (j = 0; j < N; j++) { + keep_nodes[samples[j]] = true; + tcs[0].nodes.flags[samples[j]] |= TSK_NODE_IS_SAMPLE; + } + + for (j = 0; j < num_chroms; j++) { + edge_child = tcs[j].edges.child; + edge_parent = tcs[j].edges.parent; + num_edges = tcs[j].edges.num_rows; + for (k = 0; k < num_edges; k++) { + keep_nodes[edge_child[k]] = true; + keep_nodes[edge_parent[k]] = true; + } + } + tsk_node_table_keep_rows(&tcs[0].nodes, keep_nodes, 0, node_id_map); + printf("\tdone: %lld nodes\n", (long long) tcs[0].nodes.num_rows); + + /* Remap node references */ + for (j = 0; j < num_chroms; j++) { + edge_child = tcs[j].edges.child; + edge_parent = tcs[j].edges.parent; + num_edges = tcs[j].edges.num_rows; + for (k = 0; k < num_edges; k++) { + edge_child[k] = node_id_map[edge_child[k]]; + edge_parent[k] = node_id_map[edge_parent[k]]; + } + ret = tsk_table_collection_check_integrity(&tcs[j], 0); + check_tsk_error(ret); + } + for (j = 0; j < N; j++) { + samples[j] = node_id_map[samples[j]]; + } + free(keep_nodes); + free(node_id_map); +} + +void +simulate( + tsk_table_collection_t *tcs, int num_chroms, int N, int T, int simplify_interval) +{ + tsk_id_t *buffer, *parents, *children, child, left_parent, right_parent; + bool left_is_first; + double chunk_left, chunk_right; + int ret, j, t, b, k; + + assert(simplify_interval != 0); // leads to division by zero + buffer = malloc(2 * N * sizeof(tsk_id_t)); + if (buffer == NULL) { + errx(EXIT_FAILURE, "Out of memory"); + } + for (k = 0; k < num_chroms; k++) { + tcs[k].sequence_length = num_chroms; + } + parents = buffer; + for (j = 0; j < N; j++) { + parents[j] + = tsk_node_table_add_row(&tcs[0].nodes, 0, T, TSK_NULL, TSK_NULL, NULL, 0); + check_tsk_error(parents[j]); + } + b = 0; + for (t = T - 1; t >= 0; t--) { + /* Alternate between using the first and last N values in the buffer */ + parents = buffer + (b * N); + b = (b + 1) % 2; + children = buffer + (b * N); + for (j = 0; j < N; j++) { + child = tsk_node_table_add_row( + &tcs[0].nodes, 0, t, TSK_NULL, TSK_NULL, NULL, 0); + check_tsk_error(child); + /* NOTE: the use of rand() is discouraged for + * research code and proper random number generator + * libraries should be preferred. + */ + left_parent = parents[(size_t)((rand() / (1. + RAND_MAX)) * N)]; + right_parent = parents[(size_t)((rand() / (1. + RAND_MAX)) * N)]; + left_is_first = rand() < 0.5; + chunk_left = 0.0; + for (k = 0; k < num_chroms; k++) { + chunk_right = chunk_left + rand() / (1. + RAND_MAX); + /* a very tiny chance that right and left are equal */ + if (chunk_right > chunk_left) { + ret = tsk_edge_table_add_row(&tcs[k].edges, chunk_left, chunk_right, + left_is_first ? left_parent : right_parent, child, NULL, 0); + check_tsk_error(ret); + } + chunk_left += 1.0; + if (chunk_right < chunk_left) { + ret = tsk_edge_table_add_row(&tcs[k].edges, chunk_right, chunk_left, + left_is_first ? right_parent : left_parent, child, NULL, 0); + check_tsk_error(ret); + } + } + children[j] = child; + } + if (t % simplify_interval == 0) { + simplify_tables(tcs, num_chroms, children, N); + } + } + /* Set the sample flags for final generation */ + for (j = 0; j < N; j++) { + tcs[0].nodes.flags[children[j]] = TSK_NODE_IS_SAMPLE; + } + free(buffer); +} + +int +main(int argc, char **argv) +{ + int ret; + int num_chroms; + + if (argc != 7) { + errx(EXIT_FAILURE, "usage: N T simplify-interval output seed num-chroms"); + } + + num_chroms = atoi(argv[6]); + tsk_table_collection_t tcs[num_chroms]; + + srand((unsigned) atoi(argv[5])); + init_tables(tcs, num_chroms); + simulate(tcs, num_chroms, atoi(argv[1]), atoi(argv[2]), atoi(argv[3])); + join_tables(tcs, num_chroms); + ret = tsk_table_collection_dump(&tcs[0], argv[4], 0); + check_tsk_error(ret); + free_tables(tcs, num_chroms); + + return 0; +} diff --git a/c/examples/multichrom_wright_fisher_singlethreaded.c b/c/examples/multichrom_wright_fisher_singlethreaded.c new file mode 100644 index 0000000000..649cd408d3 --- /dev/null +++ b/c/examples/multichrom_wright_fisher_singlethreaded.c @@ -0,0 +1,120 @@ +#include +#include +#include +#include +#include + +#include + +#define check_tsk_error(val) \ + if (val < 0) { \ + errx(EXIT_FAILURE, "line %d: %s\n", __LINE__, tsk_strerror(val)); \ + } + +void +simulate( + tsk_table_collection_t *tables, int num_chroms, int N, int T, int simplify_interval) +{ + tsk_id_t *buffer, *parents, *children, child, left_parent, right_parent; + bool left_is_first; + double chunk_left, chunk_right; + int ret, j, t, b, k; + + assert(simplify_interval != 0); // leads to division by zero + buffer = malloc(2 * N * sizeof(tsk_id_t)); + if (buffer == NULL) { + errx(EXIT_FAILURE, "Out of memory"); + } + tables->sequence_length = num_chroms; + parents = buffer; + for (j = 0; j < N; j++) { + parents[j] + = tsk_node_table_add_row(&tables->nodes, 0, T, TSK_NULL, TSK_NULL, NULL, 0); + check_tsk_error(parents[j]); + } + b = 0; + for (t = T - 1; t >= 0; t--) { + /* Alternate between using the first and last N values in the buffer */ + parents = buffer + (b * N); + b = (b + 1) % 2; + children = buffer + (b * N); + for (j = 0; j < N; j++) { + child = tsk_node_table_add_row( + &tables->nodes, 0, t, TSK_NULL, TSK_NULL, NULL, 0); + check_tsk_error(child); + /* NOTE: the use of rand() is discouraged for + * research code and proper random number generator + * libraries should be preferred. + */ + left_parent = parents[(size_t)((rand() / (1. + RAND_MAX)) * N)]; + right_parent = parents[(size_t)((rand() / (1. + RAND_MAX)) * N)]; + left_is_first = rand() < 0.5; + chunk_left = 0.0; + for (k = 0; k < num_chroms; k++) { + chunk_right = chunk_left + rand() / (1. + RAND_MAX); + /* a very tiny chance that right and left are equal */ + if (chunk_right > chunk_left) { + ret = tsk_edge_table_add_row(&tables->edges, chunk_left, chunk_right, + left_is_first ? left_parent : right_parent, child, NULL, 0); + check_tsk_error(ret); + } + chunk_left += 1.0; + if (chunk_right < chunk_left) { + ret = tsk_edge_table_add_row(&tables->edges, chunk_right, chunk_left, + left_is_first ? right_parent : left_parent, child, NULL, 0); + check_tsk_error(ret); + } + } + children[j] = child; + } + if (t % simplify_interval == 0) { + printf("Simplify at generation %lld: (%lld nodes %lld edges)", + (long long) t, + (long long) tables->nodes.num_rows, + (long long) tables->edges.num_rows); + /* Note: Edges must be sorted for simplify to work, and we use a brute force + * approach of sorting each time here for simplicity. This is inefficient. */ + ret = tsk_table_collection_sort(tables, NULL, 0); + check_tsk_error(ret); + ret = tsk_table_collection_simplify(tables, children, N, 0, NULL); + check_tsk_error(ret); + printf(" -> (%lld nodes %lld edges)\n", + (long long) tables->nodes.num_rows, + (long long) tables->edges.num_rows); + for (j = 0; j < N; j++) { + children[j] = j; + } + } + } + /* Set the sample flags for final generation */ + for (j = 0; j < N; j++) { + tables->nodes.flags[children[j]] = TSK_NODE_IS_SAMPLE; + } + free(buffer); +} + +int +main(int argc, char **argv) +{ + int ret; + tsk_table_collection_t tables; + + if (argc != 7) { + errx(EXIT_FAILURE, "usage: N T simplify-interval output seed num-chroms"); + } + ret = tsk_table_collection_init(&tables, 0); + check_tsk_error(ret); + srand((unsigned)atoi(argv[5])); + simulate(&tables, atoi(argv[6]), atoi(argv[1]), atoi(argv[2]), atoi(argv[3])); + + /* Sort and index so that the result can be opened as a tree sequence */ + ret = tsk_table_collection_sort(&tables, NULL, 0); + check_tsk_error(ret); + ret = tsk_table_collection_build_index(&tables, 0); + check_tsk_error(ret); + ret = tsk_table_collection_dump(&tables, argv[4], 0); + check_tsk_error(ret); + + tsk_table_collection_free(&tables); + return 0; +} diff --git a/c/meson.build b/c/meson.build index 91df7e539d..f5c1a0f585 100644 --- a/c/meson.build +++ b/c/meson.build @@ -122,5 +122,13 @@ if not meson.is_subproject() executable('haploid_wright_fisher', sources: ['examples/haploid_wright_fisher.c'], link_with: [tskit_lib], dependencies: lib_deps) + executable('multichrom_wright_fisher_singlethreaded', + sources: ['examples/multichrom_wright_fisher_singlethreaded.c'], + link_with: [tskit_lib], dependencies: lib_deps) + + thread_dep = dependency('threads') + executable('multichrom_wright_fisher', + sources: ['examples/multichrom_wright_fisher.c'], + link_with: [tskit_lib], dependencies: [m_dep, kastore_dep, thread_dep]) endif endif diff --git a/docs/c-api.rst b/docs/c-api.rst index 460ce74b58..83f19ca9c1 100644 --- a/docs/c-api.rst +++ b/docs/c-api.rst @@ -920,3 +920,31 @@ stored in the file ``no_mutations.trees``:: $ ./build/streaming < no_mutations.trees > /dev/null Tree sequence 0 had 0 mutations Tree sequence 1 had 0 mutations + +------------------------------------ +Parallel, multichromosome simulation +------------------------------------ + +A substantial bottleneck in forwards simulations using tree sequences +is *simplification*. This is therefore a natural target for parallelization. +The potential for breaking up a chromosome into discrete chunks that +are separately parallelized is limited, however, since any edge +that extends across the boundary between two chunks is split; +thus creating more work. +However, distinct chromosomes provide a natural target: +the edge tables describing inheritance for each chromosome can be +independently simplified, as long as the fact that they all refer to +the same set of nodes. +This simulation keeps each chromosome in a separate tree sequence, +but they essentially share a common node table; +the :c:macro:`TSK_SIMPLIFY_NO_FILTER_NODES` flag is used so that +each call to :c:func:`tsk_table_collection_simplify` does not +change the common node table. +Afterwards, we iterate though the edge tables to determine which +nodes need to be retained, and use +:c:func:`tsk_node_table_keep_rows` to remove unused nodes. + + +.. literalinclude:: ../c/examples/multichrom_wright_fisher.c + :language: c +