tskit/c/examples/cpp_sorting_example.cpp at 7b0cf066809c92c2853bddc8c301d83633a3485c · benjeffery/tskit · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
#include <cstddef>
#include <vector>
#include <algorithm>
#include <stdexcept>
#include <sstream>
#include <iostream>
#include <type_traits>
#include <tskit.h>

static void
handle_tskit_return_code(int code)
{
    if (code != 0) {
        std::ostringstream o;
        o << tsk_strerror(code);
        throw std::runtime_error(o.str());
    }
}

struct edge_plus_time {
    double time;
    tsk_id_t parent, child;
    double left, right;
};

int
sort_edges(tsk_table_sorter_t *sorter, tsk_size_t start)
{
    if (sorter->tables->edges.metadata_length != 0) {
        throw std::invalid_argument(
            "the sorter does not currently handle edge metadata");
    }
    if (start != 0) {
        throw std::invalid_argument("the sorter requires start==0");
    }

    std::vector<edge_plus_time> temp;
    temp.reserve(static_cast<std::size_t>(sorter->tables->edges.num_rows));

    auto edges = &sorter->tables->edges;
    auto nodes = &sorter->tables->nodes;

    for (tsk_size_t i = 0; i < sorter->tables->edges.num_rows; ++i) {
        temp.push_back(edge_plus_time{ nodes->time[edges->parent[i]], edges->parent[i],
            edges->child[i], edges->left[i], edges->right[i] });
    }

    std::sort(begin(temp), end(temp),
        [](const edge_plus_time &lhs, const edge_plus_time &rhs) {
            if (lhs.time == rhs.time) {
                if (lhs.parent == rhs.parent) {
                    if (lhs.child == rhs.child) {
                        return lhs.left < rhs.left;
                    }
                    return lhs.child < rhs.child;
                }
                return lhs.parent < rhs.parent;
            }
            return lhs.time < rhs.time;
        });

    for (std::size_t i = 0; i < temp.size(); ++i) {
        edges->left[i] = temp[i].left;
        edges->right[i] = temp[i].right;
        edges->parent[i] = temp[i].parent;
        edges->child[i] = temp[i].child;
    }

    return 0;
}

int
main(int argc, char **argv)
{
    if (argc != 3) {
        std::cerr << "Usage: " << argv[0] << " input.arg output.arg\n";
        std::exit(0);
    }
    const char *infile = argv[1];
    const char *outfile = argv[2];

    tsk_table_collection_t tables;
    auto ret = tsk_table_collection_load(&tables, infile, 0);
    handle_tskit_return_code(ret);

    tsk_table_sorter_t sorter;
    ret = tsk_table_sorter_init(&sorter, &tables, 0);
    handle_tskit_return_code(ret);
    sorter.sort_edges = sort_edges;
    try {
        ret = tsk_table_sorter_run(&sorter, NULL);
    } catch (std::exception &e) {
        std::cerr << e.what() << '\n';
        std::exit(1);
    }
    handle_tskit_return_code(ret);
    ret = tsk_table_collection_dump(&tables, outfile, 0);
    handle_tskit_return_code(ret);
    ret = tsk_table_collection_free(&tables);
    handle_tskit_return_code(ret);
}