Basic multichrom simulation running

tskit-dev · Jan 10, 2023 · 1ecd671 · 1ecd671
1 parent 76bb39c
commit 1ecd671
Showing 1 changed file with 29 additions and 13 deletions.
diff --git a/c/examples/multichrom_wright_fisher.c b/c/examples/multichrom_wright_fisher.c
@@ -8,7 +8,7 @@
 
 #define check_tsk_error(val)                                                            \
     if (val < 0) {                                                                      \
-        errx(EXIT_FAILURE, "line %d: %s", __LINE__, tsk_strerror(val));                 \
+        errx(EXIT_FAILURE, "line %d: %s\n", __LINE__, tsk_strerror(val));               \
     }
 
 static void
@@ -42,9 +42,9 @@ free_tables(tsk_table_collection_t *tcs, int num_chroms)
 static void
 dump_tables(tsk_table_collection_t *tcs, int num_chroms, const char *prefix)
 {
-
     int j, ret;
 
+    /* FIXME need to write to prefix_%d.ts files, not just overwrite same file */
     for (j = 0; j < num_chroms; j++) {
         ret = tsk_table_collection_dump(&tcs[j], prefix, 0);
         check_tsk_error(ret);
@@ -58,6 +58,7 @@ simplify_tables(tsk_table_collection_t *tcs, int num_chroms, int *samples, int N
     const tsk_size_t num_nodes = tcs[0].nodes.num_rows;
     bool *delete_nodes = malloc(num_nodes * sizeof(*delete_nodes));
     tsk_id_t *node_id_map = malloc(num_nodes * sizeof(*node_id_map));
+    tsk_id_t *edge_child, *edge_parent;
 
     if (delete_nodes == NULL || node_id_map == NULL) {
         errx(EXIT_FAILURE, "Out of memory");
@@ -74,30 +75,42 @@ simplify_tables(tsk_table_collection_t *tcs, int num_chroms, int *samples, int N
         /* tsk_table_collection_print_state(&tcs[j], stdout); */
         check_tsk_error(ret);
         ret = tsk_table_collection_simplify(&tcs[j], samples, N,
-                TSK_SIMPLIFY_NO_FILTER_NODES|TSK_SIMPLIFY_NO_UPDATE_SAMPLE_FLAGS,
-                NULL);
+            TSK_SIMPLIFY_NO_FILTER_NODES | TSK_SIMPLIFY_NO_UPDATE_SAMPLE_FLAGS, NULL);
         check_tsk_error(ret);
         printf(" -> %lld\n", (long long) tcs[j].edges.num_rows);
     }
 
     for (j = 0; j < num_nodes; j++) {
         delete_nodes[j] = true;
     }
+    for (j = 0; j < N; j++) {
+        delete_nodes[samples[j]] = false;
+    }
 
     for (j = 0; j < num_chroms; j++) {
+        edge_child = tcs[j].edges.child;
+        edge_parent = tcs[j].edges.parent;
         for (k = 0; k < tcs[j].edges.num_rows; k++) {
-            delete_nodes[tcs[j].edges.child[k]] = false;
-            delete_nodes[tcs[j].edges.parent[k]] = false;
+            delete_nodes[edge_child[k]] = false;
+            delete_nodes[edge_parent[k]] = false;
         }
     }
-    /* tsk_node_table_delete_rows(&tcs[0].nodes, delete_nodes, 0, node_id_map); */
-
-    /* for (j = 0; j < tcs[0].nodes.num_rows; j++) { */
-    /*     printf("%d\t%d\n", j, keep_nodes[j]); */
-    /* } */
+    tsk_node_table_delete_rows(&tcs[0].nodes, delete_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;
+        for (k = 0; k < tcs[j].edges.num_rows; 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[j];
+        samples[j] = node_id_map[samples[j]];
     }
     free(delete_nodes);
     free(node_id_map);
@@ -135,6 +148,9 @@ simulate(
             child = tsk_node_table_add_row(
                 &tcs[0].nodes, 0, t, TSK_NULL, TSK_NULL, NULL, 0);
             check_tsk_error(child);
+            /* NOTE: this is a pretty unbiological idea of multiple
+             * chromosomes because each chromosome picks its parents
+             * independently.*/
             for (k = 0; k < num_chroms; k++) {
                 /* NOTE: the use of rand() is discouraged for
                  * research code and proper random number generator
@@ -164,7 +180,7 @@ simulate(
 int
 main(int argc, char **argv)
 {
-    const int num_chroms = 2;
+    const int num_chroms = 8;
     tsk_table_collection_t tcs[num_chroms];
 
     if (argc != 6) {