Decoding haplotypes without trees

c/examples/Makefile

@@ -23,7 +23,7 @@ TSKIT_SOURCE=../tskit/*.c ../subprojects/kastore/kastore.c
 targets = api_structure error_handling \
 	haploid_wright_fisher streaming \
 	tree_iteration tree_traversal \
-	take_ownership
+	take_ownership haplotype_benchmark
 
 all: $(targets)
 
@@ -32,4 +32,3 @@ $(targets): %: %.c
 
 clean:
 	rm -f $(targets)
-

c/examples/haplotype_benchmark.c

@@ -0,0 +1,90 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <stdint.h>
+#include <time.h>
+
+#include <tskit.h>
+#include <tskit/genotypes.h>
+#include <tskit/tables.h>
+
+#define CHECK_TSK(err)                                                                  \
+    do {                                                                                \
+        if ((err) < 0) {                                                                \
+            fprintf(stderr, "Error: line %d: %s\n", __LINE__, tsk_strerror(err));       \
+            exit(EXIT_FAILURE);                                                         \
+        }                                                                               \
+    } while (0)
+
+#define NUM_ITERATIONS 1
+#define MAX_BENCHMARK_NODES 500
+
+int
+main(int argc, char **argv)
+{
+    int ret;
+    tsk_table_collection_t tables;
+    tsk_treeseq_t treeseq;
+    tsk_haplotype_t haplotype_decoder;
+    int8_t *haplotype = NULL;
+    double elapsed_seconds;
+    clock_t start_clock, end_clock;
+    uint64_t checksum = 0;
+
+    const char *filename = "../../simulated_chrom_21_100k.ts";
+    if (argc > 1) {
+        filename = argv[1];
+    }
+
+    ret = tsk_table_collection_init(&tables, 0);
+    CHECK_TSK(ret);
+
+    ret = tsk_table_collection_load(&tables, filename, 0);
+    CHECK_TSK(ret);
+
+    ret = tsk_treeseq_init(&treeseq, &tables, 0);
+    CHECK_TSK(ret);
+
+    tsk_size_t num_nodes = tsk_treeseq_get_num_nodes(&treeseq);
+    tsk_size_t num_sites = tsk_treeseq_get_num_sites(&treeseq);
+    if (num_sites == 0) {
+        fprintf(stderr, "Tree sequence has no sites\n");
+        exit(EXIT_FAILURE);
+    }
+
+    tsk_id_t node_limit
+        = (tsk_id_t) (num_nodes < MAX_BENCHMARK_NODES ? num_nodes : MAX_BENCHMARK_NODES);
+
+    ret = tsk_haplotype_init(&haplotype_decoder, &treeseq, 0, (tsk_id_t) num_sites);
+    CHECK_TSK(ret);
+
+    haplotype = malloc(num_sites * sizeof(*haplotype));
+    if (haplotype == NULL) {
+        fprintf(stderr, "Failed to allocate haplotype buffer\n");
+        exit(EXIT_FAILURE);
+    }
+
+    start_clock = clock();
+    for (int iter = 0; iter < NUM_ITERATIONS; iter++) {
+        for (tsk_id_t node = 0; node < node_limit; node++) {
+            ret = tsk_haplotype_decode(&haplotype_decoder, node, haplotype);
+            CHECK_TSK(ret);
+            for (tsk_id_t site = 0; site < (tsk_id_t) num_sites; site++) {
+                checksum += (uint64_t) haplotype[site];
+            }
+        }
+    }
+    end_clock = clock();
+
+    elapsed_seconds = (double) (end_clock - start_clock) / CLOCKS_PER_SEC;
+
+    printf("Loaded tree sequence from %s\n", filename);
+    printf("Decoded %d iterations over %lld nodes × %lld sites in %.3f seconds\n",
+        NUM_ITERATIONS, (long long) node_limit, (long long) num_sites, elapsed_seconds);
+    printf("Checksummed haplotypes: %llu\n", (unsigned long long) checksum);
+
+    free(haplotype);
+    tsk_haplotype_free(&haplotype_decoder);
+    tsk_treeseq_free(&treeseq);
+    tsk_table_collection_free(&tables);
+    return EXIT_SUCCESS;
+}

c/meson.build

@@ -113,6 +113,9 @@ if not meson.is_subproject()
       executable('tree_traversal',
           sources: ['examples/tree_traversal.c'], 
           link_with: [tskit_lib], dependencies: lib_deps)
+      executable('haplotype_benchmark',
+          sources: ['examples/haplotype_benchmark.c'],
+          link_with: [tskit_lib], dependencies: lib_deps)
       executable('streaming',
           sources: ['examples/streaming.c'], 
           link_with: [tskit_lib], dependencies: lib_deps)