Join edges

jeromekelleher · Jan 19, 2023 · dd99990 · dd99990
1 parent 38fccba
commit dd99990
Showing 1 changed file with 28 additions and 16 deletions.
diff --git a/c/examples/multichrom_wright_fisher.c b/c/examples/multichrom_wright_fisher.c
@@ -41,15 +41,25 @@ free_tables(tsk_table_collection_t *tcs, int num_chunks)
 }
 
 static void
-dump_tables(tsk_table_collection_t *tcs, int num_chunks, const char *prefix)
+join_tables(tsk_table_collection_t *tcs, int num_chunks)
 {
     int j, ret;
 
-    /* FIXME need to write to prefix_%d.ts files, not just overwrite same file */
-    for (j = 0; j < num_chunks; j++) {
-        ret = tsk_table_collection_dump(&tcs[j], prefix, 0);
+    for (j = 1; j < num_chunks; 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 {
@@ -207,22 +217,22 @@ simulate(
             for (k = 0; k < num_chunks; k++) {
                 chunk_right = (k + 1) / (double) num_chunks;
                 /* FIXME not certain about the boundaries here, check! */
-                if (chunk_right < breakpoint ) {
-                    ret = tsk_edge_table_add_row(
-                        &tcs[k].edges, chunk_left, chunk_right, left_parent, child, NULL, 0);
+                if (chunk_right < breakpoint) {
+                    ret = tsk_edge_table_add_row(&tcs[k].edges, chunk_left, chunk_right,
+                        left_parent, child, NULL, 0);
                     check_tsk_error(ret);
                 } else if (breakpoint <= chunk_left) {
-                    ret = tsk_edge_table_add_row(
-                        &tcs[k].edges, chunk_left, chunk_right, right_parent, child, NULL, 0);
+                    ret = tsk_edge_table_add_row(&tcs[k].edges, chunk_left, chunk_right,
+                        right_parent, child, NULL, 0);
                     check_tsk_error(ret);
                 } else {
                     assert(chunk_left < breakpoint && breakpoint < chunk_right);
                     /* Breakpoint is within this chunk so need two edges */
-                    ret = tsk_edge_table_add_row(
-                        &tcs[k].edges, chunk_left, breakpoint, left_parent, child, NULL, 0);
+                    ret = tsk_edge_table_add_row(&tcs[k].edges, chunk_left, breakpoint,
+                        left_parent, child, NULL, 0);
                     check_tsk_error(ret);
-                    ret = tsk_edge_table_add_row(
-                        &tcs[k].edges, breakpoint, chunk_right, right_parent, child, NULL, 0);
+                    ret = tsk_edge_table_add_row(&tcs[k].edges, breakpoint, chunk_right,
+                        right_parent, child, NULL, 0);
                     check_tsk_error(ret);
                 }
                 chunk_left = chunk_right;
@@ -239,17 +249,19 @@ simulate(
 int
 main(int argc, char **argv)
 {
+    int ret;
     const int num_chunks = 8;
     tsk_table_collection_t tcs[num_chunks];
 
     if (argc != 6) {
-        errx(EXIT_FAILURE, "usage: N T simplify-interval output-prefix seed");
+        errx(EXIT_FAILURE, "usage: N T simplify-interval output seed");
     }
     srand((unsigned) atoi(argv[5]));
     init_tables(tcs, num_chunks);
     simulate(tcs, num_chunks, atoi(argv[1]), atoi(argv[2]), atoi(argv[3]));
-    /* TODO this should be join_tables */
-    dump_tables(tcs, num_chunks, argv[4]);
+    join_tables(tcs, num_chunks);
+    ret = tsk_table_collection_dump(&tcs[0], argv[4], 0);
+    check_tsk_error(ret);
     free_tables(tcs, num_chunks);
 
     return 0;