General function for reordering nodes in a tree sequence

[hyanwong]

There have been a number of occasions when I have wanted to reorder the nodes in a tree sequence, for instance, to have the oldest nodes first (this e.g. allows younger nodes to be culled by truncating the table, implementing delete_younger. It's possible to do that using ts.subset, but that can also have unintended consequences for individual IDs etc. Here's a nice function that simply reorders the nodes. In the long term it could be nice to have this in core, but I'm putting it here until that time:

def reorder_nodes(tables, order):
    # Reorder the nodes in a table collection, but keep everything else identical
    # The passed in `order` is simply an array of the same length as the number of nodes
    # which can have repeated values. Each node will be ordered based on the value
    # (the node with the lowest value will have the lowest ID: order in the case of ties is preserved)
    node_map = np.argsort(order, kind='stable').astype(tables.mutations.node.dtype)
    inverse_node_map = np.empty_like(node_map)
    inverse_node_map[node_map] = np.arange(len(node_map), dtype=node_map.dtype)
    tables.nodes.replace_with(tables.nodes[node_map])
    tables.edges.parent = inverse_node_map[tables.edges.parent]
    tables.edges.child = inverse_node_map[tables.edges.child]
    tables.mutations.node = inverse_node_map[tables.mutations.node]
    tables.migrations.node = inverse_node_map[tables.migrations.node]

def reorder_ts_nodes(ts, order):
    # equivalent function for a tree sequence
    tables = ts.dump_tables()
    reorder_nodes(tables, order)
    tables.sort()
    return tables.tree_sequence()
    


# example usage: reorder in descending time order (oldest nodes first)
reordered_ts = reorder_ts_nodes(ts, order=-ts.nodes_time)